dale_yxc/ragflow_python
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Refactor API for document and session (#2819)

Browse Source 
### What problem does this PR solve?

Refactor API for document and session.

### Type of change


- [x] Refactoring

---------

Co-authored-by: liuhua <10215101452@stu.ecun.edu.cn>
This commit is contained in:
liuhua
2024-10-12 19:35:19 +08:00
committed by GitHub
parent 7d80fc474c
commit 6eed115723
9 changed files with 821 additions and 745 deletions
Download Patch File Download Diff File Expand all files Collapse all files

24
sdk/python/ragflow/modules/chat.py
View File

@@ -51,28 +51,28 @@ class Chat(Base):
def create_session(self, name: str = "New session") -> Session:
res = self.post("/session/save", {"name": name, "assistant_id": self.id})
res = self.post(f"/chat/{self.id}/session", {"name": name})
res = res.json()
if res.get("retmsg") == "success":
if res.get("code") == 0:
return Session(self.rag, res['data'])
raise Exception(res["retmsg"])
raise Exception(res["message"])
def list_session(self) -> List[Session]:
res = self.get('/session/list', {"assistant_id": self.id})
def list_sessions(self,page: int = 1, page_size: int = 1024, orderby: str = "create_time", desc: bool = True,
id: str = None, name: str = None) -> List[Session]:
res = self.get(f'/chat/{self.id}/session',{"page": page, "page_size": page_size, "orderby": orderby, "desc": desc, "id": id, "name": name} )
res = res.json()
if res.get("retmsg") == "success":
if res.get("code") == 0:
result_list = []
for data in res["data"]:
result_list.append(Session(self.rag, data))
return result_list
raise Exception(res["retmsg"])
raise Exception(res["message"])
def get_session(self, id) -> Session:
res = self.get("/session/get", {"id": id,"assistant_id":self.id})
def delete_sessions(self,ids):
res = self.rm(f"/chat/{self.id}/session", {"ids": ids})
res = res.json()
if res.get("retmsg") == "success":
return Session(self.rag, res["data"])
raise Exception(res["retmsg"])
if res.get("code") != 0:
raise Exception(res.get("message"))
def get_prologue(self):
return self.prompt.opener

27
sdk/python/ragflow/modules/session.py
View File

@@ -8,20 +8,20 @@ class Session(Base):
self.id = None
self.name = "New session"
self.messages = [{"role": "assistant", "content": "Hi! I am your assistant，can I help you?"}]
self.assistant_id = None
self.chat_id = None
super().__init__(rag, res_dict)
def chat(self, question: str, stream: bool = False):
def ask(self, question: str, stream: bool = False):
for message in self.messages:
if "reference" in message:
message.pop("reference")
res = self.post("/session/completion",
{"session_id": self.id, "question": question, "stream": True}, stream=stream)
res = self.post(f"/chat/{self.chat_id}/session/{self.id}/completion",
{"question": question, "stream": True}, 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["retmsg"])
raise Exception(json_data["message"])
if line.startswith("data:"):
json_data = json.loads(line[5:])
if json_data["data"] != True:
@@ -52,19 +52,12 @@ class Session(Base):
message = Message(self.rag, temp_dict)
yield message
def save(self):
res = self.post("/session/save",
{"id": self.id, "assistant_id": self.assistant_id, "name": self.name})
def update(self,update_message):
res = self.put(f"/chat/{self.chat_id}/session/{self.id}",
update_message)
res = res.json()
if res.get("retmsg") == "success": return True
raise Exception(res.get("retmsg"))
def delete(self):
res = self.rm("/session/delete", {"id": self.id})
res = res.json()
if res.get("retmsg") == "success": return True
raise Exception(res.get("retmsg"))
if res.get("code") != 0:
raise Exception(res.get("message"))
class Message(Base):
def __init__(self, rag, res_dict):