wip: reimplement model tests

Signed-off-by: Diwank Tomer <[email protected]>

agents-api/agents_api/common/utils/types.py

@@ -0,0 +1,22 @@
+from typing import Type
+
+from beartype.vale import Is
+from beartype.vale._core._valecore import BeartypeValidator
+from pydantic import BaseModel
+
+
+def dict_like(pydantic_model_class: Type[BaseModel]) -> BeartypeValidator:
+    required_fields_set: set[str] = set(
+        [
+            field
+            for field, info in pydantic_model_class.model_fields.items()
+            if info.is_required()
+        ]
+    )
+
+    validator = Is[
+        lambda x: isinstance(x, pydantic_model_class)
+        or required_fields_set.issubset(set(x.keys()))
+    ]
+
+    return validator

agents-api/agents_api/models/agent/test_agent_queries.py

@@ -1,11 +1,12 @@
 # Tests for agent queries
 from uuid import uuid4
 
-from cozo_migrate.api import init, apply
+from cozo_migrate.api import apply, init
 from pycozo import Client
 from ward import test
 
 from agents_api.autogen.openapi_model import Agent
+
 from .create_agent import create_agent
 from .delete_agent import delete_agent
 from .get_agent import get_agent
@@ -69,9 +70,7 @@ def _():
     agent_id = uuid4()
     developer_id = uuid4()
 
-    result = get_agent(
-        agent_id=agent_id, developer_id=developer_id, client=client
-    )
+    result = get_agent(agent_id=agent_id, developer_id=developer_id, client=client)
 
     assert result is None
 
@@ -94,9 +93,7 @@ def _():
         client=client,
     )
 
-    result = get_agent(
-        agent_id=agent_id, developer_id=developer_id, client=client
-    )
+    result = get_agent(agent_id=agent_id, developer_id=developer_id, client=client)
 
     assert result is not None
     assert isinstance(result, Agent)
@@ -122,14 +119,10 @@ def _():
     )
 
     # Delete the agent
-    delete_agent(
-        agent_id=agent_id, developer_id=developer_id, client=client
-    )
+    delete_agent(agent_id=agent_id, developer_id=developer_id, client=client)
 
     # Check that the agent is deleted
-    result = get_agent(
-        agent_id=agent_id, developer_id=developer_id, client=client
-    )
+    result = get_agent(agent_id=agent_id, developer_id=developer_id, client=client)
 
     assert result is None
 

agents-api/agents_api/models/docs/test_docs_queries.py

@@ -1,15 +1,15 @@
 # Tests for entry queries
 from uuid import uuid4
 
-from cozo_migrate.api import init, apply
+from cozo_migrate.api import apply, init
 from pycozo import Client
 from ward import test
 
 from .create_doc import create_doc
 from .delete_doc import delete_doc
+from .embed_snippets import embed_snippets
 from .get_doc import get_doc
 from .list_docs import list_docs
-from .embed_snippets import embed_snippets
 from .search_docs import search_docs_by_embedding
 
 EMBEDDING_SIZE: int = 1024

agents-api/agents_api/models/entry/test_entry_queries.py

@@ -6,22 +6,20 @@
 # Tests for entry queries
 from uuid import uuid4
 
-from cozo_migrate.api import init, apply
+from cozo_migrate.api import apply, init
 from pycozo import Client
 from ward import test
 
-from ...autogen.openapi_model import FunctionDef
-from ...common.protocol.entries import Entry
-from ..docs.create_docs import create_doc
-from ..docs.embed_snippets import embed_snippets
+from ...autogen.openapi_model import Entry, FunctionDef
 from ..agent.create_agent import create_agent
+from ..docs.create_doc import create_doc
+from ..docs.embed_snippets import embed_snippets
 from ..session.create_session import create_session
-from ..tools.create_tools import create_tool
-from ..tools.embed_tools import embed_tools
+from ..tools.create_tools import create_tools
 from ..user.create_user import create_user
 from .create_entries import create_entries
-from .list_entries import list_entries
 from .get_history import get_history
+from .list_entries import list_entries
 
 MODEL = "julep-ai/samantha-1-turbo"
 
@@ -174,14 +172,16 @@ def _():
             },
             client=client,
         ),
-        create_tool(
+        create_tools(
             developer_id=developer_id,
             agent_id=agent_id,
-            data={
-                "name": test_function.name,
-                "description": test_function.description,
-                "parameters": test_function.parameters,
-            },
+            data=[
+                {
+                    "name": test_function.name,
+                    "description": test_function.description,
+                    "parameters": test_function.parameters,
+                }
+            ],
             client=client,
         ),
         create_doc(
@@ -206,13 +206,6 @@ def _():
             },
             client=client,
         ),
-        embed_tools(
-            developer_id=developer_id,
-            agent_id=agent_id,
-            tool_ids=[tool_id],
-            embeddings=[[1.0] * 768],
-            client=client,
-        ),
         embed_snippets(
             developer_id=developer_id,
             doc_id=agent_doc_id,

agents-api/agents_api/models/execution/create_execution.py

@@ -1,3 +1,4 @@
+from typing import Annotated
 from uuid import UUID, uuid4
 
 from beartype import beartype
@@ -7,6 +8,7 @@
 
 from ...autogen.openapi_model import CreateExecutionRequest, Execution
 from ...common.utils.cozo import cozo_process_mutate_data
+from ...common.utils.types import dict_like
 from ..utils import (
     cozo_query,
     partialclass,
@@ -36,16 +38,20 @@ def create_execution(
     developer_id: UUID,
     task_id: UUID,
     execution_id: UUID | None = None,
-    data: CreateExecutionRequest,
+    data: Annotated[CreateExecutionRequest | dict, dict_like(CreateExecutionRequest)],
 ) -> tuple[list[str], dict]:
     execution_id = execution_id or uuid4()
 
     developer_id = str(developer_id)
     task_id = str(task_id)
     execution_id = str(execution_id)
 
-    data.metadata = data.metadata or {}
-    execution_data = data.model_dump()
+    if isinstance(data, CreateExecutionRequest):
+        data.metadata = data.metadata or {}
+        execution_data = data.model_dump()
+    else:
+        data["metadata"] = data.get("metadata", {})
+        execution_data = data
 
     columns, values = cozo_process_mutate_data(
         {

agents-api/agents_api/models/execution/test_execution_queries.py

@@ -1,19 +1,20 @@
 # Tests for execution queries
 from uuid import uuid4
 
-from cozo_migrate.api import init, apply
+from cozo_migrate.api import apply, init
 from pycozo import Client
 from ward import test
 
 from agents_api.autogen.openapi_model import Execution, Transition
+
 from ..agent.create_agent import create_agent
 from ..task.create_task import create_task
 from .create_execution import create_execution
-from .get_execution import get_execution
-from .list_executions import list_executions
 from .create_execution_transition import create_execution_transition
+from .get_execution import get_execution
 from .get_execution_transition import get_execution_transition
 from .list_execution_transitions import list_execution_transitions
+from .list_executions import list_executions
 
 MODEL = "julep-ai/samantha-1-turbo"
 
@@ -140,7 +141,12 @@ def _():
         developer_id=developer_id,
         execution_id=execution_id,
         transition_id=transition_id,
-        data={"type": "step", "from": "test", "to": "test", "outputs": {"input": "test"}},
+        data={
+            "type": "step",
+            "from": "test",
+            "to": "test",
+            "outputs": {"input": "test"},
+        },
         client=client,
     )
 
@@ -156,7 +162,12 @@ def _():
         developer_id=developer_id,
         execution_id=execution_id,
         transition_id=transition_id,
-        data={"type": "step", "from": "test", "to": "test", "outputs": {"input": "test"}},
+        data={
+            "type": "step",
+            "from": "test",
+            "to": "test",
+            "outputs": {"input": "test"},
+        },
         client=client,
     )
 
@@ -182,7 +193,12 @@ def _():
         developer_id=developer_id,
         execution_id=execution_id,
         transition_id=transition_id,
-        data={"type": "step", "from": "test", "to": "test", "outputs": {"input": "test"}},
+        data={
+            "type": "step",
+            "from": "test",
+            "to": "test",
+            "outputs": {"input": "test"},
+        },
         client=client,
     )
 

agents-api/agents_api/models/session/test_session_queries.py

@@ -1,11 +1,12 @@
 # Tests for session queries
 from uuid import uuid4
 
-from cozo_migrate.api import init, apply
+from cozo_migrate.api import apply, init
 from pycozo import Client
 from ward import test
 
 from agents_api.autogen.openapi_model import Session
+
 from ..agent.create_agent import create_agent
 from ..user.create_user import create_user
 from .create_session import create_session

agents-api/agents_api/models/task/test_task_queries.py

@@ -1,11 +1,12 @@
 # Tests for task queries
 from uuid import uuid4
 
-from cozo_migrate.api import init, apply
+from cozo_migrate.api import apply, init
 from pycozo import Client
 from ward import test
 
 from agents_api.autogen.openapi_model import Task
+
 from .create_task import create_task
 from .delete_task import delete_task
 from .get_task import get_task

agents-api/agents_api/models/tools/test_tool_queries.py

@@ -1,11 +1,12 @@
 # Tests for tool queries
 from uuid import uuid4
 
-from cozo_migrate.api import init, apply
+from cozo_migrate.api import apply, init
 from pycozo import Client
 from ward import test
 
 from agents_api.autogen.openapi_model import FunctionDef, Tool
+
 from .create_tools import create_tools
 from .delete_tool import delete_tool
 from .get_tool import get_tool

agents-api/agents_api/models/user/test_user_queries.py

@@ -3,11 +3,12 @@
 # Tests for user queries
 from uuid import uuid4
 
-from cozo_migrate.api import init, apply
+from cozo_migrate.api import apply, init
 from pycozo import Client
 from ward import raises, test
 
 from agents_api.autogen.openapi_model import User
+
 from .create_user import create_user
 from .get_user import get_user
 from .list_users import list_users