pengjunjie/ragflow
1
0
Fork 1
mirror of https://github.com/infiniflow/ragflow.git synced 2025-12-08 20:42:30 +08:00
Code Issues Packages Projects Releases Wiki Activity

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

Browse Source 
### 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
Download Patch File Download Diff File Expand all files Collapse all files

34
sdk/python/ragflow_sdk/modules/agent.py
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

38
sdk/python/ragflow_sdk/modules/session.py
View File

@ -17,7 +17,7 @@ class Session(Base):
self.__session_type = "agent"
super().__init__(rag, res_dict)
def ask(self, question,stream=True,**kwargs):
def ask(self, question="",stream=True,**kwargs):
if self.__session_type == "agent":
res=self._ask_agent(question,stream)
elif self.__session_type == "chat":
@ -27,23 +27,22 @@ class Session(Base):
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(self.rag, temp_dict)
yield message
if not line.startswith("data:"):
continue
json_data = json.loads(line[5:])
if json_data["data"] is True or json_data["data"].get("running_status"):
continue
answer = json_data["data"]["answer"]
reference = json_data["data"].get("reference", {})
temp_dict = {
"content": answer,
"role": "assistant"
}
if reference and "chunks" in reference:
chunks = reference["chunks"]
temp_dict["reference"] = chunks
message = Message(self.rag, temp_dict)
yield message
def _ask_chat(self, question: str, stream: bool,**kwargs):
json_data={"question": question, "stream": True,"session_id":self.id}
@ -51,6 +50,7 @@ class Session(Base):
res = self.post(f"/chats/{self.chat_id}/completions",
json_data, stream=stream)
return res
def _ask_agent(self,question:str,stream:bool):
res = self.post(f"/agents/{self.agent_id}/completions",
{"question": question, "stream": True,"session_id":self.id}, stream=stream)
@ -70,4 +70,4 @@ class Message(Base):
self.role = "assistant"
self.prompt = None
self.id = None
super().__init__(rag, res_dict)
super().__init__(rag, res_dict)