F/add user queries

.gitignore

@@ -10,3 +10,4 @@ ngrok*
 */node_modules/
 .aider*
 .vscode/
+schema.sql

agents-api/agents_api/activities/execute_system.py

@@ -21,7 +21,7 @@
 from ..common.protocol.tasks import ExecutionInput, StepContext
 from ..common.storage_handler import auto_blob_store, load_from_blob_store_if_remote
 from ..env import testing
-from ..models.developer import get_developer
+from ..queries.developer import get_developer
 from .utils import get_handler
 
 # For running synchronous code in the background
@@ -94,7 +94,7 @@ async def execute_system(
 
         # Handle chat operations
         if system.operation == "chat" and system.resource == "session":
-            developer = get_developer(developer_id=arguments.get("developer_id"))
+            developer = await get_developer(developer_id=arguments.get("developer_id"))
             session_id = arguments.get("session_id")
             x_custom_api_key = arguments.get("x_custom_api_key", None)
             chat_input = ChatInput(**arguments)

agents-api/agents_api/autogen/Agents.py

@@ -25,16 +25,17 @@ class Agent(BaseModel):
     """
     When this resource was updated as UTC date-time
     """
-    name: Annotated[
-        str,
-        Field(
-            max_length=120,
-            pattern="^[\\p{L}\\p{Nl}\\p{Pattern_Syntax}\\p{Pattern_White_Space}]+[\\p{ID_Start}\\p{Mn}\\p{Mc}\\p{Nd}\\p{Pc}\\p{Pattern_Syntax}\\p{Pattern_White_Space}]*$",
-        ),
-    ] = ""
+    name: Annotated[str, Field(max_length=255, min_length=1)]
     """
     Name of the agent
     """
+    canonical_name: Annotated[
+        str | None,
+        Field(max_length=255, min_length=1, pattern="^[a-zA-Z][a-zA-Z0-9_]*$"),
+    ] = None
+    """
+    Canonical name of the agent
+    """
     about: str = ""
     """
     About the agent
@@ -62,16 +63,17 @@ class CreateAgentRequest(BaseModel):
         populate_by_name=True,
     )
     metadata: dict[str, Any] | None = None
-    name: Annotated[
-        str,
-        Field(
-            max_length=120,
-            pattern="^[\\p{L}\\p{Nl}\\p{Pattern_Syntax}\\p{Pattern_White_Space}]+[\\p{ID_Start}\\p{Mn}\\p{Mc}\\p{Nd}\\p{Pc}\\p{Pattern_Syntax}\\p{Pattern_White_Space}]*$",
-        ),
-    ] = ""
+    name: Annotated[str, Field(max_length=255, min_length=1)]
     """
     Name of the agent
     """
+    canonical_name: Annotated[
+        str | None,
+        Field(max_length=255, min_length=1, pattern="^[a-zA-Z][a-zA-Z0-9_]*$"),
+    ] = None
+    """
+    Canonical name of the agent
+    """
     about: str = ""
     """
     About the agent
@@ -96,16 +98,17 @@ class CreateOrUpdateAgentRequest(CreateAgentRequest):
     )
     id: UUID
     metadata: dict[str, Any] | None = None
-    name: Annotated[
-        str,
-        Field(
-            max_length=120,
-            pattern="^[\\p{L}\\p{Nl}\\p{Pattern_Syntax}\\p{Pattern_White_Space}]+[\\p{ID_Start}\\p{Mn}\\p{Mc}\\p{Nd}\\p{Pc}\\p{Pattern_Syntax}\\p{Pattern_White_Space}]*$",
-        ),
-    ] = ""
+    name: Annotated[str, Field(max_length=255, min_length=1)]
     """
     Name of the agent
     """
+    canonical_name: Annotated[
+        str | None,
+        Field(max_length=255, min_length=1, pattern="^[a-zA-Z][a-zA-Z0-9_]*$"),
+    ] = None
+    """
+    Canonical name of the agent
+    """
     about: str = ""
     """
     About the agent
@@ -133,16 +136,17 @@ class PatchAgentRequest(BaseModel):
         populate_by_name=True,
     )
     metadata: dict[str, Any] | None = None
-    name: Annotated[
-        str,
-        Field(
-            max_length=120,
-            pattern="^[\\p{L}\\p{Nl}\\p{Pattern_Syntax}\\p{Pattern_White_Space}]+[\\p{ID_Start}\\p{Mn}\\p{Mc}\\p{Nd}\\p{Pc}\\p{Pattern_Syntax}\\p{Pattern_White_Space}]*$",
-        ),
-    ] = ""
+    name: Annotated[str | None, Field(max_length=255, min_length=1)] = None
     """
     Name of the agent
     """
+    canonical_name: Annotated[
+        str | None,
+        Field(max_length=255, min_length=1, pattern="^[a-zA-Z][a-zA-Z0-9_]*$"),
+    ] = None
+    """
+    Canonical name of the agent
+    """
     about: str = ""
     """
     About the agent
@@ -170,16 +174,17 @@ class UpdateAgentRequest(BaseModel):
         populate_by_name=True,
     )
     metadata: dict[str, Any] | None = None
-    name: Annotated[
-        str,
-        Field(
-            max_length=120,
-            pattern="^[\\p{L}\\p{Nl}\\p{Pattern_Syntax}\\p{Pattern_White_Space}]+[\\p{ID_Start}\\p{Mn}\\p{Mc}\\p{Nd}\\p{Pc}\\p{Pattern_Syntax}\\p{Pattern_White_Space}]*$",
-        ),
-    ] = ""
+    name: Annotated[str, Field(max_length=255, min_length=1)]
     """
     Name of the agent
     """
+    canonical_name: Annotated[
+        str | None,
+        Field(max_length=255, min_length=1, pattern="^[a-zA-Z][a-zA-Z0-9_]*$"),
+    ] = None
+    """
+    Canonical name of the agent
+    """
     about: str = ""
     """
     About the agent

agents-api/agents_api/clients/pg.py

@@ -0,0 +1,12 @@
+import asyncpg
+
+from ..env import db_dsn
+from ..web import app
+
+
+async def get_pg_client():
+    client = getattr(app.state, "pg_client", await asyncpg.connect(db_dsn))
+    if not hasattr(app.state, "pg_client"):
+        app.state.pg_client = client
+
+    return client

agents-api/agents_api/dependencies/developer_id.py

@@ -5,7 +5,7 @@
 
 from ..common.protocol.developers import Developer
 from ..env import multi_tenant_mode
-from ..models.developer.get_developer import get_developer, verify_developer
+from ..queries.developer.get_developer import get_developer, verify_developer
 from .exceptions import InvalidHeaderFormat
 
 
@@ -36,7 +36,9 @@ async def get_developer_data(
         assert (
             not x_developer_id
         ), "X-Developer-Id header not allowed in multi-tenant mode"
-        return get_developer(developer_id=UUID("00000000-0000-0000-0000-000000000000"))
+        return await get_developer(
+            developer_id=UUID("00000000-0000-0000-0000-000000000000")
+        )
 
     if not x_developer_id:
         raise InvalidHeaderFormat("X-Developer-Id header required")
@@ -47,6 +49,6 @@ async def get_developer_data(
         except ValueError as e:
             raise InvalidHeaderFormat("X-Developer-Id must be a valid UUID") from e
 
-    developer = get_developer(developer_id=x_developer_id)
+    developer = await get_developer(developer_id=x_developer_id)
 
     return developer

agents-api/agents_api/env.py

@@ -59,6 +59,13 @@
     "DO_VERIFY_DEVELOPER_OWNS_RESOURCE", default=True
 )
 
+# PostgreSQL
+# ----
+db_dsn: str = env.str(
+    "DB_DSN",
+    default="postgres://postgres:[email protected]:5432/postgres?sslmode=disable",
+)
+
 
 # Auth
 # ----

agents-api/agents_api/models/utils.py

@@ -7,6 +7,7 @@
 from uuid import UUID
 
 import pandas as pd
+from asyncpg import Record
 from fastapi import HTTPException
 from httpcore import ConnectError, NetworkError, TimeoutException
 from httpx import ConnectError as HttpxConnectError
@@ -463,12 +464,12 @@ def wrap_in_class(
     transform: Callable[[dict], dict] | None = None,
     _kind: str | None = None,
 ):
-    def _return_data(df: pd.DataFrame):
+    def _return_data(rec: Record):
         # Convert df to list of dicts
-        if _kind:
-            df = df[df["_kind"] == _kind]
+        # if _kind:
+        #     rec = rec[rec["_kind"] == _kind]
 
-        data = df.to_dict(orient="records")
+        data = list(rec.items())
 
         nonlocal transform
         transform = transform or (lambda x: x)

agents-api/agents_api/queries/agent/__init__.py

@@ -0,0 +1,21 @@
+"""
+The `agent` module within the `queries` package provides a comprehensive suite of SQL query functions for managing agents in the PostgreSQL database. This includes:
+
+- Creating new agents
+- Updating existing agents
+- Retrieving details about specific agents
+- Listing agents with filtering and pagination
+- Deleting agents from the database
+
+Each function in this module constructs and returns SQL queries along with their parameters for database operations.
+"""
+
+# ruff: noqa: F401, F403, F405
+
+from .create_agent import create_agent
+from .create_or_update_agent import create_or_update_agent_query
+from .delete_agent import delete_agent_query
+from .get_agent import get_agent_query
+from .list_agents import list_agents_query
+from .patch_agent import patch_agent_query
+from .update_agent import update_agent_query

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

@@ -0,0 +1,140 @@
+"""
+This module contains the functionality for creating agents in the PostgreSQL database.
+It includes functions to construct and execute SQL queries for inserting new agent records.
+"""
+
+from typing import Any, TypeVar
+from uuid import UUID
+
+from beartype import beartype
+from fastapi import HTTPException
+from psycopg import errors as psycopg_errors
+from pydantic import ValidationError
+from uuid_extensions import uuid7
+
+from ...autogen.openapi_model import Agent, CreateAgentRequest
+from ...metrics.counters import increase_counter
+from ..utils import (
+    generate_canonical_name,
+    partialclass,
+    pg_query,
+    rewrap_exceptions,
+    wrap_in_class,
+)
+
+ModelT = TypeVar("ModelT", bound=Any)
+T = TypeVar("T")
+
+
+@rewrap_exceptions(
+    {
+        psycopg_errors.ForeignKeyViolation: partialclass(
+            HTTPException,
+            status_code=404,
+            detail="The specified developer does not exist.",
+        ),
+        psycopg_errors.UniqueViolation: partialclass(
+            HTTPException,
+            status_code=409,
+            detail="An agent with this canonical name already exists for this developer.",
+        ),
+        psycopg_errors.CheckViolation: partialclass(
+            HTTPException,
+            status_code=400,
+            detail="The provided data violates one or more constraints. Please check the input values.",
+        ),
+        ValidationError: partialclass(
+            HTTPException,
+            status_code=400,
+            detail="Input validation failed. Please check the provided data.",
+        ),
+        TypeError: partialclass(
+            HTTPException,
+            status_code=400,
+            detail="A type mismatch occurred. Please review the input.",
+        ),
+    }
+)
+@wrap_in_class(
+    Agent,
+    one=True,
+    transform=lambda d: {"id": d["agent_id"], **d},
+    _kind="inserted",
+)
+@pg_query
+@increase_counter("create_agent")
+@beartype
+def create_agent(
+    *,
+    developer_id: UUID,
+    agent_id: UUID | None = None,
+    data: CreateAgentRequest,
+) -> tuple[str, dict]:
+    """
+    Constructs and executes a SQL query to create a new agent in the database.
+
+    Parameters:
+        agent_id (UUID | None): The unique identifier for the agent.
+        developer_id (UUID): The unique identifier for the developer creating the agent.
+        data (CreateAgentRequest): The data for the new agent.
+
+    Returns:
+        tuple[str, dict]: SQL query and parameters for creating the agent.
+    """
+    agent_id = agent_id or uuid7()
+
+    # Ensure instructions is a list
+    data.instructions = (
+        data.instructions
+        if isinstance(data.instructions, list)
+        else [data.instructions]
+    )
+
+    # Convert default_settings to dict if it exists
+    default_settings = (
+        data.default_settings.model_dump() if data.default_settings else None
+    )
+
+    # Set default values
+    data.metadata = data.metadata or None
+    data.canonical_name = data.canonical_name or generate_canonical_name(data.name)
+
+    query = """
+    INSERT INTO agents (
+        developer_id,
+        agent_id,
+        canonical_name,
+        name,
+        about,
+        instructions,
+        model,
+        metadata,
+        default_settings
+    )
+    VALUES (
+        %(developer_id)s,
+        %(agent_id)s,
+        %(canonical_name)s,
+        %(name)s,
+        %(about)s,
+        %(instructions)s,
+        %(model)s,
+        %(metadata)s,
+        %(default_settings)s
+    )
+    RETURNING *;
+    """
+
+    params = {
+        "developer_id": developer_id,
+        "agent_id": agent_id,
+        "canonical_name": data.canonical_name,
+        "name": data.name,
+        "about": data.about,
+        "instructions": data.instructions,
+        "model": data.model,
+        "metadata": data.metadata,
+        "default_settings": default_settings,
+    }
+
+    return query, params