Merge pull request #971 from julep-ai/f/test-pg-queries

feat(agents-api): implement agent queries
julep-ai · Dec 18, 2024 · 94637ed · 94637ed
2 parents bd83d4f + 451a88f
commit 94637ed
Show file tree

Hide file tree

Showing 11 changed files with 567 additions and 504 deletions.
diff --git a/agents-api/agents_api/queries/agents/__init__.py b/agents-api/agents_api/queries/agents/__init__.py
@@ -13,9 +13,9 @@
 # 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
+from .create_or_update_agent import create_or_update_agent
+from .delete_agent import delete_agent
+from .get_agent import get_agent
+from .list_agents import list_agents
+from .patch_agent import patch_agent
+from .update_agent import update_agent
diff --git a/agents-api/agents_api/queries/agents/create_agent.py b/agents-api/agents_api/queries/agents/create_agent.py
@@ -8,14 +8,14 @@
 
 from beartype import beartype
 from fastapi import HTTPException
-from psycopg import errors as psycopg_errors
 from pydantic import ValidationError
+from sqlglot import parse_one
 from uuid_extensions import uuid7
 
 from ...autogen.openapi_model import Agent, CreateAgentRequest
 from ...metrics.counters import increase_counter
 from ..utils import (
-    # generate_canonical_name,
+    generate_canonical_name,
     partialclass,
     pg_query,
     rewrap_exceptions,
@@ -25,51 +25,78 @@
 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.",
-        ),
-    }
+raw_query = """
+INSERT INTO agents (
+    developer_id,
+    agent_id,
+    canonical_name,
+    name,
+    about,
+    instructions,
+    model,
+    metadata,
+    default_settings
+)
+VALUES (
+    $1,
+    $2,
+    $3,
+    $4,
+    $5,
+    $6,
+    $7,
+    $8,
+    $9
 )
+RETURNING *;
+"""
+
+query = parse_one(raw_query).sql(pretty=True)
+
+
+# @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",
 )
+@increase_counter("create_agent")
 @pg_query
-# @increase_counter("create_agent")
 @beartype
-def create_agent(
+async def create_agent(
     *,
     developer_id: UUID,
     agent_id: UUID | None = None,
     data: CreateAgentRequest,
-) -> tuple[str, dict]:
+) -> tuple[str, list]:
     """
     Constructs and executes a SQL query to create a new agent in the database.
 
@@ -92,49 +119,23 @@ def create_agent(
 
     # Convert default_settings to dict if it exists
     default_settings = (
-        data.default_settings.model_dump() if data.default_settings else None
+        data.default_settings.model_dump() if data.default_settings else {}
     )
 
     # Set default values
-    data.metadata = data.metadata or None
-    # data.canonical_name = data.canonical_name or generate_canonical_name(data.name)
+    data.metadata = data.metadata or {}
+    data.canonical_name = data.canonical_name or generate_canonical_name(data.name)
 
-    query = """
-    INSERT INTO agents (
+    params = [
         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,
-    }
+        data.canonical_name,
+        data.name,
+        data.about,
+        data.instructions,
+        data.model,
+        data.metadata,
+        default_settings,
+    ]
 
     return query, params
diff --git a/agents-api/agents_api/queries/agents/create_or_update_agent.py b/agents-api/agents_api/queries/agents/create_or_update_agent.py
@@ -8,12 +8,13 @@
 
 from beartype import beartype
 from fastapi import HTTPException
-from psycopg import errors as psycopg_errors
+from sqlglot import parse_one
+from sqlglot.optimizer import optimize
 
 from ...autogen.openapi_model import Agent, CreateOrUpdateAgentRequest
 from ...metrics.counters import increase_counter
 from ..utils import (
-    # generate_canonical_name,
+    generate_canonical_name,
     partialclass,
     pg_query,
     rewrap_exceptions,
@@ -23,28 +24,55 @@
 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.",
-        )
-    }
+raw_query = """
+INSERT INTO agents (
+    developer_id,
+    agent_id,
+    canonical_name,
+    name,
+    about,
+    instructions,
+    model,
+    metadata,
+    default_settings
+)
+VALUES (
+    $1,
+    $2,
+    $3,
+    $4,
+    $5,
+    $6,
+    $7,
+    $8,
+    $9
 )
+RETURNING *;
+"""
+
+query = parse_one(raw_query).sql(pretty=True)
+
+
+# @rewrap_exceptions(
+#     {
+#         psycopg_errors.ForeignKeyViolation: partialclass(
+#             HTTPException,
+#             status_code=404,
+#             detail="The specified developer does not exist.",
+#         )
+#     }
+# )
 @wrap_in_class(
     Agent,
     one=True,
     transform=lambda d: {"id": d["agent_id"], **d},
-    _kind="inserted",
 )
+@increase_counter("create_or_update_agent")
 @pg_query
-# @increase_counter("create_or_update_agent1")
 @beartype
-def create_or_update_agent_query(
+async def create_or_update_agent(
     *, agent_id: UUID, developer_id: UUID, data: CreateOrUpdateAgentRequest
-) -> tuple[list[str], dict]:
+) -> tuple[str, list]:
     """
     Constructs the SQL queries to create a new agent or update an existing agent's details.
 
@@ -66,49 +94,23 @@ def create_or_update_agent_query(
 
     # Convert default_settings to dict if it exists
     default_settings = (
-        data.default_settings.model_dump() if data.default_settings else None
+        data.default_settings.model_dump() if data.default_settings else {}
     )
 
     # Set default values
-    data.metadata = data.metadata or None
-    # data.canonical_name = data.canonical_name or generate_canonical_name(data.name)
+    data.metadata = data.metadata or {}
+    data.canonical_name = data.canonical_name or generate_canonical_name(data.name)
 
-    query = """
-    INSERT INTO agents (
+    params = [
         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,
-    }
+        data.canonical_name,
+        data.name,
+        data.about,
+        data.instructions,
+        data.model,
+        data.metadata,
+        default_settings,
+    ]
 
     return (query, params)