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