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

pygeai/tests/lab/agents/test_clients.py

@@ -15,11 +15,12 @@ class TestAgentClient(unittest.TestCase):
     """
 
     def setUp(self):
-        self.client = AgentClient()
         self.project_id = "test-project-id"
+        self.client = AgentClient(api_key="test_key", base_url="https://test.url", project_id=self.project_id)
         self.agent_id = "test-agent-id"
         self.agent_data_prompt = {"instructions": "Do this task"}
         self.agent_data_llm_config = {"maxTokens": 100, "timeout": 30}
+        self.agent_data_strategy_name = "default"
         self.agent_data_models = [{"name": "gpt-4o"}]
         self.agent_data_resource_pools = [{"name": "pool1", "tools": [{"name": "tool1"}]}]
 
@@ -30,7 +31,6 @@ class TestAgentClient(unittest.TestCase):
         mock_response.json.return_value = expected_response
 
         result = self.client.list_agents(
-            project_id=self.project_id,
             status="active",
             start=0,
             count=10,
@@ -63,7 +63,7 @@ class TestAgentClient(unittest.TestCase):
         mock_response.text = "Invalid JSON response"
 
         with self.assertRaises(InvalidAPIResponseException) as context:
-            self.client.list_agents(project_id=self.project_id)
+            self.client.list_agents()
 
         self.assertEqual(str(context.exception), f"Unable to list agents for project {self.project_id}: Invalid JSON response")
         mock_get.assert_called_once()
@@ -75,7 +75,6 @@ class TestAgentClient(unittest.TestCase):
         mock_response.json.return_value = expected_response
 
         result = self.client.create_agent(
-            project_id=self.project_id,
             name="Test Agent",
             access_scope="private",
             public_name="test-agent",
@@ -84,6 +83,7 @@ class TestAgentClient(unittest.TestCase):
             description="Agent Description",
             agent_data_prompt=self.agent_data_prompt,
             agent_data_llm_config=self.agent_data_llm_config,
+            agent_data_strategy_name=self.agent_data_strategy_name,
             agent_data_models=self.agent_data_models,
             agent_data_resource_pools=self.agent_data_resource_pools,
             automatic_publish=True
@@ -116,7 +116,6 @@ class TestAgentClient(unittest.TestCase):
         mock_response.json.return_value = expected_response
 
         result = self.client.create_agent(
-            project_id=self.project_id,
             name="Test Agent",
             access_scope="private",
             public_name="test-agent",
@@ -125,6 +124,7 @@ class TestAgentClient(unittest.TestCase):
             description="Agent Description",
             agent_data_prompt=self.agent_data_prompt,
             agent_data_llm_config=self.agent_data_llm_config,
+            agent_data_strategy_name=self.agent_data_strategy_name,
             agent_data_models=self.agent_data_models,
             automatic_publish=False
         )
@@ -137,7 +137,6 @@ class TestAgentClient(unittest.TestCase):
     def test_create_agent_invalid_access_scope(self):
         with self.assertRaises(ValueError) as context:
             self.client.create_agent(
-                project_id=self.project_id,
                 name="Test Agent",
                 access_scope="invalid_scope",
                 public_name="test-agent",
@@ -146,6 +145,7 @@ class TestAgentClient(unittest.TestCase):
                 description="Agent Description",
                 agent_data_prompt=self.agent_data_prompt,
                 agent_data_llm_config=self.agent_data_llm_config,
+            agent_data_strategy_name=self.agent_data_strategy_name,
                 agent_data_models=self.agent_data_models
             )
         self.assertEqual(str(context.exception), f"Access scope must be one of {', '.join(VALID_ACCESS_SCOPES)}.")
@@ -159,7 +159,6 @@ class TestAgentClient(unittest.TestCase):
 
         with self.assertRaises(InvalidAPIResponseException) as context:
             self.client.create_agent(
-                project_id=self.project_id,
                 name="Test Agent",
                 access_scope="private",
                 public_name="test-agent",
@@ -168,6 +167,7 @@ class TestAgentClient(unittest.TestCase):
                 description="Agent Description",
                 agent_data_prompt=self.agent_data_prompt,
                 agent_data_llm_config=self.agent_data_llm_config,
+            agent_data_strategy_name=self.agent_data_strategy_name,
                 agent_data_models=self.agent_data_models
             )
 
@@ -181,7 +181,6 @@ class TestAgentClient(unittest.TestCase):
         mock_response.json.return_value = expected_response
 
         result = self.client.get_agent(
-            project_id=self.project_id,
             agent_id=self.agent_id,
             revision="1",
             version=2,
@@ -201,14 +200,9 @@ class TestAgentClient(unittest.TestCase):
         headers = mock_get.call_args[1]['headers']
         self.assertEqual(headers['ProjectId'], self.project_id)
 
-    def test_get_agent_missing_project_id(self):
-        with self.assertRaises(MissingRequirementException) as context:
-            self.client.get_agent(project_id="", agent_id=self.agent_id)
-        self.assertEqual(str(context.exception), "Cannot retrieve agent without specifying project_id")
-
     def test_get_agent_missing_agent_id(self):
         with self.assertRaises(MissingRequirementException) as context:
-            self.client.get_agent(project_id=self.project_id, agent_id="")
+            self.client.get_agent(agent_id="")
         self.assertEqual(str(context.exception), "agent_id must be specified in order to retrieve the agent")
 
     @patch("pygeai.core.services.rest.ApiService.get")
@@ -219,7 +213,7 @@ class TestAgentClient(unittest.TestCase):
         mock_response.text = "Invalid JSON response"
 
         with self.assertRaises(InvalidAPIResponseException) as context:
-            self.client.get_agent(project_id=self.project_id, agent_id=self.agent_id)
+            self.client.get_agent(agent_id=self.agent_id)
 
         self.assertEqual(str(context.exception), f"Unable to retrieve agent {self.agent_id} for project {self.project_id}: Invalid JSON response")
         mock_get.assert_called_once()
@@ -231,7 +225,6 @@ class TestAgentClient(unittest.TestCase):
         mock_response.json.return_value = expected_response
 
         result = self.client.create_sharing_link(
-            project_id=self.project_id,
             agent_id=self.agent_id
         )
 
@@ -244,14 +237,9 @@ class TestAgentClient(unittest.TestCase):
         headers = mock_get.call_args[1]['headers']
         self.assertEqual(headers['ProjectId'], self.project_id)
 
-    def test_create_sharing_link_missing_project_id(self):
-        with self.assertRaises(MissingRequirementException) as context:
-            self.client.create_sharing_link(project_id="", agent_id=self.agent_id)
-        self.assertEqual(str(context.exception), "Cannot create sharing link without specifying project_id")
-
     def test_create_sharing_link_missing_agent_id(self):
         with self.assertRaises(MissingRequirementException) as context:
-            self.client.create_sharing_link(project_id=self.project_id, agent_id="")
+            self.client.create_sharing_link(agent_id="")
         self.assertEqual(str(context.exception), "agent_id must be specified in order to create sharing link")
 
     @patch("pygeai.core.services.rest.ApiService.get")
@@ -263,7 +251,6 @@ class TestAgentClient(unittest.TestCase):
 
         with self.assertRaises(InvalidAPIResponseException) as context:
             self.client.create_sharing_link(
-                project_id=self.project_id,
                 agent_id=self.agent_id
             )
 
@@ -278,7 +265,6 @@ class TestAgentClient(unittest.TestCase):
         mock_response.json.return_value = expected_response
 
         result = self.client.publish_agent_revision(
-            project_id=self.project_id,
             agent_id=self.agent_id,
             revision=revision
         )
@@ -302,7 +288,6 @@ class TestAgentClient(unittest.TestCase):
 
         with self.assertRaises(InvalidAPIResponseException) as context:
             self.client.publish_agent_revision(
-                project_id=self.project_id,
                 agent_id=self.agent_id,
                 revision=revision
             )
@@ -312,12 +297,11 @@ class TestAgentClient(unittest.TestCase):
 
     @patch("pygeai.core.services.rest.ApiService.delete")
     def test_delete_agent_success(self, mock_delete):
-        expected_response = {"status": "deleted"}
+        expected_response = {}
         mock_response = mock_delete.return_value
-        mock_response.json.return_value = expected_response
+        mock_response.status_code = 204
 
         result = self.client.delete_agent(
-            project_id=self.project_id,
             agent_id=self.agent_id
         )
 
@@ -339,7 +323,6 @@ class TestAgentClient(unittest.TestCase):
 
         with self.assertRaises(InvalidAPIResponseException) as context:
             self.client.delete_agent(
-                project_id=self.project_id,
                 agent_id=self.agent_id
             )
 
@@ -353,7 +336,6 @@ class TestAgentClient(unittest.TestCase):
         mock_response.json.return_value = expected_response
 
         result = self.client.update_agent(
-            project_id=self.project_id,
             agent_id=self.agent_id,
             name="Updated Agent",
             access_scope="public",
@@ -363,6 +345,7 @@ class TestAgentClient(unittest.TestCase):
             description="Updated Description",
             agent_data_prompt=self.agent_data_prompt,
             agent_data_llm_config=self.agent_data_llm_config,
+            agent_data_strategy_name=self.agent_data_strategy_name,
             agent_data_models=self.agent_data_models,
             agent_data_resource_pools=self.agent_data_resource_pools,
             automatic_publish=True,
@@ -396,7 +379,6 @@ class TestAgentClient(unittest.TestCase):
         mock_response.json.return_value = expected_response
 
         result = self.client.update_agent(
-            project_id=self.project_id,
             agent_id=self.agent_id,
             name="Upserted Agent",
             access_scope="private",
@@ -406,6 +388,7 @@ class TestAgentClient(unittest.TestCase):
             description="Upserted Description",
             agent_data_prompt=self.agent_data_prompt,
             agent_data_llm_config=self.agent_data_llm_config,
+            agent_data_strategy_name=self.agent_data_strategy_name,
             agent_data_models=self.agent_data_models,
             automatic_publish=False,
             upsert=True
@@ -425,7 +408,6 @@ class TestAgentClient(unittest.TestCase):
         mock_response.json.return_value = expected_response
 
         result = self.client.update_agent(
-            project_id=self.project_id,
             agent_id=self.agent_id,
             name="Updated Agent No Pools",
             access_scope="private",
@@ -435,6 +417,7 @@ class TestAgentClient(unittest.TestCase):
             description="Updated Description",
             agent_data_prompt=self.agent_data_prompt,
             agent_data_llm_config=self.agent_data_llm_config,
+            agent_data_strategy_name=self.agent_data_strategy_name,
             agent_data_models=self.agent_data_models,
             automatic_publish=False,
             upsert=False
@@ -448,7 +431,6 @@ class TestAgentClient(unittest.TestCase):
     def test_update_agent_invalid_access_scope(self):
         with self.assertRaises(ValueError) as context:
             self.client.update_agent(
-                project_id=self.project_id,
                 agent_id=self.agent_id,
                 name="Updated Agent",
                 access_scope="invalid_scope",
@@ -458,6 +440,7 @@ class TestAgentClient(unittest.TestCase):
                 description="Updated Description",
                 agent_data_prompt=self.agent_data_prompt,
                 agent_data_llm_config=self.agent_data_llm_config,
+            agent_data_strategy_name=self.agent_data_strategy_name,
                 agent_data_models=self.agent_data_models
             )
         self.assertEqual(str(context.exception), f"Access scope must be one of {', '.join(VALID_ACCESS_SCOPES)}.")
@@ -471,7 +454,6 @@ class TestAgentClient(unittest.TestCase):
 
         with self.assertRaises(InvalidAPIResponseException) as context:
             self.client.update_agent(
-                project_id=self.project_id,
                 agent_id=self.agent_id,
                 name="Updated Agent",
                 access_scope="private",
@@ -481,6 +463,7 @@ class TestAgentClient(unittest.TestCase):
                 description="Updated Description",
                 agent_data_prompt=self.agent_data_prompt,
                 agent_data_llm_config=self.agent_data_llm_config,
+            agent_data_strategy_name=self.agent_data_strategy_name,
                 agent_data_models=self.agent_data_models
             )