Update the component of the agent API with parameters. (#4131)

### What problem does this PR solve?

Update the  component of the agent API with parameters.

### Type of change

- [x] Bug Fix (non-breaking change which fixes an issue)
- [x] New Feature (non-breaking change which adds functionality)
- [x] Refactoring

---------

Co-authored-by: liuhua <10215101452@stu.ecun.edu.cn>
Co-authored-by: Kevin Hu <kevinhu.sh@gmail.com>
Co-authored-by: writinwaters <93570324+writinwaters@users.noreply.github.com>
This commit is contained in:
liuhua
2024-12-20 17:34:16 +08:00
committed by GitHub
parent a0dc9e1bdf
commit 35580af875
8 changed files with 152 additions and 185 deletions

View File

@ -1,7 +1,6 @@
from .base import Base
from .session import Session,Message
from .session import Session
import requests
import json
class Agent(Base):
@ -52,8 +51,8 @@ class Agent(Base):
super().__init__(rag,res_dict)
@staticmethod
def create_session(id,rag) -> Session:
res = requests.post(f"{rag.api_url}/agents/{id}/sessions",headers={"Authorization": f"Bearer {rag.user_key}"},json={})
def create_session(id,rag,**kwargs) -> Session:
res = requests.post(f"{rag.api_url}/agents/{id}/sessions",headers={"Authorization": f"Bearer {rag.user_key}"},json=kwargs)
res = res.json()
if res.get("code") == 0:
return Session(rag,res.get("data"))
@ -74,30 +73,3 @@ 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