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

pygeai/tests/core/files/test_responses.py

@@ -0,0 +1,122 @@
+import unittest
+from unittest import TestCase
+
+from pygeai.core.files.models import BaseFile, File, FileList, UploadFile
+from pygeai.core.files.responses import UploadFileResponse
+
+
+class TestFileResponses(TestCase):
+    """
+    python -m unittest pygeai.tests.core.files.test_responses.TestFileResponses
+    """
+
+    def test_base_file_model_validate(self):
+        base_file_data = {
+            "id": "file123",
+            "name": "testfile.txt"
+        }
+        base_file = BaseFile.model_validate(base_file_data)
+        self.assertEqual(base_file.id, base_file_data["id"])
+        self.assertEqual(base_file.name, base_file_data["name"])
+        self.assertEqual(base_file.model_dump(), base_file_data)
+
+    def test_base_file_model_empty(self):
+        base_file_data = {}
+        base_file = BaseFile.model_validate(base_file_data)
+        self.assertIsNone(base_file.id)
+        self.assertIsNone(base_file.name)
+        self.assertEqual(base_file.model_dump(exclude_none=True), {})
+
+    def test_upload_file_model_validate(self):
+        upload_file_data = {
+            "id": "file123",
+            "name": "testfile.txt",
+            "path": "/uploads/testfile.txt",
+            "folder": "uploads"
+        }
+        upload_file = UploadFile.model_validate(upload_file_data)
+        self.assertEqual(upload_file.id, upload_file_data["id"])
+        self.assertEqual(upload_file.name, upload_file_data["name"])
+        self.assertEqual(upload_file.path, upload_file_data["path"])
+        self.assertEqual(upload_file.folder, upload_file_data["folder"])
+        self.assertEqual(upload_file.model_dump(), upload_file_data)
+
+    def test_upload_file_model_minimal(self):
+        upload_file_data = {
+            "path": "/uploads/testfile.txt"
+        }
+        upload_file = UploadFile.model_validate(upload_file_data)
+        self.assertIsNone(upload_file.id)
+        self.assertIsNone(upload_file.name)
+        self.assertEqual(upload_file.path, upload_file_data["path"])
+        self.assertIsNone(upload_file.folder)
+        self.assertEqual(upload_file.model_dump(exclude_none=True), upload_file_data)
+
+    def test_file_model_validate(self):
+        file_data = {
+            "id": "file123",
+            "name": "testfile.txt",
+            "extension": "txt",
+            "purpose": "test",
+            "size": 1024,
+            "url": "https://example.com/files/testfile.txt"
+        }
+        file_obj = File.model_validate(file_data)
+        self.assertEqual(file_obj.id, file_data["id"])
+        self.assertEqual(file_obj.name, file_data["name"])
+        self.assertEqual(file_obj.extension, file_data["extension"])
+        self.assertEqual(file_obj.purpose, file_data["purpose"])
+        self.assertEqual(file_obj.size, file_data["size"])
+        self.assertEqual(file_obj.url, file_data["url"])
+        self.assertEqual(file_obj.model_dump(), file_data)
+
+    def test_file_model_empty(self):
+        file_data = {}
+        file_obj = File.model_validate(file_data)
+        self.assertIsNone(file_obj.id)
+        self.assertIsNone(file_obj.name)
+        self.assertIsNone(file_obj.extension)
+        self.assertIsNone(file_obj.purpose)
+        self.assertEqual(file_obj.size, 0)
+        self.assertIsNone(file_obj.url)
+        self.assertEqual(file_obj.model_dump(exclude_none=True), {"size": 0})
+
+    def test_file_list_model_validate(self):
+        file_list_data = {
+            "files": [
+                {"id": "file1", "name": "file1.txt", "size": 512},
+                {"id": "file2", "name": "file2.pdf", "size": 2048, "url": "https://example.com/files/file2.pdf"}
+            ]
+        }
+        file_list = FileList.model_validate(file_list_data)
+        self.assertEqual(len(file_list.files), 2)
+        self.assertEqual(file_list.files[0].id, file_list_data["files"][0]["id"])
+        self.assertEqual(file_list.files[0].name, file_list_data["files"][0]["name"])
+        self.assertEqual(file_list.files[1].url, file_list_data["files"][1]["url"])
+        self.assertEqual(file_list.model_dump(exclude_none=True), file_list_data)
+
+    def test_file_list_model_empty(self):
+        file_list_data = {
+            "files": []
+        }
+        file_list = FileList.model_validate(file_list_data)
+        self.assertEqual(len(file_list.files), 0)
+        self.assertEqual(file_list.model_dump(exclude_none=True), file_list_data)
+
+    def test_upload_file_response_model_validate(self):
+        upload_response_data = {
+            "id": "file123",
+            "url": "https://example.com/uploaded/file123",
+            "success": True
+        }
+        upload_response = UploadFileResponse.model_validate(upload_response_data)
+        self.assertEqual(upload_response.id, upload_response_data["id"])
+        self.assertEqual(upload_response.url, upload_response_data["url"])
+        self.assertEqual(upload_response.success, upload_response_data["success"])
+        expected_dict = {
+            "dataFileId": upload_response_data["id"],
+            "dataFileUrl": upload_response_data["url"],
+            "success": upload_response_data["success"]
+        }
+        self.assertEqual(upload_response.to_dict(), expected_dict)
+        self.assertEqual(str(upload_response), str(expected_dict))

pygeai/tests/core/llm/test_clients.py

@@ -0,0 +1,142 @@
+import unittest
+from unittest.mock import patch, MagicMock
+from json import JSONDecodeError
+from pygeai.core.llm.clients import LlmClient
+from pygeai.core.common.exceptions import InvalidAPIResponseException, APIResponseError, APIResponseError
+
+
+class TestLlmClient(unittest.TestCase):
+    """
+    python -m unittest pygeai.tests.core.llm.test_clients.TestLlmClient
+    """
+    def setUp(self):
+        with patch('pygeai.core.base.clients.BaseClient.__init__', return_value=None):
+            self.client = LlmClient()
+        self.mock_response = MagicMock()
+        self.client.api_service = MagicMock()
+        self.provider_name = "test-provider"
+        self.model_name = "test-model"
+        self.model_id = "test-model-id"
+
+    def test_get_provider_list_success(self):
+        expected_response = {"providers": [{"name": "provider1"}, {"name": "provider2"}]}
+        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_provider_list()
+
+        self.client.api_service.get.assert_called_once()
+        self.assertEqual(result, expected_response)
+
+    def test_get_provider_list_json_decode_error(self):
+        self.mock_response.json.side_effect = JSONDecodeError("Invalid JSON", "", 0)
+        self.mock_response.text = "Error text"
+        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_provider_list()
+
+        self.client.api_service.get.assert_called_once()
+        self.assertIn("API returned an error", str(context.exception))  # "Unable to obtain provider list", str(context.exception))
+
+    def test_get_provider_data_success(self):
+        expected_response = {"name": self.provider_name, "details": "test details"}
+        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_provider_data(provider_name=self.provider_name)
+
+        self.client.api_service.get.assert_called_once()
+        endpoint = self.client.api_service.get.call_args[1]['endpoint']
+        self.assertIn(self.provider_name, endpoint)
+        self.assertEqual(result, expected_response)
+
+    def test_get_provider_data_json_decode_error(self):
+        self.mock_response.json.side_effect = JSONDecodeError("Invalid JSON", "", 0)
+        self.mock_response.text = "Error text"
+        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_provider_data(provider_name=self.provider_name)
+
+        self.client.api_service.get.assert_called_once()
+        self.assertIn("API returned an error", str(context.exception))  # "Unable to obtain provider data", str(context.exception))
+
+    def test_get_provider_models_success(self):
+        expected_response = {"models": [{"name": "model1"}, {"name": "model2"}]}
+        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_provider_models(provider_name=self.provider_name)
+
+        self.client.api_service.get.assert_called_once()
+        endpoint = self.client.api_service.get.call_args[1]['endpoint']
+        self.assertIn(self.provider_name, endpoint)
+        self.assertEqual(result, expected_response)
+
+    def test_get_provider_models_json_decode_error(self):
+        self.mock_response.json.side_effect = JSONDecodeError("Invalid JSON", "", 0)
+        self.mock_response.text = "Error text"
+        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_provider_models(provider_name=self.provider_name)
+
+        self.client.api_service.get.assert_called_once()
+        self.assertIn("API returned an error", str(context.exception))  # "Unable to obtain provider models", str(context.exception))
+
+    def test_get_model_data_with_model_name_success(self):
+        expected_response = {"name": self.model_name, "provider": self.provider_name}
+        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_model_data(
+            provider_name=self.provider_name,
+            model_name=self.model_name
+        )
+
+        self.client.api_service.get.assert_called_once()
+        endpoint = self.client.api_service.get.call_args[1]['endpoint']
+        self.assertIn(self.provider_name, endpoint)
+        self.assertIn(self.model_name, endpoint)
+        self.assertEqual(result, expected_response)
+
+    def test_get_model_data_with_model_id_success(self):
+        expected_response = {"id": self.model_id, "provider": self.provider_name}
+        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_model_data(
+            provider_name=self.provider_name,
+            model_id=self.model_id
+        )
+
+        self.client.api_service.get.assert_called_once()
+        endpoint = self.client.api_service.get.call_args[1]['endpoint']
+        self.assertIn(self.provider_name, endpoint)
+        self.assertIn(self.model_id, endpoint)
+        self.assertEqual(result, expected_response)
+
+    def test_get_model_data_json_decode_error(self):
+        self.mock_response.json.side_effect = JSONDecodeError("Invalid JSON", "", 0)
+        self.mock_response.text = "Error text"
+        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_model_data(
+                provider_name=self.provider_name,
+                model_name=self.model_name
+            )
+
+        self.client.api_service.get.assert_called_once()
+        self.assertIn("API returned an error", str(context.exception))  # "Unable to obtain model data", str(context.exception))
+

pygeai/tests/core/plugins/test_clients.py

@@ -0,0 +1,66 @@
+import unittest
+from unittest.mock import patch, MagicMock
+from json import JSONDecodeError
+
+from pygeai.core.plugins.clients import PluginClient
+from pygeai.core.common.exceptions import InvalidAPIResponseException, APIResponseError, APIResponseError
+
+
+class TestPluginClient(unittest.TestCase):
+    """
+    python -m unittest pygeai.tests.core.plugins.test_clients.TestPluginClient
+    """
+
+    def setUp(self):
+        self.client = PluginClient()
+        self.mock_response = MagicMock()
+
+    @patch('pygeai.core.services.rest.ApiService.get')
+    def test_list_assistants_success(self, mock_get):
+        self.mock_response.json.return_value = {
+            "assistants": [
+                {"id": "assistant-1", "name": "Assistant 1"},
+                {"id": "assistant-2", "name": "Assistant 2"}
+            ]
+        }
+        self.mock_response.status_code = 200
+        mock_get.return_value = self.mock_response
+        
+        result = self.client.list_assistants(
+            organization_id="org-123",
+            project_id="proj-456"
+        )
+        
+        mock_get.assert_called_once()
+        call_args = mock_get.call_args
+        self.assertEqual(call_args[1]['params']['organization'], "org-123")
+        self.assertEqual(call_args[1]['params']['project'], "proj-456")
+        self.assertEqual(len(result["assistants"]), 2)
+
+    @patch('pygeai.core.services.rest.ApiService.get')
+    def test_list_assistants_empty(self, mock_get):
+        self.mock_response.json.return_value = {"assistants": []}
+        self.mock_response.status_code = 200
+        mock_get.return_value = self.mock_response
+        
+        result = self.client.list_assistants(
+            organization_id="org-123",
+            project_id="proj-456"
+        )
+        
+        self.assertEqual(result["assistants"], [])
+
+    @patch('pygeai.core.services.rest.ApiService.get')
+    def test_list_assistants_json_decode_error(self, mock_get):
+        self.mock_response.json.side_effect = JSONDecodeError("error", "doc", 0)
+        self.mock_response.status_code = 500
+        self.mock_response.text = "Internal server error"
+        mock_get.return_value = self.mock_response
+        
+        with self.assertRaises(APIResponseError) as context:
+            self.client.list_assistants("org-123", "proj-456")
+        self.assertIn("API returned an error", str(context.exception))  # "Unable to list assistants for organization org-123 and project proj-456", str(context.exception))
+
+
+if __name__ == '__main__':
+    unittest.main()

pygeai/tests/core/rerank/test_clients.py

@@ -0,0 +1,76 @@
+import unittest
+from unittest.mock import MagicMock
+from pygeai.core.rerank.clients import RerankClient
+
+
+class TestRerankClient(unittest.TestCase):
+    """
+    python -m unittest pygeai.tests.core.rerank.test_clients.TestRerankClient
+    """
+
+    def setUp(self):
+        self.client = RerankClient()
+        self.client.api_service = MagicMock()
+        self.mock_response = MagicMock()
+        self.mock_response.json.return_value = {"results": [{"index": 0, "relevance_score": 0.95}]}
+        self.mock_response.status_code = 200
+
+    def test_rerank_chunks_success(self):
+        self.client.api_service.post.return_value = self.mock_response
+
+        query = "What is AI?"
+        model = "rerank-english-v3.0"
+        documents = ["AI is artificial intelligence.", "AI stands for something else.", "Random text."]
+        top_n = 2
+
+        result = self.client.rerank_chunks(
+            query=query,
+            model=model,
+            documents=documents,
+            top_n=top_n
+        )
+
+        self.client.api_service.post.assert_called_once()
+        data = self.client.api_service.post.call_args[1]['data']
+        self.assertEqual(data['query'], query)
+        self.assertEqual(data['model'], model)
+        self.assertEqual(data['documents'], documents)
+        self.assertEqual(data['top_n'], top_n)
+        self.assertEqual(result, {"results": [{"index": 0, "relevance_score": 0.95}]})
+
+    def test_rerank_chunks_default_top_n(self):
+        self.client.api_service.post.return_value = self.mock_response
+
+        query = "What is AI?"
+        model = "rerank-english-v3.0"
+        documents = ["AI is artificial intelligence.", "AI stands for something else."]
+
+        result = self.client.rerank_chunks(
+            query=query,
+            model=model,
+            documents=documents
+        )
+
+        self.client.api_service.post.assert_called_once()
+        data = self.client.api_service.post.call_args[1]['data']
+        self.assertEqual(data['top_n'], 3)  # Default value
+        self.assertEqual(result, {"results": [{"index": 0, "relevance_score": 0.95}]})
+
+    def test_rerank_chunks_empty_documents(self):
+        self.client.api_service.post.return_value = self.mock_response
+
+        query = "What is AI?"
+        model = "rerank-english-v3.0"
+        documents = []
+
+        result = self.client.rerank_chunks(
+            query=query,
+            model=model,
+            documents=documents
+        )
+
+        self.client.api_service.post.assert_called_once()
+        data = self.client.api_service.post.call_args[1]['data']
+        self.assertEqual(data['documents'], [])
+        self.assertEqual(result, {"results": [{"index": 0, "relevance_score": 0.95}]})
+

pygeai/tests/core/rerank/test_managers.py

@@ -0,0 +1,99 @@
+import unittest
+from unittest.mock import MagicMock, patch
+from pygeai.core.rerank.managers import RerankManager
+from pygeai.core.rerank.models import RerankResponse, RerankResult, RerankMetaData, ApiVersion, BilledUnits
+from pygeai.core.common.exceptions import APIError
+
+
+class TestRerankManager(unittest.TestCase):
+    """
+    python -m unittest pygeai.tests.core.rerank.test_managers.TestRerankManager
+    """
+    def setUp(self):
+        self.manager = RerankManager(api_key="test_key", base_url="test_url", alias="test_alias")
+        self.manager._RerankManager__client = MagicMock()
+
+    def test_rerank_chunks_success(self):
+        mock_response = {
+            "id": "00a3bfe3-67e6-4aab-a19a-592bda2920a7",
+            "results": [
+                {"index": 2, "relevance_score": 0.8963332},
+                {"index": 0, "relevance_score": 0.17393301},
+                {"index": 1, "relevance_score": 0.08103136}
+            ],
+            "meta": {
+                "api_version": {"version": "2"},
+                "billed_units": {"search_units": 1}
+            }
+        }
+        self.manager._RerankManager__client.rerank_chunks.return_value = mock_response
+        with patch('pygeai.core.rerank.managers.RerankResponseMapper.map_to_rerank_response') as mock_mapper:
+            mock_mapper.return_value = RerankResponse(
+                id="00a3bfe3-67e6-4aab-a19a-592bda2920a7",
+                results=[
+                    RerankResult(index=2, relevance_score=0.8963332),
+                    RerankResult(index=0, relevance_score=0.17393301),
+                    RerankResult(index=1, relevance_score=0.08103136)
+                ],
+                meta=RerankMetaData(
+                    api_version=ApiVersion(version="2"),
+                    billed_units=BilledUnits(search_units=1)
+                )
+            )
+
+            result = self.manager.rerank_chunks(
+                query="What is the Capital of the United States?",
+                model="cohere/rerank-v3.5",
+                documents=[
+                    "Carson City is the capital city of the American state of Nevada.",
+                    "The Commonwealth of the Northern Mariana Islands is a group of islands in the Pacific Ocean. Its capital is Saipan.",
+                    "Washington, D.C. is the capital of the United States.",
+                    "Capital punishment has existed in the United States since before it was a country."
+                ],
+                top_n=3
+            )
+
+            self.assertIsInstance(result, RerankResponse)
+            self.assertEqual(result.id, "00a3bfe3-67e6-4aab-a19a-592bda2920a7")
+            self.assertEqual(len(result.results), 3)
+            self.assertEqual(result.results[0].index, 2)
+            self.assertEqual(result.results[0].relevance_score, 0.8963332)
+            self.manager._RerankManager__client.rerank_chunks.assert_called_once_with(
+                query="What is the Capital of the United States?",
+                model="cohere/rerank-v3.5",
+                documents=[
+                    "Carson City is the capital city of the American state of Nevada.",
+                    "The Commonwealth of the Northern Mariana Islands is a group of islands in the Pacific Ocean. Its capital is Saipan.",
+                    "Washington, D.C. is the capital of the United States.",
+                    "Capital punishment has existed in the United States since before it was a country."
+                ],
+                top_n=3
+            )
+            mock_mapper.assert_called_once_with(mock_response)
+
+    def test_rerank_chunks_error_response(self):
+        mock_error_response = {
+            "errors": [
+                {"id": 1001, "description": "Invalid request"},
+                {"id": 1002, "description": "Model not found"}
+            ]
+        }
+        self.manager._RerankManager__client.rerank_chunks.return_value = mock_error_response
+        with patch('pygeai.core.rerank.managers.ErrorHandler.has_errors', return_value=True):
+            with patch('pygeai.core.rerank.managers.ErrorHandler.extract_error', return_value="Invalid request; Model not found"):
+                with self.assertRaises(APIError) as context:
+                    self.manager.rerank_chunks(
+                        query="Invalid query",
+                        model="invalid/model",
+                        documents=["This should fail"],
+                        top_n=3
+                    )
+
+                self.assertIn("Error received while reranking chunks", str(context.exception))
+                self.manager._RerankManager__client.rerank_chunks.assert_called_once_with(
+                    query="Invalid query",
+                    model="invalid/model",
+                    documents=["This should fail"],
+                    top_n=3
+                )
+

pygeai/tests/core/rerank/test_mappers.py

@@ -0,0 +1,54 @@
+import unittest
+
+from pygeai.core.base.mappers import ErrorMapper
+from pygeai.core.base.responses import EmptyResponse, ErrorListResponse
+from pygeai.core.rerank.mappers import RerankResponseMapper
+from pygeai.core.rerank.models import RerankResponse, RerankResult, RerankMetaData
+
+
+class TestRerankResponseMapper(unittest.TestCase):
+    """
+    python -m unittest pygeai.tests.core.rerank.test_mappers.TestRerankResponseMapper
+    """
+
+    def test_map_to_rerank_response_success(self):
+        response_data = {
+            "id": "ad884ed0-d901-4025-ad98-a26174a52dda",
+            "results": [
+                {"index": 2, "relevance_score": 0.8963332},
+                {"index": 0, "relevance_score": 0.17393301},
+                {"index": 1, "relevance_score": 0.08103136}
+            ],
+            "meta": {
+                "api_version": {"version": "2"},
+                "billed_units": {"search_units": 1}
+            }
+        }
+
+        result = RerankResponseMapper.map_to_rerank_response(response_data)
+
+        self.assertIsInstance(result, RerankResponse)
+        self.assertEqual(result.id, "ad884ed0-d901-4025-ad98-a26174a52dda")
+        self.assertEqual(len(result.results), 3)
+        self.assertEqual(result.results[0].index, 2)
+        self.assertEqual(result.results[0].relevance_score, 0.8963332)
+        self.assertEqual(result.meta.api_version.version, "2")
+        self.assertEqual(result.meta.billed_units.search_units, 1)
+
+    def test_map_to_rerank_response_empty_results(self):
+        response_data = {
+            "id": "empty-results-id",
+            "results": [],
+            "meta": {
+                "api_version": {"version": "2"},
+                "billed_units": {"search_units": 1}
+            }
+        }
+
+        result = RerankResponseMapper.map_to_rerank_response(response_data)
+
+        self.assertIsInstance(result, RerankResponse)
+        self.assertEqual(result.id, "empty-results-id")
+        self.assertEqual(len(result.results), 0)
+        self.assertEqual(result.meta.api_version.version, "2")
+        self.assertEqual(result.meta.billed_units.search_units, 1)