feat: new routes for docs, users, search

agents-api/agents_api/autogen/Chat.py

@@ -6,7 +6,7 @@
 from typing import Annotated, Literal
 from uuid import UUID
 
-from pydantic import AwareDatetime, BaseModel, ConfigDict, Field
+from pydantic import AwareDatetime, BaseModel, ConfigDict, Field, RootModel
 
 from .Docs import DocReference
 from .Entries import ChatMLMessage

agents-api/agents_api/autogen/Entries.py

@@ -6,7 +6,7 @@
 from typing import Annotated, Literal
 from uuid import UUID
 
-from pydantic import AnyUrl, AwareDatetime, BaseModel, ConfigDict, Field
+from pydantic import AnyUrl, AwareDatetime, BaseModel, ConfigDict, Field, RootModel
 
 from .Tools import ChosenToolCall, Tool, ToolResponse
 

agents-api/agents_api/autogen/openapi_model.py

@@ -1,5 +1,5 @@
 # ruff: noqa: F401, F403, F405
-from typing import Annotated
+from typing import Annotated, Generic, TypeVar
 from uuid import UUID
 
 from pydantic import AwareDatetime, Field
@@ -158,3 +158,10 @@ class UpdateTaskRequest(_UpdateTaskRequest):
             "extra": "allow",
         }
     )
+
+
+DataT = TypeVar("DataT", bound=BaseModel)
+
+
+class ListResponse(BaseModel, Generic[DataT]):
+    items: list[DataT]

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

@@ -43,18 +43,12 @@ def create_agent(
     Constructs and executes a datalog query to create a new agent in the database.
 
     Parameters:
-    - agent_id (UUID): The unique identifier for the agent.
+    - agent_id (UUID | None): The unique identifier for the agent.
     - developer_id (UUID): The unique identifier for the developer creating the agent.
-    - name (str): The name of the agent.
-    - about (str): A description of the agent.
-    - instructions (list[str], optional): A list of instructions for using the agent. Defaults to an empty list.
-    - model (str, optional): The model identifier for the agent. Defaults to "julep-ai/samantha-1-turbo".
-    - metadata (dict, optional): A dictionary of metadata for the agent. Defaults to an empty dict.
-    - default_settings (dict, optional): A dictionary of default settings for the agent. Defaults to an empty dict.
-    - client (CozoClient, optional): The CozoDB client instance to use for the query. Defaults to a preconfigured client instance.
+    - data (CreateAgentRequest): The data for the new agent.
 
     Returns:
-    Agent: The newly created agent record.
+    - Agent: The newly created agent record.
     """
 
     agent_id = agent_id or uuid4()

agents-api/agents_api/models/utils.py

@@ -1,7 +1,7 @@
 import inspect
 import re
 from functools import partialmethod, wraps
-from typing import Any, Callable, ParamSpec, Type
+from typing import Any, Callable, ParamSpec, Type, TypeVar
 from uuid import UUID
 
 import pandas as pd
@@ -11,6 +11,7 @@
 from ..common.utils.cozo import uuid_int_list_to_uuid4
 
 P = ParamSpec("P")
+T = TypeVar("T")
 
 
 def fix_uuid(
@@ -126,7 +127,9 @@ def cozo_query_dec(func: Callable[P, tuple[str | list[str], dict]]):
             from pprint import pprint
 
         @wraps(func)
-        def wrapper(*args, client=cozo_client, **kwargs) -> pd.DataFrame:
+        def wrapper(
+            *args: P.args, client=cozo_client, **kwargs: P.kwargs
+        ) -> pd.DataFrame:
             queries, variables = func(*args, **kwargs)
 
             if isinstance(queries, str):
@@ -173,9 +176,9 @@ def wrap_in_class(
     one: bool = False,
     transform: Callable[[dict], dict] | None = None,
 ):
-    def decorator(func: Callable[..., pd.DataFrame]):
+    def decorator(func: Callable[P, pd.DataFrame]):
         @wraps(func)
-        def wrapper(*args, **kwargs):
+        def wrapper(*args: P.args, **kwargs: P.kwargs):
             df = func(*args, **kwargs)
 
             # Convert df to list of dicts
@@ -206,13 +209,13 @@ def rewrap_exceptions(
     ],
     /,
 ):
-    def decorator(func: Callable[..., Any]):
+    def decorator(func: Callable[P, T]):
         @wraps(func)
-        def wrapper(*args, **kwargs):
+        def wrapper(*args: P.args, **kwargs: P.kwargs) -> T:
             nonlocal mapping
 
             try:
-                result = func(*args, **kwargs)
+                result: T = func(*args, **kwargs)
 
             except BaseException as error:
                 for check, transform in mapping.items():

agents-api/agents_api/routers/agents/__init__.py

@@ -1,7 +1,15 @@
-from .create_agent import create_agent  # noqa: F401
-from .delete_agent import delete_agent  # noqa: F401
-from .get_agent_details import get_agent_details  # noqa: F401
-from .list_agents import list_agents  # noqa: F401
-from .patch_agent import patch_agent  # noqa: F401
-from .router import router  # noqa: F401
-from .update_agent import update_agent  # noqa: F401
+# ruff: noqa: F401
+
+from .create_agent import create_agent
+from .create_agent_tool import create_agent_tool
+from .create_or_update_agent import create_or_update_agent
+from .delete_agent import delete_agent
+from .delete_agent_tool import delete_agent_tool
+from .get_agent_details import get_agent_details
+from .list_agent_tools import list_agent_tools
+from .list_agents import list_agents
+from .patch_agent import patch_agent
+from .patch_agent_tool import patch_agent_tool
+from .router import router
+from .update_agent import update_agent
+from .update_agent_tool import update_agent_tool

agents-api/agents_api/routers/agents/create_agent.py

@@ -1,37 +1,27 @@
 from typing import Annotated
-from uuid import uuid4
 
 from fastapi import Depends
 from pydantic import UUID4
 from starlette.status import HTTP_201_CREATED
 
+import agents_api.models as models
+
 from ...autogen.openapi_model import (
     CreateAgentRequest,
     ResourceCreatedResponse,
 )
 from ...dependencies.developer_id import get_developer_id
-from ...models.agent.create_agent import create_agent as create_agent_query
 from .router import router
 
 
 @router.post("/agents", status_code=HTTP_201_CREATED, tags=["agents"])
 async def create_agent(
-    request: CreateAgentRequest,
     x_developer_id: Annotated[UUID4, Depends(get_developer_id)],
+    data: CreateAgentRequest,
 ) -> ResourceCreatedResponse:
-    new_agent_id = uuid4()
-
-    _, resp = next(
-        create_agent_query(
-            developer_id=x_developer_id,
-            agent_id=new_agent_id,
-            name=request.name,
-            about=request.about,
-            instructions=request.instructions or [],
-            model=request.model,
-            default_settings=request.default_settings or {},
-            metadata=request.metadata or {},
-        ).iterrows()
+    agent = models.agent.create_agent(
+        developer_id=x_developer_id,
+        data=data,
     )
 
-    return ResourceCreatedResponse(id=new_agent_id, created_at=resp["created_at"])
+    return ResourceCreatedResponse(id=agent.id, created_at=agent.created_at)

agents-api/agents_api/routers/agents/create_agent_tools.py → agents-api/agents_api/routers/agents/create_agent_tool.py

@@ -5,29 +5,26 @@
 from pydantic import UUID4
 from starlette.status import HTTP_201_CREATED
 
+import agents_api.models as models
+
 from ...autogen.openapi_model import (
     CreateToolRequest,
     ResourceCreatedResponse,
 )
 from ...dependencies.developer_id import get_developer_id
-from ...models.tools.create_tools import create_tools as create_tools_query
 from .router import router
 
 
 @router.post("/agents/{agent_id}/tools", status_code=HTTP_201_CREATED, tags=["agents"])
-async def create_agent_tools(
+async def create_agent_tool(
     agent_id: UUID,
     x_developer_id: Annotated[UUID4, Depends(get_developer_id)],
-    data: list[CreateToolRequest],
-    ignore_existing: bool = False,
+    data: CreateToolRequest,
 ) -> ResourceCreatedResponse:
-    _, resp = next(
-        create_tools_query(
-            developer_id=x_developer_id,
-            agent_id=agent_id,
-            data=data,
-            ignore_existing=ignore_existing,
-        ).iterrows()
+    tool = models.tools.create_tools(
+        developer_id=x_developer_id,
+        agent_id=agent_id,
+        data=[data],
     )
 
-    return ResourceCreatedResponse(id=resp["tool_id"], created_at=resp["created_at"])
+    return ResourceCreatedResponse(id=tool.id, created_at=tool.created_at)

agents-api/agents_api/routers/agents/create_or_update_agent.py

@@ -5,29 +5,26 @@
 from pydantic import UUID4
 from starlette.status import HTTP_201_CREATED
 
+import agents_api.models as models
+
 from ...autogen.openapi_model import (
     CreateOrUpdateAgentRequest,
     ResourceCreatedResponse,
 )
 from ...dependencies.developer_id import get_developer_id
-from ...models.agent.create_or_update_agent import (
-    create_or_update_agent as create_or_update_agent_query,
-)
 from .router import router
 
 
 @router.post("/agents/{agent_id}", status_code=HTTP_201_CREATED, tags=["agents"])
 async def create_or_update_agent(
     agent_id: UUID,
-    x_developer_id: Annotated[UUID4, Depends(get_developer_id)],
     data: CreateOrUpdateAgentRequest,
+    x_developer_id: Annotated[UUID4, Depends(get_developer_id)],
 ) -> ResourceCreatedResponse:
-    _, resp = next(
-        create_or_update_agent_query(
-            developer_id=x_developer_id,
-            agent_id=agent_id,
-            data=data,
-        ).iterrows()
+    agent = models.agent.create_or_update_agent(
+        developer_id=x_developer_id,
+        agent_id=agent_id,
+        data=data,
     )
 
-    return ResourceCreatedResponse(id=agent_id, created_at=resp["created_at"])
+    return ResourceCreatedResponse(id=agent.id, created_at=agent.created_at)

agents-api/agents_api/routers/agents/delete_agent.py

@@ -1,12 +1,10 @@
 from typing import Annotated
 
-from fastapi import Depends, HTTPException
+from fastapi import Depends
 from pydantic import UUID4
-from starlette.status import HTTP_202_ACCEPTED, HTTP_404_NOT_FOUND
+from starlette.status import HTTP_202_ACCEPTED
 
 from ...autogen.openapi_model import ResourceDeletedResponse
-from ...common.exceptions.agents import AgentNotFoundError
-from ...common.utils.datetime import utcnow
 from ...dependencies.developer_id import get_developer_id
 from ...models.agent.delete_agent import delete_agent as delete_agent_query
 from .router import router
@@ -16,8 +14,4 @@
 async def delete_agent(
     agent_id: UUID4, x_developer_id: Annotated[UUID4, Depends(get_developer_id)]
 ) -> ResourceDeletedResponse:
-    try:
-        delete_agent_query(x_developer_id, agent_id)
-    except AgentNotFoundError as e:
-        raise HTTPException(status_code=HTTP_404_NOT_FOUND, detail=str(e))
-    return ResourceDeletedResponse(id=agent_id, deleted_at=utcnow())
+    return delete_agent_query(developer_id=x_developer_id, agent_id=agent_id)

agents-api/agents_api/routers/agents/delete_agent_tools.py → agents-api/agents_api/routers/agents/delete_agent_tool.py

@@ -11,7 +11,7 @@
 
 
 @router.delete("/agents/{agent_id}/tools/{tool_id}", tags=["agents"])
-async def delete_agent_tools(
+async def delete_agent_tool(
     agent_id: UUID,
     tool_id: UUID,
     x_developer_id: Annotated[UUID4, Depends(get_developer_id)],

agents-api/agents_api/routers/agents/get_agent_details.py

@@ -1,11 +1,9 @@
 from typing import Annotated
 
-from fastapi import Depends, HTTPException
+from fastapi import Depends
 from pydantic import UUID4
-from starlette.status import HTTP_404_NOT_FOUND
 
 from ...autogen.openapi_model import Agent
-from ...common.exceptions.agents import AgentNotFoundError
 from ...dependencies.developer_id import get_developer_id
 from ...models.agent.get_agent import get_agent as get_agent_query
 from .router import router
@@ -16,10 +14,4 @@ async def get_agent_details(
     agent_id: UUID4,
     x_developer_id: Annotated[UUID4, Depends(get_developer_id)],
 ) -> Agent:
-    try:
-        agent = get_agent_query(developer_id=x_developer_id, agent_id=agent_id)
-        if not agent:
-            raise AgentNotFoundError(x_developer_id, agent_id)
-        return Agent(**agent)
-    except AgentNotFoundError as e:
-        raise HTTPException(status_code=HTTP_404_NOT_FOUND, detail=str(e))
+    return get_agent_query(developer_id=x_developer_id, agent_id=agent_id)

agents-api/agents_api/routers/agents/list_agent_tools.py

@@ -4,28 +4,28 @@
 from fastapi import Depends
 from pydantic import UUID4
 
+from ...autogen.openapi_model import ListResponse, Tool
 from ...dependencies.developer_id import get_developer_id
 from ...models.tools.list_tools import list_tools as list_tools_query
 from .router import router
 
 
 @router.get("/agents/{agent_id}/tools", tags=["agents"])
 async def list_agent_tools(
-    agent_id: UUID,
     x_developer_id: Annotated[UUID4, Depends(get_developer_id)],
+    agent_id: UUID,
     limit: int = 100,
     offset: int = 0,
     sort_by: Literal["created_at", "updated_at"] = "created_at",
     direction: Literal["asc", "desc"] = "desc",
-) -> list[tuple[str, dict]]:
-    return [
-        row
-        for _, row in list_tools_query(
-            developer_id=x_developer_id,
-            agent_id=agent_id,
-            limit=limit,
-            offset=offset,
-            sort_by=sort_by,
-            direction=direction,
-        ).iterrows()
-    ]
+) -> ListResponse[Tool]:
+    tools = list_tools_query(
+        agent_id=agent_id,
+        developer_id=x_developer_id,
+        limit=limit,
+        offset=offset,
+        sort_by=sort_by,
+        direction=direction,
+    )
+
+    return ListResponse[Tool](items=tools)

agents-api/agents_api/routers/agents/list_agents.py

@@ -1,11 +1,11 @@
-from typing import Annotated, List
+from typing import Annotated, Literal
 
 from fastapi import Depends
 from pydantic import UUID4
 
-from ...autogen.openapi_model import Agent
+from ...autogen.openapi_model import Agent, ListResponse
 from ...dependencies.developer_id import get_developer_id
-from ...models.agent.list_agents import list_agents
+from ...models.agent.list_agents import list_agents as list_agents_query
 from .router import router
 
 
@@ -14,12 +14,17 @@ async def list_agents(
     x_developer_id: Annotated[UUID4, Depends(get_developer_id)],
     limit: int = 100,
     offset: int = 0,
+    sort_by: Literal["created_at", "updated_at"] = "created_at",
+    direction: Literal["asc", "desc"] = "desc",
     metadata_filter: str = "{}",
-) -> List[Agent]:
-    agents = list_agents(
+) -> ListResponse[Agent]:
+    agents = list_agents_query(
         developer_id=x_developer_id,
         limit=limit,
         offset=offset,
+        sort_by=sort_by,
+        direction=direction,
         metadata_filter=metadata_filter,
     )
-    return [Agent(**agent) for agent in agents]
+
+    return ListResponse[Agent](items=agents)