SDK for Assistant (#2266)

### What problem does this PR solve?

SDK for Assistant
#1102 

### Type of change

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

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

sdk/python/ragflow/__init__.py

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

sdk/python/ragflow/modules/chat_assistant.py

@@ -0,0 +1,56 @@
+from .base import Base
+
+
+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"])

sdk/python/ragflow/ragflow.py

@@ -17,6 +17,8 @@ from typing import List
 
 import requests
 
+
+from .modules.chat_assistant import Assistant
 from .modules.dataset import DataSet
 
 
@@ -78,3 +80,66 @@ class RAGFlow:
         if res.get("retmsg") == "success":
             return DataSet(self, res['data'])
         raise Exception(res["retmsg"])
+
+    def create_assistant(self, name: str = "assistant", avatar: str = "path", knowledgebases: List[DataSet] = [],
+                         llm: Assistant.LLM = None, prompt: Assistant.Prompt = None) -> Assistant:
+        datasets = []
+        for dataset in knowledgebases:
+            datasets.append(dataset.to_json())
+
+        if llm is None:
+            llm = Assistant.LLM(self, {"model_name": "deepseek-chat",
+                                       "temperature": 0.1,
+                                       "top_p": 0.3,
+                                       "presence_penalty": 0.4,
+                                       "frequency_penalty": 0.7,
+                                       "max_tokens": 512, })
+        if prompt is None:
+            prompt = Assistant.Prompt(self, {"similarity_threshold": 0.2,
+                                             "keywords_similarity_weight": 0.7,
+                                             "top_n": 8,
+                                             "variables": [{
+                                                 "key": "knowledge",
+                                                 "optional": True
+                                             }], "rerank_model": "",
+                                             "empty_response": None,
+                                             "opener": None,
+                                             "show_quote": True,
+                                             "prompt": None})
+            if prompt.opener is None:
+                prompt.opener = "Hi! I'm your assistant, what can I do for you?"
+            if prompt.prompt is None:
+                prompt.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."
+                )
+
+        temp_dict = {"name": name,
+                     "avatar": avatar,
+                     "knowledgebases": datasets,
+                     "llm": llm.to_json(),
+                     "prompt": prompt.to_json()}
+        res = self.post("/assistant/save", temp_dict)
+        res = res.json()
+        if res.get("retmsg") == "success":
+            return Assistant(self, res["data"])
+        raise Exception(res["retmsg"])
+
+    def get_assistant(self, id: str = None, name: str = None) -> Assistant:
+        res = self.get("/assistant/get", {"id": id, "name": name})
+        res = res.json()
+        if res.get("retmsg") == "success":
+            return Assistant(self, res['data'])
+        raise Exception(res["retmsg"])
+
+    def list_assistants(self) -> List[Assistant]:
+        res = self.get("/assistant/list")
+        res = res.json()
+        result_list = []
+        if res.get("retmsg") == "success":
+            for data in res['data']:
+                result_list.append(Assistant(self, data))
+            return result_list
+        raise Exception(res["retmsg"])

sdk/python/test/common.py

@@ -1,4 +1,4 @@
 
 
-API_KEY = 'ragflow-k0N2I1MzQwNjNhMzExZWY5ODg1MDI0Mm'
+API_KEY = 'ragflow-k0YzUxMGY4NjY5YTExZWY5MjI5MDI0Mm'
 HOST_ADDRESS = 'http://127.0.0.1:9380'

sdk/python/test/t_assistant.py

@@ -0,0 +1,66 @@
+from ragflow import RAGFlow, Assistant
+
+from common import API_KEY, HOST_ADDRESS
+from test_sdkbase import TestSdk
+
+
+class TestAssistant(TestSdk):
+    def test_create_assistant_with_success(self):
+        """
+        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])
+        if isinstance(assistant, Assistant):
+            assert assistant.name == "God", "Name does not match."
+        else:
+            assert False, f"Failed to create assistant, error: {assistant}"
+
+    def test_update_assistant_with_success(self):
+        """
+        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])
+        if isinstance(assistant, Assistant):
+            assert assistant.name == "ABC", "Name does not match."
+            assistant.name = 'DEF'
+            res = assistant.save()
+            assert res is True, f"Failed to update assistant, error: {res}"
+        else:
+            assert False, f"Failed to create assistant, error: {assistant}"
+
+    def test_delete_assistant_with_success(self):
+        """
+        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])
+        if isinstance(assistant, Assistant):
+            assert assistant.name == "MA", "Name does not match."
+            res = assistant.delete()
+            assert res is True, f"Failed to delete assistant, error: {res}"
+        else:
+            assert False, f"Failed to create assistant, error: {assistant}"
+
+    def test_list_assistants_with_success(self):
+        """
+        Test listing assistants with success
+        """
+        rag = RAGFlow(API_KEY, HOST_ADDRESS)
+        list_assistants = rag.list_assistants()
+        assert len(list_assistants) > 0, "Do not exist any assistant"
+        for assistant in list_assistants:
+            assert isinstance(assistant, Assistant), "Existence type is not assistant."
+
+    def test_get_detail_assistant_with_success(self):
+        """
+        Test getting an assistant's detail with success
+        """
+        rag = RAGFlow(API_KEY, HOST_ADDRESS)
+        assistant = rag.get_assistant(name="God")
+        assert isinstance(assistant, Assistant), f"Failed to get assistant, error: {assistant}."
+        assert assistant.name == "God", "Name does not match"