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 Chat API (#2804)

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

Refactor Chat API

### Type of change

- [x] Refactoring

---------

Co-authored-by: liuhua <10215101452@stu.ecun.edu.cn>
This commit is contained in:
liuhua
2024-10-12 13:48:43 +08:00
committed by GitHub
parent 2a86472b88
commit a20b82092f
14 changed files with 883 additions and 652 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
sdk/python/ragflow/__init__.py
View File

@ -4,7 +4,7 @@ __version__ = importlib.metadata.version("ragflow")
from .ragflow import RAGFlow
from .modules.dataset import DataSet
from .modules.assistant import Assistant
from .modules.chat import Chat
from .modules.session import Session
from .modules.document import Document
from .modules.chunk import Chunk

8
sdk/python/ragflow/modules/base.py
View File

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

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

@ -4,14 +4,14 @@ from .base import Base
from .session import Session
class Assistant(Base):
class Chat(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, {})
self.llm = Chat.LLM(rag, {})
self.prompt = Chat.Prompt(rag, {})
super().__init__(rag, res_dict)
class LLM(Base):
@ -42,21 +42,13 @@ class Assistant(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()
})
def update(self, update_message: dict):
res = self.put(f'/chat/{self.id}',
update_message)
res = res.json()
if res.get("retmsg") == "success": return True
raise Exception(res["retmsg"])
if res.get("code") != 0:
raise Exception(res["message"])
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})

49
sdk/python/ragflow/ragflow.py
View File

@ -17,7 +17,7 @@ from typing import List
import requests
from .modules.assistant import Assistant
from .modules.chat import Chat
from .modules.chunk import Chunk
from .modules.dataset import DataSet
from .modules.document import Document
@ -32,16 +32,16 @@ class RAGFlow:
self.api_url = f"{base_url}/api/{version}"
self.authorization_header = {"Authorization": "{} {}".format("Bearer", self.user_key)}
def post(self, path, param, stream=False):
res = requests.post(url=self.api_url + path, json=param, headers=self.authorization_header, stream=stream)
def post(self, path, json, stream=False):
res = requests.post(url=self.api_url + path, json=json, headers=self.authorization_header, stream=stream)
return res
def get(self, path, params=None):
res = requests.get(url=self.api_url + path, params=params, headers=self.authorization_header)
return res
def delete(self, path, params):
res = requests.delete(url=self.api_url + path, json=params, headers=self.authorization_header)
def delete(self, path, json):
res = requests.delete(url=self.api_url + path, json=json, headers=self.authorization_header)
return res
def put(self, path, json):
@ -68,7 +68,7 @@ class RAGFlow:
return DataSet(self, res["data"])
raise Exception(res["message"])
def delete_dataset(self, ids: List[str] = None, names: List[str] = None):
def delete_datasets(self, ids: List[str] = None, names: List[str] = None):
res = self.delete("/dataset",{"ids": ids, "names": names})
res=res.json()
if res.get("code") != 0:
@ -87,21 +87,21 @@ class RAGFlow:
return result_list
raise Exception(res["message"])
def create_assistant(self, name: str = "assistant", avatar: str = "path", knowledgebases: List[DataSet] = [],
llm: Assistant.LLM = None, prompt: Assistant.Prompt = None) -> Assistant:
def create_chat(self, name: str = "assistant", avatar: str = "path", knowledgebases: List[DataSet] = [],
llm: Chat.LLM = None, prompt: Chat.Prompt = None) -> Chat:
datasets = []
for dataset in knowledgebases:
datasets.append(dataset.to_json())
if llm is None:
llm = Assistant.LLM(self, {"model_name": None,
llm = Chat.LLM(self, {"model_name": None,
"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,
prompt = Chat.Prompt(self, {"similarity_threshold": 0.2,
"keywords_similarity_weight": 0.7,
"top_n": 8,
"variables": [{
@ -127,28 +127,29 @@ class RAGFlow:
"knowledgebases": datasets,
"llm": llm.to_json(),
"prompt": prompt.to_json()}
res = self.post("/assistant/save", temp_dict)
res = self.post("/chat", temp_dict)
res = res.json()
if res.get("retmsg") == "success":
return Assistant(self, res["data"])
raise Exception(res["retmsg"])
if res.get("code") == 0:
return Chat(self, res["data"])
raise Exception(res["message"])
def get_assistant(self, id: str = None, name: str = None) -> Assistant:
res = self.get("/assistant/get", {"id": id, "name": name})
def delete_chats(self,ids: List[str] = None,names: List[str] = None ) -> bool:
res = self.delete('/chat',
{"ids":ids, "names":names})
res = res.json()
if res.get("retmsg") == "success":
return Assistant(self, res['data'])
raise Exception(res["retmsg"])
if res.get("code") != 0:
raise Exception(res["message"])
def list_assistants(self) -> List[Assistant]:
res = self.get("/assistant/list")
def list_chats(self, page: int = 1, page_size: int = 1024, orderby: str = "create_time", desc: bool = True,
id: str = None, name: str = None) -> List[Chat]:
res = self.get("/chat",{"page": page, "page_size": page_size, "orderby": orderby, "desc": desc, "id": id, "name": name})
res = res.json()
result_list = []
if res.get("retmsg") == "success":
if res.get("code") == 0:
for data in res['data']:
result_list.append(Assistant(self, data))
result_list.append(Chat(self, data))
return result_list
raise Exception(res["retmsg"])
raise Exception(res["message"])
def create_document(self, ds: DataSet, name: str, blob: bytes) -> bool:
url = f"/doc/dataset/{ds.id}/documents/upload"

68
sdk/python/test/t_assistant.py
View File

@ -1,68 +0,0 @@
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.create_dataset(name="test_create_assistant")
assistant = rag.create_assistant("test_create", knowledgebases=[kb])
if isinstance(assistant, Assistant):
assert assistant.name == "test_create", "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.create_dataset(name="test_update_assistant")
assistant = rag.create_assistant("test_update", knowledgebases=[kb])
if isinstance(assistant, Assistant):
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:
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.create_dataset(name="test_delete_assistant")
assistant = rag.create_assistant("test_delete", knowledgebases=[kb])
if isinstance(assistant, Assistant):
assert assistant.name == "test_delete", "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)
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 == "test_get_assistant", "Name does not match"

56
sdk/python/test/t_chat.py Normal file
View File

@ -0,0 +1,56 @@
from ragflow import RAGFlow, Chat
from common import API_KEY, HOST_ADDRESS
from test_sdkbase import TestSdk
class TestChat(TestSdk):
def test_create_chat_with_success(self):
"""
Test creating an chat with success
"""
rag = RAGFlow(API_KEY, HOST_ADDRESS)
kb = rag.create_dataset(name="test_create_chat")
chat = rag.create_chat("test_create", knowledgebases=[kb])
if isinstance(chat, Chat):
assert chat.name == "test_create", "Name does not match."
else:
assert False, f"Failed to create chat, error: {chat}"
def test_update_chat_with_success(self):
"""
Test updating an chat with success.
"""
rag = RAGFlow(API_KEY, HOST_ADDRESS)
kb = rag.create_dataset(name="test_update_chat")
chat = rag.create_chat("test_update", knowledgebases=[kb])
if isinstance(chat, Chat):
assert chat.name == "test_update", "Name does not match."
res=chat.update({"name":"new_chat"})
assert res is None, f"Failed to update chat, error: {res}"
else:
assert False, f"Failed to create chat, error: {chat}"
def test_delete_chats_with_success(self):
"""
Test deleting an chat with success
"""
rag = RAGFlow(API_KEY, HOST_ADDRESS)
kb = rag.create_dataset(name="test_delete_chat")
chat = rag.create_chat("test_delete", knowledgebases=[kb])
if isinstance(chat, Chat):
assert chat.name == "test_delete", "Name does not match."
res = rag.delete_chats(ids=[chat.id])
assert res is None, f"Failed to delete chat, error: {res}"
else:
assert False, f"Failed to create chat, error: {chat}"
def test_list_chats_with_success(self):
"""
Test listing chats with success
"""
rag = RAGFlow(API_KEY, HOST_ADDRESS)
list_chats = rag.list_chats()
assert len(list_chats) > 0, "Do not exist any chat"
for chat in list_chats:
assert isinstance(chat, Chat), "Existence type is not chat."

4
sdk/python/test/t_dataset.py
View File

@ -29,7 +29,7 @@ class TestDataset(TestSdk):
else:
assert False, f"Failed to create dataset, error: {ds}"
def test_delete_dataset_with_success(self):
def test_delete_datasets_with_success(self):
"""
Test deleting a dataset with success
"""
@ -37,7 +37,7 @@ class TestDataset(TestSdk):
ds = rag.create_dataset("MA")
if isinstance(ds, DataSet):
assert ds.name == "MA", "Name does not match."
res = rag.delete_dataset(names=["MA"])
res = rag.delete_datasets(ids=[ds.id])
assert res is None, f"Failed to delete dataset, error: {res}"
else:
assert False, f"Failed to create dataset, error: {ds}"