Improve docstrings + minor code refactor

agents-api/agents_api/activities/truncation.py

@@ -3,7 +3,7 @@
 from beartype import beartype
 from temporalio import activity
 
-from agents_api.autogen.openapi_model import Entry
+from ..autogen.openapi_model import Entry
 
 # from agents_api.models.entry.entries_summarization import get_toplevel_entries_query
 

agents-api/agents_api/clients/worker/worker.py

@@ -1,6 +1,6 @@
 import httpx
 
-from agents_api.env import temporal_worker_url
+from ...env import temporal_worker_url
 
 from .types import (
     MemoryManagementTask,

agents-api/agents_api/common/exceptions/__init__.py

@@ -1,12 +1,14 @@
 """
-This module defines a structured hierarchy of custom exceptions for the agents API, aimed at handling specific error scenarios encountered across various operations. These exceptions are designed to provide clear, actionable error messages and appropriate HTTP status codes, enhancing the API's robustness and usability.
+This module defines a structured hierarchy of custom exceptions for the agents API, aimed at handling specific error scenarios encountered across various operations.
+These exceptions are designed to provide clear, actionable error messages and appropriate HTTP status codes, enhancing the API's robustness and usability.
 
 Exceptions are organized into categories based on the domain of operation, including:
 - Agent-related operations (agents.py): Exceptions such as `AgentNotFoundError` and `AgentToolNotFoundError` cater to errors specific to agent management.
 - Session management (sessions.py): Defines exceptions like `SessionNotFoundError` for handling errors related to session operations.
 - User interactions (users.py): Includes exceptions such as `UserNotFoundError` for addressing issues encountered during user-related operations.
 
-All custom exceptions extend from `BaseCommonException`, which encapsulates common attributes and behavior, including the error message and HTTP status code. This structured approach to exception handling facilitates precise and meaningful error feedback to API consumers, thereby improving the overall developer experience.
+All custom exceptions extend from `BaseCommonException`, which encapsulates common attributes and behavior, including the error message and HTTP status code.
+This structured approach to exception handling facilitates precise and meaningful error feedback to API consumers, thereby improving the overall developer experience.
 """
 
 

agents-api/agents_api/common/exceptions/agents.py

@@ -12,7 +12,12 @@ class BaseAgentException(BaseCommonException):
 
 
 class AgentNotFoundError(BaseAgentException):
-    """Exception raised when a requested agent cannot be found."""
+    """
+    Exception raised when a requested agent cannot be found.
+    Attributes:
+        developer_id (UUID | str): The ID of the developer attempting the operation.
+        agent_id (UUID | str): The ID of the agent that was not found.
+    """
 
     def __init__(self, developer_id: UUID | str, agent_id: UUID | str):
         # Initialize the exception with a message indicating the missing agent and developer ID.
@@ -23,7 +28,12 @@ def __init__(self, developer_id: UUID | str, agent_id: UUID | str):
 
 
 class AgentToolNotFoundError(BaseAgentException):
-    """Exception raised when a requested tool associated with an agent cannot be found."""
+    """
+    Exception raised when a requested tool associated with an agent cannot be found.
+    Attributes:
+        agent_id (UUID | str): The ID of the agent that was not found.
+        tool_id (UUID | str): The ID of the tool that was not found.
+    """
 
     def __init__(self, agent_id: UUID | str, tool_id: UUID | str):
         # Initialize the exception with a message indicating the missing tool and agent ID.
@@ -33,7 +43,12 @@ def __init__(self, agent_id: UUID | str, tool_id: UUID | str):
 
 
 class AgentDocNotFoundError(BaseAgentException):
-    """Exception raised when a requested document associated with an agent cannot be found."""
+    """
+    Exception raised when a requested document associated with an agent cannot be found.
+    Attributes:
+        agent_id (UUID | str): The ID of the agent that was not found.
+        doc_id (UUID | str): The ID of the document that was not found.
+    """
 
     def __init__(self, agent_id: UUID | str, doc_id: UUID | str):
         # Initialize the exception with a message indicating the missing document and agent ID.
@@ -43,6 +58,8 @@ def __init__(self, agent_id: UUID | str, doc_id: UUID | str):
 
 
 class AgentModelNotValid(BaseAgentException):
+    """Exception raised when requested model is not recognized."""
+
     def __init__(self, model: str, all_models: list[str]):
         super().__init__(
             f"Unknown model: {model}. Please provide a valid model name."
@@ -52,6 +69,8 @@ def __init__(self, model: str, all_models: list[str]):
 
 
 class MissingAgentModelAPIKeyError(BaseAgentException):
+    """Exception raised when API key for requested model is missing."""
+
     def __init__(self, model: str):
         super().__init__(
             f"API key missing for model: {model}. Please provide a valid API key in the configuration",

agents-api/agents_api/common/exceptions/sessions.py

@@ -8,29 +8,23 @@
 
 from . import BaseCommonException
 
-"""
-Base exception class for session-related errors.
-
-This class serves as a base for all session-related exceptions, allowing for a structured exception handling approach specific to session operations.
-"""
-
-
 class BaseSessionException(BaseCommonException):
-    pass
-
+    """
+    Base exception class for session-related errors.
 
-"""
-Exception raised when a session cannot be found.
+    This class serves as a base for all session-related exceptions, allowing for a structured exception handling approach specific to session operations.
+    """
 
-This exception is used to indicate that a specific session, identified by its session ID, does not exist or is not accessible for the given developer.
-"""
+    pass
 
 
 class SessionNotFoundError(BaseSessionException):
     """
-    Initializes a new instance of the SessionNotFoundError.
+    Exception raised when a session cannot be found.
+
+    This exception is used to indicate that a specific session, identified by its session ID, does not exist or is not accessible for the given developer.
 
-    Args:
+    Attributes:
         developer_id (UUID | str): The unique identifier of the developer attempting to access the session.
         session_id (UUID | str): The unique identifier of the session that was not found.
     """

agents-api/agents_api/common/exceptions/users.py

@@ -6,13 +6,18 @@
 
 
 class BaseUserException(BaseCommonException):
-    """Base exception class for user-related errors. This class serves as a parent for all user-related exceptions to facilitate catching errors specific to user operations."""
+    """
+    Base exception class for user-related errors.
+
+    This class serves as a parent for all user-related exceptions to facilitate catching errors specific to user operations.
+    """
 
     pass
 
 
 class UserNotFoundError(BaseUserException):
-    """Exception raised when a requested user cannot be found.
+    """
+    Exception raised when a requested user cannot be found.
     Attributes:
         developer_id (UUID | str): The ID of the developer attempting the operation.
         user_id (UUID | str): The ID of the user that was not found.
@@ -27,7 +32,8 @@ def __init__(self, developer_id: UUID | str, user_id: UUID | str):
 
 
 class UserDocNotFoundError(BaseUserException):
-    """Exception raised when a specific document related to a user cannot be found.
+    """
+    Exception raised when a specific document related to a user cannot be found.
     Attributes:
         user_id (UUID | str): The ID of the user associated with the document.
         doc_id (UUID | str): The ID of the document that was not found.

agents-api/agents_api/common/protocol/agents.py

@@ -4,17 +4,23 @@
 class AgentDefaultSettings(BaseModel):
     """Defines default settings for an agent. These settings control various aspects of the agent's behavior during operation."""
 
-    """Temperature setting influencing the randomness of the agent's responses. Higher values lead to more random responses."""
     temperature: float = 0.0
-    """Top-p sampling setting controlling the nucleus of the probability distribution to sample from."""
+    """Temperature setting influencing the randomness of the agent's responses. Higher values lead to more random responses."""
+
     top_p: float = 1.0
-    """Penalty applied to discourage repetition in the agent's responses."""
+    """Top-p sampling setting controlling the nucleus of the probability distribution to sample from."""
+
     repetition_penalty: float = 1.0
-    """Penalty for longer responses, encouraging more concise outputs."""
+    """Penalty applied to discourage repetition in the agent's responses."""
+
     length_penalty: float = 1.0
-    """Penalty applied based on the presence of certain words, influencing content generation."""
+    """Penalty for longer responses, encouraging more concise outputs."""
+
     presence_penalty: float = 0.0
-    """Penalty that decreases the likelihood of frequently used words in the agent's responses."""
+    """Penalty applied based on the presence of certain words, influencing content generation."""
+
     frequency_penalty: float = 0.0
-    """Minimum probability threshold for including a word in the agent's response."""
+    """Penalty that decreases the likelihood of frequently used words in the agent's responses."""
+
     min_p: float = 0.01
+    """Minimum probability threshold for including a word in the agent's response."""

agents-api/agents_api/common/utils/json.py

@@ -11,28 +11,35 @@ 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) -> None:
-        """Initializes the custom JSON encoder.
+        """
+        Initializes the custom JSON encoder.
         Parameters:
         *args: Variable length argument list.
         **kwargs: Arbitrary keyword arguments. The 'default_empty_value' keyword argument specifies the default value to use for None objects during serialization.
         """
+
         self._default_empty_value = kwargs.pop("default_empty_value")
         super().__init__(*args, **kwargs)
 
     def encode(self, o) -> str:
-        """Encodes the given object into a JSON formatted string.
+        """
+        Encodes the given object into a JSON formatted string.
         Parameters:
         o: The object to encode.
-        Returns: A JSON formatted string representing 'o'."""
+        Returns: A JSON formatted string representing 'o'.
+        """
+
         # Use the overridden default method for serialization before encoding
         return super().encode(self.default(o))
 
     def default(self, obj) -> Any:
-        """Provides a default serialization for objects that the standard JSON encoder cannot serialize.
+        """
+        Provides a default serialization for objects that the standard JSON encoder cannot serialize.
         Parameters:
         obj: The object to serialize.
         Returns: A serializable object or raises a TypeError if the object is not serializable.
         """
+
         if obj is None:
             return self._default_empty_value
 
@@ -46,12 +53,15 @@ def default(self, obj) -> Any:
 
 
 def dumps(obj: Any, default_empty_value="", cls=None) -> str:
-    """Serializes an object to a JSON formatted string using the custom JSON encoder.
+    """
+    Serializes an object to a JSON formatted string using the custom JSON encoder.
     Parameters:
     obj: The object to serialize.
     default_empty_value: The default value to use for None objects.
     cls: The custom encoder class to use, defaults to CustomJSONEncoder.
-    Returns: A JSON formatted string."""
+    Returns: A JSON formatted string.
+    """
+
     return json.dumps(
         obj, cls=cls or CustomJSONEncoder, default_empty_value=default_empty_value
     )

agents-api/agents_api/common/utils/messages.py

@@ -1,7 +1,7 @@
 import json
 from typing import cast
 
-from agents_api.autogen.openapi_model import (
+from ...autogen.openapi_model import (
     ChatMLImageContentPart,
     ChatMLTextContentPart,
 )

agents-api/agents_api/exceptions.py

@@ -3,17 +3,23 @@ class AgentsBaseException(Exception):
 
 
 class ModelNotSupportedError(AgentsBaseException):
+    """Exception raised when model is not supported."""
+
     def __init__(self, model_name) -> None:
         super().__init__(f"model {model_name} is not supported")
 
 
 class PromptTooBigError(AgentsBaseException):
+    """Exception raised when prompt is too big."""
+
     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):
+    """Exception raised when tokenizer is unknown."""
+
     def __init__(self) -> None:
         super().__init__("unknown tokenizer")

agents-api/agents_api/models/agent/create_agent.py

@@ -55,12 +55,12 @@ def create_agent(
     Constructs and executes a datalog query to create a new agent in the database.
 
     Parameters:
-    - agent_id (UUID | None): The unique identifier for the agent.
-    - developer_id (UUID): The unique identifier for the developer creating the agent.
-    - data (CreateAgentRequest): The data for the new agent.
+        agent_id (UUID | None): The unique identifier for the agent.
+        developer_id (UUID): The unique identifier for the developer creating the agent.
+        data (CreateAgentRequest): The data for the new agent.
 
     Returns:
-    - Agent: The newly created agent record.
+        Agent: The newly created agent record.
     """
 
     agent_id = agent_id or uuid4()

agents-api/agents_api/models/agent/create_or_update_agent.py

@@ -47,18 +47,18 @@ def create_or_update_agent(
     Constructs and executes a datalog query to create a new agent in the database.
 
     Parameters:
-    - agent_id (UUID): The unique identifier for the agent.
-    - developer_id (UUID): The unique identifier for the developer creating the agent.
-    - name (str): The name of the agent.
-    - about (str): A description of the agent.
-    - instructions (list[str], optional): A list of instructions for using the agent. Defaults to an empty list.
-    - model (str, optional): The model identifier for the agent. Defaults to "gpt-4o".
-    - metadata (dict, optional): A dictionary of metadata for the agent. Defaults to an empty dict.
-    - default_settings (dict, optional): A dictionary of default settings for the agent. Defaults to an empty dict.
-    - client (CozoClient, optional): The CozoDB client instance to use for the query. Defaults to a preconfigured client instance.
+        agent_id (UUID): The unique identifier for the agent.
+        developer_id (UUID): The unique identifier for the developer creating the agent.
+        name (str): The name of the agent.
+        about (str): A description of the agent.
+        instructions (list[str], optional): A list of instructions for using the agent. Defaults to an empty list.
+        model (str, optional): The model identifier for the agent. Defaults to "gpt-4o".
+        metadata (dict, optional): A dictionary of metadata for the agent. Defaults to an empty dict.
+        default_settings (dict, optional): A dictionary of default settings for the agent. Defaults to an empty dict.
+        client (CozoClient, optional): The CozoDB client instance to use for the query. Defaults to a preconfigured client instance.
 
     Returns:
-    Agent: The newly created agent record.
+        Agent: The newly created agent record.
     """
 
     # Extract the agent data from the payload

agents-api/agents_api/models/agent/delete_agent.py

@@ -54,12 +54,12 @@ def delete_agent(*, developer_id: UUID, agent_id: UUID) -> tuple[list[str], dict
     Constructs and returns a datalog query for deleting an agent and its default settings from the database.
 
     Parameters:
-    - developer_id (UUID): The UUID of the developer owning the agent.
-    - agent_id (UUID): The UUID of the agent to be deleted.
-    - client (CozoClient, optional): An instance of the CozoClient to execute the query.
+        developer_id (UUID): The UUID of the developer owning the agent.
+        agent_id (UUID): The UUID of the agent to be deleted.
+        client (CozoClient, optional): An instance of the CozoClient to execute the query.
 
     Returns:
-    - ResourceDeletedResponse: The response indicating the deletion of the agent.
+        ResourceDeletedResponse: The response indicating the deletion of the agent.
     """
 
     queries = [

agents-api/agents_api/models/agent/get_agent.py

@@ -46,12 +46,12 @@ def get_agent(*, developer_id: UUID, agent_id: UUID) -> tuple[list[str], dict]:
     This function constructs and executes a datalog query to retrieve information about a specific agent, including its default settings, based on the provided agent_id and developer_id.
 
     Parameters:
-    - developer_id (UUID): The unique identifier for the developer.
-    - agent_id (UUID): The unique identifier for the agent.
-    - client (CozoClient, optional): The database client used to execute the query.
+        developer_id (UUID): The unique identifier for the developer.
+        agent_id (UUID): The unique identifier for the agent.
+        client (CozoClient, optional): The database client used to execute the query.
 
     Returns:
-    - Agent
+        Agent
     """
     # Constructing a datalog query to retrieve agent details and default settings.
     # The query uses input parameters for agent_id and developer_id to filter the results.

agents-api/agents_api/models/agent/patch_agent.py

@@ -46,13 +46,13 @@ def patch_agent(
     """Patches agent data based on provided updates.
 
     Parameters:
-    agent_id (UUID): The unique identifier for the agent.
-    developer_id (UUID): The unique identifier for the developer.
-    default_settings (dict, optional): Default settings to apply to the agent.
-    **update_data: Arbitrary keyword arguments representing data to update.
+        agent_id (UUID): The unique identifier for the agent.
+        developer_id (UUID): The unique identifier for the developer.
+        default_settings (dict, optional): Default settings to apply to the agent.
+        **update_data: Arbitrary keyword arguments representing data to update.
 
     Returns:
-    ResourceUpdatedResponse: The updated agent data.
+        ResourceUpdatedResponse: The updated agent data.
     """
     update_data = data.model_dump(exclude_unset=True)