File size: 1,239 Bytes
172caf6 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 |
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."
|