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
slug: /python_api_reference
---
@ -1513,4 +1513,123 @@ while True:
for ans in session.ask(question, stream=True):
print(ans.content[len(cont):], end='', flush=True)
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 .session import Session
import requests
from typing import List
class Agent(Base):
def __init__(self,rag,res_dict):
@ -57,3 +58,18 @@ class Agent(Base):
return Session(rag,res.get("data"))
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.name = "New session"
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():
if key =="chat_id" and value is not None:
self.chat_id = None
self.__session_type = "chat"
if key == "agent_id" and value is not None:
self.agent_id = None
self.__session_type = "agent"
super().__init__(rag, res_dict)

View File

@ -18,6 +18,7 @@ import requests
from .modules.chat import Chat
from .modules.chunk import Chunk
from .modules.dataset import DataSet
from .modules.agent import Agent
class RAGFlow:
@ -86,8 +87,10 @@ class RAGFlow:
return result_list
raise Exception(res["message"])
def create_chat(self, name: str, avatar: str = "", dataset_ids: list[str] = [],
llm: Chat.LLM | None = None, prompt: Chat.Prompt | None = None) -> Chat:
def create_chat(self, name: str, avatar: str = "", dataset_ids=None,
llm: Chat.LLM | None = None, prompt: Chat.Prompt | None = None) -> Chat:
if dataset_ids is None:
dataset_ids = []
dataset_list = []
for id in dataset_ids:
dataset_list.append(id)
@ -176,3 +179,15 @@ class RAGFlow:
chunks.append(chunk)
return chunks
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):
API_KEY = get_api_key_fixture
rag = RAGFlow(API_KEY, HOST_ADDRESS)
rag.create_dataset("test_list_datasets")
rag.list_datasets()

View File

@ -107,4 +107,11 @@ def test_create_agent_conversation_with_success(get_api_key_fixture):
API_KEY = "ragflow-BkOGNhYjIyN2JiODExZWY5MzVhMDI0Mm"
rag = RAGFlow(API_KEY,HOST_ADDRESS)
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)