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"