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

SDK for session (#2312)

Browse Source 
### 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>
This commit is contained in:
LiuHua
2024-09-09 17:18:08 +08:00
committed by GitHub
parent ceae4df889
commit 336a639164
8 changed files with 325 additions and 35 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@ -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"

27
sdk/python/test/t_session.py Normal file
View File

@ -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."