mirror of
https://git.mirrors.martin98.com/https://github.com/infiniflow/ragflow.git
synced 2025-04-23 06:30:00 +08:00
336a639164
### 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>
69 lines
2.9 KiB
Python
69 lines
2.9 KiB
Python
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"
|