From 7cd1c321f208b6354b18fa2aeb0a40386d39a523 Mon Sep 17 00:00:00 2001 From: Diwank Tomer Date: Wed, 29 May 2024 19:31:06 -0700 Subject: [PATCH 1/7] wip(agents-api): cozo queries for tasks Signed-off-by: Diwank Tomer --- .../agents_api/models/execution/__init__.py | 0 agents-api/agents_api/models/task/__init__.py | 0 agents-api/agents_api/models/task/get_task.py | 28 ++++++++++++++-- .../agents_api/models/task/list_tasks.py | 32 +++++++++++++++++-- 4 files changed, 55 insertions(+), 5 deletions(-) create mode 100644 agents-api/agents_api/models/execution/__init__.py create mode 100644 agents-api/agents_api/models/task/__init__.py diff --git a/agents-api/agents_api/models/execution/__init__.py b/agents-api/agents_api/models/execution/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/agents-api/agents_api/models/task/__init__.py b/agents-api/agents_api/models/task/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/agents-api/agents_api/models/task/get_task.py b/agents-api/agents_api/models/task/get_task.py index be9509353..f89da19ad 100644 --- a/agents-api/agents_api/models/task/get_task.py +++ b/agents-api/agents_api/models/task/get_task.py @@ -4,6 +4,28 @@ @cozo_query -def get_task_query(developer_id: UUID, task_id: UUID) -> tuple[str, dict]: - query = """""" - return (query, {"developer_id": str(developer_id), "task_id": str(task_id)}) +def get_task_query(agent_id: UUID, task_id: UUID) -> tuple[str, dict]: + query = """ + ?[ + name, + description, + input_schema, + tools_available, + workflows, + created_at, + updated_at, + ] := *tasks { + agent_id: to_uuid($agent_id), + task_id: to_uuid($task_id), + updated_at_ms, + name, + description, + input_schema, + tools_available, + workflows, + created_at, + @ 'NOW' + }, updated_at = to_int(updated_at_ms) / 1000 + """ + + return (query, {"agent_id": str(agent_id), "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 index 33fbf3ae0..64f3d8010 100644 --- a/agents-api/agents_api/models/task/list_tasks.py +++ b/agents-api/agents_api/models/task/list_tasks.py @@ -7,7 +7,7 @@ @cozo_query def list_tasks_query( - developer_id: UUID, + agent_id: UUID, limit: int = 100, offset: int = 0, # metadata_filter: dict[str, Any] = {}, @@ -21,4 +21,32 @@ def list_tasks_query( Returns: pd.DataFrame: A DataFrame containing the queried task data. """ - pass + + query = """ + ?[ + task_id, + name, + description, + input_schema, + tools_available, + workflows, + created_at, + updated_at, + ] := *tasks { + agent_id: to_uuid($agent_id), + task_id, + updated_at_ms, + name, + description, + input_schema, + tools_available, + workflows, + created_at, + @ 'NOW' + }, updated_at = to_int(updated_at_ms) / 1000 + + :limit $limit + :offset $offset + """ + + return (query, {"agent_id": str(agent_id), "limit": limit, "offset": offset}) From 5406a2b884fe36687578d40a07985409d46958ac Mon Sep 17 00:00:00 2001 From: Diwank Tomer Date: Wed, 29 May 2024 19:44:03 -0700 Subject: [PATCH 2/7] doc: Add documentation for the new endpoints --- agents-api/agents_api/routers/tasks/routers.py | 10 ++++++++++ 1 file changed, 10 insertions(+) diff --git a/agents-api/agents_api/routers/tasks/routers.py b/agents-api/agents_api/routers/tasks/routers.py index 8f84dc7d0..0922d82b0 100644 --- a/agents-api/agents_api/routers/tasks/routers.py +++ b/agents-api/agents_api/routers/tasks/routers.py @@ -58,6 +58,9 @@ async def create_task( x_developer_id: Annotated[UUID4, Depends(get_developer_id)], ) -> ResourceCreatedResponse: task_id = uuid4() + + # TODO: Do thorough validation of the task spec + resp: pd.DataFrame = create_task_query( task_id=task_id, developer_id=x_developer_id, @@ -73,6 +76,9 @@ async def create_task_execution( request: CreateExecution, x_developer_id: Annotated[UUID4, Depends(get_developer_id)], ) -> ResourceCreatedResponse: + # TODO: Do thorough validation of the input against task input schema + # DO NOT let the user specify the status + resp = create_execution_query() return ResourceCreatedResponse( id=resp["execution_id"][0], created_at=resp["created_at"][0] @@ -98,3 +104,7 @@ async def get_execution_transition( resp = get_execution_transition_query(execution_id, transition_id, x_developer_id) pass + +# @router.get("/executions/{execution_id}") # -> get the execution object +# @router.put("/executions/{execution_id}/transitions/{transition_id}") # -> update a waiting transition +# @router.get("/executions/{execution_id}/transitions") # -> get all transitions for an execution From 05ec664b7989fcd61a6612905d56dfce968db2f2 Mon Sep 17 00:00:00 2001 From: Diwank Tomer Date: Wed, 29 May 2024 19:47:27 -0700 Subject: [PATCH 3/7] wip: Boilerplate for execution transitions queries --- .../models/execution/get_execution.py | 25 +++++++++++++++++++ .../execution/list_execution_transitions.py | 19 ++++++++++++++ .../execution/update_execution_transition.py | 19 ++++++++++++++ 3 files changed, 63 insertions(+) create mode 100644 agents-api/agents_api/models/execution/get_execution.py create mode 100644 agents-api/agents_api/models/execution/list_execution_transitions.py create mode 100644 agents-api/agents_api/models/execution/update_execution_transition.py diff --git a/agents-api/agents_api/models/execution/get_execution.py b/agents-api/agents_api/models/execution/get_execution.py new file mode 100644 index 000000000..4abccb114 --- /dev/null +++ b/agents-api/agents_api/models/execution/get_execution.py @@ -0,0 +1,25 @@ +from uuid import UUID + +from ..utils import cozo_query +from typing import Literal, Dict, Any + + +@cozo_query +def create_execution_query( + task_id: UUID, + execution_id: UUID, + status: Literal[ + "queued", "starting", "running", "waiting-for-input", "success", "failed" + ] = "queued", + arguments: Dict[str, Any] = {}, +) -> tuple[str, dict]: + query = """""" + return ( + query, + { + "task_id": str(task_id), + "execution_id": str(execution_id), + "status": status, + "arguments": arguments, + }, + ) diff --git a/agents-api/agents_api/models/execution/list_execution_transitions.py b/agents-api/agents_api/models/execution/list_execution_transitions.py new file mode 100644 index 000000000..5b0c2f942 --- /dev/null +++ b/agents-api/agents_api/models/execution/list_execution_transitions.py @@ -0,0 +1,19 @@ +from uuid import UUID + +from ..utils import cozo_query + + +@cozo_query +def get_execution_transition_query( + execution_id: UUID, transition_id: UUID, developer_id: UUID +) -> tuple[str, dict]: + + query = """""" + return ( + query, + { + "execution_id": str(execution_id), + "transition_id": str(transition_id), + "developer_id": str(developer_id), + }, + ) diff --git a/agents-api/agents_api/models/execution/update_execution_transition.py b/agents-api/agents_api/models/execution/update_execution_transition.py new file mode 100644 index 000000000..5b0c2f942 --- /dev/null +++ b/agents-api/agents_api/models/execution/update_execution_transition.py @@ -0,0 +1,19 @@ +from uuid import UUID + +from ..utils import cozo_query + + +@cozo_query +def get_execution_transition_query( + execution_id: UUID, transition_id: UUID, developer_id: UUID +) -> tuple[str, dict]: + + query = """""" + return ( + query, + { + "execution_id": str(execution_id), + "transition_id": str(transition_id), + "developer_id": str(developer_id), + }, + ) From e630a558276776c74bcd591889ae6bde5cf22015 Mon Sep 17 00:00:00 2001 From: Diwank Tomer Date: Wed, 29 May 2024 19:50:28 -0700 Subject: [PATCH 4/7] wip: add boilerplate execution transitions queries --- .../models/execution/update_execution.py | 25 +++++++++++++++++++ .../execution/update_execution_status.py | 25 +++++++++++++++++++ .../agents_api/models/task/update_task.py | 24 ++++++++++++++++++ 3 files changed, 74 insertions(+) create mode 100644 agents-api/agents_api/models/execution/update_execution.py create mode 100644 agents-api/agents_api/models/execution/update_execution_status.py create mode 100644 agents-api/agents_api/models/task/update_task.py diff --git a/agents-api/agents_api/models/execution/update_execution.py b/agents-api/agents_api/models/execution/update_execution.py new file mode 100644 index 000000000..4abccb114 --- /dev/null +++ b/agents-api/agents_api/models/execution/update_execution.py @@ -0,0 +1,25 @@ +from uuid import UUID + +from ..utils import cozo_query +from typing import Literal, Dict, Any + + +@cozo_query +def create_execution_query( + task_id: UUID, + execution_id: UUID, + status: Literal[ + "queued", "starting", "running", "waiting-for-input", "success", "failed" + ] = "queued", + arguments: Dict[str, Any] = {}, +) -> tuple[str, dict]: + query = """""" + return ( + query, + { + "task_id": str(task_id), + "execution_id": str(execution_id), + "status": status, + "arguments": arguments, + }, + ) diff --git a/agents-api/agents_api/models/execution/update_execution_status.py b/agents-api/agents_api/models/execution/update_execution_status.py new file mode 100644 index 000000000..4abccb114 --- /dev/null +++ b/agents-api/agents_api/models/execution/update_execution_status.py @@ -0,0 +1,25 @@ +from uuid import UUID + +from ..utils import cozo_query +from typing import Literal, Dict, Any + + +@cozo_query +def create_execution_query( + task_id: UUID, + execution_id: UUID, + status: Literal[ + "queued", "starting", "running", "waiting-for-input", "success", "failed" + ] = "queued", + arguments: Dict[str, Any] = {}, +) -> tuple[str, dict]: + query = """""" + return ( + query, + { + "task_id": str(task_id), + "execution_id": str(execution_id), + "status": status, + "arguments": arguments, + }, + ) diff --git a/agents-api/agents_api/models/task/update_task.py b/agents-api/agents_api/models/task/update_task.py new file mode 100644 index 000000000..0984c7b59 --- /dev/null +++ b/agents-api/agents_api/models/task/update_task.py @@ -0,0 +1,24 @@ +""" +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 uuid import UUID +from typing import List, Optional, Dict, Any + + +from ..utils import cozo_query + + +@cozo_query +def create_task_query( + task_id: UUID, + developer_id: UUID, + agent_id: UUID, + name: str, + description: str, + input_schema: Dict[str, any], + tools_available: List[UUID] = [], + workflows: List[Dict[str, Any]] = [], +) -> tuple[str, dict]: + pass From 6f9dab5104f29a3e04b0f5d22270a0d7dc327ce2 Mon Sep 17 00:00:00 2001 From: Siddharth Balyan Date: Fri, 31 May 2024 01:44:03 +0530 Subject: [PATCH 5/7] wip: task feature routes and models --- .../agents_api/autogen/openapi_model.py | 4 +- .../models/execution/create_execution.py | 22 ++- .../models/execution/get_execution.py | 25 ++- .../models/execution/get_execution_status.py | 19 +- .../execution/get_execution_transition.py | 16 +- .../execution/list_execution_transitions.py | 23 ++- .../models/execution/list_executions.py | 35 ++++ .../models/execution/update_execution.py | 25 --- .../execution/update_execution_status.py | 23 ++- .../execution/update_execution_transition.py | 36 +++- .../agents_api/models/task/create_task.py | 39 +++- agents-api/agents_api/models/task/get_task.py | 8 +- .../agents_api/models/task/list_tasks.py | 5 +- .../agents_api/models/task/update_task.py | 5 +- .../agents_api/models/user/patch_user.py | 1 + .../agents_api/routers/tasks/__init__.py | 1 - .../agents_api/routers/tasks/routers.py | 167 +++++++++++++++--- .../migrate_1716939839_task_relations.py | 3 +- agents-api/poetry.lock | 18 +- mock_openapi.yaml | 87 +++++++-- openapi.yaml | 87 +++++++-- sdks/postman/collection.json | 152 +++++++++++----- sdks/python/julep/api/client.py | 116 ++++++++---- sdks/python/poetry.lock | 6 +- sdks/ts/src/api/services/DefaultService.ts | 91 ++++++++-- 25 files changed, 779 insertions(+), 235 deletions(-) create mode 100644 agents-api/agents_api/models/execution/list_executions.py delete mode 100644 agents-api/agents_api/models/execution/update_execution.py diff --git a/agents-api/agents_api/autogen/openapi_model.py b/agents-api/agents_api/autogen/openapi_model.py index c51f423bc..e269cb3cf 100644 --- a/agents-api/agents_api/autogen/openapi_model.py +++ b/agents-api/agents_api/autogen/openapi_model.py @@ -1,6 +1,6 @@ # generated by datamodel-codegen: # filename: openapi.yaml -# timestamp: 2024-05-30T00:31:23+00:00 +# timestamp: 2024-05-30T20:05:30+00:00 from __future__ import annotations @@ -805,7 +805,7 @@ class ImageUrl(BaseModel): """ URL or base64 data url (e.g. `data:image/jpeg;base64,`) """ - detail: Detail | None = "auto" + detail: Detail | None = "auto" # pytype: disable=annotation-type-mismatch """ image detail to feed into the model can be low | high | auto """ diff --git a/agents-api/agents_api/models/execution/create_execution.py b/agents-api/agents_api/models/execution/create_execution.py index 4abccb114..a75a99022 100644 --- a/agents-api/agents_api/models/execution/create_execution.py +++ b/agents-api/agents_api/models/execution/create_execution.py @@ -6,14 +6,34 @@ @cozo_query def create_execution_query( + agent_id: UUID, task_id: UUID, execution_id: UUID, + developer_id: UUID, status: Literal[ "queued", "starting", "running", "waiting-for-input", "success", "failed" ] = "queued", arguments: Dict[str, Any] = {}, ) -> tuple[str, dict]: - query = """""" + # TODO: Check for agent in developer ID; Assert whether dev can access agent and by relation the task + + query = """ +{ + ?[task_id, execution_id, status, arguments] <- [[ + to_uuid($task_id), + to_uuid($execution_id), + $status, + $arguments + ]] + + :insert executions { + task_id, + execution_id, + status, + arguments + } +} +""" return ( query, { diff --git a/agents-api/agents_api/models/execution/get_execution.py b/agents-api/agents_api/models/execution/get_execution.py index 4abccb114..1870a878f 100644 --- a/agents-api/agents_api/models/execution/get_execution.py +++ b/agents-api/agents_api/models/execution/get_execution.py @@ -1,25 +1,32 @@ from uuid import UUID from ..utils import cozo_query -from typing import Literal, Dict, Any + +from beartype import beartype @cozo_query -def create_execution_query( +@beartype +def get_execution_query( task_id: UUID, execution_id: UUID, - status: Literal[ - "queued", "starting", "running", "waiting-for-input", "success", "failed" - ] = "queued", - arguments: Dict[str, Any] = {}, ) -> tuple[str, dict]: - query = """""" + query = """ +{ + ?[status, arguments, created_at, updated_at] := *executions { + task_id: to_uuid($task_id), + execution_id: to_uuid($execution_id), + status, + arguments, + created_at, + updated_at, + } +} +""" return ( query, { "task_id": str(task_id), "execution_id": str(execution_id), - "status": status, - "arguments": arguments, }, ) diff --git a/agents-api/agents_api/models/execution/get_execution_status.py b/agents-api/agents_api/models/execution/get_execution_status.py index 66a4baea3..d4d355377 100644 --- a/agents-api/agents_api/models/execution/get_execution_status.py +++ b/agents-api/agents_api/models/execution/get_execution_status.py @@ -1,12 +1,21 @@ -from typing import Literal from uuid import UUID +from beartype import beartype from ..utils import cozo_query @cozo_query -def get_execution_status_query(task_id: UUID, developer_id: UUID) -> tuple[str, dict]: +@beartype +def get_execution_status_query(task_id: UUID, execution_id: UUID) -> tuple[str, dict]: task_id = str(task_id) - developer_id = str(developer_id) - query = """""" - return (query, {"task_id": task_id, "developer_id": developer_id}) + execution_id = str(execution_id) + query = """ +{ + ?[status] := *executions { + task_id: to_uuid($task_id), + execution_id: to_uuid($execution_id), + status, + } +} +""" + return (query, {"task_id": task_id, "execution_id": 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 index 5b0c2f942..06935528c 100644 --- a/agents-api/agents_api/models/execution/get_execution_transition.py +++ b/agents-api/agents_api/models/execution/get_execution_transition.py @@ -5,15 +5,25 @@ @cozo_query def get_execution_transition_query( - execution_id: UUID, transition_id: UUID, developer_id: UUID + execution_id: UUID, transition_id: UUID ) -> tuple[str, dict]: - query = """""" + query = """ +{ + ?[type, from, to, output] := *transitions { + execution_id: to_uuid($execution_id), + transition_id: to_uuid($transition_id), + type, + from, + to, + output + } +} +""" return ( query, { "execution_id": str(execution_id), "transition_id": str(transition_id), - "developer_id": str(developer_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 index 5b0c2f942..4b87e79ac 100644 --- a/agents-api/agents_api/models/execution/list_execution_transitions.py +++ b/agents-api/agents_api/models/execution/list_execution_transitions.py @@ -4,16 +4,27 @@ @cozo_query -def get_execution_transition_query( - execution_id: UUID, transition_id: UUID, developer_id: UUID -) -> tuple[str, dict]: +def list_execution_transition_query(execution_id: UUID) -> tuple[str, dict]: - query = """""" + query = """ +{ + ?[transition_id, type, from, to, output, updated_at, created_at] := *transitions { + execution_id: to_uuid($execution_id), + transition_id, + type, + from, + to, + output, + updated_at, + created_at, + } + :limit 100 + :offset 0 +} +""" return ( query, { "execution_id": str(execution_id), - "transition_id": str(transition_id), - "developer_id": str(developer_id), }, ) diff --git a/agents-api/agents_api/models/execution/list_executions.py b/agents-api/agents_api/models/execution/list_executions.py new file mode 100644 index 000000000..058b79290 --- /dev/null +++ b/agents-api/agents_api/models/execution/list_executions.py @@ -0,0 +1,35 @@ +from uuid import UUID + +from beartype import beartype + +from ..utils import cozo_query + + +@cozo_query +@beartype +def list_task_executions_query( + agent_id: UUID, task_id: UUID, developer_id: UUID +) -> tuple[str, dict]: + # TODO: Check for agent in developer ID; Assert whether dev can access agent and by relation the task + query = """ +{ + ?[ + execution_id, + status, + arguments, + created_at, + updated_at, + ] := *executions { + task_id: to_uuid($task_id), + execution_id, + status, + arguments, + created_at, + updated_at, + } + :limit 10 + :offset 0 + +} +""" + return (query, {"task_id": str(task_id)}) 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 4abccb114..000000000 --- a/agents-api/agents_api/models/execution/update_execution.py +++ /dev/null @@ -1,25 +0,0 @@ -from uuid import UUID - -from ..utils import cozo_query -from typing import Literal, Dict, Any - - -@cozo_query -def create_execution_query( - task_id: UUID, - execution_id: UUID, - status: Literal[ - "queued", "starting", "running", "waiting-for-input", "success", "failed" - ] = "queued", - arguments: Dict[str, Any] = {}, -) -> tuple[str, dict]: - query = """""" - return ( - query, - { - "task_id": str(task_id), - "execution_id": str(execution_id), - "status": status, - "arguments": arguments, - }, - ) diff --git a/agents-api/agents_api/models/execution/update_execution_status.py b/agents-api/agents_api/models/execution/update_execution_status.py index 4abccb114..bd8783a2b 100644 --- a/agents-api/agents_api/models/execution/update_execution_status.py +++ b/agents-api/agents_api/models/execution/update_execution_status.py @@ -5,15 +5,32 @@ @cozo_query -def create_execution_query( +def update_execution_status_query( task_id: UUID, execution_id: UUID, status: Literal[ "queued", "starting", "running", "waiting-for-input", "success", "failed" - ] = "queued", + ], arguments: Dict[str, Any] = {}, ) -> tuple[str, dict]: - query = """""" + query = """ +{ + ?[execution_id, task_id, status, updated_at] <- [[ + to_uuid($execution_id), + to_uuid($task_id), + $status, + now() + ]] + + :update executions { + execution_id, + task_id, + status, + updated_at + } +} + +""" return ( query, { diff --git a/agents-api/agents_api/models/execution/update_execution_transition.py b/agents-api/agents_api/models/execution/update_execution_transition.py index 5b0c2f942..2feb3ca31 100644 --- a/agents-api/agents_api/models/execution/update_execution_transition.py +++ b/agents-api/agents_api/models/execution/update_execution_transition.py @@ -1,19 +1,47 @@ from uuid import UUID +from agents_api.common.utils.datetime import utcnow + from ..utils import cozo_query +from ...common.utils.cozo import cozo_process_mutate_data @cozo_query -def get_execution_transition_query( - execution_id: UUID, transition_id: UUID, developer_id: UUID +def update_execution_transition_query( + execution_id: UUID, transition_id: UUID, **update_data ) -> tuple[str, dict]: - query = """""" + transition_update_cols, transition_update_vals = cozo_process_mutate_data( + { + **{k: v for k, v, in update_data.items() if v is not None}, + "execution_id": str(execution_id), + "transition_id": str(transition_id), + "updated_at": utcnow().timestamp(), + } + ) + query = f""" + {{ + input[{transition_update_cols}] <- $transition_update_vals + + ?[{transition_update_cols}] := input[{transition_update_cols}], + *transitions {{ + execution_id: to_uuid($execution_id), + transition_id: to_uuid($transition_id), + }} + + :update transitions {{ + {transition_update_cols} + }} + :returning + }} + +""" + return ( query, { + "transitioon_update_vals": transition_update_vals, "execution_id": str(execution_id), "transition_id": str(transition_id), - "developer_id": str(developer_id), }, ) diff --git a/agents-api/agents_api/models/task/create_task.py b/agents-api/agents_api/models/task/create_task.py index 0984c7b59..70d88c76b 100644 --- a/agents-api/agents_api/models/task/create_task.py +++ b/agents-api/agents_api/models/task/create_task.py @@ -4,21 +4,52 @@ """ from uuid import UUID -from typing import List, Optional, Dict, Any - +from typing import List, Dict, Any +from beartype import beartype from ..utils import cozo_query @cozo_query +@beartype def create_task_query( + agent_id: UUID, task_id: UUID, developer_id: UUID, - agent_id: UUID, name: str, description: str, input_schema: Dict[str, any], tools_available: List[UUID] = [], workflows: List[Dict[str, Any]] = [], ) -> tuple[str, dict]: - pass + # TODO: Check for agent in developer ID; Assert whether dev can access agent and by relation the task + query = """ +{ + ?[agent_id, task_id, name, description, input_schema, tools_available, workflows] <- [[ + to_uuid($agent_id), to_uuid($task_id), $name, $description, $input_schema, [to_uuid($tools_available)], $workflows + ]] + + :insert tasks { + agent_id, + task_id, + name, + description, + input_schema, + tools_available, + workflows + } +} +""" + + return ( + query, + { + "agent_id": str(agent_id), + "task_id": str(task_id), + "name": name, + "description": description, + "input_schema": input_schema, + "tools_available": tools_available, + "workflows": workflows, + }, + ) diff --git a/agents-api/agents_api/models/task/get_task.py b/agents-api/agents_api/models/task/get_task.py index f89da19ad..e9be458d8 100644 --- a/agents-api/agents_api/models/task/get_task.py +++ b/agents-api/agents_api/models/task/get_task.py @@ -1,10 +1,16 @@ from uuid import UUID +from beartype import beartype + from ..utils import cozo_query @cozo_query -def get_task_query(agent_id: UUID, task_id: UUID) -> tuple[str, dict]: +@beartype +def get_task_query( + agent_id: UUID, task_id: UUID, developer_id: UUID +) -> tuple[str, dict]: + # TODO: Check for agent in developer ID; Assert whether dev can access agent and by relation the task query = """ ?[ name, diff --git a/agents-api/agents_api/models/task/list_tasks.py b/agents-api/agents_api/models/task/list_tasks.py index 64f3d8010..f48b424f6 100644 --- a/agents-api/agents_api/models/task/list_tasks.py +++ b/agents-api/agents_api/models/task/list_tasks.py @@ -1,13 +1,12 @@ -from typing import Any from uuid import UUID -from ...common.utils import json from ..utils import cozo_query @cozo_query def list_tasks_query( agent_id: UUID, + developer_id: UUID, limit: int = 100, offset: int = 0, # metadata_filter: dict[str, Any] = {}, @@ -21,7 +20,7 @@ def list_tasks_query( Returns: pd.DataFrame: A DataFrame containing the queried task data. """ - + # TODO: Accepts developer ID. Checks if the developer can get this agent, by relation can get the tasks. Assert that the agent exists under the developer. query = """ ?[ task_id, diff --git a/agents-api/agents_api/models/task/update_task.py b/agents-api/agents_api/models/task/update_task.py index 0984c7b59..42e89ecbf 100644 --- a/agents-api/agents_api/models/task/update_task.py +++ b/agents-api/agents_api/models/task/update_task.py @@ -4,14 +4,14 @@ """ from uuid import UUID -from typing import List, Optional, Dict, Any +from typing import List, Dict, Any from ..utils import cozo_query @cozo_query -def create_task_query( +def cupdate_task_query( task_id: UUID, developer_id: UUID, agent_id: UUID, @@ -21,4 +21,5 @@ def create_task_query( tools_available: List[UUID] = [], workflows: List[Dict[str, Any]] = [], ) -> tuple[str, dict]: + # NOT TO IMPLEMENT FOR NOW pass diff --git a/agents-api/agents_api/models/user/patch_user.py b/agents-api/agents_api/models/user/patch_user.py index 089a291cc..48d30ff2f 100644 --- a/agents-api/agents_api/models/user/patch_user.py +++ b/agents-api/agents_api/models/user/patch_user.py @@ -39,6 +39,7 @@ def patch_user_query( ) # Construct the datalog query for updating user information. + # TODO: Modify user update to modify `updated_at` field query = f""" # update the user input[{user_update_cols}] <- $user_update_vals diff --git a/agents-api/agents_api/routers/tasks/__init__.py b/agents-api/agents_api/routers/tasks/__init__.py index afd8af03a..e69de29bb 100644 --- a/agents-api/agents_api/routers/tasks/__init__.py +++ b/agents-api/agents_api/routers/tasks/__init__.py @@ -1 +0,0 @@ -from .routers import router diff --git a/agents-api/agents_api/routers/tasks/routers.py b/agents-api/agents_api/routers/tasks/routers.py index 0922d82b0..780bf2a06 100644 --- a/agents-api/agents_api/routers/tasks/routers.py +++ b/agents-api/agents_api/routers/tasks/routers.py @@ -1,18 +1,22 @@ -from typing import Annotated, Literal +from typing import Annotated from uuid import uuid4 from agents_api.models.execution.create_execution import create_execution_query -from agents_api.models.execution.get_execution_status import get_execution_status_query +from agents_api.models.execution.get_execution import get_execution_query from agents_api.models.execution.get_execution_transition import ( get_execution_transition_query, ) +from agents_api.models.execution.list_execution_transitions import ( + list_execution_transition_query, +) +from agents_api.models.execution.list_executions import list_task_executions_query from agents_api.models.task.create_task import create_task_query from agents_api.models.task.get_task import get_task_query from agents_api.models.task.list_tasks import list_tasks_query -from fastapi import APIRouter, HTTPException, status, BackgroundTasks, Depends +from fastapi import APIRouter, HTTPException, status, Depends import pandas as pd from pydantic import BaseModel from pydantic import UUID4 -from starlette.status import HTTP_201_CREATED, HTTP_202_ACCEPTED +from starlette.status import HTTP_201_CREATED from pycozo.client import QueryException from agents_api.autogen.openapi_model import ( @@ -31,30 +35,35 @@ class TaskList(BaseModel): items: list[Task] +class ExecutionList(BaseModel): + items: list[Execution] + + +class ExecutionTransitionList(BaseModel): + items: list[ExecutionTransition] + + router = APIRouter() +# TODO: Fix arguments (named or positional) +# TODO: Add :limit 1 to all get queries from cozo + -@router.get("/tasks", tags=["tasks"]) +@router.get("/agents/{agent_id}/tasks", tags=["tasks"]) async def list_tasks( + agent_id: UUID4, x_developer_id: Annotated[UUID4, Depends(get_developer_id)], ) -> TaskList: - query_results = list_tasks_query(x_developer_id) + query_results = list_tasks_query(agent_id, x_developer_id) return TaskList( items=[Task(**row.to_dict()) for _, row in query_results.iterrows()] ) -@router.get("/tasks/{task_id}", tags=["tasks"]) -async def get_task( - task_id: UUID4, x_developer_id: Annotated[UUID4, Depends(get_developer_id)] -) -> Task: - query_results = get_task_query(x_developer_id, task_id) - pass - - -@router.post("/tasks", status_code=HTTP_201_CREATED, tags=["tasks"]) +@router.post("/agents/{agent_id}/tasks", status_code=HTTP_201_CREATED, tags=["tasks"]) async def create_task( request: CreateTask, + agent_id: UUID4, x_developer_id: Annotated[UUID4, Depends(get_developer_id)], ) -> ResourceCreatedResponse: task_id = uuid4() @@ -62,16 +71,51 @@ async def create_task( # TODO: Do thorough validation of the task spec resp: pd.DataFrame = create_task_query( + agent_id=agent_id, task_id=task_id, developer_id=x_developer_id, + name=request.name, + description=request.description, + input_schema=request.input_schema, + tools_available=request.tools_available, + workflows=request.workflows, ) return ResourceCreatedResponse( id=resp["task_id"][0], created_at=resp["created_at"][0] ) -@router.post("/tasks/{task_id}/executions", status_code=HTTP_201_CREATED, tags=["tasks"]) +@router.get("/agents/{agent_id}/tasks/{task_id}", tags=["tasks"]) +async def get_task( + task_id: UUID4, + agent_id: UUID4, + x_developer_id: Annotated[UUID4, Depends(get_developer_id)], +) -> Task: + try: + resp = [ + row.to_dict() + for _, row in get_task_query(agent_id, task_id, x_developer_id).iterrows() + ][0] + return Task(**resp) + except (IndexError, KeyError): + raise HTTPException( + status_code=status.HTTP_404_NOT_FOUND, + detail="Task not found", + ) + except QueryException as e: + if e.code == "transact::assertion_failure": + raise HTTPException( + status_code=status.HTTP_404_NOT_FOUND, detail="Task not found" + ) + + +@router.post( + "/agents/{agent_id}/tasks/{task_id}/executions", + status_code=HTTP_201_CREATED, + tags=["tasks"], +) async def create_task_execution( + agent_id: UUID4, task_id: UUID4, request: CreateExecution, x_developer_id: Annotated[UUID4, Depends(get_developer_id)], @@ -79,32 +123,99 @@ async def create_task_execution( # TODO: Do thorough validation of the input against task input schema # DO NOT let the user specify the status - resp = create_execution_query() + resp = create_execution_query( + agent_id=agent_id, + task_id=task_id, + execution_id=uuid4(), + developer_id=x_developer_id, + status=request.status, + arguments=request.arguments, + ) return ResourceCreatedResponse( id=resp["execution_id"][0], created_at=resp["created_at"][0] ) -@router.get("/tasks/{task_id}/executions", tags=["tasks"]) -async def get_task_execution_status( +@router.get("/agents/{agent_id}/tasks/{task_id}/executions", tags=["tasks"]) +async def list_task_executions( + agent_id: UUID4, task_id: UUID4, x_developer_id: Annotated[UUID4, Depends(get_developer_id)], -) -> Literal["queued", "starting", "running", "waiting_for_input", "success", "failed"]: +) -> ExecutionList: + res = list_task_executions_query(agent_id, task_id, x_developer_id) + return ExecutionList( + items=[Execution(**row.to_dict()) for _, row in res.iterrows()] + ) - resp = get_execution_status_query(task_id, x_developer_id) - pass + +@router.get("/tasks/{task_id}/executions/{execution_id}") +async def get_execution(task_id: UUID4, execution_id: UUID4) -> Execution: + try: + res = [ + row.to_dict() + for _, row in get_execution_query(task_id, execution_id).iterrows() + ][0] + return Execution(**res) + except (IndexError, KeyError): + raise HTTPException( + status_code=status.HTTP_404_NOT_FOUND, + detail="Execution not found", + ) + + +# @router.get("/tasks/{task_id}/executions/{execution_id}/status", tags=["tasks"]) +# async def get_task_execution_status( +# task_id: UUID4, +# execution_id: UUID4, +# ) -> Literal["queued", "starting", "running", "waiting_for_input", "success", "failed"]: +# try: +# res = [ +# row.to_dict() +# for _, row in get_execution_status_query(task_id, execution_id).iterrows() +# ][0] +# return res["status"] +# except (IndexError, KeyError): +# raise HTTPException( +# status_code=status.HTTP_404_NOT_FOUND, +# detail="Execution not found", +# ) @router.get("/executions/{execution_id}/transitions/{transition_id}") async def get_execution_transition( execution_id: UUID4, transition_id: UUID4, - x_developer_id: Annotated[UUID4, Depends(get_developer_id)], ) -> ExecutionTransition: - resp = get_execution_transition_query(execution_id, transition_id, x_developer_id) - + try: + res = [ + row.to_dict() + for _, row in get_execution_transition_query( + execution_id, transition_id + ).iterrows() + ][0] + return ExecutionTransition(**res) + except (IndexError, KeyError): + raise HTTPException( + status_code=status.HTTP_404_NOT_FOUND, + detail="Transition not found", + ) + + +# TODO: Later; for resuming waiting transitions +# TODO: Ask for a task token to resume a waiting transition +@router.put("/executions/{execution_id}/transitions/{transition_id}") +async def update_execution_transition( + execution_id: UUID4, transition_id: UUID4, request +) -> ResourceUpdatedResponse: + # try: + # resp = update_execution_transition_query(execution_id, transition_id, request) pass + # OpenAPI Model doesn't have update execution transition + -# @router.get("/executions/{execution_id}") # -> get the execution object -# @router.put("/executions/{execution_id}/transitions/{transition_id}") # -> update a waiting transition -# @router.get("/executions/{execution_id}/transitions") # -> get all transitions for an execution +@router.get("/executions/{execution_id}/transitions") +async def list_execution_transitions(execution_id: UUID4) -> ExecutionTransitionList: + res = list_execution_transition_query(execution_id) + return ExecutionTransitionList( + items=[ExecutionTransition(**row.to_dict()) for _, row in res.iterrows()] + ) diff --git a/agents-api/migrations/migrate_1716939839_task_relations.py b/agents-api/migrations/migrate_1716939839_task_relations.py index 6ea3fab6e..c7397a164 100644 --- a/agents-api/migrations/migrate_1716939839_task_relations.py +++ b/agents-api/migrations/migrate_1716939839_task_relations.py @@ -37,7 +37,7 @@ def run(client, queries): execution_id: Uuid, => status: String default 'queued', - # one of: "queued", "starting", "running", "waiting-for-input", "success", "failed" + # one of: "queued", "starting", "running", "waiting_for_input", "success", "failed" arguments: Json, created_at: Float default now(), @@ -60,6 +60,7 @@ def run(client, queries): to: (String, Int)?, output: Json, created_at: Float default now(), + updated_at: Float default now(), } """, down="::remove transitions", diff --git a/agents-api/poetry.lock b/agents-api/poetry.lock index 4f6c97f80..df96d5e95 100644 --- a/agents-api/poetry.lock +++ b/agents-api/poetry.lock @@ -1606,13 +1606,13 @@ dev = ["Sphinx (>=5.1.1)", "black (==23.12.1)", "build (>=0.10.0)", "coverage (> [[package]] name = "litellm" -version = "1.39.2" +version = "1.39.4" description = "Library to easily interface with LLM API providers" optional = false python-versions = "!=2.7.*,!=3.0.*,!=3.1.*,!=3.2.*,!=3.3.*,!=3.4.*,!=3.5.*,!=3.6.*,!=3.7.*,>=3.8" files = [ - {file = "litellm-1.39.2-py3-none-any.whl", hash = "sha256:843cb9a4d45c89ba6da95529815ec83ee7e4b7fe07aa0ed633102f600fddd9ad"}, - {file = "litellm-1.39.2.tar.gz", hash = "sha256:96c4f3d522ccf32817357b1e9f5f63fa36a4a884f336314e1f6d66c0576d689e"}, + {file = "litellm-1.39.4-py3-none-any.whl", hash = "sha256:3edd2b153153e8902770ada641b93bcdeaba8d23cb579e599919331b52741040"}, + {file = "litellm-1.39.4.tar.gz", hash = "sha256:f7ec8ef44257235de10c8e0d326fff0083b48de4bc71531d78d8c6778af9d401"}, ] [package.dependencies] @@ -2021,13 +2021,13 @@ files = [ [[package]] name = "openai" -version = "1.30.4" +version = "1.30.5" description = "The official Python library for the openai API" optional = false python-versions = ">=3.7.1" files = [ - {file = "openai-1.30.4-py3-none-any.whl", hash = "sha256:fb2635efd270efaf9fac2e07558d7948373b940637d3ae3ab624c1a983d4f03f"}, - {file = "openai-1.30.4.tar.gz", hash = "sha256:f3488d9a1c4e0d332b019377d27d7cb4b3d6103fd5d0a416c7ceac780d1d9b88"}, + {file = "openai-1.30.5-py3-none-any.whl", hash = "sha256:2ad95e926de0d2e09cde632a9204b0a6dca4a03c2cdcc84329b01f355784355a"}, + {file = "openai-1.30.5.tar.gz", hash = "sha256:5366562eb2c5917e6116ae0391b7ae6e3acd62b0ae3f565ada32b35d8fcfa106"}, ] [package.dependencies] @@ -3542,13 +3542,13 @@ test = ["argcomplete (>=3.0.3)", "mypy (>=1.7.0)", "pre-commit", "pytest (>=7.0, [[package]] name = "transformers" -version = "4.41.1" +version = "4.41.2" description = "State-of-the-art Machine Learning for JAX, PyTorch and TensorFlow" optional = false python-versions = ">=3.8.0" files = [ - {file = "transformers-4.41.1-py3-none-any.whl", hash = "sha256:f0680e0b1a01067eccd11f62f0522409422c7d6f91d532fe0f50b136a406129d"}, - {file = "transformers-4.41.1.tar.gz", hash = "sha256:fa859e4c66f0896633a3bf534e0d9a29a9a88478a49f94c5d8270537dc61cc42"}, + {file = "transformers-4.41.2-py3-none-any.whl", hash = "sha256:05555d20e43f808de1ef211ab64803cdb513170cef70d29a888b589caebefc67"}, + {file = "transformers-4.41.2.tar.gz", hash = "sha256:80a4db216533d573e9cc7388646c31ed9480918feb7c55eb211249cb23567f87"}, ] [package.dependencies] diff --git a/mock_openapi.yaml b/mock_openapi.yaml index bb347ab1f..de183b650 100644 --- a/mock_openapi.yaml +++ b/mock_openapi.yaml @@ -1134,7 +1134,7 @@ paths: type: string format: uuid required: true - /api/tasks: + /api/agents/{agent_id}/tasks: get: summary: Get a list of tasks description: '' @@ -1170,7 +1170,14 @@ paths: $ref: '#/components/schemas/ResourceCreatedResponse' security: - api-key: [] - /api/tasks/{task_id}/execution: + parameters: + - in: path + name: agent_id + description: '' + schema: &ref_0 + type: string + required: true + /api/agents/{agent_id}/tasks/{task_id}/executions: post: summary: Create (or start) an execution of a Task description: '' @@ -1191,28 +1198,35 @@ paths: security: - api-key: [] get: - summary: Get execution (status) of a Task + summary: List Executions of a Task description: '' operationId: GetTaskExecution - responses: + responses: &ref_1 '200': description: '' headers: {} content: application/json: schema: - $ref: '#/components/schemas/Execution' - security: + type: array + items: + $ref: '#/components/schemas/Execution' + security: &ref_2 - api-key: [] parameters: + - in: path + name: agent_id + description: '' + schema: *ref_0 + required: true - in: path name: task_id description: '' - schema: + schema: &ref_3 type: string format: uuid required: true - /api/tasks/{task_id}: + /api/agents/{agent_id}/tasks/{task_id}: get: summary: Get a Task by ID description: '' @@ -1224,10 +1238,15 @@ paths: content: application/json: schema: - $ref: '#/components/schemas/Execution' + $ref: '#/components/schemas/Task' security: - api-key: [] parameters: + - in: path + name: agent_id + description: '' + schema: *ref_0 + required: true - in: path name: task_id description: '' @@ -1235,24 +1254,24 @@ paths: type: string format: uuid required: true - /api/execution/{execution_id}/transition/{transition_id}: + /api/executions/{execution_id}/transitions/{transition_id}: get: description: '' operationId: GetExecutionTransition responses: '200': description: '' - headers: {} + headers: &ref_4 {} content: application/json: schema: - type: array - items: - $ref: '#/components/schemas/ExecutionTransition' - security: + $ref: '#/components/schemas/ExecutionTransition' + security: &ref_5 - api-key: [] + summary: Get an Execution Transition parameters: - - in: path + - &ref_6 + in: path name: execution_id description: '' schema: @@ -1291,6 +1310,42 @@ paths: $ref: '#/components/schemas/ResourceUpdatedResponse' security: - api-key: [] + /api/tasks/{task_id}/executions/{execution_id}: + get: + summary: Get Task Execution By ID + description: '' + operationId: GetTaskExecution + responses: *ref_1 + security: *ref_2 + parameters: + - in: path + name: task_id + description: '' + schema: *ref_0 + required: true + - in: path + name: execution_id + description: '' + schema: *ref_3 + required: true + /api/executions/{execution_id}/transitions/: + get: + description: '' + operationId: GetExecutionTransition + responses: + '200': + description: '' + headers: *ref_4 + content: + application/json: + schema: + type: array + items: + $ref: '#/components/schemas/ExecutionTransition' + security: *ref_5 + summary: List Transitions of an Execution + parameters: + - *ref_6 components: schemas: User: diff --git a/openapi.yaml b/openapi.yaml index 70580de29..ef8d44371 100644 --- a/openapi.yaml +++ b/openapi.yaml @@ -1134,7 +1134,7 @@ paths: type: string format: uuid required: true - /tasks: + /agents/{agent_id}/tasks: get: summary: Get a list of tasks description: '' @@ -1170,7 +1170,14 @@ paths: $ref: '#/components/schemas/ResourceCreatedResponse' security: - api-key: [] - /tasks/{task_id}/execution: + parameters: + - in: path + name: agent_id + description: '' + schema: &ref_0 + type: string + required: true + /agents/{agent_id}/tasks/{task_id}/executions: post: summary: Create (or start) an execution of a Task description: '' @@ -1191,28 +1198,35 @@ paths: security: - api-key: [] get: - summary: Get execution (status) of a Task + summary: List Executions of a Task description: '' operationId: GetTaskExecution - responses: + responses: &ref_1 '200': description: '' headers: {} content: application/json: schema: - $ref: '#/components/schemas/Execution' - security: + type: array + items: + $ref: '#/components/schemas/Execution' + security: &ref_2 - api-key: [] parameters: + - in: path + name: agent_id + description: '' + schema: *ref_0 + required: true - in: path name: task_id description: '' - schema: + schema: &ref_3 type: string format: uuid required: true - /tasks/{task_id}: + /agents/{agent_id}/tasks/{task_id}: get: summary: Get a Task by ID description: '' @@ -1224,10 +1238,15 @@ paths: content: application/json: schema: - $ref: '#/components/schemas/Execution' + $ref: '#/components/schemas/Task' security: - api-key: [] parameters: + - in: path + name: agent_id + description: '' + schema: *ref_0 + required: true - in: path name: task_id description: '' @@ -1235,24 +1254,24 @@ paths: type: string format: uuid required: true - /execution/{execution_id}/transition/{transition_id}: + /executions/{execution_id}/transitions/{transition_id}: get: description: '' operationId: GetExecutionTransition responses: '200': description: '' - headers: {} + headers: &ref_4 {} content: application/json: schema: - type: array - items: - $ref: '#/components/schemas/ExecutionTransition' - security: + $ref: '#/components/schemas/ExecutionTransition' + security: &ref_5 - api-key: [] + summary: Get an Execution Transition parameters: - - in: path + - &ref_6 + in: path name: execution_id description: '' schema: @@ -1291,6 +1310,42 @@ paths: $ref: '#/components/schemas/ResourceUpdatedResponse' security: - api-key: [] + /tasks/{task_id}/executions/{execution_id}: + get: + summary: Get Task Execution By ID + description: '' + operationId: GetTaskExecution + responses: *ref_1 + security: *ref_2 + parameters: + - in: path + name: task_id + description: '' + schema: *ref_0 + required: true + - in: path + name: execution_id + description: '' + schema: *ref_3 + required: true + /executions/{execution_id}/transitions/: + get: + description: '' + operationId: GetExecutionTransition + responses: + '200': + description: '' + headers: *ref_4 + content: + application/json: + schema: + type: array + items: + $ref: '#/components/schemas/ExecutionTransition' + security: *ref_5 + summary: List Transitions of an Execution + parameters: + - *ref_6 components: schemas: User: diff --git a/sdks/postman/collection.json b/sdks/postman/collection.json index b6bf22fef..9ca70e71e 100644 --- a/sdks/postman/collection.json +++ b/sdks/postman/collection.json @@ -3497,15 +3497,23 @@ "request": { "description": "", "url": { - "raw": "{{baseUrl}}/tasks", + "raw": "{{baseUrl}}/agents/:agent_id/tasks", "host": [ "{{baseUrl}}" ], "path": [ + "agents", + ":agent_id", "tasks" ], "query": [], - "variable": [] + "variable": [ + { + "key": "agent_id", + "description": "", + "value": "agent_id" + } + ] }, "header": [ { @@ -3532,15 +3540,23 @@ "originalRequest": { "description": "", "url": { - "raw": "{{baseUrl}}/tasks", + "raw": "{{baseUrl}}/agents/:agent_id/tasks", "host": [ "{{baseUrl}}" ], "path": [ + "agents", + ":agent_id", "tasks" ], "query": [], - "variable": [] + "variable": [ + { + "key": "agent_id", + "description": "", + "value": "agent_id" + } + ] }, "header": [ { @@ -3571,15 +3587,23 @@ "request": { "description": "", "url": { - "raw": "{{baseUrl}}/tasks", + "raw": "{{baseUrl}}/agents/:agent_id/tasks", "host": [ "{{baseUrl}}" ], "path": [ + "agents", + ":agent_id", "tasks" ], "query": [], - "variable": [] + "variable": [ + { + "key": "agent_id", + "description": "", + "value": "agent_id" + } + ] }, "header": [ { @@ -3614,15 +3638,23 @@ "originalRequest": { "description": "", "url": { - "raw": "{{baseUrl}}/tasks", + "raw": "{{baseUrl}}/agents/:agent_id/tasks", "host": [ "{{baseUrl}}" ], "path": [ + "agents", + ":agent_id", "tasks" ], "query": [], - "variable": [] + "variable": [ + { + "key": "agent_id", + "description": "", + "value": "agent_id" + } + ] }, "header": [ { @@ -3657,18 +3689,19 @@ }, { "_type": "endpoint", - "name": "Get execution (status) of a Task", + "name": "Get Task Execution By ID", "request": { "description": "", "url": { - "raw": "{{baseUrl}}/tasks/:task_id/execution", + "raw": "{{baseUrl}}/tasks/:task_id/executions/:execution_id", "host": [ "{{baseUrl}}" ], "path": [ "tasks", ":task_id", - "execution" + "executions", + ":execution_id" ], "query": [], "variable": [ @@ -3676,6 +3709,11 @@ "key": "task_id", "description": "", "value": "task_id" + }, + { + "key": "execution_id", + "description": "", + "value": "execution_id" } ] }, @@ -3704,14 +3742,15 @@ "originalRequest": { "description": "", "url": { - "raw": "{{baseUrl}}/tasks/:task_id/execution", + "raw": "{{baseUrl}}/tasks/:task_id/executions/:execution_id", "host": [ "{{baseUrl}}" ], "path": [ "tasks", ":task_id", - "execution" + "executions", + ":execution_id" ], "query": [], "variable": [ @@ -3719,6 +3758,11 @@ "key": "task_id", "description": "", "value": "task_id" + }, + { + "key": "execution_id", + "description": "", + "value": "execution_id" } ] }, @@ -3740,7 +3784,7 @@ "body": null }, "description": "", - "body": "{\n \"id\": \"id\",\n \"task_id\": \"task_id\",\n \"created_at\": \"created_at\",\n \"arguments\": {\n \"arguments\": {\n \"key\": \"value\"\n }\n },\n \"status\": \"status\"\n}", + "body": "[\n {\n \"id\": \"id\",\n \"task_id\": \"task_id\",\n \"created_at\": \"created_at\",\n \"arguments\": {\n \"arguments\": {\n \"key\": \"value\"\n }\n },\n \"status\": \"status\"\n }\n]", "_postman_previewlanguage": "json" } ] @@ -3751,17 +3795,24 @@ "request": { "description": "", "url": { - "raw": "{{baseUrl}}/tasks/:task_id/execution", + "raw": "{{baseUrl}}/agents/:agent_id/tasks/:task_id/executions", "host": [ "{{baseUrl}}" ], "path": [ + "agents", + ":agent_id", "tasks", ":task_id", - "execution" + "executions" ], "query": [], "variable": [ + { + "key": "agent_id", + "description": "", + "value": "agent_id" + }, { "key": "task_id", "description": "", @@ -3802,17 +3853,24 @@ "originalRequest": { "description": "", "url": { - "raw": "{{baseUrl}}/tasks/:task_id/execution", + "raw": "{{baseUrl}}/agents/:agent_id/tasks/:task_id/executions", "host": [ "{{baseUrl}}" ], "path": [ + "agents", + ":agent_id", "tasks", ":task_id", - "execution" + "executions" ], "query": [], "variable": [ + { + "key": "agent_id", + "description": "", + "value": "agent_id" + }, { "key": "task_id", "description": "", @@ -3857,16 +3915,23 @@ "request": { "description": "", "url": { - "raw": "{{baseUrl}}/tasks/:task_id", + "raw": "{{baseUrl}}/agents/:agent_id/tasks/:task_id", "host": [ "{{baseUrl}}" ], "path": [ + "agents", + ":agent_id", "tasks", ":task_id" ], "query": [], "variable": [ + { + "key": "agent_id", + "description": "", + "value": "agent_id" + }, { "key": "task_id", "description": "", @@ -3899,16 +3964,23 @@ "originalRequest": { "description": "", "url": { - "raw": "{{baseUrl}}/tasks/:task_id", + "raw": "{{baseUrl}}/agents/:agent_id/tasks/:task_id", "host": [ "{{baseUrl}}" ], "path": [ + "agents", + ":agent_id", "tasks", ":task_id" ], "query": [], "variable": [ + { + "key": "agent_id", + "description": "", + "value": "agent_id" + }, { "key": "task_id", "description": "", @@ -3934,26 +4006,25 @@ "body": null }, "description": "", - "body": "{\n \"id\": \"id\",\n \"task_id\": \"task_id\",\n \"created_at\": \"created_at\",\n \"arguments\": {\n \"arguments\": {\n \"key\": \"value\"\n }\n },\n \"status\": \"status\"\n}", + "body": "{\n \"name\": \"name\",\n \"description\": \"description\",\n \"tools_available\": [\n \"tools_available\"\n ],\n \"input_schema\": {\n \"input_schema\": {\n \"key\": \"value\"\n }\n },\n \"main\": [\n {\n \"prompt\": [\n {\n \"role\": \"user\",\n \"content\": \"content\"\n }\n ],\n \"settings\": {\n \"min_p\": 0.01\n }\n }\n ],\n \"id\": \"id\",\n \"created_at\": \"2024-01-15T09:30:00Z\",\n \"agent_id\": \"agent_id\"\n}", "_postman_previewlanguage": "json" } ] }, { "_type": "endpoint", - "name": "Get Execution Transition", + "name": "List Transitions of an Execution", "request": { "description": "", "url": { - "raw": "{{baseUrl}}/execution/:execution_id/transition/:transition_id", + "raw": "{{baseUrl}}/executions/:execution_id/transitions", "host": [ "{{baseUrl}}" ], "path": [ - "execution", + "executions", ":execution_id", - "transition", - ":transition_id" + "transitions" ], "query": [], "variable": [ @@ -3961,11 +4032,6 @@ "key": "execution_id", "description": "", "value": "execution_id" - }, - { - "key": "transition_id", - "description": "", - "value": "transition_id" } ] }, @@ -3994,15 +4060,14 @@ "originalRequest": { "description": "", "url": { - "raw": "{{baseUrl}}/execution/:execution_id/transition/:transition_id", + "raw": "{{baseUrl}}/executions/:execution_id/transitions", "host": [ "{{baseUrl}}" ], "path": [ - "execution", + "executions", ":execution_id", - "transition", - ":transition_id" + "transitions" ], "query": [], "variable": [ @@ -4010,11 +4075,6 @@ "key": "execution_id", "description": "", "value": "execution_id" - }, - { - "key": "transition_id", - "description": "", - "value": "transition_id" } ] }, @@ -4047,14 +4107,14 @@ "request": { "description": "", "url": { - "raw": "{{baseUrl}}/execution/:execution_id/transition/:transition_id", + "raw": "{{baseUrl}}/executions/:execution_id/transitions/:transition_id", "host": [ "{{baseUrl}}" ], "path": [ - "execution", + "executions", ":execution_id", - "transition", + "transitions", ":transition_id" ], "query": [], @@ -4104,14 +4164,14 @@ "originalRequest": { "description": "", "url": { - "raw": "{{baseUrl}}/execution/:execution_id/transition/:transition_id", + "raw": "{{baseUrl}}/executions/:execution_id/transitions/:transition_id", "host": [ "{{baseUrl}}" ], "path": [ - "execution", + "executions", ":execution_id", - "transition", + "transitions", ":transition_id" ], "query": [], diff --git a/sdks/python/julep/api/client.py b/sdks/python/julep/api/client.py index 19585251d..ae3110093 100644 --- a/sdks/python/julep/api/client.py +++ b/sdks/python/julep/api/client.py @@ -1834,21 +1834,27 @@ def get_job_status(self, job_id: str) -> JobStatus: raise ApiError(status_code=_response.status_code, body=_response.text) raise ApiError(status_code=_response.status_code, body=_response_json) - def list_tasks(self) -> typing.List[Task]: + def list_tasks(self, agent_id: str) -> typing.List[Task]: """ + Parameters: + - agent_id: str. --- from julep.client import JulepApi client = JulepApi( api_key="YOUR_API_KEY", ) - client.list_tasks() + client.list_tasks( + agent_id="agent_id", + ) """ _response = self._client_wrapper.httpx_client.request( "GET", - urllib.parse.urljoin(f"{self._client_wrapper.get_base_url()}/", "tasks"), + urllib.parse.urljoin( + f"{self._client_wrapper.get_base_url()}/", f"agents/{agent_id}/tasks" + ), headers=self._client_wrapper.get_headers(), timeout=300, ) @@ -1862,6 +1868,7 @@ def list_tasks(self) -> typing.List[Task]: def create_task( self, + agent_id: str, *, name: str, description: typing.Optional[str] = OMIT, @@ -1873,6 +1880,8 @@ def create_task( Parameters: + - agent_id: str. + - name: str. Name of the Task - description: typing.Optional[str]. Optional Description of the Task @@ -1889,6 +1898,7 @@ def create_task( api_key="YOUR_API_KEY", ) client.create_task( + agent_id="agent_id", name="name", ) """ @@ -1901,7 +1911,9 @@ def create_task( _request["input_schema"] = input_schema _response = self._client_wrapper.httpx_client.request( "POST", - urllib.parse.urljoin(f"{self._client_wrapper.get_base_url()}/", "tasks"), + urllib.parse.urljoin( + f"{self._client_wrapper.get_base_url()}/", f"agents/{agent_id}/tasks" + ), json=jsonable_encoder(_request), headers=self._client_wrapper.get_headers(), timeout=300, @@ -1914,12 +1926,16 @@ def create_task( raise ApiError(status_code=_response.status_code, body=_response.text) raise ApiError(status_code=_response.status_code, body=_response_json) - def get_task_execution(self, task_id: str) -> Execution: + def get_task_execution( + self, task_id: str, execution_id: str + ) -> typing.List[Execution]: """ Parameters: - task_id: str. + + - execution_id: str. --- from julep.client import JulepApi @@ -1928,18 +1944,20 @@ def get_task_execution(self, task_id: str) -> Execution: ) client.get_task_execution( task_id="task_id", + execution_id="execution_id", ) """ _response = self._client_wrapper.httpx_client.request( "GET", urllib.parse.urljoin( - f"{self._client_wrapper.get_base_url()}/", f"tasks/{task_id}/execution" + f"{self._client_wrapper.get_base_url()}/", + f"tasks/{task_id}/executions/{execution_id}", ), headers=self._client_wrapper.get_headers(), timeout=300, ) if 200 <= _response.status_code < 300: - return pydantic.parse_obj_as(Execution, _response.json()) # type: ignore + return pydantic.parse_obj_as(typing.List[Execution], _response.json()) # type: ignore try: _response_json = _response.json() except JSONDecodeError: @@ -1948,6 +1966,7 @@ def get_task_execution(self, task_id: str) -> Execution: def start_task_execution( self, + agent_id: str, task_id: str, *, create_execution_task_id: str, @@ -1958,6 +1977,8 @@ def start_task_execution( Parameters: + - agent_id: str. + - task_id: str. - create_execution_task_id: str. @@ -1972,6 +1993,7 @@ def start_task_execution( api_key="YOUR_API_KEY", ) client.start_task_execution( + agent_id="agent_id", task_id="task_id", create_execution_task_id="task_id", arguments={}, @@ -1981,7 +2003,8 @@ def start_task_execution( _response = self._client_wrapper.httpx_client.request( "POST", urllib.parse.urljoin( - f"{self._client_wrapper.get_base_url()}/", f"tasks/{task_id}/execution" + f"{self._client_wrapper.get_base_url()}/", + f"agents/{agent_id}/tasks/{task_id}/executions", ), json=jsonable_encoder( { @@ -2001,11 +2024,13 @@ def start_task_execution( raise ApiError(status_code=_response.status_code, body=_response.text) raise ApiError(status_code=_response.status_code, body=_response_json) - def get_task(self, task_id: str) -> Execution: + def get_task(self, agent_id: str, task_id: str) -> Task: """ Parameters: + - agent_id: str. + - task_id: str. --- from julep.client import JulepApi @@ -2014,19 +2039,21 @@ def get_task(self, task_id: str) -> Execution: api_key="YOUR_API_KEY", ) client.get_task( + agent_id="agent_id", task_id="task_id", ) """ _response = self._client_wrapper.httpx_client.request( "GET", urllib.parse.urljoin( - f"{self._client_wrapper.get_base_url()}/", f"tasks/{task_id}" + f"{self._client_wrapper.get_base_url()}/", + f"agents/{agent_id}/tasks/{task_id}", ), headers=self._client_wrapper.get_headers(), timeout=300, ) if 200 <= _response.status_code < 300: - return pydantic.parse_obj_as(Execution, _response.json()) # type: ignore + return pydantic.parse_obj_as(Task, _response.json()) # type: ignore try: _response_json = _response.json() except JSONDecodeError: @@ -2034,15 +2061,13 @@ def get_task(self, task_id: str) -> Execution: raise ApiError(status_code=_response.status_code, body=_response_json) def get_execution_transition( - self, execution_id: str, transition_id: str + self, execution_id: str ) -> typing.List[ExecutionTransition]: """ Parameters: - execution_id: str. - - - transition_id: str. --- from julep.client import JulepApi @@ -2051,14 +2076,13 @@ def get_execution_transition( ) client.get_execution_transition( execution_id="execution_id", - transition_id="transition_id", ) """ _response = self._client_wrapper.httpx_client.request( "GET", urllib.parse.urljoin( f"{self._client_wrapper.get_base_url()}/", - f"execution/{execution_id}/transition/{transition_id}", + f"executions/{execution_id}/transitions/", ), headers=self._client_wrapper.get_headers(), timeout=300, @@ -2108,7 +2132,7 @@ def resume_tool_execution( "PUT", urllib.parse.urljoin( f"{self._client_wrapper.get_base_url()}/", - f"execution/{execution_id}/transition/{transition_id}", + f"executions/{execution_id}/transitions/{transition_id}", ), json=jsonable_encoder({"responses": responses}), headers=self._client_wrapper.get_headers(), @@ -3883,21 +3907,27 @@ async def get_job_status(self, job_id: str) -> JobStatus: raise ApiError(status_code=_response.status_code, body=_response.text) raise ApiError(status_code=_response.status_code, body=_response_json) - async def list_tasks(self) -> typing.List[Task]: + async def list_tasks(self, agent_id: str) -> typing.List[Task]: """ + Parameters: + - agent_id: str. --- from julep.client import AsyncJulepApi client = AsyncJulepApi( api_key="YOUR_API_KEY", ) - await client.list_tasks() + await client.list_tasks( + agent_id="agent_id", + ) """ _response = await self._client_wrapper.httpx_client.request( "GET", - urllib.parse.urljoin(f"{self._client_wrapper.get_base_url()}/", "tasks"), + urllib.parse.urljoin( + f"{self._client_wrapper.get_base_url()}/", f"agents/{agent_id}/tasks" + ), headers=self._client_wrapper.get_headers(), timeout=300, ) @@ -3911,6 +3941,7 @@ async def list_tasks(self) -> typing.List[Task]: async def create_task( self, + agent_id: str, *, name: str, description: typing.Optional[str] = OMIT, @@ -3922,6 +3953,8 @@ async def create_task( Parameters: + - agent_id: str. + - name: str. Name of the Task - description: typing.Optional[str]. Optional Description of the Task @@ -3938,6 +3971,7 @@ async def create_task( api_key="YOUR_API_KEY", ) await client.create_task( + agent_id="agent_id", name="name", ) """ @@ -3950,7 +3984,9 @@ async def create_task( _request["input_schema"] = input_schema _response = await self._client_wrapper.httpx_client.request( "POST", - urllib.parse.urljoin(f"{self._client_wrapper.get_base_url()}/", "tasks"), + urllib.parse.urljoin( + f"{self._client_wrapper.get_base_url()}/", f"agents/{agent_id}/tasks" + ), json=jsonable_encoder(_request), headers=self._client_wrapper.get_headers(), timeout=300, @@ -3963,12 +3999,16 @@ async def create_task( raise ApiError(status_code=_response.status_code, body=_response.text) raise ApiError(status_code=_response.status_code, body=_response_json) - async def get_task_execution(self, task_id: str) -> Execution: + async def get_task_execution( + self, task_id: str, execution_id: str + ) -> typing.List[Execution]: """ Parameters: - task_id: str. + + - execution_id: str. --- from julep.client import AsyncJulepApi @@ -3977,18 +4017,20 @@ async def get_task_execution(self, task_id: str) -> Execution: ) await client.get_task_execution( task_id="task_id", + execution_id="execution_id", ) """ _response = await self._client_wrapper.httpx_client.request( "GET", urllib.parse.urljoin( - f"{self._client_wrapper.get_base_url()}/", f"tasks/{task_id}/execution" + f"{self._client_wrapper.get_base_url()}/", + f"tasks/{task_id}/executions/{execution_id}", ), headers=self._client_wrapper.get_headers(), timeout=300, ) if 200 <= _response.status_code < 300: - return pydantic.parse_obj_as(Execution, _response.json()) # type: ignore + return pydantic.parse_obj_as(typing.List[Execution], _response.json()) # type: ignore try: _response_json = _response.json() except JSONDecodeError: @@ -3997,6 +4039,7 @@ async def get_task_execution(self, task_id: str) -> Execution: async def start_task_execution( self, + agent_id: str, task_id: str, *, create_execution_task_id: str, @@ -4007,6 +4050,8 @@ async def start_task_execution( Parameters: + - agent_id: str. + - task_id: str. - create_execution_task_id: str. @@ -4021,6 +4066,7 @@ async def start_task_execution( api_key="YOUR_API_KEY", ) await client.start_task_execution( + agent_id="agent_id", task_id="task_id", create_execution_task_id="task_id", arguments={}, @@ -4030,7 +4076,8 @@ async def start_task_execution( _response = await self._client_wrapper.httpx_client.request( "POST", urllib.parse.urljoin( - f"{self._client_wrapper.get_base_url()}/", f"tasks/{task_id}/execution" + f"{self._client_wrapper.get_base_url()}/", + f"agents/{agent_id}/tasks/{task_id}/executions", ), json=jsonable_encoder( { @@ -4050,11 +4097,13 @@ async def start_task_execution( raise ApiError(status_code=_response.status_code, body=_response.text) raise ApiError(status_code=_response.status_code, body=_response_json) - async def get_task(self, task_id: str) -> Execution: + async def get_task(self, agent_id: str, task_id: str) -> Task: """ Parameters: + - agent_id: str. + - task_id: str. --- from julep.client import AsyncJulepApi @@ -4063,19 +4112,21 @@ async def get_task(self, task_id: str) -> Execution: api_key="YOUR_API_KEY", ) await client.get_task( + agent_id="agent_id", task_id="task_id", ) """ _response = await self._client_wrapper.httpx_client.request( "GET", urllib.parse.urljoin( - f"{self._client_wrapper.get_base_url()}/", f"tasks/{task_id}" + f"{self._client_wrapper.get_base_url()}/", + f"agents/{agent_id}/tasks/{task_id}", ), headers=self._client_wrapper.get_headers(), timeout=300, ) if 200 <= _response.status_code < 300: - return pydantic.parse_obj_as(Execution, _response.json()) # type: ignore + return pydantic.parse_obj_as(Task, _response.json()) # type: ignore try: _response_json = _response.json() except JSONDecodeError: @@ -4083,15 +4134,13 @@ async def get_task(self, task_id: str) -> Execution: raise ApiError(status_code=_response.status_code, body=_response_json) async def get_execution_transition( - self, execution_id: str, transition_id: str + self, execution_id: str ) -> typing.List[ExecutionTransition]: """ Parameters: - execution_id: str. - - - transition_id: str. --- from julep.client import AsyncJulepApi @@ -4100,14 +4149,13 @@ async def get_execution_transition( ) await client.get_execution_transition( execution_id="execution_id", - transition_id="transition_id", ) """ _response = await self._client_wrapper.httpx_client.request( "GET", urllib.parse.urljoin( f"{self._client_wrapper.get_base_url()}/", - f"execution/{execution_id}/transition/{transition_id}", + f"executions/{execution_id}/transitions/", ), headers=self._client_wrapper.get_headers(), timeout=300, @@ -4157,7 +4205,7 @@ async def resume_tool_execution( "PUT", urllib.parse.urljoin( f"{self._client_wrapper.get_base_url()}/", - f"execution/{execution_id}/transition/{transition_id}", + f"executions/{execution_id}/transitions/{transition_id}", ), json=jsonable_encoder({"responses": responses}), headers=self._client_wrapper.get_headers(), diff --git a/sdks/python/poetry.lock b/sdks/python/poetry.lock index f08cda47b..9b11a81ba 100644 --- a/sdks/python/poetry.lock +++ b/sdks/python/poetry.lock @@ -1711,13 +1711,13 @@ files = [ [[package]] name = "openai" -version = "1.30.4" +version = "1.30.5" description = "The official Python library for the openai API" optional = false python-versions = ">=3.7.1" files = [ - {file = "openai-1.30.4-py3-none-any.whl", hash = "sha256:fb2635efd270efaf9fac2e07558d7948373b940637d3ae3ab624c1a983d4f03f"}, - {file = "openai-1.30.4.tar.gz", hash = "sha256:f3488d9a1c4e0d332b019377d27d7cb4b3d6103fd5d0a416c7ceac780d1d9b88"}, + {file = "openai-1.30.5-py3-none-any.whl", hash = "sha256:2ad95e926de0d2e09cde632a9204b0a6dca4a03c2cdcc84329b01f355784355a"}, + {file = "openai-1.30.5.tar.gz", hash = "sha256:5366562eb2c5917e6116ae0391b7ae6e3acd62b0ae3f565ada32b35d8fcfa106"}, ] [package.dependencies] diff --git a/sdks/ts/src/api/services/DefaultService.ts b/sdks/ts/src/api/services/DefaultService.ts index 320b90376..030caa617 100644 --- a/sdks/ts/src/api/services/DefaultService.ts +++ b/sdks/ts/src/api/services/DefaultService.ts @@ -971,10 +971,17 @@ export class DefaultService { * @returns Task * @throws ApiError */ - public listTasks(): CancelablePromise> { + public listTasks({ + agentId, + }: { + agentId: string; + }): CancelablePromise> { return this.httpRequest.request({ method: "GET", - url: "/tasks", + url: "/agents/{agent_id}/tasks", + path: { + agent_id: agentId, + }, }); } /** @@ -983,13 +990,18 @@ export class DefaultService { * @throws ApiError */ public createTask({ + agentId, requestBody, }: { + agentId: string; requestBody?: CreateTask; }): CancelablePromise { return this.httpRequest.request({ method: "POST", - url: "/tasks", + url: "/agents/{agent_id}/tasks", + path: { + agent_id: agentId, + }, body: requestBody, mediaType: "application/json", }); @@ -1000,16 +1012,19 @@ export class DefaultService { * @throws ApiError */ public startTaskExecution({ + agentId, taskId, requestBody, }: { + agentId: string; taskId: string; requestBody?: CreateExecution; }): CancelablePromise { return this.httpRequest.request({ method: "POST", - url: "/tasks/{task_id}/execution", + url: "/agents/{agent_id}/tasks/{task_id}/executions", path: { + agent_id: agentId, task_id: taskId, }, body: requestBody, @@ -1017,38 +1032,49 @@ export class DefaultService { }); } /** - * Get execution (status) of a Task + * List Executions of a Task * @returns Execution * @throws ApiError */ public getTaskExecution({ + agentId, taskId, }: { + agentId: string; taskId: string; - }): CancelablePromise { + }): CancelablePromise> { return this.httpRequest.request({ method: "GET", - url: "/tasks/{task_id}/execution", + url: "/agents/{agent_id}/tasks/{task_id}/executions", path: { + agent_id: agentId, task_id: taskId, }, }); } /** * Get a Task by ID - * @returns Execution + * @returns Task * @throws ApiError */ - public getTask({ taskId }: { taskId: string }): CancelablePromise { + public getTask({ + agentId, + taskId, + }: { + agentId: string; + taskId: string; + }): CancelablePromise { return this.httpRequest.request({ method: "GET", - url: "/tasks/{task_id}", + url: "/agents/{agent_id}/tasks/{task_id}", path: { + agent_id: agentId, task_id: taskId, }, }); } /** + * Get an Execution Transition * @returns ExecutionTransition * @throws ApiError */ @@ -1058,10 +1084,10 @@ export class DefaultService { }: { executionId: string; transitionId: string; - }): CancelablePromise> { + }): CancelablePromise { return this.httpRequest.request({ method: "GET", - url: "/execution/{execution_id}/transition/{transition_id}", + url: "/executions/{execution_id}/transitions/{transition_id}", path: { execution_id: executionId, transition_id: transitionId, @@ -1085,7 +1111,7 @@ export class DefaultService { }): CancelablePromise { return this.httpRequest.request({ method: "PUT", - url: "/execution/{execution_id}/transition/{transition_id}", + url: "/executions/{execution_id}/transitions/{transition_id}", path: { execution_id: executionId, transition_id: transitionId, @@ -1094,4 +1120,43 @@ export class DefaultService { mediaType: "application/json", }); } + /** + * Get Task Execution By ID + * @returns Execution + * @throws ApiError + */ + public getTaskExecution1({ + taskId, + executionId, + }: { + taskId: string; + executionId: string; + }): CancelablePromise> { + return this.httpRequest.request({ + method: "GET", + url: "/tasks/{task_id}/executions/{execution_id}", + path: { + task_id: taskId, + execution_id: executionId, + }, + }); + } + /** + * List Transitions of an Execution + * @returns ExecutionTransition + * @throws ApiError + */ + public getExecutionTransition1({ + executionId, + }: { + executionId: string; + }): CancelablePromise> { + return this.httpRequest.request({ + method: "GET", + url: "/executions/{execution_id}/transitions/", + path: { + execution_id: executionId, + }, + }); + } } From c39ca0a2a2f45e1e9694fac571722260267a0a9e Mon Sep 17 00:00:00 2001 From: Diwank Tomer Date: Wed, 5 Jun 2024 11:14:25 -0700 Subject: [PATCH 6/7] feat(agents-api): Add tests for task queries and fix bugs Signed-off-by: Diwank Tomer --- .../agents_api/models/task/create_task.py | 8 +- agents-api/agents_api/models/task/get_task.py | 5 +- .../agents_api/models/task/list_tasks.py | 6 +- .../models/task/test_task_queries.py | 88 +++++++++++++++++++ .../agents_api/models/task/update_task.py | 2 +- 5 files changed, 100 insertions(+), 9 deletions(-) create mode 100644 agents-api/agents_api/models/task/test_task_queries.py diff --git a/agents-api/agents_api/models/task/create_task.py b/agents-api/agents_api/models/task/create_task.py index 70d88c76b..6b33240e4 100644 --- a/agents-api/agents_api/models/task/create_task.py +++ b/agents-api/agents_api/models/task/create_task.py @@ -13,12 +13,12 @@ @cozo_query @beartype def create_task_query( + developer_id: UUID, agent_id: UUID, task_id: UUID, - developer_id: UUID, name: str, description: str, - input_schema: Dict[str, any], + input_schema: Dict[str, Any], tools_available: List[UUID] = [], workflows: List[Dict[str, Any]] = [], ) -> tuple[str, dict]: @@ -26,8 +26,8 @@ def create_task_query( query = """ { ?[agent_id, task_id, name, description, input_schema, tools_available, workflows] <- [[ - to_uuid($agent_id), to_uuid($task_id), $name, $description, $input_schema, [to_uuid($tools_available)], $workflows - ]] + to_uuid($agent_id), to_uuid($task_id), $name, $description, $input_schema, $tools_available, $workflows + ]] :insert tasks { agent_id, diff --git a/agents-api/agents_api/models/task/get_task.py b/agents-api/agents_api/models/task/get_task.py index e9be458d8..c46b98434 100644 --- a/agents-api/agents_api/models/task/get_task.py +++ b/agents-api/agents_api/models/task/get_task.py @@ -13,6 +13,7 @@ def get_task_query( # TODO: Check for agent in developer ID; Assert whether dev can access agent and by relation the task query = """ ?[ + id, name, description, input_schema, @@ -31,7 +32,9 @@ def get_task_query( workflows, created_at, @ 'NOW' - }, updated_at = to_int(updated_at_ms) / 1000 + }, + updated_at = to_int(updated_at_ms) / 1000, + id = to_uuid($task_id), """ return (query, {"agent_id": str(agent_id), "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 index f48b424f6..57ceb7b00 100644 --- a/agents-api/agents_api/models/task/list_tasks.py +++ b/agents-api/agents_api/models/task/list_tasks.py @@ -5,8 +5,8 @@ @cozo_query def list_tasks_query( - agent_id: UUID, developer_id: UUID, + agent_id: UUID, limit: int = 100, offset: int = 0, # metadata_filter: dict[str, Any] = {}, @@ -23,7 +23,7 @@ def list_tasks_query( # TODO: Accepts developer ID. Checks if the developer can get this agent, by relation can get the tasks. Assert that the agent exists under the developer. query = """ ?[ - task_id, + id, name, description, input_schema, @@ -33,7 +33,7 @@ def list_tasks_query( updated_at, ] := *tasks { agent_id: to_uuid($agent_id), - task_id, + task_id: id, updated_at_ms, name, description, diff --git a/agents-api/agents_api/models/task/test_task_queries.py b/agents-api/agents_api/models/task/test_task_queries.py new file mode 100644 index 000000000..8321aeb47 --- /dev/null +++ b/agents-api/agents_api/models/task/test_task_queries.py @@ -0,0 +1,88 @@ +# Tests for task queries +from uuid import uuid4 + +from cozo_migrate.api import init, apply +from pycozo import Client +from ward import test + +from .create_task import create_task_query +from .get_task import get_task_query +from .list_tasks import list_tasks_query +from .update_task import update_task_query + + +def cozo_client(migrations_dir: str = "./migrations"): + # Create a new client for each test + # and initialize the schema. + client = Client() + + init(client) + apply(client, migrations_dir=migrations_dir, all_=True) + + return client + + +@test("model: create task") +def _(): + client = cozo_client() + developer_id = uuid4() + agent_id = uuid4() + task_id = uuid4() + + create_task_query( + developer_id=developer_id, + agent_id=agent_id, + task_id=task_id, + name="test task", + description="test task about", + input_schema={"type": "object", "additionalProperties": True}, + client=client, + ) + +@test("model: list tasks") +def _(): + client = cozo_client() + developer_id = uuid4() + agent_id = uuid4() + + result = list_tasks_query( + developer_id=developer_id, + agent_id=agent_id, + client=client, + ) + + assert len(result["id"]) == 0 + + +@test("model: get task exists") +def _(): + client = cozo_client() + developer_id = uuid4() + agent_id = uuid4() + task_id = uuid4() + + create_task_query( + developer_id=developer_id, + agent_id=agent_id, + task_id=task_id, + name="test task", + description="test task about", + input_schema={"type": "object", "additionalProperties": True}, + client=client, + ) + + result = get_task_query( + agent_id=agent_id, task_id=task_id, developer_id=developer_id, client=client + ) + + assert len(result["id"]) == 1 + + +# @test("model: delete task") +# def _(): +# raise NotImplementedError + + +# @test("model: update task") +# def _(): +# raise NotImplementedError diff --git a/agents-api/agents_api/models/task/update_task.py b/agents-api/agents_api/models/task/update_task.py index 42e89ecbf..fafc648cd 100644 --- a/agents-api/agents_api/models/task/update_task.py +++ b/agents-api/agents_api/models/task/update_task.py @@ -11,7 +11,7 @@ @cozo_query -def cupdate_task_query( +def update_task_query( task_id: UUID, developer_id: UUID, agent_id: UUID, From 3c241b23077b555493093e4c0033561021ab1cf8 Mon Sep 17 00:00:00 2001 From: Diwank Tomer Date: Wed, 5 Jun 2024 18:47:54 -0700 Subject: [PATCH 7/7] fix(agents-api): Fix cozo queries and make tests pass Signed-off-by: Diwank Tomer --- .../models/execution/create_execution.py | 9 +- .../execution/create_execution_transition.py | 49 ++++ .../execution/get_execution_transition.py | 3 + .../execution/list_execution_transitions.py | 5 +- .../execution/test_execution_queries.py | 257 ++++++++++++++++++ .../execution/update_execution_status.py | 9 +- .../execution/update_execution_transition.py | 10 +- .../agents_api/models/task/list_tasks.py | 3 + .../models/task/test_task_queries.py | 4 +- .../agents_api/models/task/update_task.py | 6 +- .../agents_api/routers/tasks/routers.py | 7 +- 11 files changed, 351 insertions(+), 11 deletions(-) create mode 100644 agents-api/agents_api/models/execution/create_execution_transition.py create mode 100644 agents-api/agents_api/models/execution/test_execution_queries.py diff --git a/agents-api/agents_api/models/execution/create_execution.py b/agents-api/agents_api/models/execution/create_execution.py index a75a99022..b098e722c 100644 --- a/agents-api/agents_api/models/execution/create_execution.py +++ b/agents-api/agents_api/models/execution/create_execution.py @@ -1,17 +1,20 @@ +from typing import Literal, Dict, Any from uuid import UUID +from beartype import beartype + from ..utils import cozo_query -from typing import Literal, Dict, Any @cozo_query +@beartype def create_execution_query( + developer_id: UUID, agent_id: UUID, task_id: UUID, execution_id: UUID, - developer_id: UUID, status: Literal[ - "queued", "starting", "running", "waiting-for-input", "success", "failed" + "queued", "starting", "running", "waiting_for_input", "success", "failed" ] = "queued", arguments: Dict[str, Any] = {}, ) -> tuple[str, dict]: diff --git a/agents-api/agents_api/models/execution/create_execution_transition.py b/agents-api/agents_api/models/execution/create_execution_transition.py new file mode 100644 index 000000000..41cb2477b --- /dev/null +++ b/agents-api/agents_api/models/execution/create_execution_transition.py @@ -0,0 +1,49 @@ +from typing import Literal, Dict, Any +from uuid import UUID + +from beartype import beartype + +from ..utils import cozo_query + + +@cozo_query +@beartype +def create_execution_transition_query( + developer_id: UUID, + execution_id: UUID, + transition_id: UUID, + type_: Literal["finished", "waiting", "error", "step"], + from_: tuple[str, int], + to: tuple[str, int] | None, + output: Dict[str, Any], +) -> tuple[str, dict]: + # TODO: Check for agent in developer ID; Assert whether dev can access agent and by relation the task + # TODO: Check for task and execution + + query = """ +{ + ?[execution_id, transition_id, type, from, to, output] <- [[ + to_uuid($execution_id), + to_uuid($transition_id), + $type, + $from, + $to, + $output, + ]] + + :insert transitions { + execution_id, transition_id, type, from, to, output + } +} +""" + return ( + query, + { + "execution_id": str(execution_id), + "transition_id": str(transition_id), + "type": type_, + "from": from_, + "to": to, + "output": output, + }, + ) diff --git a/agents-api/agents_api/models/execution/get_execution_transition.py b/agents-api/agents_api/models/execution/get_execution_transition.py index 06935528c..1e4842e0c 100644 --- a/agents-api/agents_api/models/execution/get_execution_transition.py +++ b/agents-api/agents_api/models/execution/get_execution_transition.py @@ -1,9 +1,12 @@ from uuid import UUID +from beartype import beartype + from ..utils import cozo_query @cozo_query +@beartype def get_execution_transition_query( execution_id: UUID, transition_id: UUID ) -> tuple[str, dict]: diff --git a/agents-api/agents_api/models/execution/list_execution_transitions.py b/agents-api/agents_api/models/execution/list_execution_transitions.py index 4b87e79ac..dd8549364 100644 --- a/agents-api/agents_api/models/execution/list_execution_transitions.py +++ b/agents-api/agents_api/models/execution/list_execution_transitions.py @@ -1,10 +1,13 @@ from uuid import UUID +from beartype import beartype + from ..utils import cozo_query @cozo_query -def list_execution_transition_query(execution_id: UUID) -> tuple[str, dict]: +@beartype +def list_execution_transitions_query(execution_id: UUID) -> tuple[str, dict]: query = """ { diff --git a/agents-api/agents_api/models/execution/test_execution_queries.py b/agents-api/agents_api/models/execution/test_execution_queries.py new file mode 100644 index 000000000..cf66511c6 --- /dev/null +++ b/agents-api/agents_api/models/execution/test_execution_queries.py @@ -0,0 +1,257 @@ +# Tests for execution queries +from uuid import uuid4 + +from cozo_migrate.api import init, apply +from pycozo import Client +from ward import test + +from .create_execution import create_execution_query +from .get_execution_status import get_execution_status_query +from .get_execution import get_execution_query +from .list_executions import list_task_executions_query +from .update_execution_status import update_execution_status_query +from .create_execution_transition import create_execution_transition_query +from .get_execution_transition import get_execution_transition_query +from .list_execution_transitions import list_execution_transitions_query +from .update_execution_transition import update_execution_transition_query + + +def cozo_client(migrations_dir: str = "./migrations"): + # Create a new client for each test + # and initialize the schema. + client = Client() + + init(client) + apply(client, migrations_dir=migrations_dir, all_=True) + + return client + + +@test("model: create execution") +def _(): + client = cozo_client() + developer_id = uuid4() + agent_id = uuid4() + task_id = uuid4() + execution_id = uuid4() + + create_execution_query( + developer_id=developer_id, + agent_id=agent_id, + task_id=task_id, + execution_id=execution_id, + arguments={"input": "test"}, + client=client, + ) + + +@test("model: get execution") +def _(): + client = cozo_client() + developer_id = uuid4() + agent_id = uuid4() + task_id = uuid4() + execution_id = uuid4() + + create_execution_query( + developer_id=developer_id, + agent_id=agent_id, + task_id=task_id, + execution_id=execution_id, + arguments={"input": "test"}, + client=client, + ) + + result = get_execution_query( + task_id=task_id, execution_id=execution_id, client=client + ) + + assert len(result["status"]) == 1 + assert result["status"][0] == "queued" + + +@test("model: get execution status") +def _(): + client = cozo_client() + developer_id = uuid4() + agent_id = uuid4() + task_id = uuid4() + execution_id = uuid4() + + create_execution_query( + developer_id=developer_id, + agent_id=agent_id, + task_id=task_id, + execution_id=execution_id, + arguments={"input": "test"}, + client=client, + ) + + result = get_execution_status_query( + task_id=task_id, execution_id=execution_id, client=client + ) + + assert len(result["status"]) == 1 + assert result["status"][0] == "queued" + + +@test("model: list executions empty") +def _(): + client = cozo_client() + developer_id = uuid4() + agent_id = uuid4() + task_id = uuid4() + + result = list_task_executions_query( + task_id=task_id, agent_id=agent_id, developer_id=developer_id, client=client + ) + + assert len(result) == 0 + + +@test("model: list executions") +def _(): + client = cozo_client() + developer_id = uuid4() + agent_id = uuid4() + task_id = uuid4() + execution_id = uuid4() + + create_execution_query( + developer_id=developer_id, + agent_id=agent_id, + task_id=task_id, + execution_id=execution_id, + arguments={"input": "test"}, + client=client, + ) + + result = list_task_executions_query( + task_id=task_id, agent_id=agent_id, developer_id=developer_id, client=client + ) + + assert len(result["status"]) == 1 + assert result["status"][0] == "queued" + + +@test("model: update execution status") +def _(): + client = cozo_client() + developer_id = uuid4() + agent_id = uuid4() + task_id = uuid4() + execution_id = uuid4() + + create_execution_query( + developer_id=developer_id, + agent_id=agent_id, + task_id=task_id, + execution_id=execution_id, + arguments={"input": "test"}, + client=client, + ) + + result = update_execution_status_query( + task_id=task_id, execution_id=execution_id, status="running", client=client + ) + + updated_rows = result[result["_kind"] == "inserted"].reset_index() + assert len(updated_rows) == 1 + assert updated_rows["status"][0] == "running" + + +@test("model: create execution transition") +def _(): + client = cozo_client() + developer_id = uuid4() + execution_id = uuid4() + transition_id = uuid4() + + create_execution_transition_query( + developer_id=developer_id, + execution_id=execution_id, + transition_id=transition_id, + type_="step", + from_=("test", 1), + to=("test", 2), + output={"input": "test"}, + client=client, + ) + + +@test("model: get execution transition") +def _(): + client = cozo_client() + developer_id = uuid4() + execution_id = uuid4() + transition_id = uuid4() + + create_execution_transition_query( + developer_id=developer_id, + execution_id=execution_id, + transition_id=transition_id, + type_="step", + from_=("test", 1), + to=("test", 2), + output={"input": "test"}, + client=client, + ) + + result = get_execution_transition_query( + execution_id=execution_id, transition_id=transition_id, client=client + ) + + assert len(result["type"]) == 1 + + +@test("model: list execution transitions") +def _(): + client = cozo_client() + developer_id = uuid4() + execution_id = uuid4() + transition_id = uuid4() + + create_execution_transition_query( + developer_id=developer_id, + execution_id=execution_id, + transition_id=transition_id, + type_="step", + from_=("test", 1), + to=("test", 2), + output={"input": "test"}, + client=client, + ) + + result = list_execution_transitions_query(execution_id=execution_id, client=client) + + assert len(result["type"]) == 1 + + +@test("model: update execution transitions") +def _(): + client = cozo_client() + developer_id = uuid4() + execution_id = uuid4() + transition_id = uuid4() + + create_execution_transition_query( + developer_id=developer_id, + execution_id=execution_id, + transition_id=transition_id, + type_="step", + from_=("test", 1), + to=("test", 2), + output={"input": "test"}, + client=client, + ) + + result = update_execution_transition_query( + execution_id=execution_id, + transition_id=transition_id, + type="finished", + client=client, + ) + + updated_rows = result[result["_kind"] == "inserted"].reset_index() + assert len(updated_rows) == 1 + assert updated_rows["type"][0] == "finished" diff --git a/agents-api/agents_api/models/execution/update_execution_status.py b/agents-api/agents_api/models/execution/update_execution_status.py index bd8783a2b..0b5c0f692 100644 --- a/agents-api/agents_api/models/execution/update_execution_status.py +++ b/agents-api/agents_api/models/execution/update_execution_status.py @@ -1,15 +1,18 @@ +from typing import Literal, Dict, Any from uuid import UUID +from beartype import beartype + from ..utils import cozo_query -from typing import Literal, Dict, Any @cozo_query +@beartype def update_execution_status_query( task_id: UUID, execution_id: UUID, status: Literal[ - "queued", "starting", "running", "waiting-for-input", "success", "failed" + "queued", "starting", "running", "waiting_for_input", "success", "failed" ], arguments: Dict[str, Any] = {}, ) -> tuple[str, dict]: @@ -28,6 +31,8 @@ def update_execution_status_query( status, updated_at } + + :returning } """ diff --git a/agents-api/agents_api/models/execution/update_execution_transition.py b/agents-api/agents_api/models/execution/update_execution_transition.py index 2feb3ca31..f9988fc0e 100644 --- a/agents-api/agents_api/models/execution/update_execution_transition.py +++ b/agents-api/agents_api/models/execution/update_execution_transition.py @@ -1,5 +1,7 @@ from uuid import UUID +from beartype import beartype + from agents_api.common.utils.datetime import utcnow from ..utils import cozo_query @@ -7,9 +9,15 @@ @cozo_query +@beartype def update_execution_transition_query( execution_id: UUID, transition_id: UUID, **update_data ) -> tuple[str, dict]: + # Only output and type can be updated + assert update_data.keys() <= { + "output", + "type", + }, "Only output and type can be updated for a transition" transition_update_cols, transition_update_vals = cozo_process_mutate_data( { @@ -40,7 +48,7 @@ def update_execution_transition_query( return ( query, { - "transitioon_update_vals": transition_update_vals, + "transition_update_vals": transition_update_vals, "execution_id": str(execution_id), "transition_id": str(transition_id), }, diff --git a/agents-api/agents_api/models/task/list_tasks.py b/agents-api/agents_api/models/task/list_tasks.py index 57ceb7b00..730772a17 100644 --- a/agents-api/agents_api/models/task/list_tasks.py +++ b/agents-api/agents_api/models/task/list_tasks.py @@ -1,9 +1,12 @@ from uuid import UUID +from beartype import beartype + from ..utils import cozo_query @cozo_query +@beartype def list_tasks_query( developer_id: UUID, agent_id: UUID, diff --git a/agents-api/agents_api/models/task/test_task_queries.py b/agents-api/agents_api/models/task/test_task_queries.py index 8321aeb47..9e2a48ab5 100644 --- a/agents-api/agents_api/models/task/test_task_queries.py +++ b/agents-api/agents_api/models/task/test_task_queries.py @@ -8,7 +8,6 @@ from .create_task import create_task_query from .get_task import get_task_query from .list_tasks import list_tasks_query -from .update_task import update_task_query def cozo_client(migrations_dir: str = "./migrations"): @@ -39,6 +38,7 @@ def _(): client=client, ) + @test("model: list tasks") def _(): client = cozo_client() @@ -80,9 +80,11 @@ def _(): # @test("model: delete task") # def _(): +# TODO: Implement this test # raise NotImplementedError # @test("model: update task") # def _(): +# TODO: Implement this test # raise NotImplementedError diff --git a/agents-api/agents_api/models/task/update_task.py b/agents-api/agents_api/models/task/update_task.py index fafc648cd..3e324fc22 100644 --- a/agents-api/agents_api/models/task/update_task.py +++ b/agents-api/agents_api/models/task/update_task.py @@ -6,20 +6,22 @@ from uuid import UUID from typing import List, Dict, Any +from beartype import beartype from ..utils import cozo_query @cozo_query +@beartype def update_task_query( task_id: UUID, developer_id: UUID, agent_id: UUID, name: str, description: str, - input_schema: Dict[str, any], + input_schema: Dict[str, Any], tools_available: List[UUID] = [], workflows: List[Dict[str, Any]] = [], ) -> tuple[str, dict]: # NOT TO IMPLEMENT FOR NOW - pass + raise NotImplementedError("Not implemented yet") diff --git a/agents-api/agents_api/routers/tasks/routers.py b/agents-api/agents_api/routers/tasks/routers.py index 780bf2a06..3f189f660 100644 --- a/agents-api/agents_api/routers/tasks/routers.py +++ b/agents-api/agents_api/routers/tasks/routers.py @@ -96,6 +96,7 @@ async def get_task( row.to_dict() for _, row in get_task_query(agent_id, task_id, x_developer_id).iterrows() ][0] + return Task(**resp) except (IndexError, KeyError): raise HTTPException( @@ -108,6 +109,8 @@ async def get_task( status_code=status.HTTP_404_NOT_FOUND, detail="Task not found" ) + raise + @router.post( "/agents/{agent_id}/tasks/{task_id}/executions", @@ -209,9 +212,11 @@ async def update_execution_transition( ) -> ResourceUpdatedResponse: # try: # resp = update_execution_transition_query(execution_id, transition_id, request) - pass + # OpenAPI Model doesn't have update execution transition + raise NotImplementedError("Not implemented yet") + @router.get("/executions/{execution_id}/transitions") async def list_execution_transitions(execution_id: UUID4) -> ExecutionTransitionList: