Fix bugs in agent api and update api document (#3996)

### What problem does this PR solve?

Fix bugs in agent api and update api document

### Type of change

- [x] Bug Fix (non-breaking change which fixes an issue)
- [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-13 10:25:52 +08:00
committed by GitHub
parent 68d46b2a1e
commit 1ecb687c51
9 changed files with 351 additions and 57 deletions

View File

@ -1,7 +1,9 @@
from .base import Base
from .session import Session
from .session import Session,Message
import requests
from typing import List
import json
class Agent(Base):
def __init__(self,rag,res_dict):
@ -73,3 +75,30 @@ class Agent(Base):
result_list.append(temp_agent)
return result_list
raise Exception(res.get("message"))
@staticmethod
def ask(agent_id,rag,stream=True,**kwargs):
url = f"{rag.api_url}/agents/{agent_id}/completions"
headers = {"Authorization": f"Bearer {rag.user_key}"}
res = requests.post(url=url, headers=headers, json=kwargs,stream=stream)
for line in res.iter_lines():
line = line.decode("utf-8")
if line.startswith("{"):
json_data = json.loads(line)
raise Exception(json_data["message"])
if line.startswith("data:"):
json_data = json.loads(line[5:])
if json_data["data"] is not True:
if json_data["data"].get("running_status"):
continue
answer = json_data["data"]["answer"]
reference = json_data["data"]["reference"]
temp_dict = {
"content": answer,
"role": "assistant"
}
if "chunks" in reference:
chunks = reference["chunks"]
temp_dict["reference"] = chunks
message = Message(rag, temp_dict)
yield message

View File

@ -29,7 +29,7 @@ class Session(Base):
raise Exception(json_data["message"])
if line.startswith("data:"):
json_data = json.loads(line[5:])
if not json_data["data"]:
if json_data["data"] is not True:
answer = json_data["data"]["answer"]
reference = json_data["data"]["reference"]
temp_dict = {

View File

@ -1,4 +1,4 @@
from ragflow_sdk import RAGFlow
from ragflow_sdk import RAGFlow,Agent
from common import HOST_ADDRESS
import pytest
@ -6,4 +6,14 @@ import pytest
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()
rag.list_agents()
@pytest.mark.skip(reason="")
def test_converse_with_agent_with_success(get_api_key_fixture):
API_KEY = "ragflow-BkOGNhYjIyN2JiODExZWY5MzVhMDI0Mm"
agent_id = "ebfada2eb2bc11ef968a0242ac120006"
rag = RAGFlow(API_KEY,HOST_ADDRESS)
lang = "Chinese"
file = "How is the weather tomorrow?"
Agent.ask(agent_id=agent_id,rag=rag,lang=lang,file=file)