From b11247f219f1d0ca8b38c4f30843216afdf5bb11 Mon Sep 17 00:00:00 2001 From: Ahmad-mtos Date: Sat, 14 Dec 2024 19:15:56 +0000 Subject: [PATCH] refactor: Lint agents-api (CI) --- .../agents_api/queries/agent/create_agent.py | 6 ++-- .../queries/agent/create_or_update_agent.py | 24 +++++++------- .../agents_api/queries/agent/delete_agent.py | 18 +++++----- .../agents_api/queries/agent/get_agent.py | 15 +++++---- .../agents_api/queries/agent/list_agents.py | 33 +++++++++---------- .../agents_api/queries/agent/patch_agent.py | 21 ++++++------ .../agents_api/queries/agent/update_agent.py | 23 ++++++------- agents-api/agents_api/queries/utils.py | 22 +++++++------ 8 files changed, 83 insertions(+), 79 deletions(-) diff --git a/agents-api/agents_api/queries/agent/create_agent.py b/agents-api/agents_api/queries/agent/create_agent.py index 30d73d179..52a0a22f8 100644 --- a/agents-api/agents_api/queries/agent/create_agent.py +++ b/agents-api/agents_api/queries/agent/create_agent.py @@ -16,8 +16,8 @@ from ...metrics.counters import increase_counter from ..utils import ( generate_canonical_name, - pg_query, partialclass, + pg_query, rewrap_exceptions, wrap_in_class, ) @@ -91,7 +91,9 @@ def create_agent( ) # Convert default_settings to dict if it exists - default_settings = data.default_settings.model_dump() if data.default_settings else None + default_settings = ( + data.default_settings.model_dump() if data.default_settings else None + ) # Set default values data.metadata = data.metadata or None diff --git a/agents-api/agents_api/queries/agent/create_or_update_agent.py b/agents-api/agents_api/queries/agent/create_or_update_agent.py index e403c7bcf..c93a965a5 100644 --- a/agents-api/agents_api/queries/agent/create_or_update_agent.py +++ b/agents-api/agents_api/queries/agent/create_or_update_agent.py @@ -6,29 +6,30 @@ from typing import Any, TypeVar from uuid import UUID -from ...autogen.openapi_model import Agent, CreateOrUpdateAgentRequest +from beartype import beartype from fastapi import HTTPException +from psycopg import errors as psycopg_errors + +from ...autogen.openapi_model import Agent, CreateOrUpdateAgentRequest from ...metrics.counters import increase_counter from ..utils import ( generate_canonical_name, - pg_query, partialclass, + pg_query, rewrap_exceptions, wrap_in_class, ) -from beartype import beartype -from psycopg import errors as psycopg_errors - ModelT = TypeVar("ModelT", bound=Any) T = TypeVar("T") + @rewrap_exceptions( { psycopg_errors.ForeignKeyViolation: partialclass( - HTTPException, + HTTPException, status_code=404, - detail="The specified developer does not exist." + detail="The specified developer does not exist.", ) } ) @@ -42,10 +43,7 @@ @increase_counter("create_or_update_agent") @beartype def create_or_update_agent_query( - *, - agent_id: UUID, - developer_id: UUID, - data: CreateOrUpdateAgentRequest + *, agent_id: UUID, developer_id: UUID, data: CreateOrUpdateAgentRequest ) -> tuple[list[str], dict]: """ Constructs the SQL queries to create a new agent or update an existing agent's details. @@ -67,7 +65,9 @@ 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 + default_settings = ( + data.default_settings.model_dump() if data.default_settings else None + ) # Set default values data.metadata = data.metadata or None diff --git a/agents-api/agents_api/queries/agent/delete_agent.py b/agents-api/agents_api/queries/agent/delete_agent.py index 4bd14f8ec..1d01daa20 100644 --- a/agents-api/agents_api/queries/agent/delete_agent.py +++ b/agents-api/agents_api/queries/agent/delete_agent.py @@ -6,28 +6,30 @@ 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 ...autogen.openapi_model import ResourceDeletedResponse +from ...common.utils.datetime import utcnow from ...metrics.counters import increase_counter from ..utils import ( - pg_query, partialclass, + pg_query, rewrap_exceptions, wrap_in_class, ) -from beartype import beartype -from psycopg import errors as psycopg_errors -from ...autogen.openapi_model import ResourceDeletedResponse -from ...common.utils.datetime import utcnow ModelT = TypeVar("ModelT", bound=Any) T = TypeVar("T") + @rewrap_exceptions( { psycopg_errors.ForeignKeyViolation: partialclass( - HTTPException, + HTTPException, status_code=404, - detail="The specified developer does not exist." + detail="The specified developer does not exist.", ) } # TODO: Add more exceptions @@ -83,7 +85,7 @@ def delete_agent_query(*, agent_id: UUID, developer_id: UUID) -> tuple[list[str] -- Delete the agent DELETE FROM agents WHERE agent_id = %(agent_id)s AND developer_id = %(developer_id)s; - """ + """, ] params = { diff --git a/agents-api/agents_api/queries/agent/get_agent.py b/agents-api/agents_api/queries/agent/get_agent.py index e5368eea1..982849f3a 100644 --- a/agents-api/agents_api/queries/agent/get_agent.py +++ b/agents-api/agents_api/queries/agent/get_agent.py @@ -6,28 +6,29 @@ 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 ...autogen.openapi_model import Agent from ...metrics.counters import increase_counter from ..utils import ( - pg_query, partialclass, + pg_query, rewrap_exceptions, wrap_in_class, ) -from beartype import beartype -from psycopg import errors as psycopg_errors - -from ...autogen.openapi_model import Agent ModelT = TypeVar("ModelT", bound=Any) T = TypeVar("T") + @rewrap_exceptions( { psycopg_errors.ForeignKeyViolation: partialclass( - HTTPException, + HTTPException, status_code=404, - detail="The specified developer does not exist." + detail="The specified developer does not exist.", ) } # TODO: Add more exceptions diff --git a/agents-api/agents_api/queries/agent/list_agents.py b/agents-api/agents_api/queries/agent/list_agents.py index db46704cf..a4332372f 100644 --- a/agents-api/agents_api/queries/agent/list_agents.py +++ b/agents-api/agents_api/queries/agent/list_agents.py @@ -6,28 +6,29 @@ from typing import Any, Literal, TypeVar from uuid import UUID +from beartype import beartype from fastapi import HTTPException +from psycopg import errors as psycopg_errors + +from ...autogen.openapi_model import Agent from ...metrics.counters import increase_counter from ..utils import ( - pg_query, partialclass, + pg_query, rewrap_exceptions, wrap_in_class, ) -from beartype import beartype -from psycopg import errors as psycopg_errors - -from ...autogen.openapi_model import Agent ModelT = TypeVar("ModelT", bound=Any) T = TypeVar("T") + @rewrap_exceptions( { psycopg_errors.ForeignKeyViolation: partialclass( - HTTPException, + HTTPException, status_code=404, - detail="The specified developer does not exist." + detail="The specified developer does not exist.", ) } # TODO: Add more exceptions @@ -47,7 +48,7 @@ def list_agents_query( ) -> tuple[str, dict]: """ Constructs query to list agents for a developer with pagination. - + Args: developer_id: UUID of the developer limit: Maximum number of records to return @@ -55,7 +56,7 @@ def list_agents_query( sort_by: Field to sort by direction: Sort direction ('asc' or 'desc') metadata_filter: Optional metadata filters - + Returns: Tuple of (query, params) """ @@ -67,7 +68,7 @@ def list_agents_query( metadata_clause = "" if metadata_filter: metadata_clause = "AND metadata @> %(metadata_filter)s::jsonb" - + query = f""" SELECT agent_id, @@ -87,14 +88,10 @@ def list_agents_query( ORDER BY {sort_by} {direction} LIMIT %(limit)s OFFSET %(offset)s; """ - - params = { - "developer_id": developer_id, - "limit": limit, - "offset": offset - } - + + params = {"developer_id": developer_id, "limit": limit, "offset": offset} + if metadata_filter: params["metadata_filter"] = metadata_filter - + return query, params diff --git a/agents-api/agents_api/queries/agent/patch_agent.py b/agents-api/agents_api/queries/agent/patch_agent.py index 5f935d49b..74be99df8 100644 --- a/agents-api/agents_api/queries/agent/patch_agent.py +++ b/agents-api/agents_api/queries/agent/patch_agent.py @@ -6,27 +6,29 @@ from typing import Any, TypeVar from uuid import UUID -from ...autogen.openapi_model import PatchAgentRequest, ResourceUpdatedResponse +from beartype import beartype from fastapi import HTTPException +from psycopg import errors as psycopg_errors + +from ...autogen.openapi_model import PatchAgentRequest, ResourceUpdatedResponse from ...metrics.counters import increase_counter from ..utils import ( - pg_query, partialclass, + pg_query, rewrap_exceptions, wrap_in_class, ) -from beartype import beartype -from psycopg import errors as psycopg_errors ModelT = TypeVar("ModelT", bound=Any) T = TypeVar("T") + @rewrap_exceptions( { psycopg_errors.ForeignKeyViolation: partialclass( - HTTPException, + HTTPException, status_code=404, - detail="The specified developer does not exist." + detail="The specified developer does not exist.", ) } # TODO: Add more exceptions @@ -41,10 +43,7 @@ @increase_counter("patch_agent") @beartype def patch_agent_query( - *, - agent_id: UUID, - developer_id: UUID, - data: PatchAgentRequest + *, agent_id: UUID, developer_id: UUID, data: PatchAgentRequest ) -> tuple[str, dict]: """ Constructs the SQL query to partially update an agent's details. @@ -67,7 +66,7 @@ def patch_agent_query( params[key] = value set_clause = ", ".join(set_clauses) - + query = f""" UPDATE agents SET {set_clause} diff --git a/agents-api/agents_api/queries/agent/update_agent.py b/agents-api/agents_api/queries/agent/update_agent.py index e26667874..e0ed4a46d 100644 --- a/agents-api/agents_api/queries/agent/update_agent.py +++ b/agents-api/agents_api/queries/agent/update_agent.py @@ -6,27 +6,29 @@ from typing import Any, TypeVar from uuid import UUID -from ...autogen.openapi_model import ResourceUpdatedResponse, UpdateAgentRequest +from beartype import beartype from fastapi import HTTPException +from psycopg import errors as psycopg_errors + +from ...autogen.openapi_model import ResourceUpdatedResponse, UpdateAgentRequest from ...metrics.counters import increase_counter from ..utils import ( - pg_query, partialclass, + pg_query, rewrap_exceptions, wrap_in_class, ) -from beartype import beartype -from psycopg import errors as psycopg_errors ModelT = TypeVar("ModelT", bound=Any) T = TypeVar("T") + @rewrap_exceptions( { psycopg_errors.ForeignKeyViolation: partialclass( - HTTPException, + HTTPException, status_code=404, - detail="The specified developer does not exist." + detail="The specified developer does not exist.", ) } # TODO: Add more exceptions @@ -41,10 +43,7 @@ @increase_counter("update_agent") @beartype def update_agent_query( - *, - agent_id: UUID, - developer_id: UUID, - data: UpdateAgentRequest + *, agent_id: UUID, developer_id: UUID, data: UpdateAgentRequest ) -> tuple[str, dict]: """ Constructs the SQL query to fully update an agent's details. @@ -57,7 +56,9 @@ def update_agent_query( Returns: tuple[str, dict]: A tuple containing the SQL query and its parameters. """ - fields = ", ".join([f"{key} = %({key})s" for key in data.model_dump(exclude_unset=True).keys()]) + fields = ", ".join( + [f"{key} = %({key})s" for key in data.model_dump(exclude_unset=True).keys()] + ) params = {key: value for key, value in data.model_dump(exclude_unset=True).items()} query = f""" diff --git a/agents-api/agents_api/queries/utils.py b/agents-api/agents_api/queries/utils.py index 704085a76..ba0e50fc0 100644 --- a/agents-api/agents_api/queries/utils.py +++ b/agents-api/agents_api/queries/utils.py @@ -1,23 +1,23 @@ +import inspect import re import time -from typing import Awaitable, Callable, ParamSpec, Type, TypeVar -import inspect -from fastapi import HTTPException -import pandas as pd -from pydantic import BaseModel from functools import partialmethod, wraps +from typing import Any, Awaitable, Callable, ParamSpec, Type, TypeVar + +import pandas as pd +import sqlglot from asyncpg import Record -from requests.exceptions import ConnectionError, Timeout +from fastapi import HTTPException from httpcore import NetworkError, TimeoutException from httpx import RequestError -import sqlglot - -from typing import Any +from pydantic import BaseModel +from requests.exceptions import ConnectionError, Timeout P = ParamSpec("P") T = TypeVar("T") ModelT = TypeVar("ModelT", bound=BaseModel) + def generate_canonical_name(name: str) -> str: """Convert a display name to a canonical name. Example: "My Cool Agent!" -> "my_cool_agent" @@ -32,6 +32,7 @@ def generate_canonical_name(name: str) -> str: return canonical + def partialclass(cls, *args, **kwargs): cls_signature = inspect.signature(cls) bound = cls_signature.bind_partial(*args, **kwargs) @@ -145,6 +146,7 @@ def wrapper(*args: P.args, **kwargs: P.kwargs) -> T: return decorator + def pg_query( func: Callable[P, tuple[str | list[str | None], dict]] | None = None, debug: bool | None = None, @@ -251,4 +253,4 @@ async def wrapper( if func is not None and callable(func): return pg_query_dec(func) - return pg_query_dec \ No newline at end of file + return pg_query_dec