diff --git a/agents-api/agents_api/activities/demo.py b/agents-api/agents_api/activities/demo.py
index 797ef6c90..f6d63f206 100644
--- a/agents-api/agents_api/activities/demo.py
+++ b/agents-api/agents_api/activities/demo.py
@@ -1,3 +1,5 @@
+from typing import Callable
+
from temporalio import activity
from ..env import testing
@@ -12,6 +14,6 @@ async def mock_demo_activity(a: int, b: int) -> int:
return a + b
-demo_activity = activity.defn(name="demo_activity")(
+demo_activity: Callable[[int, int], int] = activity.defn(name="demo_activity")(
demo_activity if not testing else mock_demo_activity
)
diff --git a/agents-api/agents_api/activities/embed_docs.py b/agents-api/agents_api/activities/embed_docs.py
index 1dd3793b0..3222ff9e7 100644
--- a/agents-api/agents_api/activities/embed_docs.py
+++ b/agents-api/agents_api/activities/embed_docs.py
@@ -1,8 +1,8 @@
from beartype import beartype
from temporalio import activity
+from ..clients import cozo
from ..clients import embed as embedder
-from ..clients.cozo import get_cozo_client
from ..env import testing
from ..models.docs.embed_snippets import embed_snippets as embed_snippets_query
from .types import EmbedDocsPayload
@@ -28,7 +28,7 @@ async def embed_docs(payload: EmbedDocsPayload, cozo_client=None) -> None:
doc_id=payload.doc_id,
snippet_indices=indices,
embeddings=embeddings,
- client=cozo_client or get_cozo_client(),
+ client=cozo_client or cozo.get_cozo_client(),
)
diff --git a/agents-api/agents_api/activities/logger.py b/agents-api/agents_api/activities/logger.py
index f3c31fece..ed18019dc 100644
--- a/agents-api/agents_api/activities/logger.py
+++ b/agents-api/agents_api/activities/logger.py
@@ -1,7 +1,8 @@
import logging
+from typing import TextIO
-logger = logging.getLogger(__name__)
-h = logging.StreamHandler()
-fmt = logging.Formatter("[%(asctime)s/%(levelname)s] - %(message)s")
+logger: logging.Logger = logging.getLogger(__name__)
+h: logging.StreamHandler[TextIO] = logging.StreamHandler()
+fmt: logging.Formatter = logging.Formatter("[%(asctime)s/%(levelname)s] - %(message)s")
h.setFormatter(fmt)
logger.addHandler(h)
diff --git a/agents-api/agents_api/activities/summarization.py b/agents-api/agents_api/activities/summarization.py
index 90a89fde5..a02237bf7 100644
--- a/agents-api/agents_api/activities/summarization.py
+++ b/agents-api/agents_api/activities/summarization.py
@@ -12,11 +12,11 @@
# TODO: remove stubs
-def entries_summarization_query(*args, **kwargs):
+def entries_summarization_query(*args, **kwargs) -> pd.DataFrame:
return pd.DataFrame()
-def get_toplevel_entries_query(*args, **kwargs):
+def get_toplevel_entries_query(*args, **kwargs) -> pd.DataFrame:
return pd.DataFrame()
diff --git a/agents-api/agents_api/activities/task_steps/yield_step.py b/agents-api/agents_api/activities/task_steps/yield_step.py
index 03462c18f..dfdb6fa1d 100644
--- a/agents-api/agents_api/activities/task_steps/yield_step.py
+++ b/agents-api/agents_api/activities/task_steps/yield_step.py
@@ -1,3 +1,5 @@
+from typing import Callable
+
from beartype import beartype
from temporalio import activity
@@ -33,8 +35,8 @@ async def yield_step(context: StepContext[YieldStep]) -> StepOutcome:
# Note: This is here just for clarity. We could have just imported yield_step directly
# They do the same thing, so we dont need to mock the yield_step function
-mock_yield_step = yield_step
+mock_yield_step: Callable[[StepContext], StepOutcome] = yield_step
-yield_step = activity.defn(name="yield_step")(
+yield_step: Callable[[StepContext], StepOutcome] = activity.defn(name="yield_step")(
yield_step if not testing else mock_yield_step
)
diff --git a/agents-api/agents_api/clients/cozo.py b/agents-api/agents_api/clients/cozo.py
index d582e5db4..e2991c9d8 100644
--- a/agents-api/agents_api/clients/cozo.py
+++ b/agents-api/agents_api/clients/cozo.py
@@ -1,14 +1,16 @@
+from typing import Any, Dict
+
from pycozo.client import Client
from ..env import cozo_auth, cozo_host
from ..web import app
-options = {"host": cozo_host}
+options: Dict[str, str] = {"host": cozo_host}
if cozo_auth:
options.update({"auth": cozo_auth})
-def get_cozo_client():
+def get_cozo_client() -> Any:
client = getattr(app.state, "cozo_client", Client("http", options=options))
if not hasattr(app.state, "cozo_client"):
app.state.cozo_client = client
diff --git a/agents-api/agents_api/clients/litellm.py b/agents-api/agents_api/clients/litellm.py
index 5c3a2e28e..4c78e2876 100644
--- a/agents-api/agents_api/clients/litellm.py
+++ b/agents-api/agents_api/clients/litellm.py
@@ -1,11 +1,14 @@
from functools import wraps
+from typing import List, TypeVar
from litellm import acompletion as _acompletion
from litellm.utils import CustomStreamWrapper, ModelResponse
from ..env import litellm_master_key, litellm_url
-__all__ = ["acompletion"]
+_RWrapped = TypeVar("_RWrapped")
+
+__all__: List[str] = ["acompletion"]
@wraps(_acompletion)
diff --git a/agents-api/agents_api/common/exceptions/__init__.py b/agents-api/agents_api/common/exceptions/__init__.py
index fa0016b4e..a38a546a2 100644
--- a/agents-api/agents_api/common/exceptions/__init__.py
+++ b/agents-api/agents_api/common/exceptions/__init__.py
@@ -11,6 +11,6 @@
class BaseCommonException(Exception):
- def __init__(self, msg: str, http_code: int):
+ def __init__(self, msg: str, http_code: int) -> None:
super().__init__(msg)
self.http_code = http_code
diff --git a/agents-api/agents_api/common/utils/cozo.py b/agents-api/agents_api/common/utils/cozo.py
index db8f336f0..a8195a1ba 100644
--- a/agents-api/agents_api/common/utils/cozo.py
+++ b/agents-api/agents_api/common/utils/cozo.py
@@ -8,7 +8,7 @@
from pycozo import Client
# Define a mock client for testing purposes, simulating Cozo API client behavior.
-_fake_client = SimpleNamespace()
+_fake_client: SimpleNamespace = SimpleNamespace()
# Lambda function to process and mutate data dictionaries using the Cozo client's internal method. This is a workaround to access protected member functions for testing.
_fake_client._process_mutate_data_dict = lambda data: (
Client._process_mutate_data_dict(_fake_client, data)
@@ -20,5 +20,5 @@
)
-def uuid_int_list_to_uuid4(data):
+def uuid_int_list_to_uuid4(data) -> UUID:
return UUID(bytes=b"".join([i.to_bytes(1, "big") for i in data]))
diff --git a/agents-api/agents_api/common/utils/json.py b/agents-api/agents_api/common/utils/json.py
index 9beff6049..3157af9c8 100644
--- a/agents-api/agents_api/common/utils/json.py
+++ b/agents-api/agents_api/common/utils/json.py
@@ -10,7 +10,7 @@
class CustomJSONEncoder(json.JSONEncoder):
"""A custom JSON encoder subclass that handles None values and UUIDs for JSON serialization. It allows specifying a default value for None objects during initialization."""
- def __init__(self, *args, **kwargs):
+ def __init__(self, *args, **kwargs) -> None:
"""Initializes the custom JSON encoder.
Parameters:
*args: Variable length argument list.
@@ -19,7 +19,7 @@ def __init__(self, *args, **kwargs):
self._default_empty_value = kwargs.pop("default_empty_value")
super().__init__(*args, **kwargs)
- def encode(self, o):
+ def encode(self, o) -> str:
"""Encodes the given object into a JSON formatted string.
Parameters:
o: The object to encode.
@@ -27,7 +27,7 @@ def encode(self, o):
# Use the overridden default method for serialization before encoding
return super().encode(self.default(o))
- def default(self, obj):
+ def default(self, obj) -> Any:
"""Provides a default serialization for objects that the standard JSON encoder cannot serialize.
Parameters:
obj: The object to serialize.
diff --git a/agents-api/agents_api/common/utils/template.py b/agents-api/agents_api/common/utils/template.py
index d806ddb6d..1fcc143b3 100644
--- a/agents-api/agents_api/common/utils/template.py
+++ b/agents-api/agents_api/common/utils/template.py
@@ -1,14 +1,16 @@
+from typing import List
+
import arrow
from jinja2.sandbox import ImmutableSandboxedEnvironment
from jinja2schema import infer, to_json_schema
from jsonschema import validate
-__all__ = [
+__all__: List[str] = [
"render_template",
]
# jinja environment
-jinja_env = ImmutableSandboxedEnvironment(
+jinja_env: ImmutableSandboxedEnvironment = ImmutableSandboxedEnvironment(
autoescape=False,
trim_blocks=True,
lstrip_blocks=True,
diff --git a/agents-api/agents_api/dependencies/auth.py b/agents-api/agents_api/dependencies/auth.py
index 0054cb1cc..e5e22995b 100644
--- a/agents-api/agents_api/dependencies/auth.py
+++ b/agents-api/agents_api/dependencies/auth.py
@@ -1,13 +1,18 @@
+from typing import Any
+
from fastapi import HTTPException, Security
from fastapi.security.api_key import APIKeyHeader
from starlette.status import HTTP_403_FORBIDDEN
from ..env import api_key, api_key_header_name
-api_key_header = APIKeyHeader(name=api_key_header_name, auto_error=False)
+api_key_header: Any = APIKeyHeader(name=api_key_header_name, auto_error=False)
-async def get_api_key(user_api_key: str = Security(api_key_header)):
+async def get_api_key(
+ user_api_key: str = Security(api_key_header),
+) -> str:
+ user_api_key = str(user_api_key)
user_api_key = (user_api_key or "").replace("Bearer ", "").strip()
if user_api_key != api_key:
diff --git a/agents-api/agents_api/env.py b/agents-api/agents_api/env.py
index afdc81e1e..64d9082ef 100644
--- a/agents-api/agents_api/env.py
+++ b/agents-api/agents_api/env.py
@@ -5,11 +5,12 @@
import random
from pprint import pprint
+from typing import Any, Dict
from environs import Env
# Initialize the Env object for environment variable parsing.
-env = Env()
+env: Any = Env()
# Debug
@@ -30,7 +31,7 @@
# Auth
# ----
-_random_generated_key = "".join(str(random.randint(0, 9)) for _ in range(32))
+_random_generated_key: str = "".join(str(random.randint(0, 9)) for _ in range(32))
api_key: str = env.str("AGENTS_API_KEY", _random_generated_key)
if api_key == _random_generated_key:
@@ -65,12 +66,12 @@
temporal_namespace: str = env.str("TEMPORAL_NAMESPACE", default="default")
temporal_client_cert: str = env.str("TEMPORAL_CLIENT_CERT", default=None)
temporal_private_key: str = env.str("TEMPORAL_PRIVATE_KEY", default=None)
-temporal_endpoint = env.str("TEMPORAL_ENDPOINT", default="localhost:7233")
-temporal_task_queue = env.str("TEMPORAL_TASK_QUEUE", default="julep-task-queue")
+temporal_endpoint: Any = env.str("TEMPORAL_ENDPOINT", default="localhost:7233")
+temporal_task_queue: Any = env.str("TEMPORAL_TASK_QUEUE", default="julep-task-queue")
# Consolidate environment variables
-environment = dict(
+environment: Dict[str, Any] = dict(
debug=debug,
cozo_host=cozo_host,
cozo_auth=cozo_auth,
diff --git a/agents-api/agents_api/exceptions.py b/agents-api/agents_api/exceptions.py
index 2ccc5a67f..fbb8f00f8 100644
--- a/agents-api/agents_api/exceptions.py
+++ b/agents-api/agents_api/exceptions.py
@@ -3,17 +3,17 @@ class AgentsBaseException(Exception):
class ModelNotSupportedError(AgentsBaseException):
- def __init__(self, model_name):
+ def __init__(self, model_name) -> None:
super().__init__(f"model {model_name} is not supported")
class PromptTooBigError(AgentsBaseException):
- def __init__(self, token_count, max_tokens):
+ def __init__(self, token_count, max_tokens) -> None:
super().__init__(
f"prompt is too big, {token_count} tokens provided, exceeds maximum of {max_tokens}"
)
class UnknownTokenizerError(AgentsBaseException):
- def __init__(self):
+ def __init__(self) -> None:
super().__init__("unknown tokenizer")
diff --git a/agents-api/agents_api/model_registry.py b/agents-api/agents_api/model_registry.py
index 99ae66ea3..0120cc205 100644
--- a/agents-api/agents_api/model_registry.py
+++ b/agents-api/agents_api/model_registry.py
@@ -66,7 +66,7 @@
}
-DISCONTINUED_MODELS = {
+DISCONTINUED_MODELS: Dict[str, int] = {
"code-davinci-002": 8001,
"code-davinci-001": 8001,
"code-cushman-002": 2048,
@@ -84,9 +84,14 @@
"claude-3-haiku-20240307": 180000,
}
-OPENAI_MODELS = {**GPT4_MODELS, **TURBO_MODELS, **GPT3_5_MODELS, **GPT3_MODELS}
+OPENAI_MODELS: Dict[str, int] = {
+ **GPT4_MODELS,
+ **TURBO_MODELS,
+ **GPT3_5_MODELS,
+ **GPT3_MODELS,
+}
-LOCAL_MODELS = {
+LOCAL_MODELS: Dict[str, int] = {
"gpt-4o": 32768,
"gpt-4o-awq": 32768,
"TinyLlama/TinyLlama_v1.1": 2048,
@@ -95,13 +100,13 @@
"OpenPipe/Hermes-2-Theta-Llama-3-8B-32k": 32768,
}
-LOCAL_MODELS_WITH_TOOL_CALLS = {
+LOCAL_MODELS_WITH_TOOL_CALLS: Dict[str, int] = {
"OpenPipe/Hermes-2-Theta-Llama-3-8B-32k": 32768,
"julep-ai/Hermes-2-Theta-Llama-3-8B": 8192,
}
-OLLAMA_MODELS = {
+OLLAMA_MODELS: Dict[str, int] = {
"llama2": 4096,
}
-CHAT_MODELS = {**GPT4_MODELS, **TURBO_MODELS, **CLAUDE_MODELS}
+CHAT_MODELS: Dict[str, int] = {**GPT4_MODELS, **TURBO_MODELS, **CLAUDE_MODELS}
diff --git a/agents-api/agents_api/models/agent/create_agent.py b/agents-api/agents_api/models/agent/create_agent.py
index a4b408c78..73be4ec77 100644
--- a/agents-api/agents_api/models/agent/create_agent.py
+++ b/agents-api/agents_api/models/agent/create_agent.py
@@ -3,6 +3,7 @@
It includes functions to construct and execute datalog queries for inserting new agent records.
"""
+from typing import Any, TypeVar
from uuid import UUID, uuid4
from beartype import beartype
@@ -20,6 +21,9 @@
wrap_in_class,
)
+ModelT = TypeVar("ModelT", bound=Any)
+T = TypeVar("T")
+
@rewrap_exceptions(
{
diff --git a/agents-api/agents_api/models/agent/create_or_update_agent.py b/agents-api/agents_api/models/agent/create_or_update_agent.py
index fb80f6dd7..64b008d44 100644
--- a/agents-api/agents_api/models/agent/create_or_update_agent.py
+++ b/agents-api/agents_api/models/agent/create_or_update_agent.py
@@ -3,6 +3,7 @@
It includes functions to construct and execute datalog queries for inserting new agent records.
"""
+from typing import Any, TypeVar
from uuid import UUID
from beartype import beartype
@@ -20,6 +21,9 @@
wrap_in_class,
)
+ModelT = TypeVar("ModelT", bound=Any)
+T = TypeVar("T")
+
@rewrap_exceptions(
{
diff --git a/agents-api/agents_api/models/agent/delete_agent.py b/agents-api/agents_api/models/agent/delete_agent.py
index e1efd7333..409c755d3 100644
--- a/agents-api/agents_api/models/agent/delete_agent.py
+++ b/agents-api/agents_api/models/agent/delete_agent.py
@@ -2,6 +2,7 @@
This module contains the implementation of the delete_agent_query function, which is responsible for deleting an agent and its related default settings from the CozoDB database.
"""
+from typing import Any, TypeVar
from uuid import UUID
from beartype import beartype
@@ -20,6 +21,9 @@
wrap_in_class,
)
+ModelT = TypeVar("ModelT", bound=Any)
+T = TypeVar("T")
+
@rewrap_exceptions(
{
diff --git a/agents-api/agents_api/models/agent/get_agent.py b/agents-api/agents_api/models/agent/get_agent.py
index ceef527f3..c977fa614 100644
--- a/agents-api/agents_api/models/agent/get_agent.py
+++ b/agents-api/agents_api/models/agent/get_agent.py
@@ -1,3 +1,4 @@
+from typing import Any, TypeVar
from uuid import UUID
from beartype import beartype
@@ -15,6 +16,9 @@
wrap_in_class,
)
+ModelT = TypeVar("ModelT", bound=Any)
+T = TypeVar("T")
+
@rewrap_exceptions(
{
diff --git a/agents-api/agents_api/models/agent/list_agents.py b/agents-api/agents_api/models/agent/list_agents.py
index 5d291f654..5266659d7 100644
--- a/agents-api/agents_api/models/agent/list_agents.py
+++ b/agents-api/agents_api/models/agent/list_agents.py
@@ -1,4 +1,4 @@
-from typing import Any, Literal
+from typing import Any, Literal, TypeVar
from uuid import UUID
from beartype import beartype
@@ -16,6 +16,9 @@
wrap_in_class,
)
+ModelT = TypeVar("ModelT", bound=Any)
+T = TypeVar("T")
+
@rewrap_exceptions(
{
diff --git a/agents-api/agents_api/models/agent/patch_agent.py b/agents-api/agents_api/models/agent/patch_agent.py
index 336e33a91..72fdc7811 100644
--- a/agents-api/agents_api/models/agent/patch_agent.py
+++ b/agents-api/agents_api/models/agent/patch_agent.py
@@ -1,3 +1,4 @@
+from typing import Any, TypeVar
from uuid import UUID
from beartype import beartype
@@ -17,6 +18,9 @@
wrap_in_class,
)
+ModelT = TypeVar("ModelT", bound=Any)
+T = TypeVar("T")
+
@rewrap_exceptions(
{
diff --git a/agents-api/agents_api/models/agent/update_agent.py b/agents-api/agents_api/models/agent/update_agent.py
index 4dd489333..be7e9ea21 100644
--- a/agents-api/agents_api/models/agent/update_agent.py
+++ b/agents-api/agents_api/models/agent/update_agent.py
@@ -1,3 +1,4 @@
+from typing import Any, TypeVar
from uuid import UUID
from beartype import beartype
@@ -16,6 +17,9 @@
wrap_in_class,
)
+ModelT = TypeVar("ModelT", bound=Any)
+T = TypeVar("T")
+
@rewrap_exceptions(
{
diff --git a/agents-api/agents_api/models/chat/gather_messages.py b/agents-api/agents_api/models/chat/gather_messages.py
index 2a3c0eca1..f8e08632d 100644
--- a/agents-api/agents_api/models/chat/gather_messages.py
+++ b/agents-api/agents_api/models/chat/gather_messages.py
@@ -1,3 +1,4 @@
+from typing import TypeVar
from uuid import UUID
from beartype import beartype
@@ -18,6 +19,8 @@
rewrap_exceptions,
)
+T = TypeVar("T")
+
@rewrap_exceptions(
{
diff --git a/agents-api/agents_api/models/chat/prepare_chat_context.py b/agents-api/agents_api/models/chat/prepare_chat_context.py
index 5ca93b8c0..742038535 100644
--- a/agents-api/agents_api/models/chat/prepare_chat_context.py
+++ b/agents-api/agents_api/models/chat/prepare_chat_context.py
@@ -1,3 +1,4 @@
+from typing import Any, TypeVar
from uuid import UUID
from beartype import beartype
@@ -17,6 +18,9 @@
wrap_in_class,
)
+ModelT = TypeVar("ModelT", bound=Any)
+T = TypeVar("T")
+
@rewrap_exceptions(
{
diff --git a/agents-api/agents_api/models/developer/get_developer.py b/agents-api/agents_api/models/developer/get_developer.py
index c78517613..31ade5334 100644
--- a/agents-api/agents_api/models/developer/get_developer.py
+++ b/agents-api/agents_api/models/developer/get_developer.py
@@ -1,5 +1,6 @@
"""Module for retrieving document snippets from the CozoDB based on document IDs."""
+from typing import Any, TypeVar
from uuid import UUID
from beartype import beartype
@@ -16,6 +17,9 @@
wrap_in_class,
)
+ModelT = TypeVar("ModelT", bound=Any)
+T = TypeVar("T")
+
@rewrap_exceptions({QueryException: partialclass(HTTPException, status_code=401)})
@cozo_query
diff --git a/agents-api/agents_api/models/docs/create_doc.py b/agents-api/agents_api/models/docs/create_doc.py
index 3a86ad0e5..ee26df484 100644
--- a/agents-api/agents_api/models/docs/create_doc.py
+++ b/agents-api/agents_api/models/docs/create_doc.py
@@ -1,4 +1,4 @@
-from typing import Literal
+from typing import Any, Literal, TypeVar
from uuid import UUID, uuid4
from beartype import beartype
@@ -17,6 +17,9 @@
wrap_in_class,
)
+ModelT = TypeVar("ModelT", bound=Any)
+T = TypeVar("T")
+
@rewrap_exceptions(
{
diff --git a/agents-api/agents_api/models/docs/delete_doc.py b/agents-api/agents_api/models/docs/delete_doc.py
index e7a02f3d9..c02705756 100644
--- a/agents-api/agents_api/models/docs/delete_doc.py
+++ b/agents-api/agents_api/models/docs/delete_doc.py
@@ -1,3 +1,4 @@
+from typing import Any, TypeVar
from uuid import UUID
from beartype import beartype
@@ -16,6 +17,9 @@
wrap_in_class,
)
+ModelT = TypeVar("ModelT", bound=Any)
+T = TypeVar("T")
+
@rewrap_exceptions(
{
diff --git a/agents-api/agents_api/models/docs/embed_snippets.py b/agents-api/agents_api/models/docs/embed_snippets.py
index a701ef1be..e810d0379 100644
--- a/agents-api/agents_api/models/docs/embed_snippets.py
+++ b/agents-api/agents_api/models/docs/embed_snippets.py
@@ -1,5 +1,6 @@
"""Module for embedding documents in the cozodb database. Contains functions to update document embeddings."""
+from typing import Any, TypeVar
from uuid import UUID
from beartype import beartype
@@ -18,6 +19,9 @@
wrap_in_class,
)
+ModelT = TypeVar("ModelT", bound=Any)
+T = TypeVar("T")
+
@rewrap_exceptions(
{
diff --git a/agents-api/agents_api/models/docs/get_doc.py b/agents-api/agents_api/models/docs/get_doc.py
index 6c612fef7..84cd181ec 100644
--- a/agents-api/agents_api/models/docs/get_doc.py
+++ b/agents-api/agents_api/models/docs/get_doc.py
@@ -1,5 +1,6 @@
"""Module for retrieving document snippets from the CozoDB based on document IDs."""
+from typing import Any, TypeVar
from uuid import UUID
from beartype import beartype
@@ -16,6 +17,9 @@
wrap_in_class,
)
+ModelT = TypeVar("ModelT", bound=Any)
+T = TypeVar("T")
+
@rewrap_exceptions(
{
diff --git a/agents-api/agents_api/models/docs/list_docs.py b/agents-api/agents_api/models/docs/list_docs.py
index d785cc331..afdf06c2d 100644
--- a/agents-api/agents_api/models/docs/list_docs.py
+++ b/agents-api/agents_api/models/docs/list_docs.py
@@ -1,7 +1,7 @@
"""This module contains functions for querying document-related data from the 'cozodb' database using datalog queries."""
import json
-from typing import Any, Literal
+from typing import Any, Literal, TypeVar
from uuid import UUID
from beartype import beartype
@@ -19,6 +19,9 @@
wrap_in_class,
)
+ModelT = TypeVar("ModelT", bound=Any)
+T = TypeVar("T")
+
@rewrap_exceptions(
{
diff --git a/agents-api/agents_api/models/docs/search_docs_by_embedding.py b/agents-api/agents_api/models/docs/search_docs_by_embedding.py
index 3f7114a23..acebd09cd 100644
--- a/agents-api/agents_api/models/docs/search_docs_by_embedding.py
+++ b/agents-api/agents_api/models/docs/search_docs_by_embedding.py
@@ -1,6 +1,6 @@
"""This module contains functions for searching documents in the CozoDB based on embedding queries."""
-from typing import Literal
+from typing import Any, Literal, TypeVar
from uuid import UUID
from beartype import beartype
@@ -18,6 +18,9 @@
wrap_in_class,
)
+ModelT = TypeVar("ModelT", bound=Any)
+T = TypeVar("T")
+
@rewrap_exceptions(
{
diff --git a/agents-api/agents_api/models/docs/search_docs_by_text.py b/agents-api/agents_api/models/docs/search_docs_by_text.py
index a5e379f24..0662aa84d 100644
--- a/agents-api/agents_api/models/docs/search_docs_by_text.py
+++ b/agents-api/agents_api/models/docs/search_docs_by_text.py
@@ -1,6 +1,6 @@
"""This module contains functions for searching documents in the CozoDB based on embedding queries."""
-from typing import Literal
+from typing import Any, Literal, TypeVar
from uuid import UUID
from beartype import beartype
@@ -18,6 +18,9 @@
wrap_in_class,
)
+ModelT = TypeVar("ModelT", bound=Any)
+T = TypeVar("T")
+
@rewrap_exceptions(
{
diff --git a/agents-api/agents_api/models/entry/create_entries.py b/agents-api/agents_api/models/entry/create_entries.py
index e71a81c7a..e227714d1 100644
--- a/agents-api/agents_api/models/entry/create_entries.py
+++ b/agents-api/agents_api/models/entry/create_entries.py
@@ -1,3 +1,4 @@
+from typing import Any, TypeVar
from uuid import UUID, uuid4
from beartype import beartype
@@ -19,6 +20,9 @@
wrap_in_class,
)
+ModelT = TypeVar("ModelT", bound=Any)
+T = TypeVar("T")
+
@rewrap_exceptions(
{
diff --git a/agents-api/agents_api/models/entry/delete_entries.py b/agents-api/agents_api/models/entry/delete_entries.py
index 1d7cf0386..48c37cd25 100644
--- a/agents-api/agents_api/models/entry/delete_entries.py
+++ b/agents-api/agents_api/models/entry/delete_entries.py
@@ -1,3 +1,4 @@
+from typing import Any, TypeVar
from uuid import UUID
from beartype import beartype
@@ -17,6 +18,9 @@
wrap_in_class,
)
+ModelT = TypeVar("ModelT", bound=Any)
+T = TypeVar("T")
+
@rewrap_exceptions(
{
diff --git a/agents-api/agents_api/models/entry/get_history.py b/agents-api/agents_api/models/entry/get_history.py
index fed658ea5..68fe05979 100644
--- a/agents-api/agents_api/models/entry/get_history.py
+++ b/agents-api/agents_api/models/entry/get_history.py
@@ -1,3 +1,4 @@
+from typing import Any, TypeVar
from uuid import UUID
from beartype import beartype
@@ -16,6 +17,9 @@
wrap_in_class,
)
+ModelT = TypeVar("ModelT", bound=Any)
+T = TypeVar("T")
+
@rewrap_exceptions(
{
diff --git a/agents-api/agents_api/models/entry/list_entries.py b/agents-api/agents_api/models/entry/list_entries.py
index da2341c4c..d3081a9b0 100644
--- a/agents-api/agents_api/models/entry/list_entries.py
+++ b/agents-api/agents_api/models/entry/list_entries.py
@@ -1,4 +1,4 @@
-from typing import Literal
+from typing import Any, Literal, TypeVar
from uuid import UUID
from beartype import beartype
@@ -16,6 +16,9 @@
wrap_in_class,
)
+ModelT = TypeVar("ModelT", bound=Any)
+T = TypeVar("T")
+
@rewrap_exceptions(
{
diff --git a/agents-api/agents_api/models/execution/create_execution.py b/agents-api/agents_api/models/execution/create_execution.py
index c564c343d..b4918065b 100644
--- a/agents-api/agents_api/models/execution/create_execution.py
+++ b/agents-api/agents_api/models/execution/create_execution.py
@@ -1,4 +1,4 @@
-from typing import Annotated
+from typing import Annotated, Any, TypeVar
from uuid import UUID, uuid4
from beartype import beartype
@@ -18,6 +18,9 @@
wrap_in_class,
)
+ModelT = TypeVar("ModelT", bound=Any)
+T = TypeVar("T")
+
@rewrap_exceptions(
{
diff --git a/agents-api/agents_api/models/execution/create_temporal_lookup.py b/agents-api/agents_api/models/execution/create_temporal_lookup.py
index 1867b64b9..21c3005dd 100644
--- a/agents-api/agents_api/models/execution/create_temporal_lookup.py
+++ b/agents-api/agents_api/models/execution/create_temporal_lookup.py
@@ -1,3 +1,4 @@
+from typing import TypeVar
from uuid import UUID, uuid4
from beartype import beartype
@@ -15,6 +16,8 @@
verify_developer_owns_resource_query,
)
+T = TypeVar("T")
+
@rewrap_exceptions(
{
diff --git a/agents-api/agents_api/models/execution/get_execution.py b/agents-api/agents_api/models/execution/get_execution.py
index a0bbe550c..cf9df2e09 100644
--- a/agents-api/agents_api/models/execution/get_execution.py
+++ b/agents-api/agents_api/models/execution/get_execution.py
@@ -1,3 +1,4 @@
+from typing import Any, TypeVar
from uuid import UUID
from beartype import beartype
@@ -13,6 +14,9 @@
wrap_in_class,
)
+ModelT = TypeVar("ModelT", bound=Any)
+T = TypeVar("T")
+
@rewrap_exceptions(
{
diff --git a/agents-api/agents_api/models/execution/get_execution_transition.py b/agents-api/agents_api/models/execution/get_execution_transition.py
index 39d30278c..d418ef5f4 100644
--- a/agents-api/agents_api/models/execution/get_execution_transition.py
+++ b/agents-api/agents_api/models/execution/get_execution_transition.py
@@ -1,3 +1,4 @@
+from typing import Any, TypeVar
from uuid import UUID
from beartype import beartype
@@ -14,6 +15,9 @@
wrap_in_class,
)
+ModelT = TypeVar("ModelT", bound=Any)
+T = TypeVar("T")
+
@rewrap_exceptions(
{
diff --git a/agents-api/agents_api/models/execution/get_paused_execution_token.py b/agents-api/agents_api/models/execution/get_paused_execution_token.py
index 1b1f6e803..d8b0945c1 100644
--- a/agents-api/agents_api/models/execution/get_paused_execution_token.py
+++ b/agents-api/agents_api/models/execution/get_paused_execution_token.py
@@ -1,3 +1,4 @@
+from typing import Any, TypeVar
from uuid import UUID
from beartype import beartype
@@ -13,6 +14,9 @@
wrap_in_class,
)
+ModelT = TypeVar("ModelT", bound=Any)
+T = TypeVar("T")
+
@rewrap_exceptions(
{
diff --git a/agents-api/agents_api/models/execution/get_temporal_workflow_data.py b/agents-api/agents_api/models/execution/get_temporal_workflow_data.py
index 104918f22..bb0a462ef 100644
--- a/agents-api/agents_api/models/execution/get_temporal_workflow_data.py
+++ b/agents-api/agents_api/models/execution/get_temporal_workflow_data.py
@@ -1,3 +1,4 @@
+from typing import Any, TypeVar
from uuid import UUID
from beartype import beartype
@@ -12,6 +13,9 @@
wrap_in_class,
)
+ModelT = TypeVar("ModelT", bound=Any)
+T = TypeVar("T")
+
@rewrap_exceptions(
{
diff --git a/agents-api/agents_api/models/execution/list_execution_transitions.py b/agents-api/agents_api/models/execution/list_execution_transitions.py
index 9b103acf0..45aca935a 100644
--- a/agents-api/agents_api/models/execution/list_execution_transitions.py
+++ b/agents-api/agents_api/models/execution/list_execution_transitions.py
@@ -1,4 +1,4 @@
-from typing import Literal
+from typing import Any, Literal, TypeVar
from uuid import UUID
from beartype import beartype
@@ -9,6 +9,9 @@
from ...autogen.openapi_model import Transition
from ..utils import cozo_query, partialclass, rewrap_exceptions, wrap_in_class
+ModelT = TypeVar("ModelT", bound=Any)
+T = TypeVar("T")
+
@rewrap_exceptions(
{
diff --git a/agents-api/agents_api/models/execution/list_executions.py b/agents-api/agents_api/models/execution/list_executions.py
index e497ec4fe..09194cdbd 100644
--- a/agents-api/agents_api/models/execution/list_executions.py
+++ b/agents-api/agents_api/models/execution/list_executions.py
@@ -1,4 +1,4 @@
-from typing import Literal
+from typing import Any, Literal, TypeVar
from uuid import UUID
from beartype import beartype
@@ -16,6 +16,9 @@
wrap_in_class,
)
+ModelT = TypeVar("ModelT", bound=Any)
+T = TypeVar("T")
+
@rewrap_exceptions(
{
diff --git a/agents-api/agents_api/models/execution/prepare_execution_input.py b/agents-api/agents_api/models/execution/prepare_execution_input.py
index 8f146da24..c858bc6a0 100644
--- a/agents-api/agents_api/models/execution/prepare_execution_input.py
+++ b/agents-api/agents_api/models/execution/prepare_execution_input.py
@@ -1,3 +1,4 @@
+from typing import Any, TypeVar
from uuid import UUID
from beartype import beartype
@@ -20,6 +21,9 @@
)
from .get_execution import get_execution
+ModelT = TypeVar("ModelT", bound=Any)
+T = TypeVar("T")
+
@rewrap_exceptions(
{
diff --git a/agents-api/agents_api/models/execution/update_execution.py b/agents-api/agents_api/models/execution/update_execution.py
index f3c1e528e..90a8cb1cc 100644
--- a/agents-api/agents_api/models/execution/update_execution.py
+++ b/agents-api/agents_api/models/execution/update_execution.py
@@ -1,3 +1,4 @@
+from typing import Any, TypeVar
from uuid import UUID
from beartype import beartype
@@ -22,6 +23,9 @@
wrap_in_class,
)
+ModelT = TypeVar("ModelT", bound=Any)
+T = TypeVar("T")
+
@rewrap_exceptions(
{
diff --git a/agents-api/agents_api/models/session/create_or_update_session.py b/agents-api/agents_api/models/session/create_or_update_session.py
index 7b2e39e74..60c0b7f71 100644
--- a/agents-api/agents_api/models/session/create_or_update_session.py
+++ b/agents-api/agents_api/models/session/create_or_update_session.py
@@ -1,3 +1,4 @@
+from typing import Any, TypeVar
from uuid import UUID
from beartype import beartype
@@ -19,6 +20,9 @@
wrap_in_class,
)
+ModelT = TypeVar("ModelT", bound=Any)
+T = TypeVar("T")
+
@rewrap_exceptions(
{
diff --git a/agents-api/agents_api/models/session/create_session.py b/agents-api/agents_api/models/session/create_session.py
index 7bb0576c9..a83837ffd 100644
--- a/agents-api/agents_api/models/session/create_session.py
+++ b/agents-api/agents_api/models/session/create_session.py
@@ -3,6 +3,7 @@
It constructs and executes a datalog query to insert session data.
"""
+from typing import Any, TypeVar
from uuid import UUID, uuid4
from beartype import beartype
@@ -20,6 +21,9 @@
wrap_in_class,
)
+ModelT = TypeVar("ModelT", bound=Any)
+T = TypeVar("T")
+
@rewrap_exceptions(
{
diff --git a/agents-api/agents_api/models/session/delete_session.py b/agents-api/agents_api/models/session/delete_session.py
index 71f153fb1..af9e331c7 100644
--- a/agents-api/agents_api/models/session/delete_session.py
+++ b/agents-api/agents_api/models/session/delete_session.py
@@ -1,5 +1,6 @@
"""This module contains the implementation for deleting sessions from the 'cozodb' database using datalog queries."""
+from typing import Any, TypeVar
from uuid import UUID
from beartype import beartype
@@ -18,6 +19,9 @@
wrap_in_class,
)
+ModelT = TypeVar("ModelT", bound=Any)
+T = TypeVar("T")
+
@rewrap_exceptions(
{
diff --git a/agents-api/agents_api/models/session/get_session.py b/agents-api/agents_api/models/session/get_session.py
index 8e7d6adb5..0a365df2f 100644
--- a/agents-api/agents_api/models/session/get_session.py
+++ b/agents-api/agents_api/models/session/get_session.py
@@ -1,3 +1,4 @@
+from typing import Any, TypeVar
from uuid import UUID
from beartype import beartype
@@ -15,6 +16,9 @@
wrap_in_class,
)
+ModelT = TypeVar("ModelT", bound=Any)
+T = TypeVar("T")
+
@rewrap_exceptions(
{
diff --git a/agents-api/agents_api/models/session/list_sessions.py b/agents-api/agents_api/models/session/list_sessions.py
index dba36ab98..fa1097e5e 100644
--- a/agents-api/agents_api/models/session/list_sessions.py
+++ b/agents-api/agents_api/models/session/list_sessions.py
@@ -1,6 +1,6 @@
"""This module contains functions for querying session data from the 'cozodb' database."""
-from typing import Any, Literal
+from typing import Any, Literal, TypeVar
from uuid import UUID
from beartype import beartype
@@ -18,6 +18,9 @@
wrap_in_class,
)
+ModelT = TypeVar("ModelT", bound=Any)
+T = TypeVar("T")
+
@rewrap_exceptions(
{
diff --git a/agents-api/agents_api/models/session/patch_session.py b/agents-api/agents_api/models/session/patch_session.py
index feafcc679..e6e0e5750 100644
--- a/agents-api/agents_api/models/session/patch_session.py
+++ b/agents-api/agents_api/models/session/patch_session.py
@@ -1,5 +1,6 @@
"""This module contains functions for patching session data in the 'cozodb' database using datalog queries."""
+from typing import Any, List, TypeVar
from uuid import UUID
from beartype import beartype
@@ -18,7 +19,10 @@
wrap_in_class,
)
-_fields = [
+ModelT = TypeVar("ModelT", bound=Any)
+T = TypeVar("T")
+
+_fields: List[str] = [
"situation",
"summary",
"created_at",
diff --git a/agents-api/agents_api/models/session/prepare_session_data.py b/agents-api/agents_api/models/session/prepare_session_data.py
index 754b0538c..9a936b183 100644
--- a/agents-api/agents_api/models/session/prepare_session_data.py
+++ b/agents-api/agents_api/models/session/prepare_session_data.py
@@ -1,3 +1,4 @@
+from typing import Any, TypeVar
from uuid import UUID
from beartype import beartype
@@ -15,6 +16,9 @@
wrap_in_class,
)
+ModelT = TypeVar("ModelT", bound=Any)
+T = TypeVar("T")
+
@rewrap_exceptions(
{
diff --git a/agents-api/agents_api/models/session/update_session.py b/agents-api/agents_api/models/session/update_session.py
index 14b989de1..99688bd98 100644
--- a/agents-api/agents_api/models/session/update_session.py
+++ b/agents-api/agents_api/models/session/update_session.py
@@ -1,3 +1,4 @@
+from typing import Any, List, TypeVar
from uuid import UUID
from beartype import beartype
@@ -16,7 +17,10 @@
wrap_in_class,
)
-_fields = [
+ModelT = TypeVar("ModelT", bound=Any)
+T = TypeVar("T")
+
+_fields: List[str] = [
"situation",
"summary",
"metadata",
diff --git a/agents-api/agents_api/models/task/create_or_update_task.py b/agents-api/agents_api/models/task/create_or_update_task.py
index a3014a64f..af7e258d9 100644
--- a/agents-api/agents_api/models/task/create_or_update_task.py
+++ b/agents-api/agents_api/models/task/create_or_update_task.py
@@ -3,6 +3,7 @@
It constructs and executes a datalog query to insert Task data.
"""
+from typing import Any, TypeVar
from uuid import UUID
from beartype import beartype
@@ -26,6 +27,9 @@
wrap_in_class,
)
+ModelT = TypeVar("ModelT", bound=Any)
+T = TypeVar("T")
+
@rewrap_exceptions(
{
diff --git a/agents-api/agents_api/models/task/create_task.py b/agents-api/agents_api/models/task/create_task.py
index 61f0bca72..a44146c34 100644
--- a/agents-api/agents_api/models/task/create_task.py
+++ b/agents-api/agents_api/models/task/create_task.py
@@ -3,6 +3,7 @@
It constructs and executes a datalog query to insert Task data.
"""
+from typing import Any, TypeVar
from uuid import UUID, uuid4
from beartype import beartype
@@ -24,6 +25,9 @@
wrap_in_class,
)
+ModelT = TypeVar("ModelT", bound=Any)
+T = TypeVar("T")
+
@rewrap_exceptions(
{
diff --git a/agents-api/agents_api/models/task/delete_task.py b/agents-api/agents_api/models/task/delete_task.py
index 1eb780057..60d6f2756 100644
--- a/agents-api/agents_api/models/task/delete_task.py
+++ b/agents-api/agents_api/models/task/delete_task.py
@@ -1,3 +1,4 @@
+from typing import Any, TypeVar
from uuid import UUID
from beartype import beartype
@@ -16,6 +17,9 @@
wrap_in_class,
)
+ModelT = TypeVar("ModelT", bound=Any)
+T = TypeVar("T")
+
@rewrap_exceptions(
{
diff --git a/agents-api/agents_api/models/task/get_task.py b/agents-api/agents_api/models/task/get_task.py
index 975da28cd..076936b6c 100644
--- a/agents-api/agents_api/models/task/get_task.py
+++ b/agents-api/agents_api/models/task/get_task.py
@@ -1,3 +1,4 @@
+from typing import Any, TypeVar
from uuid import UUID
from beartype import beartype
@@ -15,6 +16,9 @@
)
from .create_task import spec_to_task
+ModelT = TypeVar("ModelT", bound=Any)
+T = TypeVar("T")
+
@rewrap_exceptions(
{
diff --git a/agents-api/agents_api/models/task/list_tasks.py b/agents-api/agents_api/models/task/list_tasks.py
index 1c6e16efd..573c1404e 100644
--- a/agents-api/agents_api/models/task/list_tasks.py
+++ b/agents-api/agents_api/models/task/list_tasks.py
@@ -1,4 +1,4 @@
-from typing import Literal
+from typing import Any, Literal, TypeVar
from uuid import UUID
from beartype import beartype
@@ -16,6 +16,9 @@
)
from .create_task import spec_to_task
+ModelT = TypeVar("ModelT", bound=Any)
+T = TypeVar("T")
+
@rewrap_exceptions(
{
diff --git a/agents-api/agents_api/models/task/patch_task.py b/agents-api/agents_api/models/task/patch_task.py
index dc32f83d2..1837064c7 100644
--- a/agents-api/agents_api/models/task/patch_task.py
+++ b/agents-api/agents_api/models/task/patch_task.py
@@ -3,6 +3,7 @@
It constructs and executes a datalog query to insert Task data.
"""
+from typing import Any, TypeVar
from uuid import UUID
from beartype import beartype
@@ -22,6 +23,9 @@
wrap_in_class,
)
+ModelT = TypeVar("ModelT", bound=Any)
+T = TypeVar("T")
+
@rewrap_exceptions(
{
diff --git a/agents-api/agents_api/models/task/update_task.py b/agents-api/agents_api/models/task/update_task.py
index 151d4fb4d..9cfb04357 100644
--- a/agents-api/agents_api/models/task/update_task.py
+++ b/agents-api/agents_api/models/task/update_task.py
@@ -3,6 +3,7 @@
It constructs and executes a datalog query to insert Task data.
"""
+from typing import Any, TypeVar
from uuid import UUID
from beartype import beartype
@@ -22,6 +23,9 @@
wrap_in_class,
)
+ModelT = TypeVar("ModelT", bound=Any)
+T = TypeVar("T")
+
@rewrap_exceptions(
{
diff --git a/agents-api/agents_api/models/tools/create_tools.py b/agents-api/agents_api/models/tools/create_tools.py
index 0c034d4d2..dd8397797 100644
--- a/agents-api/agents_api/models/tools/create_tools.py
+++ b/agents-api/agents_api/models/tools/create_tools.py
@@ -1,5 +1,6 @@
"""This module contains functions for creating tools in the CozoDB database."""
+from typing import Any, TypeVar
from uuid import UUID, uuid4
from beartype import beartype
@@ -17,6 +18,9 @@
wrap_in_class,
)
+ModelT = TypeVar("ModelT", bound=Any)
+T = TypeVar("T")
+
@rewrap_exceptions(
{
diff --git a/agents-api/agents_api/models/tools/delete_tool.py b/agents-api/agents_api/models/tools/delete_tool.py
index e6d00498a..c79cdfd29 100644
--- a/agents-api/agents_api/models/tools/delete_tool.py
+++ b/agents-api/agents_api/models/tools/delete_tool.py
@@ -1,3 +1,4 @@
+from typing import Any, TypeVar
from uuid import UUID
from beartype import beartype
@@ -16,6 +17,9 @@
wrap_in_class,
)
+ModelT = TypeVar("ModelT", bound=Any)
+T = TypeVar("T")
+
@rewrap_exceptions(
{
diff --git a/agents-api/agents_api/models/tools/get_tool.py b/agents-api/agents_api/models/tools/get_tool.py
index f3e6a52c3..5ea009064 100644
--- a/agents-api/agents_api/models/tools/get_tool.py
+++ b/agents-api/agents_api/models/tools/get_tool.py
@@ -1,3 +1,4 @@
+from typing import Any, TypeVar
from uuid import UUID
from beartype import beartype
@@ -15,6 +16,9 @@
wrap_in_class,
)
+ModelT = TypeVar("ModelT", bound=Any)
+T = TypeVar("T")
+
@rewrap_exceptions(
{
diff --git a/agents-api/agents_api/models/tools/list_tools.py b/agents-api/agents_api/models/tools/list_tools.py
index e1636fdd4..4b44fc1e0 100644
--- a/agents-api/agents_api/models/tools/list_tools.py
+++ b/agents-api/agents_api/models/tools/list_tools.py
@@ -1,4 +1,4 @@
-from typing import Literal
+from typing import Any, Literal, TypeVar
from uuid import UUID
from beartype import beartype
@@ -16,6 +16,9 @@
wrap_in_class,
)
+ModelT = TypeVar("ModelT", bound=Any)
+T = TypeVar("T")
+
@rewrap_exceptions(
{
diff --git a/agents-api/agents_api/models/tools/patch_tool.py b/agents-api/agents_api/models/tools/patch_tool.py
index 0b2777030..5bbfe1c91 100644
--- a/agents-api/agents_api/models/tools/patch_tool.py
+++ b/agents-api/agents_api/models/tools/patch_tool.py
@@ -1,3 +1,4 @@
+from typing import Any, TypeVar
from uuid import UUID
from beartype import beartype
@@ -16,6 +17,9 @@
wrap_in_class,
)
+ModelT = TypeVar("ModelT", bound=Any)
+T = TypeVar("T")
+
@rewrap_exceptions(
{
diff --git a/agents-api/agents_api/models/tools/update_tool.py b/agents-api/agents_api/models/tools/update_tool.py
index 3e91b7562..d1676e984 100644
--- a/agents-api/agents_api/models/tools/update_tool.py
+++ b/agents-api/agents_api/models/tools/update_tool.py
@@ -1,3 +1,4 @@
+from typing import Any, TypeVar
from uuid import UUID
from beartype import beartype
@@ -19,6 +20,9 @@
wrap_in_class,
)
+ModelT = TypeVar("ModelT", bound=Any)
+T = TypeVar("T")
+
@rewrap_exceptions(
{
diff --git a/agents-api/agents_api/models/user/create_or_update_user.py b/agents-api/agents_api/models/user/create_or_update_user.py
index 5784db880..9e9045e74 100644
--- a/agents-api/agents_api/models/user/create_or_update_user.py
+++ b/agents-api/agents_api/models/user/create_or_update_user.py
@@ -3,6 +3,7 @@
It includes functions to construct and execute datalog queries for inserting new user records.
"""
+from typing import Any, TypeVar
from uuid import UUID
from beartype import beartype
@@ -19,6 +20,9 @@
wrap_in_class,
)
+ModelT = TypeVar("ModelT", bound=Any)
+T = TypeVar("T")
+
@rewrap_exceptions(
{
diff --git a/agents-api/agents_api/models/user/create_user.py b/agents-api/agents_api/models/user/create_user.py
index 4115b3326..9dd036c57 100644
--- a/agents-api/agents_api/models/user/create_user.py
+++ b/agents-api/agents_api/models/user/create_user.py
@@ -3,6 +3,7 @@
It defines a query for inserting user data into the 'users' relation.
"""
+from typing import Any, TypeVar
from uuid import UUID, uuid4
from beartype import beartype
@@ -19,6 +20,9 @@
wrap_in_class,
)
+ModelT = TypeVar("ModelT", bound=Any)
+T = TypeVar("T")
+
@rewrap_exceptions(
{
diff --git a/agents-api/agents_api/models/user/delete_user.py b/agents-api/agents_api/models/user/delete_user.py
index b9ebc2db7..b5fcb8424 100644
--- a/agents-api/agents_api/models/user/delete_user.py
+++ b/agents-api/agents_api/models/user/delete_user.py
@@ -2,6 +2,7 @@
This module contains the implementation of the delete_user_query function, which is responsible for deleting an user and its related default settings from the CozoDB database.
"""
+from typing import Any, TypeVar
from uuid import UUID
from beartype import beartype
@@ -20,6 +21,9 @@
wrap_in_class,
)
+ModelT = TypeVar("ModelT", bound=Any)
+T = TypeVar("T")
+
@rewrap_exceptions(
{
diff --git a/agents-api/agents_api/models/user/get_user.py b/agents-api/agents_api/models/user/get_user.py
index d4556a365..181bf05f0 100644
--- a/agents-api/agents_api/models/user/get_user.py
+++ b/agents-api/agents_api/models/user/get_user.py
@@ -1,3 +1,4 @@
+from typing import Any, TypeVar
from uuid import UUID
from beartype import beartype
@@ -15,6 +16,9 @@
wrap_in_class,
)
+ModelT = TypeVar("ModelT", bound=Any)
+T = TypeVar("T")
+
@rewrap_exceptions(
{
diff --git a/agents-api/agents_api/models/user/list_users.py b/agents-api/agents_api/models/user/list_users.py
index 5fdc62ef0..57dc9b8c8 100644
--- a/agents-api/agents_api/models/user/list_users.py
+++ b/agents-api/agents_api/models/user/list_users.py
@@ -1,4 +1,4 @@
-from typing import Any, Literal
+from typing import Any, Literal, TypeVar
from uuid import UUID
from beartype import beartype
@@ -16,6 +16,9 @@
wrap_in_class,
)
+ModelT = TypeVar("ModelT", bound=Any)
+T = TypeVar("T")
+
@rewrap_exceptions(
{
diff --git a/agents-api/agents_api/models/user/patch_user.py b/agents-api/agents_api/models/user/patch_user.py
index d8dfeb2ad..faf38298c 100644
--- a/agents-api/agents_api/models/user/patch_user.py
+++ b/agents-api/agents_api/models/user/patch_user.py
@@ -1,5 +1,6 @@
"""Module for generating datalog queries to update user information in the 'cozodb' database."""
+from typing import Any, TypeVar
from uuid import UUID
from beartype import beartype
@@ -19,6 +20,9 @@
wrap_in_class,
)
+ModelT = TypeVar("ModelT", bound=Any)
+T = TypeVar("T")
+
@rewrap_exceptions(
{
diff --git a/agents-api/agents_api/models/user/update_user.py b/agents-api/agents_api/models/user/update_user.py
index 8411ecc3e..9a13d9369 100644
--- a/agents-api/agents_api/models/user/update_user.py
+++ b/agents-api/agents_api/models/user/update_user.py
@@ -1,3 +1,4 @@
+from typing import Any, TypeVar
from uuid import UUID
from beartype import beartype
@@ -16,6 +17,9 @@
wrap_in_class,
)
+ModelT = TypeVar("ModelT", bound=Any)
+T = TypeVar("T")
+
@rewrap_exceptions(
{
diff --git a/agents-api/agents_api/models/utils.py b/agents-api/agents_api/models/utils.py
index ee2ba3fdd..bc36af8d2 100644
--- a/agents-api/agents_api/models/utils.py
+++ b/agents-api/agents_api/models/utils.py
@@ -195,10 +195,10 @@ def wrapper(*args: P.args, client=None, **kwargs: P.kwargs) -> pd.DataFrame:
)
# Run the query
- from ..clients.cozo import get_cozo_client
+ from ..clients import cozo
try:
- client = client or get_cozo_client()
+ client = client or cozo.get_cozo_client()
result = client.run(query, variables)
except Exception as e:
diff --git a/agents-api/agents_api/rec_sum/data.py b/agents-api/agents_api/rec_sum/data.py
index e7857a037..23474c995 100644
--- a/agents-api/agents_api/rec_sum/data.py
+++ b/agents-api/agents_api/rec_sum/data.py
@@ -1,24 +1,25 @@
import json
from pathlib import Path
+from typing import Any
-module_directory = Path(__file__).parent
+module_directory: Path = Path(__file__).parent
with open(f"{module_directory}/entities_example_chat.json", "r") as _f:
- entities_example_chat = json.load(_f)
+ entities_example_chat: Any = json.load(_f)
with open(f"{module_directory}/trim_example_chat.json", "r") as _f:
- trim_example_chat = json.load(_f)
+ trim_example_chat: Any = json.load(_f)
with open(f"{module_directory}/trim_example_result.json", "r") as _f:
- trim_example_result = json.load(_f)
+ trim_example_result: Any = json.load(_f)
with open(f"{module_directory}/summarize_example_chat.json", "r") as _f:
- summarize_example_chat = json.load(_f)
+ summarize_example_chat: Any = json.load(_f)
with open(f"{module_directory}/summarize_example_result.json", "r") as _f:
- summarize_example_result = json.load(_f)
+ summarize_example_result: Any = json.load(_f)
diff --git a/agents-api/agents_api/rec_sum/entities.py b/agents-api/agents_api/rec_sum/entities.py
index 9992063ff..01b29951b 100644
--- a/agents-api/agents_api/rec_sum/entities.py
+++ b/agents-api/agents_api/rec_sum/entities.py
@@ -12,21 +12,21 @@
## Entities ##
##############
-entities_example_plan = """\
+entities_example_plan: str = """\
Thinking step by step:
- To add context for future entries, let's outline the main entities in the session above.
- In this session, as mentioned in the first message metadata, the user's name is Camille and the assistant's name is JaneBot.
- They talk about Elon Musk and the banana tattoo on Camille's arm briefly."""
-entities_example_result = """\
+entities_example_result: str = """\
1. Camille (The user): Humorous, creative, and enjoys playful banter.
2. JaneBot (The assistant): Engages in lighthearted conversation and tries to guess user's thoughts.
3. Elon Musk: Camille and JaneBot discuss the polarizing tech and space industry figure.
4. Banana Tattoo: Camille has a tattoo of a banana on their arm."""
-entities_instructions = """\
+entities_instructions: str = """\
Your goal is to identify the main entities in the session. Entities should include:
- Characters in the conversation: Assistant, User1, User2
- People references or spoken about
diff --git a/agents-api/agents_api/rec_sum/summarize.py b/agents-api/agents_api/rec_sum/summarize.py
index f98f35094..46a6662a3 100644
--- a/agents-api/agents_api/rec_sum/summarize.py
+++ b/agents-api/agents_api/rec_sum/summarize.py
@@ -1,4 +1,5 @@
import json
+from typing import List
from tenacity import retry, stop_after_attempt
@@ -10,7 +11,7 @@
## summarize ##
##########
-summarize_example_plan = """\
+summarize_example_plan: str = """\
Planning step by step:
- We can replace entries 1,2,3,4 with a summary of those messages.
- We can replace entries 5,6,7,8 similarly.
@@ -23,7 +24,7 @@
- We can safely summarize message 34's essay into just the salient points only."""
-summarize_instructions = """\
+summarize_instructions: str = """\
Your goal is to compactify the history by coalescing redundant information in messages into their summary in order to reduce its size and save costs.
Instructions:
@@ -34,7 +35,9 @@
- VERY IMPORTANT: Add the indices of messages that are being summarized so that those messages can then be removed from the session otherwise, there'll be no way to identify which messages to remove. See example for more details."""
-def make_summarize_prompt(session, user="a user", assistant="gpt-4-turbo", **_):
+def make_summarize_prompt(
+ session, user="a user", assistant="gpt-4-turbo", **_
+) -> List[str]:
return [
f"You are given a session history of a chat between {user or 'a user'} and {assistant or 'gpt-4-turbo'}. The session is formatted in the ChatML JSON format (from OpenAI).\n\n{summarize_instructions}\n\n\n{json.dumps(add_indices(summarize_example_chat), indent=2)}\n\n\n\n{summarize_example_plan}\n\n\n\n{json.dumps(summarize_example_result, indent=2)}\n",
f"Begin! Write the summarized messages as a json list just like the example above. First write your plan inside and then your answer between . Don't forget to add the indices of the messages being summarized alongside each summary.\n\n\n{json.dumps(add_indices(session), indent=2)}\n\n",
diff --git a/agents-api/agents_api/rec_sum/trim.py b/agents-api/agents_api/rec_sum/trim.py
index 7e902859c..ee4025ea0 100644
--- a/agents-api/agents_api/rec_sum/trim.py
+++ b/agents-api/agents_api/rec_sum/trim.py
@@ -1,4 +1,5 @@
import json
+from typing import List
from tenacity import retry, stop_after_attempt
@@ -10,7 +11,7 @@
## Trim ##
##########
-trim_example_plan = """\
+trim_example_plan: str = """\
Thinking step by step:
- To trim the context, let's examine the messages in the session above.
- Messages 1, 2, and 3 are succinct and do not need trimming.
@@ -18,7 +19,7 @@
- Message 7 is short enough and doesn't need any edits."""
-trim_instructions = """\
+trim_instructions: str = """\
Your goal is to identify messages in the session that are needlessly verbose and then trim them in length without losing any meaning or changing the tone of the message.
Instructions:
@@ -32,7 +33,7 @@
# It is important to make keep the tone, setting and flow of the conversation consistent while trimming the messages.
-def make_trim_prompt(session, user="a user", assistant="gpt-4-turbo", **_):
+def make_trim_prompt(session, user="a user", assistant="gpt-4-turbo", **_) -> List[str]:
return [
f"You are given a session history of a chat between {user or 'a user'} and {assistant or 'gpt-4-turbo'}. The session is formatted in the ChatML JSON format (from OpenAI).\n\n{trim_instructions}\n\n\n{json.dumps(add_indices(trim_example_chat), indent=2)}\n\n\n\n{trim_example_plan}\n\n\n\n{json.dumps(trim_example_result, indent=2)}\n",
f"Begin! Write the trimmed messages as a json list. First write your plan inside and then your answer between .\n\n\n{json.dumps(add_indices(session), indent=2)}\n\n",
diff --git a/agents-api/agents_api/rec_sum/utils.py b/agents-api/agents_api/rec_sum/utils.py
index 596174e08..c674a4d44 100644
--- a/agents-api/agents_api/rec_sum/utils.py
+++ b/agents-api/agents_api/rec_sum/utils.py
@@ -3,9 +3,14 @@
###########
+from typing import Any, Dict, List, TypeVar
+
+_T2 = TypeVar("_T2")
+
+
class chatml:
@staticmethod
- def make(content, role="system", name=None, **_):
+ def make(content, role="system", name: _T2 = None, **_) -> Dict[str, _T2]:
return {
key: value
for key, value in dict(role=role, name=name, content=content).items()
@@ -13,39 +18,39 @@ def make(content, role="system", name=None, **_):
}
@staticmethod
- def user(content, name=None):
+ def user(content, name=None) -> Any:
return chatml.make(role="user", content=content, name=name)
@staticmethod
- def assistant(content, name=None):
+ def assistant(content, name=None) -> Any:
return chatml.make(role="assistant", content=content, name=name)
@staticmethod
- def system(content, name=None):
+ def system(content, name=None) -> Any:
return chatml.make(content, name=name)
@staticmethod
- def thought(content, name=None):
+ def thought(content, name=None) -> Any:
return chatml.make(content, name="thought")
@staticmethod
- def information(content):
+ def information(content) -> Any:
return chatml.system(content, name="information")
@staticmethod
- def summary(content):
+ def summary(content) -> Any:
return chatml.system(content, name="summary")
@staticmethod
- def entities(content):
+ def entities(content) -> Any:
return chatml.system(content, name="entity")
-def add_indices(list_of_dicts, idx_name="index"):
+def add_indices(list_of_dicts, idx_name="index") -> List[dict]:
return [{idx_name: i, **msg} for i, msg in enumerate(list_of_dicts)]
-def get_names_from_session(session):
+def get_names_from_session(session) -> Dict[str, Any]:
return {
role: next(
(msg.get("name", None) for msg in session if msg["role"] == role), None
diff --git a/agents-api/agents_api/routers/agents/router.py b/agents-api/agents_api/routers/agents/router.py
index af9233c56..6582baa6f 100644
--- a/agents-api/agents_api/routers/agents/router.py
+++ b/agents-api/agents_api/routers/agents/router.py
@@ -1,3 +1,5 @@
+from typing import Any
+
from fastapi import APIRouter
-router = APIRouter()
+router: Any = APIRouter()
diff --git a/agents-api/agents_api/routers/docs/router.py b/agents-api/agents_api/routers/docs/router.py
index af9233c56..6582baa6f 100644
--- a/agents-api/agents_api/routers/docs/router.py
+++ b/agents-api/agents_api/routers/docs/router.py
@@ -1,3 +1,5 @@
+from typing import Any
+
from fastapi import APIRouter
-router = APIRouter()
+router: Any = APIRouter()
diff --git a/agents-api/agents_api/routers/docs/search_docs.py b/agents-api/agents_api/routers/docs/search_docs.py
index ad19a3178..f2864164e 100644
--- a/agents-api/agents_api/routers/docs/search_docs.py
+++ b/agents-api/agents_api/routers/docs/search_docs.py
@@ -1,5 +1,5 @@
import time
-from typing import Annotated
+from typing import Annotated, Any, Dict, List, Optional, Tuple, Union
from fastapi import Depends
from pydantic import UUID4
@@ -17,7 +17,11 @@
from .router import router
-def get_search_fn_and_params(search_params):
+def get_search_fn_and_params(
+ search_params,
+) -> Tuple[
+ Any, Optional[Dict[str, Union[float, int, str, Dict[str, float], List[float]]]]
+]:
search_fn, params = None, None
match search_params:
diff --git a/agents-api/agents_api/routers/jobs/routers.py b/agents-api/agents_api/routers/jobs/routers.py
index 1ec60c2b4..5f7d3ef67 100644
--- a/agents-api/agents_api/routers/jobs/routers.py
+++ b/agents-api/agents_api/routers/jobs/routers.py
@@ -1,4 +1,4 @@
-from typing import Literal
+from typing import Any, Literal
from fastapi import APIRouter
from pydantic import UUID4
@@ -7,7 +7,7 @@
from agents_api.autogen.openapi_model import JobStatus
from agents_api.clients.temporal import get_client
-router = APIRouter()
+router: Any = APIRouter()
def map_job_status(
diff --git a/agents-api/agents_api/routers/sessions/exceptions.py b/agents-api/agents_api/routers/sessions/exceptions.py
index add4b79cb..b7a5bb971 100644
--- a/agents-api/agents_api/routers/sessions/exceptions.py
+++ b/agents-api/agents_api/routers/sessions/exceptions.py
@@ -3,7 +3,7 @@ class BaseSessionException(Exception):
class InputTooBigError(BaseSessionException):
- def __init__(self, actual_tokens, required_tokens):
+ def __init__(self, actual_tokens, required_tokens) -> None:
super().__init__(
f"Input is too big, {actual_tokens} tokens provided, but only {required_tokens} tokens are allowed."
)
diff --git a/agents-api/agents_api/routers/sessions/router.py b/agents-api/agents_api/routers/sessions/router.py
index af9233c56..6582baa6f 100644
--- a/agents-api/agents_api/routers/sessions/router.py
+++ b/agents-api/agents_api/routers/sessions/router.py
@@ -1,3 +1,5 @@
+from typing import Any
+
from fastapi import APIRouter
-router = APIRouter()
+router: Any = APIRouter()
diff --git a/agents-api/agents_api/routers/tasks/create_task_execution.py b/agents-api/agents_api/routers/tasks/create_task_execution.py
index 988af074f..6202baa93 100644
--- a/agents-api/agents_api/routers/tasks/create_task_execution.py
+++ b/agents-api/agents_api/routers/tasks/create_task_execution.py
@@ -30,7 +30,7 @@
from ...models.task.get_task import get_task as get_task_query
from .router import router
-logger = logging.getLogger(__name__)
+logger: logging.Logger = logging.getLogger(__name__)
logger.setLevel(logging.DEBUG)
diff --git a/agents-api/agents_api/routers/tasks/router.py b/agents-api/agents_api/routers/tasks/router.py
index af9233c56..6582baa6f 100644
--- a/agents-api/agents_api/routers/tasks/router.py
+++ b/agents-api/agents_api/routers/tasks/router.py
@@ -1,3 +1,5 @@
+from typing import Any
+
from fastapi import APIRouter
-router = APIRouter()
+router: Any = APIRouter()
diff --git a/agents-api/agents_api/routers/users/router.py b/agents-api/agents_api/routers/users/router.py
index af9233c56..6582baa6f 100644
--- a/agents-api/agents_api/routers/users/router.py
+++ b/agents-api/agents_api/routers/users/router.py
@@ -1,3 +1,5 @@
+from typing import Any
+
from fastapi import APIRouter
-router = APIRouter()
+router: Any = APIRouter()
diff --git a/agents-api/agents_api/web.py b/agents-api/agents_api/web.py
index 2b85ece5f..2a24dcebb 100644
--- a/agents-api/agents_api/web.py
+++ b/agents-api/agents_api/web.py
@@ -3,6 +3,7 @@
"""
import logging
+from typing import Any, Callable
import fire
import sentry_sdk
@@ -37,10 +38,10 @@
)
-logger = logging.getLogger(__name__)
+logger: logging.Logger = logging.getLogger(__name__)
-def make_exception_handler(status: int):
+def make_exception_handler(status: int) -> Callable[[Any, Any], Any]:
"""
Creates a custom exception handler for the application.
@@ -60,7 +61,7 @@ async def _handler(request: Request, exc):
return _handler
-def register_exceptions(app: FastAPI):
+def register_exceptions(app: FastAPI) -> None:
"""
Registers custom exception handlers for the FastAPI application.
@@ -77,7 +78,7 @@ def register_exceptions(app: FastAPI):
)
-app = FastAPI(dependencies=[Depends(get_api_key)])
+app: Any = FastAPI(dependencies=[Depends(get_api_key)])
app.add_middleware(
CORSMiddleware,
@@ -147,7 +148,7 @@ def main(
timeout_keep_alive=30,
workers=None,
log_level="info",
-):
+) -> None:
uvicorn.run(
app,
host=host,
diff --git a/agents-api/agents_api/worker/codec.py b/agents-api/agents_api/worker/codec.py
index abf3ddfe5..cdd7e3448 100644
--- a/agents-api/agents_api/worker/codec.py
+++ b/agents-api/agents_api/worker/codec.py
@@ -19,7 +19,7 @@
from agents_api.common.utils.json import dumps as json_dumps
# Map of model name to class so that we can look up the class when deserializing
-model_class_map = {
+model_class_map: dict = {
subclass.__module__ + "." + subclass.__name__: subclass
for subclass in {
# All the models we want to support
@@ -83,7 +83,7 @@ def __init__(self) -> None:
# Use the default data converter, but change the payload converter.
-pydantic_data_converter = dataclasses.replace(
+pydantic_data_converter: Any = dataclasses.replace(
temporalio.converter.default(),
payload_converter_class=PydanticPayloadConverter,
)
diff --git a/agents-api/agents_api/worker/worker.py b/agents-api/agents_api/worker/worker.py
index d3821c902..a5e5016d3 100644
--- a/agents-api/agents_api/worker/worker.py
+++ b/agents-api/agents_api/worker/worker.py
@@ -1,10 +1,11 @@
from datetime import timedelta
+from typing import Any
from temporalio.client import Client
from temporalio.worker import Worker
-def create_worker(client: Client):
+def create_worker(client: Client) -> Any:
"""
Initializes the Temporal client and worker with TLS configuration (if provided),
then create a worker to listen for tasks on the configured task queue.
diff --git a/agents-api/agents_api/workflows/task_execution.py b/agents-api/agents_api/workflows/task_execution.py
index 5e5eb8ce6..7e8498828 100644
--- a/agents-api/agents_api/workflows/task_execution.py
+++ b/agents-api/agents_api/workflows/task_execution.py
@@ -163,4 +163,4 @@ async def run(
# Otherwise, recurse to the next step
workflow.continue_as_new(
args=[execution_input, next_target, previous_inputs + [final_output]]
- )
\ No newline at end of file
+ )