fix(agents-api): Update TODO stuff

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

agents-api/agents_api/__init__.py

@@ -9,5 +9,3 @@
 
 with workflow.unsafe.imports_passed_through():
     import msgpack as msgpack
-
-import os

agents-api/agents_api/app.py

@@ -1,16 +1,13 @@
 import os
 from contextlib import asynccontextmanager
-from typing import Any, Callable, Coroutine
 
 from aiobotocore.session import get_session
-from fastapi import APIRouter, FastAPI, Request, Response
-from fastapi.params import Depends
+from fastapi import APIRouter, FastAPI
 from prometheus_fastapi_instrumentator import Instrumentator
 from scalar_fastapi import get_scalar_api_reference
 
 from .clients.pg import create_db_pool
-from .dependencies.content_length import valid_content_length
-from .env import api_prefix, hostname, max_payload_size, protocol, public_port
+from .env import api_prefix, hostname, protocol, public_port
 
 
 # TODO: This currently doesn't use .env variables, but we should move to using them

agents-api/agents_api/queries/executions/count_executions.py

@@ -1,4 +1,4 @@
-from typing import Any, Literal
+from typing import Literal
 from uuid import UUID
 
 import asyncpg

agents-api/agents_api/queries/executions/create_execution.py

@@ -1,4 +1,4 @@
-from typing import Annotated, Any, TypeVar
+from typing import Annotated
 from uuid import UUID
 
 import asyncpg

agents-api/agents_api/queries/executions/create_execution_transition.py

@@ -52,6 +52,7 @@
 """).sql(pretty=True)
 
 
+# FIXME: Remove this function
 def validate_transition_targets(data: CreateTransitionRequest) -> None:
     # Make sure the current/next targets are valid
     match data.type:

agents-api/agents_api/queries/executions/create_temporal_lookup.py

@@ -54,7 +54,7 @@
 @beartype
 async def create_temporal_lookup(
     *,
-    developer_id: UUID,  # TODO: what to do with this parameter?
+    developer_id: UUID,  # FIXME: Remove this parameter
     execution_id: UUID,
     workflow_handle: WorkflowHandle,
 ) -> tuple[str, list]:

agents-api/agents_api/queries/executions/get_execution_transition.py

@@ -1,4 +1,4 @@
-from typing import Any, Literal, TypeVar
+from typing import Literal
 from uuid import UUID
 
 import asyncpg
@@ -14,6 +14,7 @@
     wrap_in_class,
 )
 
+# FIXME: Use latest_transitions instead of transitions
 # Query to get an execution transition
 get_execution_transition_query = parse_one("""
 SELECT * FROM transitions
@@ -59,7 +60,7 @@ def _transform(d):
 @beartype
 async def get_execution_transition(
     *,
-    developer_id: UUID,  # TODO: what to do with this parameter?
+    developer_id: UUID,  # FIXME: Remove this parameter
     transition_id: UUID | None = None,
     task_token: str | None = None,
 ) -> tuple[str, list, Literal["fetch", "fetchmany", "fetchrow"]]:

agents-api/agents_api/queries/executions/get_paused_execution_token.py

@@ -1,10 +1,9 @@
-from typing import Any, Literal, TypeVar
+from typing import Literal
 from uuid import UUID
 
 import asyncpg
 from beartype import beartype
 from fastapi import HTTPException
-from sqlglot import parse_one
 
 from ..utils import (
     partialclass,
@@ -14,14 +13,14 @@
 )
 
 # Query to get a paused execution token
-get_paused_execution_token_query = parse_one("""
-SELECT * FROM transitions
+get_paused_execution_token_query = """
+SELECT * FROM latest_transitions
 WHERE
     execution_id = $1
         AND type = 'wait'
     ORDER BY created_at DESC
     LIMIT 1;
-""").sql(pretty=True)
+"""
 
 
 @rewrap_exceptions(
@@ -53,20 +52,6 @@ async def get_paused_execution_token(
     """
     execution_id = str(execution_id)
 
-    # TODO: what to do with this query?
-    # check_status_query = """
-    # ?[execution_id, status] :=
-    #     *executions:execution_id_status_idx {
-    #         execution_id,
-    #         status,
-    #     },
-    #     execution_id = to_uuid($execution_id),
-    #     status = "awaiting_input"
-
-    # :limit 1
-    # :assert some
-    # """
-
     return (
         get_paused_execution_token_query,
         [execution_id],

agents-api/agents_api/queries/executions/list_executions.py

@@ -1,4 +1,4 @@
-from typing import Any, Literal, TypeVar
+from typing import Literal
 from uuid import UUID
 
 import asyncpg

agents-api/agents_api/queries/executions/lookup_temporal_data.py

@@ -1,4 +1,4 @@
-from typing import Any, Literal, TypeVar
+from typing import Literal
 from uuid import UUID
 
 import asyncpg
@@ -8,11 +8,20 @@
 
 from ..utils import partialclass, pg_query, rewrap_exceptions, wrap_in_class
 
+# FIXME: Check if this query is correct
+
+
 # Query to lookup temporal data
 lookup_temporal_data_query = parse_one("""
-SELECT * FROM temporal_executions_lookup
+SELECT t.*
+FROM
+    temporal_executions_lookup t,
+    executions e
 WHERE
-    execution_id = $1
+    t.execution_id = e.execution_id
+    AND t.developer_id = e.developer_id
+    AND e.execution_id = $1
+    AND e.developer_id = $2
 LIMIT 1;
 """).sql(pretty=True)
 
@@ -31,7 +40,7 @@
 @beartype
 async def lookup_temporal_data(
     *,
-    developer_id: UUID,  # TODO: what to do with this parameter?
+    developer_id: UUID,
     execution_id: UUID,
 ) -> tuple[str, list, Literal["fetch", "fetchmany", "fetchrow"]]:
     """
@@ -49,6 +58,6 @@ async def lookup_temporal_data(
 
     return (
         lookup_temporal_data_query,
-        [execution_id],
+        [execution_id, developer_id],
         "fetchrow",
     )

agents-api/agents_api/queries/executions/prepare_execution_input.py

@@ -1,4 +1,3 @@
-from typing import Any
 from uuid import UUID
 
 from beartype import beartype

agents-api/agents_api/routers/docs/create_doc.py

@@ -3,15 +3,9 @@
 
 from fastapi import BackgroundTasks, Depends
 from starlette.status import HTTP_201_CREATED
-from temporalio.client import Client as TemporalClient
-from uuid_extensions import uuid7
 
-from ...activities.types import EmbedDocsPayload
 from ...autogen.openapi_model import CreateDocRequest, Doc, ResourceCreatedResponse
-from ...clients import temporal
-from ...common.retry_policies import DEFAULT_RETRY_POLICY
 from ...dependencies.developer_id import get_developer_id
-from ...env import temporal_task_queue, testing
 from ...queries.docs.create_doc import create_doc as create_doc_query
 from .router import router
 

agents-api/agents_api/routers/tasks/create_task_execution.py

@@ -15,7 +15,6 @@
     CreateTransitionRequest,
     Execution,
     ResourceCreatedResponse,
-    UpdateExecutionRequest,
 )
 from ...clients.temporal import run_task_execution_workflow
 from ...common.protocol.developers import Developer

agents-api/tests/fixtures.py

@@ -36,7 +36,6 @@
 from agents_api.queries.sessions.create_session import create_session
 from agents_api.queries.tasks.create_task import create_task
 from agents_api.queries.tools.create_tools import create_tools
-from agents_api.queries.tools.delete_tool import delete_tool
 from agents_api.queries.users.create_user import create_user
 from agents_api.web import app