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

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):

34
sdk/python/test/t_session.py
View File

@ -7,52 +7,44 @@ class TestSession:
def test_create_session(self):
rag = RAGFlow(API_KEY, HOST_ADDRESS)
kb = rag.create_dataset(name="test_create_session")
assistant = rag.create_assistant(name="test_create_session", knowledgebases=[kb])
assistant = rag.create_chat(name="test_create_session", knowledgebases=[kb])
session = assistant.create_session()
assert isinstance(session,Session), "Failed to create a session."
def test_create_chat_with_success(self):
rag = RAGFlow(API_KEY, HOST_ADDRESS)
kb = rag.create_dataset(name="test_create_chat")
assistant = rag.create_assistant(name="test_create_chat", knowledgebases=[kb])
assistant = rag.create_chat(name="test_create_chat", knowledgebases=[kb])
session = assistant.create_session()
question = "What is AI"
for ans in session.chat(question, stream=True):
for ans in session.ask(question, stream=True):
pass
assert not ans.content.startswith("**ERROR**"), "Please check this error."
def test_delete_session_with_success(self):
def test_delete_sessions_with_success(self):
rag = RAGFlow(API_KEY, HOST_ADDRESS)
kb = rag.create_dataset(name="test_delete_session")
assistant = rag.create_assistant(name="test_delete_session",knowledgebases=[kb])
assistant = rag.create_chat(name="test_delete_session",knowledgebases=[kb])
session=assistant.create_session()
res=session.delete()
assert res, "Failed to delete the dataset."
res=assistant.delete_sessions(ids=[session.id])
assert res is None, "Failed to delete the dataset."
def test_update_session_with_success(self):
rag=RAGFlow(API_KEY,HOST_ADDRESS)
kb=rag.create_dataset(name="test_update_session")
assistant = rag.create_assistant(name="test_update_session",knowledgebases=[kb])
assistant = rag.create_chat(name="test_update_session",knowledgebases=[kb])
session=assistant.create_session(name="old session")
session.name="new session"
res=session.save()
assert res,"Failed to update the session"
res=session.update({"name":"new session"})
assert res is None,"Failed to update the session"
def test_get_session_with_success(self):
rag=RAGFlow(API_KEY,HOST_ADDRESS)
kb=rag.create_dataset(name="test_get_session")
assistant = rag.create_assistant(name="test_get_session",knowledgebases=[kb])
session = assistant.create_session()
session_2= assistant.get_session(id=session.id)
assert session.to_json()==session_2.to_json(),"Failed to get the session"
def test_list_session_with_success(self):
def test_list_sessions_with_success(self):
rag=RAGFlow(API_KEY,HOST_ADDRESS)
kb=rag.create_dataset(name="test_list_session")
assistant=rag.create_assistant(name="test_list_session",knowledgebases=[kb])
assistant=rag.create_chat(name="test_list_session",knowledgebases=[kb])
assistant.create_session("test_1")
assistant.create_session("test_2")
sessions=assistant.list_session()
sessions=assistant.list_sessions()
if isinstance(sessions,list):
for session in sessions:
assert isinstance(session,Session),"Non-Session elements exist in the list"