SDK for session (#2312)

### What problem does this PR solve?

SDK for session
#1102 

### Type of change

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

---------

Co-authored-by: Feiue <10215101452@stu.ecun.edu.cn>
Co-authored-by: Kevin Hu <kevinhu.sh@gmail.com>

sdk/python/ragflow/modules/chat_assistant.py

@@ -1,9 +1,12 @@
+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.id = ""
         self.name = "assistant"
         self.avatar = "path/to/avatar"
         self.knowledgebases = ["kb1"]
@@ -41,8 +44,8 @@ class Assistant(Base):
 
     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()
+                        {"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
@@ -54,3 +57,15 @@ class Assistant(Base):
         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

sdk/python/ragflow/modules/session.py

@@ -0,0 +1,64 @@
+import json
+
+from .base import Base
+
+
+class Session(Base):
+    def __init__(self, rag, res_dict):
+        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):
+        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:
+            if line.startswith("data:"):
+                json_data = json.loads(line[5:])
+                if json_data["data"] != True:
+                    answer = json_data["data"]["answer"]
+                    reference = json_data["data"]["reference"]
+                    temp_dict = {
+                        "content": answer,
+                        "role": "assistant",
+                        "reference": reference
+                    }
+                    message = Message(self.rag, temp_dict)
+                    message_list.append(message)
+        return message_list
+
+    def save(self):
+        res = self.post("/session/save",
+                        {"id": self.id, "dialog_id": self.assistant_id, "name": self.name, "message": self.messages})
+        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.role = "assistant"
+        self.prompt=None
+        super().__init__(rag, res_dict)
+
+
+class Chunk(Base):
+    def __init__(self, rag, res_dict):
+        self.id = None
+        self.content = None
+        self.document_id = None
+        self.document_name = None
+        self.knowledgebase_id = None
+        self.image_id = None
+        self.similarity = None
+        self.vector_similarity = None
+        self.term_similarity = None
+        self.positions = None
+        super().__init__(rag, res_dict)

sdk/python/ragflow/ragflow.py

@@ -17,7 +17,6 @@ from typing import List
 
 import requests
 
-
 from .modules.chat_assistant import Assistant
 from .modules.dataset import DataSet
 
@@ -88,7 +87,7 @@ class RAGFlow:
             datasets.append(dataset.to_json())
 
         if llm is None:
-            llm = Assistant.LLM(self, {"model_name": "deepseek-chat",
+            llm = Assistant.LLM(self, {"model_name": None,
                                        "temperature": 0.1,
                                        "top_p": 0.3,
                                        "presence_penalty": 0.4,
@@ -142,4 +141,4 @@ class RAGFlow:
             for data in res['data']:
                 result_list.append(Assistant(self, data))
             return result_list
-        raise Exception(res["retmsg"])
+        raise Exception(res["retmsg"])

sdk/python/test/t_assistant.py

@@ -10,10 +10,10 @@ class TestAssistant(TestSdk):
         Test creating an assistant with success
         """
         rag = RAGFlow(API_KEY, HOST_ADDRESS)
-        kb = rag.get_dataset(name="God")
-        assistant = rag.create_assistant("God",knowledgebases=[kb])
+        kb = rag.create_dataset(name="test_create_assistant")
+        assistant = rag.create_assistant("test_create", knowledgebases=[kb])
         if isinstance(assistant, Assistant):
-            assert assistant.name == "God", "Name does not match."
+            assert assistant.name == "test_create", "Name does not match."
         else:
             assert False, f"Failed to create assistant, error: {assistant}"
 
@@ -22,11 +22,11 @@ class TestAssistant(TestSdk):
         Test updating an assistant with success.
         """
         rag = RAGFlow(API_KEY, HOST_ADDRESS)
-        kb = rag.get_dataset(name="God")
-        assistant = rag.create_assistant("ABC",knowledgebases=[kb])
+        kb = rag.create_dataset(name="test_update_assistant")
+        assistant = rag.create_assistant("test_update", knowledgebases=[kb])
         if isinstance(assistant, Assistant):
-            assert assistant.name == "ABC", "Name does not match."
-            assistant.name = 'DEF'
+            assert assistant.name == "test_update", "Name does not match."
+            assistant.name = 'new_assistant'
             res = assistant.save()
             assert res is True, f"Failed to update assistant, error: {res}"
         else:
@@ -37,10 +37,10 @@ class TestAssistant(TestSdk):
         Test deleting an assistant with success
         """
         rag = RAGFlow(API_KEY, HOST_ADDRESS)
-        kb = rag.get_dataset(name="God")
-        assistant = rag.create_assistant("MA",knowledgebases=[kb])
+        kb = rag.create_dataset(name="test_delete_assistant")
+        assistant = rag.create_assistant("test_delete", knowledgebases=[kb])
         if isinstance(assistant, Assistant):
-            assert assistant.name == "MA", "Name does not match."
+            assert assistant.name == "test_delete", "Name does not match."
             res = assistant.delete()
             assert res is True, f"Failed to delete assistant, error: {res}"
         else:
@@ -61,6 +61,8 @@ class TestAssistant(TestSdk):
         Test getting an assistant's detail with success
         """
         rag = RAGFlow(API_KEY, HOST_ADDRESS)
-        assistant = rag.get_assistant(name="God")
+        kb = rag.create_dataset(name="test_get_assistant")
+        rag.create_assistant("test_get_assistant", knowledgebases=[kb])
+        assistant = rag.get_assistant(name="test_get_assistant")
         assert isinstance(assistant, Assistant), f"Failed to get assistant, error: {assistant}."
-        assert assistant.name == "God", "Name does not match"
+        assert assistant.name == "test_get_assistant", "Name does not match"

sdk/python/test/t_session.py

@@ -0,0 +1,27 @@
+from ragflow import RAGFlow
+
+from common import API_KEY, HOST_ADDRESS
+
+
+class TestChatSession:
+    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."
+
+    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."