refactor: Lint agents-api (CI)

julep-ai · Dec 25, 2024 · 88ca072 · 88ca072
1 parent eaada88
commit 88ca072
Show file tree

Hide file tree

Showing 16 changed files with 100 additions and 107 deletions.
diff --git a/agents-api/agents_api/queries/executions/__init__.py b/agents-api/agents_api/queries/executions/__init__.py
@@ -20,7 +20,6 @@
 from .lookup_temporal_data import lookup_temporal_data
 from .prepare_execution_input import prepare_execution_input
 
-
 __all__ = [
     "count_executions",
     "create_execution",

diff --git a/agents-api/agents_api/queries/executions/count_executions.py b/agents-api/agents_api/queries/executions/count_executions.py
@@ -1,15 +1,16 @@
 from typing import Any, Literal
 from uuid import UUID
 
+import asyncpg
 from beartype import beartype
 from fastapi import HTTPException
 from sqlglot import parse_one
-import asyncpg
+
 from ..utils import (
+    partialclass,
     pg_query,
-    wrap_in_class,
     rewrap_exceptions,
-    partialclass,
+    wrap_in_class,
 )
 
 # Query to count executions for a given task
@@ -22,18 +23,18 @@
 
 
 @rewrap_exceptions(
-{
-    asyncpg.NoDataFoundError: partialclass(
-        HTTPException, 
-        status_code=404,
-        detail="No executions found for the specified task"
-    ),
-    asyncpg.ForeignKeyViolationError: partialclass(
-        HTTPException,
-        status_code=404,
-        detail="The specified developer or task does not exist"
-    ),
-}
+    {
+        asyncpg.NoDataFoundError: partialclass(
+            HTTPException,
+            status_code=404,
+            detail="No executions found for the specified task",
+        ),
+        asyncpg.ForeignKeyViolationError: partialclass(
+            HTTPException,
+            status_code=404,
+            detail="The specified developer or task does not exist",
+        ),
+    }
 )
 @wrap_in_class(dict, one=True)
 @pg_query

diff --git a/agents-api/agents_api/queries/executions/create_execution.py b/agents-api/agents_api/queries/executions/create_execution.py
@@ -1,25 +1,24 @@
 from typing import Annotated, Any, TypeVar
 from uuid import UUID
 
+import asyncpg
 from beartype import beartype
+from fastapi import HTTPException
+from sqlglot import parse_one
 from uuid_extensions import uuid7
 
 from ...autogen.openapi_model import CreateExecutionRequest, Execution
 from ...common.utils.datetime import utcnow
 from ...common.utils.types import dict_like
 from ...metrics.counters import increase_counter
-from sqlglot import parse_one
-import asyncpg
-from fastapi import HTTPException
 from ..utils import (
+    partialclass,
     pg_query,
-    wrap_in_class,
     rewrap_exceptions,
-    partialclass,
+    wrap_in_class,
 )
 from .constants import OUTPUT_UNNEST_KEY
 
-
 create_execution_query = parse_one("""
 INSERT INTO executions
 (
@@ -44,18 +43,18 @@
 
 
 @rewrap_exceptions(
-{
-    asyncpg.NoDataFoundError: partialclass(
-        HTTPException, 
-        status_code=404,
-        detail="No executions found for the specified task"
-    ),
-    asyncpg.ForeignKeyViolationError: partialclass(
-        HTTPException,
-        status_code=404,
-        detail="The specified developer or task does not exist"
-    ),
-}
+    {
+        asyncpg.NoDataFoundError: partialclass(
+            HTTPException,
+            status_code=404,
+            detail="No executions found for the specified task",
+        ),
+        asyncpg.ForeignKeyViolationError: partialclass(
+            HTTPException,
+            status_code=404,
+            detail="The specified developer or task does not exist",
+        ),
+    }
 )
 @wrap_in_class(
     Execution,

diff --git a/agents-api/agents_api/queries/executions/create_execution_transition.py b/agents-api/agents_api/queries/executions/create_execution_transition.py
@@ -1,23 +1,23 @@
 from typing import Literal
 from uuid import UUID
 
+import asyncpg
 from beartype import beartype
+from fastapi import HTTPException
+from sqlglot import parse_one
 from uuid_extensions import uuid7
 
 from ...autogen.openapi_model import (
     CreateTransitionRequest,
     Transition,
 )
-import asyncpg
-from fastapi import HTTPException
-from sqlglot import parse_one
 from ...common.utils.datetime import utcnow
 from ...metrics.counters import increase_counter
 from ..utils import (
+    partialclass,
     pg_query,
-    wrap_in_class,
     rewrap_exceptions,
-    partialclass,
+    wrap_in_class,
 )
 
 # Query to create a transition
@@ -87,18 +87,18 @@ def validate_transition_targets(data: CreateTransitionRequest) -> None:
 
 
 @rewrap_exceptions(
-{
-    asyncpg.NoDataFoundError: partialclass(
-        HTTPException, 
-        status_code=404,
-        detail="No executions found for the specified task"
-    ),
-    asyncpg.ForeignKeyViolationError: partialclass(
-        HTTPException,
-        status_code=404,
-        detail="The specified developer or task does not exist"
-    ),
-}
+    {
+        asyncpg.NoDataFoundError: partialclass(
+            HTTPException,
+            status_code=404,
+            detail="No executions found for the specified task",
+        ),
+        asyncpg.ForeignKeyViolationError: partialclass(
+            HTTPException,
+            status_code=404,
+            detail="The specified developer or task does not exist",
+        ),
+    }
 )
 @wrap_in_class(
     Transition,

diff --git a/agents-api/agents_api/queries/executions/create_temporal_lookup.py b/agents-api/agents_api/queries/executions/create_temporal_lookup.py
@@ -1,16 +1,16 @@
+from uuid import UUID
 
-from beartype import beartype
-from temporalio.client import WorkflowHandle
-from sqlglot import parse_one
 import asyncpg
+from beartype import beartype
 from fastapi import HTTPException
-from uuid import UUID
+from sqlglot import parse_one
+from temporalio.client import WorkflowHandle
 
 from ...metrics.counters import increase_counter
 from ..utils import (
+    partialclass,
     pg_query,
     rewrap_exceptions,
-    partialclass,
 )
 
 # Query to create a temporal lookup
@@ -36,18 +36,18 @@
 
 
 @rewrap_exceptions(
-{
-    asyncpg.NoDataFoundError: partialclass(
-        HTTPException, 
-        status_code=404,
-        detail="No executions found for the specified task"
-    ),
-    asyncpg.ForeignKeyViolationError: partialclass(
-        HTTPException,
-        status_code=404,
-        detail="The specified developer or task does not exist"
-    ),
-}
+    {
+        asyncpg.NoDataFoundError: partialclass(
+            HTTPException,
+            status_code=404,
+            detail="No executions found for the specified task",
+        ),
+        asyncpg.ForeignKeyViolationError: partialclass(
+            HTTPException,
+            status_code=404,
+            detail="The specified developer or task does not exist",
+        ),
+    }
 )
 @pg_query
 @increase_counter("create_temporal_lookup")

diff --git a/agents-api/agents_api/queries/executions/get_execution.py b/agents-api/agents_api/queries/executions/get_execution.py
@@ -2,9 +2,9 @@
 from uuid import UUID
 
 import asyncpg
+from beartype import beartype
 from fastapi import HTTPException
 from sqlglot import parse_one
-from beartype import beartype
 
 from ...autogen.openapi_model import Execution
 from ..utils import (
@@ -27,9 +27,9 @@
 @rewrap_exceptions(
     {
         asyncpg.NoDataFoundError: partialclass(
-            HTTPException, 
+            HTTPException,
             status_code=404,
-            detail="No executions found for the specified task"
+            detail="No executions found for the specified task",
         ),
     }
 )

diff --git a/agents-api/agents_api/queries/executions/get_execution_transition.py b/agents-api/agents_api/queries/executions/get_execution_transition.py
@@ -5,6 +5,7 @@
 from beartype import beartype
 from fastapi import HTTPException
 from sqlglot import parse_one
+
 from ...autogen.openapi_model import Transition
 from ..utils import (
     partialclass,
@@ -40,18 +41,18 @@ def _transform(d):
 
 
 @rewrap_exceptions(
-{
-    asyncpg.NoDataFoundError: partialclass(
-        HTTPException, 
-        status_code=404,
-        detail="No executions found for the specified task"
-    ),
-    asyncpg.ForeignKeyViolationError: partialclass(
-        HTTPException,
-        status_code=404,
-        detail="The specified developer or task does not exist"
-    ),
-}
+    {
+        asyncpg.NoDataFoundError: partialclass(
+            HTTPException,
+            status_code=404,
+            detail="No executions found for the specified task",
+        ),
+        asyncpg.ForeignKeyViolationError: partialclass(
+            HTTPException,
+            status_code=404,
+            detail="The specified developer or task does not exist",
+        ),
+    }
 )
 @wrap_in_class(Transition, one=True, transform=_transform)
 @pg_query

diff --git a/agents-api/agents_api/queries/executions/get_paused_execution_token.py b/agents-api/agents_api/queries/executions/get_paused_execution_token.py
@@ -27,9 +27,9 @@
 @rewrap_exceptions(
     {
         asyncpg.NoDataFoundError: partialclass(
-            HTTPException, 
+            HTTPException,
             status_code=404,
-            detail="No paused executions found for the specified task"
+            detail="No paused executions found for the specified task",
         ),
     }
 )
@@ -72,4 +72,3 @@ async def get_paused_execution_token(
         [execution_id],
         "fetchrow",
     )
-
diff --git a/agents-api/agents_api/queries/executions/get_temporal_workflow_data.py b/agents-api/agents_api/queries/executions/get_temporal_workflow_data.py
@@ -25,9 +25,9 @@
 @rewrap_exceptions(
     {
         asyncpg.NoDataFoundError: partialclass(
-            HTTPException, 
+            HTTPException,
             status_code=404,
-            detail="No temporal workflow data found for the specified execution"
+            detail="No temporal workflow data found for the specified execution",
         ),
     }
 )

diff --git a/agents-api/agents_api/queries/executions/list_execution_transitions.py b/agents-api/agents_api/queries/executions/list_execution_transitions.py
@@ -42,14 +42,10 @@ def _transform(d):
 @rewrap_exceptions(
     {
         asyncpg.InvalidRowCountInLimitClauseError: partialclass(
-            HTTPException, 
-            status_code=400,
-            detail="Invalid limit clause"
+            HTTPException, status_code=400, detail="Invalid limit clause"
         ),
         asyncpg.InvalidRowCountInResultOffsetClauseError: partialclass(
-            HTTPException, 
-            status_code=400,
-            detail="Invalid offset clause"
+            HTTPException, status_code=400, detail="Invalid offset clause"
         ),
     }
 )

diff --git a/agents-api/agents_api/queries/executions/list_executions.py b/agents-api/agents_api/queries/executions/list_executions.py
@@ -5,6 +5,7 @@
 from beartype import beartype
 from fastapi import HTTPException
 from sqlglot import parse_one
+
 from ...autogen.openapi_model import Execution
 from ..utils import (
     partialclass,
@@ -32,14 +33,10 @@
 @rewrap_exceptions(
     {
         asyncpg.InvalidRowCountInLimitClauseError: partialclass(
-            HTTPException, 
-            status_code=400, 
-            detail="Invalid limit clause"
+            HTTPException, status_code=400, detail="Invalid limit clause"
         ),
         asyncpg.InvalidRowCountInResultOffsetClauseError: partialclass(
-            HTTPException, 
-            status_code=400, 
-            detail="Invalid offset clause"
+            HTTPException, status_code=400, detail="Invalid offset clause"
         ),
     }
 )

diff --git a/agents-api/agents_api/queries/executions/lookup_temporal_data.py b/agents-api/agents_api/queries/executions/lookup_temporal_data.py
@@ -5,6 +5,7 @@
 from beartype import beartype
 from fastapi import HTTPException
 from sqlglot import parse_one
+
 from ..utils import partialclass, pg_query, rewrap_exceptions, wrap_in_class
 
 # Query to lookup temporal data
@@ -19,9 +20,9 @@
 @rewrap_exceptions(
     {
         asyncpg.NoDataFoundError: partialclass(
-            HTTPException, 
+            HTTPException,
             status_code=404,
-            detail="No temporal data found for the specified execution"
+            detail="No temporal data found for the specified execution",
         ),
     }
 )

diff --git a/agents-api/agents_api/queries/executions/prepare_execution_input.py b/agents-api/agents_api/queries/executions/prepare_execution_input.py
@@ -3,6 +3,7 @@
 
 from beartype import beartype
 from sqlglot import parse_one
+
 from ...common.protocol.tasks import ExecutionInput
 from ..utils import (
     pg_query,