SDK for session (#2354)

### What problem does this PR solve?

Includes SDK for creating, updating sessions, getting sessions, listing
sessions, and dialogues
#1102 
### Type of change

- [x] New Feature (non-breaking change which adds functionality)

---------

Co-authored-by: liuhua <10215101452@stu.ecun.edu.cn>

sdk/python/ragflow/__init__.py

@@ -4,4 +4,5 @@ __version__ = importlib.metadata.version("ragflow")
 
 from .ragflow import RAGFlow
 from .modules.dataset import DataSet
-from .modules.chat_assistant import Assistant
+from .modules.assistant import Assistant
+from .modules.session import Session

sdk/python/ragflow/modules/assistant.py

@@ -1,71 +1,86 @@
-from typing import List
-
-from .base import Base
-from .session import Session, Message
-
-
-class Assistant(Base):
-    def __init__(self, rag, res_dict):
-        self.id = ""
-        self.name = "assistant"
-        self.avatar = "path/to/avatar"
-        self.knowledgebases = ["kb1"]
-        self.llm = Assistant.LLM(rag, {})
-        self.prompt = Assistant.Prompt(rag, {})
-        super().__init__(rag, res_dict)
-
-    class LLM(Base):
-        def __init__(self, rag, res_dict):
-            self.model_name = "deepseek-chat"
-            self.temperature = 0.1
-            self.top_p = 0.3
-            self.presence_penalty = 0.4
-            self.frequency_penalty = 0.7
-            self.max_tokens = 512
-            super().__init__(rag, res_dict)
-
-    class Prompt(Base):
-        def __init__(self, rag, res_dict):
-            self.similarity_threshold = 0.2
-            self.keywords_similarity_weight = 0.7
-            self.top_n = 8
-            self.variables = [{"key": "knowledge", "optional": True}]
-            self.rerank_model = None
-            self.empty_response = None
-            self.opener = "Hi! I'm your assistant, what can I do for you?"
-            self.show_quote = True
-            self.prompt = (
-                "You are an intelligent assistant. Please summarize the content of the knowledge base to answer the question. "
-                "Please list the data in the knowledge base and answer in detail. When all knowledge base content is irrelevant to the question, "
-                "your answer must include the sentence 'The answer you are looking for is not found in the knowledge base!' "
-                "Answers need to consider chat history.\nHere is the knowledge base:\n{knowledge}\nThe above is the knowledge base."
-            )
-            super().__init__(rag, res_dict)
-
-    def save(self) -> bool:
-        res = self.post('/assistant/save',
-                        {"id": self.id, "name": self.name, "avatar": self.avatar, "knowledgebases": self.knowledgebases,
-                         "llm": self.llm.to_json(), "prompt": self.prompt.to_json()
-                         })
-        res = res.json()
-        if res.get("retmsg") == "success": return True
-        raise Exception(res["retmsg"])
-
-    def delete(self) -> bool:
-        res = self.rm('/assistant/delete',
-                      {"id": self.id})
-        res = res.json()
-        if res.get("retmsg") == "success": return True
-        raise Exception(res["retmsg"])
-
-    def create_session(self, name: str = "New session", messages: List[Message] = [
-        {"role": "assistant", "reference": [],
-         "content": "您好，我是您的助手小樱，长得可爱又善良，can I help you?"}]) -> Session:
-        res = self.post("/session/save", {"name": name, "messages": messages, "assistant_id": self.id, })
-        res = res.json()
-        if res.get("retmsg") == "success":
-            return Session(self.rag, res['data'])
-        raise Exception(res["retmsg"])
-
-    def get_prologue(self):
-        return self.prompt.opener
+from typing import List
+
+from .base import Base
+from .session import Session
+
+
+class Assistant(Base):
+    def __init__(self, rag, res_dict):
+        self.id = ""
+        self.name = "assistant"
+        self.avatar = "path/to/avatar"
+        self.knowledgebases = ["kb1"]
+        self.llm = Assistant.LLM(rag, {})
+        self.prompt = Assistant.Prompt(rag, {})
+        super().__init__(rag, res_dict)
+
+    class LLM(Base):
+        def __init__(self, rag, res_dict):
+            self.model_name = "deepseek-chat"
+            self.temperature = 0.1
+            self.top_p = 0.3
+            self.presence_penalty = 0.4
+            self.frequency_penalty = 0.7
+            self.max_tokens = 512
+            super().__init__(rag, res_dict)
+
+    class Prompt(Base):
+        def __init__(self, rag, res_dict):
+            self.similarity_threshold = 0.2
+            self.keywords_similarity_weight = 0.7
+            self.top_n = 8
+            self.variables = [{"key": "knowledge", "optional": True}]
+            self.rerank_model = None
+            self.empty_response = None
+            self.opener = "Hi! I'm your assistant, what can I do for you?"
+            self.show_quote = True
+            self.prompt = (
+                "You are an intelligent assistant. Please summarize the content of the knowledge base to answer the question. "
+                "Please list the data in the knowledge base and answer in detail. When all knowledge base content is irrelevant to the question, "
+                "your answer must include the sentence 'The answer you are looking for is not found in the knowledge base!' "
+                "Answers need to consider chat history.\nHere is the knowledge base:\n{knowledge}\nThe above is the knowledge base."
+            )
+            super().__init__(rag, res_dict)
+
+    def save(self) -> bool:
+        res = self.post('/assistant/save',
+                        {"id": self.id, "name": self.name, "avatar": self.avatar, "knowledgebases": self.knowledgebases,
+                         "llm": self.llm.to_json(), "prompt": self.prompt.to_json()
+                         })
+        res = res.json()
+        if res.get("retmsg") == "success": return True
+        raise Exception(res["retmsg"])
+
+    def delete(self) -> bool:
+        res = self.rm('/assistant/delete',
+                      {"id": self.id})
+        res = res.json()
+        if res.get("retmsg") == "success": return True
+        raise Exception(res["retmsg"])
+
+    def create_session(self, name: str = "New session") -> Session:
+        res = self.post("/session/save", {"name": name, "assistant_id": self.id})
+        res = res.json()
+        if res.get("retmsg") == "success":
+            return Session(self.rag, res['data'])
+        raise Exception(res["retmsg"])
+
+    def list_session(self) -> List[Session]:
+        res = self.get('/session/list', {"assistant_id": self.id})
+        res = res.json()
+        if res.get("retmsg") == "success":
+            result_list = []
+            for data in res["data"]:
+                result_list.append(Session(self.rag, data))
+            return result_list
+        raise Exception(res["retmsg"])
+
+    def get_session(self, id) -> Session:
+        res = self.get("/session/get", {"id": id})
+        res = res.json()
+        if res.get("retmsg") == "success":
+            return Session(self.rag, res["data"])
+        raise Exception(res["retmsg"])
+
+    def get_prologue(self):
+        return self.prompt.opener

sdk/python/ragflow/modules/base.py

@@ -18,8 +18,8 @@ class Base(object):
                     pr[name] = value
         return pr
 
-    def post(self, path, param):
-        res = self.rag.post(path, param)
+    def post(self, path, param, stream=False):
+        res = self.rag.post(path, param, stream=stream)
         return res
 
     def get(self, path, params):

sdk/python/ragflow/modules/session.py

@@ -8,17 +8,17 @@ 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
         super().__init__(rag, res_dict)
 
     def chat(self, question: str, stream: bool = False):
+        for message in self.messages:
+            if "reference" in message:
+                message.pop("reference")
         res = self.post("/session/completion",
-                        {"id": self.id, "question": question, "stream": stream, "messages": self.messages})
-        res = res.text
-        response_lines = res.splitlines()
-        message_list = []
-        for line in response_lines:
+                        {"id": self.id, "question": question, "stream": stream}, stream=True)
+        for line in res.iter_lines():
+            line = line.decode("utf-8")
             if line.startswith("data:"):
                 json_data = json.loads(line[5:])
                 if json_data["data"] != True:
@@ -26,26 +26,49 @@ class Session(Base):
                     reference = json_data["data"]["reference"]
                     temp_dict = {
                         "content": answer,
-                        "role": "assistant",
-                        "reference": reference
+                        "role": "assistant"
                     }
+                    if "chunks" in reference:
+                        chunks = reference["chunks"]
+                        chunk_list = []
+                        for chunk in chunks:
+                            new_chunk = {
+                                "id": chunk["chunk_id"],
+                                "content": chunk["content_with_weight"],
+                                "document_id": chunk["doc_id"],
+                                "document_name": chunk["docnm_kwd"],
+                                "knowledgebase_id": chunk["kb_id"],
+                                "image_id": chunk["img_id"],
+                                "similarity": chunk["similarity"],
+                                "vector_similarity": chunk["vector_similarity"],
+                                "term_similarity": chunk["term_similarity"],
+                                "positions": chunk["positions"],
+                            }
+                            chunk_list.append(new_chunk)
+                        temp_dict["reference"] = chunk_list
                     message = Message(self.rag, temp_dict)
-                    message_list.append(message)
-        return message_list
+                    yield message
 
     def save(self):
         res = self.post("/session/save",
-                        {"id": self.id, "dialog_id": self.assistant_id, "name": self.name, "message": self.messages})
+                        {"id": self.id, "assistant_id": self.assistant_id, "name": self.name})
         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"))
+
+
 class Message(Base):
     def __init__(self, rag, res_dict):
-        self.content = "您好，我是您的助手小樱，长得可爱又善良，can I help you?"
-        self.reference = []
+        self.content = "Hi! I am your assistant，can I help you?"
+        self.reference = None
         self.role = "assistant"
-        self.prompt=None
+        self.prompt = None
         super().__init__(rag, res_dict)
 
 

sdk/python/ragflow/ragflow.py

@@ -17,7 +17,7 @@ from typing import List
 
 import requests
 
-from .modules.chat_assistant import Assistant
+from .modules.assistant import Assistant
 from .modules.dataset import DataSet
 
 
@@ -30,8 +30,8 @@ class RAGFlow:
         self.api_url = f"{base_url}/api/{version}"
         self.authorization_header = {"Authorization": "{} {}".format("Bearer", self.user_key)}
 
-    def post(self, path, param):
-        res = requests.post(url=self.api_url + path, json=param, headers=self.authorization_header)
+    def post(self, path, param, stream=False):
+        res = requests.post(url=self.api_url + path, json=param, headers=self.authorization_header, stream=stream)
         return res
 
     def get(self, path, params=None):

sdk/python/test/t_session.py

@@ -1,27 +1,60 @@
-from ragflow import RAGFlow
+from ragflow import RAGFlow,Session
 
 from common import API_KEY, HOST_ADDRESS
 
 
-class TestChatSession:
+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])
         session = assistant.create_session()
-        assert assistant is not None, "Failed to get the assistant."
-        assert session is not None, "Failed to create a 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])
         session = assistant.create_session()
-        assert session is not None, "Failed to create a session."
-        prologue = assistant.get_prologue()
-        assert isinstance(prologue, str), "Prologue is not a string."
-        assert len(prologue) > 0, "Prologue is empty."
         question = "What is AI"
-        ans = session.chat(question, stream=True)
-        response = ans[-1].content
-        assert len(response) > 0, "Assistant did not return any response."
+        for ans in session.chat(question, stream=True):
+            pass
+        assert ans.content!="\n**ERROR**", "Please check this error."
+
+    def test_delete_session_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])
+        session=assistant.create_session()
+        res=session.delete()
+        assert res, "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])
+        session=assistant.create_session(name="old session")
+        session.name="new session"
+        res=session.save()
+        assert res,"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):
+        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.create_session("test_1")
+        assistant.create_session("test_2")
+        sessions=assistant.list_session()
+        if isinstance(sessions,list):
+            for session in sessions:
+                assert isinstance(session,Session),"Non-Session elements exist in the list"
+        else :
+            assert False,"Failed to retrieve the session list."