wip

Author: harel-tussi

ai21/clients/studio/resources/agents/__init__.py

@@ -3,7 +3,7 @@
 
 __all__ = [
     "Agents",
-    "AsyncAgents", 
+    "AsyncAgents",
     "AgentRun",
     "AsyncAgentRun",
-]
+]

ai21/clients/studio/resources/agents/agents.py

@@ -24,7 +24,7 @@ def create(self, *, request: CreateAgentRequest) -> Agent:
         # Map agent_type to assistant_type for API compatibility
         if "agent_type" in body:
             body["assistant_type"] = body.pop("agent_type")
-        
+
         result = self._post(path=f"/{self._module_name}", body=body, response_cls=Agent)
         assert result is not None  # response_cls is provided, so result should never be None
         return result
@@ -65,7 +65,7 @@ async def create(self, *, request: CreateAgentRequest) -> Agent:
         body = request.model_dump(exclude_none=True)
         if "agent_type" in body:
             body["assistant_type"] = body.pop("agent_type")
-        
+
         result = await self._post(path=f"/{self._module_name}", body=body, response_cls=Agent)
         assert result is not None  # response_cls is provided, so result should never be None
         return result
@@ -93,4 +93,4 @@ async def delete(self, agent_id: str) -> DeleteAgentResponse:
         """Delete an agent"""
         result = await self._delete(path=f"/{self._module_name}/{agent_id}", response_cls=DeleteAgentResponse)
         assert result is not None  # response_cls is provided, so result should never be None
-        return result
+        return result

ai21/models/agents/__init__.py

@@ -24,4 +24,4 @@
     "RunAgentRequest",
     "RunResponse",
     "Visibility",
-]
+]

ai21/models/agents/agent.py

@@ -105,4 +105,4 @@ class RunResponse(AI21BaseModel):
     status: str  # "completed", "failed", "in_progress", "requires_action"
     result: Optional[Any] = None
     data_sources: Optional[Dict[str, Any]] = None
-    requirements_result: Optional[Dict[str, Any]] = None
+    requirements_result: Optional[Dict[str, Any]] = None

tests/integration_tests/clients/studio/test_agents.py

@@ -8,25 +8,25 @@
 async def test_agents_create_and_delete():
     """Test creating and deleting an agent"""
     client = AsyncAI21Client()
-    
+
     # Create an agent
     create_request = CreateAgentRequest(
         name="Test Agent for Integration",
         description="This is a test agent created by integration tests",
         budget=BudgetLevel.LOW,
         agent_type=AgentType.DEFAULT,
     )
-    
+
     agent = await client.beta.agents.create(request=create_request)
-    
+
     # Verify agent was created
     assert agent.id is not None
     assert agent.name == "Test Agent for Integration"
     assert agent.description == "This is a test agent created by integration tests"
     assert agent.budget == BudgetLevel.LOW
     assert agent.agent_type == AgentType.DEFAULT
     assert agent.object == "agent"
-    
+
     # Clean up - delete the agent
     delete_response = await client.beta.agents.delete(agent.id)
     assert delete_response.deleted is True
@@ -37,39 +37,39 @@ async def test_agents_create_and_delete():
 async def test_agents_crud_operations():
     """Test full CRUD operations for agents"""
     client = AsyncAI21Client()
-    
+
     # Create
     create_request = CreateAgentRequest(
         name="CRUD Test Agent",
         description="Testing CRUD operations",
         budget=BudgetLevel.MEDIUM,
     )
-    
+
     agent = await client.beta.agents.create(request=create_request)
     agent_id = agent.id
-    
+
     try:
         # Read - Get specific agent
         retrieved_agent = await client.beta.agents.get(agent_id)
         assert retrieved_agent.id == agent_id
         assert retrieved_agent.name == "CRUD Test Agent"
-        
+
         # Read - List agents (should include our agent)
         agents_list = await client.beta.agents.list()
         agent_ids = [a.id for a in agents_list.results]
         assert agent_id in agent_ids
-        
+
         # Update
         modify_request = ModifyAgentRequest(
             name="Modified CRUD Test Agent",
             description="Updated description for testing",
         )
-        
+
         updated_agent = await client.beta.agents.modify(agent_id, request=modify_request)
         assert updated_agent.id == agent_id
         assert updated_agent.name == "Modified CRUD Test Agent"
         assert updated_agent.description == "Updated description for testing"
-        
+
     finally:
         # Delete
         delete_response = await client.beta.agents.delete(agent_id)
@@ -80,38 +80,38 @@ async def test_agents_crud_operations():
 async def test_agent_run_basic():
     """Test running an agent with basic input"""
     client = AsyncAI21Client()
-    
+
     # Create a test agent
     create_request = CreateAgentRequest(
         name="Test Run Agent",
         description="Agent for testing runs",
         budget=BudgetLevel.LOW,
     )
-    
+
     agent = await client.beta.agents.create(request=create_request)
     agent_id = agent.id
-    
+
     try:
         # Test agent run
         input_messages = [{"role": "user", "content": "What is 2+2?"}]
-        
+
         run_response = await client.beta.agents.runs.create_and_poll(
-            agent_id, 
+            agent_id,
             input=input_messages,
             poll_timeout_sec=120,  # 2 minutes timeout
         )
-        
+
         assert run_response.id is not None
         assert run_response.status in ["completed", "failed"]
-        
+
         if run_response.status == "completed":
             assert run_response.result is not None
-            
+
         # Test retrieving the run
         retrieved_run = await client.beta.agents.runs.retrieve(str(run_response.id))
         assert retrieved_run.id == run_response.id
         assert retrieved_run.status == run_response.status
-        
+
     finally:
         # Clean up
         await client.beta.agents.delete(agent_id)
@@ -121,21 +121,21 @@ async def test_agent_run_basic():
 async def test_agent_run_with_options():
     """Test running an agent with various options"""
     client = AsyncAI21Client()
-    
+
     # Create a test agent
     create_request = CreateAgentRequest(
         name="Test Options Agent",
         description="Agent for testing run options",
         budget=BudgetLevel.MEDIUM,
     )
-    
+
     agent = await client.beta.agents.create(request=create_request)
     agent_id = agent.id
-    
+
     try:
         # Test agent run with options
         input_messages = [{"role": "user", "content": "Tell me about AI"}]
-        
+
         run_response = await client.beta.agents.runs.create_and_poll(
             agent_id,
             input=input_messages,
@@ -144,33 +144,33 @@ async def test_agent_run_with_options():
             response_language="english",
             poll_timeout_sec=120,
         )
-        
+
         assert run_response.id is not None
         assert run_response.status in ["completed", "failed"]
-        
+
         if run_response.status == "completed":
             assert run_response.result is not None
-            
+
     finally:
         # Clean up
         await client.beta.agents.delete(agent_id)
 
 
-@pytest.mark.asyncio 
+@pytest.mark.asyncio
 async def test_agents_list_empty_or_populated():
     """Test listing agents when there may be none or some"""
     client = AsyncAI21Client()
-    
+
     # List all agents
     agents_list = await client.beta.agents.list()
-    
+
     # Should return a valid response
-    assert hasattr(agents_list, 'results')
+    assert hasattr(agents_list, "results")
     assert isinstance(agents_list.results, list)
-    
+
     # Each result should be a proper Agent object
     for agent in agents_list.results:
-        assert hasattr(agent, 'id')
-        assert hasattr(agent, 'name')
-        assert hasattr(agent, 'created_at')
-        assert hasattr(agent, 'updated_at')
+        assert hasattr(agent, "id")
+        assert hasattr(agent, "name")
+        assert hasattr(agent, "created_at")
+        assert hasattr(agent, "updated_at")