Supoprt agent.create and agent.execute (#113)

Author: spillai

tests/conftest.py

@@ -98,6 +98,7 @@ def __init__(self, api_key=None, base_url=None):
             self.document = self.DocumentPredictions(self)
             self.audio = self.AudioPredictions(self)
             self.feedback = self.Feedback(self)
+            self.agent = self.Agent(self)
 
         class FineTuning:
             def __init__(self, client):
@@ -469,6 +470,126 @@ def get(self, request_id, limit=10, offset=0):
                     ],
                 )
 
+        class Agent:
+            def __init__(self, client):
+                self._client = client
+
+            def get(self, name=None, version=None, id=None):
+                from vlmrun.client.types import AgentInfo
+                from datetime import datetime
+
+                if id and name:
+                    raise ValueError("Only one of `id` or `name` can be provided.")
+                if not id and not name:
+                    raise ValueError("Either `id` or `name` must be provided.")
+
+                if id:
+                    agent_id = id
+                    agent_name = f"agent-{id}"
+                else:
+                    agent_id = f"agent-{name}-{version or 'latest'}"
+                    agent_name = name
+
+                return AgentInfo(
+                    id=agent_id,
+                    name=agent_name,
+                    version=version or "latest",
+                    description="Test agent description",
+                    prompt="Test agent prompt",
+                    json_schema={
+                        "type": "object",
+                        "properties": {"result": {"type": "string"}},
+                    },
+                    json_sample={"result": "test result"},
+                    created_at=datetime.fromisoformat("2024-01-01T00:00:00+00:00"),
+                    updated_at=datetime.fromisoformat("2024-01-01T00:00:00+00:00"),
+                    status="completed",
+                )
+
+            def list(self):
+                from vlmrun.client.types import AgentInfo
+                from datetime import datetime
+
+                return [
+                    AgentInfo(
+                        id="agent-1",
+                        name="test-agent-1",
+                        version="1.0.0",
+                        description="First test agent",
+                        prompt="Test prompt 1",
+                        json_schema={
+                            "type": "object",
+                            "properties": {"result": {"type": "string"}},
+                        },
+                        json_sample={"result": "test result 1"},
+                        created_at=datetime.fromisoformat("2024-01-01T00:00:00+00:00"),
+                        updated_at=datetime.fromisoformat("2024-01-01T00:00:00+00:00"),
+                        status="completed",
+                    ),
+                    AgentInfo(
+                        id="agent-2",
+                        name="test-agent-2",
+                        version="2.0.0",
+                        description="Second test agent",
+                        prompt="Test prompt 2",
+                        json_schema={
+                            "type": "object",
+                            "properties": {"output": {"type": "string"}},
+                        },
+                        json_sample={"output": "test output 2"},
+                        created_at=datetime.fromisoformat("2024-01-02T00:00:00+00:00"),
+                        updated_at=datetime.fromisoformat("2024-01-02T00:00:00+00:00"),
+                        status="completed",
+                    ),
+                ]
+
+            def create(self, config, name=None, inputs=None, callback_url=None):
+                from vlmrun.client.types import AgentCreationResponse
+                from datetime import datetime
+
+                if config.prompt is None:
+                    raise ValueError(
+                        "Prompt is not provided as a request parameter, please provide a prompt."
+                    )
+
+                return AgentCreationResponse(
+                    id=f"agent-{name or 'created'}",
+                    name=name or "created-agent",
+                    version="1.0.0",
+                    created_at=datetime.fromisoformat("2024-01-01T00:00:00+00:00"),
+                    updated_at=datetime.fromisoformat("2024-01-01T00:00:00+00:00"),
+                    status="pending",
+                )
+
+            def execute(
+                self,
+                name,
+                version=None,
+                inputs=None,
+                batch=True,
+                config=None,
+                metadata=None,
+                callback_url=None,
+            ):
+                from vlmrun.client.types import AgentExecutionResponse, CreditUsage
+                from datetime import datetime
+
+                if not batch:
+                    raise NotImplementedError(
+                        "Batch mode is required for agent execution"
+                    )
+
+                return AgentExecutionResponse(
+                    id=f"execution-{name}",
+                    name=name,
+                    version=version or "latest",
+                    status="completed",
+                    created_at=datetime.fromisoformat("2024-01-01T00:00:00+00:00"),
+                    completed_at=datetime.fromisoformat("2024-01-01T00:00:01+00:00"),
+                    response={"result": "execution result"},
+                    usage=CreditUsage(credits_used=50),
+                )
+
     monkeypatch.setattr("vlmrun.cli.cli.VLMRun", MockVLMRun)
     return MockVLMRun()
 

tests/test_agent.py

@@ -1,7 +1,15 @@
 """Tests for the Agent resource."""
 
+import pytest
 from datetime import datetime
-from vlmrun.client.types import AgentCreationResponse
+from vlmrun.client.types import (
+    AgentCreationResponse,
+    AgentCreationConfig,
+    AgentExecutionConfig,
+    AgentExecutionResponse,
+    AgentInfo,
+    CreditUsage,
+)
 
 
 class TestAgentCreationResponse:
@@ -13,74 +21,183 @@ def test_agent_creation_response_creation(self):
             "id": "test-agent-123",
             "name": "Test Agent",
             "version": "1.0.0",
-            "description": "A test agent for testing purposes",
             "created_at": datetime.now(),
             "updated_at": datetime.now(),
-            "output_json_schema": {
-                "type": "object",
-                "properties": {"result": {"type": "string"}},
-            },
-            "output_json_sample": {"result": "test result"},
-            "input_type": "document",
-            "input_json_schema": None,
+            "status": "completed",
         }
 
         response = AgentCreationResponse(**response_data)
 
         assert response.id == "test-agent-123"
         assert response.name == "Test Agent"
         assert response.version == "1.0.0"
-        assert response.description == "A test agent for testing purposes"
-        assert response.input_type == "document"
-        assert response.input_json_schema is None
-        assert "result" in response.output_json_schema["properties"]
-        assert response.output_json_sample["result"] == "test result"
-
-    def test_agent_creation_response_with_input_schema(self):
-        """Test creating an AgentCreationResponse with input schema."""
-        response_data = {
-            "id": "test-agent-456",
-            "name": "Test Agent with Input Schema",
-            "version": "2.0.0",
-            "description": "A test agent with input schema",
-            "created_at": datetime.now(),
-            "updated_at": datetime.now(),
-            "output_json_schema": {
+        assert response.status == "completed"
+
+
+class TestAgentMethods:
+    """Test the Agent API methods."""
+
+    def test_agent_get_by_name_and_version(self, mock_client):
+        """Test getting an agent by name and version."""
+        client = mock_client
+        response = client.agent.get(name="test-agent", version="1.0.0")
+
+        assert isinstance(response, AgentInfo)
+        assert response.name == "test-agent"
+        assert response.version == "1.0.0"
+        assert response.description == "Test agent description"
+        assert response.prompt == "Test agent prompt"
+        assert response.status == "completed"
+
+    def test_agent_get_by_id(self, mock_client):
+        """Test getting an agent by ID."""
+        client = mock_client
+        response = client.agent.get(id="agent-123")
+
+        assert isinstance(response, AgentInfo)
+        assert response.id == "agent-123"
+        assert response.name == "agent-agent-123"
+        assert response.version == "latest"
+
+    def test_agent_get_validation_error(self, mock_client):
+        """Test that get method validates input parameters."""
+        client = mock_client
+
+        with pytest.raises(
+            ValueError, match="Only one of `id` or `name` can be provided."
+        ):
+            client.agent.get(id="agent-123", name="test-agent")
+
+        with pytest.raises(ValueError, match="Either `id` or `name` must be provided."):
+            client.agent.get()
+
+    def test_agent_list(self, mock_client):
+        """Test listing all agents."""
+        client = mock_client
+        response = client.agent.list()
+
+        assert isinstance(response, list)
+        assert len(response) == 2
+        assert all(isinstance(agent, AgentInfo) for agent in response)
+        assert response[0].name == "test-agent-1"
+        assert response[1].name == "test-agent-2"
+
+    def test_agent_create(self, mock_client):
+        """Test creating an agent."""
+        client = mock_client
+        config = AgentCreationConfig(
+            prompt="Create a test agent",
+            json_schema={
+                "type": "object",
+                "properties": {"result": {"type": "string"}},
+            },
+        )
+
+        response = client.agent.create(
+            config=config, name="new-agent", inputs={"test": "input"}
+        )
+
+        assert isinstance(response, AgentCreationResponse)
+        assert response.name == "new-agent"
+        assert response.version == "1.0.0"
+        assert response.status == "pending"
+
+    def test_agent_create_validation_error(self, mock_client):
+        """Test that create method validates config has prompt."""
+        client = mock_client
+        config = AgentCreationConfig()
+
+        with pytest.raises(
+            ValueError,
+            match="Prompt is not provided as a request parameter, please provide a prompt.",
+        ):
+            client.agent.create(config=config)
+
+    def test_agent_execute(self, mock_client):
+        """Test executing an agent."""
+        client = mock_client
+        config = AgentExecutionConfig(
+            prompt="Execute this agent",
+            json_schema={
                 "type": "object",
                 "properties": {"output": {"type": "string"}},
             },
-            "output_json_sample": {"output": "test output"},
-            "input_type": "mixed",
-            "input_json_schema": {
+        )
+
+        response = client.agent.execute(
+            name="test-agent",
+            version="1.0.0",
+            inputs={"input": "test data"},
+            config=config,
+        )
+
+        assert isinstance(response, AgentExecutionResponse)
+        assert response.name == "test-agent"
+        assert response.version == "1.0.0"
+        assert response.status == "completed"
+        assert response.response == {"result": "execution result"}
+        assert isinstance(response.usage, CreditUsage)
+        assert response.usage.credits_used == 50
+
+    def test_agent_execute_without_version(self, mock_client):
+        """Test executing an agent without specifying version."""
+        client = mock_client
+
+        response = client.agent.execute(name="test-agent")
+
+        assert isinstance(response, AgentExecutionResponse)
+        assert response.name == "test-agent"
+        assert response.version == "latest"
+
+    def test_agent_execute_batch_mode_required(self, mock_client):
+        """Test that execute method requires batch mode."""
+        client = mock_client
+
+        with pytest.raises(
+            NotImplementedError, match="Batch mode is required for agent execution"
+        ):
+            client.agent.execute(name="test-agent", batch=False)
+
+
+class TestAgentConfigModels:
+    """Test the Agent configuration models."""
+
+    def test_agent_creation_config(self):
+        """Test AgentCreationConfig model."""
+        config = AgentCreationConfig(
+            prompt="Test prompt",
+            json_schema={
+                "type": "object",
+                "properties": {"result": {"type": "string"}},
+            },
+        )
+
+        assert config.prompt == "Test prompt"
+        assert config.json_schema is not None
+        assert config.response_model is None
+
+    def test_agent_execution_config(self):
+        """Test AgentExecutionConfig model."""
+        config = AgentExecutionConfig(
+            prompt="Execute prompt",
+            json_schema={
                 "type": "object",
-                "properties": {"input": {"type": "string"}},
+                "properties": {"output": {"type": "string"}},
             },
-        }
+        )
 
-        response = AgentCreationResponse(**response_data)
+        assert config.prompt == "Execute prompt"
+        assert config.json_schema is not None
+        assert config.response_model is None
+
+    def test_config_validation_error(self):
+        """Test that config validates response_model and json_schema are not both provided."""
+        from pydantic import BaseModel
+
+        class TestModel(BaseModel):
+            result: str
 
-        assert response.id == "test-agent-456"
-        assert response.input_type == "mixed"
-        assert response.input_json_schema is not None
-        assert "input" in response.input_json_schema["properties"]
-
-    def test_agent_creation_response_input_types(self):
-        """Test that all input types are accepted."""
-        valid_input_types = ["text", "document", "image", "video", "audio", "mixed"]
-
-        for input_type in valid_input_types:
-            response_data = {
-                "id": f"test-agent-{input_type}",
-                "name": f"Test Agent {input_type}",
-                "version": "1.0.0",
-                "description": f"A test agent for {input_type}",
-                "created_at": datetime.now(),
-                "updated_at": datetime.now(),
-                "output_json_schema": {"type": "object", "properties": {}},
-                "output_json_sample": {},
-                "input_type": input_type,
-                "input_json_schema": None,
-            }
-
-            response = AgentCreationResponse(**response_data)
-            assert response.input_type == input_type
+        with pytest.raises(ValueError, match="cannot be used together"):
+            AgentCreationConfig(
+                prompt="Test", response_model=TestModel, json_schema={"type": "object"}
+            )