From 2f836d67fdfe301b8d168ebdd5f4b9d75f379a35 Mon Sep 17 00:00:00 2001 From: Ahmad Haidar Date: Tue, 24 Dec 2024 20:13:26 +0300 Subject: [PATCH 1/4] chore(agents-api): remove cozo related stuff --- .../agents_api/activities/embed_docs.py | 73 --- .../activities/task_steps/__init__.py | 2 +- .../activities/task_steps/cozo_query_step.py | 28 - .../activities/task_steps/pg_query_step.py | 37 ++ agents-api/agents_api/activities/utils.py | 44 +- agents-api/agents_api/clients/__init__.py | 2 +- .../agents_api/common/utils/__init__.py | 2 +- agents-api/agents_api/common/utils/cozo.py | 26 - agents-api/agents_api/models/__init__.py | 20 - .../agents_api/models/agent/__init__.py | 22 - .../agents_api/models/agent/create_agent.py | 148 ----- .../models/agent/create_or_update_agent.py | 186 ------ .../agents_api/models/agent/delete_agent.py | 134 ---- .../agents_api/models/agent/get_agent.py | 117 ---- .../agents_api/models/agent/list_agents.py | 122 ---- .../agents_api/models/agent/patch_agent.py | 132 ---- .../agents_api/models/agent/update_agent.py | 149 ----- agents-api/agents_api/models/docs/__init__.py | 25 - .../agents_api/models/docs/create_doc.py | 141 ----- .../agents_api/models/docs/delete_doc.py | 102 ---- .../agents_api/models/docs/embed_snippets.py | 102 ---- agents-api/agents_api/models/docs/get_doc.py | 103 ---- .../agents_api/models/docs/list_docs.py | 141 ----- agents-api/agents_api/models/docs/mmr.py | 109 ---- .../models/docs/search_docs_by_embedding.py | 369 ----------- .../models/docs/search_docs_by_text.py | 206 ------- .../models/docs/search_docs_hybrid.py | 138 ----- .../agents_api/models/entry/__init__.py | 19 - .../agents_api/models/entry/create_entries.py | 128 ---- .../agents_api/models/entry/delete_entries.py | 153 ----- .../agents_api/models/entry/get_history.py | 150 ----- .../agents_api/models/entry/list_entries.py | 112 ---- .../agents_api/models/execution/__init__.py | 15 - .../agents_api/models/execution/constants.py | 5 - .../models/execution/count_executions.py | 61 -- .../models/execution/create_execution.py | 98 --- .../execution/create_execution_transition.py | 259 -------- .../execution/create_temporal_lookup.py | 72 --- .../models/execution/get_execution.py | 78 --- .../execution/get_execution_transition.py | 80 --- .../execution/get_paused_execution_token.py | 77 --- .../execution/get_temporal_workflow_data.py | 57 -- .../execution/list_execution_transitions.py | 69 --- .../models/execution/list_executions.py | 95 --- .../models/execution/lookup_temporal_data.py | 66 -- .../execution/prepare_execution_input.py | 223 ------- .../models/execution/update_execution.py | 130 ---- .../agents_api/models/files/__init__.py | 3 - .../agents_api/models/files/create_file.py | 122 ---- .../agents_api/models/files/delete_file.py | 97 --- .../agents_api/models/files/get_file.py | 116 ---- .../agents_api/models/session/__init__.py | 22 - .../models/session/count_sessions.py | 64 -- .../session/create_or_update_session.py | 158 ----- .../models/session/create_session.py | 154 ----- .../models/session/delete_session.py | 125 ---- .../agents_api/models/session/get_session.py | 116 ---- .../models/session/list_sessions.py | 131 ---- .../models/session/patch_session.py | 127 ---- .../models/session/prepare_session_data.py | 235 ------- .../models/session/update_session.py | 127 ---- agents-api/agents_api/models/task/__init__.py | 9 - .../models/task/create_or_update_task.py | 109 ---- .../agents_api/models/task/create_task.py | 118 ---- .../agents_api/models/task/delete_task.py | 91 --- agents-api/agents_api/models/task/get_task.py | 120 ---- .../agents_api/models/task/list_tasks.py | 130 ---- .../agents_api/models/task/patch_task.py | 133 ---- .../agents_api/models/task/update_task.py | 129 ---- agents-api/agents_api/models/user/__init__.py | 18 - .../models/user/create_or_update_user.py | 125 ---- .../agents_api/models/user/create_user.py | 116 ---- .../agents_api/models/user/delete_user.py | 116 ---- agents-api/agents_api/models/user/get_user.py | 107 ---- .../agents_api/models/user/list_users.py | 116 ---- .../agents_api/models/user/patch_user.py | 121 ---- .../agents_api/models/user/update_user.py | 118 ---- agents-api/agents_api/models/utils.py | 578 ------------------ agents-api/agents_api/queries/__init__.py | 21 + .../queries/developers/get_developer.py | 4 +- .../agents_api/queries/tools/create_tools.py | 5 +- .../agents_api/queries/tools/patch_tool.py | 3 +- agents-api/agents_api/worker/worker.py | 4 - agents-api/agents_api/workflows/embed_docs.py | 27 - 84 files changed, 90 insertions(+), 8452 deletions(-) delete mode 100644 agents-api/agents_api/activities/embed_docs.py delete mode 100644 agents-api/agents_api/activities/task_steps/cozo_query_step.py create mode 100644 agents-api/agents_api/activities/task_steps/pg_query_step.py delete mode 100644 agents-api/agents_api/common/utils/cozo.py delete mode 100644 agents-api/agents_api/models/__init__.py delete mode 100644 agents-api/agents_api/models/agent/__init__.py delete mode 100644 agents-api/agents_api/models/agent/create_agent.py delete mode 100644 agents-api/agents_api/models/agent/create_or_update_agent.py delete mode 100644 agents-api/agents_api/models/agent/delete_agent.py delete mode 100644 agents-api/agents_api/models/agent/get_agent.py delete mode 100644 agents-api/agents_api/models/agent/list_agents.py delete mode 100644 agents-api/agents_api/models/agent/patch_agent.py delete mode 100644 agents-api/agents_api/models/agent/update_agent.py delete mode 100644 agents-api/agents_api/models/docs/__init__.py delete mode 100644 agents-api/agents_api/models/docs/create_doc.py delete mode 100644 agents-api/agents_api/models/docs/delete_doc.py delete mode 100644 agents-api/agents_api/models/docs/embed_snippets.py delete mode 100644 agents-api/agents_api/models/docs/get_doc.py delete mode 100644 agents-api/agents_api/models/docs/list_docs.py delete mode 100644 agents-api/agents_api/models/docs/mmr.py delete mode 100644 agents-api/agents_api/models/docs/search_docs_by_embedding.py delete mode 100644 agents-api/agents_api/models/docs/search_docs_by_text.py delete mode 100644 agents-api/agents_api/models/docs/search_docs_hybrid.py delete mode 100644 agents-api/agents_api/models/entry/__init__.py delete mode 100644 agents-api/agents_api/models/entry/create_entries.py delete mode 100644 agents-api/agents_api/models/entry/delete_entries.py delete mode 100644 agents-api/agents_api/models/entry/get_history.py delete mode 100644 agents-api/agents_api/models/entry/list_entries.py delete mode 100644 agents-api/agents_api/models/execution/__init__.py delete mode 100644 agents-api/agents_api/models/execution/constants.py delete mode 100644 agents-api/agents_api/models/execution/count_executions.py delete mode 100644 agents-api/agents_api/models/execution/create_execution.py delete mode 100644 agents-api/agents_api/models/execution/create_execution_transition.py delete mode 100644 agents-api/agents_api/models/execution/create_temporal_lookup.py delete mode 100644 agents-api/agents_api/models/execution/get_execution.py delete mode 100644 agents-api/agents_api/models/execution/get_execution_transition.py delete mode 100644 agents-api/agents_api/models/execution/get_paused_execution_token.py delete mode 100644 agents-api/agents_api/models/execution/get_temporal_workflow_data.py delete mode 100644 agents-api/agents_api/models/execution/list_execution_transitions.py delete mode 100644 agents-api/agents_api/models/execution/list_executions.py delete mode 100644 agents-api/agents_api/models/execution/lookup_temporal_data.py delete mode 100644 agents-api/agents_api/models/execution/prepare_execution_input.py delete mode 100644 agents-api/agents_api/models/execution/update_execution.py delete mode 100644 agents-api/agents_api/models/files/__init__.py delete mode 100644 agents-api/agents_api/models/files/create_file.py delete mode 100644 agents-api/agents_api/models/files/delete_file.py delete mode 100644 agents-api/agents_api/models/files/get_file.py delete mode 100644 agents-api/agents_api/models/session/__init__.py delete mode 100644 agents-api/agents_api/models/session/count_sessions.py delete mode 100644 agents-api/agents_api/models/session/create_or_update_session.py delete mode 100644 agents-api/agents_api/models/session/create_session.py delete mode 100644 agents-api/agents_api/models/session/delete_session.py delete mode 100644 agents-api/agents_api/models/session/get_session.py delete mode 100644 agents-api/agents_api/models/session/list_sessions.py delete mode 100644 agents-api/agents_api/models/session/patch_session.py delete mode 100644 agents-api/agents_api/models/session/prepare_session_data.py delete mode 100644 agents-api/agents_api/models/session/update_session.py delete mode 100644 agents-api/agents_api/models/task/__init__.py delete mode 100644 agents-api/agents_api/models/task/create_or_update_task.py delete mode 100644 agents-api/agents_api/models/task/create_task.py delete mode 100644 agents-api/agents_api/models/task/delete_task.py delete mode 100644 agents-api/agents_api/models/task/get_task.py delete mode 100644 agents-api/agents_api/models/task/list_tasks.py delete mode 100644 agents-api/agents_api/models/task/patch_task.py delete mode 100644 agents-api/agents_api/models/task/update_task.py delete mode 100644 agents-api/agents_api/models/user/__init__.py delete mode 100644 agents-api/agents_api/models/user/create_or_update_user.py delete mode 100644 agents-api/agents_api/models/user/create_user.py delete mode 100644 agents-api/agents_api/models/user/delete_user.py delete mode 100644 agents-api/agents_api/models/user/get_user.py delete mode 100644 agents-api/agents_api/models/user/list_users.py delete mode 100644 agents-api/agents_api/models/user/patch_user.py delete mode 100644 agents-api/agents_api/models/user/update_user.py delete mode 100644 agents-api/agents_api/models/utils.py create mode 100644 agents-api/agents_api/queries/__init__.py delete mode 100644 agents-api/agents_api/workflows/embed_docs.py diff --git a/agents-api/agents_api/activities/embed_docs.py b/agents-api/agents_api/activities/embed_docs.py deleted file mode 100644 index a9a7cae44..000000000 --- a/agents-api/agents_api/activities/embed_docs.py +++ /dev/null @@ -1,73 +0,0 @@ -import asyncio -import operator -from functools import reduce -from itertools import batched - -from beartype import beartype -from temporalio import activity - -from ..clients import cozo, litellm -from ..env import testing -from ..models.docs.embed_snippets import embed_snippets as embed_snippets_query -from .types import EmbedDocsPayload - - -@beartype -async def embed_docs( - payload: EmbedDocsPayload, cozo_client=None, max_batch_size: int = 100 -) -> None: - # Create batches of both indices and snippets together - indexed_snippets = list(enumerate(payload.content)) - # Batch snippets into groups of max_batch_size for parallel processing - batched_indexed_snippets = list(batched(indexed_snippets, max_batch_size)) - # Get embedding instruction and title from payload, defaulting to empty strings - embed_instruction: str = payload.embed_instruction or "" - title: str = payload.title or "" - - # Helper function to embed a batch of snippets - async def embed_batch(indexed_batch): - # Split indices and snippets for the batch - batch_indices, batch_snippets = zip(*indexed_batch) - embeddings = await litellm.aembedding( - inputs=[ - ((title + "\n\n" + snippet) if title else snippet).strip() - for snippet in batch_snippets - ], - embed_instruction=embed_instruction, - ) - return list(zip(batch_indices, embeddings)) - - # Gather embeddings with their corresponding indices - indexed_embeddings = reduce( - operator.add, - await asyncio.gather( - *[embed_batch(batch) for batch in batched_indexed_snippets] - ), - ) - - # Split indices and embeddings after all batches are processed - indices, embeddings = zip(*indexed_embeddings) - - # Convert to lists since embed_snippets_query expects list types - indices = list(indices) - embeddings = list(embeddings) - - embed_snippets_query( - developer_id=payload.developer_id, - doc_id=payload.doc_id, - snippet_indices=indices, - embeddings=embeddings, - client=cozo_client or cozo.get_cozo_client(), - ) - - -async def mock_embed_docs( - payload: EmbedDocsPayload, cozo_client=None, max_batch_size=100 -) -> None: - # Does nothing - return None - - -embed_docs = activity.defn(name="embed_docs")( - embed_docs if not testing else mock_embed_docs -) diff --git a/agents-api/agents_api/activities/task_steps/__init__.py b/agents-api/agents_api/activities/task_steps/__init__.py index 573884629..5d02db858 100644 --- a/agents-api/agents_api/activities/task_steps/__init__.py +++ b/agents-api/agents_api/activities/task_steps/__init__.py @@ -1,7 +1,7 @@ # ruff: noqa: F401, F403, F405 from .base_evaluate import base_evaluate -from .cozo_query_step import cozo_query_step +from .pg_query_step import pg_query_step from .evaluate_step import evaluate_step from .for_each_step import for_each_step from .get_value_step import get_value_step diff --git a/agents-api/agents_api/activities/task_steps/cozo_query_step.py b/agents-api/agents_api/activities/task_steps/cozo_query_step.py deleted file mode 100644 index 8d28d83c9..000000000 --- a/agents-api/agents_api/activities/task_steps/cozo_query_step.py +++ /dev/null @@ -1,28 +0,0 @@ -from typing import Any - -from beartype import beartype -from temporalio import activity - -from ... import models -from ...env import testing - - -@beartype -async def cozo_query_step( - query_name: str, - values: dict[str, Any], -) -> Any: - (module_name, name) = query_name.split(".") - - module = getattr(models, module_name) - query = getattr(module, name) - return query(**values) - - -# Note: This is here just for clarity. We could have just imported cozo_query_step directly -# They do the same thing, so we dont need to mock the cozo_query_step function -mock_cozo_query_step = cozo_query_step - -cozo_query_step = activity.defn(name="cozo_query_step")( - cozo_query_step if not testing else mock_cozo_query_step -) diff --git a/agents-api/agents_api/activities/task_steps/pg_query_step.py b/agents-api/agents_api/activities/task_steps/pg_query_step.py new file mode 100644 index 000000000..bfddc716f --- /dev/null +++ b/agents-api/agents_api/activities/task_steps/pg_query_step.py @@ -0,0 +1,37 @@ +from typing import Any + +from async_lru import alru_cache +from beartype import beartype +from temporalio import activity + +from ... import queries +from ...env import testing, db_dsn + +from ...clients.pg import create_db_pool + +@alru_cache(maxsize=1) +async def get_db_pool(dsn: str): + return await create_db_pool(dsn=dsn) + +@beartype +async def pg_query_step( + query_name: str, + values: dict[str, Any], + dsn: str = db_dsn, +) -> Any: + pool = await get_db_pool(dsn=dsn) + + (module_name, name) = query_name.split(".") + + module = getattr(queries, module_name) + query = getattr(module, name) + return await query(**values, connection_pool=pool) + + +# Note: This is here just for clarity. We could have just imported pg_query_step directly +# They do the same thing, so we dont need to mock the pg_query_step function +mock_pg_query_step = pg_query_step + +pg_query_step = activity.defn(name="pg_query_step")( + pg_query_step if not testing else mock_pg_query_step +) diff --git a/agents-api/agents_api/activities/utils.py b/agents-api/agents_api/activities/utils.py index d9ad1840c..9b97f5f71 100644 --- a/agents-api/agents_api/activities/utils.py +++ b/agents-api/agents_api/activities/utils.py @@ -296,28 +296,28 @@ def get_handler(system: SystemDef) -> Callable: The base handler function. """ - from ..models.agent.create_agent import create_agent as create_agent_query - from ..models.agent.delete_agent import delete_agent as delete_agent_query - from ..models.agent.get_agent import get_agent as get_agent_query - from ..models.agent.list_agents import list_agents as list_agents_query - from ..models.agent.update_agent import update_agent as update_agent_query - from ..models.docs.delete_doc import delete_doc as delete_doc_query - from ..models.docs.list_docs import list_docs as list_docs_query - from ..models.session.create_session import create_session as create_session_query - from ..models.session.delete_session import delete_session as delete_session_query - from ..models.session.get_session import get_session as get_session_query - from ..models.session.list_sessions import list_sessions as list_sessions_query - from ..models.session.update_session import update_session as update_session_query - from ..models.task.create_task import create_task as create_task_query - from ..models.task.delete_task import delete_task as delete_task_query - from ..models.task.get_task import get_task as get_task_query - from ..models.task.list_tasks import list_tasks as list_tasks_query - from ..models.task.update_task import update_task as update_task_query - from ..models.user.create_user import create_user as create_user_query - from ..models.user.delete_user import delete_user as delete_user_query - from ..models.user.get_user import get_user as get_user_query - from ..models.user.list_users import list_users as list_users_query - from ..models.user.update_user import update_user as update_user_query + from ..queries.agents.create_agent import create_agent as create_agent_query + from ..queries.agents.delete_agent import delete_agent as delete_agent_query + from ..queries.agents.get_agent import get_agent as get_agent_query + from ..queries.agents.list_agents import list_agents as list_agents_query + from ..queries.agents.update_agent import update_agent as update_agent_query + from ..queries.docs.delete_doc import delete_doc as delete_doc_query + from ..queries.docs.list_docs import list_docs as list_docs_query + from ..queries.sessions.create_session import create_session as create_session_query + from ..queries.sessions.delete_session import delete_session as delete_session_query + from ..queries.sessions.get_session import get_session as get_session_query + from ..queries.sessions.list_sessions import list_sessions as list_sessions_query + from ..queries.sessions.update_session import update_session as update_session_query + from ..queries.tasks.create_task import create_task as create_task_query + from ..queries.tasks.delete_task import delete_task as delete_task_query + from ..queries.tasks.get_task import get_task as get_task_query + from ..queries.tasks.list_tasks import list_tasks as list_tasks_query + from ..queries.tasks.update_task import update_task as update_task_query + from ..queries.users.create_user import create_user as create_user_query + from ..queries.users.delete_user import delete_user as delete_user_query + from ..queries.users.get_user import get_user as get_user_query + from ..queries.users.list_users import list_users as list_users_query + from ..queries.users.update_user import update_user as update_user_query from ..routers.docs.create_doc import create_agent_doc, create_user_doc from ..routers.docs.search_docs import search_agent_docs, search_user_docs from ..routers.sessions.chat import chat diff --git a/agents-api/agents_api/clients/__init__.py b/agents-api/agents_api/clients/__init__.py index 43a17ab08..714cc5294 100644 --- a/agents-api/agents_api/clients/__init__.py +++ b/agents-api/agents_api/clients/__init__.py @@ -1,7 +1,7 @@ """ The `clients` module contains client classes and functions for interacting with various external services and APIs, abstracting the complexity of HTTP requests and API interactions to provide a simplified interface for the rest of the application. -- `cozo.py`: Handles communication with the Cozo service, facilitating operations such as retrieving product information. +- `pg.py`: Handles communication with the PostgreSQL service, facilitating operations such as retrieving product information. - `embed.py`: Manages requests to an Embedding Service for text embedding functionalities. - `openai.py`: Facilitates interaction with OpenAI's API for natural language processing tasks. - `temporal.py`: Provides functionality for connecting to Temporal workflows, enabling asynchronous task execution and management. diff --git a/agents-api/agents_api/common/utils/__init__.py b/agents-api/agents_api/common/utils/__init__.py index 891594c02..fbe7d490c 100644 --- a/agents-api/agents_api/common/utils/__init__.py +++ b/agents-api/agents_api/common/utils/__init__.py @@ -1,7 +1,7 @@ """ The `utils` module within the `agents-api` project offers a collection of utility functions designed to support various aspects of the application. This includes: -- `cozo.py`: Utilities for interacting with the Cozo API client, including data mutation processes. +- `pg.py`: Utilities for interacting with the PostgreSQL API client, including data mutation processes. - `datetime.py`: Functions for handling date and time operations, ensuring consistent use of time zones and formats across the application. - `json.py`: Custom JSON utilities, including a custom JSON encoder for handling specific object types like UUIDs, and a utility function for JSON serialization with support for default values for None objects. diff --git a/agents-api/agents_api/common/utils/cozo.py b/agents-api/agents_api/common/utils/cozo.py deleted file mode 100644 index f342ba617..000000000 --- a/agents-api/agents_api/common/utils/cozo.py +++ /dev/null @@ -1,26 +0,0 @@ -#!/usr/bin/env python3 - -"""This module provides utility functions for interacting with the Cozo API client, including data mutation processes.""" - -from types import SimpleNamespace -from uuid import UUID - -from beartype import beartype -from pycozo import Client - -# Define a mock client for testing purposes, simulating Cozo API client behavior. -_fake_client: SimpleNamespace = SimpleNamespace() -# Lambda function to process and mutate data dictionaries using the Cozo client's internal method. This is a workaround to access protected member functions for testing. -_fake_client._process_mutate_data_dict = lambda data: ( - Client._process_mutate_data_dict(_fake_client, data) -) - -# Public interface to process and mutate data using the Cozo client. It wraps the client's internal processing method for external use. -cozo_process_mutate_data = _fake_client._process_mutate_data = lambda data: ( - Client._process_mutate_data(_fake_client, data) -) - - -@beartype -def uuid_int_list_to_uuid(data: list[int]) -> UUID: - return UUID(bytes=b"".join([i.to_bytes(1, "big") for i in data])) diff --git a/agents-api/agents_api/models/__init__.py b/agents-api/agents_api/models/__init__.py deleted file mode 100644 index e59b5b01c..000000000 --- a/agents-api/agents_api/models/__init__.py +++ /dev/null @@ -1,20 +0,0 @@ -""" -The `models` module of the agents API is designed to encapsulate all data interactions with the CozoDB database. It provides a structured way to perform CRUD (Create, Read, Update, Delete) operations and other specific data manipulations across various entities such as agents, documents, entries, sessions, tools, and users. - -Each sub-module within this module corresponds to a specific entity and contains functions and classes that implement datalog queries for interacting with the database. These interactions include creating new records, updating existing ones, retrieving data for specific conditions, and deleting records. The operations are crucial for the functionality of the agents API, enabling it to manage and process data effectively for each entity. - -This module also integrates with the `common` module for exception handling and utility functions, ensuring robust error management and providing reusable components for data processing and query construction. -""" - -# ruff: noqa: F401, F403, F405 - -from . import agent as agent -from . import developer as developer -from . import docs as docs -from . import entry as entry -from . import execution as execution -from . import files as files -from . import session as session -from . import task as task -from . import tools as tools -from . import user as user diff --git a/agents-api/agents_api/models/agent/__init__.py b/agents-api/agents_api/models/agent/__init__.py deleted file mode 100644 index 2beaf8166..000000000 --- a/agents-api/agents_api/models/agent/__init__.py +++ /dev/null @@ -1,22 +0,0 @@ -""" -The `agent` module within the `agents-api` package provides a comprehensive suite of functionalities for managing agents in the CozoDB database. This includes: - -- Creating new agents and their associated tools. -- Updating existing agents and their settings. -- Retrieving details about specific agents or a list of agents. -- Deleting agents from the database. - -Additionally, the module supports operations related to agent tools, including creating, updating, and patching tools associated with agents. - -This module serves as the backbone for agent management within the CozoDB ecosystem, facilitating a wide range of operations necessary for the effective handling of agent data. -""" - -# ruff: noqa: F401, F403, F405 - -from .create_agent import create_agent -from .create_or_update_agent import create_or_update_agent -from .delete_agent import delete_agent -from .get_agent import get_agent -from .list_agents import list_agents -from .patch_agent import patch_agent -from .update_agent import update_agent diff --git a/agents-api/agents_api/models/agent/create_agent.py b/agents-api/agents_api/models/agent/create_agent.py deleted file mode 100644 index 1872a6f36..000000000 --- a/agents-api/agents_api/models/agent/create_agent.py +++ /dev/null @@ -1,148 +0,0 @@ -""" -This module contains the functionality for creating agents in the CozoDB database. -It includes functions to construct and execute datalog queries for inserting new agent records. -""" - -from typing import Any, TypeVar -from uuid import UUID - -from beartype import beartype -from fastapi import HTTPException -from pycozo.client import QueryException -from pydantic import ValidationError -from uuid_extensions import uuid7 - -from ...autogen.openapi_model import Agent, CreateAgentRequest -from ...common.utils.cozo import cozo_process_mutate_data -from ...metrics.counters import increase_counter -from ..utils import ( - cozo_query, - partialclass, - rewrap_exceptions, - verify_developer_id_query, - wrap_in_class, -) - -ModelT = TypeVar("ModelT", bound=Any) -T = TypeVar("T") - - -@rewrap_exceptions( - { - lambda e: isinstance(e, QueryException) - and "asserted to return some results, but returned none" - in str(e): lambda *_: HTTPException( - 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.", - ), - } -) -@wrap_in_class( - Agent, - one=True, - transform=lambda d: {"id": UUID(d.pop("agent_id")), **d}, - _kind="inserted", -) -@cozo_query -@increase_counter("create_agent") -@beartype -def create_agent( - *, - developer_id: UUID, - agent_id: UUID | None = None, - data: CreateAgentRequest, -) -> tuple[list[str], dict]: - """ - 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. - - Returns: - Agent: The newly created agent record. - """ - - agent_id = agent_id or uuid7() - - # Extract the agent data from the payload - data.metadata = data.metadata or {} - data.default_settings = data.default_settings or {} - - data.instructions = ( - data.instructions - if isinstance(data.instructions, list) - else [data.instructions] - ) - - agent_data = data.model_dump() - default_settings = agent_data.pop("default_settings") - - settings_cols, settings_vals = cozo_process_mutate_data( - { - **default_settings, - "agent_id": str(agent_id), - } - ) - - # Create default agent settings - # Construct a query to insert default settings for the new agent - default_settings_query = f""" - ?[{settings_cols}] <- $settings_vals - - :insert agent_default_settings {{ - {settings_cols} - }} - """ - # create the agent - # Construct a query to insert the new agent record into the agents table - agent_query = """ - ?[agent_id, developer_id, model, name, about, metadata, instructions, created_at, updated_at] <- [ - [$agent_id, $developer_id, $model, $name, $about, $metadata, $instructions, now(), now()] - ] - - :insert agents { - developer_id, - agent_id => - model, - name, - about, - metadata, - instructions, - created_at, - updated_at, - } - :returning - """ - - queries = [ - verify_developer_id_query(developer_id), - default_settings_query, - agent_query, - ] - - return ( - queries, - { - "settings_vals": settings_vals, - "agent_id": str(agent_id), - "developer_id": str(developer_id), - **agent_data, - }, - ) 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 deleted file mode 100644 index 9a1feb717..000000000 --- a/agents-api/agents_api/models/agent/create_or_update_agent.py +++ /dev/null @@ -1,186 +0,0 @@ -""" -This module contains the functionality for creating agents in the CozoDB database. -It includes functions to construct and execute datalog queries for inserting new agent records. -""" - -from typing import Any, TypeVar -from uuid import UUID - -from beartype import beartype -from fastapi import HTTPException -from pycozo.client import QueryException -from pydantic import ValidationError - -from ...autogen.openapi_model import Agent, CreateOrUpdateAgentRequest -from ...common.utils.cozo import cozo_process_mutate_data -from ...metrics.counters import increase_counter -from ..utils import ( - cozo_query, - partialclass, - rewrap_exceptions, - verify_developer_id_query, - wrap_in_class, -) - -ModelT = TypeVar("ModelT", bound=Any) -T = TypeVar("T") - - -@rewrap_exceptions( - { - 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( - Agent, one=True, transform=lambda d: {"id": UUID(d.pop("agent_id")), **d} -) -@cozo_query -@increase_counter("create_or_update_agent") -@beartype -def create_or_update_agent( - *, - developer_id: UUID, - agent_id: UUID, - data: CreateOrUpdateAgentRequest, -) -> tuple[list[str | None], dict]: - """ - 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. - - Returns: - Agent: The newly created agent record. - """ - - # Extract the agent data from the payload - data.metadata = data.metadata or {} - data.instructions = ( - data.instructions - if isinstance(data.instructions, list) - else [data.instructions] - ) - data.default_settings = data.default_settings or {} - - agent_data = data.model_dump() - default_settings = ( - data.default_settings.model_dump(exclude_none=True) - if data.default_settings - else {} - ) - - settings_cols, settings_vals = cozo_process_mutate_data( - { - **default_settings, - "agent_id": str(agent_id), - } - ) - - # TODO: remove this - ### # Create default agent settings - ### # Construct a query to insert default settings for the new agent - ### default_settings_query = f""" - ### %if {{ - ### len[count(agent_id)] := - ### *agent_default_settings{{agent_id}}, - ### agent_id = to_uuid($agent_id) - - ### ?[should_create] := len[count], count > 0 - ### }} - ### %then {{ - ### ?[{settings_cols}] <- $settings_vals - - ### :put agent_default_settings {{ - ### {settings_cols} - ### }} - ### }} - ### """ - - # FIXME: This create or update query will overwrite the settings - # Need to find a way to only run the insert query if the agent_default_settings - - # Create default agent settings - # Construct a query to insert default settings for the new agent - default_settings_query = f""" - ?[{settings_cols}] <- $settings_vals - - :put agent_default_settings {{ - {settings_cols} - }} - """ - - # create the agent - # Construct a query to insert the new agent record into the agents table - agent_query = """ - input[agent_id, developer_id, model, name, about, metadata, instructions, updated_at] <- [ - [$agent_id, $developer_id, $model, $name, $about, $metadata, $instructions, now()] - ] - - ?[agent_id, developer_id, model, name, about, metadata, instructions, created_at, updated_at] := - input[_agent_id, developer_id, model, name, about, metadata, instructions, updated_at], - *agents{ - agent_id, - developer_id, - created_at, - }, - agent_id = to_uuid(_agent_id), - - ?[agent_id, developer_id, model, name, about, metadata, instructions, created_at, updated_at] := - input[_agent_id, developer_id, model, name, about, metadata, instructions, updated_at], - not *agents{ - agent_id, - developer_id, - }, created_at = now(), - agent_id = to_uuid(_agent_id), - - :put agents { - developer_id, - agent_id => - model, - name, - about, - metadata, - instructions, - created_at, - updated_at, - } - :returning - """ - - queries = [ - verify_developer_id_query(developer_id), - default_settings_query, - agent_query, - ] - - return ( - queries, - { - "settings_vals": settings_vals, - "agent_id": str(agent_id), - "developer_id": str(developer_id), - **agent_data, - }, - ) diff --git a/agents-api/agents_api/models/agent/delete_agent.py b/agents-api/agents_api/models/agent/delete_agent.py deleted file mode 100644 index 60de66292..000000000 --- a/agents-api/agents_api/models/agent/delete_agent.py +++ /dev/null @@ -1,134 +0,0 @@ -""" -This module contains the implementation of the delete_agent_query function, which is responsible for deleting an agent and its related default settings from the CozoDB database. -""" - -from typing import Any, TypeVar -from uuid import UUID - -from beartype import beartype -from fastapi import HTTPException -from pycozo.client import QueryException -from pydantic import ValidationError - -from ...autogen.openapi_model import ResourceDeletedResponse -from ...common.utils.datetime import utcnow -from ..utils import ( - cozo_query, - partialclass, - rewrap_exceptions, - verify_developer_id_query, - verify_developer_owns_resource_query, - wrap_in_class, -) - -ModelT = TypeVar("ModelT", bound=Any) -T = TypeVar("T") - - -@rewrap_exceptions( - { - lambda e: isinstance(e, QueryException) - and "Developer does not own resource" - in e.resp["display"]: lambda *_: HTTPException( - 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.", - ), - } -) -@wrap_in_class( - ResourceDeletedResponse, - one=True, - transform=lambda d: { - "id": UUID(d.pop("agent_id")), - "deleted_at": utcnow(), - "jobs": [], - }, - _kind="deleted", -) -@cozo_query -@beartype -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. - - Returns: - ResourceDeletedResponse: The response indicating the deletion of the agent. - """ - - queries = [ - verify_developer_id_query(developer_id), - verify_developer_owns_resource_query(developer_id, "agents", agent_id=agent_id), - """ - # Delete docs - ?[owner_id, owner_type, doc_id] := - *docs{ - owner_type, - owner_id, - doc_id, - }, - owner_id = to_uuid($agent_id), - owner_type = "agent" - - :delete docs { - owner_type, - owner_id, - doc_id - } - :returning - """, - """ - # Delete tools - ?[agent_id, tool_id] := - *tools{ - agent_id, - tool_id, - }, agent_id = to_uuid($agent_id) - - :delete tools { - agent_id, - tool_id - } - :returning - """, - """ - # Delete default agent settings - ?[agent_id] <- [[$agent_id]] - - :delete agent_default_settings { - agent_id - } - :returning - """, - """ - # Delete the agent - ?[agent_id, developer_id] <- [[$agent_id, $developer_id]] - - :delete agents { - developer_id, - agent_id - } - :returning - """, - ] - - return (queries, {"agent_id": str(agent_id), "developer_id": str(developer_id)}) diff --git a/agents-api/agents_api/models/agent/get_agent.py b/agents-api/agents_api/models/agent/get_agent.py deleted file mode 100644 index 008e39454..000000000 --- a/agents-api/agents_api/models/agent/get_agent.py +++ /dev/null @@ -1,117 +0,0 @@ -from typing import Any, TypeVar -from uuid import UUID - -from beartype import beartype -from fastapi import HTTPException -from pycozo.client import QueryException -from pydantic import ValidationError - -from ...autogen.openapi_model import Agent -from ..utils import ( - cozo_query, - partialclass, - rewrap_exceptions, - verify_developer_id_query, - verify_developer_owns_resource_query, - wrap_in_class, -) - -ModelT = TypeVar("ModelT", bound=Any) -T = TypeVar("T") - - -@rewrap_exceptions( - { - lambda e: isinstance(e, QueryException) - and "Developer not found" in str(e): lambda *_: HTTPException( - detail="Developer does not exist", status_code=403 - ), - lambda e: isinstance(e, QueryException) - and "Developer does not own resource" - in e.resp["display"]: lambda *_: HTTPException( - detail="Developer does not own resource", status_code=404 - ), - QueryException: partialclass(HTTPException, status_code=400), - ValidationError: partialclass(HTTPException, status_code=400), - TypeError: partialclass(HTTPException, status_code=400), - } -) -@wrap_in_class(Agent, one=True) -@cozo_query -@beartype -def get_agent(*, developer_id: UUID, agent_id: UUID) -> tuple[list[str], dict]: - """ - Fetches agent details and default settings from the database. - - 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. - - Returns: - 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. - # It joins the 'agents' and 'agent_default_settings' relations to fetch comprehensive details. - get_query = """ - input[agent_id] <- [[to_uuid($agent_id)]] - - ?[ - id, - model, - name, - about, - created_at, - updated_at, - metadata, - default_settings, - instructions, - ] := input[id], - *agents { - developer_id: to_uuid($developer_id), - agent_id: id, - model, - name, - about, - created_at, - updated_at, - metadata, - instructions, - }, - *agent_default_settings { - agent_id: id, - frequency_penalty, - presence_penalty, - length_penalty, - repetition_penalty, - top_p, - temperature, - min_p, - preset, - }, - default_settings = { - "frequency_penalty": frequency_penalty, - "presence_penalty": presence_penalty, - "length_penalty": length_penalty, - "repetition_penalty": repetition_penalty, - "top_p": top_p, - "temperature": temperature, - "min_p": min_p, - "preset": preset, - } - - :limit 1 - """ - - queries = [ - verify_developer_id_query(developer_id), - verify_developer_owns_resource_query(developer_id, "agents", agent_id=agent_id), - get_query, - ] - - # Execute the constructed datalog query using the provided CozoClient. - # The result is returned as a pandas DataFrame. - return (queries, {"agent_id": str(agent_id), "developer_id": str(developer_id)}) diff --git a/agents-api/agents_api/models/agent/list_agents.py b/agents-api/agents_api/models/agent/list_agents.py deleted file mode 100644 index 882b6c8c6..000000000 --- a/agents-api/agents_api/models/agent/list_agents.py +++ /dev/null @@ -1,122 +0,0 @@ -from typing import Any, Literal, TypeVar -from uuid import UUID - -from beartype import beartype -from fastapi import HTTPException -from pycozo.client import QueryException -from pydantic import ValidationError - -from ...autogen.openapi_model import Agent -from ...common.utils import json -from ..utils import ( - cozo_query, - partialclass, - rewrap_exceptions, - verify_developer_id_query, - wrap_in_class, -) - -ModelT = TypeVar("ModelT", bound=Any) -T = TypeVar("T") - - -@rewrap_exceptions( - { - QueryException: partialclass(HTTPException, status_code=400), - ValidationError: partialclass(HTTPException, status_code=400), - TypeError: partialclass(HTTPException, status_code=400), - } -) -@wrap_in_class(Agent) -@cozo_query -@beartype -def list_agents( - *, - developer_id: UUID, - limit: int = 100, - offset: int = 0, - sort_by: Literal["created_at", "updated_at"] = "created_at", - direction: Literal["asc", "desc"] = "desc", - metadata_filter: dict[str, Any] = {}, -) -> tuple[list[str], dict]: - """ - Constructs and executes a datalog query to list agents from the 'cozodb' database. - - Parameters: - developer_id: UUID of the developer. - limit: Maximum number of agents to return. - offset: Number of agents to skip before starting to collect the result set. - metadata_filter: Dictionary to filter agents based on metadata. - client: Instance of CozoClient to execute the query. - """ - # Transforms the metadata_filter dictionary into a string representation for the datalog query. - metadata_filter_str = ", ".join( - [ - f"metadata->{json.dumps(k)} == {json.dumps(v)}" - for k, v in metadata_filter.items() - ] - ) - - sort = f"{'-' if direction == 'desc' else ''}{sort_by}" - - # Datalog query to retrieve agent information based on filters, sorted by creation date in descending order. - queries = [ - verify_developer_id_query(developer_id), - f""" - input[developer_id] <- [[to_uuid($developer_id)]] - - ?[ - id, - model, - name, - about, - created_at, - updated_at, - metadata, - default_settings, - instructions, - ] := input[developer_id], - *agents {{ - developer_id, - agent_id: id, - model, - name, - about, - created_at, - updated_at, - metadata, - instructions, - }}, - *agent_default_settings {{ - agent_id: id, - frequency_penalty, - presence_penalty, - length_penalty, - repetition_penalty, - top_p, - temperature, - min_p, - preset, - }}, - default_settings = {{ - "frequency_penalty": frequency_penalty, - "presence_penalty": presence_penalty, - "length_penalty": length_penalty, - "repetition_penalty": repetition_penalty, - "top_p": top_p, - "temperature": temperature, - "min_p": min_p, - "preset": preset, - }}, - {metadata_filter_str} - - :limit $limit - :offset $offset - :sort {sort} - """, - ] - - return ( - queries, - {"developer_id": str(developer_id), "limit": limit, "offset": offset}, - ) diff --git a/agents-api/agents_api/models/agent/patch_agent.py b/agents-api/agents_api/models/agent/patch_agent.py deleted file mode 100644 index 99d4e3553..000000000 --- a/agents-api/agents_api/models/agent/patch_agent.py +++ /dev/null @@ -1,132 +0,0 @@ -from typing import Any, TypeVar -from uuid import UUID - -from beartype import beartype -from fastapi import HTTPException -from pycozo.client import QueryException -from pydantic import ValidationError - -from ...autogen.openapi_model import PatchAgentRequest, ResourceUpdatedResponse -from ...common.utils.cozo import cozo_process_mutate_data -from ...common.utils.datetime import utcnow -from ...metrics.counters import increase_counter -from ..utils import ( - cozo_query, - partialclass, - rewrap_exceptions, - verify_developer_id_query, - verify_developer_owns_resource_query, - wrap_in_class, -) - -ModelT = TypeVar("ModelT", bound=Any) -T = TypeVar("T") - - -@rewrap_exceptions( - { - QueryException: partialclass(HTTPException, status_code=400), - ValidationError: partialclass(HTTPException, status_code=400), - TypeError: partialclass(HTTPException, status_code=400), - } -) -@wrap_in_class( - ResourceUpdatedResponse, - one=True, - transform=lambda d: {"id": d["agent_id"], "jobs": [], **d}, - _kind="inserted", -) -@cozo_query -@increase_counter("patch_agent") -@beartype -def patch_agent( - *, - agent_id: UUID, - developer_id: UUID, - data: PatchAgentRequest, -) -> tuple[list[str], dict]: - """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. - - Returns: - ResourceUpdatedResponse: The updated agent data. - """ - update_data = data.model_dump(exclude_unset=True) - - # Construct the query for updating agent information in the database. - # Agent update query - metadata = update_data.pop("metadata", {}) or {} - default_settings = update_data.pop("default_settings", {}) or {} - agent_update_cols, agent_update_vals = cozo_process_mutate_data( - { - **{k: v for k, v in update_data.items() if v is not None}, - "agent_id": str(agent_id), - "developer_id": str(developer_id), - "updated_at": utcnow().timestamp(), - } - ) - - update_query = f""" - # update the agent - input[{agent_update_cols}] <- $agent_update_vals - - ?[{agent_update_cols}, metadata] := - input[{agent_update_cols}], - *agents {{ - agent_id: to_uuid($agent_id), - metadata: md, - }}, - metadata = concat(md, $metadata) - - :update agents {{ - {agent_update_cols}, - metadata, - }} - :returning - """ - - # Construct the query for updating agent's default settings in the database. - # Settings update query - settings_cols, settings_vals = cozo_process_mutate_data( - { - **default_settings, - "agent_id": str(agent_id), - } - ) - - settings_update_query = f""" - # update the agent settings - ?[{settings_cols}] <- $settings_vals - - :update agent_default_settings {{ - {settings_cols} - }} - """ - - # Combine agent and settings update queries if default settings are provided. - # Combine the queries - queries = [update_query] - - if len(default_settings) != 0: - queries.insert(0, settings_update_query) - - queries = [ - verify_developer_id_query(developer_id), - verify_developer_owns_resource_query(developer_id, "agents", agent_id=agent_id), - *queries, - ] - - return ( - queries, - { - "agent_update_vals": agent_update_vals, - "settings_vals": settings_vals, - "metadata": metadata, - "agent_id": str(agent_id), - }, - ) diff --git a/agents-api/agents_api/models/agent/update_agent.py b/agents-api/agents_api/models/agent/update_agent.py deleted file mode 100644 index b36f687eb..000000000 --- a/agents-api/agents_api/models/agent/update_agent.py +++ /dev/null @@ -1,149 +0,0 @@ -from typing import Any, TypeVar -from uuid import UUID - -from beartype import beartype -from fastapi import HTTPException -from pycozo.client import QueryException -from pydantic import ValidationError - -from ...autogen.openapi_model import ResourceUpdatedResponse, UpdateAgentRequest -from ...common.utils.cozo import cozo_process_mutate_data -from ...metrics.counters import increase_counter -from ..utils import ( - cozo_query, - partialclass, - rewrap_exceptions, - verify_developer_id_query, - verify_developer_owns_resource_query, - wrap_in_class, -) - -ModelT = TypeVar("ModelT", bound=Any) -T = TypeVar("T") - - -@rewrap_exceptions( - { - QueryException: partialclass(HTTPException, status_code=400), - ValidationError: partialclass(HTTPException, status_code=400), - TypeError: partialclass(HTTPException, status_code=400), - } -) -@wrap_in_class( - ResourceUpdatedResponse, - one=True, - transform=lambda d: {"id": d["agent_id"], "jobs": [], **d}, - _kind="inserted", -) -@cozo_query -@increase_counter("update_agent") -@beartype -def update_agent( - *, - agent_id: UUID, - developer_id: UUID, - data: UpdateAgentRequest, -) -> tuple[list[str], dict]: - """ - 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. - - Returns: - ResourceUpdatedResponse: The updated agent data. - """ - default_settings = ( - data.default_settings.model_dump(exclude_none=True) - if data.default_settings - else {} - ) - update_data = data.model_dump() - - # Remove default settings from the agent update data - update_data.pop("default_settings", None) - - agent_id = str(agent_id) - developer_id = str(developer_id) - update_data["instructions"] = update_data.get("instructions", []) - update_data["instructions"] = ( - update_data["instructions"] - if isinstance(update_data["instructions"], list) - else [update_data["instructions"]] - ) - - # Construct the agent update part of the query with dynamic columns and values based on `update_data`. - # Agent update query - agent_update_cols, agent_update_vals = cozo_process_mutate_data( - { - **{k: v for k, v in update_data.items() if v is not None}, - "agent_id": agent_id, - "developer_id": developer_id, - } - ) - - update_query = f""" - # update the agent - input[{agent_update_cols}] <- $agent_update_vals - original[created_at] := *agents{{ - developer_id: to_uuid($developer_id), - agent_id: to_uuid($agent_id), - created_at, - }}, - - ?[created_at, updated_at, {agent_update_cols}] := - input[{agent_update_cols}], - original[created_at], - updated_at = now(), - - :put agents {{ - created_at, - updated_at, - {agent_update_cols} - }} - :returning - """ - - # Construct the settings update part of the query if `default_settings` are provided. - # Settings update query - settings_cols, settings_vals = cozo_process_mutate_data( - { - **default_settings, - "agent_id": agent_id, - } - ) - - settings_update_query = f""" - # update the agent settings - ?[{settings_cols}] <- $settings_vals - - :put agent_default_settings {{ - {settings_cols} - }} - """ - - # Combine agent and settings update queries into a single query string. - # Combine the queries - queries = [update_query] - - if len(default_settings) != 0: - queries.insert(0, settings_update_query) - - queries = [ - verify_developer_id_query(developer_id), - verify_developer_owns_resource_query(developer_id, "agents", agent_id=agent_id), - *queries, - ] - - return ( - queries, - { - "agent_update_vals": agent_update_vals, - "settings_vals": settings_vals, - "agent_id": agent_id, - "developer_id": developer_id, - }, - ) diff --git a/agents-api/agents_api/models/docs/__init__.py b/agents-api/agents_api/models/docs/__init__.py deleted file mode 100644 index 0ba3db0d4..000000000 --- a/agents-api/agents_api/models/docs/__init__.py +++ /dev/null @@ -1,25 +0,0 @@ -""" -Module: agents_api/models/docs - -This module is responsible for managing document-related operations within the application, particularly for agents and possibly other entities. It serves as a core component of the document management system, enabling features such as document creation, listing, deletion, and embedding of snippets for enhanced search and retrieval capabilities. - -Main functionalities include: -- Creating new documents and associating them with agents or users. -- Listing documents based on various criteria, including ownership and metadata filters. -- Deleting documents by their unique identifiers. -- Embedding document snippets for retrieval purposes. - -The module interacts with other parts of the application, such as the agents and users modules, to provide a comprehensive document management system. Its role is crucial in enabling document search, retrieval, and management features within the context of agents and users. - -This documentation aims to provide clear, concise, and sufficient context for new developers or contributors to understand the module's role without needing to dive deep into the code immediately. -""" - -# ruff: noqa: F401, F403, F405 - -from .create_doc import create_doc -from .delete_doc import delete_doc -from .embed_snippets import embed_snippets -from .get_doc import get_doc -from .list_docs import list_docs -from .search_docs_by_embedding import search_docs_by_embedding -from .search_docs_by_text import search_docs_by_text diff --git a/agents-api/agents_api/models/docs/create_doc.py b/agents-api/agents_api/models/docs/create_doc.py deleted file mode 100644 index ceb8b5fe0..000000000 --- a/agents-api/agents_api/models/docs/create_doc.py +++ /dev/null @@ -1,141 +0,0 @@ -from typing import Any, Literal, TypeVar -from uuid import UUID - -from beartype import beartype -from fastapi import HTTPException -from pycozo.client import QueryException -from pydantic import ValidationError -from uuid_extensions import uuid7 - -from ...autogen.openapi_model import CreateDocRequest, Doc -from ...common.utils.cozo import cozo_process_mutate_data -from ...metrics.counters import increase_counter -from ..utils import ( - cozo_query, - partialclass, - rewrap_exceptions, - verify_developer_id_query, - verify_developer_owns_resource_query, - wrap_in_class, -) - -ModelT = TypeVar("ModelT", bound=Any) -T = TypeVar("T") - - -@rewrap_exceptions( - { - QueryException: partialclass(HTTPException, status_code=400), - ValidationError: partialclass(HTTPException, status_code=400), - TypeError: partialclass(HTTPException, status_code=400), - } -) -@wrap_in_class( - Doc, - one=True, - transform=lambda d: { - "id": UUID(d["doc_id"]), - **d, - }, -) -@cozo_query -@increase_counter("create_doc") -@beartype -def create_doc( - *, - developer_id: UUID, - owner_type: Literal["user", "agent"], - owner_id: UUID, - doc_id: UUID | None = None, - data: CreateDocRequest, -) -> tuple[list[str], dict]: - """ - 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. - doc_id (UUID): The UUID of the document to be created. - data (CreateDocRequest): The content of the document. - """ - - doc_id = str(doc_id or uuid7()) - owner_id = str(owner_id) - - if isinstance(data.content, str): - data.content = [data.content] - - data.metadata = data.metadata or {} - - doc_data = data.model_dump() - doc_data.pop("embed_instruction", None) - content = doc_data.pop("content") - - doc_data["owner_type"] = owner_type - doc_data["owner_id"] = owner_id - doc_data["doc_id"] = doc_id - - doc_cols, doc_rows = cozo_process_mutate_data(doc_data) - - snippet_cols, snippet_rows = "", [] - - # Process each content snippet and prepare data for the datalog query. - for snippet_idx, snippet in enumerate(content): - snippet_cols, new_snippet_rows = cozo_process_mutate_data( - dict( - doc_id=doc_id, - index=snippet_idx, - content=snippet, - ) - ) - - snippet_rows += new_snippet_rows - - create_snippets_query = f""" - ?[{snippet_cols}] <- $snippet_rows - - :create _snippets {{ {snippet_cols} }} - }} {{ - ?[{snippet_cols}] <- $snippet_rows - :insert snippets {{ {snippet_cols} }} - :returning - """ - - # Construct the datalog query for creating the document and its snippets. - create_doc_query = f""" - ?[{doc_cols}] <- $doc_rows - - :create _docs {{ {doc_cols} }} - }} {{ - ?[{doc_cols}] <- $doc_rows - :insert docs {{ {doc_cols} }} - :returning - }} {{ - snippet_rows[collect(content)] := - *_snippets {{ - content - }} - - ?[{doc_cols}, content, created_at] := - *_docs {{ {doc_cols} }}, - snippet_rows[content], - created_at = now() - """ - - queries = [ - verify_developer_id_query(developer_id), - verify_developer_owns_resource_query( - developer_id, f"{owner_type}s", **{f"{owner_type}_id": owner_id} - ), - create_snippets_query, - create_doc_query, - ] - - # Execute the constructed datalog query and return the results as a DataFrame. - return ( - queries, - { - "doc_rows": doc_rows, - "snippet_rows": snippet_rows, - }, - ) diff --git a/agents-api/agents_api/models/docs/delete_doc.py b/agents-api/agents_api/models/docs/delete_doc.py deleted file mode 100644 index c02705756..000000000 --- a/agents-api/agents_api/models/docs/delete_doc.py +++ /dev/null @@ -1,102 +0,0 @@ -from typing import Any, TypeVar -from uuid import UUID - -from beartype import beartype -from fastapi import HTTPException -from pycozo.client import QueryException -from pydantic import ValidationError - -from ...autogen.openapi_model import ResourceDeletedResponse -from ...common.utils.datetime import utcnow -from ..utils import ( - cozo_query, - partialclass, - rewrap_exceptions, - verify_developer_id_query, - verify_developer_owns_resource_query, - wrap_in_class, -) - -ModelT = TypeVar("ModelT", bound=Any) -T = TypeVar("T") - - -@rewrap_exceptions( - { - QueryException: partialclass(HTTPException, status_code=400), - ValidationError: partialclass(HTTPException, status_code=400), - TypeError: partialclass(HTTPException, status_code=400), - } -) -@wrap_in_class( - ResourceDeletedResponse, - one=True, - transform=lambda d: { - "id": UUID(d.pop("doc_id")), - "deleted_at": utcnow(), - "jobs": [], - }, - _kind="deleted", -) -@cozo_query -@beartype -def delete_doc( - *, - developer_id: UUID, - owner_id: UUID, - owner_type: str, - doc_id: UUID, -) -> tuple[list[str], dict]: - """Constructs and returns a datalog query for deleting documents and associated information snippets. - - This function targets the 'cozodb' database, allowing for the removal of documents and their related information snippets based on the provided document ID and owner (user or agent). - - Parameters: - doc_id (UUID): The UUID of the document to be deleted. - client (CozoClient): An instance of the CozoClient to execute the query. - - Returns: - pd.DataFrame: The result of the executed datalog query. - """ - # Convert UUID parameters to string format for use in the datalog query - doc_id = str(doc_id) - owner_id = str(owner_id) - - # The following query is divided into two main parts: - # 1. Deleting information snippets associated with the document - # 2. Deleting the document itself - delete_snippets_query = """ - # This section constructs the subquery for identifying and deleting all information snippets associated with the given document ID. - # Delete snippets - input[doc_id] <- [[to_uuid($doc_id)]] - ?[doc_id, index] := - input[doc_id], - *snippets { - doc_id, - index, - } - - :delete snippets { - doc_id, - index - } - """ - - delete_doc_query = """ - # Delete the docs - ?[doc_id, owner_type, owner_id] <- [[ to_uuid($doc_id), $owner_type, to_uuid($owner_id) ]] - - :delete docs { doc_id, owner_type, owner_id } - :returning - """ - - queries = [ - verify_developer_id_query(developer_id), - verify_developer_owns_resource_query( - developer_id, f"{owner_type}s", **{f"{owner_type}_id": owner_id} - ), - delete_snippets_query, - delete_doc_query, - ] - - return (queries, {"doc_id": doc_id, "owner_type": owner_type, "owner_id": owner_id}) diff --git a/agents-api/agents_api/models/docs/embed_snippets.py b/agents-api/agents_api/models/docs/embed_snippets.py deleted file mode 100644 index 8d8ae1e62..000000000 --- a/agents-api/agents_api/models/docs/embed_snippets.py +++ /dev/null @@ -1,102 +0,0 @@ -"""Module for embedding documents in the cozodb database. Contains functions to update document embeddings.""" - -from typing import Any, TypeVar -from uuid import UUID - -from beartype import beartype -from fastapi import HTTPException -from pycozo.client import QueryException -from pydantic import ValidationError - -from ...autogen.openapi_model import ResourceUpdatedResponse -from ...common.utils.cozo import cozo_process_mutate_data -from ...common.utils.datetime import utcnow -from ..utils import ( - cozo_query, - partialclass, - rewrap_exceptions, - verify_developer_id_query, - wrap_in_class, -) - -ModelT = TypeVar("ModelT", bound=Any) -T = TypeVar("T") - - -@rewrap_exceptions( - { - QueryException: partialclass(HTTPException, status_code=400), - ValidationError: partialclass(HTTPException, status_code=400), - TypeError: partialclass(HTTPException, status_code=400), - } -) -@wrap_in_class( - ResourceUpdatedResponse, - one=True, - transform=lambda d: {"id": d["doc_id"], "updated_at": utcnow(), "jobs": []}, - _kind="inserted", -) -@cozo_query -@beartype -def embed_snippets( - *, - developer_id: UUID, - doc_id: UUID, - snippet_indices: list[int] | tuple[int, ...], - embeddings: list[list[float]], - embedding_size: int = 1024, -) -> tuple[list[str], dict]: - """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 = str(doc_id) - - # Ensure the number of snippet indices matches the number of embeddings. - assert len(snippet_indices) == len(embeddings) - assert all(len(embedding) == embedding_size for embedding in embeddings) - assert min(snippet_indices) >= 0 - - # Ensure all embeddings are non-zero. - assert all(sum(embedding) for embedding in embeddings) - - # Create a list of records to update the document snippet embeddings in the database. - records = [ - {"doc_id": doc_id, "index": snippet_idx, "embedding": embedding} - for snippet_idx, embedding in zip(snippet_indices, embeddings) - ] - - cols, vals = cozo_process_mutate_data(records) - - # Ensure that index is present in the records. - check_indices_query = f""" - ?[index] := - *snippets {{ - doc_id: $doc_id, - index, - }}, - index > {max(snippet_indices)} - - :limit 1 - :assert none - """ - - # Define the datalog query for updating document snippet embeddings in the database. - embed_query = f""" - ?[{cols}] <- $vals - - :update snippets {{ {cols} }} - :returning - """ - - queries = [ - verify_developer_id_query(developer_id), - check_indices_query, - embed_query, - ] - - return (queries, {"vals": vals, "doc_id": doc_id}) diff --git a/agents-api/agents_api/models/docs/get_doc.py b/agents-api/agents_api/models/docs/get_doc.py deleted file mode 100644 index d47cc80a8..000000000 --- a/agents-api/agents_api/models/docs/get_doc.py +++ /dev/null @@ -1,103 +0,0 @@ -"""Module for retrieving document snippets from the CozoDB based on document IDs.""" - -from typing import Any, TypeVar -from uuid import UUID - -from beartype import beartype -from fastapi import HTTPException -from pycozo.client import QueryException -from pydantic import ValidationError - -from ...autogen.openapi_model import Doc -from ..utils import ( - cozo_query, - partialclass, - rewrap_exceptions, - verify_developer_id_query, - wrap_in_class, -) - -ModelT = TypeVar("ModelT", bound=Any) -T = TypeVar("T") - - -@rewrap_exceptions( - { - lambda e: isinstance(e, AssertionError) - and "Expected one result" in repr(e): partialclass( - HTTPException, status_code=404 - ), - QueryException: partialclass(HTTPException, status_code=400), - ValidationError: partialclass(HTTPException, status_code=400), - TypeError: partialclass(HTTPException, status_code=400), - } -) -@wrap_in_class( - Doc, - 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]) - if s[2] is not None - ], - **d, - }, -) -@cozo_query -@beartype -def get_doc( - *, - developer_id: UUID, - doc_id: UUID, -) -> tuple[list[str], dict]: - """ - Retrieves snippets of documents by their ID from the CozoDB. - - Parameters: - doc_id (UUID): The unique identifier of the document. - client (CozoClient, optional): The CozoDB client instance. Defaults to a pre-configured client. - - Returns: - pd.DataFrame: A DataFrame containing the document snippets and related metadata. - """ - - doc_id = str(doc_id) - - get_query = """ - input[doc_id] <- [[to_uuid($doc_id)]] - snippets[collect(snippet_data)] := - input[doc_id], - *snippets { - doc_id, - index, - content, - embedding, - }, - snippet_data = [index, content, embedding] - - ?[ - id, - title, - snippet_data, - created_at, - metadata, - ] := input[id], - *docs { - doc_id: id, - title, - created_at, - metadata, - }, - snippets[snippet_data] - - :limit 1 - """ - - queries = [ - verify_developer_id_query(developer_id), - get_query, - ] - - return (queries, {"doc_id": doc_id}) diff --git a/agents-api/agents_api/models/docs/list_docs.py b/agents-api/agents_api/models/docs/list_docs.py deleted file mode 100644 index dd389d58c..000000000 --- a/agents-api/agents_api/models/docs/list_docs.py +++ /dev/null @@ -1,141 +0,0 @@ -"""This module contains functions for querying document-related data from the 'cozodb' database using datalog queries.""" - -import json -from typing import Any, Literal, TypeVar -from uuid import UUID - -from beartype import beartype -from fastapi import HTTPException -from pycozo.client import QueryException -from pydantic import ValidationError - -from ...autogen.openapi_model import Doc -from ..utils import ( - cozo_query, - partialclass, - rewrap_exceptions, - verify_developer_id_query, - verify_developer_owns_resource_query, - wrap_in_class, -) - -ModelT = TypeVar("ModelT", bound=Any) -T = TypeVar("T") - - -@rewrap_exceptions( - { - QueryException: partialclass(HTTPException, status_code=400), - ValidationError: partialclass(HTTPException, status_code=400), - TypeError: partialclass(HTTPException, status_code=400), - } -) -@wrap_in_class( - 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]) - if s[2] is not None - ], - **d, - }, -) -@cozo_query -@beartype -def list_docs( - *, - developer_id: UUID, - owner_type: Literal["user", "agent"], - owner_id: UUID, - limit: int = 100, - offset: int = 0, - sort_by: Literal["created_at"] = "created_at", - direction: Literal["asc", "desc"] = "desc", - metadata_filter: dict[str, Any] = {}, - include_without_embeddings: bool = False, -) -> tuple[list[str], dict]: - """ - Constructs and returns a datalog query for listing documents and their associated information snippets. - - Parameters: - developer_id (UUID): The unique identifier of the developer associated with the documents. - owner_id (UUID): The unique identifier of the owner (user or agent) associated with the documents. - owner_type (Literal["user", "agent"]): The type of owner associated with the documents. - limit (int): The maximum number of documents to return. - offset (int): The number of documents to skip before returning the results. - sort_by (Literal["created_at"]): The field to sort the documents by. - direction (Literal["asc", "desc"]): The direction to sort the documents in. - metadata_filter (dict): A dictionary of metadata filters to apply to the documents. - include_without_embeddings (bool): Whether to include documents without embeddings in the results. - - Returns: - Doc[] - """ - - # Transforms the metadata_filter dictionary into a string representation for the datalog query. - metadata_filter_str = ", ".join( - [ - f"metadata->{json.dumps(k)} == {json.dumps(v)}" - for k, v in metadata_filter.items() - ] - ) - - owner_id = str(owner_id) - sort = f"{'-' if direction == 'desc' else ''}{sort_by}" - - get_query = f""" - snippets[id, collect(snippet_data)] := - *snippets {{ - doc_id: id, - index, - content, - embedding, - }}, - {"" if include_without_embeddings else "not is_null(embedding),"} - snippet_data = [index, content, embedding] - - ?[ - owner_type, - id, - title, - snippet_data, - created_at, - metadata, - ] := - owner_type = $owner_type, - owner_id = to_uuid($owner_id), - *docs {{ - owner_type, - owner_id, - doc_id: id, - title, - created_at, - metadata, - }}, - snippets[id, snippet_data], - {metadata_filter_str} - - :limit $limit - :offset $offset - :sort {sort} - """ - - queries = [ - verify_developer_id_query(developer_id), - verify_developer_owns_resource_query( - developer_id, f"{owner_type}s", **{f"{owner_type}_id": owner_id} - ), - get_query, - ] - - return ( - queries, - { - "owner_id": owner_id, - "owner_type": owner_type, - "limit": limit, - "offset": offset, - }, - ) diff --git a/agents-api/agents_api/models/docs/mmr.py b/agents-api/agents_api/models/docs/mmr.py deleted file mode 100644 index d214e8c04..000000000 --- a/agents-api/agents_api/models/docs/mmr.py +++ /dev/null @@ -1,109 +0,0 @@ -from __future__ import annotations - -import logging -from typing import Union - -import numpy as np - -Matrix = Union[list[list[float]], list[np.ndarray], np.ndarray] - -logger = logging.getLogger(__name__) - - -def _cosine_similarity(x: Matrix, y: Matrix) -> np.ndarray: - """Row-wise cosine similarity between two equal-width matrices. - - Args: - x: A matrix of shape (n, m). - y: A matrix of shape (k, m). - - Returns: - A matrix of shape (n, k) where each element (i, j) is the cosine similarity - between the ith row of X and the jth row of Y. - - Raises: - ValueError: If the number of columns in X and Y are not the same. - ImportError: If numpy is not installed. - """ - - if len(x) == 0 or len(y) == 0: - return np.array([]) - - x = [xx for xx in x if xx is not None] - y = [yy for yy in y if yy is not None] - - x = np.array(x) - y = np.array(y) - if x.shape[1] != y.shape[1]: - msg = ( - f"Number of columns in X and Y must be the same. X has shape {x.shape} " - f"and Y has shape {y.shape}." - ) - raise ValueError(msg) - try: - import simsimd as simd # type: ignore - - x = np.array(x, dtype=np.float32) - y = np.array(y, dtype=np.float32) - z = 1 - np.array(simd.cdist(x, y, metric="cosine")) - return z - except ImportError: - logger.debug( - "Unable to import simsimd, defaulting to NumPy implementation. If you want " - "to use simsimd please install with `pip install simsimd`." - ) - x_norm = np.linalg.norm(x, axis=1) - y_norm = np.linalg.norm(y, axis=1) - # Ignore divide by zero errors run time warnings as those are handled below. - with np.errstate(divide="ignore", invalid="ignore"): - similarity = np.dot(x, y.T) / np.outer(x_norm, y_norm) - similarity[np.isnan(similarity) | np.isinf(similarity)] = 0.0 - return similarity - - -def maximal_marginal_relevance( - query_embedding: np.ndarray, - embedding_list: list, - lambda_mult: float = 0.5, - k: int = 4, -) -> list[int]: - """Calculate maximal marginal relevance. - - Args: - query_embedding: The query embedding. - embedding_list: A list of embeddings. - lambda_mult: The lambda parameter for MMR. Default is 0.5. - k: The number of embeddings to return. Default is 4. - - Returns: - A list of indices of the embeddings to return. - - Raises: - ImportError: If numpy is not installed. - """ - - if min(k, len(embedding_list)) <= 0: - return [] - if query_embedding.ndim == 1: - query_embedding = np.expand_dims(query_embedding, axis=0) - similarity_to_query = _cosine_similarity(query_embedding, embedding_list)[0] - most_similar = int(np.argmax(similarity_to_query)) - idxs = [most_similar] - selected = np.array([embedding_list[most_similar]]) - while len(idxs) < min(k, len(embedding_list)): - best_score = -np.inf - idx_to_add = -1 - similarity_to_selected = _cosine_similarity(embedding_list, selected) - for i, query_score in enumerate(similarity_to_query): - if i in idxs: - continue - redundant_score = max(similarity_to_selected[i]) - equation_score = ( - lambda_mult * query_score - (1 - lambda_mult) * redundant_score - ) - if equation_score > best_score: - best_score = equation_score - idx_to_add = i - idxs.append(idx_to_add) - selected = np.append(selected, [embedding_list[idx_to_add]], axis=0) - return idxs 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 deleted file mode 100644 index 992e12f9d..000000000 --- a/agents-api/agents_api/models/docs/search_docs_by_embedding.py +++ /dev/null @@ -1,369 +0,0 @@ -"""This module contains functions for searching documents in the CozoDB based on embedding queries.""" - -import json -from typing import Any, Literal, TypeVar -from uuid import UUID - -from beartype import beartype -from fastapi import HTTPException -from pycozo.client import QueryException -from pydantic import ValidationError - -from ...autogen.openapi_model import DocReference -from ..utils import ( - cozo_query, - partialclass, - rewrap_exceptions, - verify_developer_id_query, - verify_developer_owns_resource_query, - wrap_in_class, -) - -ModelT = TypeVar("ModelT", bound=Any) -T = TypeVar("T") - - -@rewrap_exceptions( - { - QueryException: partialclass(HTTPException, status_code=400), - ValidationError: partialclass(HTTPException, status_code=400), - TypeError: partialclass(HTTPException, status_code=400), - } -) -@wrap_in_class( - DocReference, - transform=lambda d: { - "owner": { - "id": d["owner_id"], - "role": d["owner_type"], - }, - "metadata": d.get("metadata", {}), - **d, - }, -) -@cozo_query -@beartype -def search_docs_by_embedding( - *, - developer_id: UUID, - owners: list[tuple[Literal["user", "agent"], UUID]], - query_embedding: list[float], - k: int = 3, - confidence: float = 0.5, - ef: int = 50, - embedding_size: int = 1024, - ann_threshold: int = 1_000_000, - metadata_filter: dict[str, Any] = {}, -) -> tuple[str, dict]: - """ - 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. - embedding_size (int): Embedding vector length - metadata_filter (dict[str, Any]): Dictionary to filter agents based on metadata. - """ - - assert len(query_embedding) == embedding_size - assert sum(query_embedding) - - metadata_filter_str = ", ".join( - [ - f"metadata->{json.dumps(k)} == {json.dumps(v)}" - for k, v in metadata_filter.items() - ] - ) - - owners: list[list[str]] = [ - [owner_type, str(owner_id)] for owner_type, owner_id in owners - ] - - # Calculate the search radius based on confidence level - radius: float = 1.0 - confidence - - determine_knn_ann_query = f""" - owners[owner_type, owner_id] <- $owners - snippet_counter[count(item)] := - owners[owner_type, owner_id_str], - owner_id = to_uuid(owner_id_str), - *docs {{ - owner_type, - owner_id, - doc_id: item, - metadata, - }} - {', ' + metadata_filter_str if metadata_filter_str.strip() else ''} - - ?[use_ann] := - snippet_counter[count], - count > {ann_threshold}, - use_ann = true - - :limit 1 - :create _determine_knn_ann {{ - use_ann - }} - """ - - # Construct the datalog query for searching document snippets - search_query = f""" - # %debug _determine_knn_ann - %if {{ - ?[use_ann] := *_determine_knn_ann{{ use_ann }} - }} - - %then {{ - owners[owner_type, owner_id] <- $owners - input[ - owner_type, - owner_id, - query_embedding, - ] := - owners[owner_type, owner_id_str], - owner_id = to_uuid(owner_id_str), - query_embedding = vec($query_embedding) - - # Search for documents by owner - ?[ - doc_id, - index, - title, - content, - distance, - embedding, - metadata, - ] := - # Get input values - input[owner_type, owner_id, query], - - # Restrict the search to all documents that match the owner - *docs {{ - owner_type, - owner_id, - doc_id, - title, - metadata, - }}, - - # Search for snippets in the embedding space - ~snippets:embedding_space {{ - doc_id, - index, - content - | - query: query, - k: {k}, - ef: {ef}, - radius: {radius}, - bind_distance: distance, - bind_vector: embedding, - }} - - :sort distance - :limit {k} - - :create _search_result {{ - doc_id, - index, - title, - content, - distance, - embedding, - metadata, - }} - }} - - %else {{ - owners[owner_type, owner_id] <- $owners - input[ - owner_type, - owner_id, - query_embedding, - ] := - owners[owner_type, owner_id_str], - owner_id = to_uuid(owner_id_str), - query_embedding = vec($query_embedding) - - # Search for documents by owner - ?[ - doc_id, - index, - title, - content, - distance, - embedding, - metadata, - ] := - # Get input values - input[owner_type, owner_id, query], - - # Restrict the search to all documents that match the owner - *docs {{ - owner_type, - owner_id, - doc_id, - title, - metadata, - }}, - - # Search for snippets in the embedding space - *snippets {{ - doc_id, - index, - content, - embedding, - }}, - !is_null(embedding), - distance = cos_dist(query, embedding), - distance <= {radius} - - :sort distance - :limit {k} - - :create _search_result {{ - doc_id, - index, - title, - content, - distance, - embedding, - metadata, - }} - }} - %end - """ - - normal_interim_query = f""" - owners[owner_type, owner_id] <- $owners - - ?[ - owner_type, - owner_id, - doc_id, - snippet_data, - distance, - title, - embedding, - metadata, - ] := - owners[owner_type, owner_id_str], - owner_id = to_uuid(owner_id_str), - *_search_result{{ doc_id, index, title, content, distance, embedding, metadata }}, - snippet_data = [index, content] - - :sort distance - :limit {k} - - :create _interim {{ - owner_type, - owner_id, - doc_id, - snippet_data, - distance, - title, - embedding, - metadata, - }} - """ - - collect_query = """ - n[ - doc_id, - owner_type, - owner_id, - unique(snippet_data), - distance, - title, - embedding, - metadata, - ] := - *_interim { - owner_type, - owner_id, - doc_id, - snippet_data, - distance, - title, - embedding, - metadata, - } - - m[ - doc_id, - owner_type, - owner_id, - snippet, - distance, - title, - metadata, - ] := - n[ - doc_id, - owner_type, - owner_id, - snippet_data, - distance, - title, - embedding, - metadata, - ], - snippet = { - "index": snippet_datum->0, - "content": snippet_datum->1, - "embedding": embedding, - }, - snippet_datum in snippet_data - - ?[ - id, - owner_type, - owner_id, - snippet, - distance, - title, - metadata, - ] := m[ - id, - owner_type, - owner_id, - snippet, - distance, - title, - metadata, - ] - - :sort distance - """ - - verify_query = "}\n\n{".join( - [ - verify_developer_id_query(developer_id), - *[ - verify_developer_owns_resource_query( - developer_id, f"{owner_type}s", **{f"{owner_type}_id": owner_id} - ) - for owner_type, owner_id in owners - ], - ] - ) - - query = f""" - {{ {verify_query} }} - {{ {determine_knn_ann_query} }} - {search_query} - {{ {normal_interim_query} }} - {{ {collect_query} }} - """ - - return ( - query, - { - "owners": owners, - "query_embedding": query_embedding, - }, - ) 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 deleted file mode 100644 index ac1a9f54f..000000000 --- a/agents-api/agents_api/models/docs/search_docs_by_text.py +++ /dev/null @@ -1,206 +0,0 @@ -"""This module contains functions for searching documents in the CozoDB based on embedding queries.""" - -import json -import re -from typing import Any, Literal, TypeVar -from uuid import UUID - -from beartype import beartype -from fastapi import HTTPException -from pycozo.client import QueryException -from pydantic import ValidationError - -from ...autogen.openapi_model import DocReference -from ...common.nlp import paragraph_to_custom_queries -from ..utils import ( - cozo_query, - partialclass, - rewrap_exceptions, - verify_developer_id_query, - verify_developer_owns_resource_query, - wrap_in_class, -) - -ModelT = TypeVar("ModelT", bound=Any) -T = TypeVar("T") - - -@rewrap_exceptions( - { - QueryException: partialclass(HTTPException, status_code=400), - ValidationError: partialclass(HTTPException, status_code=400), - TypeError: partialclass(HTTPException, status_code=400), - } -) -@wrap_in_class( - DocReference, - transform=lambda d: { - "owner": { - "id": d["owner_id"], - "role": d["owner_type"], - }, - "metadata": d.get("metadata", {}), - **d, - }, -) -@cozo_query -@beartype -def search_docs_by_text( - *, - developer_id: UUID, - owners: list[tuple[Literal["user", "agent"], UUID]], - query: str, - k: int = 3, - metadata_filter: dict[str, Any] = {}, -) -> tuple[list[str], dict]: - """ - 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. - metadata_filter (dict[str, Any]): Dictionary to filter agents based on metadata. - """ - metadata_filter_str = ", ".join( - [ - f"metadata->{json.dumps(k)} == {json.dumps(v)}" - for k, v in metadata_filter.items() - ] - ) - - owners: list[list[str]] = [ - [owner_type, str(owner_id)] for owner_type, owner_id in owners - ] - - # See: https://docs.cozodb.org/en/latest/vector.html#full-text-search-fts - fts_queries = paragraph_to_custom_queries(query) or [ - re.sub(r"[^\w\s\-_]+", "", query) - ] - - # Construct the datalog query for searching document snippets - search_query = f""" - owners[owner_type, owner_id] <- $owners - input[ - owner_type, - owner_id, - ] := - owners[owner_type, owner_id_str], - owner_id = to_uuid(owner_id_str) - - candidate[doc_id] := - input[owner_type, owner_id], - *docs {{ - owner_type, - owner_id, - doc_id, - metadata, - }} - {', ' + metadata_filter_str if metadata_filter_str.strip() else ''} - - # search_result[ - # doc_id, - # snippet_data, - # distance, - # ] := - # candidate[doc_id], - # ~snippets:lsh {{ - # doc_id, - # index, - # content - # | - # query: $query, - # k: {k}, - # }}, - # distance = 10000000, # Very large distance to depict no valid distance - # snippet_data = [index, content] - - search_result[ - doc_id, - snippet_data, - distance, - ] := - candidate[doc_id], - ~snippets:fts {{ - doc_id, - index, - content - | - query: query, - k: {k}, - score_kind: 'tf_idf', - bind_score: score, - }}, - query in $fts_queries, - distance = -score, - snippet_data = [index, content] - - m[ - doc_id, - snippet, - distance, - title, - owner_type, - owner_id, - metadata, - ] := - candidate[doc_id], - *docs {{ - owner_type, - owner_id, - doc_id, - title, - metadata, - }}, - search_result [ - doc_id, - snippet_data, - distance, - ], - snippet = {{ - "index": snippet_data->0, - "content": snippet_data->1, - }} - - - ?[ - id, - owner_type, - owner_id, - snippet, - distance, - title, - metadata, - ] := - candidate[id], - input[owner_type, owner_id], - m[ - id, - snippet, - distance, - title, - owner_type, - owner_id, - metadata, - ] - - # Sort the results by distance to find the closest matches - :sort distance - :limit {k} - """ - - queries = [ - verify_developer_id_query(developer_id), - *[ - verify_developer_owns_resource_query( - developer_id, f"{owner_type}s", **{f"{owner_type}_id": owner_id} - ) - for owner_type, owner_id in owners - ], - search_query, - ] - - return ( - queries, - {"owners": owners, "query": query, "fts_queries": fts_queries}, - ) diff --git a/agents-api/agents_api/models/docs/search_docs_hybrid.py b/agents-api/agents_api/models/docs/search_docs_hybrid.py deleted file mode 100644 index c43f8c97b..000000000 --- a/agents-api/agents_api/models/docs/search_docs_hybrid.py +++ /dev/null @@ -1,138 +0,0 @@ -"""This module contains functions for searching documents in the CozoDB based on embedding queries.""" - -from statistics import mean, stdev -from typing import Any, Literal -from uuid import UUID - -from beartype import beartype - -from ...autogen.openapi_model import DocReference -from ..utils import run_concurrently -from .search_docs_by_embedding import search_docs_by_embedding -from .search_docs_by_text import search_docs_by_text - - -# Distribution based score normalization -# https://medium.com/plain-simple-software/distribution-based-score-fusion-dbsf-a-new-approach-to-vector-search-ranking-f87c37488b18 -def dbsf_normalize(scores: list[float]) -> list[float]: - """ - Scores scaled using minmax scaler with our custom feature range - (extremes indicated as 3 standard deviations from the mean) - """ - if len(scores) < 2: - return scores - - sd = stdev(scores) - if sd == 0: - return scores - - m = mean(scores) - m3d = 3 * sd + m - m_3d = m - 3 * sd - - return [(s - m_3d) / (m3d - m_3d) for s in scores] - - -def dbsf_fuse( - text_results: list[DocReference], - embedding_results: list[DocReference], - alpha: float = 0.7, # Weight of the embedding search results (this is a good default) -) -> list[DocReference]: - """ - Weighted reciprocal-rank fusion of text and embedding search results - """ - all_docs = {doc.id: doc for doc in text_results + embedding_results} - - text_scores: dict[UUID, float] = { - doc.id: -(doc.distance or 0.0) for doc in text_results - } - - # Because these are cosine distances, we need to invert them - embedding_scores: dict[UUID, float] = { - doc.id: 1.0 - doc.distance for doc in embedding_results - } - - # normalize the scores - text_scores_normalized = dbsf_normalize(list(text_scores.values())) - text_scores = { - doc_id: score - for doc_id, score in zip(text_scores.keys(), text_scores_normalized) - } - - embedding_scores_normalized = dbsf_normalize(list(embedding_scores.values())) - embedding_scores = { - doc_id: score - for doc_id, score in zip(embedding_scores.keys(), embedding_scores_normalized) - } - - # Combine the scores - text_weight: float = 1 - alpha - embedding_weight: float = alpha - - combined_scores = [] - - for id in all_docs.keys(): - text_score = text_weight * text_scores.get(id, 0) - embedding_score = embedding_weight * embedding_scores.get(id, 0) - - combined_scores.append((id, text_score + embedding_score)) - - # Sort by the combined score - combined_scores = sorted(combined_scores, key=lambda x: x[1], reverse=True) - - # Rank the results - ranked_results = [] - for id, score in combined_scores: - doc = all_docs[id].model_copy() - doc.distance = 1.0 - score - ranked_results.append(doc) - - return ranked_results - - -@beartype -def search_docs_hybrid( - *, - developer_id: UUID, - owners: list[tuple[Literal["user", "agent"], UUID]], - query: str, - query_embedding: list[float], - k: int = 3, - alpha: float = 0.7, # Weight of the embedding search results (this is a good default) - embed_search_options: dict = {}, - text_search_options: dict = {}, - metadata_filter: dict[str, Any] = {}, -) -> list[DocReference]: - # Parallelize the text and embedding search queries - fns = [ - search_docs_by_text if bool(query.strip()) else lambda: [], - search_docs_by_embedding if bool(sum(query_embedding)) else lambda: [], - ] - - kwargs_list = [ - { - "developer_id": developer_id, - "owners": owners, - "query": query, - "k": k, - "metadata_filter": metadata_filter, - **text_search_options, - } - if bool(query.strip()) - else {}, - { - "developer_id": developer_id, - "owners": owners, - "query_embedding": query_embedding, - "k": k, - "metadata_filter": metadata_filter, - **embed_search_options, - } - if bool(sum(query_embedding)) - else {}, - ] - - results = run_concurrently(fns, kwargs_list=kwargs_list) - text_results, embedding_results = results - - return dbsf_fuse(text_results, embedding_results, alpha)[:k] diff --git a/agents-api/agents_api/models/entry/__init__.py b/agents-api/agents_api/models/entry/__init__.py deleted file mode 100644 index 32231c364..000000000 --- a/agents-api/agents_api/models/entry/__init__.py +++ /dev/null @@ -1,19 +0,0 @@ -""" -The `entry` module is responsible for managing entries related to agents' activities and interactions within the 'cozodb' database. It provides a comprehensive set of functionalities for adding, deleting, summarizing, and retrieving entries, as well as processing them to retrieve memory context based on embeddings. - -Key functionalities include: -- Adding entries to the database. -- Deleting entries from the database based on session IDs. -- Summarizing entries and managing their relationships. -- Retrieving entries from the database, including top-level entries and entries based on session IDs. -- Processing entries to retrieve memory context based on embeddings. - -The module utilizes pandas DataFrames for handling query results and integrates with the CozoClient for database operations, ensuring efficient and effective management of entries. -""" - -# ruff: noqa: F401, F403, F405 - -from .create_entries import create_entries -from .delete_entries import delete_entries -from .get_history import get_history -from .list_entries import list_entries diff --git a/agents-api/agents_api/models/entry/create_entries.py b/agents-api/agents_api/models/entry/create_entries.py deleted file mode 100644 index 140a5696b..000000000 --- a/agents-api/agents_api/models/entry/create_entries.py +++ /dev/null @@ -1,128 +0,0 @@ -from typing import Any, TypeVar -from uuid import UUID - -from beartype import beartype -from fastapi import HTTPException -from pycozo.client import QueryException -from pydantic import ValidationError -from uuid_extensions import uuid7 - -from ...autogen.openapi_model import CreateEntryRequest, Entry, Relation -from ...common.utils.cozo import cozo_process_mutate_data -from ...common.utils.datetime import utcnow -from ...common.utils.messages import content_to_json -from ...metrics.counters import increase_counter -from ..utils import ( - cozo_query, - mark_session_updated_query, - partialclass, - rewrap_exceptions, - verify_developer_id_query, - verify_developer_owns_resource_query, - wrap_in_class, -) - -ModelT = TypeVar("ModelT", bound=Any) -T = TypeVar("T") - - -@rewrap_exceptions( - { - QueryException: partialclass(HTTPException, status_code=400), - ValidationError: partialclass(HTTPException, status_code=400), - TypeError: partialclass(HTTPException, status_code=400), - } -) -@wrap_in_class( - Entry, - transform=lambda d: { - "id": UUID(d.pop("entry_id")), - **d, - }, - _kind="inserted", -) -@cozo_query -@increase_counter("create_entries") -@beartype -def create_entries( - *, - developer_id: UUID, - session_id: UUID, - data: list[CreateEntryRequest], - mark_session_as_updated: bool = True, -) -> tuple[list[str], dict]: - developer_id = str(developer_id) - session_id = str(session_id) - - data_dicts = [item.model_dump(mode="json") for item in data] - - for item in data_dicts: - item["content"] = content_to_json(item["content"] or []) - item["session_id"] = session_id - item["entry_id"] = item.pop("id", None) or str(uuid7()) - item["created_at"] = (item.get("created_at") or utcnow()).timestamp() - - cols, rows = cozo_process_mutate_data(data_dicts) - - # Construct a datalog query to insert the processed entries into the 'cozodb' database. - # Refer to the schema for the 'entries' relation in the README.md for column names and types. - create_query = f""" - ?[{cols}] <- $rows - - :insert entries {{ - {cols} - }} - - :returning - """ - - queries = [ - verify_developer_id_query(developer_id), - verify_developer_owns_resource_query( - developer_id, "sessions", session_id=session_id - ), - mark_session_updated_query(developer_id, session_id) - if mark_session_as_updated - else "", - create_query, - ] - - return (queries, {"rows": rows}) - - -@rewrap_exceptions( - { - QueryException: partialclass(HTTPException, status_code=400), - ValidationError: partialclass(HTTPException, status_code=400), - TypeError: partialclass(HTTPException, status_code=400), - } -) -@wrap_in_class(Relation, _kind="inserted") -@cozo_query -@beartype -def add_entry_relations( - *, - developer_id: UUID, - data: list[Relation], -) -> tuple[list[str], dict]: - developer_id = str(developer_id) - - data_dicts = [item.model_dump(mode="json") for item in data] - cols, rows = cozo_process_mutate_data(data_dicts) - - create_query = f""" - ?[{cols}] <- $rows - - :insert relations {{ - {cols} - }} - - :returning - """ - - queries = [ - verify_developer_id_query(developer_id), - create_query, - ] - - return (queries, {"rows": rows}) diff --git a/agents-api/agents_api/models/entry/delete_entries.py b/agents-api/agents_api/models/entry/delete_entries.py deleted file mode 100644 index c98b6c7d2..000000000 --- a/agents-api/agents_api/models/entry/delete_entries.py +++ /dev/null @@ -1,153 +0,0 @@ -from typing import Any, TypeVar -from uuid import UUID - -from beartype import beartype -from fastapi import HTTPException -from pycozo.client import QueryException -from pydantic import ValidationError - -from ...autogen.openapi_model import ResourceDeletedResponse -from ...common.utils.datetime import utcnow -from ..utils import ( - cozo_query, - mark_session_updated_query, - partialclass, - rewrap_exceptions, - verify_developer_id_query, - verify_developer_owns_resource_query, - wrap_in_class, -) - -ModelT = TypeVar("ModelT", bound=Any) -T = TypeVar("T") - - -@rewrap_exceptions( - { - QueryException: partialclass(HTTPException, status_code=400), - ValidationError: partialclass(HTTPException, status_code=400), - TypeError: partialclass(HTTPException, status_code=400), - IndexError: partialclass(HTTPException, status_code=404), - } -) -@wrap_in_class( - ResourceDeletedResponse, - one=True, - transform=lambda d: { - "id": UUID(d.pop("session_id")), # Only return session cleared - "deleted_at": utcnow(), - "jobs": [], - }, - _kind="deleted", -) -@cozo_query -@beartype -def delete_entries_for_session( - *, developer_id: UUID, session_id: UUID, mark_session_as_updated: bool = True -) -> tuple[list[str], dict]: - """ - 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. - """ - - delete_query = """ - input[session_id] <- [[ - to_uuid($session_id), - ]] - - ?[ - session_id, - entry_id, - source, - role, - ] := input[session_id], - *entries{ - session_id, - entry_id, - source, - role, - } - - :delete entries { - session_id, - entry_id, - source, - role, - } - - :returning - """ - - queries = [ - verify_developer_id_query(developer_id), - verify_developer_owns_resource_query( - developer_id, "sessions", session_id=session_id - ), - mark_session_updated_query(developer_id, session_id) - if mark_session_as_updated - else "", - delete_query, - ] - - return (queries, {"session_id": str(session_id)}) - - -@rewrap_exceptions( - { - QueryException: partialclass(HTTPException, status_code=400), - ValidationError: partialclass(HTTPException, status_code=400), - TypeError: partialclass(HTTPException, status_code=400), - } -) -@wrap_in_class( - ResourceDeletedResponse, - transform=lambda d: { - "id": UUID(d.pop("entry_id")), - "deleted_at": utcnow(), - "jobs": [], - }, -) -@cozo_query -@beartype -def delete_entries( - *, developer_id: UUID, session_id: UUID, entry_ids: list[UUID] -) -> tuple[list[str], dict]: - delete_query = """ - input[entry_id_str] <- $entry_ids - - ?[ - entry_id, - session_id, - source, - role, - ] := - input[entry_id_str], - entry_id = to_uuid(entry_id_str), - *entries { - session_id, - entry_id, - source, - role, - } - - :delete entries { - session_id, - entry_id, - source, - role, - } - - :returning - """ - - queries = [ - verify_developer_id_query(developer_id), - verify_developer_owns_resource_query( - developer_id, "sessions", session_id=session_id - ), - delete_query, - ] - - return (queries, {"entry_ids": [[str(id)] for id in entry_ids]}) diff --git a/agents-api/agents_api/models/entry/get_history.py b/agents-api/agents_api/models/entry/get_history.py deleted file mode 100644 index bb12b1c5b..000000000 --- a/agents-api/agents_api/models/entry/get_history.py +++ /dev/null @@ -1,150 +0,0 @@ -from typing import Any, TypeVar -from uuid import UUID - -from beartype import beartype -from fastapi import HTTPException -from pycozo.client import QueryException -from pydantic import ValidationError - -from ...autogen.openapi_model import History -from ...common.utils.cozo import uuid_int_list_to_uuid as fix_uuid -from ..utils import ( - cozo_query, - partialclass, - rewrap_exceptions, - verify_developer_id_query, - verify_developer_owns_resource_query, - wrap_in_class, -) - -ModelT = TypeVar("ModelT", bound=Any) -T = TypeVar("T") - - -@rewrap_exceptions( - { - QueryException: partialclass(HTTPException, status_code=400), - ValidationError: partialclass(HTTPException, status_code=400), - TypeError: partialclass(HTTPException, status_code=400), - } -) -@wrap_in_class( - History, - one=True, - transform=lambda d: { - "relations": [ - { - # This is needed because cozo has a bug: - # https://github.com/cozodb/cozo/issues/269 - "head": fix_uuid(r["head"]), - "relation": r["relation"], - "tail": fix_uuid(r["tail"]), - } - for r in d.pop("relations") - ], - # TODO: Remove this once we sort the entries in the cozo query - # Sort entries by created_at - "entries": sorted(d.pop("entries"), key=lambda entry: entry["created_at"]), - **d, - }, -) -@cozo_query -@beartype -def get_history( - *, - developer_id: UUID, - session_id: UUID, - allowed_sources: list[str] = ["api_request", "api_response"], -) -> tuple[list[str], dict]: - developer_id = str(developer_id) - session_id = str(session_id) - - history_query = """ - session_entries[collect(entry)] := - *entries { - session_id, - entry_id, - role, - name, - content, - source, - token_count, - tokenizer, - created_at, - tool_calls, - timestamp, - tool_call_id, - }, - source in $allowed_sources, - session_id = to_uuid($session_id), - entry = { - "session_id": session_id, - "id": entry_id, - "role": role, - "name": name, - "content": content, - "source": source, - "token_count": token_count, - "tokenizer": tokenizer, - "created_at": created_at, - "timestamp": timestamp, - "tool_calls": tool_calls, - "tool_call_id": tool_call_id, - } - - session_relations[unique(item)] := - session_id = to_uuid($session_id), - *entries { - session_id, - entry_id: head - }, - - *relations { - head, - relation, - tail - }, - - item = { - "head": head, - "relation": relation, - "tail": tail - } - - session_relations[unique(item)] := - session_id = to_uuid($session_id), - *entries { - session_id, - entry_id: tail - }, - - *relations { - head, - relation, - tail - }, - - item = { - "head": head, - "relation": relation, - "tail": tail - } - - ?[entries, relations, session_id, created_at] := - session_entries[entries], - session_relations[relations], - session_id = to_uuid($session_id), - created_at = now() - - :limit 1 - """ - - queries = [ - verify_developer_id_query(developer_id), - verify_developer_owns_resource_query( - developer_id, "sessions", session_id=session_id - ), - history_query, - ] - - return (queries, {"session_id": session_id, "allowed_sources": allowed_sources}) diff --git a/agents-api/agents_api/models/entry/list_entries.py b/agents-api/agents_api/models/entry/list_entries.py deleted file mode 100644 index d3081a9b0..000000000 --- a/agents-api/agents_api/models/entry/list_entries.py +++ /dev/null @@ -1,112 +0,0 @@ -from typing import Any, Literal, TypeVar -from uuid import UUID - -from beartype import beartype -from fastapi import HTTPException -from pycozo.client import QueryException -from pydantic import ValidationError - -from ...autogen.openapi_model import Entry -from ..utils import ( - cozo_query, - partialclass, - rewrap_exceptions, - verify_developer_id_query, - verify_developer_owns_resource_query, - wrap_in_class, -) - -ModelT = TypeVar("ModelT", bound=Any) -T = TypeVar("T") - - -@rewrap_exceptions( - { - QueryException: partialclass(HTTPException, status_code=400), - ValidationError: partialclass(HTTPException, status_code=400), - TypeError: partialclass(HTTPException, status_code=400), - } -) -@wrap_in_class(Entry) -@cozo_query -@beartype -def list_entries( - *, - developer_id: UUID, - session_id: UUID, - allowed_sources: list[str] = ["api_request", "api_response"], - limit: int = -1, - offset: int = 0, - sort_by: Literal["created_at", "timestamp"] = "timestamp", - direction: Literal["asc", "desc"] = "asc", - exclude_relations: list[str] = [], -) -> tuple[list[str], dict]: - """ - Constructs and executes a query to retrieve entries from the 'cozodb' database. - """ - - developer_id = str(developer_id) - session_id = str(session_id) - - sort = f"{'-' if direction == 'desc' else ''}{sort_by}" - - exclude_relations_query = """ - not *relations { - relation, - tail: id, - }, - relation in $exclude_relations, - # !is_in(relation, $exclude_relations), - """ - - list_query = f""" - ?[ - session_id, - id, - role, - name, - content, - source, - token_count, - tokenizer, - created_at, - timestamp, - ] := *entries {{ - session_id, - entry_id: id, - role, - name, - content, - source, - token_count, - tokenizer, - created_at, - timestamp, - }}, - {exclude_relations_query if exclude_relations else ''} - source in $allowed_sources, - session_id = to_uuid($session_id), - - :sort {sort} - """ - - if limit > 0: - list_query += f"\n:limit {limit}" - list_query += f"\n:offset {offset}" - - queries = [ - verify_developer_id_query(developer_id), - verify_developer_owns_resource_query( - developer_id, "sessions", session_id=session_id - ), - list_query, - ] - - return ( - queries, - { - "session_id": session_id, - "allowed_sources": allowed_sources, - "exclude_relations": exclude_relations, - }, - ) diff --git a/agents-api/agents_api/models/execution/__init__.py b/agents-api/agents_api/models/execution/__init__.py deleted file mode 100644 index abd3c7e47..000000000 --- a/agents-api/agents_api/models/execution/__init__.py +++ /dev/null @@ -1,15 +0,0 @@ -# ruff: noqa: F401, F403, F405 - -from .count_executions import count_executions -from .create_execution import create_execution -from .create_execution_transition import ( - create_execution_transition, - create_execution_transition_async, -) -from .get_execution import get_execution -from .get_execution_transition import get_execution_transition -from .list_execution_transitions import list_execution_transitions -from .list_executions import list_executions -from .lookup_temporal_data import lookup_temporal_data -from .prepare_execution_input import prepare_execution_input -from .update_execution import update_execution diff --git a/agents-api/agents_api/models/execution/constants.py b/agents-api/agents_api/models/execution/constants.py deleted file mode 100644 index 8d4568ba2..000000000 --- a/agents-api/agents_api/models/execution/constants.py +++ /dev/null @@ -1,5 +0,0 @@ -########## -# Consts # -########## - -OUTPUT_UNNEST_KEY = "$$e7w_unnest$$" diff --git a/agents-api/agents_api/models/execution/count_executions.py b/agents-api/agents_api/models/execution/count_executions.py deleted file mode 100644 index d130f0359..000000000 --- a/agents-api/agents_api/models/execution/count_executions.py +++ /dev/null @@ -1,61 +0,0 @@ -from typing import Any, TypeVar -from uuid import UUID - -from beartype import beartype -from fastapi import HTTPException -from pycozo.client import QueryException -from pydantic import ValidationError - -from ..utils import ( - cozo_query, - partialclass, - rewrap_exceptions, - verify_developer_id_query, - verify_developer_owns_resource_query, - wrap_in_class, -) - -ModelT = TypeVar("ModelT", bound=Any) -T = TypeVar("T") - - -@rewrap_exceptions( - { - QueryException: partialclass(HTTPException, status_code=400), - ValidationError: partialclass(HTTPException, status_code=400), - TypeError: partialclass(HTTPException, status_code=400), - } -) -@wrap_in_class(dict, one=True) -@cozo_query -@beartype -def count_executions( - *, - developer_id: UUID, - task_id: UUID, -) -> tuple[list[str], dict]: - count_query = """ - input[task_id] <- [[to_uuid($task_id)]] - - counter[count(id)] := - input[task_id], - *executions:task_id_execution_id_idx { - task_id, - execution_id: id, - } - - ?[count] := counter[count] - """ - - queries = [ - verify_developer_id_query(developer_id), - verify_developer_owns_resource_query( - developer_id, - "tasks", - task_id=task_id, - parents=[("agents", "agent_id")], - ), - count_query, - ] - - return (queries, {"task_id": str(task_id)}) diff --git a/agents-api/agents_api/models/execution/create_execution.py b/agents-api/agents_api/models/execution/create_execution.py deleted file mode 100644 index 59efd7ac3..000000000 --- a/agents-api/agents_api/models/execution/create_execution.py +++ /dev/null @@ -1,98 +0,0 @@ -from typing import Annotated, Any, TypeVar -from uuid import UUID - -from beartype import beartype -from fastapi import HTTPException -from pycozo.client import QueryException -from pydantic import ValidationError -from uuid_extensions import uuid7 - -from ...autogen.openapi_model import CreateExecutionRequest, Execution -from ...common.utils.cozo import cozo_process_mutate_data -from ...common.utils.types import dict_like -from ...metrics.counters import increase_counter -from ..utils import ( - cozo_query, - partialclass, - rewrap_exceptions, - verify_developer_id_query, - verify_developer_owns_resource_query, - wrap_in_class, -) -from .constants import OUTPUT_UNNEST_KEY - -ModelT = TypeVar("ModelT", bound=Any) -T = TypeVar("T") - - -@rewrap_exceptions( - { - QueryException: partialclass(HTTPException, status_code=400), - ValidationError: partialclass(HTTPException, status_code=400), - TypeError: partialclass(HTTPException, status_code=400), - } -) -@wrap_in_class( - Execution, - one=True, - transform=lambda d: {"id": d["execution_id"], **d}, - _kind="inserted", -) -@cozo_query -@increase_counter("create_execution") -@beartype -def create_execution( - *, - developer_id: UUID, - task_id: UUID, - execution_id: UUID | None = None, - data: Annotated[CreateExecutionRequest | dict, dict_like(CreateExecutionRequest)], -) -> tuple[list[str], dict]: - execution_id = execution_id or uuid7() - - developer_id = str(developer_id) - task_id = str(task_id) - execution_id = str(execution_id) - - if isinstance(data, CreateExecutionRequest): - data.metadata = data.metadata or {} - execution_data = data.model_dump() - else: - data["metadata"] = data.get("metadata", {}) - execution_data = data - - if execution_data["output"] is not None and not isinstance( - execution_data["output"], dict - ): - execution_data["output"] = {OUTPUT_UNNEST_KEY: execution_data["output"]} - - columns, values = cozo_process_mutate_data( - { - **execution_data, - "task_id": task_id, - "execution_id": execution_id, - } - ) - - insert_query = f""" - ?[{columns}] <- $values - - :insert executions {{ - {columns} - }} - - :returning - """ - - queries = [ - verify_developer_id_query(developer_id), - verify_developer_owns_resource_query( - developer_id, - "tasks", - task_id=task_id, - parents=[("agents", "agent_id")], - ), - insert_query, - ] - - return (queries, {"values": values}) diff --git a/agents-api/agents_api/models/execution/create_execution_transition.py b/agents-api/agents_api/models/execution/create_execution_transition.py deleted file mode 100644 index 5cbcb97bc..000000000 --- a/agents-api/agents_api/models/execution/create_execution_transition.py +++ /dev/null @@ -1,259 +0,0 @@ -from uuid import UUID - -from beartype import beartype -from fastapi import HTTPException -from pycozo.client import QueryException -from pydantic import ValidationError -from uuid_extensions import uuid7 - -from ...autogen.openapi_model import ( - CreateTransitionRequest, - Transition, - UpdateExecutionRequest, -) -from ...common.protocol.tasks import transition_to_execution_status, valid_transitions -from ...common.utils.cozo import cozo_process_mutate_data -from ...metrics.counters import increase_counter -from ..utils import ( - cozo_query, - cozo_query_async, - partialclass, - rewrap_exceptions, - verify_developer_id_query, - verify_developer_owns_resource_query, - wrap_in_class, -) -from .update_execution import update_execution - - -@beartype -def _create_execution_transition( - *, - developer_id: UUID, - execution_id: UUID, - data: CreateTransitionRequest, - # Only one of these needed - transition_id: UUID | None = None, - task_token: str | None = None, - # Only required for updating the execution status as well - update_execution_status: bool = False, - task_id: UUID | None = None, -) -> tuple[list[str | None], dict]: - transition_id = transition_id or uuid7() - data.metadata = data.metadata or {} - data.execution_id = execution_id - - # Dump to json - if isinstance(data.output, list): - data.output = [ - item.model_dump(mode="json") if hasattr(item, "model_dump") else item - for item in data.output - ] - - elif hasattr(data.output, "model_dump"): - data.output = data.output.model_dump(mode="json") - - # TODO: This is a hack to make sure the transition is valid - # (parallel transitions are whack, we should do something better) - is_parallel = data.current.workflow.startswith("PAR:") - - # Prepare the transition data - transition_data = data.model_dump(exclude_unset=True, exclude={"id"}) - - # Parse the current and next targets - validate_transition_targets(data) - current_target = transition_data.pop("current") - next_target = transition_data.pop("next") - - transition_data["current"] = (current_target["workflow"], current_target["step"]) - transition_data["next"] = next_target and ( - next_target["workflow"], - next_target["step"], - ) - - columns, transition_values = cozo_process_mutate_data( - { - **transition_data, - "task_token": str(task_token), # Converting to str for JSON serialisation - "transition_id": str(transition_id), - "execution_id": str(execution_id), - } - ) - - # Make sure the transition is valid - check_last_transition_query = f""" - valid_transition[start, end] <- [ - {", ".join(f'["{start}", "{end}"]' for start, ends in valid_transitions.items() for end in ends)} - ] - - last_transition_type[min_cost(type_created_at)] := - *transitions:execution_id_type_created_at_idx {{ - execution_id: to_uuid("{str(execution_id)}"), - type, - created_at, - }}, - type_created_at = [type, -created_at] - - matched[collect(last_type)] := - last_transition_type[data], - last_type_data = first(data), - last_type = if(is_null(last_type_data), "init", last_type_data), - valid_transition[last_type, $next_type] - - ?[valid] := - matched[prev_transitions], - found = length(prev_transitions), - valid = if($next_type == "init", found == 0, found > 0), - assert(valid, "Invalid transition"), - - :limit 1 - """ - - # Prepare the insert query - insert_query = f""" - ?[{columns}] <- $transition_values - - :insert transitions {{ - {columns} - }} - - :returning - """ - - validate_status_query, update_execution_query, update_execution_params = ( - "", - "", - {}, - ) - - if update_execution_status: - assert ( - task_id is not None - ), "task_id is required for updating the execution status" - - # Prepare the execution update query - [*_, validate_status_query, update_execution_query], update_execution_params = ( - update_execution.__wrapped__( - developer_id=developer_id, - task_id=task_id, - execution_id=execution_id, - data=UpdateExecutionRequest( - status=transition_to_execution_status[data.type] - ), - output=data.output if data.type != "error" else None, - error=str(data.output) - if data.type == "error" and data.output - else None, - ) - ) - - queries = [ - verify_developer_id_query(developer_id), - verify_developer_owns_resource_query( - developer_id, - "executions", - execution_id=execution_id, - parents=[("agents", "agent_id"), ("tasks", "task_id")], - ), - validate_status_query if not is_parallel else None, - update_execution_query if not is_parallel else None, - check_last_transition_query if not is_parallel else None, - insert_query, - ] - - return ( - queries, - { - "transition_values": transition_values, - "next_type": data.type, - "valid_transitions": valid_transitions, - **update_execution_params, - }, - ) - - -def validate_transition_targets(data: CreateTransitionRequest) -> None: - # Make sure the current/next targets are valid - match data.type: - case "finish_branch": - pass # TODO: Implement - case "finish" | "error" | "cancelled": - pass - - ### FIXME: HACK: Fix this and uncomment - - ### assert ( - ### data.next is None - ### ), "Next target must be None for finish/finish_branch/error/cancelled" - - case "init_branch" | "init": - assert ( - data.next and data.current.step == data.next.step == 0 - ), "Next target must be same as current for init_branch/init and step 0" - - case "wait": - assert data.next is None, "Next target must be None for wait" - - case "resume" | "step": - assert data.next is not None, "Next target must be provided for resume/step" - - if data.next.workflow == data.current.workflow: - assert ( - data.next.step > data.current.step - ), "Next step must be greater than current" - - case _: - raise ValueError(f"Invalid transition type: {data.type}") - - -create_execution_transition = rewrap_exceptions( - { - QueryException: partialclass(HTTPException, status_code=400), - ValidationError: partialclass(HTTPException, status_code=400), - TypeError: partialclass(HTTPException, status_code=400), - } -)( - wrap_in_class( - Transition, - transform=lambda d: { - **d, - "id": d["transition_id"], - "current": {"workflow": d["current"][0], "step": d["current"][1]}, - "next": d["next"] and {"workflow": d["next"][0], "step": d["next"][1]}, - }, - one=True, - _kind="inserted", - )( - cozo_query( - increase_counter("create_execution_transition")( - _create_execution_transition - ) - ) - ) -) - -create_execution_transition_async = rewrap_exceptions( - { - QueryException: partialclass(HTTPException, status_code=400), - ValidationError: partialclass(HTTPException, status_code=400), - TypeError: partialclass(HTTPException, status_code=400), - } -)( - wrap_in_class( - Transition, - transform=lambda d: { - **d, - "id": d["transition_id"], - "current": {"workflow": d["current"][0], "step": d["current"][1]}, - "next": d["next"] and {"workflow": d["next"][0], "step": d["next"][1]}, - }, - one=True, - _kind="inserted", - )( - cozo_query_async( - increase_counter("create_execution_transition_async")( - _create_execution_transition - ) - ) - ) -) diff --git a/agents-api/agents_api/models/execution/create_temporal_lookup.py b/agents-api/agents_api/models/execution/create_temporal_lookup.py deleted file mode 100644 index e47a505db..000000000 --- a/agents-api/agents_api/models/execution/create_temporal_lookup.py +++ /dev/null @@ -1,72 +0,0 @@ -from typing import TypeVar -from uuid import UUID - -from beartype import beartype -from fastapi import HTTPException -from pycozo.client import QueryException -from pydantic import ValidationError -from temporalio.client import WorkflowHandle - -from ...common.utils.cozo import cozo_process_mutate_data -from ...metrics.counters import increase_counter -from ..utils import ( - cozo_query, - partialclass, - rewrap_exceptions, - verify_developer_id_query, - verify_developer_owns_resource_query, -) - -T = TypeVar("T") - - -@rewrap_exceptions( - { - AssertionError: partialclass(HTTPException, status_code=404), - QueryException: partialclass(HTTPException, status_code=400), - ValidationError: partialclass(HTTPException, status_code=400), - TypeError: partialclass(HTTPException, status_code=400), - } -) -@cozo_query -@increase_counter("create_temporal_lookup") -@beartype -def create_temporal_lookup( - *, - developer_id: UUID, - execution_id: UUID, - workflow_handle: WorkflowHandle, -) -> tuple[list[str], dict]: - developer_id = str(developer_id) - execution_id = str(execution_id) - - temporal_columns, temporal_values = cozo_process_mutate_data( - { - "execution_id": execution_id, - "id": workflow_handle.id, - "run_id": workflow_handle.run_id, - "first_execution_run_id": workflow_handle.first_execution_run_id, - "result_run_id": workflow_handle.result_run_id, - } - ) - - temporal_executions_lookup_query = f""" - ?[{temporal_columns}] <- $temporal_values - - :insert temporal_executions_lookup {{ - {temporal_columns} - }} - """ - - queries = [ - verify_developer_id_query(developer_id), - verify_developer_owns_resource_query( - developer_id, - "executions", - execution_id=execution_id, - parents=[("agents", "agent_id"), ("tasks", "task_id")], - ), - temporal_executions_lookup_query, - ] - - return (queries, {"temporal_values": temporal_values}) diff --git a/agents-api/agents_api/models/execution/get_execution.py b/agents-api/agents_api/models/execution/get_execution.py deleted file mode 100644 index db0279b1f..000000000 --- a/agents-api/agents_api/models/execution/get_execution.py +++ /dev/null @@ -1,78 +0,0 @@ -from typing import Any, TypeVar -from uuid import UUID - -from beartype import beartype -from fastapi import HTTPException -from pycozo.client import QueryException -from pydantic import ValidationError - -from ...autogen.openapi_model import Execution -from ..utils import ( - cozo_query, - partialclass, - rewrap_exceptions, - wrap_in_class, -) -from .constants import OUTPUT_UNNEST_KEY - -ModelT = TypeVar("ModelT", bound=Any) -T = TypeVar("T") - - -@rewrap_exceptions( - { - AssertionError: partialclass(HTTPException, status_code=404), - QueryException: partialclass(HTTPException, status_code=400), - ValidationError: partialclass(HTTPException, status_code=400), - TypeError: partialclass(HTTPException, status_code=400), - } -) -@wrap_in_class( - Execution, - one=True, - transform=lambda d: { - **d, - "output": d["output"][OUTPUT_UNNEST_KEY] - if isinstance(d["output"], dict) and OUTPUT_UNNEST_KEY in d["output"] - else d["output"], - }, -) -@cozo_query -@beartype -def get_execution( - *, - execution_id: UUID, -) -> tuple[str, dict]: - # Executions are allowed direct GET access if they have execution_id - - # NOTE: Do not remove outer curly braces - query = """ - { - input[execution_id] <- [[to_uuid($execution_id)]] - - ?[id, task_id, status, input, output, error, session_id, metadata, created_at, updated_at] := - input[execution_id], - *executions { - task_id, - execution_id, - status, - input, - output, - error, - session_id, - metadata, - created_at, - updated_at, - }, - id = execution_id - - :limit 1 - } - """ - - return ( - query, - { - "execution_id": str(execution_id), - }, - ) diff --git a/agents-api/agents_api/models/execution/get_execution_transition.py b/agents-api/agents_api/models/execution/get_execution_transition.py deleted file mode 100644 index e2b38789a..000000000 --- a/agents-api/agents_api/models/execution/get_execution_transition.py +++ /dev/null @@ -1,80 +0,0 @@ -from typing import Any, TypeVar -from uuid import UUID - -from beartype import beartype -from fastapi import HTTPException -from pycozo.client import QueryException -from pydantic import ValidationError - -from ...autogen.openapi_model import Transition -from ..utils import ( - cozo_query, - partialclass, - rewrap_exceptions, - verify_developer_id_query, - wrap_in_class, -) - -ModelT = TypeVar("ModelT", bound=Any) -T = TypeVar("T") - - -@rewrap_exceptions( - { - QueryException: partialclass(HTTPException, status_code=400), - ValidationError: partialclass(HTTPException, status_code=400), - TypeError: partialclass(HTTPException, status_code=400), - AssertionError: partialclass(HTTPException, status_code=500), - } -) -@wrap_in_class(Transition, one=True) -@cozo_query -@beartype -def get_execution_transition( - *, - developer_id: UUID, - transition_id: UUID | None = None, - task_token: str | None = None, -) -> tuple[list[str], dict]: - # At least one of `transition_id` or `task_token` must be provided - assert ( - transition_id or task_token - ), "At least one of `transition_id` or `task_token` must be provided." - - if transition_id: - transition_id = str(transition_id) - filter = "id = to_uuid($transition_id)" - - else: - filter = "task_token = $task_token" - - get_query = """ - ?[id, type, current, next, output, metadata, updated_at, created_at] := - *transitions { - transition_id: id, - type, - current: current_tuple, - next: next_tuple, - output, - metadata, - updated_at, - created_at, - }, - current = {"workflow": current_tuple->0, "step": current_tuple->1}, - next = if( - is_null(next_tuple), - null, - {"workflow": next_tuple->0, "step": next_tuple->1}, - ) - - :limit 1 - """ - - get_query += filter - - queries = [ - verify_developer_id_query(developer_id), - get_query, - ] - - return (queries, {"task_token": task_token, "transition_id": transition_id}) 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 deleted file mode 100644 index 6c32c7692..000000000 --- a/agents-api/agents_api/models/execution/get_paused_execution_token.py +++ /dev/null @@ -1,77 +0,0 @@ -from typing import Any, TypeVar -from uuid import UUID - -from beartype import beartype -from fastapi import HTTPException -from pycozo.client import QueryException -from pydantic import ValidationError - -from ..utils import ( - cozo_query, - partialclass, - rewrap_exceptions, - verify_developer_id_query, - wrap_in_class, -) - -ModelT = TypeVar("ModelT", bound=Any) -T = TypeVar("T") - - -@rewrap_exceptions( - { - QueryException: partialclass(HTTPException, status_code=400), - ValidationError: partialclass(HTTPException, status_code=400), - TypeError: partialclass(HTTPException, status_code=400), - AssertionError: partialclass(HTTPException, status_code=500), - } -) -@wrap_in_class(dict, one=True) -@cozo_query -@beartype -def get_paused_execution_token( - *, - developer_id: UUID, - execution_id: UUID, -) -> tuple[list[str], dict]: - execution_id = str(execution_id) - - check_status_query = """ - ?[execution_id, status] := - *executions:execution_id_status_idx { - execution_id, - status, - }, - execution_id = to_uuid($execution_id), - status = "awaiting_input" - - :limit 1 - :assert some - """ - - get_query = """ - ?[task_token, created_at, metadata] := - execution_id = to_uuid($execution_id), - *executions { - execution_id, - }, - *transitions { - execution_id, - created_at, - task_token, - type, - metadata, - }, - type = "wait" - - :sort -created_at - :limit 1 - """ - - queries = [ - verify_developer_id_query(developer_id), - check_status_query, - get_query, - ] - - return (queries, {"execution_id": execution_id}) 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 deleted file mode 100644 index 8b1bf4604..000000000 --- a/agents-api/agents_api/models/execution/get_temporal_workflow_data.py +++ /dev/null @@ -1,57 +0,0 @@ -from typing import Any, TypeVar -from uuid import UUID - -from beartype import beartype -from fastapi import HTTPException -from pycozo.client import QueryException -from pydantic import ValidationError - -from ..utils import ( - cozo_query, - partialclass, - rewrap_exceptions, - wrap_in_class, -) - -ModelT = TypeVar("ModelT", bound=Any) -T = TypeVar("T") - - -@rewrap_exceptions( - { - QueryException: partialclass(HTTPException, status_code=400), - ValidationError: partialclass(HTTPException, status_code=400), - TypeError: partialclass(HTTPException, status_code=400), - } -) -@wrap_in_class(dict, one=True) -@cozo_query -@beartype -def get_temporal_workflow_data( - *, - execution_id: UUID, -) -> tuple[str, dict]: - # Executions are allowed direct GET access if they have execution_id - - query = """ - input[execution_id] <- [[to_uuid($execution_id)]] - - ?[id, run_id, result_run_id, first_execution_run_id] := - input[execution_id], - *temporal_executions_lookup { - execution_id, - id, - run_id, - result_run_id, - first_execution_run_id, - } - - :limit 1 - """ - - return ( - query, - { - "execution_id": str(execution_id), - }, - ) diff --git a/agents-api/agents_api/models/execution/list_execution_transitions.py b/agents-api/agents_api/models/execution/list_execution_transitions.py deleted file mode 100644 index 8931676f6..000000000 --- a/agents-api/agents_api/models/execution/list_execution_transitions.py +++ /dev/null @@ -1,69 +0,0 @@ -from typing import Any, Literal, TypeVar -from uuid import UUID - -from beartype import beartype -from fastapi import HTTPException -from pycozo.client import QueryException -from pydantic import ValidationError - -from ...autogen.openapi_model import Transition -from ..utils import cozo_query, partialclass, rewrap_exceptions, wrap_in_class - -ModelT = TypeVar("ModelT", bound=Any) -T = TypeVar("T") - - -@rewrap_exceptions( - { - QueryException: partialclass(HTTPException, status_code=400), - ValidationError: partialclass(HTTPException, status_code=400), - TypeError: partialclass(HTTPException, status_code=400), - } -) -@wrap_in_class(Transition) -@cozo_query -@beartype -def list_execution_transitions( - *, - execution_id: UUID, - limit: int = 100, - offset: int = 0, - sort_by: Literal["created_at", "updated_at"] = "created_at", - direction: Literal["asc", "desc"] = "desc", -) -> tuple[str, dict]: - sort = f"{'-' if direction == 'desc' else ''}{sort_by}" - - query = f""" - ?[id, execution_id, type, current, next, output, metadata, updated_at, created_at] := - *transitions {{ - execution_id, - transition_id: id, - type, - current: current_tuple, - next: next_tuple, - output, - metadata, - updated_at, - created_at, - }}, - current = {{"workflow": current_tuple->0, "step": current_tuple->1}}, - next = if( - is_null(next_tuple), - null, - {{"workflow": next_tuple->0, "step": next_tuple->1}}, - ), - execution_id = to_uuid($execution_id) - - :limit $limit - :offset $offset - :sort {sort} - """ - - return ( - query, - { - "execution_id": str(execution_id), - "limit": limit, - "offset": offset, - }, - ) diff --git a/agents-api/agents_api/models/execution/list_executions.py b/agents-api/agents_api/models/execution/list_executions.py deleted file mode 100644 index 64add074f..000000000 --- a/agents-api/agents_api/models/execution/list_executions.py +++ /dev/null @@ -1,95 +0,0 @@ -from typing import Any, Literal, TypeVar -from uuid import UUID - -from beartype import beartype -from fastapi import HTTPException -from pycozo.client import QueryException -from pydantic import ValidationError - -from ...autogen.openapi_model import Execution -from ..utils import ( - cozo_query, - partialclass, - rewrap_exceptions, - verify_developer_id_query, - verify_developer_owns_resource_query, - wrap_in_class, -) -from .constants import OUTPUT_UNNEST_KEY - -ModelT = TypeVar("ModelT", bound=Any) -T = TypeVar("T") - - -@rewrap_exceptions( - { - QueryException: partialclass(HTTPException, status_code=400), - ValidationError: partialclass(HTTPException, status_code=400), - TypeError: partialclass(HTTPException, status_code=400), - } -) -@wrap_in_class( - Execution, - transform=lambda d: { - **d, - "output": d["output"][OUTPUT_UNNEST_KEY] - if isinstance(d.get("output"), dict) and OUTPUT_UNNEST_KEY in d["output"] - else d.get("output"), - }, -) -@cozo_query -@beartype -def list_executions( - *, - developer_id: UUID, - task_id: UUID, - limit: int = 100, - offset: int = 0, - sort_by: Literal["created_at", "updated_at"] = "created_at", - direction: Literal["asc", "desc"] = "desc", -) -> tuple[list[str], dict]: - sort = f"{'-' if direction == 'desc' else ''}{sort_by}" - - list_query = f""" - input[task_id] <- [[to_uuid($task_id)]] - - ?[ - id, - task_id, - status, - input, - output, - session_id, - metadata, - created_at, - updated_at, - ] := input[task_id], - *executions {{ - task_id, - execution_id: id, - status, - input, - output, - session_id, - metadata, - created_at, - updated_at, - }} - - :limit {limit} - :offset {offset} - :sort {sort} - """ - - queries = [ - verify_developer_id_query(developer_id), - verify_developer_owns_resource_query( - developer_id, - "tasks", - task_id=task_id, - parents=[("agents", "agent_id")], - ), - list_query, - ] - - return (queries, {"task_id": str(task_id), "limit": limit, "offset": offset}) diff --git a/agents-api/agents_api/models/execution/lookup_temporal_data.py b/agents-api/agents_api/models/execution/lookup_temporal_data.py deleted file mode 100644 index 35f09129b..000000000 --- a/agents-api/agents_api/models/execution/lookup_temporal_data.py +++ /dev/null @@ -1,66 +0,0 @@ -from typing import Any, TypeVar -from uuid import UUID - -from beartype import beartype -from fastapi import HTTPException -from pycozo.client import QueryException -from pydantic import ValidationError - -from ..utils import ( - cozo_query, - partialclass, - rewrap_exceptions, - verify_developer_id_query, - verify_developer_owns_resource_query, - wrap_in_class, -) - -ModelT = TypeVar("ModelT", bound=Any) -T = TypeVar("T") - - -@rewrap_exceptions( - { - QueryException: partialclass(HTTPException, status_code=400), - ValidationError: partialclass(HTTPException, status_code=400), - TypeError: partialclass(HTTPException, status_code=400), - } -) -@wrap_in_class(dict, one=True) -@cozo_query -@beartype -def lookup_temporal_data( - *, - developer_id: UUID, - execution_id: UUID, -) -> tuple[list[str], dict]: - developer_id = str(developer_id) - execution_id = str(execution_id) - - temporal_query = """ - ?[id] := - execution_id = to_uuid($execution_id), - *temporal_executions_lookup { - id, execution_id, run_id, first_execution_run_id, result_run_id - } - - :limit 1 - """ - - queries = [ - verify_developer_id_query(developer_id), - verify_developer_owns_resource_query( - developer_id, - "executions", - execution_id=execution_id, - parents=[("agents", "agent_id"), ("tasks", "task_id")], - ), - temporal_query, - ] - - return ( - queries, - { - "execution_id": str(execution_id), - }, - ) diff --git a/agents-api/agents_api/models/execution/prepare_execution_input.py b/agents-api/agents_api/models/execution/prepare_execution_input.py deleted file mode 100644 index 5e841b9f2..000000000 --- a/agents-api/agents_api/models/execution/prepare_execution_input.py +++ /dev/null @@ -1,223 +0,0 @@ -from typing import Any, TypeVar -from uuid import UUID - -from beartype import beartype -from fastapi import HTTPException -from pycozo.client import QueryException -from pydantic import ValidationError - -from ...common.protocol.tasks import ExecutionInput -from ..agent.get_agent import get_agent -from ..task.get_task import get_task -from ..tools.list_tools import list_tools -from ..utils import ( - cozo_query, - fix_uuid_if_present, - make_cozo_json_query, - partialclass, - rewrap_exceptions, - verify_developer_id_query, - verify_developer_owns_resource_query, - wrap_in_class, -) -from .get_execution import get_execution - -ModelT = TypeVar("ModelT", bound=Any) -T = TypeVar("T") - - -@rewrap_exceptions( - { - QueryException: partialclass(HTTPException, status_code=400), - ValidationError: partialclass(HTTPException, status_code=400), - TypeError: partialclass(HTTPException, status_code=400), - AssertionError: lambda e: HTTPException( - status_code=429, - detail=str(e), - headers={"x-should-retry": "true"}, - ), - } -) -@wrap_in_class( - ExecutionInput, - one=True, - transform=lambda d: { - **d, - "task": { - "tools": [*map(fix_uuid_if_present, d["task"].pop("tools"))], - **d["task"], - }, - "agent_tools": [ - {tool["type"]: tool.pop("spec"), **tool} - for tool in map(fix_uuid_if_present, d["tools"]) - ], - }, -) -@cozo_query -@beartype -def prepare_execution_input( - *, - developer_id: UUID, - task_id: UUID, - execution_id: UUID, -) -> tuple[list[str], dict]: - execution_query, execution_params = get_execution.__wrapped__( - execution_id=execution_id - ) - - # Remove the outer curly braces - execution_query = execution_query.strip()[1:-1] - - execution_fields = ( - "id", - "task_id", - "status", - "input", - "session_id", - "metadata", - "created_at", - "updated_at", - ) - execution_query += f""" - :create _execution {{ - {", ".join(execution_fields)} - }} - """ - - task_query, task_params = get_task.__wrapped__( - developer_id=developer_id, task_id=task_id - ) - - # Remove the outer curly braces - task_query = task_query[-1].strip() - - task_fields = ( - "id", - "agent_id", - "name", - "description", - "input_schema", - "tools", - "inherit_tools", - "workflows", - "created_at", - "updated_at", - "metadata", - ) - task_query += f""" - :create _task {{ - {", ".join(task_fields)} - }} - """ - - dummy_agent_id = UUID(int=0) - - [*_, agent_query], agent_params = get_agent.__wrapped__( - developer_id=developer_id, - agent_id=dummy_agent_id, # We will replace this with value from the query - ) - agent_params.pop("agent_id") - agent_query = agent_query.replace( - "<- [[to_uuid($agent_id)]]", ":= *_task { agent_id }" - ) - - agent_fields = ( - "id", - "name", - "model", - "about", - "metadata", - "default_settings", - "instructions", - "created_at", - "updated_at", - ) - - agent_query += f""" - :create _agent {{ - {", ".join(agent_fields)} - }} - """ - - [*_, tools_query], tools_params = list_tools.__wrapped__( - developer_id=developer_id, - agent_id=dummy_agent_id, # We will replace this with value from the query - ) - tools_params.pop("agent_id") - tools_query = tools_query.replace( - "<- [[to_uuid($agent_id)]]", ":= *_task { agent_id }" - ) - - tools_fields = ( - "id", - "agent_id", - "name", - "type", - "spec", - "description", - "created_at", - "updated_at", - ) - tools_query += f""" - :create _tools {{ - {", ".join(tools_fields)} - }} - """ - - combine_query = f""" - collected_tools[collect(tool)] := - *_tools {{ {', '.join(tools_fields)} }}, - tool = {{ {make_cozo_json_query(tools_fields)} }} - - agent_json[agent] := - *_agent {{ {', '.join(agent_fields)} }}, - agent = {{ {make_cozo_json_query(agent_fields)} }} - - task_json[task] := - *_task {{ {', '.join(task_fields)} }}, - task = {{ {make_cozo_json_query(task_fields)} }} - - execution_json[execution] := - *_execution {{ {', '.join(execution_fields)} }}, - execution = {{ {make_cozo_json_query(execution_fields)} }} - - ?[developer_id, execution, task, agent, user, session, tools, arguments] := - developer_id = to_uuid($developer_id), - - agent_json[agent], - task_json[task], - execution_json[execution], - collected_tools[tools], - - # TODO: Enable these later - user = null, - session = null, - arguments = execution->"input" - - :limit 1 - """ - - queries = [ - verify_developer_id_query(developer_id), - verify_developer_owns_resource_query( - developer_id, "tasks", task_id=task_id, parents=[("agents", "agent_id")] - ), - execution_query, - task_query, - agent_query, - tools_query, - combine_query, - ] - - return ( - queries, - { - "developer_id": str(developer_id), - "task_id": str(task_id), - "execution_id": str(execution_id), - **execution_params, - **task_params, - **agent_params, - **tools_params, - }, - ) diff --git a/agents-api/agents_api/models/execution/update_execution.py b/agents-api/agents_api/models/execution/update_execution.py deleted file mode 100644 index f33368412..000000000 --- a/agents-api/agents_api/models/execution/update_execution.py +++ /dev/null @@ -1,130 +0,0 @@ -from typing import Any, TypeVar -from uuid import UUID - -from beartype import beartype -from fastapi import HTTPException -from pycozo.client import QueryException -from pydantic import ValidationError - -from ...autogen.openapi_model import ( - ResourceUpdatedResponse, - UpdateExecutionRequest, -) -from ...common.protocol.tasks import ( - valid_previous_statuses as valid_previous_statuses_map, -) -from ...common.utils.cozo import cozo_process_mutate_data -from ...metrics.counters import increase_counter -from ..utils import ( - cozo_query, - partialclass, - rewrap_exceptions, - verify_developer_id_query, - verify_developer_owns_resource_query, - wrap_in_class, -) -from .constants import OUTPUT_UNNEST_KEY - -ModelT = TypeVar("ModelT", bound=Any) -T = TypeVar("T") - - -@rewrap_exceptions( - { - QueryException: partialclass(HTTPException, status_code=400), - ValidationError: partialclass(HTTPException, status_code=400), - TypeError: partialclass(HTTPException, status_code=400), - } -) -@wrap_in_class( - ResourceUpdatedResponse, - one=True, - transform=lambda d: {"id": d["execution_id"], **d}, - _kind="inserted", -) -@cozo_query -@increase_counter("update_execution") -@beartype -def update_execution( - *, - developer_id: UUID, - task_id: UUID, - execution_id: UUID, - data: UpdateExecutionRequest, - output: dict | Any | None = None, - error: str | None = None, -) -> tuple[list[str], dict]: - developer_id = str(developer_id) - task_id = str(task_id) - execution_id = str(execution_id) - - valid_previous_statuses: list[str] | None = valid_previous_statuses_map.get( - data.status, None - ) - - execution_data: dict = data.model_dump(exclude_none=True) - - if output is not None and not isinstance(output, dict): - output: dict = {OUTPUT_UNNEST_KEY: output} - - columns, values = cozo_process_mutate_data( - { - **execution_data, - "task_id": task_id, - "execution_id": execution_id, - "output": output, - "error": error, - } - ) - - validate_status_query = """ - valid_status[count(status)] := - *executions { - status, - execution_id: to_uuid($execution_id), - task_id: to_uuid($task_id), - }, - status in $valid_previous_statuses - - ?[num] := - valid_status[num], - assert(num > 0, 'Invalid status') - - :limit 1 - """ - - update_query = f""" - input[{columns}] <- $values - ?[{columns}, updated_at] := - input[{columns}], - updated_at = now() - - :update executions {{ - updated_at, - {columns} - }} - - :returning - """ - - queries = [ - verify_developer_id_query(developer_id), - verify_developer_owns_resource_query( - developer_id, - "executions", - execution_id=execution_id, - parents=[("agents", "agent_id"), ("tasks", "task_id")], - ), - validate_status_query if valid_previous_statuses is not None else "", - update_query, - ] - - return ( - queries, - { - "values": values, - "valid_previous_statuses": valid_previous_statuses, - "execution_id": str(execution_id), - "task_id": task_id, - }, - ) diff --git a/agents-api/agents_api/models/files/__init__.py b/agents-api/agents_api/models/files/__init__.py deleted file mode 100644 index 444c0a6eb..000000000 --- a/agents-api/agents_api/models/files/__init__.py +++ /dev/null @@ -1,3 +0,0 @@ -from .create_file import create_file as create_file -from .delete_file import delete_file as delete_file -from .get_file import get_file as get_file diff --git a/agents-api/agents_api/models/files/create_file.py b/agents-api/agents_api/models/files/create_file.py deleted file mode 100644 index 58948038b..000000000 --- a/agents-api/agents_api/models/files/create_file.py +++ /dev/null @@ -1,122 +0,0 @@ -""" -This module contains the functionality for creating a new user in the CozoDB database. -It defines a query for inserting user data into the 'users' relation. -""" - -import base64 -import hashlib -from typing import Any, TypeVar -from uuid import UUID - -from beartype import beartype -from fastapi import HTTPException -from pycozo.client import QueryException -from pydantic import ValidationError -from uuid_extensions import uuid7 - -from ...autogen.openapi_model import CreateFileRequest, File -from ...metrics.counters import increase_counter -from ..utils import ( - cozo_query, - partialclass, - rewrap_exceptions, - verify_developer_id_query, - wrap_in_class, -) - -ModelT = TypeVar("ModelT", bound=Any) -T = TypeVar("T") - - -@rewrap_exceptions( - { - lambda e: isinstance(e, QueryException) - and "asserted to return some results, but returned none" - in str(e): lambda *_: HTTPException( - 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.", - ), - } -) -@wrap_in_class( - File, - one=True, - transform=lambda d: { - **d, - "id": d["file_id"], - "content": "DUMMY: NEED TO FETCH CONTENT FROM BLOB STORAGE", - }, - _kind="inserted", -) -@cozo_query -@increase_counter("create_file") -@beartype -def create_file( - *, - developer_id: UUID, - file_id: UUID | None = None, - data: CreateFileRequest, -) -> tuple[list[str], dict]: - """ - Constructs and executes a datalog query to create a new file in the CozoDB database. - - Parameters: - user_id (UUID): The unique identifier for the user. - developer_id (UUID): The unique identifier for the developer creating the file. - """ - - file_id = file_id or uuid7() - file_data = data.model_dump(exclude={"content"}) - - content_bytes = base64.b64decode(data.content) - size = len(content_bytes) - hash = hashlib.sha256(content_bytes).hexdigest() - - create_query = """ - # Then create the file - ?[file_id, developer_id, name, description, mime_type, size, hash] <- [ - [to_uuid($file_id), to_uuid($developer_id), $name, $description, $mime_type, $size, $hash] - ] - - :insert files { - developer_id, - file_id => - name, - description, - mime_type, - size, - hash, - } - :returning - """ - - queries = [ - verify_developer_id_query(developer_id), - create_query, - ] - - return ( - queries, - { - "file_id": str(file_id), - "developer_id": str(developer_id), - "size": size, - "hash": hash, - **file_data, - }, - ) diff --git a/agents-api/agents_api/models/files/delete_file.py b/agents-api/agents_api/models/files/delete_file.py deleted file mode 100644 index 053402e2f..000000000 --- a/agents-api/agents_api/models/files/delete_file.py +++ /dev/null @@ -1,97 +0,0 @@ -""" -This module contains the implementation of the delete_user_query function, which is responsible for deleting an user and its related default settings from the CozoDB database. -""" - -from typing import Any, TypeVar -from uuid import UUID - -from beartype import beartype -from fastapi import HTTPException -from pycozo.client import QueryException -from pydantic import ValidationError - -from ...autogen.openapi_model import ResourceDeletedResponse -from ...common.utils.datetime import utcnow -from ..utils import ( - cozo_query, - partialclass, - rewrap_exceptions, - verify_developer_id_query, - verify_developer_owns_resource_query, - wrap_in_class, -) - -ModelT = TypeVar("ModelT", bound=Any) -T = TypeVar("T") - - -@rewrap_exceptions( - { - lambda e: isinstance(e, QueryException) - and "Developer does not exist" in str(e): lambda *_: HTTPException( - detail="The specified developer does not exist.", - status_code=403, - ), - lambda e: isinstance(e, QueryException) - and "Developer does not own resource" - in e.resp["display"]: lambda *_: HTTPException( - 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.", - ), - } -) -@wrap_in_class( - ResourceDeletedResponse, - one=True, - transform=lambda d: { - "id": UUID(d.pop("file_id")), - "deleted_at": utcnow(), - "jobs": [], - }, - _kind="deleted", -) -@cozo_query -@beartype -def delete_file(*, developer_id: UUID, file_id: UUID) -> tuple[list[str], dict]: - """ - Constructs and returns a datalog query for deleting an file from the database. - - Parameters: - developer_id (UUID): The UUID of the developer owning the file. - file_id (UUID): The UUID of the file 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. - """ - - queries = [ - verify_developer_id_query(developer_id), - verify_developer_owns_resource_query(developer_id, "files", file_id=file_id), - """ - ?[file_id, developer_id] <- [[$file_id, $developer_id]] - - :delete files { - developer_id, - file_id - } - :returning - """, - ] - - return (queries, {"file_id": str(file_id), "developer_id": str(developer_id)}) diff --git a/agents-api/agents_api/models/files/get_file.py b/agents-api/agents_api/models/files/get_file.py deleted file mode 100644 index f3b85c2f7..000000000 --- a/agents-api/agents_api/models/files/get_file.py +++ /dev/null @@ -1,116 +0,0 @@ -from typing import Any, TypeVar -from uuid import UUID - -from beartype import beartype -from fastapi import HTTPException -from pycozo.client import QueryException -from pydantic import ValidationError - -from ...autogen.openapi_model import File -from ..utils import ( - cozo_query, - partialclass, - rewrap_exceptions, - verify_developer_id_query, - verify_developer_owns_resource_query, - wrap_in_class, -) - -ModelT = TypeVar("ModelT", bound=Any) -T = TypeVar("T") - - -@rewrap_exceptions( - { - lambda e: isinstance(e, QueryException) - and "Developer does not exist" in str(e): lambda *_: HTTPException( - detail="The specified developer does not exist.", - status_code=403, - ), - lambda e: isinstance(e, QueryException) - and "Developer does not own resource" - in e.resp["display"]: lambda *_: HTTPException( - 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.", - ), - } -) -@wrap_in_class( - File, - one=True, - transform=lambda d: { - **d, - "content": "DUMMY: NEED TO FETCH CONTENT FROM BLOB STORAGE", - }, -) -@cozo_query -@beartype -def get_file( - *, - developer_id: UUID, - file_id: UUID, -) -> tuple[list[str], dict]: - """ - Retrieves a file by their unique identifier. - - - Parameters: - developer_id (UUID): The unique identifier of the developer associated with the file. - file_id (UUID): The unique identifier of the file to retrieve. - - Returns: - File: The retrieved file. - """ - - # Convert UUIDs to strings for query compatibility. - file_id = str(file_id) - developer_id = str(developer_id) - - get_query = """ - input[developer_id, file_id] <- [[to_uuid($developer_id), to_uuid($file_id)]] - - ?[ - id, - name, - description, - mime_type, - size, - hash, - created_at, - ] := input[developer_id, id], - *files { - file_id: id, - developer_id, - name, - description, - mime_type, - size, - hash, - created_at, - } - - :limit 1 - """ - - queries = [ - verify_developer_id_query(developer_id), - verify_developer_owns_resource_query(developer_id, "files", file_id=file_id), - get_query, - ] - - return (queries, {"developer_id": developer_id, "file_id": file_id}) diff --git a/agents-api/agents_api/models/session/__init__.py b/agents-api/agents_api/models/session/__init__.py deleted file mode 100644 index bf80c9f4b..000000000 --- a/agents-api/agents_api/models/session/__init__.py +++ /dev/null @@ -1,22 +0,0 @@ -"""The session module is responsible for managing session data in the 'cozodb' database. It provides functionalities to create, retrieve, list, update, and delete session information. This module utilizes the `CozoClient` for database interactions, ensuring that sessions are uniquely identified and managed through UUIDs. - -Key functionalities include: -- Creating new sessions with specific metadata. -- Retrieving session information based on developer and session IDs. -- Listing all sessions with optional filters for pagination and metadata. -- Updating session data, including situation, summary, and metadata. -- Deleting sessions and their associated data from the database. - -This module plays a crucial role in the application by facilitating the management of session data, which is essential for tracking and analyzing user interactions and behaviors within the system.""" - -# ruff: noqa: F401, F403, F405 - -from .count_sessions import count_sessions -from .create_or_update_session import create_or_update_session -from .create_session import create_session -from .delete_session import delete_session -from .get_session import get_session -from .list_sessions import list_sessions -from .patch_session import patch_session -from .prepare_session_data import prepare_session_data -from .update_session import update_session diff --git a/agents-api/agents_api/models/session/count_sessions.py b/agents-api/agents_api/models/session/count_sessions.py deleted file mode 100644 index 3599cc2fb..000000000 --- a/agents-api/agents_api/models/session/count_sessions.py +++ /dev/null @@ -1,64 +0,0 @@ -"""This module contains functions for querying session data from the 'cozodb' database.""" - -from typing import Any, TypeVar -from uuid import UUID - -from beartype import beartype -from fastapi import HTTPException -from pycozo.client import QueryException -from pydantic import ValidationError - -from ..utils import ( - cozo_query, - partialclass, - rewrap_exceptions, - verify_developer_id_query, - wrap_in_class, -) - -ModelT = TypeVar("ModelT", bound=Any) -T = TypeVar("T") - - -@rewrap_exceptions( - { - QueryException: partialclass(HTTPException, status_code=400), - ValidationError: partialclass(HTTPException, status_code=400), - TypeError: partialclass(HTTPException, status_code=400), - } -) -@wrap_in_class(dict, one=True) -@cozo_query -@beartype -def count_sessions( - *, - developer_id: UUID, -) -> tuple[list[str], dict]: - """ - Counts sessions from the 'cozodb' database. - - Parameters: - developer_id (UUID): The developer's ID to filter sessions by. - """ - - count_query = """ - input[developer_id] <- [[ - to_uuid($developer_id), - ]] - - counter[count(id)] := - input[developer_id], - *sessions{ - developer_id, - session_id: id, - } - - ?[count] := counter[count] - """ - - queries = [ - verify_developer_id_query(developer_id), - count_query, - ] - - return (queries, {"developer_id": str(developer_id)}) diff --git a/agents-api/agents_api/models/session/create_or_update_session.py b/agents-api/agents_api/models/session/create_or_update_session.py deleted file mode 100644 index e34a63ca5..000000000 --- a/agents-api/agents_api/models/session/create_or_update_session.py +++ /dev/null @@ -1,158 +0,0 @@ -from typing import Any, TypeVar -from uuid import UUID - -from beartype import beartype -from fastapi import HTTPException -from pycozo.client import QueryException -from pydantic import ValidationError - -from ...autogen.openapi_model import ( - CreateOrUpdateSessionRequest, - ResourceUpdatedResponse, -) -from ...common.utils.cozo import cozo_process_mutate_data -from ...metrics.counters import increase_counter -from ..utils import ( - cozo_query, - partialclass, - rewrap_exceptions, - verify_developer_id_query, - verify_developer_owns_resource_query, - wrap_in_class, -) - -ModelT = TypeVar("ModelT", bound=Any) -T = TypeVar("T") - - -@rewrap_exceptions( - { - AssertionError: partialclass(HTTPException, status_code=400), - QueryException: partialclass(HTTPException, status_code=400), - ValidationError: partialclass(HTTPException, status_code=400), - TypeError: partialclass(HTTPException, status_code=400), - } -) -@wrap_in_class( - ResourceUpdatedResponse, - one=True, - transform=lambda d: { - "id": d["session_id"], - "updated_at": d.pop("updated_at")[0], - "jobs": [], - **d, - }, -) -@cozo_query -@increase_counter("create_or_update_session") -@beartype -def create_or_update_session( - *, - session_id: UUID, - developer_id: UUID, - data: CreateOrUpdateSessionRequest, -) -> tuple[list[str], dict]: - data.metadata = data.metadata or {} - session_data = data.model_dump(exclude={"auto_run_tools", "disable_cache"}) - - user = session_data.pop("user") - agent = session_data.pop("agent") - users = session_data.pop("users") - agents = session_data.pop("agents") - - # Only one of agent or agents should be provided. - if agent and agents: - raise ValueError("Only one of 'agent' or 'agents' should be provided.") - - agents = agents or ([agent] if agent else []) - assert len(agents) > 0, "At least one agent must be provided." - - # Users are zero or more, so we default to an empty list if not provided. - if not (user or users): - users = [] - - else: - users = users or [user] - - participants = [ - *[("user", str(user)) for user in users], - *[("agent", str(agent)) for agent in agents], - ] - - # Construct the datalog query for creating a new session and its lookup. - clear_lookup_query = """ - input[session_id] <- [[$session_id]] - ?[session_id, participant_id, participant_type] := - input[session_id], - *session_lookup { - session_id, - participant_type, - participant_id, - }, - - :delete session_lookup { - session_id, - participant_type, - participant_id, - } - """ - - lookup_query = """ - # This section creates a new session lookup to ensure uniqueness and manage session metadata. - session[session_id] <- [[$session_id]] - participants[participant_type, participant_id] <- $participants - ?[session_id, participant_id, participant_type] := - session[session_id], - participants[participant_type, participant_id], - - :put session_lookup { - session_id, - participant_id, - participant_type, - } - """ - - session_update_cols, session_update_vals = cozo_process_mutate_data( - {k: v for k, v in session_data.items() if v is not None} - ) - - # Construct the datalog query for creating or updating session information. - update_query = f""" - input[{session_update_cols}] <- $session_update_vals - ids[session_id, developer_id] <- [[to_uuid($session_id), to_uuid($developer_id)]] - - ?[{session_update_cols}, session_id, developer_id] := - input[{session_update_cols}], - ids[session_id, developer_id], - - :put sessions {{ - {session_update_cols}, session_id, developer_id - }} - - :returning - """ - - queries = [ - verify_developer_id_query(developer_id), - *[ - verify_developer_owns_resource_query( - developer_id, - f"{participant_type}s", - **{f"{participant_type}_id": participant_id}, - ) - for participant_type, participant_id in participants - ], - clear_lookup_query, - lookup_query, - update_query, - ] - - return ( - queries, - { - "session_update_vals": session_update_vals, - "session_id": str(session_id), - "developer_id": str(developer_id), - "participants": participants, - }, - ) diff --git a/agents-api/agents_api/models/session/create_session.py b/agents-api/agents_api/models/session/create_session.py deleted file mode 100644 index a08059961..000000000 --- a/agents-api/agents_api/models/session/create_session.py +++ /dev/null @@ -1,154 +0,0 @@ -""" -This module contains the functionality for creating a new session in the 'cozodb' database. -It constructs and executes a datalog query to insert session data. -""" - -from typing import Any, TypeVar -from uuid import UUID - -from beartype import beartype -from fastapi import HTTPException -from pycozo.client import QueryException -from pydantic import ValidationError -from uuid_extensions import uuid7 - -from ...autogen.openapi_model import CreateSessionRequest, Session -from ...metrics.counters import increase_counter -from ..utils import ( - cozo_query, - partialclass, - rewrap_exceptions, - verify_developer_id_query, - verify_developer_owns_resource_query, - wrap_in_class, -) - -ModelT = TypeVar("ModelT", bound=Any) -T = TypeVar("T") - - -@rewrap_exceptions( - { - AssertionError: partialclass(HTTPException, status_code=400), - QueryException: partialclass(HTTPException, status_code=400), - ValidationError: partialclass(HTTPException, status_code=400), - TypeError: partialclass(HTTPException, status_code=400), - } -) -@wrap_in_class( - Session, - one=True, - transform=lambda d: { - "id": UUID(d.pop("session_id")), - "updated_at": (d.pop("updated_at")[0]), - **d, - }, - _kind="inserted", -) -@cozo_query -@increase_counter("create_session") -@beartype -def create_session( - *, - developer_id: UUID, - session_id: UUID | None = None, - data: CreateSessionRequest, -) -> tuple[list[str], dict]: - """ - Constructs and executes a datalog query to create a new session in the database. - """ - - session_id = session_id or uuid7() - - data.metadata = data.metadata or {} - session_data = data.model_dump(exclude={"auto_run_tools", "disable_cache"}) - - user = session_data.pop("user") - agent = session_data.pop("agent") - users = session_data.pop("users") - agents = session_data.pop("agents") - - # Only one of agent or agents should be provided. - if agent and agents: - raise ValueError("Only one of 'agent' or 'agents' should be provided.") - - agents = agents or ([agent] if agent else []) - assert len(agents) > 0, "At least one agent must be provided." - - # Users are zero or more, so we default to an empty list if not provided. - if not (user or users): - users = [] - - else: - users = users or [user] - - participants = [ - *[("user", str(user)) for user in users], - *[("agent", str(agent)) for agent in agents], - ] - - # Construct the datalog query for creating a new session and its lookup. - lookup_query = """ - # This section creates a new session lookup to ensure uniqueness and manage session metadata. - session[session_id] <- [[$session_id]] - participants[participant_type, participant_id] <- $participants - ?[session_id, participant_id, participant_type] := - session[session_id], - participants[participant_type, participant_id], - - :insert session_lookup { - session_id, - participant_id, - participant_type, - } - """ - - create_query = """ - # Insert the new session data into the 'session' table with the specified columns. - ?[session_id, developer_id, situation, metadata, render_templates, token_budget, context_overflow] <- [[ - $session_id, - $developer_id, - $situation, - $metadata, - $render_templates, - $token_budget, - $context_overflow, - ]] - - :insert sessions { - developer_id, - session_id, - situation, - metadata, - render_templates, - token_budget, - context_overflow, - } - # Specify the data to return after the query execution, typically the newly created session's ID. - :returning - """ - - queries = [ - verify_developer_id_query(developer_id), - *[ - verify_developer_owns_resource_query( - developer_id, - f"{participant_type}s", - **{f"{participant_type}_id": participant_id}, - ) - for participant_type, participant_id in participants - ], - lookup_query, - create_query, - ] - - # Execute the constructed query with the provided parameters and return the result. - return ( - queries, - { - "session_id": str(session_id), - "developer_id": str(developer_id), - "participants": participants, - **session_data, - }, - ) diff --git a/agents-api/agents_api/models/session/delete_session.py b/agents-api/agents_api/models/session/delete_session.py deleted file mode 100644 index 81f8e1f7c..000000000 --- a/agents-api/agents_api/models/session/delete_session.py +++ /dev/null @@ -1,125 +0,0 @@ -"""This module contains the implementation for deleting sessions from the 'cozodb' database using datalog queries.""" - -from typing import Any, TypeVar -from uuid import UUID - -from beartype import beartype -from fastapi import HTTPException -from pycozo.client import QueryException -from pydantic import ValidationError - -from ...autogen.openapi_model import ResourceDeletedResponse -from ...common.utils.datetime import utcnow -from ..utils import ( - cozo_query, - partialclass, - rewrap_exceptions, - verify_developer_id_query, - verify_developer_owns_resource_query, - wrap_in_class, -) - -ModelT = TypeVar("ModelT", bound=Any) -T = TypeVar("T") - - -@rewrap_exceptions( - { - QueryException: partialclass(HTTPException, status_code=400), - ValidationError: partialclass(HTTPException, status_code=400), - TypeError: partialclass(HTTPException, status_code=400), - } -) -@wrap_in_class( - ResourceDeletedResponse, - one=True, - transform=lambda d: { - "id": UUID(d.pop("session_id")), - "deleted_at": utcnow(), - "jobs": [], - }, - _kind="deleted", -) -@cozo_query -@beartype -def delete_session( - *, - developer_id: UUID, - session_id: UUID, -) -> tuple[list[str], dict]: - """ - 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. - - Returns: - ResourceDeletedResponse: The response indicating the deletion of the session. - """ - session_id = str(session_id) - developer_id = str(developer_id) - - # Constructs and executes a datalog query to delete the specified session and its associated data based on the session_id and developer_id. - delete_lookup_query = """ - # Convert session_id to UUID format - input[session_id] <- [[ - to_uuid($session_id), - ]] - - # Select sessions based on the session_id provided - ?[ - session_id, - participant_id, - participant_type, - ] := - input[session_id], - *session_lookup{ - session_id, - participant_id, - participant_type, - } - - # Delete entries from session_lookup table matching the criteria - :delete session_lookup { - session_id, - participant_id, - participant_type, - } - """ - - delete_query = """ - # Convert developer_id and session_id to UUID format - input[developer_id, session_id] <- [[ - to_uuid($developer_id), - to_uuid($session_id), - ]] - - # Select sessions based on the developer_id and session_id provided - ?[developer_id, session_id, updated_at] := - input[developer_id, session_id], - *sessions { - developer_id, - session_id, - updated_at, - } - - # Delete entries from sessions table matching the criteria - :delete sessions { - developer_id, - session_id, - updated_at, - } - :returning - """ - - queries = [ - verify_developer_id_query(developer_id), - verify_developer_owns_resource_query( - developer_id, "sessions", session_id=session_id - ), - delete_lookup_query, - delete_query, - ] - - return (queries, {"session_id": session_id, "developer_id": developer_id}) diff --git a/agents-api/agents_api/models/session/get_session.py b/agents-api/agents_api/models/session/get_session.py deleted file mode 100644 index f99f2524c..000000000 --- a/agents-api/agents_api/models/session/get_session.py +++ /dev/null @@ -1,116 +0,0 @@ -from typing import Any, TypeVar -from uuid import UUID - -from beartype import beartype -from fastapi import HTTPException -from pycozo.client import QueryException -from pydantic import ValidationError - -from ...common.protocol.sessions import make_session -from ..utils import ( - cozo_query, - partialclass, - rewrap_exceptions, - verify_developer_id_query, - verify_developer_owns_resource_query, - wrap_in_class, -) - -ModelT = TypeVar("ModelT", bound=Any) -T = TypeVar("T") - - -@rewrap_exceptions( - { - QueryException: partialclass(HTTPException, status_code=400), - ValidationError: partialclass(HTTPException, status_code=400), - TypeError: partialclass(HTTPException, status_code=400), - } -) -@wrap_in_class(make_session, one=True) -@cozo_query -@beartype -def get_session( - *, - developer_id: UUID, - session_id: UUID, -) -> tuple[list[str], dict]: - """ - Constructs and executes a datalog query to retrieve session information from the 'cozodb' database. - - Parameters: - developer_id (UUID): The developer's unique identifier. - session_id (UUID): The session's unique identifier. - """ - session_id = str(session_id) - developer_id = str(developer_id) - - # This query retrieves session information by using `input` to pass parameters, - get_query = """ - input[developer_id, session_id] <- [[ - to_uuid($developer_id), - to_uuid($session_id), - ]] - - participants[collect(participant_id), participant_type] := - input[_, session_id], - *session_lookup{ - session_id, - participant_id, - participant_type, - } - - # We have to do this dance because users can be zero or more - users_p[users] := - participants[users, "user"] - - users_p[users] := - not participants[_, "user"], - users = [] - - ?[ - agents, - users, - id, - situation, - summary, - updated_at, - created_at, - metadata, - render_templates, - token_budget, - context_overflow, - recall_options, - forward_tool_calls, - ] := input[developer_id, id], - users_p[users], - participants[agents, "agent"], - *sessions{ - developer_id, - session_id: id, - situation, - summary, - created_at, - updated_at: validity, - metadata, - render_templates, - token_budget, - context_overflow, - recall_options, - forward_tool_calls, - @ "END" - }, - updated_at = to_int(validity) - - :limit 1 - """ - - queries = [ - verify_developer_id_query(developer_id), - verify_developer_owns_resource_query( - developer_id, "sessions", session_id=session_id - ), - get_query, - ] - - return (queries, {"session_id": session_id, "developer_id": developer_id}) diff --git a/agents-api/agents_api/models/session/list_sessions.py b/agents-api/agents_api/models/session/list_sessions.py deleted file mode 100644 index 4adb84a6c..000000000 --- a/agents-api/agents_api/models/session/list_sessions.py +++ /dev/null @@ -1,131 +0,0 @@ -"""This module contains functions for querying session data from the 'cozodb' database.""" - -from typing import Any, Literal, TypeVar -from uuid import UUID - -from beartype import beartype -from fastapi import HTTPException -from pycozo.client import QueryException -from pydantic import ValidationError - -from ...common.protocol.sessions import make_session -from ...common.utils import json -from ..utils import ( - cozo_query, - partialclass, - rewrap_exceptions, - verify_developer_id_query, - wrap_in_class, -) - -ModelT = TypeVar("ModelT", bound=Any) -T = TypeVar("T") - - -@rewrap_exceptions( - { - QueryException: partialclass(HTTPException, status_code=400), - ValidationError: partialclass(HTTPException, status_code=400), - TypeError: partialclass(HTTPException, status_code=400), - } -) -@wrap_in_class(make_session) -@cozo_query -@beartype -def list_sessions( - *, - developer_id: UUID, - limit: int = 100, - offset: int = 0, - sort_by: Literal["created_at", "updated_at"] = "created_at", - 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. - - Parameters: - developer_id (UUID): The developer's ID to filter sessions by. - limit (int): The maximum number of sessions to return. - offset (int): The offset from which to start listing sessions. - metadata_filter (dict[str, Any]): A dictionary of metadata fields to filter sessions by. - """ - metadata_filter_str = ", ".join( - [ - f"metadata->{json.dumps(k)} == {json.dumps(v)}" - for k, v in metadata_filter.items() - ] - ) - - sort = f"{'-' if direction == 'desc' else ''}{sort_by}" - - list_query = f""" - input[developer_id] <- [[ - to_uuid($developer_id), - ]] - - participants[collect(participant_id), participant_type, session_id] := - *session_lookup{{ - session_id, - participant_id, - participant_type, - }} - - # We have to do this dance because users can be zero or more - users_p[users, session_id] := - participants[users, "user", session_id] - - users_p[users, session_id] := - not participants[_, "user", session_id], - users = [] - - ?[ - agents, - users, - id, - situation, - summary, - updated_at, - created_at, - metadata, - token_budget, - context_overflow, - recall_options, - forward_tool_calls, - ] := - input[developer_id], - *sessions{{ - developer_id, - session_id: id, - situation, - summary, - created_at, - updated_at: validity, - metadata, - token_budget, - context_overflow, - recall_options, - forward_tool_calls, - @ "END" - }}, - users_p[users, id], - participants[agents, "agent", id], - updated_at = to_int(validity), - {metadata_filter_str} - - :limit $limit - :offset $offset - :sort {sort} - """ - - # Datalog query to retrieve agent information based on filters, sorted by creation date in descending order. - queries = [ - verify_developer_id_query(developer_id), - list_query, - ] - - # Execute the datalog query and return the results as a pandas DataFrame. - return ( - queries, - {"developer_id": str(developer_id), "limit": limit, "offset": offset}, - ) diff --git a/agents-api/agents_api/models/session/patch_session.py b/agents-api/agents_api/models/session/patch_session.py deleted file mode 100644 index 4a119a684..000000000 --- a/agents-api/agents_api/models/session/patch_session.py +++ /dev/null @@ -1,127 +0,0 @@ -"""This module contains functions for patching session data in the 'cozodb' database using datalog queries.""" - -from typing import Any, List, TypeVar -from uuid import UUID - -from beartype import beartype -from fastapi import HTTPException -from pycozo.client import QueryException -from pydantic import ValidationError - -from ...autogen.openapi_model import PatchSessionRequest, ResourceUpdatedResponse -from ...common.utils.cozo import cozo_process_mutate_data -from ..utils import ( - cozo_query, - partialclass, - rewrap_exceptions, - verify_developer_id_query, - verify_developer_owns_resource_query, - wrap_in_class, -) - -ModelT = TypeVar("ModelT", bound=Any) -T = TypeVar("T") - -_fields: List[str] = [ - "situation", - "summary", - "created_at", - "session_id", - "developer_id", -] - - -# TODO: Add support for updating `render_templates` field - - -@rewrap_exceptions( - { - QueryException: partialclass(HTTPException, status_code=400), - ValidationError: partialclass(HTTPException, status_code=400), - TypeError: partialclass(HTTPException, status_code=400), - } -) -@wrap_in_class( - ResourceUpdatedResponse, - one=True, - transform=lambda d: { - "id": d["session_id"], - "updated_at": d.pop("updated_at")[0], - "jobs": [], - **d, - }, - _kind="inserted", -) -@cozo_query -@beartype -def patch_session( - *, - session_id: UUID, - developer_id: UUID, - data: PatchSessionRequest, -) -> tuple[list[str], dict]: - """ - 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. - """ - - update_data = data.model_dump(exclude_unset=True) - metadata = update_data.pop("metadata", {}) or {} - - session_update_cols, session_update_vals = cozo_process_mutate_data( - {k: v for k, v in update_data.items() if v is not None} - ) - - # Prepare lists of columns for the query. - session_update_cols_lst = session_update_cols.split(",") - all_fields_lst = list(set(session_update_cols_lst).union(set(_fields))) - all_fields = ", ".join(all_fields_lst) - rest_fields = ", ".join( - list( - set(all_fields_lst) - - set([k for k, v in update_data.items() if v is not None]) - ) - ) - - # Construct the datalog query for updating session information. - update_query = f""" - input[{session_update_cols}] <- $session_update_vals - ids[session_id, developer_id] <- [[to_uuid($session_id), to_uuid($developer_id)]] - - ?[{all_fields}, metadata, updated_at] := - input[{session_update_cols}], - ids[session_id, developer_id], - *sessions{{ - {rest_fields}, metadata: md, @ "END" - }}, - updated_at = 'ASSERT', - metadata = concat(md, $metadata), - - :put sessions {{ - {all_fields}, metadata, updated_at - }} - - :returning - """ - - queries = [ - verify_developer_id_query(developer_id), - verify_developer_owns_resource_query( - developer_id, "sessions", session_id=session_id - ), - update_query, - ] - - return ( - queries, - { - "session_update_vals": session_update_vals, - "session_id": str(session_id), - "developer_id": str(developer_id), - "metadata": metadata, - }, - ) diff --git a/agents-api/agents_api/models/session/prepare_session_data.py b/agents-api/agents_api/models/session/prepare_session_data.py deleted file mode 100644 index 83ee0c219..000000000 --- a/agents-api/agents_api/models/session/prepare_session_data.py +++ /dev/null @@ -1,235 +0,0 @@ -from typing import Any, TypeVar -from uuid import UUID - -from beartype import beartype -from fastapi import HTTPException -from pycozo.client import QueryException -from pydantic import ValidationError - -from ...common.protocol.sessions import SessionData, make_session -from ..utils import ( - cozo_query, - partialclass, - rewrap_exceptions, - verify_developer_id_query, - verify_developer_owns_resource_query, - wrap_in_class, -) - -ModelT = TypeVar("ModelT", bound=Any) -T = TypeVar("T") - - -@rewrap_exceptions( - { - QueryException: partialclass(HTTPException, status_code=400), - ValidationError: partialclass(HTTPException, status_code=400), - TypeError: partialclass(HTTPException, status_code=400), - } -) -@wrap_in_class( - SessionData, - one=True, - transform=lambda d: { - "session": make_session( - **d["session"], - agents=[a["id"] for a in d["agents"]], - users=[u["id"] for u in d["users"]], - ), - }, -) -@cozo_query -@beartype -def prepare_session_data( - *, - developer_id: UUID, - session_id: UUID, -) -> tuple[list[str], dict]: - """Constructs and executes a datalog query to retrieve session data from the 'cozodb' database. - - Parameters: - developer_id (UUID): The developer's unique identifier. - session_id (UUID): The session's unique identifier. - """ - session_id = str(session_id) - developer_id = str(developer_id) - - # This query retrieves session information by using `input` to pass parameters, - get_query = """ - input[session_id, developer_id] <- [[ - to_uuid($session_id), - to_uuid($developer_id), - ]] - - participants[collect(participant_id), participant_type] := - input[session_id, developer_id], - *session_lookup{ - session_id, - participant_id, - participant_type, - } - - agents[agent_ids] := participants[agent_ids, "agent"] - - # We have to do this dance because users can be zero or more - users[user_ids] := - participants[user_ids, "user"] - - users[user_ids] := - not participants[_, "user"], - user_ids = [] - - settings_data[agent_id, settings] := - *agent_default_settings { - agent_id, - frequency_penalty, - presence_penalty, - length_penalty, - repetition_penalty, - top_p, - temperature, - min_p, - preset, - }, - settings = { - "frequency_penalty": frequency_penalty, - "presence_penalty": presence_penalty, - "length_penalty": length_penalty, - "repetition_penalty": repetition_penalty, - "top_p": top_p, - "temperature": temperature, - "min_p": min_p, - "preset": preset, - } - - agent_data[collect(record)] := - input[session_id, developer_id], - agents[agent_ids], - agent_id in agent_ids, - *agents{ - developer_id, - agent_id, - model, - name, - about, - created_at, - updated_at, - metadata, - instructions, - }, - settings_data[agent_id, default_settings], - record = { - "id": agent_id, - "name": name, - "model": model, - "about": about, - "created_at": created_at, - "updated_at": updated_at, - "metadata": metadata, - "default_settings": default_settings, - "instructions": instructions, - } - - # Version where we don't have default settings - agent_data[collect(record)] := - input[session_id, developer_id], - agents[agent_ids], - agent_id in agent_ids, - *agents{ - developer_id, - agent_id, - model, - name, - about, - created_at, - updated_at, - metadata, - instructions, - }, - not settings_data[agent_id, _], - record = { - "id": agent_id, - "name": name, - "model": model, - "about": about, - "created_at": created_at, - "updated_at": updated_at, - "metadata": metadata, - "default_settings": {}, - "instructions": instructions, - } - - user_data[collect(record)] := - input[session_id, developer_id], - users[user_ids], - user_id in user_ids, - *users{ - developer_id, - user_id, - name, - about, - created_at, - updated_at, - metadata, - }, - record = { - "id": user_id, - "name": name, - "about": about, - "created_at": created_at, - "updated_at": updated_at, - "metadata": metadata, - } - - session_data[record] := - input[session_id, developer_id], - *sessions{ - developer_id, - session_id, - situation, - summary, - created_at, - updated_at: validity, - metadata, - render_templates, - token_budget, - context_overflow, - @ "END" - }, - updated_at = to_int(validity), - record = { - "id": session_id, - "situation": situation, - "summary": summary, - "created_at": created_at, - "updated_at": updated_at, - "metadata": metadata, - "render_templates": render_templates, - "token_budget": token_budget, - "context_overflow": context_overflow, - } - - ?[ - agents, - users, - session, - ] := - session_data[session], - user_data[users], - agent_data[agents] - - :limit 1 - """ - - queries = [ - verify_developer_id_query(developer_id), - verify_developer_owns_resource_query( - developer_id, "sessions", session_id=session_id - ), - get_query, - ] - - return ( - queries, - {"developer_id": developer_id, "session_id": session_id}, - ) diff --git a/agents-api/agents_api/models/session/update_session.py b/agents-api/agents_api/models/session/update_session.py deleted file mode 100644 index cc8b61f16..000000000 --- a/agents-api/agents_api/models/session/update_session.py +++ /dev/null @@ -1,127 +0,0 @@ -from typing import Any, List, TypeVar -from uuid import UUID - -from beartype import beartype -from fastapi import HTTPException -from pycozo.client import QueryException -from pydantic import ValidationError - -from ...autogen.openapi_model import ResourceUpdatedResponse, UpdateSessionRequest -from ...common.utils.cozo import cozo_process_mutate_data -from ...metrics.counters import increase_counter -from ..utils import ( - cozo_query, - partialclass, - rewrap_exceptions, - verify_developer_id_query, - verify_developer_owns_resource_query, - wrap_in_class, -) - -ModelT = TypeVar("ModelT", bound=Any) -T = TypeVar("T") - -_fields: List[str] = [ - "situation", - "summary", - "metadata", - "created_at", - "session_id", - "developer_id", -] - -# TODO: Add support for updating `render_templates` field - - -@rewrap_exceptions( - { - QueryException: partialclass(HTTPException, status_code=400), - ValidationError: partialclass(HTTPException, status_code=400), - TypeError: partialclass(HTTPException, status_code=400), - } -) -@wrap_in_class( - ResourceUpdatedResponse, - one=True, - transform=lambda d: { - "id": d["session_id"], - "updated_at": d.pop("updated_at")[0], - "jobs": [], - **d, - }, - _kind="inserted", -) -@cozo_query -@increase_counter("update_session") -@beartype -def update_session( - *, - session_id: UUID, - developer_id: UUID, - data: UpdateSessionRequest, -) -> tuple[list[str], dict]: - """ - Updates a session with the provided data. - - Parameters: - session_id (UUID): The unique identifier of the session to update. - developer_id (UUID): The unique identifier of the developer associated with the session. - data (UpdateSessionRequest): The data to update the session with. - - Returns: - ResourceUpdatedResponse: The updated session. - """ - - update_data = data.model_dump(exclude_unset=True) - - session_update_cols, session_update_vals = cozo_process_mutate_data( - {k: v for k, v in update_data.items() if v is not None} - ) - - # Prepare lists of columns for the query. - session_update_cols_lst = session_update_cols.split(",") - all_fields_lst = list(set(session_update_cols_lst).union(set(_fields))) - all_fields = ", ".join(all_fields_lst) - rest_fields = ", ".join( - list( - set(all_fields_lst) - - set([k for k, v in update_data.items() if v is not None]) - ) - ) - - # Construct the datalog query for updating session information. - update_query = f""" - input[{session_update_cols}] <- $session_update_vals - ids[session_id, developer_id] <- [[to_uuid($session_id), to_uuid($developer_id)]] - - ?[{all_fields}, updated_at] := - input[{session_update_cols}], - ids[session_id, developer_id], - *sessions{{ - {rest_fields}, @ "END" - }}, - updated_at = 'ASSERT' - - :put sessions {{ - {all_fields}, updated_at - }} - - :returning - """ - - queries = [ - verify_developer_id_query(developer_id), - verify_developer_owns_resource_query( - developer_id, "sessions", session_id=session_id - ), - update_query, - ] - - return ( - queries, - { - "session_update_vals": session_update_vals, - "session_id": str(session_id), - "developer_id": str(developer_id), - }, - ) diff --git a/agents-api/agents_api/models/task/__init__.py b/agents-api/agents_api/models/task/__init__.py deleted file mode 100644 index 2eaff3ab3..000000000 --- a/agents-api/agents_api/models/task/__init__.py +++ /dev/null @@ -1,9 +0,0 @@ -# ruff: noqa: F401, F403, F405 - -from .create_or_update_task import create_or_update_task -from .create_task import create_task -from .delete_task import delete_task -from .get_task import get_task -from .list_tasks import list_tasks -from .patch_task import patch_task -from .update_task import update_task diff --git a/agents-api/agents_api/models/task/create_or_update_task.py b/agents-api/agents_api/models/task/create_or_update_task.py deleted file mode 100644 index 1f615a3ad..000000000 --- a/agents-api/agents_api/models/task/create_or_update_task.py +++ /dev/null @@ -1,109 +0,0 @@ -""" -This module contains the functionality for creating a new Task in the 'cozodb` database. -It constructs and executes a datalog query to insert Task data. -""" - -from typing import Any, TypeVar -from uuid import UUID - -from beartype import beartype -from fastapi import HTTPException -from pycozo.client import QueryException -from pydantic import ValidationError - -from ...autogen.openapi_model import ( - CreateOrUpdateTaskRequest, - ResourceUpdatedResponse, -) -from ...common.protocol.tasks import task_to_spec -from ...common.utils.cozo import cozo_process_mutate_data -from ...common.utils.datetime import utcnow -from ...metrics.counters import increase_counter -from ..utils import ( - cozo_query, - partialclass, - rewrap_exceptions, - verify_developer_id_query, - verify_developer_owns_resource_query, - wrap_in_class, -) - -ModelT = TypeVar("ModelT", bound=Any) -T = TypeVar("T") - - -@rewrap_exceptions( - { - QueryException: partialclass(HTTPException, status_code=400), - ValidationError: partialclass(HTTPException, status_code=400), - TypeError: partialclass(HTTPException, status_code=400), - } -) -@wrap_in_class( - ResourceUpdatedResponse, - one=True, - transform=lambda d: { - "id": d["task_id"], - "jobs": [], - "updated_at": d["updated_at_ms"][0] / 1000, - **d, - }, -) -@cozo_query -@increase_counter("create_or_update_task") -@beartype -def create_or_update_task( - *, - developer_id: UUID, - agent_id: UUID, - task_id: UUID, - data: CreateOrUpdateTaskRequest, -) -> tuple[list[str], dict]: - developer_id = str(developer_id) - agent_id = str(agent_id) - task_id = str(task_id) - - data.metadata = data.metadata or {} - data.input_schema = data.input_schema or {} - - task_data = task_to_spec(data).model_dump(exclude_none=True, mode="json") - task_data.pop("task_id", None) - task_data["created_at"] = utcnow().timestamp() - - columns, values = cozo_process_mutate_data(task_data) - - update_query = f""" - input[{columns}] <- $values - ids[agent_id, task_id] := - agent_id = to_uuid($agent_id), - task_id = to_uuid($task_id) - - ?[updated_at_ms, agent_id, task_id, {columns}] := - ids[agent_id, task_id], - input[{columns}], - updated_at_ms = [floor(now() * 1000), true] - - :put tasks {{ - agent_id, - task_id, - updated_at_ms, - {columns}, - }} - - :returning - """ - - queries = [ - verify_developer_id_query(developer_id), - verify_developer_owns_resource_query(developer_id, "agents", agent_id=agent_id), - update_query, - ] - - return ( - queries, - { - "values": values, - "agent_id": agent_id, - "task_id": task_id, - }, - ) diff --git a/agents-api/agents_api/models/task/create_task.py b/agents-api/agents_api/models/task/create_task.py deleted file mode 100644 index 7cd1e8f4a..000000000 --- a/agents-api/agents_api/models/task/create_task.py +++ /dev/null @@ -1,118 +0,0 @@ -""" -This module contains the functionality for creating a new Task in the 'cozodb` database. -It constructs and executes a datalog query to insert Task data. -""" - -from typing import Any, TypeVar -from uuid import UUID - -from beartype import beartype -from fastapi import HTTPException -from pycozo.client import QueryException -from pydantic import ValidationError -from uuid_extensions import uuid7 - -from ...autogen.openapi_model import ( - CreateTaskRequest, - ResourceCreatedResponse, -) -from ...common.protocol.tasks import task_to_spec -from ...common.utils.cozo import cozo_process_mutate_data -from ...metrics.counters import increase_counter -from ..utils import ( - cozo_query, - partialclass, - rewrap_exceptions, - verify_developer_id_query, - verify_developer_owns_resource_query, - wrap_in_class, -) - -ModelT = TypeVar("ModelT", bound=Any) -T = TypeVar("T") - - -@rewrap_exceptions( - { - QueryException: partialclass(HTTPException, status_code=400), - ValidationError: partialclass(HTTPException, status_code=400), - TypeError: partialclass(HTTPException, status_code=400), - } -) -@wrap_in_class( - ResourceCreatedResponse, - one=True, - transform=lambda d: { - "id": d["task_id"], - "jobs": [], - "created_at": d["created_at"], - **d, - }, -) -@cozo_query -@increase_counter("create_task") -@beartype -def create_task( - *, - developer_id: UUID, - agent_id: UUID, - task_id: UUID | None = None, - data: CreateTaskRequest, -) -> tuple[list[str], dict]: - """ - Creates a new task. - - Parameters: - developer_id (UUID): The unique identifier of the developer associated with the task. - agent_id (UUID): The unique identifier of the agent associated with the task. - task_id (UUID | None): The unique identifier of the task. If not provided, a new UUID will be generated. - data (CreateTaskRequest): The data to create the task with. - - Returns: - ResourceCreatedResponse: The created task. - """ - - data.metadata = data.metadata or {} - data.input_schema = data.input_schema or {} - - task_id = task_id or uuid7() - task_spec = task_to_spec(data) - - # Prepares the update data by filtering out None values and adding user_id and developer_id. - columns, values = cozo_process_mutate_data( - { - **task_spec.model_dump(exclude_none=True, mode="json"), - "task_id": str(task_id), - "agent_id": str(agent_id), - } - ) - - create_query = f""" - input[{columns}] <- $values - ?[{columns}, updated_at_ms, created_at] := - input[{columns}], - updated_at_ms = [floor(now() * 1000), true], - created_at = now(), - - :insert tasks {{ - {columns}, - updated_at_ms, - created_at, - }} - - :returning - """ - - queries = [ - verify_developer_id_query(developer_id), - verify_developer_owns_resource_query(developer_id, "agents", agent_id=agent_id), - create_query, - ] - - return ( - queries, - { - "agent_id": str(agent_id), - "values": values, - }, - ) diff --git a/agents-api/agents_api/models/task/delete_task.py b/agents-api/agents_api/models/task/delete_task.py deleted file mode 100644 index 10c377a25..000000000 --- a/agents-api/agents_api/models/task/delete_task.py +++ /dev/null @@ -1,91 +0,0 @@ -from typing import Any, TypeVar -from uuid import UUID - -from beartype import beartype -from fastapi import HTTPException -from pycozo.client import QueryException -from pydantic import ValidationError - -from ...autogen.openapi_model import ResourceDeletedResponse -from ...common.utils.datetime import utcnow -from ..utils import ( - cozo_query, - partialclass, - rewrap_exceptions, - verify_developer_id_query, - verify_developer_owns_resource_query, - wrap_in_class, -) - -ModelT = TypeVar("ModelT", bound=Any) -T = TypeVar("T") - - -@rewrap_exceptions( - { - QueryException: partialclass(HTTPException, status_code=400), - ValidationError: partialclass(HTTPException, status_code=400), - TypeError: partialclass(HTTPException, status_code=400), - } -) -@wrap_in_class( - ResourceDeletedResponse, - one=True, - transform=lambda d: { - "id": UUID(d.pop("task_id")), - "jobs": [], - "deleted_at": utcnow(), - **d, - }, - _kind="deleted", -) -@cozo_query -@beartype -def delete_task( - *, - developer_id: UUID, - agent_id: UUID, - task_id: UUID, -) -> tuple[list[str], dict]: - """ - Deletes a task. - - Parameters: - developer_id (UUID): The unique identifier of the developer associated with the task. - agent_id (UUID): The unique identifier of the agent associated with the task. - task_id (UUID): The unique identifier of the task to delete. - - Returns: - ResourceDeletedResponse: The deleted task. - """ - - delete_query = """ - input[agent_id, task_id] <- [[ - to_uuid($agent_id), - to_uuid($task_id), - ]] - - ?[agent_id, task_id, updated_at_ms] := - input[agent_id, task_id], - *tasks{ - agent_id, - task_id, - updated_at_ms, - } - - :delete tasks { - agent_id, - task_id, - updated_at_ms, - } - - :returning - """ - - queries = [ - verify_developer_id_query(developer_id), - verify_developer_owns_resource_query(developer_id, "agents", agent_id=agent_id), - delete_query, - ] - - return (queries, {"agent_id": str(agent_id), "task_id": str(task_id)}) diff --git a/agents-api/agents_api/models/task/get_task.py b/agents-api/agents_api/models/task/get_task.py deleted file mode 100644 index 460fdc38b..000000000 --- a/agents-api/agents_api/models/task/get_task.py +++ /dev/null @@ -1,120 +0,0 @@ -from typing import Any, TypeVar -from uuid import UUID - -from beartype import beartype -from fastapi import HTTPException -from pycozo.client import QueryException -from pydantic import ValidationError - -from ...common.protocol.tasks import spec_to_task -from ..utils import ( - cozo_query, - partialclass, - rewrap_exceptions, - verify_developer_id_query, - verify_developer_owns_resource_query, - wrap_in_class, -) - -ModelT = TypeVar("ModelT", bound=Any) -T = TypeVar("T") - - -@rewrap_exceptions( - { - QueryException: partialclass(HTTPException, status_code=400), - ValidationError: partialclass(HTTPException, status_code=400), - TypeError: partialclass(HTTPException, status_code=400), - } -) -@wrap_in_class(spec_to_task, one=True) -@cozo_query -@beartype -def get_task( - *, - developer_id: UUID, - task_id: UUID, -) -> tuple[list[str], dict]: - """ - Retrieves a task by its unique identifier. - - Parameters: - developer_id (UUID): The unique identifier of the developer associated with the task. - task_id (UUID): The unique identifier of the task to retrieve. - - Returns: - Task | CreateTaskRequest: The retrieved task. - """ - - get_query = """ - input[task_id] <- [[to_uuid($task_id)]] - - task_data[ - task_id, - agent_id, - name, - description, - input_schema, - tools, - inherit_tools, - workflows, - created_at, - updated_at, - metadata, - ] := - input[task_id], - *tasks { - agent_id, - task_id, - updated_at_ms, - name, - description, - input_schema, - tools, - inherit_tools, - workflows, - created_at, - metadata, - @ 'END' - }, - updated_at = to_int(updated_at_ms) / 1000 - - ?[ - id, - agent_id, - name, - description, - input_schema, - tools, - inherit_tools, - workflows, - created_at, - updated_at, - metadata, - ] := - task_data[ - id, - agent_id, - name, - description, - input_schema, - tools, - inherit_tools, - workflows, - created_at, - updated_at, - metadata, - ] - - :limit 1 - """ - - queries = [ - verify_developer_id_query(developer_id), - verify_developer_owns_resource_query( - developer_id, "tasks", task_id=task_id, parents=[("agents", "agent_id")] - ), - get_query, - ] - - return (queries, {"task_id": str(task_id)}) diff --git a/agents-api/agents_api/models/task/list_tasks.py b/agents-api/agents_api/models/task/list_tasks.py deleted file mode 100644 index d873e817e..000000000 --- a/agents-api/agents_api/models/task/list_tasks.py +++ /dev/null @@ -1,130 +0,0 @@ -from typing import Any, Literal, TypeVar -from uuid import UUID - -from beartype import beartype -from fastapi import HTTPException -from pycozo.client import QueryException -from pydantic import ValidationError - -from ...common.protocol.tasks import spec_to_task -from ..utils import ( - cozo_query, - partialclass, - rewrap_exceptions, - verify_developer_id_query, - verify_developer_owns_resource_query, - wrap_in_class, -) - -ModelT = TypeVar("ModelT", bound=Any) -T = TypeVar("T") - - -@rewrap_exceptions( - { - QueryException: partialclass(HTTPException, status_code=400), - ValidationError: partialclass(HTTPException, status_code=400), - TypeError: partialclass(HTTPException, status_code=400), - } -) -@wrap_in_class(spec_to_task) -@cozo_query -@beartype -def list_tasks( - *, - developer_id: UUID, - agent_id: UUID, - limit: int = 100, - offset: int = 0, - sort_by: Literal["created_at", "updated_at"] = "created_at", - direction: Literal["asc", "desc"] = "desc", -) -> tuple[list[str], dict]: - """ - Lists tasks for a given agent. - - Parameters: - developer_id (UUID): The unique identifier of the developer associated with the tasks. - agent_id (UUID): The unique identifier of the agent associated with the tasks. - limit (int): The maximum number of tasks to return. - offset (int): The number of tasks to skip before returning the results. - sort_by (Literal["created_at", "updated_at"]): The field to sort the tasks by. - direction (Literal["asc", "desc"]): The direction to sort the tasks in. - - Returns: - Task[] | CreateTaskRequest[]: The list of tasks. - """ - - sort = f"{'-' if direction == 'desc' else ''}{sort_by}" - - list_query = f""" - input[agent_id] <- [[to_uuid($agent_id)]] - - task_data[ - task_id, - agent_id, - name, - description, - input_schema, - tools, - inherit_tools, - workflows, - created_at, - updated_at, - metadata, - ] := - input[agent_id], - *tasks {{ - agent_id, - task_id, - updated_at_ms, - name, - description, - input_schema, - tools, - inherit_tools, - workflows, - created_at, - metadata, - @ 'END' - }}, - updated_at = to_int(updated_at_ms) / 1000 - - ?[ - task_id, - agent_id, - name, - description, - input_schema, - tools, - inherit_tools, - workflows, - created_at, - updated_at, - metadata, - ] := - task_data[ - task_id, - agent_id, - name, - description, - input_schema, - tools, - inherit_tools, - workflows, - created_at, - updated_at, - metadata, - ] - - :limit $limit - :offset $offset - :sort {sort} - """ - - queries = [ - verify_developer_id_query(developer_id), - verify_developer_owns_resource_query(developer_id, "agents", agent_id=agent_id), - list_query, - ] - - return (queries, {"agent_id": str(agent_id), "limit": limit, "offset": offset}) diff --git a/agents-api/agents_api/models/task/patch_task.py b/agents-api/agents_api/models/task/patch_task.py deleted file mode 100644 index 178b9daa3..000000000 --- a/agents-api/agents_api/models/task/patch_task.py +++ /dev/null @@ -1,133 +0,0 @@ -""" -This module contains the functionality for creating a new Task in the 'cozodb` database. -It constructs and executes a datalog query to insert Task data. -""" - -from typing import Any, TypeVar -from uuid import UUID - -from beartype import beartype -from fastapi import HTTPException -from pycozo.client import QueryException -from pydantic import ValidationError - -from ...autogen.openapi_model import PatchTaskRequest, ResourceUpdatedResponse, TaskSpec -from ...common.protocol.tasks import task_to_spec -from ...common.utils.cozo import cozo_process_mutate_data -from ...metrics.counters import increase_counter -from ..utils import ( - cozo_query, - partialclass, - rewrap_exceptions, - verify_developer_id_query, - verify_developer_owns_resource_query, - wrap_in_class, -) - -ModelT = TypeVar("ModelT", bound=Any) -T = TypeVar("T") - - -@rewrap_exceptions( - { - QueryException: partialclass(HTTPException, status_code=400), - ValidationError: partialclass(HTTPException, status_code=400), - TypeError: partialclass(HTTPException, status_code=400), - } -) -@wrap_in_class( - ResourceUpdatedResponse, - one=True, - transform=lambda d: { - "id": d["task_id"], - "jobs": [], - "updated_at": d["updated_at_ms"][0] / 1000, - **d, - }, - _kind="inserted", -) -@cozo_query -@increase_counter("patch_task") -@beartype -def patch_task( - *, - developer_id: UUID, - agent_id: UUID, - task_id: UUID, - data: PatchTaskRequest, -) -> tuple[list[str], dict]: - developer_id = str(developer_id) - agent_id = str(agent_id) - task_id = str(task_id) - - data.input_schema = data.input_schema or {} - task_data = task_to_spec(data, exclude_none=True, exclude_unset=True).model_dump( - exclude_none=True, exclude_unset=True - ) - task_data.pop("task_id", None) - - assert len(task_data), "No data provided to update task" - metadata = task_data.pop("metadata", {}) - columns, values = cozo_process_mutate_data(task_data) - - all_columns = list(TaskSpec.model_fields.keys()) - all_columns.remove("id") - all_columns.remove("main") - - missing_columns = ( - set(all_columns) - - set(columns.split(",")) - - {"metadata", "created_at", "updated_at"} - ) - missing_columns_str = ",".join(missing_columns) - - patch_query = f""" - input[{columns}] <- $values - ids[agent_id, task_id] := - agent_id = to_uuid($agent_id), - task_id = to_uuid($task_id) - - original[created_at, metadata, {missing_columns_str}] := - ids[agent_id, task_id], - *tasks{{ - agent_id, - task_id, - created_at, - metadata, - {missing_columns_str}, - }} - - ?[created_at, updated_at_ms, agent_id, task_id, metadata, {columns}, {missing_columns_str}] := - ids[agent_id, task_id], - input[{columns}], - original[created_at, _metadata, {missing_columns_str}], - updated_at_ms = [floor(now() * 1000), true], - metadata = _metadata ++ $metadata - - :put tasks {{ - agent_id, - task_id, - created_at, - updated_at_ms, - metadata, - {columns}, {missing_columns_str} - }} - - :returning - """ - - queries = [ - verify_developer_id_query(developer_id), - verify_developer_owns_resource_query(developer_id, "agents", agent_id=agent_id), - patch_query, - ] - - return ( - queries, - { - "values": values, - "agent_id": agent_id, - "task_id": task_id, - "metadata": metadata, - }, - ) diff --git a/agents-api/agents_api/models/task/update_task.py b/agents-api/agents_api/models/task/update_task.py deleted file mode 100644 index cd98d85d5..000000000 --- a/agents-api/agents_api/models/task/update_task.py +++ /dev/null @@ -1,129 +0,0 @@ -""" -This module contains the functionality for creating a new Task in the 'cozodb` database. -It constructs and executes a datalog query to insert Task data. -""" - -from typing import Any, TypeVar -from uuid import UUID - -from beartype import beartype -from fastapi import HTTPException -from pycozo.client import QueryException -from pydantic import ValidationError - -from ...autogen.openapi_model import ResourceUpdatedResponse, UpdateTaskRequest -from ...common.protocol.tasks import task_to_spec -from ...common.utils.cozo import cozo_process_mutate_data -from ...metrics.counters import increase_counter -from ..utils import ( - cozo_query, - partialclass, - rewrap_exceptions, - verify_developer_id_query, - verify_developer_owns_resource_query, - wrap_in_class, -) - -ModelT = TypeVar("ModelT", bound=Any) -T = TypeVar("T") - - -@rewrap_exceptions( - { - QueryException: partialclass(HTTPException, status_code=400), - ValidationError: partialclass(HTTPException, status_code=400), - TypeError: partialclass(HTTPException, status_code=400), - } -) -@wrap_in_class( - ResourceUpdatedResponse, - one=True, - transform=lambda d: { - "id": d["task_id"], - "jobs": [], - "updated_at": d["updated_at_ms"][0] / 1000, - **d, - }, -) -@cozo_query -@increase_counter("update_task") -@beartype -def update_task( - *, - developer_id: UUID, - agent_id: UUID, - task_id: UUID, - data: UpdateTaskRequest, -) -> tuple[list[str], dict]: - """ - Updates a task. - - Parameters: - developer_id (UUID): The unique identifier of the developer associated with the task. - agent_id (UUID): The unique identifier of the agent associated with the task. - task_id (UUID): The unique identifier of the task to update. - data (UpdateTaskRequest): The data to update the task with. - - Returns: - ResourceUpdatedResponse: The updated task. - """ - - developer_id = str(developer_id) - agent_id = str(agent_id) - task_id = str(task_id) - - data.metadata = data.metadata or {} - data.input_schema = data.input_schema or {} - - task_data = task_to_spec(data, exclude_none=True, exclude_unset=True).model_dump( - exclude_none=True, exclude_unset=True - ) - task_data.pop("task_id", None) - - columns, values = cozo_process_mutate_data(task_data) - - update_query = f""" - input[{columns}] <- $values - ids[agent_id, task_id] := - agent_id = to_uuid($agent_id), - task_id = to_uuid($task_id) - - original[created_at] := - ids[agent_id, task_id], - *tasks{{ - agent_id, - task_id, - created_at, - }} - - ?[created_at, updated_at_ms, agent_id, task_id, {columns}] := - ids[agent_id, task_id], - input[{columns}], - original[created_at], - updated_at_ms = [floor(now() * 1000), true] - - :put tasks {{ - agent_id, - task_id, - created_at, - updated_at_ms, - {columns}, - }} - - :returning - """ - - queries = [ - verify_developer_id_query(developer_id), - verify_developer_owns_resource_query(developer_id, "agents", agent_id=agent_id), - update_query, - ] - - return ( - queries, - { - "values": values, - "agent_id": agent_id, - "task_id": task_id, - }, - ) diff --git a/agents-api/agents_api/models/user/__init__.py b/agents-api/agents_api/models/user/__init__.py deleted file mode 100644 index 5ae76865f..000000000 --- a/agents-api/agents_api/models/user/__init__.py +++ /dev/null @@ -1,18 +0,0 @@ -""" -This module is responsible for managing user data in the CozoDB database. It provides functionalities to create, retrieve, list, and update user information. - -Functions: -- create_user_query: Creates a new user in the CozoDB database, accepting parameters such as user ID, developer ID, name, about, and optional metadata. -- get_user_query: Retrieves a user's information from the CozoDB database by their user ID and developer ID. -- list_users_query: Lists users associated with a specific developer, with support for pagination and metadata-based filtering. -- patch_user_query: Updates a user's information in the CozoDB database, allowing for changes to fields such as name, about, and metadata. -""" - -# ruff: noqa: F401, F403, F405 - -from .create_or_update_user import create_or_update_user -from .create_user import create_user -from .get_user import get_user -from .list_users import list_users -from .patch_user import patch_user -from .update_user import update_user 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 deleted file mode 100644 index 3e9b1f3a6..000000000 --- a/agents-api/agents_api/models/user/create_or_update_user.py +++ /dev/null @@ -1,125 +0,0 @@ -""" -This module contains the functionality for creating users in the CozoDB database. -It includes functions to construct and execute datalog queries for inserting new user records. -""" - -from typing import Any, TypeVar -from uuid import UUID - -from beartype import beartype -from fastapi import HTTPException -from pycozo.client import QueryException -from pydantic import ValidationError - -from ...autogen.openapi_model import CreateOrUpdateUserRequest, User -from ...metrics.counters import increase_counter -from ..utils import ( - cozo_query, - partialclass, - rewrap_exceptions, - verify_developer_id_query, - wrap_in_class, -) - -ModelT = TypeVar("ModelT", bound=Any) -T = TypeVar("T") - - -@rewrap_exceptions( - { - 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(User, one=True, transform=lambda d: {"id": UUID(d.pop("user_id")), **d}) -@cozo_query -@increase_counter("create_or_update_user") -@beartype -def create_or_update_user( - *, - developer_id: UUID, - user_id: UUID, - data: CreateOrUpdateUserRequest, -) -> tuple[list[str], dict]: - """ - 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. - - Returns: - User: The newly created user record. - """ - - # Extract the user data from the payload - data.metadata = data.metadata or {} - - user_data = data.model_dump() - - # Create the user - # Construct a query to insert the new user record into the users table - user_query = """ - input[user_id, developer_id, name, about, metadata, updated_at] <- [ - [$user_id, $developer_id, $name, $about, $metadata, now()] - ] - - ?[user_id, developer_id, name, about, metadata, created_at, updated_at] := - input[_user_id, developer_id, name, about, metadata, updated_at], - *users{ - developer_id, - user_id, - created_at, - }, - user_id = to_uuid(_user_id), - - ?[user_id, developer_id, name, about, metadata, created_at, updated_at] := - input[_user_id, developer_id, name, about, metadata, updated_at], - not *users{ - developer_id, - user_id, - }, created_at = now(), - user_id = to_uuid(_user_id), - - :put users { - developer_id, - user_id => - name, - about, - metadata, - created_at, - updated_at, - } - :returning - """ - - queries = [ - verify_developer_id_query(developer_id), - user_query, - ] - - return ( - queries, - { - "user_id": str(user_id), - "developer_id": str(developer_id), - **user_data, - }, - ) diff --git a/agents-api/agents_api/models/user/create_user.py b/agents-api/agents_api/models/user/create_user.py deleted file mode 100644 index 62975a6d4..000000000 --- a/agents-api/agents_api/models/user/create_user.py +++ /dev/null @@ -1,116 +0,0 @@ -""" -This module contains the functionality for creating a new user in the CozoDB database. -It defines a query for inserting user data into the 'users' relation. -""" - -from typing import Any, TypeVar -from uuid import UUID - -from beartype import beartype -from fastapi import HTTPException -from pycozo.client import QueryException -from pydantic import ValidationError -from uuid_extensions import uuid7 - -from ...autogen.openapi_model import CreateUserRequest, User -from ...metrics.counters import increase_counter -from ..utils import ( - cozo_query, - partialclass, - rewrap_exceptions, - verify_developer_id_query, - wrap_in_class, -) - -ModelT = TypeVar("ModelT", bound=Any) -T = TypeVar("T") - - -@rewrap_exceptions( - { - lambda e: isinstance(e, QueryException) - and "asserted to return some results, but returned none" - in str(e): lambda *_: HTTPException( - 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.", - ), - } -) -@wrap_in_class( - User, - one=True, - transform=lambda d: {"id": UUID(d.pop("user_id")), **d}, - _kind="inserted", -) -@cozo_query -@increase_counter("create_user") -@beartype -def create_user( - *, - developer_id: UUID, - user_id: UUID | None = None, - data: CreateUserRequest, -) -> tuple[list[str], dict]: - """ - Constructs and executes a datalog query to create a new user in the CozoDB 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 brief description about the user. - metadata (dict, optional): Additional metadata about the user. Defaults to an empty dict. - client (CozoClient, optional): The CozoDB client instance to run the query. Defaults to a pre-configured client instance. - - Returns: - pd.DataFrame: A DataFrame containing the result of the query execution. - """ - - user_id = user_id or uuid7() - data.metadata = data.metadata or {} - user_data = data.model_dump() - - create_query = """ - # Then create the user - ?[user_id, developer_id, name, about, metadata] <- [ - [to_uuid($user_id), to_uuid($developer_id), $name, $about, $metadata] - ] - - :insert users { - developer_id, - user_id => - name, - about, - metadata, - } - :returning - """ - - queries = [ - verify_developer_id_query(developer_id), - create_query, - ] - - return ( - queries, - { - "user_id": str(user_id), - "developer_id": str(developer_id), - **user_data, - }, - ) diff --git a/agents-api/agents_api/models/user/delete_user.py b/agents-api/agents_api/models/user/delete_user.py deleted file mode 100644 index 7f08316be..000000000 --- a/agents-api/agents_api/models/user/delete_user.py +++ /dev/null @@ -1,116 +0,0 @@ -""" -This module contains the implementation of the delete_user_query function, which is responsible for deleting an user and its related default settings from the CozoDB database. -""" - -from typing import Any, TypeVar -from uuid import UUID - -from beartype import beartype -from fastapi import HTTPException -from pycozo.client import QueryException -from pydantic import ValidationError - -from ...autogen.openapi_model import ResourceDeletedResponse -from ...common.utils.datetime import utcnow -from ..utils import ( - cozo_query, - partialclass, - rewrap_exceptions, - verify_developer_id_query, - verify_developer_owns_resource_query, - wrap_in_class, -) - -ModelT = TypeVar("ModelT", bound=Any) -T = TypeVar("T") - - -@rewrap_exceptions( - { - lambda e: isinstance(e, QueryException) - and "Developer does not exist" in str(e): lambda *_: HTTPException( - detail="The specified developer does not exist.", - status_code=403, - ), - lambda e: isinstance(e, QueryException) - and "Developer does not own resource" - in e.resp["display"]: lambda *_: HTTPException( - 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.", - ), - } -) -@wrap_in_class( - ResourceDeletedResponse, - one=True, - transform=lambda d: { - "id": UUID(d.pop("user_id")), - "deleted_at": utcnow(), - "jobs": [], - }, - _kind="deleted", -) -@cozo_query -@beartype -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. - - Returns: - ResourceDeletedResponse: The response indicating the deletion of the user. - """ - - queries = [ - verify_developer_id_query(developer_id), - verify_developer_owns_resource_query(developer_id, "users", user_id=user_id), - """ - # Delete docs - ?[owner_type, owner_id, doc_id] := - *docs{ - owner_id, - owner_type, - doc_id, - }, - owner_id = to_uuid($user_id), - owner_type = "user" - - :delete docs { - owner_type, - owner_id, - doc_id - } - :returning - """, - """ - # Delete the user - ?[user_id, developer_id] <- [[$user_id, $developer_id]] - - :delete users { - developer_id, - user_id - } - :returning - """, - ] - - return (queries, {"user_id": str(user_id), "developer_id": str(developer_id)}) diff --git a/agents-api/agents_api/models/user/get_user.py b/agents-api/agents_api/models/user/get_user.py deleted file mode 100644 index 69b3da883..000000000 --- a/agents-api/agents_api/models/user/get_user.py +++ /dev/null @@ -1,107 +0,0 @@ -from typing import Any, TypeVar -from uuid import UUID - -from beartype import beartype -from fastapi import HTTPException -from pycozo.client import QueryException -from pydantic import ValidationError - -from ...autogen.openapi_model import User -from ..utils import ( - cozo_query, - partialclass, - rewrap_exceptions, - verify_developer_id_query, - verify_developer_owns_resource_query, - wrap_in_class, -) - -ModelT = TypeVar("ModelT", bound=Any) -T = TypeVar("T") - - -@rewrap_exceptions( - { - lambda e: isinstance(e, QueryException) - and "Developer does not exist" in str(e): lambda *_: HTTPException( - detail="The specified developer does not exist.", - status_code=403, - ), - lambda e: isinstance(e, QueryException) - and "Developer does not own resource" - in e.resp["display"]: lambda *_: HTTPException( - 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.", - ), - } -) -@wrap_in_class(User, one=True) -@cozo_query -@beartype -def get_user( - *, - developer_id: UUID, - user_id: UUID, -) -> tuple[list[str], dict]: - """ - Retrieves a user by their unique identifier. - - - Parameters: - developer_id (UUID): The unique identifier of the developer associated with the user. - user_id (UUID): The unique identifier of the user to retrieve. - - Returns: - User: The retrieved user. - """ - - # Convert UUIDs to strings for query compatibility. - user_id = str(user_id) - developer_id = str(developer_id) - - get_query = """ - input[developer_id, user_id] <- [[to_uuid($developer_id), to_uuid($user_id)]] - - ?[ - id, - name, - about, - created_at, - updated_at, - metadata, - ] := input[developer_id, id], - *users { - user_id: id, - developer_id, - name, - about, - created_at, - updated_at, - metadata, - } - - :limit 1 - """ - - queries = [ - verify_developer_id_query(developer_id), - verify_developer_owns_resource_query(developer_id, "users", user_id=user_id), - get_query, - ] - - return (queries, {"developer_id": developer_id, "user_id": user_id}) diff --git a/agents-api/agents_api/models/user/list_users.py b/agents-api/agents_api/models/user/list_users.py deleted file mode 100644 index f1e06adf4..000000000 --- a/agents-api/agents_api/models/user/list_users.py +++ /dev/null @@ -1,116 +0,0 @@ -from typing import Any, Literal, TypeVar -from uuid import UUID - -from beartype import beartype -from fastapi import HTTPException -from pycozo.client import QueryException -from pydantic import ValidationError - -from ...autogen.openapi_model import User -from ...common.utils import json -from ..utils import ( - cozo_query, - partialclass, - rewrap_exceptions, - verify_developer_id_query, - wrap_in_class, -) - -ModelT = TypeVar("ModelT", bound=Any) -T = TypeVar("T") - - -@rewrap_exceptions( - { - 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(User) -@cozo_query -@beartype -def list_users( - *, - developer_id: UUID, - limit: int = 100, - offset: int = 0, - sort_by: Literal["created_at", "updated_at"] = "created_at", - direction: Literal["asc", "desc"] = "desc", - metadata_filter: dict[str, Any] = {}, -) -> tuple[list[str], dict]: - """ - 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. - sort_by (Literal["created_at", "updated_at"]): The field to sort the users by. Defaults to "created_at". - direction (Literal["asc", "desc"]): The direction to sort the users in. Defaults to "desc". - metadata_filter (dict[str, Any]): A dictionary representing filters to apply on user metadata. - - Returns: - 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( - [ - f"metadata->{json.dumps(k)} == {json.dumps(v)}" - for k, v in metadata_filter.items() - ] - ) - - sort = f"{'-' if direction == 'desc' else ''}{sort_by}" - - # Define the datalog query for retrieving user information based on the specified filters and sorting them by creation date in descending order. - list_query = f""" - input[developer_id] <- [[to_uuid($developer_id)]] - - ?[ - id, - name, - about, - created_at, - updated_at, - metadata, - ] := - input[developer_id], - *users {{ - user_id: id, - developer_id, - name, - about, - created_at, - updated_at, - metadata, - }}, - {metadata_filter_str} - - :limit $limit - :offset $offset - :sort {sort} - """ - - queries = [ - verify_developer_id_query(developer_id), - list_query, - ] - - # Execute the datalog query with the specified parameters and return the results as a DataFrame. - return ( - queries, - {"developer_id": str(developer_id), "limit": limit, "offset": offset}, - ) diff --git a/agents-api/agents_api/models/user/patch_user.py b/agents-api/agents_api/models/user/patch_user.py deleted file mode 100644 index bd3fc0246..000000000 --- a/agents-api/agents_api/models/user/patch_user.py +++ /dev/null @@ -1,121 +0,0 @@ -"""Module for generating datalog queries to update user information in the 'cozodb' database.""" - -from typing import Any, TypeVar -from uuid import UUID - -from beartype import beartype -from fastapi import HTTPException -from pycozo.client import QueryException -from pydantic import ValidationError - -from ...autogen.openapi_model import PatchUserRequest, ResourceUpdatedResponse -from ...common.utils.cozo import cozo_process_mutate_data -from ...common.utils.datetime import utcnow -from ...metrics.counters import increase_counter -from ..utils import ( - cozo_query, - partialclass, - rewrap_exceptions, - verify_developer_id_query, - verify_developer_owns_resource_query, - wrap_in_class, -) - -ModelT = TypeVar("ModelT", bound=Any) -T = TypeVar("T") - - -@rewrap_exceptions( - { - 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( - ResourceUpdatedResponse, - one=True, - transform=lambda d: {"id": d["user_id"], "jobs": [], **d}, - _kind="inserted", -) -@cozo_query -@increase_counter("patch_user") -@beartype -def patch_user( - *, - developer_id: UUID, - user_id: UUID, - data: PatchUserRequest, -) -> tuple[list[str], dict]: - """ - 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. - - Returns: - tuple[str, dict]: A pandas DataFrame containing the results of the query execution. - """ - - update_data = data.model_dump(exclude_unset=True) - - # Prepare data for mutation by filtering out None values and adding system-generated fields. - metadata = update_data.pop("metadata", {}) or {} - user_update_cols, user_update_vals = cozo_process_mutate_data( - { - **{k: v for k, v in update_data.items() if v is not None}, - "user_id": str(user_id), - "developer_id": str(developer_id), - "updated_at": utcnow().timestamp(), - } - ) - - # Construct the datalog query for updating user information. - update_query = f""" - # update the user - input[{user_update_cols}] <- $user_update_vals - - ?[{user_update_cols}, metadata] := - input[{user_update_cols}], - *users:developer_id_metadata_user_id_idx {{ - developer_id: to_uuid($developer_id), - user_id: to_uuid($user_id), - metadata: md, - }}, - metadata = concat(md, $metadata) - - :update users {{ - {user_update_cols}, metadata - }} - :returning - """ - - queries = [ - verify_developer_id_query(developer_id), - verify_developer_owns_resource_query(developer_id, "users", user_id=user_id), - update_query, - ] - - return ( - queries, - { - "user_update_vals": user_update_vals, - "metadata": metadata, - "user_id": str(user_id), - "developer_id": str(developer_id), - }, - ) diff --git a/agents-api/agents_api/models/user/update_user.py b/agents-api/agents_api/models/user/update_user.py deleted file mode 100644 index 68e6e6c25..000000000 --- a/agents-api/agents_api/models/user/update_user.py +++ /dev/null @@ -1,118 +0,0 @@ -from typing import Any, TypeVar -from uuid import UUID - -from beartype import beartype -from fastapi import HTTPException -from pycozo.client import QueryException -from pydantic import ValidationError - -from ...autogen.openapi_model import ResourceUpdatedResponse, UpdateUserRequest -from ...common.utils.cozo import cozo_process_mutate_data -from ...metrics.counters import increase_counter -from ..utils import ( - cozo_query, - partialclass, - rewrap_exceptions, - verify_developer_id_query, - verify_developer_owns_resource_query, - wrap_in_class, -) - -ModelT = TypeVar("ModelT", bound=Any) -T = TypeVar("T") - - -@rewrap_exceptions( - { - 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( - ResourceUpdatedResponse, - one=True, - transform=lambda d: {"id": d["user_id"], "jobs": [], **d}, - _kind="inserted", -) -@cozo_query -@increase_counter("update_user") -@beartype -def update_user( - *, developer_id: UUID, user_id: UUID, data: UpdateUserRequest -) -> tuple[list[str], dict]: - """ - Updates user information in the 'cozodb' database. - - Parameters: - developer_id (UUID): The developer's unique identifier. - user_id (UUID): The user's unique identifier. - client (CozoClient): The Cozo database client instance. - **update_data: Arbitrary keyword arguments representing the data to update. - - Returns: - pd.DataFrame: A DataFrame containing the result of the update operation. - """ - user_id = str(user_id) - developer_id = str(developer_id) - update_data = data.model_dump() - - # Prepares the update data by filtering out None values and adding user_id and developer_id. - user_update_cols, user_update_vals = cozo_process_mutate_data( - { - **{k: v for k, v in update_data.items() if v is not None}, - "user_id": user_id, - "developer_id": developer_id, - } - ) - - # Constructs the update operation for the user, setting new values and updating 'updated_at'. - update_query = f""" - # update the user - # This line updates the user's information based on the provided columns and values. - input[{user_update_cols}] <- $user_update_vals - original[created_at] := *users{{ - developer_id: to_uuid($developer_id), - user_id: to_uuid($user_id), - created_at, - }}, - - ?[created_at, updated_at, {user_update_cols}] := - input[{user_update_cols}], - original[created_at], - updated_at = now(), - - :put users {{ - created_at, - updated_at, - {user_update_cols} - }} - :returning - """ - - queries = [ - verify_developer_id_query(developer_id), - verify_developer_owns_resource_query(developer_id, "users", user_id=user_id), - update_query, - ] - - return ( - queries, - { - "user_update_vals": user_update_vals, - "developer_id": developer_id, - "user_id": user_id, - }, - ) diff --git a/agents-api/agents_api/models/utils.py b/agents-api/agents_api/models/utils.py deleted file mode 100644 index 08006d1c7..000000000 --- a/agents-api/agents_api/models/utils.py +++ /dev/null @@ -1,578 +0,0 @@ -import concurrent.futures -import inspect -import re -import time -from functools import partialmethod, wraps -from typing import Any, Awaitable, Callable, ParamSpec, Type, TypeVar -from uuid import UUID - -import pandas as pd -from asyncpg import Record -from fastapi import HTTPException -from httpcore import ConnectError, NetworkError, TimeoutException -from httpx import ConnectError as HttpxConnectError -from httpx import RequestError -from pydantic import BaseModel -from requests.exceptions import ConnectionError, Timeout - -from ..common.utils.cozo import uuid_int_list_to_uuid -from ..env import do_verify_developer, do_verify_developer_owns_resource - -P = ParamSpec("P") -T = TypeVar("T") -ModelT = TypeVar("ModelT", bound=BaseModel) - - -def fix_uuid( - item: dict[str, Any], attr_regex: str = r"^(?:id|.*_id)$" -) -> dict[str, Any]: - # find the attributes that are ids - id_attrs = [ - attr for attr in item.keys() if re.match(attr_regex, attr) and item[attr] - ] - - if not id_attrs: - return item - - fixed = { - **item, - **{ - attr: uuid_int_list_to_uuid(item[attr]) - for attr in id_attrs - if isinstance(item[attr], list) - }, - } - - return fixed - - -def fix_uuid_list( - items: list[dict[str, Any]], attr_regex: str = r"^(?:id|.*_id)$" -) -> list[dict[str, Any]]: - fixed = list(map(lambda item: fix_uuid(item, attr_regex), items)) - return fixed - - -def fix_uuid_if_present(item: Any, attr_regex: str = r"^(?:id|.*_id)$") -> Any: - match item: - case [dict(), *_]: - return fix_uuid_list(item, attr_regex) - - case dict(): - return fix_uuid(item, attr_regex) - - case _: - return item - - -def partialclass(cls, *args, **kwargs): - cls_signature = inspect.signature(cls) - bound = cls_signature.bind_partial(*args, **kwargs) - - # The `updated=()` argument is necessary to avoid a TypeError when using @wraps for a class - @wraps(cls, updated=()) - class NewCls(cls): - __init__ = partialmethod(cls.__init__, *bound.args, **bound.kwargs) - - return NewCls - - -def mark_session_updated_query(developer_id: UUID | str, session_id: UUID | str) -> str: - return f""" - input[developer_id, session_id] <- [[ - to_uuid("{str(developer_id)}"), - to_uuid("{str(session_id)}"), - ]] - - ?[ - developer_id, - session_id, - situation, - summary, - created_at, - metadata, - render_templates, - token_budget, - context_overflow, - updated_at, - ] := - input[developer_id, session_id], - *sessions {{ - session_id, - situation, - summary, - created_at, - metadata, - render_templates, - token_budget, - context_overflow, - @ 'END' - }}, - updated_at = [floor(now()), true] - - :put sessions {{ - developer_id, - session_id, - situation, - summary, - created_at, - metadata, - render_templates, - token_budget, - context_overflow, - updated_at, - }} - """ - - -def verify_developer_id_query(developer_id: UUID | str) -> str: - if not do_verify_developer: - return "?[exists] := exists = true" - - return f""" - matched[count(developer_id)] := - *developers{{ - developer_id, - }}, developer_id = to_uuid("{str(developer_id)}") - - ?[exists] := - matched[num], - exists = num > 0, - assert(exists, "Developer does not exist") - - :limit 1 - """ - - -def verify_developer_owns_resource_query( - developer_id: UUID | str, - resource: str, - parents: list[tuple[str, str]] | None = None, - **resource_id, -) -> str: - if not do_verify_developer_owns_resource: - return "?[exists] := exists = true" - - parents = parents or [] - resource_id_key, resource_id_value = next(iter(resource_id.items())) - - parents.append((resource, resource_id_key)) - parent_keys = ["developer_id", *map(lambda x: x[1], parents)] - - rule_head = f""" - found[count({resource_id_key})] := - developer_id = to_uuid("{str(developer_id)}"), - {resource_id_key} = to_uuid("{str(resource_id_value)}"), - """ - - rule_body = "" - for parent_key, (relation, key) in zip(parent_keys, parents): - rule_body += f""" - *{relation}{{ - {parent_key}, - {key}, - }}, - """ - - assertion = f""" - ?[exists] := - 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 - return rule - - -def make_cozo_json_query(fields): - return ", ".join(f'"{field}": {field}' for field in fields).strip() - - -def cozo_query( - func: Callable[P, tuple[str | list[str | None], dict]] | None = None, - debug: bool | None = None, - only_on_error: bool = False, - timeit: bool = False, -): - def cozo_query_dec(func: Callable[P, tuple[str | list[Any], dict]]): - """ - Decorator that wraps a function that takes arbitrary arguments, and - returns a (query string, variables) tuple. - - The wrapped function should additionally take a client keyword argument - and then run the query using the client, returning a DataFrame. - """ - - from pprint import pprint - - from tenacity import ( - retry, - retry_if_exception, - stop_after_attempt, - wait_exponential, - ) - - def is_resource_busy(e: Exception) -> bool: - return ( - isinstance(e, HTTPException) - and e.status_code == 429 - and not getattr(e, "cozo_offline", False) - ) - - @retry( - stop=stop_after_attempt(4), - wait=wait_exponential(multiplier=1, min=4, max=10), - retry=retry_if_exception(is_resource_busy), - ) - @wraps(func) - def wrapper(*args: P.args, client=None, **kwargs: P.kwargs) -> pd.DataFrame: - queries, variables = func(*args, **kwargs) - - if isinstance(queries, str): - query = queries - else: - queries = [str(query) for query in queries if query] - query = "}\n\n{\n".join(queries) - query = f"{{ {query} }}" - - not only_on_error and debug and print(query) - not only_on_error and debug and pprint( - dict( - variables=variables, - ) - ) - - # Run the query - from ..clients import cozo - - try: - client = client or cozo.get_cozo_client() - - start = timeit and time.perf_counter() - result = client.run(query, variables) - end = timeit and time.perf_counter() - - timeit and print(f"Cozo query time: {end - start:.2f} seconds") - - except Exception as e: - if only_on_error and debug: - print(query) - pprint(variables) - - debug and print(repr(e)) - - pretty_error = repr(e).lower() - cozo_busy = ("busy" in pretty_error) or ( - "when executing against relation '_" in pretty_error - ) - cozo_offline = isinstance(e, ConnectionError) and ( - ("connection refused" in pretty_error) - or ("name or service not known" in pretty_error) - ) - connection_error = isinstance( - e, - ( - ConnectionError, - Timeout, - TimeoutException, - NetworkError, - RequestError, - ), - ) - - if cozo_busy or connection_error or cozo_offline: - exc = HTTPException( - status_code=429, detail="Resource busy. Please try again later." - ) - exc.cozo_offline = cozo_offline - raise exc from e - - raise - - # Need to fix the UUIDs in the result - result = result.map(fix_uuid_if_present) - - not only_on_error and debug and pprint( - dict( - result=result.to_dict(orient="records"), - ) - ) - - return result - - # Set the wrapped function as an attribute of the wrapper, - # forwards the __wrapped__ attribute if it exists. - setattr(wrapper, "__wrapped__", getattr(func, "__wrapped__", func)) - - return wrapper - - if func is not None and callable(func): - return cozo_query_dec(func) - - return cozo_query_dec - - -def cozo_query_async( - func: Callable[ - P, - tuple[str | list[str | None], dict] - | Awaitable[tuple[str | list[str | None], dict]], - ] - | None = None, - debug: bool | None = None, - only_on_error: bool = False, - timeit: bool = False, -): - def cozo_query_dec( - func: Callable[ - P, tuple[str | list[Any], dict] | Awaitable[tuple[str | list[Any], dict]] - ], - ): - """ - Decorator that wraps a function that takes arbitrary arguments, and - returns a (query string, variables) tuple. - - The wrapped function should additionally take a client keyword argument - and then run the query using the client, returning a DataFrame. - """ - - from pprint import pprint - - from tenacity import ( - retry, - retry_if_exception, - stop_after_attempt, - wait_exponential, - ) - - def is_resource_busy(e: Exception) -> bool: - return ( - isinstance(e, HTTPException) - and e.status_code == 429 - and not getattr(e, "cozo_offline", False) - ) - - @retry( - stop=stop_after_attempt(6), - wait=wait_exponential(multiplier=1.2, min=3, max=10), - retry=retry_if_exception(is_resource_busy), - reraise=True, - ) - @wraps(func) - async def wrapper( - *args: P.args, client=None, **kwargs: P.kwargs - ) -> pd.DataFrame: - if inspect.iscoroutinefunction(func): - queries, variables = await func(*args, **kwargs) - else: - queries, variables = func(*args, **kwargs) - - if isinstance(queries, str): - query = queries - else: - queries = [str(query) for query in queries if query] - query = "}\n\n{\n".join(queries) - query = f"{{ {query} }}" - - not only_on_error and debug and print(query) - not only_on_error and debug and pprint( - dict( - variables=variables, - ) - ) - - # Run the query - from ..clients import cozo - - try: - client = client or cozo.get_async_cozo_client() - - start = timeit and time.perf_counter() - result = await client.run(query, variables) - end = timeit and time.perf_counter() - - timeit and print(f"Cozo query time: {end - start:.2f} seconds") - - except Exception as e: - if only_on_error and debug: - print(query) - pprint(variables) - - debug and print(repr(e)) - - pretty_error = repr(e).lower() - cozo_busy = ("busy" in pretty_error) or ( - "when executing against relation '_" in pretty_error - ) - cozo_offline = ( - isinstance(e, ConnectError) - or isinstance(e, HttpxConnectError) - and ( - ("all connection attempts failed" in pretty_error) - or ("name or service not known" in pretty_error) - ) - ) - connection_error = isinstance( - e, - ( - ConnectError, - HttpxConnectError, - TimeoutException, - NetworkError, - RequestError, - ), - ) - - if cozo_busy or connection_error or cozo_offline: - exc = HTTPException( - status_code=429, detail="Resource busy. Please try again later." - ) - exc.cozo_offline = cozo_offline - raise exc from e - - raise - - # Need to fix the UUIDs in the result - result = result.map(fix_uuid_if_present) - - not only_on_error and debug and pprint( - dict( - result=result.to_dict(orient="records"), - ) - ) - - return result - - # Set the wrapped function as an attribute of the wrapper, - # forwards the __wrapped__ attribute if it exists. - setattr(wrapper, "__wrapped__", getattr(func, "__wrapped__", func)) - - return wrapper - - if func is not None and callable(func): - return cozo_query_dec(func) - - return cozo_query_dec - - -def wrap_in_class( - cls: Type[ModelT] | Callable[..., ModelT], - one: bool = False, - transform: Callable[[dict], dict] | None = None, - _kind: str | None = None, -): - def _return_data(rec: Record): - # Convert df to list of dicts - # if _kind: - # rec = rec[rec["_kind"] == _kind] - - data = list(rec.items()) - - nonlocal transform - transform = transform or (lambda x: x) - - if one: - assert len(data) >= 1, "Expected one result, got none" - obj: ModelT = cls(**transform(data[0])) - return obj - - objs: list[ModelT] = [cls(**item) for item in map(transform, data)] - return objs - - def decorator(func: Callable[P, pd.DataFrame | Awaitable[pd.DataFrame]]): - @wraps(func) - def wrapper(*args: P.args, **kwargs: P.kwargs) -> ModelT | list[ModelT]: - return _return_data(func(*args, **kwargs)) - - @wraps(func) - async def async_wrapper( - *args: P.args, **kwargs: P.kwargs - ) -> ModelT | list[ModelT]: - return _return_data(await func(*args, **kwargs)) - - # Set the wrapped function as an attribute of the wrapper, - # forwards the __wrapped__ attribute if it exists. - setattr(wrapper, "__wrapped__", getattr(func, "__wrapped__", func)) - setattr(async_wrapper, "__wrapped__", getattr(func, "__wrapped__", func)) - - return async_wrapper if inspect.iscoroutinefunction(func) else wrapper - - return decorator - - -def rewrap_exceptions( - mapping: dict[ - Type[BaseException] | Callable[[BaseException], bool], - Type[BaseException] | Callable[[BaseException], BaseException], - ], - /, -): - def _check_error(error): - nonlocal mapping - - for check, transform in mapping.items(): - should_catch = ( - isinstance(error, check) if isinstance(check, type) else check(error) - ) - - if should_catch: - new_error = ( - transform(str(error)) - if isinstance(transform, type) - else transform(error) - ) - - setattr(new_error, "__cause__", error) - - raise new_error from error - - def decorator(func: Callable[P, T | Awaitable[T]]): - @wraps(func) - async def async_wrapper(*args: P.args, **kwargs: P.kwargs) -> T: - try: - result: T = await func(*args, **kwargs) - except BaseException as error: - _check_error(error) - raise - - return result - - @wraps(func) - def wrapper(*args: P.args, **kwargs: P.kwargs) -> T: - try: - result: T = func(*args, **kwargs) - except BaseException as error: - _check_error(error) - raise - - return result - - # Set the wrapped function as an attribute of the wrapper, - # forwards the __wrapped__ attribute if it exists. - setattr(wrapper, "__wrapped__", getattr(func, "__wrapped__", func)) - setattr(async_wrapper, "__wrapped__", getattr(func, "__wrapped__", func)) - - return async_wrapper if inspect.iscoroutinefunction(func) else wrapper - - return decorator - - -def run_concurrently( - fns: list[Callable[..., Any]], - *, - args_list: list[tuple] = [], - kwargs_list: list[dict] = [], -) -> list[Any]: - args_list = args_list or [tuple()] * len(fns) - kwargs_list = kwargs_list or [dict()] * len(fns) - - with concurrent.futures.ThreadPoolExecutor() as executor: - futures = [ - executor.submit(fn, *args, **kwargs) - for fn, args, kwargs in zip(fns, args_list, kwargs_list) - ] - - return [future.result() for future in concurrent.futures.as_completed(futures)] diff --git a/agents-api/agents_api/queries/__init__.py b/agents-api/agents_api/queries/__init__.py new file mode 100644 index 000000000..eabb352e5 --- /dev/null +++ b/agents-api/agents_api/queries/__init__.py @@ -0,0 +1,21 @@ +""" +The `queries` module of the agents API is designed to encapsulate all data interactions with the PostgreSQL database. It provides a structured way to perform CRUD (Create, Read, Update, Delete) operations and other specific data manipulations across various entities such as agents, documents, entries, sessions, tools, and users. + +Each sub-module within this module corresponds to a specific entity and contains functions and classes that implement SQL queries for interacting with the database. These interactions include creating new records, updating existing ones, retrieving data for specific conditions, and deleting records. The operations are crucial for the functionality of the agents API, enabling it to manage and process data effectively for each entity. + +This module also integrates with the `common` module for exception handling and utility functions, ensuring robust error management and providing reusable components for data processing and query construction. +""" + +# ruff: noqa: F401, F403, F405 + +from . import agents as agents +from . import developers as developers +from . import docs as docs +from . import entries as entries +from . import executions as executions +from . import files as files +from . import sessions as sessions +from . import tasks as tasks +from . import tools as tools +from . import users as users + diff --git a/agents-api/agents_api/queries/developers/get_developer.py b/agents-api/agents_api/queries/developers/get_developer.py index b164bad81..a02a8f914 100644 --- a/agents-api/agents_api/queries/developers/get_developer.py +++ b/agents-api/agents_api/queries/developers/get_developer.py @@ -1,4 +1,6 @@ -"""Module for retrieving document snippets from the CozoDB based on document IDs.""" +""" +Module for retrieving developer information from the PostgreSQL database. +""" from uuid import UUID diff --git a/agents-api/agents_api/queries/tools/create_tools.py b/agents-api/agents_api/queries/tools/create_tools.py index 70277ab99..4f47ee099 100644 --- a/agents-api/agents_api/queries/tools/create_tools.py +++ b/agents-api/agents_api/queries/tools/create_tools.py @@ -1,4 +1,4 @@ -"""This module contains functions for creating tools in the CozoDB database.""" +"""This module contains functions for creating tools in the PostgreSQL database.""" from typing import Any from uuid import UUID @@ -78,9 +78,10 @@ async def create_tools( ignore_existing: bool = False, # TODO: what to do with this flag? ) -> tuple[str, list, str]: """ - Constructs a datalog query for inserting tool records into the 'agent_functions' relation in the CozoDB. + Constructs an SQL query for inserting tool records into the 'tools' relation in the PostgreSQL database. Parameters: + developer_id (UUID): The unique identifier for the developer. agent_id (UUID): The unique identifier for the agent. data (list[CreateToolRequest]): A list of function definitions to be inserted. diff --git a/agents-api/agents_api/queries/tools/patch_tool.py b/agents-api/agents_api/queries/tools/patch_tool.py index b65eca481..c41d89b4e 100644 --- a/agents-api/agents_api/queries/tools/patch_tool.py +++ b/agents-api/agents_api/queries/tools/patch_tool.py @@ -50,8 +50,7 @@ async def patch_tool( *, developer_id: UUID, agent_id: UUID, tool_id: UUID, data: PatchToolRequest ) -> tuple[str, list]: """ - 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. + Updates the tool information for a given agent and tool ID in the 'PostgreSQL' database. Parameters: agent_id (UUID): The unique identifier of the agent. diff --git a/agents-api/agents_api/worker/worker.py b/agents-api/agents_api/worker/worker.py index 39eff2b54..c88fdb72b 100644 --- a/agents-api/agents_api/worker/worker.py +++ b/agents-api/agents_api/worker/worker.py @@ -21,7 +21,6 @@ def create_worker(client: Client) -> Any: from ..activities import task_steps from ..activities.demo import demo_activity - from ..activities.embed_docs import embed_docs from ..activities.excecute_api_call import execute_api_call from ..activities.execute_integration import execute_integration from ..activities.execute_system import execute_system @@ -35,7 +34,6 @@ def create_worker(client: Client) -> Any: temporal_task_queue, ) from ..workflows.demo import DemoWorkflow - from ..workflows.embed_docs import EmbedDocsWorkflow from ..workflows.mem_mgmt import MemMgmtWorkflow from ..workflows.mem_rating import MemRatingWorkflow from ..workflows.summarization import SummarizationWorkflow @@ -54,14 +52,12 @@ def create_worker(client: Client) -> Any: SummarizationWorkflow, MemMgmtWorkflow, MemRatingWorkflow, - EmbedDocsWorkflow, TaskExecutionWorkflow, TruncationWorkflow, ], activities=[ *task_activities, demo_activity, - embed_docs, execute_integration, execute_system, execute_api_call, diff --git a/agents-api/agents_api/workflows/embed_docs.py b/agents-api/agents_api/workflows/embed_docs.py deleted file mode 100644 index 9e7b43d79..000000000 --- a/agents-api/agents_api/workflows/embed_docs.py +++ /dev/null @@ -1,27 +0,0 @@ -#!/usr/bin/env python3 - - -from datetime import timedelta - -from temporalio import workflow - -with workflow.unsafe.imports_passed_through(): - from ..activities.embed_docs import embed_docs - from ..activities.types import EmbedDocsPayload - from ..common.retry_policies import DEFAULT_RETRY_POLICY - from ..env import temporal_heartbeat_timeout, temporal_schedule_to_close_timeout - - -@workflow.defn -class EmbedDocsWorkflow: - @workflow.run - async def run(self, embed_payload: EmbedDocsPayload) -> None: - await workflow.execute_activity( - embed_docs, - embed_payload, - schedule_to_close_timeout=timedelta( - seconds=temporal_schedule_to_close_timeout - ), - retry_policy=DEFAULT_RETRY_POLICY, - heartbeat_timeout=timedelta(seconds=temporal_heartbeat_timeout), - ) From c53319e710a9ecefae7fdac2b323765eee07fb48 Mon Sep 17 00:00:00 2001 From: Ahmad-mtos Date: Tue, 24 Dec 2024 17:14:27 +0000 Subject: [PATCH 2/4] refactor: Lint agents-api (CI) --- agents-api/agents_api/activities/task_steps/__init__.py | 2 +- agents-api/agents_api/activities/task_steps/pg_query_step.py | 5 +++-- agents-api/agents_api/queries/__init__.py | 1 - 3 files changed, 4 insertions(+), 4 deletions(-) diff --git a/agents-api/agents_api/activities/task_steps/__init__.py b/agents-api/agents_api/activities/task_steps/__init__.py index 5d02db858..363a4d5d0 100644 --- a/agents-api/agents_api/activities/task_steps/__init__.py +++ b/agents-api/agents_api/activities/task_steps/__init__.py @@ -1,13 +1,13 @@ # ruff: noqa: F401, F403, F405 from .base_evaluate import base_evaluate -from .pg_query_step import pg_query_step from .evaluate_step import evaluate_step from .for_each_step import for_each_step from .get_value_step import get_value_step from .if_else_step import if_else_step from .log_step import log_step from .map_reduce_step import map_reduce_step +from .pg_query_step import pg_query_step from .prompt_step import prompt_step from .raise_complete_async import raise_complete_async from .return_step import return_step diff --git a/agents-api/agents_api/activities/task_steps/pg_query_step.py b/agents-api/agents_api/activities/task_steps/pg_query_step.py index bfddc716f..dc11e3b5c 100644 --- a/agents-api/agents_api/activities/task_steps/pg_query_step.py +++ b/agents-api/agents_api/activities/task_steps/pg_query_step.py @@ -5,14 +5,15 @@ from temporalio import activity from ... import queries -from ...env import testing, db_dsn - from ...clients.pg import create_db_pool +from ...env import db_dsn, testing + @alru_cache(maxsize=1) async def get_db_pool(dsn: str): return await create_db_pool(dsn=dsn) + @beartype async def pg_query_step( query_name: str, diff --git a/agents-api/agents_api/queries/__init__.py b/agents-api/agents_api/queries/__init__.py index eabb352e5..4b00a644d 100644 --- a/agents-api/agents_api/queries/__init__.py +++ b/agents-api/agents_api/queries/__init__.py @@ -18,4 +18,3 @@ from . import tasks as tasks from . import tools as tools from . import users as users - From 7798826c01005450f179d4c94a34ae405483a892 Mon Sep 17 00:00:00 2001 From: Ahmad Haidar Date: Wed, 25 Dec 2024 19:50:17 +0300 Subject: [PATCH 3/4] chore: remove cozo completely, and integrate postgres --- agents-api/Dockerfile | 2 +- agents-api/Dockerfile.worker | 2 +- agents-api/agents_api/__init__.py | 3 + .../activities/execute_integration.py | 2 +- .../activities/task_steps/__init__.py | 2 +- .../activities/task_steps/pg_query_step.py | 4 +- .../activities/task_steps/transition_step.py | 4 +- agents-api/agents_api/app.py | 30 ++++--- agents-api/agents_api/clients/pg.py | 4 +- .../agents_api/common/protocol/tasks.py | 2 +- agents-api/agents_api/env.py | 21 ++--- .../queries/chat/prepare_chat_context.py | 9 +- .../executions/prepare_execution_input.py | 40 ++++----- .../queries/tasks/create_or_update_task.py | 6 +- .../agents_api/queries/tasks/create_task.py | 6 +- .../agents_api/queries/tasks/get_task.py | 7 +- .../agents_api/queries/tasks/patch_task.py | 2 +- .../agents_api/queries/tasks/update_task.py | 2 +- .../agents_api/queries/users/create_user.py | 3 +- agents-api/agents_api/queries/utils.py | 2 +- .../agents_api/routers/docs/create_doc.py | 68 +------------- .../agents_api/routers/tasks/__init__.py | 11 ++- .../routers/tasks/create_task_execution.py | 5 ++ .../routers/tasks/patch_execution.py | 51 +++++------ agents-api/docker-compose.yml | 3 +- agents-api/tests/fixtures.py | 2 +- deploy/simple-docker-compose.yaml | 88 ++----------------- embedding-service/docker-compose.yml | 3 +- memory-store/.gitignore | 1 - memory-store/docker-compose.yml | 18 +++- 30 files changed, 131 insertions(+), 272 deletions(-) diff --git a/agents-api/Dockerfile b/agents-api/Dockerfile index 54ae6b576..3408c38b5 100644 --- a/agents-api/Dockerfile +++ b/agents-api/Dockerfile @@ -30,4 +30,4 @@ COPY . ./ ENV PYTHONUNBUFFERED=1 ENV GUNICORN_CMD_ARGS="--capture-output --enable-stdio-inheritance" -ENTRYPOINT ["uv", "run", "gunicorn", "agents_api.web:app", "-c", "gunicorn_conf.py"] +ENTRYPOINT ["uv", "run", "--offline", "--no-sync", "gunicorn", "agents_api.web:app", "-c", "gunicorn_conf.py"] diff --git a/agents-api/Dockerfile.worker b/agents-api/Dockerfile.worker index 88f30e2d2..34538a27d 100644 --- a/agents-api/Dockerfile.worker +++ b/agents-api/Dockerfile.worker @@ -30,4 +30,4 @@ COPY . ./ ENV PYTHONUNBUFFERED=1 ENV GUNICORN_CMD_ARGS="--capture-output --enable-stdio-inheritance" -ENTRYPOINT ["uv", "run", "python", "-m", "agents_api.worker"] +ENTRYPOINT ["uv", "run", "--offline", "--no-sync", "python", "-m", "agents_api.worker"] diff --git a/agents-api/agents_api/__init__.py b/agents-api/agents_api/__init__.py index dfe10ea38..e8fb2e7ec 100644 --- a/agents-api/agents_api/__init__.py +++ b/agents-api/agents_api/__init__.py @@ -9,3 +9,6 @@ with workflow.unsafe.imports_passed_through(): import msgpack as msgpack + +import os + diff --git a/agents-api/agents_api/activities/execute_integration.py b/agents-api/agents_api/activities/execute_integration.py index d058553c4..08046498c 100644 --- a/agents-api/agents_api/activities/execute_integration.py +++ b/agents-api/agents_api/activities/execute_integration.py @@ -8,7 +8,7 @@ from ..common.exceptions.tools import IntegrationExecutionException from ..common.protocol.tasks import ExecutionInput, StepContext from ..env import testing -from ..models.tools import get_tool_args_from_metadata +from ..queries.tools import get_tool_args_from_metadata @beartype diff --git a/agents-api/agents_api/activities/task_steps/__init__.py b/agents-api/agents_api/activities/task_steps/__init__.py index 78caeafa6..c85dfc0ec 100644 --- a/agents-api/agents_api/activities/task_steps/__init__.py +++ b/agents-api/agents_api/activities/task_steps/__init__.py @@ -2,7 +2,7 @@ from .base_evaluate import base_evaluate -# from .cozo_query_step import cozo_query_step +from .pg_query_step import pg_query_step from .evaluate_step import evaluate_step from .for_each_step import for_each_step from .get_value_step import get_value_step diff --git a/agents-api/agents_api/activities/task_steps/pg_query_step.py b/agents-api/agents_api/activities/task_steps/pg_query_step.py index dc11e3b5c..b5113c89d 100644 --- a/agents-api/agents_api/activities/task_steps/pg_query_step.py +++ b/agents-api/agents_api/activities/task_steps/pg_query_step.py @@ -6,7 +6,7 @@ from ... import queries from ...clients.pg import create_db_pool -from ...env import db_dsn, testing +from ...env import pg_dsn, testing @alru_cache(maxsize=1) @@ -18,7 +18,7 @@ async def get_db_pool(dsn: str): async def pg_query_step( query_name: str, values: dict[str, Any], - dsn: str = db_dsn, + dsn: str = pg_dsn, ) -> Any: pool = await get_db_pool(dsn=dsn) diff --git a/agents-api/agents_api/activities/task_steps/transition_step.py b/agents-api/agents_api/activities/task_steps/transition_step.py index 57d594ec3..bbed37679 100644 --- a/agents-api/agents_api/activities/task_steps/transition_step.py +++ b/agents-api/agents_api/activities/task_steps/transition_step.py @@ -15,7 +15,7 @@ ) from ...exceptions import LastErrorInput, TooManyRequestsError from ...queries.executions.create_execution_transition import ( - create_execution_transition_async, + create_execution_transition, ) from ..utils import RateLimiter @@ -52,7 +52,7 @@ async def transition_step( # Create transition try: - transition = await create_execution_transition_async( + transition = await create_execution_transition( developer_id=context.execution_input.developer_id, execution_id=context.execution_input.execution.id, task_id=context.execution_input.task.id, diff --git a/agents-api/agents_api/app.py b/agents-api/agents_api/app.py index 0ce9be5e8..752a07dfd 100644 --- a/agents-api/agents_api/app.py +++ b/agents-api/agents_api/app.py @@ -17,10 +17,10 @@ @asynccontextmanager async def lifespan(app: FastAPI): # INIT POSTGRES # - db_dsn = os.environ.get("DB_DSN") + pg_dsn = os.environ.get("PG_DSN") if not getattr(app.state, "postgres_pool", None): - app.state.postgres_pool = await create_db_pool(db_dsn) + app.state.postgres_pool = await create_db_pool(pg_dsn) # INIT S3 # s3_access_key = os.environ.get("S3_ACCESS_KEY") @@ -67,7 +67,8 @@ async def lifespan(app: FastAPI): lifespan=lifespan, # # Global dependencies - dependencies=[Depends(valid_content_length)], + # FIXME: This is blocking access to scalar + # dependencies=[Depends(valid_content_length)], ) # Enable metrics @@ -92,19 +93,20 @@ async def scalar_html(): # content-length validation +# FIXME: This is blocking access to scalar # NOTE: This relies on client reporting the correct content-length header # TODO: We should use streaming for large payloads -@app.middleware("http") -async def validate_content_length( - request: Request, - call_next: Callable[[Request], Coroutine[Any, Any, Response]], -): - content_length = request.headers.get("content-length") +# @app.middleware("http") +# async def validate_content_length( +# request: Request, +# call_next: Callable[[Request], Coroutine[Any, Any, Response]], +# ): +# content_length = request.headers.get("content-length") - if not content_length: - return Response(status_code=411, content="Content-Length header is required") +# if not content_length: +# return Response(status_code=411, content="Content-Length header is required") - if int(content_length) > max_payload_size: - return Response(status_code=413, content="Payload too large") +# if int(content_length) > max_payload_size: +# return Response(status_code=413, content="Payload too large") - return await call_next(request) +# return await call_next(request) diff --git a/agents-api/agents_api/clients/pg.py b/agents-api/agents_api/clients/pg.py index acf7a2b0e..ebb1ae7f0 100644 --- a/agents-api/agents_api/clients/pg.py +++ b/agents-api/agents_api/clients/pg.py @@ -2,7 +2,7 @@ import asyncpg -from ..env import db_dsn +from ..env import pg_dsn async def _init_conn(conn): @@ -16,5 +16,5 @@ async def _init_conn(conn): async def create_db_pool(dsn: str | None = None): return await asyncpg.create_pool( - dsn if dsn is not None else db_dsn, init=_init_conn + dsn if dsn is not None else pg_dsn, init=_init_conn ) diff --git a/agents-api/agents_api/common/protocol/tasks.py b/agents-api/agents_api/common/protocol/tasks.py index f3bb81d07..31543b0be 100644 --- a/agents-api/agents_api/common/protocol/tasks.py +++ b/agents-api/agents_api/common/protocol/tasks.py @@ -139,7 +139,7 @@ class PartialTransition(create_partial_model(CreateTransitionRequest)): class ExecutionInput(BaseModel): developer_id: UUID - execution: Execution + execution: Execution | None = None task: TaskSpecDef agent: Agent agent_tools: list[Tool | CreateToolRequest] diff --git a/agents-api/agents_api/env.py b/agents-api/agents_api/env.py index 54c8a2eee..1ac4becb6 100644 --- a/agents-api/agents_api/env.py +++ b/agents-api/agents_api/env.py @@ -51,24 +51,15 @@ s3_secret_key: str | None = env.str("S3_SECRET_KEY", default=None) -# Cozo -# ---- -cozo_host: str = env.str("COZO_HOST", default="http://127.0.0.1:9070") -cozo_auth: str = env.str("COZO_AUTH_TOKEN", default=None) -summarization_model_name: str = env.str( - "SUMMARIZATION_MODEL_NAME", default="gpt-4-turbo" -) -do_verify_developer: bool = env.bool("DO_VERIFY_DEVELOPER", default=True) -do_verify_developer_owns_resource: bool = env.bool( - "DO_VERIFY_DEVELOPER_OWNS_RESOURCE", default=True -) - # PostgreSQL # ---- -db_dsn: str = env.str( - "DB_DSN", +pg_dsn: str = env.str( + "PG_DSN", default="postgres://postgres:postgres@0.0.0.0:5432/postgres?sslmode=disable", ) +summarization_model_name: str = env.str( + "SUMMARIZATION_MODEL_NAME", default="gpt-4-turbo" +) query_timeout: float = env.float("QUERY_TIMEOUT", default=90.0) @@ -156,8 +147,6 @@ def _parse_optional_int(val: str | None) -> int | None: environment: Dict[str, Any] = dict( debug=debug, multi_tenant_mode=multi_tenant_mode, - cozo_host=cozo_host, - cozo_auth=cozo_auth, sentry_dsn=sentry_dsn, temporal_endpoint=temporal_endpoint, temporal_task_queue=temporal_task_queue, diff --git a/agents-api/agents_api/queries/chat/prepare_chat_context.py b/agents-api/agents_api/queries/chat/prepare_chat_context.py index 1d9bd52fb..de532844f 100644 --- a/agents-api/agents_api/queries/chat/prepare_chat_context.py +++ b/agents-api/agents_api/queries/chat/prepare_chat_context.py @@ -1,11 +1,9 @@ from typing import Any, TypeVar from uuid import UUID -import sqlvalidator from beartype import beartype from ...common.protocol.sessions import ChatContext, make_session -from ...exceptions import InvalidSQLQuery from ..utils import ( pg_query, wrap_in_class, @@ -15,8 +13,8 @@ T = TypeVar("T") -sql_query = sqlvalidator.parse( - """SELECT * FROM +sql_query =""" +SELECT * FROM ( SELECT jsonb_agg(u) AS users FROM ( SELECT @@ -103,9 +101,6 @@ session_lookup.participant_type = 'agent' ) r ) AS toolsets""" -) -if not sql_query.is_valid(): - raise InvalidSQLQuery("prepare_chat_context") def _transform(d): diff --git a/agents-api/agents_api/queries/executions/prepare_execution_input.py b/agents-api/agents_api/queries/executions/prepare_execution_input.py index d738624f2..5940c4047 100644 --- a/agents-api/agents_api/queries/executions/prepare_execution_input.py +++ b/agents-api/agents_api/queries/executions/prepare_execution_input.py @@ -14,14 +14,17 @@ sql_query = """SELECT * FROM ( - SELECT to_jsonb(a) AS agents FROM ( + SELECT to_jsonb(a) AS agent FROM ( SELECT * FROM agents WHERE developer_id = $1 AND - agent_id = $4 + agent_id = ( + SELECT agent_id FROM tasks + WHERE developer_id = $1 AND task_id = $2 + ) LIMIT 1 ) a -) AS agents, +) AS agent, ( SELECT jsonb_agg(r) AS tools FROM ( SELECT * FROM tools @@ -31,25 +34,25 @@ ) r ) AS tools, ( - SELECT to_jsonb(t) AS tasks FROM ( + SELECT to_jsonb(t) AS task FROM ( SELECT * FROM tasks WHERE developer_id = $1 AND task_id = $2 LIMIT 1 ) t -) AS tasks, -( - SELECT to_jsonb(e) AS executions FROM ( - SELECT * FROM executions - WHERE - developer_id = $1 AND - task_id = $2 AND - execution_id = $3 - LIMIT 1 - ) e -) AS executions; +) AS task; """ +# ( +# SELECT to_jsonb(e) AS execution FROM ( +# SELECT * FROM latest_executions +# WHERE +# developer_id = $1 AND +# task_id = $2 AND +# execution_id = $3 +# LIMIT 1 +# ) e +# ) AS execution; # @rewrap_exceptions( @@ -70,7 +73,7 @@ transform=lambda d: { **d, "task": { - "tools": [*d["task"].pop("tools")], + "tools": d["tools"], **d["task"], }, "agent_tools": [ @@ -86,14 +89,11 @@ async def prepare_execution_input( task_id: UUID, execution_id: UUID, ) -> tuple[str, list]: - dummy_agent_id = UUID(int=0) - return ( sql_query, [ str(developer_id), str(task_id), - str(execution_id), - str(dummy_agent_id), + # str(execution_id), ], ) diff --git a/agents-api/agents_api/queries/tasks/create_or_update_task.py b/agents-api/agents_api/queries/tasks/create_or_update_task.py index d02814875..8be5cde84 100644 --- a/agents-api/agents_api/queries/tasks/create_or_update_task.py +++ b/agents-api/agents_api/queries/tasks/create_or_update_task.py @@ -192,13 +192,13 @@ async def create_or_update_task( tool.type, tool.name, tool.description, - getattr(tool, tool.type), # spec + getattr(tool, tool.type) and getattr(tool, tool.type).model_dump(mode="json"), # spec ] for tool in data.tools or [] ] # Generate workflows from task data using task_to_spec - workflows_spec = task_to_spec(data).model_dump(exclude_none=True, mode="json") + workflows_spec = task_to_spec(data).model_dump(mode="json") workflow_params = [] for workflow in workflows_spec.get("workflows", []): workflow_name = workflow.get("name") @@ -211,7 +211,7 @@ async def create_or_update_task( workflow_name, # $3 step_idx, # $4 step["kind_"], # $5 - step[step["kind_"]], # $6 + step, # $6 ] ) diff --git a/agents-api/agents_api/queries/tasks/create_task.py b/agents-api/agents_api/queries/tasks/create_task.py index 6deffc3d5..5c05c3666 100644 --- a/agents-api/agents_api/queries/tasks/create_task.py +++ b/agents-api/agents_api/queries/tasks/create_task.py @@ -167,13 +167,13 @@ async def create_task( tool.type, tool.name, tool.description, - getattr(tool, tool.type), # spec + getattr(tool, tool.type) and getattr(tool, tool.type).model_dump(mode="json"), # spec ] for tool in data.tools or [] ] # Generate workflows from task data using task_to_spec - workflows_spec = task_to_spec(data).model_dump(exclude_none=True, mode="json") + workflows_spec = task_to_spec(data).model_dump(mode="json") workflow_params = [] for workflow in workflows_spec.get("workflows", []): workflow_name = workflow.get("name") @@ -187,7 +187,7 @@ async def create_task( workflow_name, # $4 step_idx, # $5 step["kind_"], # $6 - step[step["kind_"]], # $7 + step, # $7 ] ) diff --git a/agents-api/agents_api/queries/tasks/get_task.py b/agents-api/agents_api/queries/tasks/get_task.py index 03da91256..1f0dd00cd 100644 --- a/agents-api/agents_api/queries/tasks/get_task.py +++ b/agents-api/agents_api/queries/tasks/get_task.py @@ -17,12 +17,7 @@ CASE WHEN w.name IS NOT NULL THEN jsonb_build_object( 'name', w.name, - 'steps', jsonb_build_array( - jsonb_build_object( - w.step_type, w.step_definition, - 'step_idx', w.step_idx -- Not sure if this is needed - ) - ) + 'steps', jsonb_build_array(w.step_definition) ) END ) FILTER (WHERE w.name IS NOT NULL), diff --git a/agents-api/agents_api/queries/tasks/patch_task.py b/agents-api/agents_api/queries/tasks/patch_task.py index 2349f87c5..48111a333 100644 --- a/agents-api/agents_api/queries/tasks/patch_task.py +++ b/agents-api/agents_api/queries/tasks/patch_task.py @@ -198,7 +198,7 @@ async def patch_task( else: workflow_query = new_workflows_query workflow_params = [] - workflows_spec = task_to_spec(data).model_dump(exclude_none=True, mode="json") + workflows_spec = task_to_spec(data).model_dump(mode="json") for workflow in workflows_spec.get("workflows", []): workflow_name = workflow.get("name") steps = workflow.get("steps", []) diff --git a/agents-api/agents_api/queries/tasks/update_task.py b/agents-api/agents_api/queries/tasks/update_task.py index 495499eb1..0379e0312 100644 --- a/agents-api/agents_api/queries/tasks/update_task.py +++ b/agents-api/agents_api/queries/tasks/update_task.py @@ -137,7 +137,7 @@ async def update_task( ] # Generate workflows from task data - workflows_spec = task_to_spec(data).model_dump(exclude_none=True, mode="json") + workflows_spec = task_to_spec(data).model_dump(mode="json") workflow_params = [] for workflow in workflows_spec.get("workflows", []): workflow_name = workflow.get("name") diff --git a/agents-api/agents_api/queries/users/create_user.py b/agents-api/agents_api/queries/users/create_user.py index e246c7255..982d7a97e 100644 --- a/agents-api/agents_api/queries/users/create_user.py +++ b/agents-api/agents_api/queries/users/create_user.py @@ -73,13 +73,14 @@ async def create_user( tuple[str, list]: A tuple containing the SQL query and its parameters. """ user_id = user_id or uuid7() + metadata = data.metadata.model_dump(mode="json") or {} params = [ developer_id, # $1 user_id, # $2 data.name, # $3 data.about, # $4 - data.metadata or {}, # $5 + metadata, # $5 ] return ( diff --git a/agents-api/agents_api/queries/utils.py b/agents-api/agents_api/queries/utils.py index 1a9ce7dc2..01652888b 100644 --- a/agents-api/agents_api/queries/utils.py +++ b/agents-api/agents_api/queries/utils.py @@ -228,7 +228,7 @@ def _return_data(rec: list[Record]): transform = transform or (lambda x: x) if one: - assert len(data) == 1, "Expected one result, got none" + assert len(data) == 1, f"Expected one result, got {len(data)}" obj: ModelT = cls(**transform(data[0])) return obj diff --git a/agents-api/agents_api/routers/docs/create_doc.py b/agents-api/agents_api/routers/docs/create_doc.py index c514fe9ee..1c9f65797 100644 --- a/agents-api/agents_api/routers/docs/create_doc.py +++ b/agents-api/agents_api/routers/docs/create_doc.py @@ -16,46 +16,6 @@ from .router import router -async def run_embed_docs_task( - *, - developer_id: UUID, - doc_id: UUID, - title: str, - content: list[str], - embed_instruction: str | None = None, - job_id: UUID, - background_tasks: BackgroundTasks, - client: TemporalClient | None = None, -): - from ...workflows.embed_docs import EmbedDocsWorkflow - - client = client or (await temporal.get_client()) - - embed_payload = EmbedDocsPayload( - developer_id=developer_id, - doc_id=doc_id, - content=content, - title=title, - # Default embed instruction for docs. See https://docs.voyageai.com/docs/embeddings - embed_instruction=embed_instruction or "Represent the document for retrieval: ", - ) - - handle = await client.start_workflow( - EmbedDocsWorkflow.run, - embed_payload, - task_queue=temporal_task_queue, - id=str(job_id), - retry_policy=DEFAULT_RETRY_POLICY, - ) - - # TODO: Remove this conditional once we have a way to run workflows in - # a test environment. - if not testing: - background_tasks.add_task(handle.result) - - return handle - - @router.post("/users/{user_id}/docs", status_code=HTTP_201_CREATED, tags=["docs"]) async def create_user_doc( user_id: UUID, @@ -83,20 +43,8 @@ async def create_user_doc( data=data, ) - embed_job_id = uuid7() - - await run_embed_docs_task( - developer_id=x_developer_id, - doc_id=doc.id, - title=doc.title, - content=doc.content, - embed_instruction=data.embed_instruction, - job_id=embed_job_id, - background_tasks=background_tasks, - ) - return ResourceCreatedResponse( - id=doc.id, created_at=doc.created_at, jobs=[embed_job_id] + id=doc.id, created_at=doc.created_at, jobs=[] ) @@ -114,18 +62,6 @@ async def create_agent_doc( data=data, ) - embed_job_id = uuid7() - - await run_embed_docs_task( - developer_id=x_developer_id, - doc_id=doc.id, - title=doc.title, - content=doc.content, - embed_instruction=data.embed_instruction, - job_id=embed_job_id, - background_tasks=background_tasks, - ) - return ResourceCreatedResponse( - id=doc.id, created_at=doc.created_at, jobs=[embed_job_id] + id=doc.id, created_at=doc.created_at, jobs=[] ) diff --git a/agents-api/agents_api/routers/tasks/__init__.py b/agents-api/agents_api/routers/tasks/__init__.py index 58b9fce54..e6933b27b 100644 --- a/agents-api/agents_api/routers/tasks/__init__.py +++ b/agents-api/agents_api/routers/tasks/__init__.py @@ -2,15 +2,14 @@ from .create_or_update_task import create_or_update_task from .create_task import create_task -# from .create_task_execution import create_task_execution -# from .get_execution_details import get_execution_details +from .create_task_execution import create_task_execution +from .get_execution_details import get_execution_details from .get_task_details import get_task_details -# from .list_execution_transitions import list_execution_transitions -# from .list_task_executions import list_task_executions +from .list_execution_transitions import list_execution_transitions +from .list_task_executions import list_task_executions from .list_tasks import list_tasks # from .patch_execution import patch_execution from .router import router -# from .stream_transitions_events import stream_transitions_events -# from .update_execution import update_execution +from .stream_transitions_events import stream_transitions_events diff --git a/agents-api/agents_api/routers/tasks/create_task_execution.py b/agents-api/agents_api/routers/tasks/create_task_execution.py index eee937b85..2af945729 100644 --- a/agents-api/agents_api/routers/tasks/create_task_execution.py +++ b/agents-api/agents_api/routers/tasks/create_task_execution.py @@ -50,6 +50,7 @@ async def start_execution( ) -> tuple[Execution, WorkflowHandle]: execution_id = uuid7() + execution = await create_execution_query( developer_id=developer_id, task_id=task_id, @@ -58,6 +59,7 @@ async def start_execution( connection_pool=connection_pool, ) + execution_input = await prepare_execution_input( developer_id=developer_id, task_id=task_id, @@ -138,12 +140,14 @@ async def create_task_execution( detail="Execution count exceeded the free tier limit", ) + execution, handle = await start_execution( developer_id=x_developer_id, task_id=task_id, data=data, ) + background_tasks.add_task( create_temporal_lookup, # @@ -152,6 +156,7 @@ async def create_task_execution( workflow_handle=handle, ) + return ResourceCreatedResponse( id=execution.id, created_at=execution.created_at, diff --git a/agents-api/agents_api/routers/tasks/patch_execution.py b/agents-api/agents_api/routers/tasks/patch_execution.py index 15b3162be..4e7d89d87 100644 --- a/agents-api/agents_api/routers/tasks/patch_execution.py +++ b/agents-api/agents_api/routers/tasks/patch_execution.py @@ -1,29 +1,30 @@ -from typing import Annotated -from uuid import UUID +# FIXME: check if this is needed +# from typing import Annotated +# from uuid import UUID -from fastapi import Depends +# from fastapi import Depends -from ...autogen.openapi_model import ( - ResourceUpdatedResponse, - UpdateExecutionRequest, -) -from ...dependencies.developer_id import get_developer_id -from ...queries.executions.update_execution import ( - update_execution as update_execution_query, -) -from .router import router +# from ...autogen.openapi_model import ( +# ResourceUpdatedResponse, +# UpdateExecutionRequest, +# ) +# from ...dependencies.developer_id import get_developer_id +# from ...queries.executions.update_execution import ( +# update_execution as update_execution_query, +# ) +# from .router import router -@router.patch("/tasks/{task_id}/executions/{execution_id}", tags=["tasks"]) -async def patch_execution( - x_developer_id: Annotated[UUID, Depends(get_developer_id)], - task_id: UUID, - execution_id: UUID, - data: UpdateExecutionRequest, -) -> ResourceUpdatedResponse: - return await update_execution_query( - developer_id=x_developer_id, - task_id=task_id, - execution_id=execution_id, - data=data, - ) +# @router.patch("/tasks/{task_id}/executions/{execution_id}", tags=["tasks"]) +# async def patch_execution( +# x_developer_id: Annotated[UUID, Depends(get_developer_id)], +# task_id: UUID, +# execution_id: UUID, +# data: UpdateExecutionRequest, +# ) -> ResourceUpdatedResponse: +# return await update_execution_query( +# developer_id=x_developer_id, +# task_id=task_id, +# execution_id=execution_id, +# data=data, +# ) diff --git a/agents-api/docker-compose.yml b/agents-api/docker-compose.yml index 67591e945..1f27ac8e2 100644 --- a/agents-api/docker-compose.yml +++ b/agents-api/docker-compose.yml @@ -8,8 +8,7 @@ x--shared-environment: &shared-environment AGENTS_API_PUBLIC_PORT: ${AGENTS_API_PUBLIC_PORT:-80} AGENTS_API_PROTOCOL: ${AGENTS_API_PROTOCOL:-http} AGENTS_API_URL: ${AGENTS_API_URL:-http://agents-api:8080} - COZO_AUTH_TOKEN: ${COZO_AUTH_TOKEN} - COZO_HOST: ${COZO_HOST:-http://memory-store:9070} + PG_DSN: ${PG_DSN:-postgres://postgres:postgres@memory-store:5432/postgres} DEBUG: ${AGENTS_API_DEBUG:-False} EMBEDDING_MODEL_ID: ${EMBEDDING_MODEL_ID:-Alibaba-NLP/gte-large-en-v1.5} INTEGRATION_SERVICE_URL: ${INTEGRATION_SERVICE_URL:-http://integrations:8000} diff --git a/agents-api/tests/fixtures.py b/agents-api/tests/fixtures.py index b35e3e5e2..9bebfd396 100644 --- a/agents-api/tests/fixtures.py +++ b/agents-api/tests/fixtures.py @@ -378,7 +378,7 @@ async def test_tool( @fixture(scope="global") def client(dsn=pg_dsn): - os.environ["DB_DSN"] = dsn + os.environ["PG_DSN"] = dsn with TestClient(app=app) as client: yield client diff --git a/deploy/simple-docker-compose.yaml b/deploy/simple-docker-compose.yaml index 0b21af407..c87e78174 100644 --- a/deploy/simple-docker-compose.yaml +++ b/deploy/simple-docker-compose.yaml @@ -13,8 +13,6 @@ services: AGENTS_API_PROTOCOL: http AGENTS_API_PUBLIC_PORT: "80" AGENTS_API_URL: http://agents-api:8080 - COZO_AUTH_TOKEN: ${COZO_AUTH_TOKEN} - COZO_HOST: http://memory-store:9070 EMBEDDING_MODEL_ID: voyage/voyage-3 INTEGRATION_SERVICE_URL: http://integrations:8000 LITELLM_MASTER_KEY: ${LITELLM_MASTER_KEY} @@ -35,32 +33,6 @@ services: published: "8080" protocol: tcp - cozo-migrate: - environment: - AGENTS_API_HOSTNAME: localhost - AGENTS_API_KEY: ${AGENTS_API_KEY} - AGENTS_API_KEY_HEADER_NAME: Authorization - AGENTS_API_PROTOCOL: http - AGENTS_API_PUBLIC_PORT: "80" - AGENTS_API_URL: http://agents-api:8080 - COZO_AUTH_TOKEN: ${COZO_AUTH_TOKEN} - COZO_HOST: http://memory-store:9070 - EMBEDDING_MODEL_ID: voyage/voyage-3 - INTEGRATION_SERVICE_URL: http://integrations:8000 - LITELLM_MASTER_KEY: ${LITELLM_MASTER_KEY} - LITELLM_URL: http://litellm:4000 - SUMMARIZATION_MODEL_NAME: gpt-4o-mini - TEMPORAL_ENDPOINT: temporal:7233 - TEMPORAL_NAMESPACE: default - TEMPORAL_TASK_QUEUE: julep-task-queue - TEMPORAL_WORKER_URL: temporal:7233 - TRUNCATE_EMBED_TEXT: "True" - WORKER_URL: temporal:7233 - image: julepai/cozo-migrate:${TAG:-dev} - networks: - default: null - restart: "no" - integrations: image: julepai/integrations:${TAG:-dev} environment: @@ -156,56 +128,11 @@ services: target: /data volume: {} - memory-store: - environment: - COZO_AUTH_TOKEN: ${COZO_AUTH_TOKEN} - COZO_BACKUP_DIR: /backup - COZO_MNT_DIR: /data - COZO_PORT: "9070" - image: julepai/memory-store:${TAG:-dev} - labels: - ofelia.enabled: "true" - ofelia.job-exec.backupcron.command: bash /app/backup.sh - ofelia.job-exec.backupcron.environment: '["COZO_PORT=9070", "COZO_AUTH_TOKEN=${COZO_AUTH_TOKEN}", "COZO_BACKUP_DIR=/backup"]' - ofelia.job-exec.backupcron.schedule: '@every 3h' - networks: - default: null - ports: - - mode: ingress - target: 9070 - published: "9070" - protocol: tcp - volumes: - - type: volume - source: cozo_data - target: /data - volume: {} - - type: volume - source: cozo_backup - target: /backup - volume: {} + # TODO: Add memory-store with postgres + # memory-store: - memory-store-backup-cron: - command: - - daemon - - --docker - - -f - - label=com.docker.compose.project=julep - depends_on: - memory-store: - condition: service_started - required: true - image: mcuadros/ofelia:latest - networks: - default: null - restart: unless-stopped - volumes: - - type: bind - source: /var/run/docker.sock - target: /var/run/docker.sock - read_only: true - bind: - create_host_path: true + # TODO: Add memory-store-backup-cron + # memory-store-backup-cron: temporal: depends_on: @@ -295,8 +222,7 @@ services: AGENTS_API_PROTOCOL: http AGENTS_API_PUBLIC_PORT: "80" AGENTS_API_URL: http://agents-api:8080 - COZO_AUTH_TOKEN: ${COZO_AUTH_TOKEN} - COZO_HOST: http://memory-store:9070 + PG_DSN: ${PG_DSN:-postgres://postgres:postgres@memory-store:5432/postgres} EMBEDDING_MODEL_ID: voyage/voyage-3 INTEGRATION_SERVICE_URL: http://integrations:8000 LITELLM_MASTER_KEY: ${LITELLM_MASTER_KEY} @@ -317,10 +243,6 @@ networks: name: julep_default volumes: - cozo_backup: - name: cozo_backup - cozo_data: - name: cozo_data litellm-db-data: name: julep_litellm-db-data litellm-redis-data: diff --git a/embedding-service/docker-compose.yml b/embedding-service/docker-compose.yml index 73df579be..a51a93e7f 100644 --- a/embedding-service/docker-compose.yml +++ b/embedding-service/docker-compose.yml @@ -17,8 +17,7 @@ x--shared-environment: &shared-environment AGENTS_API_KEY_HEADER_NAME: ${AGENTS_API_KEY_HEADER_NAME:-Authorization} AGENTS_API_HOSTNAME: ${AGENTS_API_HOSTNAME:-localhost} AGENTS_API_URL: ${AGENTS_API_URL:-http://agents-api:8080} - COZO_AUTH_TOKEN: ${COZO_AUTH_TOKEN} - COZO_HOST: ${COZO_HOST:-http://memory-store:9070} + PG_DSN: ${PG_DSN:-postgres://postgres:postgres@memory-store:5432/postgres} DEBUG: ${AGENTS_API_DEBUG:-False} EMBEDDING_MODEL_ID: ${EMBEDDING_MODEL_ID:-Alibaba-NLP/gte-large-en-v1.5} LITELLM_MASTER_KEY: ${LITELLM_MASTER_KEY} diff --git a/memory-store/.gitignore b/memory-store/.gitignore index 9383f36da..c2563b460 100644 --- a/memory-store/.gitignore +++ b/memory-store/.gitignore @@ -1,4 +1,3 @@ -cozo.db/ tmp/ *.pyc \ No newline at end of file diff --git a/memory-store/docker-compose.yml b/memory-store/docker-compose.yml index cb687142a..dafb116e1 100644 --- a/memory-store/docker-compose.yml +++ b/memory-store/docker-compose.yml @@ -1,6 +1,6 @@ name: pgai services: - db: + memory-store: image: timescale/timescaledb-ha:pg17 # For timescaledb specific options, @@ -22,10 +22,24 @@ services: vectorizer-worker: image: timescale/pgai-vectorizer-worker:v0.3.0 environment: - - PGAI_VECTORIZER_WORKER_DB_URL=postgres://postgres:${MEMORY_STORE_PASSWORD:-postgres}@db:5432/postgres + - PGAI_VECTORIZER_WORKER_DB_URL=postgres://postgres:${MEMORY_STORE_PASSWORD:-postgres}@memory-store:5432/postgres - VOYAGE_API_KEY=${VOYAGE_API_KEY} command: [ "--poll-interval", "5s" ] + migration: + image: migrate/migrate:latest + volumes: + - ./migrations:/migrations + command: [ "-path", "/migrations", "-database", "postgres://postgres:${MEMORY_STORE_PASSWORD:-postgres}@memory-store:5432/postgres?sslmode=disable" , "up"] + restart: "no" + develop: + watch: + - path: ./migrations + target: ./migrations + action: sync+restart + depends_on: + - memory-store + volumes: memory_store_data: external: true From c0acb49494aab69c55401dd464af7eaf1d7e79a7 Mon Sep 17 00:00:00 2001 From: Ahmad-mtos Date: Wed, 25 Dec 2024 16:51:55 +0000 Subject: [PATCH 4/4] refactor: Lint agents-api (CI) --- agents-api/agents_api/__init__.py | 1 - agents-api/agents_api/activities/task_steps/__init__.py | 2 -- .../agents_api/queries/chat/prepare_chat_context.py | 2 +- .../queries/executions/prepare_execution_input.py | 2 +- .../agents_api/queries/tasks/create_or_update_task.py | 3 ++- agents-api/agents_api/queries/tasks/create_task.py | 3 ++- agents-api/agents_api/routers/docs/create_doc.py | 8 ++------ agents-api/agents_api/routers/tasks/__init__.py | 2 -- .../agents_api/routers/tasks/create_task_execution.py | 5 ----- 9 files changed, 8 insertions(+), 20 deletions(-) diff --git a/agents-api/agents_api/__init__.py b/agents-api/agents_api/__init__.py index e8fb2e7ec..6c62e1f3d 100644 --- a/agents-api/agents_api/__init__.py +++ b/agents-api/agents_api/__init__.py @@ -11,4 +11,3 @@ import msgpack as msgpack import os - diff --git a/agents-api/agents_api/activities/task_steps/__init__.py b/agents-api/agents_api/activities/task_steps/__init__.py index c85dfc0ec..363a4d5d0 100644 --- a/agents-api/agents_api/activities/task_steps/__init__.py +++ b/agents-api/agents_api/activities/task_steps/__init__.py @@ -1,8 +1,6 @@ # ruff: noqa: F401, F403, F405 from .base_evaluate import base_evaluate - -from .pg_query_step import pg_query_step from .evaluate_step import evaluate_step from .for_each_step import for_each_step from .get_value_step import get_value_step diff --git a/agents-api/agents_api/queries/chat/prepare_chat_context.py b/agents-api/agents_api/queries/chat/prepare_chat_context.py index de532844f..c3a8b8ba3 100644 --- a/agents-api/agents_api/queries/chat/prepare_chat_context.py +++ b/agents-api/agents_api/queries/chat/prepare_chat_context.py @@ -13,7 +13,7 @@ T = TypeVar("T") -sql_query =""" +sql_query = """ SELECT * FROM ( SELECT jsonb_agg(u) AS users FROM ( diff --git a/agents-api/agents_api/queries/executions/prepare_execution_input.py b/agents-api/agents_api/queries/executions/prepare_execution_input.py index 5940c4047..51ddec7a6 100644 --- a/agents-api/agents_api/queries/executions/prepare_execution_input.py +++ b/agents-api/agents_api/queries/executions/prepare_execution_input.py @@ -47,7 +47,7 @@ # SELECT to_jsonb(e) AS execution FROM ( # SELECT * FROM latest_executions # WHERE -# developer_id = $1 AND +# developer_id = $1 AND # task_id = $2 AND # execution_id = $3 # LIMIT 1 diff --git a/agents-api/agents_api/queries/tasks/create_or_update_task.py b/agents-api/agents_api/queries/tasks/create_or_update_task.py index 8be5cde84..09b4a192d 100644 --- a/agents-api/agents_api/queries/tasks/create_or_update_task.py +++ b/agents-api/agents_api/queries/tasks/create_or_update_task.py @@ -192,7 +192,8 @@ async def create_or_update_task( tool.type, tool.name, tool.description, - getattr(tool, tool.type) and getattr(tool, tool.type).model_dump(mode="json"), # spec + getattr(tool, tool.type) + and getattr(tool, tool.type).model_dump(mode="json"), # spec ] for tool in data.tools or [] ] diff --git a/agents-api/agents_api/queries/tasks/create_task.py b/agents-api/agents_api/queries/tasks/create_task.py index 5c05c3666..17eabeefe 100644 --- a/agents-api/agents_api/queries/tasks/create_task.py +++ b/agents-api/agents_api/queries/tasks/create_task.py @@ -167,7 +167,8 @@ async def create_task( tool.type, tool.name, tool.description, - getattr(tool, tool.type) and getattr(tool, tool.type).model_dump(mode="json"), # spec + getattr(tool, tool.type) + and getattr(tool, tool.type).model_dump(mode="json"), # spec ] for tool in data.tools or [] ] diff --git a/agents-api/agents_api/routers/docs/create_doc.py b/agents-api/agents_api/routers/docs/create_doc.py index 1c9f65797..cbf096355 100644 --- a/agents-api/agents_api/routers/docs/create_doc.py +++ b/agents-api/agents_api/routers/docs/create_doc.py @@ -43,9 +43,7 @@ async def create_user_doc( data=data, ) - return ResourceCreatedResponse( - id=doc.id, created_at=doc.created_at, jobs=[] - ) + return ResourceCreatedResponse(id=doc.id, created_at=doc.created_at, jobs=[]) @router.post("/agents/{agent_id}/docs", status_code=HTTP_201_CREATED, tags=["docs"]) @@ -62,6 +60,4 @@ async def create_agent_doc( data=data, ) - return ResourceCreatedResponse( - id=doc.id, created_at=doc.created_at, jobs=[] - ) + return ResourceCreatedResponse(id=doc.id, created_at=doc.created_at, jobs=[]) diff --git a/agents-api/agents_api/routers/tasks/__init__.py b/agents-api/agents_api/routers/tasks/__init__.py index e6933b27b..7e61a2ba6 100644 --- a/agents-api/agents_api/routers/tasks/__init__.py +++ b/agents-api/agents_api/routers/tasks/__init__.py @@ -1,11 +1,9 @@ # ruff: noqa: F401, F403, F405 from .create_or_update_task import create_or_update_task from .create_task import create_task - from .create_task_execution import create_task_execution from .get_execution_details import get_execution_details from .get_task_details import get_task_details - from .list_execution_transitions import list_execution_transitions from .list_task_executions import list_task_executions from .list_tasks import list_tasks diff --git a/agents-api/agents_api/routers/tasks/create_task_execution.py b/agents-api/agents_api/routers/tasks/create_task_execution.py index 2af945729..eee937b85 100644 --- a/agents-api/agents_api/routers/tasks/create_task_execution.py +++ b/agents-api/agents_api/routers/tasks/create_task_execution.py @@ -50,7 +50,6 @@ async def start_execution( ) -> tuple[Execution, WorkflowHandle]: execution_id = uuid7() - execution = await create_execution_query( developer_id=developer_id, task_id=task_id, @@ -59,7 +58,6 @@ async def start_execution( connection_pool=connection_pool, ) - execution_input = await prepare_execution_input( developer_id=developer_id, task_id=task_id, @@ -140,14 +138,12 @@ async def create_task_execution( detail="Execution count exceeded the free tier limit", ) - execution, handle = await start_execution( developer_id=x_developer_id, task_id=task_id, data=data, ) - background_tasks.add_task( create_temporal_lookup, # @@ -156,7 +152,6 @@ async def create_task_execution( workflow_handle=handle, ) - return ResourceCreatedResponse( id=execution.id, created_at=execution.created_at,