diff --git a/.github/workflows/lint-and-format.yml b/.github/workflows/lint-agents-api-pr.yml similarity index 74% rename from .github/workflows/lint-and-format.yml rename to .github/workflows/lint-agents-api-pr.yml index 65e4983c4..90fdd0a4b 100644 --- a/.github/workflows/lint-and-format.yml +++ b/.github/workflows/lint-agents-api-pr.yml @@ -33,35 +33,17 @@ jobs: restore-keys: | ${{ runner.os }}-agents-api-poetry- - - name: Cache pytype - uses: actions/cache@v4 - with: - path: agents-api/.pytype - key: ${{ runner.os }}-agents-api-pytype-${{ hashFiles('agents-api/**/*.py') }} - restore-keys: | - ${{ runner.os }}-agents-api-pytype- - - name: Install dependencies run: | cd agents-api poetry install - - name: Typecheck - run: | - cd agents-api - poetry run poe typecheck - - name: Lint and format run: | cd agents-api poetry run poe format poetry run poe lint - - name: Run tests - run: | - cd agents-api - poetry run poe test - - uses: stefanzweifel/git-auto-commit-action@v4 with: commit_message: "refactor: Lint agents-api (CI)" diff --git a/.github/workflows/test-agents-api-pr.yml b/.github/workflows/test-agents-api-pr.yml new file mode 100644 index 000000000..b0f815634 --- /dev/null +++ b/.github/workflows/test-agents-api-pr.yml @@ -0,0 +1,48 @@ +name: Test agents-api +run-name: ${{ github.actor }} is testing the code + +# TODO: Fix CI github actions +# SCRUM-26 + +on: [pull_request] + +jobs: + Test: + runs-on: ubuntu-latest + + steps: + - uses: actions/checkout@v4 + + - uses: actions/setup-python@v5 + with: + python-version: "3.12" + + - name: Install and configure Poetry + uses: snok/install-poetry@v1 + + - name: Configure Poetry to use .venv + run: | + cd agents-api + poetry config virtualenvs.in-project true + + - name: Cache Poetry virtualenv + uses: actions/cache@v4 + with: + path: agents-api/.venv + key: ${{ runner.os }}-agents-api-poetry-${{ hashFiles('agents-api/poetry.lock') }} + restore-keys: | + ${{ runner.os }}-agents-api-poetry- + + - name: Install dependencies + run: | + cd agents-api + poetry install + + - name: Run tests + run: | + cd agents-api + poetry run poe test --fail-limit 1 + +concurrency: + group: ${{ github.workflow }}-${{ github.ref }} + cancel-in-progress: ${{ github.ref != 'refs/heads/main' }} diff --git a/.github/workflows/typecheck-agents-api-pr.yml b/.github/workflows/typecheck-agents-api-pr.yml new file mode 100644 index 000000000..513390883 --- /dev/null +++ b/.github/workflows/typecheck-agents-api-pr.yml @@ -0,0 +1,56 @@ +name: Typecheck agents-api +run-name: ${{ github.actor }} is typechecking the code + +# TODO: Fix CI github actions +# SCRUM-26 + +on: [pull_request] + +jobs: + Typecheck: + runs-on: ubuntu-latest + + steps: + - uses: actions/checkout@v4 + + - uses: actions/setup-python@v5 + with: + python-version: "3.12" + + - name: Install and configure Poetry + uses: snok/install-poetry@v1 + + - name: Configure Poetry to use .venv + run: | + cd agents-api + poetry config virtualenvs.in-project true + + - name: Cache Poetry virtualenv + uses: actions/cache@v4 + with: + path: agents-api/.venv + key: ${{ runner.os }}-agents-api-poetry-${{ hashFiles('agents-api/poetry.lock') }} + restore-keys: | + ${{ runner.os }}-agents-api-poetry- + + - name: Cache pytype + uses: actions/cache@v4 + with: + path: agents-api/.pytype + key: ${{ runner.os }}-agents-api-pytype-${{ github.base_ref }} + restore-keys: | + ${{ runner.os }}-agents-api-pytype- + + - name: Install dependencies + run: | + cd agents-api + poetry install + + - name: Typecheck + run: | + cd agents-api + poetry run poe typecheck + +concurrency: + group: ${{ github.workflow }}-${{ github.ref }} + cancel-in-progress: ${{ github.ref != 'refs/heads/main' }} diff --git a/agents-api/agents_api/activities/truncation.py b/agents-api/agents_api/activities/truncation.py index 4742ee6d4..afdb43da4 100644 --- a/agents-api/agents_api/activities/truncation.py +++ b/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 diff --git a/agents-api/agents_api/autogen/Docs.py b/agents-api/agents_api/autogen/Docs.py index a7023ddfc..0c1524e48 100644 --- a/agents-api/agents_api/autogen/Docs.py +++ b/agents-api/agents_api/autogen/Docs.py @@ -57,6 +57,13 @@ class Doc(BaseModel): """ Contents of the document """ + embeddings: Annotated[ + list[float] | list[list[float]] | None, + Field(None, json_schema_extra={"readOnly": True}), + ] + """ + Embeddings for the document + """ class DocOwner(BaseModel): diff --git a/agents-api/agents_api/autogen/Tools.py b/agents-api/agents_api/autogen/Tools.py index c2d4199a2..8227e5759 100644 --- a/agents-api/agents_api/autogen/Tools.py +++ b/agents-api/agents_api/autogen/Tools.py @@ -33,10 +33,6 @@ class CreateToolRequest(BaseModel): model_config = ConfigDict( populate_by_name=True, ) - type: Literal["function", "integration", "system", "api_call"] = "function" - """ - Whether this tool is a `function`, `api_call`, `system` etc. (Only `function` tool supported right now) - """ name: Annotated[str, Field(max_length=40, pattern="^[^\\W0-9]\\w*$")] """ Name of the tool (must be unique for this agent and a valid python identifier string ) @@ -168,10 +164,6 @@ class NamedToolChoice(BaseModel): model_config = ConfigDict( populate_by_name=True, ) - type: Literal["function", "integration", "system", "api_call"] - """ - Whether this tool is a `function`, `api_call`, `system` etc. (Only `function` tool supported right now) - """ function: FunctionCallOption | None = None @@ -183,10 +175,6 @@ class PatchToolRequest(BaseModel): model_config = ConfigDict( populate_by_name=True, ) - type: Literal["function", "integration", "system", "api_call"] = "function" - """ - Whether this tool is a `function`, `api_call`, `system` etc. (Only `function` tool supported right now) - """ name: Annotated[str | None, Field(None, max_length=40, pattern="^[^\\W0-9]\\w*$")] """ Name of the tool (must be unique for this agent and a valid python identifier string ) @@ -247,10 +235,6 @@ class Tool(BaseModel): model_config = ConfigDict( populate_by_name=True, ) - type: Literal["function", "integration", "system", "api_call"] = "function" - """ - Whether this tool is a `function`, `api_call`, `system` etc. (Only `function` tool supported right now) - """ name: Annotated[str, Field(max_length=40, pattern="^[^\\W0-9]\\w*$")] """ Name of the tool (must be unique for this agent and a valid python identifier string ) @@ -291,10 +275,6 @@ class UpdateToolRequest(BaseModel): model_config = ConfigDict( populate_by_name=True, ) - type: Literal["function", "integration", "system", "api_call"] = "function" - """ - Whether this tool is a `function`, `api_call`, `system` etc. (Only `function` tool supported right now) - """ name: Annotated[str, Field(max_length=40, pattern="^[^\\W0-9]\\w*$")] """ Name of the tool (must be unique for this agent and a valid python identifier string ) @@ -316,14 +296,3 @@ class ChosenFunctionCall(ChosenToolCall): """ The function to call """ - - -class NamedFunctionChoice(NamedToolChoice): - model_config = ConfigDict( - populate_by_name=True, - ) - type: Literal["function"] = "function" - function: FunctionCallOption - """ - The function to call - """ diff --git a/agents-api/agents_api/autogen/openapi_model.py b/agents-api/agents_api/autogen/openapi_model.py index 48811eb20..0d9390816 100644 --- a/agents-api/agents_api/autogen/openapi_model.py +++ b/agents-api/agents_api/autogen/openapi_model.py @@ -6,7 +6,14 @@ import jinja2 from litellm.utils import _select_tokenizer as select_tokenizer from litellm.utils import token_counter -from pydantic import AwareDatetime, Field, field_validator, model_validator, validator +from pydantic import ( + AwareDatetime, + Field, + computed_field, + field_validator, + model_validator, + validator, +) from ..common.utils.datetime import utcnow from .Agents import * @@ -70,92 +77,39 @@ class InputChatMLMessage(Message): pass -# Custom types (not generated correctly) -# -------------------------------------- - -ChatMLContent = ( - list[ChatMLTextContentPart | ChatMLImageContentPart] - | Tool - | ChosenToolCall - | str - | ToolResponse - | list[ - list[ChatMLTextContentPart | ChatMLImageContentPart] - | Tool - | ChosenToolCall - | str - | ToolResponse - ] -) - -# Extract ChatMLRole -ChatMLRole = BaseEntry.model_fields["role"].annotation - -# Extract ChatMLSource -ChatMLSource = BaseEntry.model_fields["source"].annotation - -# Extract ExecutionStatus -ExecutionStatus = Execution.model_fields["status"].annotation - -# Extract TransitionType -TransitionType = Transition.model_fields["type"].annotation - -# Assertions to ensure consistency (optional, but recommended for runtime checks) -assert ChatMLRole == BaseEntry.model_fields["role"].annotation -assert ChatMLSource == BaseEntry.model_fields["source"].annotation -assert ExecutionStatus == Execution.model_fields["status"].annotation -assert TransitionType == Transition.model_fields["type"].annotation - - -# Create models -# ------------- +# Patches +# ------- -class CreateTransitionRequest(Transition): - # The following fields are optional in this +def type_property(self: BaseModel) -> str: + return ( + "function" + if self.function + else "integration" + if self.integration + else "system" + if self.system + else "api_call" + if self.api_call + else None + ) - id: UUID | None = None - execution_id: UUID | None = None - created_at: AwareDatetime | None = None - updated_at: AwareDatetime | None = None - metadata: dict[str, Any] | None = None - task_token: str | None = None +# Patch original Tool class to add 'type' property +TaskTool.type = computed_field(property(type_property)) -class CreateEntryRequest(BaseEntry): - timestamp: Annotated[ - float, Field(ge=0.0, default_factory=lambda: utcnow().timestamp()) - ] +# Patch original Tool class to add 'type' property +Tool.type = computed_field(property(type_property)) - @classmethod - def from_model_input( - cls: Type[Self], - model: str, - *, - role: ChatMLRole, - content: ChatMLContent, - name: str | None = None, - source: ChatMLSource, - **kwargs: dict, - ) -> Self: - tokenizer: dict = select_tokenizer(model=model) - token_count = token_counter( - model=model, messages=[{"role": role, "content": content, "name": name}] - ) +# Patch original UpdateToolRequest class to add 'type' property +UpdateToolRequest.type = computed_field(property(type_property)) - return cls( - role=role, - content=content, - name=name, - source=source, - tokenizer=tokenizer["type"], - token_count=token_count, - **kwargs, - ) +# Patch original PatchToolRequest class to add 'type' property +PatchToolRequest.type = computed_field(property(type_property)) # Patch Task Workflow Steps -# -------------------------------------- +# ------------------------- def validate_python_expression(expr: str) -> tuple[bool, str]: @@ -186,145 +140,255 @@ def validate_jinja_template(template: str) -> tuple[bool, str]: return False, f"TemplateSyntaxError in '{template}': {str(e)}" -_EvaluateStep = EvaluateStep +@field_validator("evaluate") +def validate_evaluate_expressions(cls, v): + for key, expr in v.items(): + is_valid, error = validate_python_expression(expr) + if not is_valid: + raise ValueError(f"Invalid Python expression in key '{key}': {error}") + return v + + +EvaluateStep.validate_evaluate_expressions = validate_evaluate_expressions -class EvaluateStep(_EvaluateStep): - @field_validator("evaluate") - def validate_evaluate_expressions(cls, v): +@field_validator("arguments") +def validate_arguments(cls, v): + if isinstance(v, dict): for key, expr in v.items(): - is_valid, error = validate_python_expression(expr) - if not is_valid: - raise ValueError(f"Invalid Python expression in key '{key}': {error}") - return v + if isinstance(expr, str): + is_valid, error = validate_python_expression(expr) + if not is_valid: + raise ValueError( + f"Invalid Python expression in arguments key '{key}': {error}" + ) + return v -_ToolCallStep = ToolCallStep +ToolCallStep.validate_arguments = validate_arguments -class ToolCallStep(_ToolCallStep): - @field_validator("arguments") - def validate_arguments(cls, v): - if isinstance(v, dict): - for key, expr in v.items(): - if isinstance(expr, str): - is_valid, error = validate_python_expression(expr) - if not is_valid: - raise ValueError( - f"Invalid Python expression in arguments key '{key}': {error}" - ) - return v +# Add the new validator function +@field_validator("prompt") +def validate_prompt(cls, v): + if isinstance(v, str): + is_valid, error = validate_jinja_template(v) + if not is_valid: + raise ValueError(f"Invalid Jinja template in prompt: {error}") + elif isinstance(v, list): + for item in v: + if "content" in item: + is_valid, error = validate_jinja_template(item["content"]) + if not is_valid: + raise ValueError( + f"Invalid Jinja template in prompt content: {error}" + ) + return v -_PromptStep = PromptStep +# Patch the original PromptStep class to add the new validator +PromptStep.validate_prompt = validate_prompt -class PromptStep(_PromptStep): - @field_validator("prompt") - def validate_prompt(cls, v): - if isinstance(v, str): - is_valid, error = validate_jinja_template(v) - if not is_valid: - raise ValueError(f"Invalid Jinja template in prompt: {error}") - elif isinstance(v, list): - for item in v: - if "content" in item: - is_valid, error = validate_jinja_template(item["content"]) - if not is_valid: - raise ValueError( - f"Invalid Jinja template in prompt content: {error}" - ) - return v +@field_validator("set") +def validate_set_expressions(cls, v): + for key, expr in v.items(): + is_valid, error = validate_python_expression(expr) + if not is_valid: + raise ValueError(f"Invalid Python expression in set key '{key}': {error}") + return v -_SetStep = SetStep +SetStep.validate_set_expressions = validate_set_expressions -class SetStep(_SetStep): - @field_validator("set") - def validate_set_expressions(cls, v): - for key, expr in v.items(): - is_valid, error = validate_python_expression(expr) - if not is_valid: - raise ValueError( - f"Invalid Python expression in set key '{key}': {error}" - ) - return v +@field_validator("log") +def validate_log_template(cls, v): + is_valid, error = validate_jinja_template(v) + if not is_valid: + raise ValueError(f"Invalid Jinja template in log: {error}") + return v -_LogStep = LogStep +LogStep.validate_log_template = validate_log_template -class LogStep(_LogStep): - @field_validator("log") - def validate_log_template(cls, v): - is_valid, error = validate_jinja_template(v) +@field_validator("return_") +def validate_return_expressions(cls, v): + for key, expr in v.items(): + is_valid, error = validate_python_expression(expr) if not is_valid: - raise ValueError(f"Invalid Jinja template in log: {error}") - return v + raise ValueError( + f"Invalid Python expression in return key '{key}': {error}" + ) + return v -_ReturnStep = ReturnStep +ReturnStep.validate_return_expressions = validate_return_expressions -class ReturnStep(_ReturnStep): - @field_validator("return_") - def validate_return_expressions(cls, v): +@field_validator("arguments") +def validate_yield_arguments(cls, v): + if isinstance(v, dict): for key, expr in v.items(): is_valid, error = validate_python_expression(expr) if not is_valid: raise ValueError( - f"Invalid Python expression in return key '{key}': {error}" + f"Invalid Python expression in yield arguments key '{key}': {error}" ) - return v + return v -_YieldStep = YieldStep +YieldStep.validate_yield_arguments = validate_yield_arguments -class YieldStep(_YieldStep): - @field_validator("arguments") - def validate_yield_arguments(cls, v): - if isinstance(v, dict): - for key, expr in v.items(): - is_valid, error = validate_python_expression(expr) - if not is_valid: - raise ValueError( - f"Invalid Python expression in yield arguments key '{key}': {error}" - ) - return v +@field_validator("if_") +def validate_if_expression(cls, v): + is_valid, error = validate_python_expression(v) + if not is_valid: + raise ValueError(f"Invalid Python expression in if condition: {error}") + return v -_IfElseWorkflowStep = IfElseWorkflowStep +IfElseWorkflowStep.validate_if_expression = validate_if_expression -class IfElseWorkflowStep(_IfElseWorkflowStep): - @field_validator("if_") - def validate_if_expression(cls, v): +@field_validator("over") +def validate_over_expression(cls, v): + is_valid, error = validate_python_expression(v) + if not is_valid: + raise ValueError(f"Invalid Python expression in over: {error}") + return v + + +@field_validator("reduce") +def validate_reduce_expression(cls, v): + if v is not None: is_valid, error = validate_python_expression(v) if not is_valid: - raise ValueError(f"Invalid Python expression in if condition: {error}") - return v + raise ValueError(f"Invalid Python expression in reduce: {error}") + return v -_MapReduceStep = MapReduceStep +MapReduceStep.validate_over_expression = validate_over_expression +MapReduceStep.validate_reduce_expression = validate_reduce_expression -class MapReduceStep(_MapReduceStep): - @field_validator("over") - def validate_over_expression(cls, v): - is_valid, error = validate_python_expression(v) - if not is_valid: - raise ValueError(f"Invalid Python expression in over: {error}") - return v +# Patch workflow +# -------------- - @field_validator("reduce") - def validate_reduce_expression(cls, v): - if v is not None: - is_valid, error = validate_python_expression(v) - if not is_valid: - raise ValueError(f"Invalid Python expression in reduce: {error}") - return v +_CreateTaskRequest = CreateTaskRequest + +CreateTaskRequest.model_config = ConfigDict( + **{ + **_CreateTaskRequest.model_config, + "extra": "allow", + } +) + + +@model_validator(mode="after") +def validate_subworkflows(self): + subworkflows = { + k: v + for k, v in self.model_dump().items() + if k not in _CreateTaskRequest.model_fields + } + + for workflow_name, workflow_definition in subworkflows.items(): + try: + WorkflowType.model_validate(workflow_definition) + setattr(self, workflow_name, WorkflowType(workflow_definition)) + except Exception as e: + raise ValueError(f"Invalid subworkflow '{workflow_name}': {str(e)}") + return self + + +CreateTaskRequest.validate_subworkflows = validate_subworkflows + + +# Custom types (not generated correctly) +# -------------------------------------- + +ChatMLContent = ( + list[ChatMLTextContentPart | ChatMLImageContentPart] + | Tool + | ChosenToolCall + | str + | ToolResponse + | list[ + list[ChatMLTextContentPart | ChatMLImageContentPart] + | Tool + | ChosenToolCall + | str + | ToolResponse + ] +) + +# Extract ChatMLRole +ChatMLRole = BaseEntry.model_fields["role"].annotation + +# Extract ChatMLSource +ChatMLSource = BaseEntry.model_fields["source"].annotation + +# Extract ExecutionStatus +ExecutionStatus = Execution.model_fields["status"].annotation + +# Extract TransitionType +TransitionType = Transition.model_fields["type"].annotation + +# Assertions to ensure consistency (optional, but recommended for runtime checks) +assert ChatMLRole == BaseEntry.model_fields["role"].annotation +assert ChatMLSource == BaseEntry.model_fields["source"].annotation +assert ExecutionStatus == Execution.model_fields["status"].annotation +assert TransitionType == Transition.model_fields["type"].annotation + + +# Create models +# ------------- + + +class CreateTransitionRequest(Transition): + # The following fields are optional in this + + id: UUID | None = None + execution_id: UUID | None = None + created_at: AwareDatetime | None = None + updated_at: AwareDatetime | None = None + metadata: dict[str, Any] | None = None + task_token: str | None = None + + +class CreateEntryRequest(BaseEntry): + timestamp: Annotated[ + float, Field(ge=0.0, default_factory=lambda: utcnow().timestamp()) + ] + + @classmethod + def from_model_input( + cls: Type[Self], + model: str, + *, + role: ChatMLRole, + content: ChatMLContent, + name: str | None = None, + source: ChatMLSource, + **kwargs: dict, + ) -> Self: + tokenizer: dict = select_tokenizer(model=model) + token_count = token_counter( + model=model, messages=[{"role": role, "content": content, "name": name}] + ) + + return cls( + role=role, + content=content, + name=name, + source=source, + tokenizer=tokenizer["type"], + token_count=token_count, + **kwargs, + ) # Workflow related models @@ -427,34 +491,6 @@ class Task(_Task): ] -_CreateTaskRequest = CreateTaskRequest - - -class CreateTaskRequest(_CreateTaskRequest): - model_config = ConfigDict( - **{ - **_CreateTaskRequest.model_config, - "extra": "allow", - } - ) - - @model_validator(mode="after") - def validate_subworkflows(self) -> Self: - subworkflows = { - k: v - for k, v in self.model_dump().items() - if k not in _CreateTaskRequest.model_fields - } - - for workflow_name, workflow_definition in subworkflows.items(): - try: - WorkflowType.model_validate(workflow_definition) - setattr(self, workflow_name, WorkflowType(workflow_definition)) - except Exception as e: - raise ValueError(f"Invalid subworkflow '{workflow_name}': {str(e)}") - return self - - CreateOrUpdateTaskRequest = CreateTaskRequest _PatchTaskRequest = PatchTaskRequest diff --git a/agents-api/agents_api/clients/worker/worker.py b/agents-api/agents_api/clients/worker/worker.py index 1deb8d1c3..8befa3080 100644 --- a/agents-api/agents_api/clients/worker/worker.py +++ b/agents-api/agents_api/clients/worker/worker.py @@ -1,7 +1,6 @@ import httpx -from agents_api.env import temporal_worker_url - +from ...env import temporal_worker_url from .types import ( MemoryManagementTask, MemoryManagementTaskArgs, diff --git a/agents-api/agents_api/common/exceptions/__init__.py b/agents-api/agents_api/common/exceptions/__init__.py index a38a546a2..571691082 100644 --- a/agents-api/agents_api/common/exceptions/__init__.py +++ b/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. """ diff --git a/agents-api/agents_api/common/exceptions/agents.py b/agents-api/agents_api/common/exceptions/agents.py index c412eba35..e58f25104 100644 --- a/agents-api/agents_api/common/exceptions/agents.py +++ b/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", diff --git a/agents-api/agents_api/common/exceptions/sessions.py b/agents-api/agents_api/common/exceptions/sessions.py index d6c04aff5..6e9941d43 100644 --- a/agents-api/agents_api/common/exceptions/sessions.py +++ b/agents-api/agents_api/common/exceptions/sessions.py @@ -8,29 +8,24 @@ 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. """ diff --git a/agents-api/agents_api/common/exceptions/users.py b/agents-api/agents_api/common/exceptions/users.py index 2f7f58ed6..cf4e995ad 100644 --- a/agents-api/agents_api/common/exceptions/users.py +++ b/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. diff --git a/agents-api/agents_api/common/protocol/agents.py b/agents-api/agents_api/common/protocol/agents.py index 222f91f01..89f7f4a2d 100644 --- a/agents-api/agents_api/common/protocol/agents.py +++ b/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.""" diff --git a/agents-api/agents_api/common/protocol/tasks.py b/agents-api/agents_api/common/protocol/tasks.py index 6892d7098..bbb5c28d3 100644 --- a/agents-api/agents_api/common/protocol/tasks.py +++ b/agents-api/agents_api/common/protocol/tasks.py @@ -218,19 +218,28 @@ def task_to_spec( task: Task | CreateTaskRequest | UpdateTaskRequest | PatchTaskRequest, **model_opts ) -> TaskSpecDef | PartialTaskSpecDef: task_data = task.model_dump(**model_opts) - main = task_data.pop("main") - workflows = [Workflow(name="main", steps=main)] + if "tools" in task_data: + del task_data["tools"] - for k in list(task_data.keys()): - if k in TaskSpec.model_fields.keys(): - continue + tools = [] + for tool in task.tools: + tool_spec = getattr(tool, tool.type) - steps = task_data.pop(k) - workflows.append(Workflow(name=k, steps=steps)) + tools.append( + TaskToolDef( + type=tool.type, + spec=tool_spec.model_dump(), + **tool.model_dump(exclude={"type"}), + ) + ) - tools = task_data.pop("tools", []) - tools = [TaskToolDef(spec=tool.pop(tool["type"]), **tool) for tool in tools] + workflows = [Workflow(name="main", steps=task_data.pop("main"))] + + for key, steps in list(task_data.items()): + if key not in TaskSpec.model_fields: + workflows.append(Workflow(name=key, steps=steps)) + del task_data[key] cls = PartialTaskSpecDef if isinstance(task, PatchTaskRequest) else TaskSpecDef diff --git a/agents-api/agents_api/common/utils/json.py b/agents-api/agents_api/common/utils/json.py index 3157af9c8..f1b82742a 100644 --- a/agents-api/agents_api/common/utils/json.py +++ b/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 ) diff --git a/agents-api/agents_api/common/utils/messages.py b/agents-api/agents_api/common/utils/messages.py index fd971296c..c94808324 100644 --- a/agents-api/agents_api/common/utils/messages.py +++ b/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, ) diff --git a/agents-api/agents_api/exceptions.py b/agents-api/agents_api/exceptions.py index fbb8f00f8..4fbf8a2b9 100644 --- a/agents-api/agents_api/exceptions.py +++ b/agents-api/agents_api/exceptions.py @@ -3,11 +3,15 @@ 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}" @@ -15,5 +19,7 @@ def __init__(self, token_count, max_tokens) -> None: class UnknownTokenizerError(AgentsBaseException): + """Exception raised when tokenizer is unknown.""" + def __init__(self) -> None: super().__init__("unknown tokenizer") diff --git a/agents-api/agents_api/models/agent/create_agent.py b/agents-api/agents_api/models/agent/create_agent.py index 73be4ec77..98daab540 100644 --- a/agents-api/agents_api/models/agent/create_agent.py +++ b/agents-api/agents_api/models/agent/create_agent.py @@ -30,11 +30,24 @@ lambda e: isinstance(e, QueryException) and "asserted to return some results, but returned none" in str(e): lambda *_: HTTPException( - detail="developer not found", status_code=403 + detail="Developer not found. Please ensure the provided auth token (which refers to your developer_id) is valid and the developer has the necessary permissions to create an agent.", + status_code=403, + ), + QueryException: partialclass( + HTTPException, + status_code=400, + detail="A database query failed to return the expected results. This might occur if the requested resource doesn't exist or your query parameters are incorrect.", + ), + ValidationError: partialclass( + HTTPException, + status_code=400, + detail="Input validation failed. Please check the provided data for missing or incorrect fields, and ensure it matches the required format.", + ), + TypeError: partialclass( + HTTPException, + status_code=400, + detail="A type mismatch occurred. This likely means the data provided is of an incorrect type (e.g., string instead of integer). Please review the input and try again.", ), - QueryException: partialclass(HTTPException, status_code=400), - ValidationError: partialclass(HTTPException, status_code=400), - TypeError: partialclass(HTTPException, status_code=400), } ) @wrap_in_class( @@ -55,12 +68,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() 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 156c8ae61..3902c1bb5 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 @@ -27,9 +27,21 @@ @rewrap_exceptions( { - QueryException: partialclass(HTTPException, status_code=400), - ValidationError: partialclass(HTTPException, status_code=400), - TypeError: partialclass(HTTPException, status_code=400), + QueryException: partialclass( + HTTPException, + status_code=400, + detail="A database query failed to return the expected results. This might occur if the requested resource doesn't exist or your query parameters are incorrect.", + ), + ValidationError: partialclass( + HTTPException, + status_code=400, + detail="Input validation failed. Please check the provided data for missing or incorrect fields, and ensure it matches the required format.", + ), + TypeError: partialclass( + HTTPException, + status_code=400, + detail="A type mismatch occurred. This likely means the data provided is of an incorrect type (e.g., string instead of integer). Please review the input and try again.", + ), } ) @wrap_in_class( @@ -47,18 +59,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 diff --git a/agents-api/agents_api/models/agent/delete_agent.py b/agents-api/agents_api/models/agent/delete_agent.py index 926007bdf..60de66292 100644 --- a/agents-api/agents_api/models/agent/delete_agent.py +++ b/agents-api/agents_api/models/agent/delete_agent.py @@ -30,11 +30,24 @@ lambda e: isinstance(e, QueryException) and "Developer does not own resource" in e.resp["display"]: lambda *_: HTTPException( - detail="developer not found or doesnt own resource", status_code=404 + detail="The specified developer does not own the requested resource. Please verify the ownership or check if the developer ID is correct.", + status_code=404, + ), + QueryException: partialclass( + HTTPException, + status_code=400, + detail="A database query failed to return the expected results. This might occur if the requested resource doesn't exist or your query parameters are incorrect.", + ), + ValidationError: partialclass( + HTTPException, + status_code=400, + detail="Input validation failed. Please check the provided data for missing or incorrect fields, and ensure it matches the required format.", + ), + TypeError: partialclass( + HTTPException, + status_code=400, + detail="A type mismatch occurred. This likely means the data provided is of an incorrect type (e.g., string instead of integer). Please review the input and try again.", ), - QueryException: partialclass(HTTPException, status_code=400), - ValidationError: partialclass(HTTPException, status_code=400), - TypeError: partialclass(HTTPException, status_code=400), } ) @wrap_in_class( @@ -54,12 +67,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 = [ diff --git a/agents-api/agents_api/models/agent/get_agent.py b/agents-api/agents_api/models/agent/get_agent.py index 956fa46a5..bdae85fcb 100644 --- a/agents-api/agents_api/models/agent/get_agent.py +++ b/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. @@ -101,6 +101,8 @@ def get_agent(*, developer_id: UUID, agent_id: UUID) -> tuple[list[str], dict]: "min_p": min_p, "preset": preset, } + + :limit 1 """ queries = [ diff --git a/agents-api/agents_api/models/agent/patch_agent.py b/agents-api/agents_api/models/agent/patch_agent.py index 72fdc7811..364d1a974 100644 --- a/agents-api/agents_api/models/agent/patch_agent.py +++ b/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) diff --git a/agents-api/agents_api/models/agent/update_agent.py b/agents-api/agents_api/models/agent/update_agent.py index be7e9ea21..992bb9796 100644 --- a/agents-api/agents_api/models/agent/update_agent.py +++ b/agents-api/agents_api/models/agent/update_agent.py @@ -46,13 +46,13 @@ def update_agent( Constructs and executes a datalog query to update an agent and its default settings in the 'cozodb' database. Parameters: - - agent_id (UUID): The unique identifier of the agent to be updated. - - developer_id (UUID): The unique identifier of the developer associated with the agent. - - data (UpdateAgentRequest): The request payload containing the updated agent data. - - client (CozoClient, optional): The database client used to execute the query. Defaults to a pre-configured client instance. + agent_id (UUID): The unique identifier of the agent to be updated. + developer_id (UUID): The unique identifier of the developer associated with the agent. + data (UpdateAgentRequest): The request payload containing the updated agent data. + client (CozoClient, optional): The database client used to execute the query. Defaults to a pre-configured client instance. Returns: - ResourceUpdatedResponse: The updated agent data. + ResourceUpdatedResponse: The updated agent data. """ default_settings = ( data.default_settings.model_dump(exclude_none=True) diff --git a/agents-api/agents_api/models/chat/gather_messages.py b/agents-api/agents_api/models/chat/gather_messages.py index c1683653f..b6bfdc217 100644 --- a/agents-api/agents_api/models/chat/gather_messages.py +++ b/agents-api/agents_api/models/chat/gather_messages.py @@ -6,8 +6,7 @@ from pycozo.client import QueryException from pydantic import ValidationError -from agents_api.autogen.Chat import ChatInput - +from ...autogen.Chat import ChatInput from ...autogen.openapi_model import DocReference, History from ...clients import litellm from ...common.protocol.developers import Developer diff --git a/agents-api/agents_api/models/chat/get_cached_response.py b/agents-api/agents_api/models/chat/get_cached_response.py index 5440a6703..368c88567 100644 --- a/agents-api/agents_api/models/chat/get_cached_response.py +++ b/agents-api/agents_api/models/chat/get_cached_response.py @@ -9,6 +9,7 @@ def get_cached_response(key: str) -> tuple[str, dict]: query = """ input[key] <- [[$key]] ?[key, value] := input[key], *session_cache{key, value} + :limit 1 """ return (query, {"key": key}) 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 4d521acb0..9be2d64aa 100644 --- a/agents-api/agents_api/models/chat/prepare_chat_context.py +++ b/agents-api/agents_api/models/chat/prepare_chat_context.py @@ -119,6 +119,8 @@ def prepare_chat_context( ?[{', '.join(session_data_fields)}, toolsets] := *_session_data_json {{ {', '.join(session_data_fields)} }}, *_toolsets_json {{ toolsets }} + + :limit 1 """ queries = [ diff --git a/agents-api/agents_api/models/developer/get_developer.py b/agents-api/agents_api/models/developer/get_developer.py index 31ade5334..0ae5421aa 100644 --- a/agents-api/agents_api/models/developer/get_developer.py +++ b/agents-api/agents_api/models/developer/get_developer.py @@ -67,6 +67,8 @@ def get_developer( created_at, updated_at, } + + :limit 1 """ return (query, {"developer_id": developer_id}) diff --git a/agents-api/agents_api/models/docs/create_doc.py b/agents-api/agents_api/models/docs/create_doc.py index ee26df484..1c667bd51 100644 --- a/agents-api/agents_api/models/docs/create_doc.py +++ b/agents-api/agents_api/models/docs/create_doc.py @@ -50,10 +50,10 @@ def create_doc( Constructs and executes a datalog query to create a new document and its associated snippets in the 'cozodb' database. Parameters: - - owner_type (Literal["user", "agent"]): The type of the owner of the document. - - owner_id (UUID): The UUID of the document owner. - - id (UUID): The UUID of the document to be created. - - data (CreateDocRequest): The content of the document. + owner_type (Literal["user", "agent"]): The type of the owner of the document. + owner_id (UUID): The UUID of the document owner. + id (UUID): The UUID of the document to be created. + data (CreateDocRequest): The content of the document. """ doc_id = str(doc_id or uuid4()) diff --git a/agents-api/agents_api/models/docs/embed_snippets.py b/agents-api/agents_api/models/docs/embed_snippets.py index e810d0379..8d8ae1e62 100644 --- a/agents-api/agents_api/models/docs/embed_snippets.py +++ b/agents-api/agents_api/models/docs/embed_snippets.py @@ -49,9 +49,9 @@ def embed_snippets( """Embeds document snippets in the cozodb database. Parameters: - doc_id (UUID): The unique identifier for the document. - snippet_indices (list[int]): Indices of the snippets in the document. - embeddings (list[list[float]]): Embedding vectors for the snippets. + doc_id (UUID): The unique identifier for the document. + snippet_indices (list[int]): Indices of the snippets in the document. + embeddings (list[list[float]]): Embedding vectors for the snippets. """ doc_id = str(doc_id) @@ -81,6 +81,7 @@ def embed_snippets( }}, index > {max(snippet_indices)} + :limit 1 :assert none """ diff --git a/agents-api/agents_api/models/docs/get_doc.py b/agents-api/agents_api/models/docs/get_doc.py index 84cd181ec..e81084985 100644 --- a/agents-api/agents_api/models/docs/get_doc.py +++ b/agents-api/agents_api/models/docs/get_doc.py @@ -37,6 +37,7 @@ one=True, transform=lambda d: { "content": [s[1] for s in sorted(d["snippet_data"], key=lambda x: x[0])], + "embeddings": [s[2] for s in sorted(d["snippet_data"], key=lambda x: x[0])], **d, }, ) @@ -68,8 +69,9 @@ def get_doc( doc_id, index, content, + embedding, }, - snippet_data = [index, content] + snippet_data = [index, content, embedding] ?[ id, @@ -85,6 +87,8 @@ def get_doc( metadata, }, snippets[snippet_data] + + :limit 1 """ queries = [ diff --git a/agents-api/agents_api/models/docs/list_docs.py b/agents-api/agents_api/models/docs/list_docs.py index afdf06c2d..3c095c2db 100644 --- a/agents-api/agents_api/models/docs/list_docs.py +++ b/agents-api/agents_api/models/docs/list_docs.py @@ -34,6 +34,7 @@ Doc, transform=lambda d: { "content": [s[1] for s in sorted(d["snippet_data"], key=lambda x: x[0])], + "embeddings": [s[2] for s in sorted(d["snippet_data"], key=lambda x: x[0])], **d, }, ) @@ -67,8 +68,9 @@ def list_docs( doc_id: id, index, content, + embedding, }}, - snippet_data = [index, content] + snippet_data = [index, content, embedding] ?[ owner_type, 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 d6bed97fc..83418aa21 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 @@ -56,12 +56,12 @@ def search_docs_by_embedding( Searches for document snippets in CozoDB by embedding query. Parameters: - - owner_type (Literal["user", "agent"]): The type of the owner of the documents. - - owner_id (UUID): The unique identifier of the owner. - - query_embedding (list[float]): The embedding vector of the query. - - k (int, optional): The number of nearest neighbors to retrieve. Defaults to 3. - - confidence (float, optional): The confidence threshold for filtering results. Defaults to 0.8. - - mmr_lambda (float, optional): The lambda parameter for MMR. Defaults to 0.25. + owner_type (Literal["user", "agent"]): The type of the owner of the documents. + owner_id (UUID): The unique identifier of the owner. + query_embedding (list[float]): The embedding vector of the query. + k (int, optional): The number of nearest neighbors to retrieve. Defaults to 3. + confidence (float, optional): The confidence threshold for filtering results. Defaults to 0.8. + mmr_lambda (float, optional): The lambda parameter for MMR. Defaults to 0.25. """ assert len(query_embedding) == embedding_size 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 eeae8362c..bb700a494 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 @@ -53,9 +53,9 @@ def search_docs_by_text( Searches for document snippets in CozoDB by embedding query. Parameters: - - owners (list[tuple[Literal["user", "agent"], UUID]]): The type of the owner of the documents. - - query (str): The query string. - - k (int, optional): The number of nearest neighbors to retrieve. Defaults to 3. + owners (list[tuple[Literal["user", "agent"], UUID]]): The type of the owner of the documents. + query (str): The query string. + k (int, optional): The number of nearest neighbors to retrieve. Defaults to 3. """ owners: list[list[str]] = [ diff --git a/agents-api/agents_api/models/entry/delete_entries.py b/agents-api/agents_api/models/entry/delete_entries.py index 48c37cd25..c98b6c7d2 100644 --- a/agents-api/agents_api/models/entry/delete_entries.py +++ b/agents-api/agents_api/models/entry/delete_entries.py @@ -49,7 +49,7 @@ def delete_entries_for_session( Constructs and returns a datalog query for deleting entries associated with a given session ID from the 'cozodb' database. Parameters: - - session_id (UUID): The unique identifier of the session whose entries are to be deleted. + session_id (UUID): The unique identifier of the session whose entries are to be deleted. """ delete_query = """ diff --git a/agents-api/agents_api/models/entry/get_history.py b/agents-api/agents_api/models/entry/get_history.py index 8918d357f..bd3c0397a 100644 --- a/agents-api/agents_api/models/entry/get_history.py +++ b/agents-api/agents_api/models/entry/get_history.py @@ -131,7 +131,9 @@ def get_history( session_relations[relations], session_id = to_uuid($session_id), created_at = now() - """ + + :limit 1 + """ queries = [ verify_developer_id_query(developer_id), diff --git a/agents-api/agents_api/models/execution/create_execution_transition.py b/agents-api/agents_api/models/execution/create_execution_transition.py index 42c6f1c22..f40395126 100644 --- a/agents-api/agents_api/models/execution/create_execution_transition.py +++ b/agents-api/agents_api/models/execution/create_execution_transition.py @@ -141,6 +141,8 @@ def create_execution_transition( found = length(prev_transitions), valid = if($next_type == "init", found == 0, found > 0), assert(valid, "Invalid transition"), + + :limit 1 """ # Prepare the insert query 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 d3ebf3783..e2b38789a 100644 --- a/agents-api/agents_api/models/execution/get_execution_transition.py +++ b/agents-api/agents_api/models/execution/get_execution_transition.py @@ -65,7 +65,9 @@ def get_execution_transition( is_null(next_tuple), null, {"workflow": next_tuple->0, "step": next_tuple->1}, - ), + ) + + :limit 1 """ get_query += filter 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 b4c9f9081..44eb8a4da 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 @@ -45,6 +45,7 @@ def get_paused_execution_token( execution_id = to_uuid($execution_id), status = "awaiting_input" + :limit 1 :assert some """ 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 bb0a462ef..8b1bf4604 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 @@ -45,6 +45,8 @@ def get_temporal_workflow_data( result_run_id, first_execution_run_id, } + + :limit 1 """ return ( diff --git a/agents-api/agents_api/models/execution/lookup_temporal_data.py b/agents-api/agents_api/models/execution/lookup_temporal_data.py index e3c369b94..35f09129b 100644 --- a/agents-api/agents_api/models/execution/lookup_temporal_data.py +++ b/agents-api/agents_api/models/execution/lookup_temporal_data.py @@ -43,6 +43,8 @@ def lookup_temporal_data( *temporal_executions_lookup { id, execution_id, run_id, first_execution_run_id, result_run_id } + + :limit 1 """ queries = [ 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 b763a7508..5f30e7f83 100644 --- a/agents-api/agents_api/models/execution/prepare_execution_input.py +++ b/agents-api/agents_api/models/execution/prepare_execution_input.py @@ -187,6 +187,8 @@ def prepare_execution_input( user = null, session = null, arguments = execution->"input" + + :limit 1 """ queries = [ diff --git a/agents-api/agents_api/models/execution/update_execution.py b/agents-api/agents_api/models/execution/update_execution.py index c2f40c7ff..4d0367ae4 100644 --- a/agents-api/agents_api/models/execution/update_execution.py +++ b/agents-api/agents_api/models/execution/update_execution.py @@ -82,6 +82,8 @@ def update_execution( ?[num] := valid_status[num], assert(num > 0, 'Invalid status') + + :limit 1 """ update_query = f""" diff --git a/agents-api/agents_api/models/session/delete_session.py b/agents-api/agents_api/models/session/delete_session.py index af9e331c7..81f8e1f7c 100644 --- a/agents-api/agents_api/models/session/delete_session.py +++ b/agents-api/agents_api/models/session/delete_session.py @@ -51,11 +51,11 @@ def delete_session( Deletes a session and its related data from the 'cozodb' database. Parameters: - - developer_id (UUID): The unique identifier for the developer. - - session_id (UUID): The unique identifier for the session to be deleted. + developer_id (UUID): The unique identifier for the developer. + session_id (UUID): The unique identifier for the session to be deleted. Returns: - - ResourceDeletedResponse: The response indicating the deletion of the session. + ResourceDeletedResponse: The response indicating the deletion of the session. """ session_id = str(session_id) developer_id = str(developer_id) diff --git a/agents-api/agents_api/models/session/get_session.py b/agents-api/agents_api/models/session/get_session.py index 2d8956eb7..45d971c3f 100644 --- a/agents-api/agents_api/models/session/get_session.py +++ b/agents-api/agents_api/models/session/get_session.py @@ -95,7 +95,10 @@ def get_session( token_budget, context_overflow, @ "END" - }, updated_at = to_int(validity) + }, + updated_at = to_int(validity) + + :limit 1 """ queries = [ diff --git a/agents-api/agents_api/models/session/list_sessions.py b/agents-api/agents_api/models/session/list_sessions.py index 14f4cad0d..b7d5c0049 100644 --- a/agents-api/agents_api/models/session/list_sessions.py +++ b/agents-api/agents_api/models/session/list_sessions.py @@ -41,7 +41,8 @@ def list_sessions( direction: Literal["asc", "desc"] = "desc", metadata_filter: dict[str, Any] = {}, ) -> tuple[list[str], dict]: - """Lists sessions from the 'cozodb' database based on the provided filters. + """ + Lists sessions from the 'cozodb' database based on the provided filters. Parameters: developer_id (UUID): The developer's ID to filter sessions by. diff --git a/agents-api/agents_api/models/session/patch_session.py b/agents-api/agents_api/models/session/patch_session.py index a9f3121b9..4a119a684 100644 --- a/agents-api/agents_api/models/session/patch_session.py +++ b/agents-api/agents_api/models/session/patch_session.py @@ -60,12 +60,13 @@ def patch_session( developer_id: UUID, data: PatchSessionRequest, ) -> tuple[list[str], dict]: - """Patch session data in the 'cozodb' database. + """ + Patch session data in the 'cozodb' database. Parameters: - - session_id (UUID): The unique identifier for the session to be updated. - - developer_id (UUID): The unique identifier for the developer making the update. - - data (PatchSessionRequest): The request payload containing the updates to apply. + session_id (UUID): The unique identifier for the session to be updated. + developer_id (UUID): The unique identifier for the developer making the update. + data (PatchSessionRequest): The request payload containing the updates to apply. """ update_data = data.model_dump(exclude_unset=True) 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 e8cfb7fc7..bbbd9c4cd 100644 --- a/agents-api/agents_api/models/session/prepare_session_data.py +++ b/agents-api/agents_api/models/session/prepare_session_data.py @@ -209,6 +209,8 @@ def prepare_session_data( session_data[session], user_data[users], agent_data[agents] + + :limit 1 """ queries = [ diff --git a/agents-api/agents_api/models/tools/create_tools.py b/agents-api/agents_api/models/tools/create_tools.py index dd8397797..fe7e28228 100644 --- a/agents-api/agents_api/models/tools/create_tools.py +++ b/agents-api/agents_api/models/tools/create_tools.py @@ -51,11 +51,11 @@ def create_tools( Constructs a datalog query for inserting tool records into the 'agent_functions' relation in the CozoDB. Parameters: - - agent_id (UUID): The unique identifier for the agent. - - data (list[CreateToolRequest]): A list of function definitions to be inserted. + agent_id (UUID): The unique identifier for the agent. + data (list[CreateToolRequest]): A list of function definitions to be inserted. Returns: - list[Tool] + list[Tool] """ tools_data = [ @@ -80,6 +80,7 @@ def create_tools( name, } + :limit 1 :assert none """ diff --git a/agents-api/agents_api/models/tools/get_tool.py b/agents-api/agents_api/models/tools/get_tool.py index 5ea009064..465fd2efe 100644 --- a/agents-api/agents_api/models/tools/get_tool.py +++ b/agents-api/agents_api/models/tools/get_tool.py @@ -68,6 +68,8 @@ def get_tool( updated_at, created_at, } + + :limit 1 """ queries = [ diff --git a/agents-api/agents_api/models/tools/get_tool_args_from_metadata.py b/agents-api/agents_api/models/tools/get_tool_args_from_metadata.py index ade296a29..2cdb92cb9 100644 --- a/agents-api/agents_api/models/tools/get_tool_args_from_metadata.py +++ b/agents-api/agents_api/models/tools/get_tool_args_from_metadata.py @@ -46,6 +46,8 @@ def tool_args_for_task( # Right values overwrite left values # See: https://docs.cozodb.org/en/latest/functions.html#Func.Vector.concat values = concat(agent_{arg_type}, task_{arg_type}), + + :limit 1 """ queries = [ @@ -88,6 +90,8 @@ def tool_args_for_session( # Right values overwrite left values # See: https://docs.cozodb.org/en/latest/functions.html#Func.Vector.concat values = concat(agent_{arg_type}, session_{arg_type}), + + :limit 1 """ queries = [ diff --git a/agents-api/agents_api/models/tools/patch_tool.py b/agents-api/agents_api/models/tools/patch_tool.py index 5bbfe1c91..0d8304d7d 100644 --- a/agents-api/agents_api/models/tools/patch_tool.py +++ b/agents-api/agents_api/models/tools/patch_tool.py @@ -40,16 +40,16 @@ def patch_tool( *, developer_id: UUID, agent_id: UUID, tool_id: UUID, data: PatchToolRequest ) -> tuple[list[str], dict]: """ - # Execute the datalog query and return the results as a DataFrame + Execute the datalog query and return the results as a DataFrame Updates the tool information for a given agent and tool ID in the 'cozodb' database. Parameters: - - agent_id (UUID): The unique identifier of the agent. - - tool_id (UUID): The unique identifier of the tool to be updated. - - data (PatchToolRequest): The request payload containing the updated tool information. + agent_id (UUID): The unique identifier of the agent. + tool_id (UUID): The unique identifier of the tool to be updated. + data (PatchToolRequest): The request payload containing the updated tool information. Returns: - - ResourceUpdatedResponse: The updated tool data. + ResourceUpdatedResponse: The updated tool data. """ agent_id = str(agent_id) 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 9e9045e74..97db913c5 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 @@ -44,15 +44,15 @@ def create_or_update_user( Constructs and executes a datalog query to create a new user in the database. Parameters: - - user_id (UUID): The unique identifier for the user. - - developer_id (UUID): The unique identifier for the developer creating the user. - - name (str): The name of the user. - - about (str): A description of the user. - - metadata (dict, optional): A dictionary of metadata for the user. Defaults to an empty dict. - - client (CozoClient, optional): The CozoDB client instance to use for the query. Defaults to a preconfigured client instance. + user_id (UUID): The unique identifier for the user. + developer_id (UUID): The unique identifier for the developer creating the user. + name (str): The name of the user. + about (str): A description of the user. + metadata (dict, optional): A dictionary of metadata for the user. Defaults to an empty dict. + client (CozoClient, optional): The CozoDB client instance to use for the query. Defaults to a preconfigured client instance. Returns: - User: The newly created user record. + User: The newly created user record. """ # Extract the user data from the payload diff --git a/agents-api/agents_api/models/user/delete_user.py b/agents-api/agents_api/models/user/delete_user.py index b5fcb8424..0532f5cfa 100644 --- a/agents-api/agents_api/models/user/delete_user.py +++ b/agents-api/agents_api/models/user/delete_user.py @@ -49,12 +49,12 @@ def delete_user(*, developer_id: UUID, user_id: UUID) -> tuple[list[str], dict]: Constructs and returns a datalog query for deleting an user and its default settings from the database. Parameters: - - developer_id (UUID): The UUID of the developer owning the user. - - user_id (UUID): The UUID of the user 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 user. + user_id (UUID): The UUID of the user to be deleted. + client (CozoClient, optional): An instance of the CozoClient to execute the query. Returns: - - ResourceDeletedResponse: The response indicating the deletion of the user. + ResourceDeletedResponse: The response indicating the deletion of the user. """ queries = [ diff --git a/agents-api/agents_api/models/user/get_user.py b/agents-api/agents_api/models/user/get_user.py index cc7c6f970..2b4f59c83 100644 --- a/agents-api/agents_api/models/user/get_user.py +++ b/agents-api/agents_api/models/user/get_user.py @@ -67,7 +67,10 @@ def get_user( created_at, updated_at, metadata, - }""" + } + + :limit 1 + """ queries = [ verify_developer_id_query(developer_id), diff --git a/agents-api/agents_api/models/user/list_users.py b/agents-api/agents_api/models/user/list_users.py index 57dc9b8c8..2a810b8e0 100644 --- a/agents-api/agents_api/models/user/list_users.py +++ b/agents-api/agents_api/models/user/list_users.py @@ -43,13 +43,13 @@ def list_users( Queries the 'cozodb' database to list users associated with a specific developer. Parameters: - - developer_id (UUID): The unique identifier of the developer. - - limit (int): The maximum number of users to return. Defaults to 100. - - offset (int): The number of users to skip before starting to collect the result set. Defaults to 0. - - metadata_filter (dict[str, Any]): A dictionary representing filters to apply on user metadata. + developer_id (UUID): The unique identifier of the developer. + limit (int): The maximum number of users to return. Defaults to 100. + offset (int): The number of users to skip before starting to collect the result set. Defaults to 0. + metadata_filter (dict[str, Any]): A dictionary representing filters to apply on user metadata. Returns: - - pd.DataFrame: A DataFrame containing the queried user data. + pd.DataFrame: A DataFrame containing the queried user data. """ # Construct a filter string for the metadata based on the provided dictionary. metadata_filter_str = ", ".join( diff --git a/agents-api/agents_api/models/user/patch_user.py b/agents-api/agents_api/models/user/patch_user.py index faf38298c..4498c6ded 100644 --- a/agents-api/agents_api/models/user/patch_user.py +++ b/agents-api/agents_api/models/user/patch_user.py @@ -49,12 +49,12 @@ def patch_user( Generates a datalog query for updating a user's information. Parameters: - - developer_id (UUID): The UUID of the developer. - - user_id (UUID): The UUID of the user to be updated. - - **update_data: Arbitrary keyword arguments representing the data to be updated. + developer_id (UUID): The UUID of the developer. + user_id (UUID): The UUID of the user to be updated. + **update_data: Arbitrary keyword arguments representing the data to be updated. Returns: - - tuple[str, dict]: A pandas DataFrame containing the results of the query execution. + tuple[str, dict]: A pandas DataFrame containing the results of the query execution. """ update_data = data.model_dump(exclude_unset=True) diff --git a/agents-api/agents_api/models/user/update_user.py b/agents-api/agents_api/models/user/update_user.py index 9a13d9369..fd8e7e2c8 100644 --- a/agents-api/agents_api/models/user/update_user.py +++ b/agents-api/agents_api/models/user/update_user.py @@ -39,7 +39,8 @@ def update_user( *, developer_id: UUID, user_id: UUID, data: UpdateUserRequest ) -> tuple[list[str], dict]: - """Updates user information in the 'cozodb' database. + """ + Updates user information in the 'cozodb' database. Parameters: developer_id (UUID): The developer's unique identifier. diff --git a/agents-api/agents_api/models/utils.py b/agents-api/agents_api/models/utils.py index c97d92451..c163642c0 100644 --- a/agents-api/agents_api/models/utils.py +++ b/agents-api/agents_api/models/utils.py @@ -127,6 +127,8 @@ def verify_developer_id_query(developer_id: UUID | str) -> str: matched[num], exists = num > 0, assert(exists, "Developer does not exist") + + :limit 1 """ @@ -162,6 +164,8 @@ def verify_developer_owns_resource_query( found[num], exists = num > 0, assert(exists, "Developer does not own resource {resource} with {resource_id_key} {resource_id_value}") + + :limit 1 """ rule = rule_head + rule_body + assertion diff --git a/agents-api/agents_api/routers/jobs/routers.py b/agents-api/agents_api/routers/jobs/routers.py index 8bff6c7cc..dff4bed7b 100644 --- a/agents-api/agents_api/routers/jobs/routers.py +++ b/agents-api/agents_api/routers/jobs/routers.py @@ -4,8 +4,8 @@ from fastapi import APIRouter from temporalio.client import WorkflowExecutionStatus -from agents_api.autogen.openapi_model import JobStatus -from agents_api.clients.temporal import get_client +from ...autogen.openapi_model import JobStatus +from ...clients.temporal import get_client router: APIRouter = APIRouter() diff --git a/agents-api/agents_api/routers/tasks/create_or_update_task.py b/agents-api/agents_api/routers/tasks/create_or_update_task.py index e5245670a..50dbf19d9 100644 --- a/agents-api/agents_api/routers/tasks/create_or_update_task.py +++ b/agents-api/agents_api/routers/tasks/create_or_update_task.py @@ -6,15 +6,14 @@ from jsonschema.exceptions import SchemaError, ValidationError from starlette.status import HTTP_201_CREATED -from agents_api.autogen.openapi_model import ( +from ...autogen.openapi_model import ( CreateOrUpdateTaskRequest, ResourceUpdatedResponse, ) -from agents_api.dependencies.developer_id import get_developer_id -from agents_api.models.task.create_or_update_task import ( +from ...dependencies.developer_id import get_developer_id +from ...models.task.create_or_update_task import ( create_or_update_task as create_or_update_task_query, ) - from .router import router diff --git a/agents-api/agents_api/routers/tasks/get_execution_details.py b/agents-api/agents_api/routers/tasks/get_execution_details.py index 8da7b98c8..95bccbc07 100644 --- a/agents-api/agents_api/routers/tasks/get_execution_details.py +++ b/agents-api/agents_api/routers/tasks/get_execution_details.py @@ -1,12 +1,11 @@ from uuid import UUID -from agents_api.autogen.openapi_model import ( +from ...autogen.openapi_model import ( Execution, ) -from agents_api.models.execution.get_execution import ( +from ...models.execution.get_execution import ( get_execution as get_execution_query, ) - from .router import router diff --git a/agents-api/agents_api/routers/tasks/list_execution_transitions.py b/agents-api/agents_api/routers/tasks/list_execution_transitions.py index fd7be992a..9ce169509 100644 --- a/agents-api/agents_api/routers/tasks/list_execution_transitions.py +++ b/agents-api/agents_api/routers/tasks/list_execution_transitions.py @@ -1,14 +1,13 @@ from typing import Literal from uuid import UUID -from agents_api.autogen.openapi_model import ( +from ...autogen.openapi_model import ( ListResponse, Transition, ) -from agents_api.models.execution.list_execution_transitions import ( +from ...models.execution.list_execution_transitions import ( list_execution_transitions as list_execution_transitions_query, ) - from .router import router diff --git a/agents-api/agents_api/routers/tasks/list_task_executions.py b/agents-api/agents_api/routers/tasks/list_task_executions.py index f2961e54a..72cbd9b40 100644 --- a/agents-api/agents_api/routers/tasks/list_task_executions.py +++ b/agents-api/agents_api/routers/tasks/list_task_executions.py @@ -3,15 +3,14 @@ from fastapi import Depends -from agents_api.autogen.openapi_model import ( +from ...autogen.openapi_model import ( Execution, ListResponse, ) -from agents_api.dependencies.developer_id import get_developer_id -from agents_api.models.execution.list_executions import ( +from ...dependencies.developer_id import get_developer_id +from ...models.execution.list_executions import ( list_executions as list_task_executions_query, ) - from .router import router diff --git a/agents-api/agents_api/routers/tasks/list_tasks.py b/agents-api/agents_api/routers/tasks/list_tasks.py index 43a7e9158..a53983006 100644 --- a/agents-api/agents_api/routers/tasks/list_tasks.py +++ b/agents-api/agents_api/routers/tasks/list_tasks.py @@ -3,13 +3,12 @@ from fastapi import Depends -from agents_api.autogen.openapi_model import ( +from ...autogen.openapi_model import ( ListResponse, Task, ) -from agents_api.dependencies.developer_id import get_developer_id -from agents_api.models.task.list_tasks import list_tasks as list_tasks_query - +from ...dependencies.developer_id import get_developer_id +from ...models.task.list_tasks import list_tasks as list_tasks_query from .router import router diff --git a/agents-api/agents_api/routers/tasks/patch_execution.py b/agents-api/agents_api/routers/tasks/patch_execution.py index 0eb159c83..3cc45ee37 100644 --- a/agents-api/agents_api/routers/tasks/patch_execution.py +++ b/agents-api/agents_api/routers/tasks/patch_execution.py @@ -3,15 +3,14 @@ from fastapi import Depends -from agents_api.autogen.openapi_model import ( +from ...autogen.openapi_model import ( ResourceUpdatedResponse, UpdateExecutionRequest, ) -from agents_api.dependencies.developer_id import get_developer_id -from agents_api.models.execution.update_execution import ( +from ...dependencies.developer_id import get_developer_id +from ...models.execution.update_execution import ( update_execution as update_execution_query, ) - from .router import router diff --git a/agents-api/agents_api/routers/tasks/update_execution.py b/agents-api/agents_api/routers/tasks/update_execution.py index 968b6bdfb..d887c455d 100644 --- a/agents-api/agents_api/routers/tasks/update_execution.py +++ b/agents-api/agents_api/routers/tasks/update_execution.py @@ -4,19 +4,18 @@ from fastapi import Depends, HTTPException -from agents_api.autogen.openapi_model import ( +from ...autogen.openapi_model import ( ResumeExecutionRequest, StopExecutionRequest, ) -from agents_api.clients.temporal import get_client -from agents_api.dependencies.developer_id import get_developer_id -from agents_api.models.execution.get_paused_execution_token import ( +from ...clients.temporal import get_client +from ...dependencies.developer_id import get_developer_id +from ...models.execution.get_paused_execution_token import ( get_paused_execution_token, ) -from agents_api.models.execution.get_temporal_workflow_data import ( +from ...models.execution.get_temporal_workflow_data import ( get_temporal_workflow_data, ) - from .router import router diff --git a/agents-api/tests/fixtures.py b/agents-api/tests/fixtures.py index d5b032311..9ae198c78 100644 --- a/agents-api/tests/fixtures.py +++ b/agents-api/tests/fixtures.py @@ -36,8 +36,7 @@ from agents_api.models.user.create_user import create_user from agents_api.models.user.delete_user import delete_user from agents_api.web import app - -from .utils import patch_embed_acompletion as patch_embed_acompletion_ctx +from tests.utils import patch_embed_acompletion as patch_embed_acompletion_ctx EMBEDDING_SIZE: int = 1024 diff --git a/agents-api/tests/test_activities.py b/agents-api/tests/test_activities.py index 98dfc97b5..a2f15d179 100644 --- a/agents-api/tests/test_activities.py +++ b/agents-api/tests/test_activities.py @@ -7,13 +7,8 @@ from agents_api.clients import temporal from agents_api.env import temporal_task_queue from agents_api.workflows.demo import DemoWorkflow - -from .fixtures import ( - cozo_client, - test_developer_id, - test_doc, -) -from .utils import patch_testing_temporal +from tests.fixtures import cozo_client, test_developer_id, test_doc +from tests.utils import patch_testing_temporal @test("activity: call direct embed_docs") diff --git a/agents-api/tests/test_docs_routes.py b/agents-api/tests/test_docs_routes.py index d4b677d05..d61bfbcb7 100644 --- a/agents-api/tests/test_docs_routes.py +++ b/agents-api/tests/test_docs_routes.py @@ -1,6 +1,6 @@ from ward import test -from .fixtures import ( +from tests.fixtures import ( make_request, patch_embed_acompletion, test_agent, @@ -8,7 +8,7 @@ test_user, test_user_doc, ) -from .utils import patch_testing_temporal +from tests.utils import patch_testing_temporal @test("route: create user doc") diff --git a/agents-api/tests/test_execution_queries.py b/agents-api/tests/test_execution_queries.py index 29ddcbd86..42904776d 100644 --- a/agents-api/tests/test_execution_queries.py +++ b/agents-api/tests/test_execution_queries.py @@ -15,8 +15,7 @@ from agents_api.models.execution.get_execution import get_execution from agents_api.models.execution.list_executions import list_executions from agents_api.models.execution.lookup_temporal_data import lookup_temporal_data - -from .fixtures import ( +from tests.fixtures import ( cozo_client, test_developer_id, test_execution, diff --git a/agents-api/tests/test_execution_workflow.py b/agents-api/tests/test_execution_workflow.py index b6394f1bc..3df23e5cd 100644 --- a/agents-api/tests/test_execution_workflow.py +++ b/agents-api/tests/test_execution_workflow.py @@ -14,9 +14,8 @@ ) from agents_api.models.task.create_task import create_task from agents_api.routers.tasks.create_task_execution import start_execution - -from .fixtures import cozo_client, test_agent, test_developer_id -from .utils import patch_integration_service, patch_testing_temporal +from tests.fixtures import cozo_client, test_agent, test_developer_id +from tests.utils import patch_integration_service, patch_testing_temporal EMBEDDING_SIZE: int = 1024 diff --git a/agents-api/tests/test_task_queries.py b/agents-api/tests/test_task_queries.py index 2399416db..e61489df8 100644 --- a/agents-api/tests/test_task_queries.py +++ b/agents-api/tests/test_task_queries.py @@ -15,8 +15,7 @@ from agents_api.models.task.get_task import get_task from agents_api.models.task.list_tasks import list_tasks from agents_api.models.task.update_task import update_task - -from .fixtures import cozo_client, test_agent, test_developer_id, test_task +from tests.fixtures import cozo_client, test_agent, test_developer_id, test_task @test("model: create task") diff --git a/agents-api/tests/test_task_routes.py b/agents-api/tests/test_task_routes.py index 4ab708560..5d3c2f998 100644 --- a/agents-api/tests/test_task_routes.py +++ b/agents-api/tests/test_task_routes.py @@ -4,14 +4,14 @@ from ward import test -from .fixtures import ( +from tests.fixtures import ( client, make_request, test_agent, test_execution, test_task, ) -from .utils import patch_testing_temporal +from tests.utils import patch_testing_temporal @test("route: unauthorized should fail") diff --git a/agents-api/tests/test_workflow_routes.py b/agents-api/tests/test_workflow_routes.py index 34aa0101c..d1538535d 100644 --- a/agents-api/tests/test_workflow_routes.py +++ b/agents-api/tests/test_workflow_routes.py @@ -4,8 +4,8 @@ from ward import test -from .fixtures import cozo_client, test_agent, test_developer_id -from .utils import patch_http_client_with_temporal +from tests.fixtures import cozo_client, test_agent, test_developer_id +from tests.utils import patch_http_client_with_temporal @test("workflow route: evaluate step single") diff --git a/agents-api/tests/utils.py b/agents-api/tests/utils.py index 3f0f1f94a..dc1007d13 100644 --- a/agents-api/tests/utils.py +++ b/agents-api/tests/utils.py @@ -10,6 +10,7 @@ from agents_api.worker.codec import pydantic_data_converter from agents_api.worker.worker import create_worker +# Replicated here to prevent circular import EMBEDDING_SIZE: int = 1024 diff --git a/typespec/docs/models.tsp b/typespec/docs/models.tsp index ee283fe02..f01b379c6 100644 --- a/typespec/docs/models.tsp +++ b/typespec/docs/models.tsp @@ -23,6 +23,10 @@ model Doc { /** Contents of the document */ content: string | string[]; + + /** Embeddings for the document */ + @visibility("read") + embeddings?: float32[] | float32[][]; } /** Payload for creating a doc */ diff --git a/typespec/tools/models.tsp b/typespec/tools/models.tsp index be0fb4f68..8a8cead44 100644 --- a/typespec/tools/models.tsp +++ b/typespec/tools/models.tsp @@ -74,9 +74,6 @@ model SystemDef { // TODO: We should use this model for all tools, not just functions and discriminate on the type model Tool { - /** Whether this tool is a `function`, `api_call`, `system` etc. (Only `function` tool supported right now) */ - type: ToolType = ToolType.function; - /** Name of the tool (must be unique for this agent and a valid python identifier string )*/ name: validPythonIdentifier; @@ -95,24 +92,13 @@ model FunctionCallOption { name: string; } -@discriminator("type") model NamedToolChoice { - /** Whether this tool is a `function`, `api_call`, `system` etc. (Only `function` tool supported right now) */ - type: ToolType; - function?: FunctionCallOption; integration?: never; // TODO: Implement system?: never; // TODO: Implement api_call?: never; // TODO: Implement } -model NamedFunctionChoice extends NamedToolChoice { - type: ToolType.function; - - /** The function to call */ - function: FunctionCallOption; -} - model ToolResponse { @key id: uuid;