refactor: Lint agents-api (CI)

julep-ai · Nov 29, 2024 · 6e612c1 · 6e612c1
1 parent be87865
commit 6e612c1
Show file tree

Hide file tree

Showing 6 changed files with 22 additions and 17 deletions.
diff --git a/agents-api/agents_api/activities/sync_items_remote.py b/agents-api/agents_api/activities/sync_items_remote.py
@@ -1,23 +1,28 @@
+import asyncio
 from typing import Any
 
 from beartype import beartype
 from temporalio import activity
 
 from ..common.protocol.remote import RemoteObject
-import asyncio
+
 
 @beartype
 async def save_inputs_remote_fn(inputs: list[Any]) -> list[Any | RemoteObject]:
     from ..common.storage_handler import store_in_blob_store_if_large
 
-    return await asyncio.gather(*[store_in_blob_store_if_large(input) for input in inputs])
+    return await asyncio.gather(
+        *[store_in_blob_store_if_large(input) for input in inputs]
+    )
 
 
 @beartype
 async def load_inputs_remote_fn(inputs: list[Any | RemoteObject]) -> list[Any]:
     from ..common.storage_handler import load_from_blob_store_if_remote
 
-    return await asyncio.gather(*[load_from_blob_store_if_remote(input) for input in inputs])
+    return await asyncio.gather(
+        *[load_from_blob_store_if_remote(input) for input in inputs]
+    )
 
 
 save_inputs_remote = activity.defn(name="save_inputs_remote")(save_inputs_remote_fn)

diff --git a/agents-api/agents_api/activities/task_steps/prompt_step.py b/agents-api/agents_api/activities/task_steps/prompt_step.py
@@ -5,11 +5,9 @@
 from langchain_core.tools import BaseTool
 from langchain_core.tools.convert import tool as tool_decorator
 from litellm.types.utils import ModelResponse
-from litellm.types.utils import ModelResponse
 from temporalio import activity
 from temporalio.exceptions import ApplicationError
 
-from ...autogen.openapi_model import Tool
 from ...autogen.openapi_model import Tool
 from ...clients import (
     litellm,  # We dont directly import `acompletion` so we can mock it

diff --git a/agents-api/agents_api/activities/task_steps/transition_step.py b/agents-api/agents_api/activities/task_steps/transition_step.py
@@ -14,7 +14,9 @@ async def transition_step(
     transition_info: CreateTransitionRequest,
 ) -> Transition:
     # Load output from blob store if it is a remote object
-    transition_info.output = await load_from_blob_store_if_remote(transition_info.output)
+    transition_info.output = await load_from_blob_store_if_remote(
+        transition_info.output
+    )
 
     # Create transition
     transition = create_execution_transition(

diff --git a/agents-api/agents_api/common/storage_handler.py b/agents-api/agents_api/common/storage_handler.py
@@ -105,7 +105,9 @@ async def load_args(
                                 setattr(
                                     arg,
                                     field,
-                                    await load_from_blob_store_if_remote(getattr(arg, field)),
+                                    await load_from_blob_store_if_remote(
+                                        getattr(arg, field)
+                                    ),
                                 )
                             elif isinstance(getattr(arg, field), RemoteList):
                                 setattr(
@@ -147,7 +149,9 @@ async def load_args(
                                 setattr(
                                     v,
                                     field,
-                                    await load_from_blob_store_if_remote(getattr(v, field)),
+                                    await load_from_blob_store_if_remote(
+                                        getattr(v, field)
+                                    ),
                                 )
                             elif isinstance(getattr(v, field), RemoteList):
                                 setattr(

diff --git a/agents-api/agents_api/common/sync_storage_handler.py b/agents-api/agents_api/common/sync_storage_handler.py
@@ -12,6 +12,7 @@
 )
 from ..worker.codec import deserialize, serialize
 
+
 def sync_store_in_blob_store_if_large(x: Any) -> RemoteObject | Any:
     if not use_blob_store_for_temporal:
         return x
@@ -27,6 +28,7 @@ def sync_store_in_blob_store_if_large(x: Any) -> RemoteObject | Any:
 
     return x
 
+
 def sync_load_from_blob_store_if_remote(x: Any | RemoteObject) -> Any:
     if not use_blob_store_for_temporal:
         return x
@@ -46,13 +48,12 @@ def sync_load_from_blob_store_if_remote(x: Any | RemoteObject) -> Any:
 
     return x
 
+
 def sync_load_args(
     deep: bool, args: list | tuple, kwargs: dict[str, Any]
 ) -> tuple[list | tuple, dict[str, Any]]:
     new_args = [sync_load_from_blob_store_if_remote(arg) for arg in args]
-    new_kwargs = {
-        k: sync_load_from_blob_store_if_remote(v) for k, v in kwargs.items()
-    }
+    new_kwargs = {k: sync_load_from_blob_store_if_remote(v) for k, v in kwargs.items()}
 
     if deep:
         args = new_args
@@ -141,4 +142,4 @@ def sync_load_args(
             else:
                 new_kwargs[k] = v
 
-    return new_args, new_kwargs
+    return new_args, new_kwargs
diff --git a/agents-api/agents_api/routers/sessions/chat.py b/agents-api/agents_api/routers/sessions/chat.py
@@ -1,10 +1,7 @@
 from datetime import datetime
 from typing import Annotated, Callable, Optional
-from datetime import datetime
-from typing import Annotated, Callable, Optional
 from uuid import UUID, uuid4
 
-from fastapi import BackgroundTasks, Depends, Header, HTTPException, status
 from fastapi import BackgroundTasks, Depends, Header, HTTPException, status
 from starlette.status import HTTP_201_CREATED
 
@@ -22,12 +19,10 @@
 from ...common.utils.template import render_template
 from ...dependencies.developer_id import get_developer_data
 from ...env import max_free_sessions
-from ...env import max_free_sessions
 from ...models.chat.gather_messages import gather_messages
 from ...models.chat.prepare_chat_context import prepare_chat_context
 from ...models.entry.create_entries import create_entries
 from ...models.session.count_sessions import count_sessions as count_sessions_query
-from ...models.session.count_sessions import count_sessions as count_sessions_query
 from .metrics import total_tokens_per_user
 from .router import router