pygeai 0.1.6__py3-none-any.whl → 0.6.0b15__py3-none-any.whl

pygeai/tests/assistants/rag/test_models.py

@@ -0,0 +1,292 @@
+import unittest
+from datetime import datetime
+
+from pydantic import BaseModel
+
+from pygeai.assistant.rag.models import (
+    DocumentMetadata,
+    Document,
+    Search,
+    RetrieverOptions,
+    ChainOptions,
+    EmbeddingsOptions,
+    IngestionOptions,
+    SearchOptions,
+    ChildOptions,
+    ChildDocumentOptions,
+    ChunkOptions,
+    IndexOptions,
+    RAGAssistant,
+    UploadDocument,
+    UploadType
+)
+from pygeai.core.models import LlmSettings, ChatVariableList, Assistant
+
+
+class TestRagModels(unittest.TestCase):
+    """
+    python -m unittest pygeai.tests.assistants.rag.test_models.TestRagModels
+    """
+    def test_document_metadata_to_dict(self):
+        metadata = DocumentMetadata(key="test_key", value="test_value")
+        result = metadata.to_dict()
+        self.assertEqual(result, {"key": "test_key", "value": "test_value"})
+
+    def test_document_metadata_str(self):
+        metadata = DocumentMetadata(key="test_key", value="test_value")
+        result = str(metadata)
+        self.assertEqual(result, "{'key': 'test_key', 'value': 'test_value'}")
+
+    def test_document_to_dict(self):
+        doc = Document(
+            id="doc1",
+            chunks="chunk_data",
+            name="test_doc",
+            extension="pdf",
+            index_status="indexed",
+            metadata=[DocumentMetadata(key="key1", value="value1")],
+            timestamp=datetime(2023, 1, 1),
+            url="http://example.com"
+        )
+        result = doc.to_dict()
+        expected = {
+            "id": "doc1",
+            "chunks": "chunk_data",
+            "name": "test_doc",
+            "extension": "pdf",
+            "indexStatus": "indexed",
+            "metadata": [{"key": "key1", "value": "value1"}],
+            "timestamp": datetime(2023, 1, 1),
+            "url": "http://example.com"
+        }
+        self.assertEqual(result["id"], expected["id"])
+        self.assertEqual(result["chunks"], expected["chunks"])
+        self.assertEqual(result["extension"], expected["extension"])
+        self.assertEqual(result["indexStatus"], expected["indexStatus"])
+        self.assertEqual(result["url"], expected["url"])
+        self.assertEqual(result["timestamp"], expected["timestamp"])
+        self.assertEqual(result["name"], expected["name"])
+        self.assertEqual(result["metadata"], expected["metadata"])
+
+    def test_document_str(self):
+        doc = Document(
+            id="doc1",
+            chunks="chunk_data",
+            extension="pdf",
+            index_status="indexed",
+            url="http://example.com"
+        )
+        result = str(doc)
+        self.assertIn("'id': 'doc1'", result)
+        self.assertIn("'url': 'http://example.com'", result)
+
+    def test_search_to_dict_with_mmr(self):
+        search = Search(
+            k=5,
+            type="mmr",
+            fetch_k=10.0,
+            lambda_=0.5,
+            prompt="test prompt",
+            return_source_documents=True,
+            score_threshold=0.8,
+            template="test template"
+        )
+        result = search.to_dict()
+        expected = {
+            "k": 5,
+            "type": "mmr",
+            "fetchK": 10.0,
+            "lambda": 0.5,
+            "prompt": "test prompt",
+            "returnSourceDocuments": True,
+            "scoreThreshold": 0.8,
+            "template": "test template"
+        }
+        self.assertEqual(result, expected)
+
+    def test_search_to_dict_similarity(self):
+        search = Search(
+            k=3,
+            type="similarity",
+            prompt="test prompt",
+            return_source_documents=False,
+            score_threshold=0.5,
+            template="test template"
+        )
+        result = search.to_dict()
+        expected = {
+            "k": 3,
+            "type": "similarity",
+            "prompt": "test prompt",
+            "returnSourceDocuments": False,
+            "scoreThreshold": 0.5,
+            "template": "test template"
+        }
+        self.assertEqual(result, expected)
+
+    def test_retriever_options_to_dict(self):
+        options = RetrieverOptions(
+            type="multiQuery",
+            search_type="hybrid",
+            step="documents",
+            prompt="custom prompt"
+        )
+        result = options.to_dict()
+        expected = {
+            "type": "multiQuery",
+            "searchType": "hybrid",
+            "step": "documents",
+            "prompt": "custom prompt"
+        }
+        self.assertEqual(result, expected)
+
+    def test_embeddings_options_initialization(self):
+        embeddings = EmbeddingsOptions(
+            dimensions=512,
+            model_name="test-model",
+            provider="test-provider",
+            use_proxy=True
+        )
+        self.assertEqual(embeddings.dimensions, 512)
+        self.assertEqual(embeddings.model_name, "test-model")
+        self.assertEqual(embeddings.provider, "test-provider")
+        self.assertTrue(embeddings.use_proxy)
+
+    def test_ingestion_options_initialization(self):
+        ingestion = IngestionOptions(
+            geai_options={"opt1": "val1"},
+            llama_parse_options={"opt2": "val2"},
+            provider="test-provider"
+        )
+        self.assertEqual(ingestion.geai_options, {"opt1": "val1"})
+        self.assertEqual(ingestion.llama_parse_options, {"opt2": "val2"})
+        self.assertEqual(ingestion.provider, "test-provider")
+
+    def test_search_options_to_dict_full(self):
+        search_opts = SearchOptions(
+            history_count=5,
+            llm=LlmSettings(provider_name="test", model_name="model"),
+            search=Search(k=3, prompt="prompt", return_source_documents=True, score_threshold=0.5, template="template"),
+            retriever=RetrieverOptions(type="vectorStore"),
+            chain=ChainOptions(type="test-chain"),
+            embeddings=EmbeddingsOptions(dimensions=512, model_name="model", provider="provider"),
+            ingestion=IngestionOptions(geai_options={}, llama_parse_options={}, provider="provider"),
+            options={"opt": "val"},
+            rerank={"rerank_opt": "val"},
+            variables=ChatVariableList(variables=[]),
+            vector_store={"store": "data"}
+        )
+        result = search_opts.to_dict()
+        self.assertEqual(result["historyCount"], 5)
+        self.assertIn("llm", result)
+        self.assertIn("search", result)
+        self.assertIn("retriever", result)
+        self.assertIn("chain", result)
+        self.assertIn("embeddings", result)
+        self.assertIn("ingestion", result)
+        self.assertIn("options", result)
+        self.assertIn("rerank", result)
+        self.assertIn("vectorStore", result)
+        # Do not assert "variables" since it may be excluded if empty or None after to_dict()
+
+    def test_child_options_to_dict(self):
+        child_opts = ChildOptions(chunk_size=100.0, chunk_overlap=20.0, content_processing="clean")
+        result = child_opts.to_dict()
+        expected = {
+            "chunkSize": 100.0,
+            "chunkOverlap": 20.0,
+            "contentProcessing": "clean"
+        }
+        self.assertEqual(result, expected)
+
+    def test_child_document_options_to_dict(self):
+        child_doc_opts = ChildDocumentOptions(
+            child_k=5.0,
+            child=ChildOptions(chunk_size=100.0, chunk_overlap=20.0)
+        )
+        result = child_doc_opts.to_dict()
+        expected = {
+            "childK": 5.0,
+            "child": {"chunkSize": 100.0, "chunkOverlap": 20.0, "contentProcessing": ""}
+        }
+        self.assertEqual(result, expected)
+
+    def test_chunk_options_to_dict(self):
+        chunk_opts = ChunkOptions(chunk_overlap=50, chunk_size=500)
+        result = chunk_opts.to_dict()
+        expected = {"chunkOverlap": 50, "chunkSize": 500}
+        self.assertEqual(result, expected)
+
+    def test_index_options_to_dict_with_child(self):
+        index_opts = IndexOptions(
+            chunks=ChunkOptions(chunk_overlap=50, chunk_size=500),
+            use_parent_document=True,
+            child_document=ChildDocumentOptions(
+                child_k=5.0,
+                child=ChildOptions(chunk_size=100.0, chunk_overlap=20.0)
+            )
+        )
+        result = index_opts.to_dict()
+        expected = {
+            "chunks": {"chunkOverlap": 50, "chunkSize": 500},
+            "useParentDocument": True,
+            "childDocument": {
+                "childK": 5.0,
+                "child": {"chunkSize": 100.0, "chunkOverlap": 20.0, "contentProcessing": ""}
+            }
+        }
+        self.assertEqual(result, expected)
+
+    def test_rag_assistant_to_dict(self):
+        rag_assistant = RAGAssistant(
+            id="rag1",
+            name="RAG Assistant",
+            type="rag",
+            template="rag template",
+            search_options=SearchOptions(
+                history_count=5,
+                llm=LlmSettings(provider_name="test", model_name="model"),
+                search=Search(k=3, prompt="prompt", return_source_documents=True, score_threshold=0.5, template="template"),
+                retriever=RetrieverOptions(type="vectorStore")
+            ),
+            index_options=IndexOptions(chunks=ChunkOptions(chunk_overlap=50, chunk_size=500))
+        )
+        result = rag_assistant.to_dict()
+        self.assertIn("template", result)
+        self.assertIn("searchOptions", result)
+        self.assertIn("indexOptions", result)
+        self.assertEqual(result["template"], "rag template")
+
+    def test_rag_assistant_str(self):
+        rag_assistant = RAGAssistant(
+            id="rag1",
+            name="RAG Assistant",
+            type="rag",
+            template="rag template",
+            search_options=SearchOptions(
+                history_count=5,
+                llm=LlmSettings(provider_name="test", model_name="model"),
+                search=Search(k=3, prompt="prompt", return_source_documents=True, score_threshold=0.5, template="template"),
+                retriever=RetrieverOptions(type="vectorStore")
+            ),
+            index_options=IndexOptions(chunks=ChunkOptions(chunk_overlap=50, chunk_size=500))
+        )
+        result = str(rag_assistant)
+        self.assertIn("'template': 'rag template'", result)
+
+    def test_upload_document_initialization(self):
+        upload_doc = UploadDocument(
+            path="/path/to/doc",
+            upload_type="binary",
+            metadata={"key": "value"},
+            content_type="application/pdf"
+        )
+        self.assertEqual(upload_doc.path, "/path/to/doc")
+        self.assertEqual(upload_doc.upload_type, "binary")
+        self.assertEqual(upload_doc.metadata, {"key": "value"})
+        self.assertEqual(upload_doc.content_type, "application/pdf")
+
+    def test_upload_type_constants(self):
+        self.assertEqual(UploadType.BINARY, "binary")
+        self.assertEqual(UploadType.MULTIPART, "multipart")
+

pygeai/tests/assistants/test_clients.py

@@ -1,99 +1,195 @@
 import unittest
-from unittest import TestCase
-from unittest.mock import patch
-
+from unittest.mock import patch, MagicMock
+from json import JSONDecodeError
 from pygeai.assistant.clients import AssistantClient
-from pygeai.core.base.session import get_session
-
-session = get_session()
+from pygeai.core.common.exceptions import InvalidAPIResponseException, APIResponseError, APIResponseError
 
 
-class TestAssistantClient(TestCase):
+class TestAssistantClient(unittest.TestCase):
     """
     python -m unittest pygeai.tests.assistants.test_clients.TestAssistantClient
     """
+    def setUp(self):
+        with patch('pygeai.core.base.clients.BaseClient.__init__', return_value=None):
+            self.client = AssistantClient()
+        self.mock_response = MagicMock()
+        self.client.api_service = MagicMock()
+        self.assistant_id = "test_id"
+        self.assistant_name = "Test Assistant"
+        self.request_id = "req_123"
+
+    def test_get_assistant_data_successful_json_response(self):
+        detail = "summary"
+        expected_response = {"id": self.assistant_id, "name": "Test Assistant"}
+        self.mock_response.json.return_value = expected_response
+        self.mock_response.status_code = 200
+        self.client.api_service.get.return_value = self.mock_response
+
+        result = self.client.get_assistant_data(self.assistant_id, detail)
+
+        self.client.api_service.get.assert_called_once()
+        self.assertEqual(result, expected_response)
+
+    def test_get_assistant_data_json_decode_error(self):
+        detail = "full"
+        self.mock_response.json.side_effect = JSONDecodeError("Invalid JSON", "", 0)
+        self.mock_response.text = "Invalid JSON response"
+        self.mock_response.status_code = 500
+        self.client.api_service.get.return_value = self.mock_response
+
+        with self.assertRaises(APIResponseError) as context:
+            self.client.get_assistant_data(self.assistant_id, detail)
+
+        self.client.api_service.get.assert_called_once()
+        self.assertIn("API returned an error", str(context.exception))  # "Unable to get assistant data", str(context.exception))
+
+    def test_create_assistant_successful_json_response(self):
+        assistant_type = "text"
+        name = "New Assistant"
+        prompt = "Help with text"
+        expected_response = {"id": "new_id", "name": name}
+        self.mock_response.json.return_value = expected_response
+        self.mock_response.status_code = 200
+        self.client.api_service.post.return_value = self.mock_response
+
+        result = self.client.create_assistant(assistant_type, name, prompt)
+
+        self.client.api_service.post.assert_called_once()
+        self.assertEqual(result, expected_response)
+
+    def test_create_assistant_json_decode_error(self):
+        assistant_type = "chat"
+        name = "Chat Assistant"
+        prompt = "Chat prompt"
+        self.mock_response.json.side_effect = JSONDecodeError("Invalid JSON", "", 0)
+        self.mock_response.text = "Error in response"
+        self.mock_response.status_code = 500
+        self.client.api_service.post.return_value = self.mock_response
+
+        with self.assertRaises(APIResponseError) as context:
+            self.client.create_assistant(assistant_type, name, prompt)
+
+        self.client.api_service.post.assert_called_once()
+        self.assertIn("API returned an error", str(context.exception))  # "Unable to create assistant", str(context.exception))
+
+    def test_update_assistant_successful_json_response(self):
+        status = 1
+        action = "save"
+        expected_response = {"id": self.assistant_id, "status": status}
+        self.mock_response.json.return_value = expected_response
+        self.mock_response.status_code = 200
+        self.client.api_service.put.return_value = self.mock_response
+
+        result = self.client.update_assistant(self.assistant_id, status, action)
+
+        self.client.api_service.put.assert_called_once()
+        self.assertEqual(result, expected_response)
+
+    def test_update_assistant_json_decode_error(self):
+        status = 2
+        action = "saveNewRevision"
+        self.mock_response.json.side_effect = JSONDecodeError("Invalid JSON", "", 0)
+        self.mock_response.text = "Update failed"
+        self.mock_response.status_code = 500
+        self.client.api_service.put.return_value = self.mock_response
+
+        with self.assertRaises(APIResponseError) as context:
+            self.client.update_assistant(self.assistant_id, status, action)
+
+        self.client.api_service.put.assert_called_once()
+        self.assertIn("API returned an error", str(context.exception))  # "Unable to update assistant", str(context.exception))
+
+    def test_delete_assistant_successful_json_response(self):
+        expected_response = {"message": "Assistant deleted"}
+        self.mock_response.json.return_value = expected_response
+        self.mock_response.status_code = 200
+        self.client.api_service.delete.return_value = self.mock_response
+
+        result = self.client.delete_assistant(self.assistant_id)
+
+        self.client.api_service.delete.assert_called_once()
+        self.assertEqual(result, expected_response)
+
+    def test_delete_assistant_json_decode_error(self):
+        self.mock_response.json.side_effect = JSONDecodeError("Invalid JSON", "", 0)
+        self.mock_response.text = "Delete error"
+        self.mock_response.status_code = 500
+        self.client.api_service.delete.return_value = self.mock_response
+
+        with self.assertRaises(APIResponseError) as context:
+            self.client.delete_assistant(self.assistant_id)
+
+        self.client.api_service.delete.assert_called_once()
+        self.assertIn("API returned an error", str(context.exception))  # "Unable to delete assistant", str(context.exception))
+
+    def test_send_chat_request_successful_json_response(self):
+        messages = [{"role": "user", "content": "Hello"}]
+        expected_response = {"response": "Hi there!"}
+        self.mock_response.json.return_value = expected_response
+        self.mock_response.status_code = 200
+        self.client.api_service.post.return_value = self.mock_response
+
+        result = self.client.send_chat_request(self.assistant_name, messages)
+
+        self.client.api_service.post.assert_called_once()
+        self.assertEqual(result, expected_response)
+
+    def test_send_chat_request_json_decode_error(self):
+        messages = [{"role": "user", "content": "Hello"}]
+        self.mock_response.json.side_effect = JSONDecodeError("Invalid JSON", "", 0)
+        self.mock_response.text = "Chat error"
+        self.mock_response.status_code = 500
+        self.client.api_service.post.return_value = self.mock_response
+
+        with self.assertRaises(APIResponseError) as context:
+            self.client.send_chat_request(self.assistant_name, messages)
+
+        self.client.api_service.post.assert_called_once()
+        self.assertIn("API returned an error", str(context.exception))  # "Unable to send chat request", str(context.exception))
 
-    @patch("pygeai.core.services.rest.ApiService.get")
-    def test_get_assistant_data(self, mock_get):
-        mock_get.return_value.content = '{"assistant": {"id": "123", "name": "assistant1"}}'
-
-        client = AssistantClient()
-        result = client.get_assistant_data("123")
-
-        self.assertIsNotNone(result)
-        self.assertEqual(result['assistant']['name'], "assistant1")
-
-    @patch("pygeai.core.services.rest.ApiService.post")
-    def test_create_assistant(self, mock_post):
-        mock_post.return_value.content = '{"assistant": {"id": "123", "name": "assistant1"}}'
-
-        client = AssistantClient()
-        result = client.create_assistant(
-            assistant_type="text",
-            name="assistant1",
-            prompt="Hello",
-            description="Test assistant"
-        )
-
-        self.assertIsNotNone(result)
-        self.assertEqual(result['assistant']['name'], "assistant1")
-
-    @patch("pygeai.core.services.rest.ApiService.put")
-    def test_update_assistant(self, mock_put):
-        mock_put.return_value.content = '{"assistant": {"id": "123", "name": "updated_assistant"}}'
-
-        client = AssistantClient()
-        result = client.update_assistant(
-            assistant_id="123",
-            status=1,
-            action="save",
-            name="updated_assistant"
-        )
-
-        self.assertIsNotNone(result)
-        self.assertEqual(result['assistant']['name'], "updated_assistant")
-
-    @patch("pygeai.core.services.rest.ApiService.delete")
-    def test_delete_assistant(self, mock_delete):
-        mock_delete.return_value.content = '{"status": "deleted"}'
+    def test_get_request_status_successful_json_response(self):
+        expected_response = {"status": "completed"}
+        self.mock_response.json.return_value = expected_response
+        self.mock_response.status_code = 200
+        self.client.api_service.get.return_value = self.mock_response
 
-        client = AssistantClient()
-        result = client.delete_assistant("123")
+        result = self.client.get_request_status(self.request_id)
 
-        self.assertIsNotNone(result)
-        self.assertEqual(result['status'], "deleted")
+        self.client.api_service.get.assert_called_once()
+        self.assertEqual(result, expected_response)
+
+    def test_get_request_status_json_decode_error(self):
+        self.mock_response.json.side_effect = JSONDecodeError("Invalid JSON", "", 0)
+        self.mock_response.text = "Status error"
+        self.mock_response.status_code = 500
+        self.client.api_service.get.return_value = self.mock_response
 
-    @patch("pygeai.core.services.rest.ApiService.post")
-    def test_send_chat_request(self, mock_post):
-        mock_post.return_value.content = '{"status": "success"}'
+        with self.assertRaises(APIResponseError) as context:
+            self.client.get_request_status(self.request_id)
 
-        client = AssistantClient()
-        result = client.send_chat_request(
-            assistant_name="assistant1",
-            messages=[{"role": "user", "content": "Hello"}],
-            revision=1,
-            revision_name="v1"
-        )
+        self.client.api_service.get.assert_called_once()
+        self.assertIn("API returned an error", str(context.exception))  # "Unable to get request status", str(context.exception))
 
-        self.assertIsNotNone(result)
-        self.assertEqual(result['status'], "success")
+    def test_cancel_request_successful_json_response(self):
+        expected_response = {"message": "Request cancelled"}
+        self.mock_response.json.return_value = expected_response
+        self.mock_response.status_code = 200
+        self.client.api_service.post.return_value = self.mock_response
 
-    @patch("pygeai.core.services.rest.ApiService.get")
-    def test_get_request_status(self, mock_get):
-        mock_get.return_value.content = '{"status": "completed"}'
+        result = self.client.cancel_request(self.request_id)
 
-        client = AssistantClient()
-        result = client.get_request_status("req123")
+        self.client.api_service.post.assert_called_once()
+        self.assertEqual(result, expected_response)
 
-        self.assertIsNotNone(result)
-        self.assertEqual(result['status'], "completed")
+    def test_cancel_request_json_decode_error(self):
+        self.mock_response.json.side_effect = JSONDecodeError("Invalid JSON", "", 0)
+        self.mock_response.text = "Cancel error"
+        self.mock_response.status_code = 500
+        self.client.api_service.post.return_value = self.mock_response
 
-    @patch("pygeai.core.services.rest.ApiService.post")
-    def test_cancel_request(self, mock_post):
-        mock_post.return_value.content = '{"status": "canceled"}'
+        with self.assertRaises(APIResponseError) as context:
+            self.client.cancel_request(self.request_id)
 
-        client = AssistantClient()
-        result = client.cancel_request("req123")
+        self.client.api_service.post.assert_called_once()
+        self.assertIn("API returned an error", str(context.exception))  # "Unable to cancel request", str(context.exception))
 
-        self.assertIsNotNone(result)
-        self.assertEqual(result['status'], "canceled")