mirror of
https://github.com/infiniflow/ragflow.git
synced 2025-07-04 15:41:09 +00:00
28 lines
1.2 KiB
Python
28 lines
1.2 KiB
Python
|
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."
|