From 59462d2bdd655f57e28b9b3cf5ea355839eacf9e Mon Sep 17 00:00:00 2001 From: Diwank Tomer Date: Mon, 22 Jul 2024 16:43:17 -0400 Subject: [PATCH] wip Signed-off-by: Diwank Tomer --- .../agents_api/models/agent/create_agent.py | 44 +++++---- .../agents_api/models/agent/create_tools.py | 73 +++++++------- .../agents_api/models/agent/delete_agent.py | 83 ++++++++++++---- .../agents_api/models/agent/get_agent.py | 29 ++++-- .../agents_api/models/agent/list_agents.py | 23 ++++- .../agents_api/models/agent/patch_agent.py | 47 ++++++--- .../agents_api/models/agent/patch_tool.py | 21 ++-- .../agents_api/models/agent/update_agent.py | 67 +++++++------ agents-api/agents_api/models/utils.py | 91 +++++++++++++++--- typespec/agents/models.tsp | 9 +- typespec/common/interfaces.tsp | 96 ++++++++++++++++++- typespec/executions/endpoints.tsp | 8 +- typespec/main.tsp | 10 +- typespec/sessions/models.tsp | 5 + typespec/tasks/endpoints.tsp | 20 ++-- typespec/tasks/models.tsp | 5 + typespec/tools/endpoints.tsp | 10 +- typespec/users/models.tsp | 5 + 18 files changed, 477 insertions(+), 169 deletions(-) diff --git a/agents-api/agents_api/models/agent/create_agent.py b/agents-api/agents_api/models/agent/create_agent.py index 0dfde252e..20571061f 100644 --- a/agents-api/agents_api/models/agent/create_agent.py +++ b/agents-api/agents_api/models/agent/create_agent.py @@ -3,23 +3,23 @@ It includes functions to construct and execute datalog queries for inserting new agent records. """ -from uuid import UUID +from uuid import UUID, uuid4 +from beartype import beartype +from ...autogen.openapi_model import Agent, CreateAgentRequest from ...common.utils.cozo import cozo_process_mutate_data -from ..utils import cozo_query +from ..utils import cozo_query, verify_developer_id_query, wrap_in_class +@wrap_in_class(Agent, one=True, transform=lambda d: {"id": UUID(d.pop("agent_id")), **d}) @cozo_query +@beartype def create_agent_query( - agent_id: UUID, + *, developer_id: UUID, - name: str, - about: str, - model: str, - instructions: list[str] = [], - metadata: dict = {}, - default_settings: dict = {}, + agent_id: UUID | None = None, + create_agent: CreateAgentRequest, ) -> tuple[str, dict]: """ Constructs and executes a datalog query to create a new agent in the database. @@ -36,11 +36,18 @@ def create_agent_query( - client (CozoClient, optional): The CozoDB client instance to use for the query. Defaults to a preconfigured client instance. Returns: - pd.DataFrame: A DataFrame containing the results of the query execution. + Agent: The newly created agent record. """ - preset = default_settings.get("preset") - default_settings["preset"] = getattr(preset, "value", preset) + agent_id = agent_id or uuid4() + + # Extract the agent data from the payload + create_agent.metadata = create_agent.metadata or {} + create_agent.instructions = create_agent.instructions if isinstance(create_agent.instructions, list) else [create_agent.instructions] + create_agent.default_settings = create_agent.default_settings or {} + + agent_data = create_agent.model_dump() + default_settings = agent_data.pop("default_settings") settings_cols, settings_vals = cozo_process_mutate_data( { @@ -61,8 +68,8 @@ def create_agent_query( # 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] <- [ - [$agent_id, $developer_id, $model, $name, $about, $metadata, $instructions] + ?[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 { @@ -73,11 +80,14 @@ def create_agent_query( about, metadata, instructions, + created_at, + updated_at, } :returning """ queries = [ + verify_developer_id_query(developer_id), default_settings_query, agent_query, ] @@ -91,10 +101,6 @@ def create_agent_query( "settings_vals": settings_vals, "agent_id": str(agent_id), "developer_id": str(developer_id), - "model": model, - "name": name, - "about": about, - "metadata": metadata, - "instructions": instructions, + **agent_data, }, ) diff --git a/agents-api/agents_api/models/agent/create_tools.py b/agents-api/agents_api/models/agent/create_tools.py index 42332b8dd..f559d57a2 100644 --- a/agents-api/agents_api/models/agent/create_tools.py +++ b/agents-api/agents_api/models/agent/create_tools.py @@ -2,65 +2,68 @@ from uuid import UUID, uuid4 +from beartype import beartype -from ...autogen.openapi_model import FunctionDef +from ...autogen.openapi_model import CreateToolRequest, Tool -from ..utils import cozo_query +from ..utils import cozo_query, verify_developer_id_query, verify_developer_owns_resource_query, wrap_in_class +@wrap_in_class( + Tool, + transform=lambda d: {"id": UUID(d.pop("tool_id")), d["tool_type"]: d.pop("spec"), "type": d.pop("tool_type"), **d}, +) @cozo_query +@beartype def create_tools_query( + *, + developer_id: UUID, agent_id: UUID, - functions: list[FunctionDef], - embeddings: list[list[float]], + create_tools: list[CreateToolRequest], ) -> tuple[str, dict]: """ Constructs a datalog query for inserting tool records into the 'agent_functions' relation in the CozoDB. Parameters: - agent_id (UUID): The unique identifier for the agent. - - functions (list[FunctionDef]): A list of function definitions to be inserted. - - embeddings (list[list[float]]): A list of embeddings corresponding to each function. + - create_tools (list[CreateToolRequest]): A list of function definitions to be inserted. Returns: - - pd.DataFrame: A DataFrame containing the results of the query execution. + list[Tool] """ - # Ensure the number of functions matches the number of embeddings - assert len(functions) == len(embeddings) - # Construct the input records for the datalog query - functions_input: list[list] = [] - - for function, embedding in zip(functions, embeddings): - parameters = function.parameters.model_dump() - functions_input.append( - [ - str(agent_id), - str(uuid4()), - function.name, - function.description or "", - parameters, - embedding, - ] - ) + tools_data = [ + [ + str(agent_id), + str(uuid4()), + tool.type, + tool.name, + getattr(tool, tool.type).dict(), + ] + for tool in create_tools + ] # Datalog query for inserting new tool records into the 'agent_functions' relation - query = """ - input[agent_id, tool_id, name, description, parameters, embedding] <- $records - ?[agent_id, tool_id, name, description, parameters, embedding, updated_at] := - input[agent_id, tool_id, name, description, parameters, embedding], - updated_at = now(), + create_tools_query = """ + ?[agent_id, tool_id, tool_type, name, spec] <- $records - :insert agent_functions { + :insert tools { agent_id, tool_id, + tool_type, name, - description, - parameters, - embedding, - updated_at, + spec, } :returning """ - return (query, {"records": functions_input}) + queries = [ + verify_developer_id_query(developer_id), + verify_developer_owns_resource_query(developer_id, "agents", agent_id=agent_id), + create_tools_query, + ] + + query = "}\n\n{\n".join(queries) + query = f"{{ {query} }}" + + return (query, {"records": tools_data}) diff --git a/agents-api/agents_api/models/agent/delete_agent.py b/agents-api/agents_api/models/agent/delete_agent.py index 6557c63d6..071bd6b98 100644 --- a/agents-api/agents_api/models/agent/delete_agent.py +++ b/agents-api/agents_api/models/agent/delete_agent.py @@ -4,33 +4,79 @@ from uuid import UUID +from beartype import beartype -from ..utils import cozo_query +from ...autogen.openapi_model import ResourceDeletedResponse +from ...common.utils.datetime import utcnow +from ..utils import ( + cozo_query, + verify_developer_id_query, + verify_developer_owns_resource_query, + wrap_in_class, +) -""" -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. +@wrap_in_class( + ResourceDeletedResponse, + one=True, + transform=lambda d: {"id": UUID(d.pop("agent_id")), "deleted_at": utcnow(), "jobs": []}, +) +@cozo_query +@beartype +def delete_agent_query(*, developer_id: UUID, agent_id: UUID) -> tuple[str, dict]: + """ + Constructs and returns a datalog query for deleting an agent and its default settings from the database. -Returns: -- tuple[str, dict]: A DataFrame containing the results of the deletion query. -""" + 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. + """ -@cozo_query -def delete_agent_query(developer_id: UUID, agent_id: UUID) -> tuple[str, dict]: - query = """ - { + queries = [ + verify_developer_id_query(developer_id), + verify_developer_owns_resource_query(developer_id, "agents", agent_id=agent_id), + """ + # Delete docs + ?[agent_id, doc_id] := + *agent_docs{ + agent_id, + doc_id, + }, agent_id = to_uuid($agent_id) + + :delete agent_docs { + agent_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]] @@ -38,6 +84,11 @@ def delete_agent_query(developer_id: UUID, agent_id: UUID) -> tuple[str, dict]: developer_id, agent_id } - }""" + :returning + """, + ] + + query = "}\n\n{\n".join(queries) + query = f"{{ {query} }}" return (query, {"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 index 325df7930..281283836 100644 --- a/agents-api/agents_api/models/agent/get_agent.py +++ b/agents-api/agents_api/models/agent/get_agent.py @@ -1,11 +1,21 @@ from uuid import UUID +from beartype import beartype -from ..utils import cozo_query +from ...autogen.openapi_model import Agent +from ..utils import ( + cozo_query, + verify_developer_id_query, + verify_developer_owns_resource_query, + wrap_in_class, +) + +@wrap_in_class(Agent, one=True) @cozo_query -def get_agent_query(developer_id: UUID, agent_id: UUID) -> tuple[str, dict]: +@beartype +def get_agent_query(*, developer_id: UUID, agent_id: UUID) -> tuple[str, dict]: """ Fetches agent details and default settings from the database. @@ -17,13 +27,12 @@ def get_agent_query(developer_id: UUID, agent_id: UUID) -> tuple[str, dict]: - client (CozoClient, optional): The database client used to execute the query. Returns: - - pd.DataFrame: A DataFrame containing the agent details and default settings. + - 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. - query = """ - { + get_query = """ input[agent_id, developer_id] <- [[to_uuid($agent_id), to_uuid($developer_id)]] ?[ @@ -69,9 +78,17 @@ def get_agent_query(developer_id: UUID, agent_id: UUID) -> tuple[str, dict]: "min_p": min_p, "preset": preset, } - } """ + queries = [ + verify_developer_id_query(developer_id), + verify_developer_owns_resource_query(developer_id, "agents", agent_id=agent_id), + get_query, + ] + + query = "}\n\n{\n".join(queries) + query = f"{{ {query} }}" + # Execute the constructed datalog query using the provided CozoClient. # The result is returned as a pandas DataFrame. return (query, {"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 index 838b7acdb..4fdb8403c 100644 --- a/agents-api/agents_api/models/agent/list_agents.py +++ b/agents-api/agents_api/models/agent/list_agents.py @@ -1,13 +1,23 @@ from typing import Any from uuid import UUID +from beartype import beartype + +from ...autogen.openapi_model import Agent from ...common.utils import json -from ..utils import cozo_query +from ..utils import ( + cozo_query, + verify_developer_id_query, + wrap_in_class, +) +@wrap_in_class(Agent) @cozo_query +@beartype def list_agents_query( + *, developer_id: UUID, limit: int = 100, offset: int = 0, @@ -35,8 +45,9 @@ def list_agents_query( ) # Datalog query to retrieve agent information based on filters, sorted by creation date in descending order. - query = f""" - {{ + queries = [ + verify_developer_id_query(developer_id), + f""" input[developer_id] <- [[to_uuid($developer_id)]] ?[ @@ -65,7 +76,11 @@ def list_agents_query( :limit $limit :offset $offset :sort -created_at - }}""" + """, + ] + + query = "}\n\n{\n".join(queries) + query = f"{{ {query} }}" return ( query, diff --git a/agents-api/agents_api/models/agent/patch_agent.py b/agents-api/agents_api/models/agent/patch_agent.py index ee2ad7103..ea2b43355 100644 --- a/agents-api/agents_api/models/agent/patch_agent.py +++ b/agents-api/agents_api/models/agent/patch_agent.py @@ -1,17 +1,30 @@ from uuid import UUID +from beartype import beartype +from ...autogen.openapi_model import PatchAgentRequest, ResourceUpdatedResponse from ...common.utils.cozo import cozo_process_mutate_data -from ..utils import cozo_query from ...common.utils.datetime import utcnow - - +from ..utils import ( + cozo_query, + verify_developer_id_query, + verify_developer_owns_resource_query, + wrap_in_class, +) + + +@wrap_in_class( + ResourceUpdatedResponse, + one=True, + transform=lambda d: {"id": d["agent_id"], "jobs": [], **d}, +) @cozo_query +@beartype def patch_agent_query( + *, agent_id: UUID, developer_id: UUID, - default_settings: dict = {}, - **update_data, + patch_agent: PatchAgentRequest, ) -> tuple[str, dict]: """Patches agent data based on provided updates. @@ -22,11 +35,14 @@ def patch_agent_query( **update_data: Arbitrary keyword arguments representing data to update. Returns: - pd.DataFrame: The result of the query execution. + ResourceUpdatedResponse: The updated agent data. """ + update_data = patch_agent.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}, @@ -37,7 +53,6 @@ def patch_agent_query( ) agent_update_query = f""" - {{ # update the agent input[{agent_update_cols}] <- $agent_update_vals @@ -47,13 +62,14 @@ def patch_agent_query( agent_id: to_uuid($agent_id), metadata: md, }}, + # agent_id = to_uuid($agent_id), metadata = concat(md, $metadata) :update agents {{ - {agent_update_cols}, metadata, + {agent_update_cols}, + metadata, }} :returning - }} """ # Construct the query for updating agent's default settings in the database. @@ -66,14 +82,12 @@ def patch_agent_query( ) 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. @@ -83,10 +97,17 @@ def patch_agent_query( if len(default_settings) != 0: queries.insert(0, settings_update_query) - combined_query = "\n".join(queries) + queries = [ + verify_developer_id_query(developer_id), + verify_developer_owns_resource_query(developer_id, "agents", agent_id=agent_id), + *queries, + ] + + query = "}\n\n{\n".join(queries) + query = f"{{ {query} }}" return ( - combined_query, + query, { "agent_update_vals": agent_update_vals, "settings_vals": settings_vals, diff --git a/agents-api/agents_api/models/agent/patch_tool.py b/agents-api/agents_api/models/agent/patch_tool.py index a565d9f7b..5462d3630 100644 --- a/agents-api/agents_api/models/agent/patch_tool.py +++ b/agents-api/agents_api/models/agent/patch_tool.py @@ -1,27 +1,34 @@ from uuid import UUID -from ...autogen.openapi_model import FunctionDef +from beartype import beartype -from ..utils import cozo_query +from ...autogen.openapi_model import PatchToolRequest, ResourceUpdatedResponse +from ..utils import cozo_query, verify_developer_id_query, wrap_in_class + +@wrap_in_class( + ResourceUpdatedResponse, + one=True, + transform=lambda d: {"id": d["agent_id"], "jobs": [], **d}, +) @cozo_query +@beartype def patch_tool_by_id_query( - agent_id: UUID, tool_id: UUID, function: FunctionDef, embedding: list[float] + *, developer_id: UUID, tool_id: UUID, update_tool: PatchToolRequest ) -> tuple[str, dict]: """ # Execute the datalog query and return the results as a DataFrame Updates the tool information for a given agent and tool ID in the 'cozodb' database. Parameters: - - agent_id (UUID): The unique identifier of the agent. + - developer_id (UUID): The unique identifier of the developer. - tool_id (UUID): The unique identifier of the tool to be updated. - - function (FunctionDef): The function definition containing the new tool information. - - embedding (list[float]): The embedding vector associated with the tool. + - update_tool (PatchToolRequest): The request payload containing the updated tool information. Returns: - - pd.DataFrame: A DataFrame containing the result of the update operation. + - ResourceUpdatedResponse: The updated tool data. """ # Agent update query # Convert the function definition to a dictionary for easier manipulation diff --git a/agents-api/agents_api/models/agent/update_agent.py b/agents-api/agents_api/models/agent/update_agent.py index c1b950b79..cdda45ae3 100644 --- a/agents-api/agents_api/models/agent/update_agent.py +++ b/agents-api/agents_api/models/agent/update_agent.py @@ -5,17 +5,31 @@ from uuid import UUID +from beartype import beartype +from ...autogen.openapi_model import UpdateAgentRequest, ResourceUpdatedResponse from ...common.utils.cozo import cozo_process_mutate_data -from ..utils import cozo_query - - +from ...common.utils.datetime import utcnow +from ..utils import ( + cozo_query, + verify_developer_id_query, + verify_developer_owns_resource_query, + wrap_in_class, +) + + +@wrap_in_class( + ResourceUpdatedResponse, + one=True, + transform=lambda d: {"id": d["agent_id"], "jobs": [], **d}, +) @cozo_query +@beartype def update_agent_query( + *, agent_id: UUID, developer_id: UUID, - default_settings: dict = {}, - **update_data, + update_agent: UpdateAgentRequest, ) -> tuple[str, dict]: """ Constructs and executes a datalog query to update an agent and its default settings in the 'cozodb' database. @@ -23,30 +37,26 @@ def update_agent_query( 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. - - default_settings (dict, optional): A dictionary of default settings to be updated for the agent. Defaults to an empty dict. + - update_agent (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. - - **update_data: Variable keyword arguments representing additional agent data to be updated. Returns: - - pd.DataFrame: A DataFrame containing the result of the update operation. + ResourceUpdatedResponse: The updated agent data. """ + default_settings = update_agent.default_settings.model_dump(exclude_none=True) + update_data = update_agent.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", []) - - # Assertion query to check if the agent exists - assertion_query = """ - ?[developer_id, agent_id] := - *agents { - developer_id, - agent_id, - }, - developer_id = to_uuid($developer_id), - agent_id = to_uuid($agent_id), - # Assertion to ensure the agent exists before updating. - :assert some - """ + 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 @@ -59,7 +69,6 @@ def update_agent_query( ) agent_update_query = f""" - {{ # update the agent input[{agent_update_cols}] <- $agent_update_vals original[created_at] := *agents{{ @@ -79,7 +88,6 @@ def update_agent_query( {agent_update_cols} }} :returning - }} """ # Construct the settings update part of the query if `default_settings` are provided. @@ -92,14 +100,12 @@ def update_agent_query( ) 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. @@ -110,10 +116,17 @@ def update_agent_query( queries.insert(0, settings_update_query) # Combine the assertion query with the update queries - combined_query = "{" + assertion_query + "} " + "\n".join(queries) + queries = [ + verify_developer_id_query(developer_id), + verify_developer_owns_resource_query(developer_id, "agents", agent_id=agent_id), + *queries, + ] + + query = "}\n\n{\n".join(queries) + query = f"{{ {query} }}" return ( - combined_query, + query, { "agent_update_vals": agent_update_vals, "settings_vals": settings_vals, diff --git a/agents-api/agents_api/models/utils.py b/agents-api/agents_api/models/utils.py index fe269b17c..e6abd5189 100644 --- a/agents-api/agents_api/models/utils.py +++ b/agents-api/agents_api/models/utils.py @@ -1,6 +1,9 @@ from functools import wraps -from typing import Callable, ParamSpec +from typing import Callable, ParamSpec, Type +from uuid import UUID +from beartype import beartype +from pydantic import BaseModel import pandas as pd from ..clients.cozo import client as cozo_client @@ -9,18 +12,84 @@ P = ParamSpec("P") -def cozo_query(func: Callable[P, tuple[str, dict]]): +def verify_developer_id_query(developer_id: UUID | str) -> str: + return f""" + ?[developer_id] := + *developers{{ + developer_id, + }}, developer_id = to_uuid("{str(developer_id)}") + + :assert some """ - 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. + +def verify_developer_owns_resource_query(developer_id: UUID | str, resource: str, **resource_id: dict) -> str: + + resource_id_key, resource_id_value = next(iter(resource_id.items())) + + return f""" + ?[{resource_id_key}] := + *{resource}{{ + developer_id, + {resource_id_key}, + }}, developer_id = to_uuid("{str(developer_id)}"), + {resource_id_key} = to_uuid("{str(resource_id_value)}") + + :assert some """ - @wraps(func) - def wrapper(*args, client=cozo_client, **kwargs) -> pd.DataFrame: - query, variables = func(*args, **kwargs) - return client.run(query, variables) - return wrapper +def cozo_query(func: Callable[P, tuple[str, dict]] | None = None, debug: bool | None = None): + def cozo_query_dec(func: Callable[P, tuple[str, 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. + """ + + @wraps(func) + def wrapper(*args, client=cozo_client, **kwargs) -> pd.DataFrame: + query, variables = func(*args, **kwargs) + + if debug: + from pprint import pprint + pprint(dict(query=query, variables=variables)) + + return client.run(query, variables) + + 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[BaseModel], + one: bool = False, + transform: Callable[[dict], dict] | None = None, +): + + @beartype + def decorator(func: Callable[..., pd.DataFrame]): + @wraps(func) + @beartype + def wrapper(*args, **kwargs) -> list[cls] | cls: + df = func(*args, **kwargs) + + # Convert df to list of dicts + data = df.to_dict(orient="records") + + nonlocal transform + transform = transform or (lambda x: x) + if one: + return cls(**transform(data[0])) + + return [cls(**item) for item in map(transform, data)] + + setattr(cls, func.__name__, wrapper) + return wrapper + + return decorator diff --git a/typespec/agents/models.tsp b/typespec/agents/models.tsp index 7a925fa1f..3d8642d2e 100644 --- a/typespec/agents/models.tsp +++ b/typespec/agents/models.tsp @@ -1,6 +1,10 @@ +import "@typespec/http"; + import "../common"; import "../chat"; +using TypeSpec.Http; + using Common; using Chat; @@ -44,13 +48,10 @@ model PatchAgentRequest is UpdateAgentRequest {} /** Payload for creating a agent (and associated documents) */ model CreateAgentRequest { ...UpdateAgentRequest; - - /** Documents to index for this agent. (Max: 100 items) */ - @maxItems(100) - docs: unknown[] = #[]; } model CreateOrUpdateAgentRequest { + @path id: uuid; ...UpdateAgentRequest; } \ No newline at end of file diff --git a/typespec/common/interfaces.tsp b/typespec/common/interfaces.tsp index 2d2c5fb21..d83cf81ab 100644 --- a/typespec/common/interfaces.tsp +++ b/typespec/common/interfaces.tsp @@ -41,7 +41,7 @@ interface CreateOrUpdateEndpoint< CreateOrUpdateType, DocString extends valueof string = "Create or update a resource (ID is required in payload; existing resource will be overwritten)" > { - @put + @post @doc(DocString) createOrUpdate(@header contentType: yaml | json, ...CreateOrUpdateType): { @statusCode _: "200"; @@ -167,10 +167,33 @@ interface ChildCreateEndpoint< }; } +interface ChildCreateOrUpdateEndpoint< + CreateOrUpdateType, + DocString extends valueof string = "Create or update a resource owned by the given parent" +> { + @post + @doc(DocString) + createOrUpdate( + @header contentType: yaml | json, + + @path + @doc("ID of parent resource") + parent_id: uuid, + + ...CreateOrUpdateType, + ): { + @statusCode _: "200"; + + @body + @doc("Details of the resource updated along with ID") + body: ResourceUpdatedResponse; + }; +} + interface ChildDeleteEndpoint< DocString extends valueof string = "Delete a resource owned by the given parent using its id" > { - @post + @delete @doc( DocString ) @@ -179,13 +202,13 @@ interface ChildDeleteEndpoint< @doc( "ID of parent resource", ) - parent_id: uuid, + id: uuid, @path @doc( "ID of the resource to be deleted", ) - id: uuid, + child_id: uuid, ): { @statusCode _: "202"; @@ -196,3 +219,68 @@ interface ChildDeleteEndpoint< body: ResourceDeletedResponse; }; } + +interface ChildUpdateEndpoint< + UpdateType, + DocString extends valueof string = "Update a resource owned by the given parent using its id" +> { + @put + @doc( + DocString + ) + update( + @path + @doc( + "ID of parent resource", + ) + id: uuid, + + @path + @doc( + "ID of the resource to be updated", + ) + child_id: uuid, + + ...UpdateType, + ): { + @statusCode _: "200"; + + @body + @doc( + "Details of the resource updated along with ID", + ) + body: ResourceUpdatedResponse; + }; +} + +interface ChildPatchEndpoint< + PatchType, + DocString extends valueof string = "Patch a resource owned by the given parent using its id" +> { + @patch + @doc( + DocString + ) + patch( + @path + @doc( + "ID of parent resource", + ) + id: uuid, + + @path + @doc( + "ID of the resource to be patched", + ) + child_id: uuid, + ...PatchType, + ): { + @statusCode _: "200"; + + @body + @doc( + "Details of the resource patched along with ID", + ) + body: ResourceUpdatedResponse; + }; +} diff --git a/typespec/executions/endpoints.tsp b/typespec/executions/endpoints.tsp index 0fa7b43e5..8869b7022 100644 --- a/typespec/executions/endpoints.tsp +++ b/typespec/executions/endpoints.tsp @@ -10,11 +10,13 @@ namespace Executions; // interface Endpoints - extends UpdateEndpoint< + extends GetEndpoint {} + +interface TaskEndpoints + extends ChildUpdateEndpoint< UpdateExecutionRequest, "Update an existing Execution" - >, - GetEndpoint {} + > {} interface TransitionEndpoints extends ChildLimitOffsetPagination< diff --git a/typespec/main.tsp b/typespec/main.tsp index 6698fac76..93f1a472d 100644 --- a/typespec/main.tsp +++ b/typespec/main.tsp @@ -92,11 +92,14 @@ namespace Api { @route("/docs") interface IndividualDocsRoute extends Docs.IndividualDocEndpoints {} - @route("/tasks") + @route("/agents/{id}/tasks") interface TasksRoute extends Tasks.Endpoints {} + @route("/agents/{parent_id}/tasks") + interface TasksCreateOrUpdateRoute extends Tasks.CreateOrUpdateEndpoints {} + @route("/tasks/{id}/executions") - interface TaskExecutionsRoute extends Tasks.ExecutionEndpoints {} + interface TaskExecutionsRoute extends Tasks.ExecutionEndpoints, Executions.TaskEndpoints {} @route("/executions") interface ExecutionsRoute extends Executions.Endpoints {} @@ -104,9 +107,6 @@ namespace Api { @route("/executions/{id}/transitions") interface ExecutionTransitionsRoute extends Executions.TransitionEndpoints {} - @route("/tools") - interface ToolRoute extends Tools.Endpoints {} - @route("/jobs") interface JobRoute extends Jobs.Endpoints {} } diff --git a/typespec/sessions/models.tsp b/typespec/sessions/models.tsp index f19ec1d99..ed58f55c6 100644 --- a/typespec/sessions/models.tsp +++ b/typespec/sessions/models.tsp @@ -1,5 +1,9 @@ +import "@typespec/http"; + import "../common"; +using TypeSpec.Http; + using Common; namespace Sessions; @@ -147,6 +151,7 @@ model CreateSessionRequest { @withVisibility("create", "update") model CreateOrUpdateSessionRequest { + @path id: uuid; ...CreateSessionRequest; diff --git a/typespec/tasks/endpoints.tsp b/typespec/tasks/endpoints.tsp index 289dc7c11..5d656c78a 100644 --- a/typespec/tasks/endpoints.tsp +++ b/typespec/tasks/endpoints.tsp @@ -15,22 +15,24 @@ namespace Tasks; // TASK ENDPOINTS // +interface CreateOrUpdateEndpoints + extends ChildCreateOrUpdateEndpoint< + CreateOrUpdateTaskRequest, + "Create or update a task" + > {} + interface Endpoints - extends UpdateEndpoint< + extends ChildUpdateEndpoint< UpdateTaskRequest, "Update an existing task (overwrite existing values)" >, - PatchEndpoint< + ChildPatchEndpoint< PatchTaskRequest, "Update an existing task (merges with existing values)" >, - DeleteEndpoint<"Delete a task by its id">, - LimitOffsetPagination, - CreateOrUpdateEndpoint< - CreateOrUpdateTaskRequest, - "Create or update a task" - >, - CreateEndpoint {} + ChildDeleteEndpoint<"Delete a task by its id">, + ChildLimitOffsetPagination, + ChildCreateEndpoint {} interface ExecutionEndpoints extends ChildCreateEndpoint< diff --git a/typespec/tasks/models.tsp b/typespec/tasks/models.tsp index 0135da392..7e7aab2ed 100644 --- a/typespec/tasks/models.tsp +++ b/typespec/tasks/models.tsp @@ -1,9 +1,13 @@ +import "@typespec/http"; + import "../agents"; import "../common"; import "../chat"; import "../entries"; import "../tools"; +using TypeSpec.Http; + using Agents; using Chat; using Common; @@ -168,6 +172,7 @@ model CreateTaskRequest { } model CreateOrUpdateTaskRequest { + @path id: uuid; ...CreateTaskRequest; } diff --git a/typespec/tools/endpoints.tsp b/typespec/tools/endpoints.tsp index 30e03c6d9..819af96c7 100644 --- a/typespec/tools/endpoints.tsp +++ b/typespec/tools/endpoints.tsp @@ -11,11 +11,9 @@ namespace Tools; // TOOL ENDPOINTS // -interface Endpoints - extends UpdateEndpoint, - PatchEndpoint, - DeleteEndpoint<"Delete an existing tool by id"> {} - interface AgentEndpoints extends ChildLimitOffsetPagination, - ChildCreateEndpoint {} + ChildCreateEndpoint, + ChildUpdateEndpoint, + ChildPatchEndpoint, + ChildDeleteEndpoint<"Delete an existing tool by id"> {} diff --git a/typespec/users/models.tsp b/typespec/users/models.tsp index 1d53311d4..218663cff 100644 --- a/typespec/users/models.tsp +++ b/typespec/users/models.tsp @@ -1,5 +1,9 @@ +import "@typespec/http"; + import "../common"; +using TypeSpec.Http; + using Common; namespace Users; @@ -40,6 +44,7 @@ model CreateUserRequest { } model CreateOrUpdateUserRequest { + @path id: uuid; ...UpdateUserRequest; } \ No newline at end of file