Add sdk for list agents and sessions (#3848)

### What problem does this PR solve?

Add sdk for list agents and sessions

### Type of change

- [x] New Feature (non-breaking change which adds functionality)

Co-authored-by: liuhua <10215101452@stu.ecun.edu.cn>
This commit is contained in:
liuhua 2024-12-04 16:23:22 +08:00 committed by GitHub
parent 289f4f1916
commit 1b589609a4
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
8 changed files with 1332 additions and 7 deletions

File diff suppressed because it is too large Load Diff

View File

@ -1,4 +1,4 @@
--- from sdk.python.ragflow_sdk import Agentfrom docutils.nodes import titlefrom sdk.python.ragflow_sdk import Agent---
sidebar_position: 2 sidebar_position: 2
slug: /python_api_reference slug: /python_api_reference
--- ---
@ -1514,3 +1514,122 @@ while True:
print(ans.content[len(cont):], end='', flush=True) print(ans.content[len(cont):], end='', flush=True)
cont = ans.content cont = ans.content
``` ```
---
## List agent sessions
```python
Agent.list_sessions(
agent_id,
rag
page: int = 1,
page_size: int = 30,
orderby: str = "update_time",
desc: bool = True,
id: str = None
) -> List[Session]
```
Lists sessions associated with the current agent.
### Parameters
#### page: `int`
Specifies the page on which the sessions will be displayed. Defaults to `1`.
#### page_size: `int`
The number of sessions on each page. Defaults to `30`.
#### orderby: `str`
The field by which sessions should be sorted. Available options:
- `"create_time"`
- `"update_time"`(default)
#### desc: `bool`
Indicates whether the retrieved sessions should be sorted in descending order. Defaults to `True`.
#### id: `str`
The ID of the agent session to retrieve. Defaults to `None`.
### Returns
- Success: A list of `Session` objects associated with the current agent.
- Failure: `Exception`.
### Examples
```python
from ragflow_sdk import RAGFlow
rag_object = RAGFlow(api_key="<YOUR_API_KEY>", base_url="http://<YOUR_BASE_URL>:9380")
agent_id = "2710f2269b4611ef8fdf0242ac120006"
sessions=Agent.list_sessions(agent_id,rag_object)
for session in sessions:
print(session)
```
---
## List agents
```python
RAGFlow.list_agents(
page: int = 1,
page_size: int = 30,
orderby: str = "create_time",
desc: bool = True,
id: str = None,
title: str = None
) -> List[Agent]
```
Lists agents.
### Parameters
#### page: `int`
Specifies the page on which the agents will be displayed. Defaults to `1`.
#### page_size: `int`
The number of agents on each page. Defaults to `30`.
#### orderby: `str`
The attribute by which the results are sorted. Available options:
- `"create_time"` (default)
- `"update_time"`
#### desc: `bool`
Indicates whether the retrieved agents should be sorted in descending order. Defaults to `True`.
#### id: `str`
The ID of the agent to retrieve. Defaults to `None`.
#### name: `str`
The name of the agent to retrieve. Defaults to `None`.
### Returns
- Success: A list of `Agent` objects.
- Failure: `Exception`.
### Examples
```python
from ragflow_sdk import RAGFlow
rag_object = RAGFlow(api_key="<YOUR_API_KEY>", base_url="http://<YOUR_BASE_URL>:9380")
for agent in rag_object.list_agents():
print(agent)
```

View File

@ -1,6 +1,7 @@
from .base import Base from .base import Base
from .session import Session from .session import Session
import requests import requests
from typing import List
class Agent(Base): class Agent(Base):
def __init__(self,rag,res_dict): def __init__(self,rag,res_dict):
@ -57,3 +58,18 @@ class Agent(Base):
return Session(rag,res.get("data")) return Session(rag,res.get("data"))
raise Exception(res.get("message")) raise Exception(res.get("message"))
@staticmethod
def list_sessions(agent_id,rag,page: int = 1, page_size: int = 30, orderby: str = "create_time", desc: bool = True,
id: str = None) -> List[Session]:
url = f"{rag.api_url}/agents/{agent_id}/sessions"
headers = {"Authorization": f"Bearer {rag.user_key}"}
params = {"page": page, "page_size": page_size, "orderby": orderby, "desc": desc, "id": id}
res = requests.get(url=url,headers=headers,params=params)
res = res.json()
if res.get("code") == 0:
result_list = []
for data in res.get("data"):
temp_agent = Session(rag,data)
result_list.append(temp_agent)
return result_list
raise Exception(res.get("message"))

View File

@ -8,12 +8,12 @@ class Session(Base):
self.id = None self.id = None
self.name = "New session" self.name = "New session"
self.messages = [{"role": "assistant", "content": "Hi! I am your assistantcan I help you?"}] self.messages = [{"role": "assistant", "content": "Hi! I am your assistantcan I help you?"}]
self.chat_id = None
self.agent_id = None
for key,value in res_dict.items(): for key,value in res_dict.items():
if key =="chat_id" and value is not None: if key =="chat_id" and value is not None:
self.chat_id = None
self.__session_type = "chat" self.__session_type = "chat"
if key == "agent_id" and value is not None: if key == "agent_id" and value is not None:
self.agent_id = None
self.__session_type = "agent" self.__session_type = "agent"
super().__init__(rag, res_dict) super().__init__(rag, res_dict)

View File

@ -18,6 +18,7 @@ import requests
from .modules.chat import Chat from .modules.chat import Chat
from .modules.chunk import Chunk from .modules.chunk import Chunk
from .modules.dataset import DataSet from .modules.dataset import DataSet
from .modules.agent import Agent
class RAGFlow: class RAGFlow:
@ -86,8 +87,10 @@ class RAGFlow:
return result_list return result_list
raise Exception(res["message"]) raise Exception(res["message"])
def create_chat(self, name: str, avatar: str = "", dataset_ids: list[str] = [], def create_chat(self, name: str, avatar: str = "", dataset_ids=None,
llm: Chat.LLM | None = None, prompt: Chat.Prompt | None = None) -> Chat: llm: Chat.LLM | None = None, prompt: Chat.Prompt | None = None) -> Chat:
if dataset_ids is None:
dataset_ids = []
dataset_list = [] dataset_list = []
for id in dataset_ids: for id in dataset_ids:
dataset_list.append(id) dataset_list.append(id)
@ -176,3 +179,15 @@ class RAGFlow:
chunks.append(chunk) chunks.append(chunk)
return chunks return chunks
raise Exception(res.get("message")) raise Exception(res.get("message"))
def list_agents(self, page: int = 1, page_size: int = 30, orderby: str = "update_time", desc: bool = True,
id: str | None = None, title: str | None = None) -> list[Agent]:
res = self.get("/agents",{"page": page, "page_size": page_size, "orderby": orderby, "desc": desc, "id": id, "title": title})
res = res.json()
result_list = []
if res.get("code") == 0:
for data in res['data']:
result_list.append(Agent(self, data))
return result_list
raise Exception(res["message"])

View File

@ -0,0 +1,9 @@
from ragflow_sdk import RAGFlow,Agent
from common import HOST_ADDRESS
import pytest
@pytest.mark.skip(reason="")
def test_list_agents_with_success(get_api_key_fixture):
API_KEY=get_api_key_fixture
rag = RAGFlow(API_KEY,HOST_ADDRESS)
rag.list_agents()

View File

@ -51,4 +51,5 @@ def test_delete_datasets_with_success(get_api_key_fixture):
def test_list_datasets_with_success(get_api_key_fixture): def test_list_datasets_with_success(get_api_key_fixture):
API_KEY = get_api_key_fixture API_KEY = get_api_key_fixture
rag = RAGFlow(API_KEY, HOST_ADDRESS) rag = RAGFlow(API_KEY, HOST_ADDRESS)
rag.create_dataset("test_list_datasets")
rag.list_datasets() rag.list_datasets()

View File

@ -108,3 +108,10 @@ def test_create_agent_conversation_with_success(get_api_key_fixture):
rag = RAGFlow(API_KEY,HOST_ADDRESS) rag = RAGFlow(API_KEY,HOST_ADDRESS)
session = Agent.create_session("2e45b5209c1011efa3e90242ac120006", rag) session = Agent.create_session("2e45b5209c1011efa3e90242ac120006", rag)
session.ask("What is this job") session.ask("What is this job")
@pytest.mark.skip(reason="")
def test_list_agent_sessions_with_success(get_api_key_fixture):
API_KEY = "ragflow-BkOGNhYjIyN2JiODExZWY5MzVhMDI0Mm"
agent_id = "2710f2269b4611ef8fdf0242ac120006"
rag = RAGFlow(API_KEY,HOST_ADDRESS)
Agent.list_sessions(agent_id,rag)