pygeai 0.4.0b12__py3-none-any.whl → 0.6.0b3__py3-none-any.whl

pygeai/tests/evaluation/dataset/test_clients.py

@@ -0,0 +1,263 @@
+import unittest
+from unittest.mock import patch, MagicMock, mock_open
+import json
+from pathlib import Path
+
+from pygeai.evaluation.dataset.clients import EvaluationDatasetClient
+
+
+class TestEvaluationDatasetClient(unittest.TestCase):
+    """
+    python -m unittest pygeai.tests.evaluation.dataset.test_clients.TestEvaluationDatasetClient
+    """
+
+    def setUp(self):
+        self.client = EvaluationDatasetClient()
+        self.mock_response = MagicMock()
+        self.mock_response.json.return_value = {"id": "dataset-123", "status": "success"}
+
+    @patch('pygeai.core.services.rest.ApiService.get')
+    def test_list_datasets(self, mock_get):
+        mock_get.return_value = self.mock_response
+        result = self.client.list_datasets()
+        mock_get.assert_called_once()
+        self.assertEqual(result, {"id": "dataset-123", "status": "success"})
+
+    @patch('pygeai.core.services.rest.ApiService.post')
+    def test_create_dataset_minimal(self, mock_post):
+        mock_post.return_value = self.mock_response
+        result = self.client.create_dataset(
+            dataset_name="Test Dataset",
+            dataset_description="Test Description",
+            dataset_type="T"
+        )
+        mock_post.assert_called_once()
+        call_args = mock_post.call_args
+        self.assertEqual(call_args[1]['data']['dataSetName'], "Test Dataset")
+        self.assertEqual(call_args[1]['data']['dataSetDescription'], "Test Description")
+        self.assertEqual(call_args[1]['data']['dataSetType'], "T")
+        self.assertTrue(call_args[1]['data']['dataSetActive'])
+
+    @patch('pygeai.core.services.rest.ApiService.post')
+    def test_create_dataset_with_rows(self, mock_post):
+        mock_post.return_value = self.mock_response
+        rows = [{"dataSetRowInput": "input1", "dataSetRowExpectedAnswer": "answer1"}]
+        result = self.client.create_dataset(
+            dataset_name="Test Dataset",
+            dataset_description="Test Description",
+            dataset_type="T",
+            dataset_active=False,
+            rows=rows
+        )
+        mock_post.assert_called_once()
+        call_args = mock_post.call_args
+        self.assertEqual(call_args[1]['data']['rows'], rows)
+        self.assertFalse(call_args[1]['data']['dataSetActive'])
+
+    @patch('pathlib.Path.is_file')
+    @patch('pathlib.Path.open')
+    @patch('pygeai.core.services.rest.ApiService.post_file_binary')
+    def test_create_dataset_from_file_success(self, mock_post, mock_open_file, mock_is_file):
+        mock_is_file.return_value = True
+        mock_file = MagicMock()
+        mock_open_file.return_value = mock_file
+        mock_post.return_value = self.mock_response
+
+        result = self.client.create_dataset_from_file("/path/to/dataset.json")
+        
+        mock_is_file.assert_called_once()
+        mock_open_file.assert_called_once_with("rb")
+        mock_post.assert_called_once()
+        mock_file.close.assert_called_once()
+        self.assertEqual(result, {"id": "dataset-123", "status": "success"})
+
+    @patch('pathlib.Path.is_file')
+    def test_create_dataset_from_file_not_found(self, mock_is_file):
+        mock_is_file.return_value = False
+        with self.assertRaises(FileNotFoundError):
+            self.client.create_dataset_from_file("/path/to/nonexistent.json")
+
+    @patch('pygeai.core.services.rest.ApiService.get')
+    def test_get_dataset(self, mock_get):
+        mock_get.return_value = self.mock_response
+        result = self.client.get_dataset("dataset-123")
+        mock_get.assert_called_once()
+        self.assertEqual(result, {"id": "dataset-123", "status": "success"})
+
+    @patch('pygeai.core.services.rest.ApiService.put')
+    def test_update_dataset(self, mock_put):
+        mock_put.return_value = self.mock_response
+        rows = [{"dataSetRowInput": "updated"}]
+        result = self.client.update_dataset(
+            dataset_id="dataset-123",
+            dataset_name="Updated Dataset",
+            dataset_description="Updated Description",
+            dataset_type="E",
+            dataset_active=False,
+            rows=rows
+        )
+        mock_put.assert_called_once()
+        call_args = mock_put.call_args
+        self.assertEqual(call_args[1]['data']['dataSetName'], "Updated Dataset")
+        self.assertEqual(call_args[1]['data']['rows'], rows)
+
+    @patch('pygeai.core.services.rest.ApiService.delete')
+    def test_delete_dataset(self, mock_delete):
+        mock_delete.return_value = self.mock_response
+        result = self.client.delete_dataset("dataset-123")
+        mock_delete.assert_called_once()
+        self.assertEqual(result, {"id": "dataset-123", "status": "success"})
+
+    @patch('pygeai.core.services.rest.ApiService.post')
+    def test_create_dataset_row(self, mock_post):
+        mock_post.return_value = self.mock_response
+        row = {
+            "dataSetRowInput": "input",
+            "dataSetRowExpectedAnswer": "answer",
+            "dataSetRowContextDocument": "context"
+        }
+        result = self.client.create_dataset_row("dataset-123", row)
+        mock_post.assert_called_once()
+        call_args = mock_post.call_args
+        self.assertEqual(call_args[1]['data'], row)
+
+    @patch('pygeai.core.services.rest.ApiService.get')
+    def test_list_dataset_rows(self, mock_get):
+        mock_get.return_value = self.mock_response
+        result = self.client.list_dataset_rows("dataset-123")
+        mock_get.assert_called_once()
+        self.assertEqual(result, {"id": "dataset-123", "status": "success"})
+
+    @patch('pygeai.core.services.rest.ApiService.get')
+    def test_get_dataset_row(self, mock_get):
+        mock_get.return_value = self.mock_response
+        result = self.client.get_dataset_row("dataset-123", "row-456")
+        mock_get.assert_called_once()
+        self.assertEqual(result, {"id": "dataset-123", "status": "success"})
+
+    @patch('pygeai.core.services.rest.ApiService.put')
+    def test_update_dataset_row(self, mock_put):
+        mock_put.return_value = self.mock_response
+        row = {"dataSetRowInput": "updated input"}
+        result = self.client.update_dataset_row("dataset-123", "row-456", row)
+        mock_put.assert_called_once()
+
+    @patch('pygeai.core.services.rest.ApiService.delete')
+    def test_delete_dataset_row(self, mock_delete):
+        mock_delete.return_value = self.mock_response
+        result = self.client.delete_dataset_row("dataset-123", "row-456")
+        mock_delete.assert_called_once()
+
+    @patch('pygeai.core.services.rest.ApiService.post')
+    def test_create_dataset_row_expected_source(self, mock_post):
+        mock_post.return_value = self.mock_response
+        result = self.client.create_expected_source(
+            dataset_id="dataset-123",
+            dataset_row_id="row-456",
+            expected_source_name="source1",
+            expected_source_value="value1",
+            expected_source_extension="txt"
+        )
+        mock_post.assert_called_once()
+
+    @patch('pygeai.core.services.rest.ApiService.get')
+    def test_list_dataset_row_expected_sources(self, mock_get):
+        mock_get.return_value = self.mock_response
+        result = self.client.list_expected_sources("dataset-123", "row-456")
+        mock_get.assert_called_once()
+
+    @patch('pygeai.core.services.rest.ApiService.get')
+    def test_get_expected_source(self, mock_get):
+        mock_get.return_value = self.mock_response
+        result = self.client.get_expected_source("dataset-123", "row-456", "source-789")
+        mock_get.assert_called_once()
+
+    @patch('pygeai.core.services.rest.ApiService.put')
+    def test_update_expected_source(self, mock_put):
+        mock_put.return_value = self.mock_response
+        result = self.client.update_expected_source(
+            dataset_id="dataset-123",
+            dataset_row_id="row-456",
+            expected_source_id="source-789",
+            expected_source_name="updated source",
+            expected_source_value="updated value",
+            expected_source_extension="pdf"
+        )
+        mock_put.assert_called_once()
+
+    @patch('pygeai.core.services.rest.ApiService.delete')
+    def test_delete_expected_source(self, mock_delete):
+        mock_delete.return_value = self.mock_response
+        result = self.client.delete_expected_source("dataset-123", "row-456", "source-789")
+        mock_delete.assert_called_once()
+
+    @patch('pygeai.core.services.rest.ApiService.post')
+    def test_create_dataset_row_filter_variable(self, mock_post):
+        mock_post.return_value = self.mock_response
+        result = self.client.create_filter_variable(
+            dataset_id="dataset-123",
+            dataset_row_id="row-456",
+            metadata_type="type1",
+            filter_variable_key="key1",
+            filter_variable_value="value1",
+            filter_variable_operator="eq"
+        )
+        mock_post.assert_called_once()
+
+    @patch('pygeai.core.services.rest.ApiService.get')
+    def test_list_dataset_row_filter_variables(self, mock_get):
+        mock_get.return_value = self.mock_response
+        result = self.client.list_filter_variables("dataset-123", "row-456")
+        mock_get.assert_called_once()
+
+    @patch('pygeai.core.services.rest.ApiService.get')
+    def test_get_filter_variable(self, mock_get):
+        mock_get.return_value = self.mock_response
+        result = self.client.get_filter_variable("dataset-123", "row-456", "var-789")
+        mock_get.assert_called_once()
+
+    @patch('pygeai.core.services.rest.ApiService.put')
+    def test_update_filter_variable(self, mock_put):
+        mock_put.return_value = self.mock_response
+        result = self.client.update_filter_variable(
+            dataset_id="dataset-123",
+            dataset_row_id="row-456",
+            filter_variable_id="var-789",
+            metadata_type="type2",
+            filter_variable_key="updated key",
+            filter_variable_value="updated value",
+            filter_variable_operator="ne"
+        )
+        mock_put.assert_called_once()
+
+    @patch('pygeai.core.services.rest.ApiService.delete')
+    def test_delete_filter_variable(self, mock_delete):
+        mock_delete.return_value = self.mock_response
+        result = self.client.delete_filter_variable("dataset-123", "row-456", "var-789")
+        mock_delete.assert_called_once()
+
+    @patch('pathlib.Path.is_file')
+    @patch('pathlib.Path.open')
+    @patch('pygeai.core.services.rest.ApiService.post_file_binary')
+    def test_upload_dataset_rows_file_success(self, mock_post, mock_open_file, mock_is_file):
+        mock_is_file.return_value = True
+        mock_file = MagicMock()
+        mock_open_file.return_value = mock_file
+        mock_post.return_value = self.mock_response
+
+        result = self.client.upload_dataset_rows_file("dataset-123", "/path/to/rows.json")
+        
+        mock_is_file.assert_called_once()
+        mock_open_file.assert_called_once_with("rb")
+        mock_post.assert_called_once()
+        mock_file.close.assert_called_once()
+
+    @patch('pathlib.Path.is_file')
+    def test_upload_dataset_rows_file_not_found(self, mock_is_file):
+        mock_is_file.return_value = False
+        with self.assertRaises(FileNotFoundError):
+            self.client.upload_dataset_rows_file("dataset-123", "/path/to/nonexistent.json")
+
+
+if __name__ == '__main__':
+    unittest.main()

pygeai/tests/evaluation/plan/test_clients.py

@@ -0,0 +1,193 @@
+import unittest
+from unittest.mock import patch, MagicMock
+import json
+
+from pygeai.evaluation.plan.clients import EvaluationPlanClient
+
+
+class TestEvaluationPlanClient(unittest.TestCase):
+    """
+    python -m unittest pygeai.tests.evaluation.plan.test_clients.TestEvaluationPlanClient
+    """
+
+    def setUp(self):
+        self.client = EvaluationPlanClient()
+        self.mock_response = MagicMock()
+        self.mock_response.content = json.dumps({"id": "plan-123", "status": "success"})
+
+    @patch('pygeai.core.services.rest.ApiService.get')
+    def test_list_evaluation_plans(self, mock_get):
+        mock_get.return_value = self.mock_response
+        result = self.client.list_evaluation_plans()
+        mock_get.assert_called_once()
+        self.assertEqual(result, {"id": "plan-123", "status": "success"})
+
+    @patch('pygeai.core.services.rest.ApiService.post')
+    def test_create_evaluation_plan_minimal(self, mock_post):
+        mock_post.return_value = self.mock_response
+        result = self.client.create_evaluation_plan(
+            name="Test Plan",
+            type="TextPromptAssistant"
+        )
+        mock_post.assert_called_once()
+        call_args = mock_post.call_args
+        self.assertEqual(call_args[1]['data']['evaluationPlanName'], "Test Plan")
+        self.assertEqual(call_args[1]['data']['evaluationPlanType'], "TextPromptAssistant")
+
+    @patch('pygeai.core.services.rest.ApiService.post')
+    def test_create_evaluation_plan_with_assistant(self, mock_post):
+        mock_post.return_value = self.mock_response
+        result = self.client.create_evaluation_plan(
+            name="Test Plan",
+            type="TextPromptAssistant",
+            assistant_id="asst-123",
+            assistant_name="Test Assistant",
+            assistant_revision="1.0"
+        )
+        mock_post.assert_called_once()
+        call_args = mock_post.call_args
+        self.assertEqual(call_args[1]['data']['evaluationPlanAssistantId'], "asst-123")
+        self.assertEqual(call_args[1]['data']['evaluationPlanAssistantName'], "Test Assistant")
+        self.assertEqual(call_args[1]['data']['evaluationPlanAssistantRevision'], "1.0")
+
+    @patch('pygeai.core.services.rest.ApiService.post')
+    def test_create_evaluation_plan_with_all_params(self, mock_post):
+        mock_post.return_value = self.mock_response
+        system_metrics = [{"systemMetricId": "metric-1", "systemMetricWeight": 0.5}]
+        result = self.client.create_evaluation_plan(
+            name="Test Plan",
+            type="RAG Assistant",
+            profile_name="Test Profile",
+            dataset_id="dataset-123",
+            system_metrics=system_metrics
+        )
+        mock_post.assert_called_once()
+        call_args = mock_post.call_args
+        self.assertEqual(call_args[1]['data']['evaluationPlanProfileName'], "Test Profile")
+        self.assertEqual(call_args[1]['data']['dataSetId'], "dataset-123")
+        self.assertEqual(call_args[1]['data']['systemMetrics'], system_metrics)
+
+    @patch('pygeai.core.services.rest.ApiService.get')
+    def test_get_evaluation_plan(self, mock_get):
+        mock_get.return_value = self.mock_response
+        result = self.client.get_evaluation_plan("plan-123")
+        mock_get.assert_called_once()
+        self.assertEqual(result, {"id": "plan-123", "status": "success"})
+
+    @patch('pygeai.core.services.rest.ApiService.put')
+    def test_update_evaluation_plan_name_only(self, mock_put):
+        mock_put.return_value = self.mock_response
+        result = self.client.update_evaluation_plan(
+            evaluation_plan_id="plan-123",
+            name="Updated Plan"
+        )
+        mock_put.assert_called_once()
+        call_args = mock_put.call_args
+        self.assertEqual(call_args[1]['data']['evaluationPlanName'], "Updated Plan")
+        self.assertEqual(len(call_args[1]['data']), 1)
+
+    @patch('pygeai.core.services.rest.ApiService.put')
+    def test_update_evaluation_plan_multiple_fields(self, mock_put):
+        mock_put.return_value = self.mock_response
+        system_metrics = [{"systemMetricId": "metric-1", "systemMetricWeight": 0.8}]
+        result = self.client.update_evaluation_plan(
+            evaluation_plan_id="plan-123",
+            name="Updated Plan",
+            type="RAG Assistant",
+            assistant_id="asst-456",
+            assistant_name="Updated Assistant",
+            assistant_revision="2.0",
+            profile_name="Updated Profile",
+            dataset_id="dataset-456",
+            system_metrics=system_metrics
+        )
+        mock_put.assert_called_once()
+        call_args = mock_put.call_args
+        self.assertEqual(call_args[1]['data']['evaluationPlanName'], "Updated Plan")
+        self.assertEqual(call_args[1]['data']['evaluationPlanType'], "RAG Assistant")
+        self.assertEqual(call_args[1]['data']['evaluationPlanAssistantId'], "asst-456")
+        self.assertEqual(call_args[1]['data']['systemMetrics'], system_metrics)
+
+    @patch('pygeai.core.services.rest.ApiService.delete')
+    def test_delete_evaluation_plan(self, mock_delete):
+        mock_delete.return_value = self.mock_response
+        result = self.client.delete_evaluation_plan("plan-123")
+        mock_delete.assert_called_once()
+        self.assertEqual(result, {"id": "plan-123", "status": "success"})
+
+    @patch('pygeai.core.services.rest.ApiService.get')
+    def test_list_evaluation_plan_system_metrics(self, mock_get):
+        mock_get.return_value = self.mock_response
+        result = self.client.list_evaluation_plan_system_metrics("plan-123")
+        mock_get.assert_called_once()
+        self.assertEqual(result, {"id": "plan-123", "status": "success"})
+
+    @patch('pygeai.core.services.rest.ApiService.post')
+    def test_add_evaluation_plan_system_metric(self, mock_post):
+        mock_post.return_value = self.mock_response
+        result = self.client.add_evaluation_plan_system_metric(
+            evaluation_plan_id="plan-123",
+            system_metric_id="metric-456",
+            system_metric_weight=0.75
+        )
+        mock_post.assert_called_once()
+        call_args = mock_post.call_args
+        self.assertEqual(call_args[1]['data']['systemMetricId'], "metric-456")
+        self.assertEqual(call_args[1]['data']['systemMetricWeight'], 0.75)
+
+    @patch('pygeai.core.services.rest.ApiService.get')
+    def test_get_evaluation_plan_system_metric(self, mock_get):
+        mock_get.return_value = self.mock_response
+        result = self.client.get_evaluation_plan_system_metric(
+            evaluation_plan_id="plan-123",
+            system_metric_id="metric-456"
+        )
+        mock_get.assert_called_once()
+        self.assertEqual(result, {"id": "plan-123", "status": "success"})
+
+    @patch('pygeai.core.services.rest.ApiService.put')
+    def test_update_evaluation_plan_system_metric(self, mock_put):
+        mock_put.return_value = self.mock_response
+        result = self.client.update_evaluation_plan_system_metric(
+            evaluation_plan_id="plan-123",
+            system_metric_id="metric-456",
+            system_metric_weight=0.9
+        )
+        mock_put.assert_called_once()
+        call_args = mock_put.call_args
+        self.assertEqual(call_args[1]['data']['systemMetricWeight'], 0.9)
+
+    @patch('pygeai.core.services.rest.ApiService.delete')
+    def test_delete_evaluation_plan_system_metric(self, mock_delete):
+        mock_delete.return_value = self.mock_response
+        result = self.client.delete_evaluation_plan_system_metric(
+            evaluation_plan_id="plan-123",
+            system_metric_id="metric-456"
+        )
+        mock_delete.assert_called_once()
+        self.assertEqual(result, {"id": "plan-123", "status": "success"})
+
+    @patch('pygeai.core.services.rest.ApiService.get')
+    def test_list_system_metrics(self, mock_get):
+        mock_get.return_value = self.mock_response
+        result = self.client.list_system_metrics()
+        mock_get.assert_called_once()
+        self.assertEqual(result, {"id": "plan-123", "status": "success"})
+
+    @patch('pygeai.core.services.rest.ApiService.get')
+    def test_get_system_metric(self, mock_get):
+        mock_get.return_value = self.mock_response
+        result = self.client.get_system_metric("metric-456")
+        mock_get.assert_called_once()
+        self.assertEqual(result, {"id": "plan-123", "status": "success"})
+
+    @patch('pygeai.core.services.rest.ApiService.post')
+    def test_execute_evaluation_plan(self, mock_post):
+        mock_post.return_value = self.mock_response
+        result = self.client.execute_evaluation_plan("plan-123")
+        mock_post.assert_called_once()
+        self.assertEqual(result, {"id": "plan-123", "status": "success"})
+
+
+if __name__ == '__main__':
+    unittest.main()

pygeai/tests/evaluation/result/test_clients.py

@@ -0,0 +1,64 @@
+import unittest
+from unittest.mock import patch, MagicMock
+
+from pygeai.evaluation.result.clients import EvaluationResultClient
+
+
+class TestEvaluationResultClient(unittest.TestCase):
+    """
+    python -m unittest pygeai.tests.evaluation.result.test_clients.TestEvaluationResultClient
+    """
+
+    def setUp(self):
+        self.client = EvaluationResultClient()
+        self.mock_response = MagicMock()
+
+    @patch('pygeai.core.services.rest.ApiService.get')
+    def test_list_evaluation_results_success(self, mock_get):
+        self.mock_response.content = b'[{"id": "result-1", "status": "completed"}, {"id": "result-2", "status": "pending"}]'
+        mock_get.return_value = self.mock_response
+        
+        result = self.client.list_evaluation_results("plan-123")
+        
+        mock_get.assert_called_once()
+        self.assertEqual(len(result), 2)
+        self.assertEqual(result[0]["id"], "result-1")
+        self.assertEqual(result[1]["id"], "result-2")
+
+    @patch('pygeai.core.services.rest.ApiService.get')
+    def test_list_evaluation_results_empty(self, mock_get):
+        self.mock_response.content = b'[]'
+        mock_get.return_value = self.mock_response
+        
+        result = self.client.list_evaluation_results("plan-456")
+        
+        mock_get.assert_called_once()
+        self.assertEqual(result, [])
+
+    @patch('pygeai.core.services.rest.ApiService.get')
+    def test_get_evaluation_result_success(self, mock_get):
+        self.mock_response.content = b'{"id": "result-123", "status": "completed", "score": 0.95}'
+        mock_get.return_value = self.mock_response
+        
+        result = self.client.get_evaluation_result("result-123")
+        
+        mock_get.assert_called_once()
+        self.assertEqual(result["id"], "result-123")
+        self.assertEqual(result["status"], "completed")
+        self.assertEqual(result["score"], 0.95)
+
+    @patch('pygeai.core.services.rest.ApiService.get')
+    def test_get_evaluation_result_with_details(self, mock_get):
+        self.mock_response.content = b'{"id": "result-789", "status": "failed", "error": "Test error"}'
+        mock_get.return_value = self.mock_response
+        
+        result = self.client.get_evaluation_result("result-789")
+        
+        mock_get.assert_called_once()
+        self.assertEqual(result["id"], "result-789")
+        self.assertEqual(result["status"], "failed")
+        self.assertEqual(result["error"], "Test error")
+
+
+if __name__ == '__main__':
+    unittest.main()

pygeai/tests/integration/assistants/rag/test_create_rag.py

@@ -1,4 +1,5 @@
 from unittest import TestCase
+import uuid
 from pygeai.assistant.managers import AssistantManager
 from pygeai.core.models import WelcomeData, LlmSettings
 from pygeai.assistant.rag.models import (
@@ -10,10 +11,28 @@ from pygeai.assistant.rag.models import (
     RAGAssistant,
 )
 
-assistant_manager = AssistantManager(alias="beta")
 
 class TestAssistantCreateRagIntegration(TestCase):
 
+    def setUp(self):
+        """
+        Set up the test environment.
+        """
+        self.assistant_manager = AssistantManager()
+
+        self.new_rag = self.__load_rag()
+        self.created_rag: RAGAssistant = None
+
+
+    def tearDown(self):
+        """
+        Clean up after each test if necessary.
+        This can be used to delete the created tool
+        """
+        if isinstance(self.created_rag, RAGAssistant):
+            self.assistant_manager.delete_assistant(assistant_name=self.created_rag.name)            
+
+
     def __load_rag(self):
         llm_options = LlmSettings(
             cache=False,
@@ -57,7 +76,7 @@ class TestAssistantCreateRagIntegration(TestCase):
         )
 
         return RAGAssistant(
-            name="Testaasdfasdfsasdfdfsdfsdf123in123g12355",
+            name=str(uuid.uuid4()),
             description="Test Profile with WelcomeData",
             search_options=search_options,
             index_options=index_options,
@@ -67,6 +86,6 @@ class TestAssistantCreateRagIntegration(TestCase):
 
     def test_create_rag_assistant(self):
         rag_assistant = self.__load_rag()
-        response = assistant_manager.create_assistant(rag_assistant)
-        
-        self.assertIsInstance(response, RAGAssistant, "Failed to create RAG assistant")
+        self.created_rag = self.assistant_manager.create_assistant(rag_assistant)
+
+        self.assertIsInstance(self.created_rag, RAGAssistant, "Failed to create RAG assistant")

pygeai/tests/integration/chat/test_generate_image.py

@@ -6,7 +6,7 @@ chat_client: ChatClient
 class TestChatGenerateImageIntegration(TestCase):    
 
     def setUp(self):
-        self.chat_client = ChatClient(alias="beta")
+        self.chat_client = ChatClient()
         self.new_image = self.__load_image()
 
 
@@ -56,13 +56,9 @@ class TestChatGenerateImageIntegration(TestCase):
         self.new_image["model"] = ""
         created_image = self.__generate_image()
 
-        self.assertEqual(
-            created_image["error"]["code"], 400,
-            "Expected a 400 code for no model"
-        )
         self.assertEqual(
             created_image["error"]["message"],
-            "Invalid 'model' name. Must follow pattern {provider}/{modelName}", 
+            "LLM Provider NOT provided. Pass in the LLM provider you are trying to call", 
             "Expected an error message when no model is provided"
         )
 

pygeai/tests/integration/lab/agents/test_agents_list.py

@@ -13,7 +13,7 @@ class TestAILabAgentsListIntegration(TestCase):
     ]
 
     def setUp(self):
-        self.ai_lab_manager = AILabManager(alias="beta")
+        self.ai_lab_manager = AILabManager()
         self.filter_settings = FilterSettings(
             allow_external=False,
             allow_drafts=True,

pygeai/tests/integration/lab/agents/test_create_agent.py

@@ -12,7 +12,7 @@ class TestAILabCreateAgentIntegration(TestCase):
         """
         Set up the test environment.
         """
-        self.ai_lab_manager = AILabManager(alias="beta")
+        self.ai_lab_manager = AILabManager()
         self.new_agent = self.__load_agent()
         self.created_agent: Agent = None
 
@@ -51,7 +51,7 @@ class TestAILabCreateAgentIntegration(TestCase):
                     timeout=0,
                     sampling=Sampling(temperature=0.3, top_k=0, top_p=0)
                 ),
-                models=[Model(name="gpt-4-turbo-preview")]
+                models=[Model(name="gpt-4o")]
             )
         )
 
@@ -118,7 +118,7 @@ class TestAILabCreateAgentIntegration(TestCase):
                     max_tokens=1800,
                     timeout=0
                 ),
-                models=[Model(name="gpt-4-turbo-preview")]
+                models=[Model(name="gpt-4o")]
             )
         )
         self.created_agent = self.__create_agent()
@@ -141,7 +141,6 @@ class TestAILabCreateAgentIntegration(TestCase):
         self.assertEqual(prompt.inputs, agent.agent_data.prompt.inputs)
 
 
-    @unittest.skip("Skipping test due to mismatch in expected behavior")
     def test_create_agent_without_required_data(self):
         test_params = [ True, False ]
 
@@ -151,12 +150,19 @@ class TestAILabCreateAgentIntegration(TestCase):
                 self.new_agent = Agent(
                     name=str(uuid.uuid4())
                 )
-                with self.assertRaises(APIError) as exception:
-                    self.__create_agent(automatic_publish=auto_publish)
-
-                #TODO: Change validation error to a more specific one
-                self.assertIn("400",str(exception.exception))    
-                self.assertIn("Bad Request",str(exception.exception))
+                if auto_publish:
+                    
+                    with self.assertRaises(APIError) as exception:
+                        self.__create_agent(automatic_publish=auto_publish)
+
+                    #TODO: Change validation error to a more specific one
+                    self.assertIn(
+                        "A valid prompt is required. To be valid, it must provide clear instructions to the model",
+                        str(exception.exception)
+                    )   
+                else:
+                    created_agent = self.__create_agent(automatic_publish=auto_publish)
+                    self.assertTrue(isinstance(created_agent, Agent), "Expected a created agent")
 
 
     def test_create_agent_no_name(self):