From 48439d459af9ede3817b59515dee432de99d5f3f Mon Sep 17 00:00:00 2001 From: Dmitry Paramonov Date: Wed, 18 Dec 2024 15:39:35 +0300 Subject: [PATCH 01/16] chore: Move ti queries directory --- .../models/chat/get_cached_response.py | 15 --------------- .../models/chat/set_cached_response.py | 19 ------------------- .../{models => queries}/chat/__init__.py | 2 -- .../chat/gather_messages.py | 0 .../chat/prepare_chat_context.py | 15 +++++++-------- 5 files changed, 7 insertions(+), 44 deletions(-) delete mode 100644 agents-api/agents_api/models/chat/get_cached_response.py delete mode 100644 agents-api/agents_api/models/chat/set_cached_response.py rename agents-api/agents_api/{models => queries}/chat/__init__.py (92%) rename agents-api/agents_api/{models => queries}/chat/gather_messages.py (100%) rename agents-api/agents_api/{models => queries}/chat/prepare_chat_context.py (92%) diff --git a/agents-api/agents_api/models/chat/get_cached_response.py b/agents-api/agents_api/models/chat/get_cached_response.py deleted file mode 100644 index 368c88567..000000000 --- a/agents-api/agents_api/models/chat/get_cached_response.py +++ /dev/null @@ -1,15 +0,0 @@ -from beartype import beartype - -from ..utils import cozo_query - - -@cozo_query -@beartype -def get_cached_response(key: str) -> tuple[str, dict]: - query = """ - input[key] <- [[$key]] - ?[key, value] := input[key], *session_cache{key, value} - :limit 1 - """ - - return (query, {"key": key}) diff --git a/agents-api/agents_api/models/chat/set_cached_response.py b/agents-api/agents_api/models/chat/set_cached_response.py deleted file mode 100644 index 8625f3f1b..000000000 --- a/agents-api/agents_api/models/chat/set_cached_response.py +++ /dev/null @@ -1,19 +0,0 @@ -from beartype import beartype - -from ..utils import cozo_query - - -@cozo_query -@beartype -def set_cached_response(key: str, value: dict) -> tuple[str, dict]: - query = """ - ?[key, value] <- [[$key, $value]] - - :insert session_cache { - key => value - } - - :returning - """ - - return (query, {"key": key, "value": value}) diff --git a/agents-api/agents_api/models/chat/__init__.py b/agents-api/agents_api/queries/chat/__init__.py similarity index 92% rename from agents-api/agents_api/models/chat/__init__.py rename to agents-api/agents_api/queries/chat/__init__.py index 428b72572..2c05b4f8b 100644 --- a/agents-api/agents_api/models/chat/__init__.py +++ b/agents-api/agents_api/queries/chat/__init__.py @@ -17,6 +17,4 @@ # ruff: noqa: F401, F403, F405 from .gather_messages import gather_messages -from .get_cached_response import get_cached_response from .prepare_chat_context import prepare_chat_context -from .set_cached_response import set_cached_response diff --git a/agents-api/agents_api/models/chat/gather_messages.py b/agents-api/agents_api/queries/chat/gather_messages.py similarity index 100% rename from agents-api/agents_api/models/chat/gather_messages.py rename to agents-api/agents_api/queries/chat/gather_messages.py diff --git a/agents-api/agents_api/models/chat/prepare_chat_context.py b/agents-api/agents_api/queries/chat/prepare_chat_context.py similarity index 92% rename from agents-api/agents_api/models/chat/prepare_chat_context.py rename to agents-api/agents_api/queries/chat/prepare_chat_context.py index f77686d7a..4731618f8 100644 --- a/agents-api/agents_api/models/chat/prepare_chat_context.py +++ b/agents-api/agents_api/queries/chat/prepare_chat_context.py @@ -3,7 +3,6 @@ from beartype import beartype from fastapi import HTTPException -from pycozo.client import QueryException from pydantic import ValidationError from ...common.protocol.sessions import ChatContext, make_session @@ -22,13 +21,13 @@ T = TypeVar("T") -@rewrap_exceptions( - { - QueryException: partialclass(HTTPException, status_code=400), - ValidationError: partialclass(HTTPException, status_code=400), - TypeError: partialclass(HTTPException, status_code=400), - } -) +# TODO: implement this part +# @rewrap_exceptions( +# { +# ValidationError: partialclass(HTTPException, status_code=400), +# TypeError: partialclass(HTTPException, status_code=400), +# } +# ) @wrap_in_class( ChatContext, one=True, From 7b6c502eb1de5580c93dd9e38f59ab3bf5512878 Mon Sep 17 00:00:00 2001 From: Dmitry Paramonov Date: Thu, 19 Dec 2024 15:34:37 +0300 Subject: [PATCH 02/16] feat: Add prepare chat context query --- .../queries/chat/gather_messages.py | 12 +- .../queries/chat/prepare_chat_context.py | 225 ++++++++++-------- 2 files changed, 129 insertions(+), 108 deletions(-) diff --git a/agents-api/agents_api/queries/chat/gather_messages.py b/agents-api/agents_api/queries/chat/gather_messages.py index 28dc6607f..34a7c564f 100644 --- a/agents-api/agents_api/queries/chat/gather_messages.py +++ b/agents-api/agents_api/queries/chat/gather_messages.py @@ -3,18 +3,17 @@ from beartype import beartype from fastapi import HTTPException -from pycozo.client import QueryException from pydantic import ValidationError from ...autogen.openapi_model import ChatInput, DocReference, History from ...clients import litellm from ...common.protocol.developers import Developer from ...common.protocol.sessions import ChatContext -from ..docs.search_docs_by_embedding import search_docs_by_embedding -from ..docs.search_docs_by_text import search_docs_by_text -from ..docs.search_docs_hybrid import search_docs_hybrid -from ..entry.get_history import get_history -from ..session.get_session import get_session +# from ..docs.search_docs_by_embedding import search_docs_by_embedding +# from ..docs.search_docs_by_text import search_docs_by_text +# from ..docs.search_docs_hybrid import search_docs_hybrid +# from ..entry.get_history import get_history +from ..sessions.get_session import get_session from ..utils import ( partialclass, rewrap_exceptions, @@ -25,7 +24,6 @@ @rewrap_exceptions( { - QueryException: partialclass(HTTPException, status_code=400), ValidationError: partialclass(HTTPException, status_code=400), TypeError: partialclass(HTTPException, status_code=400), } diff --git a/agents-api/agents_api/queries/chat/prepare_chat_context.py b/agents-api/agents_api/queries/chat/prepare_chat_context.py index 4731618f8..23926ea4c 100644 --- a/agents-api/agents_api/queries/chat/prepare_chat_context.py +++ b/agents-api/agents_api/queries/chat/prepare_chat_context.py @@ -2,18 +2,10 @@ from uuid import UUID from beartype import beartype -from fastapi import HTTPException -from pydantic import ValidationError from ...common.protocol.sessions import ChatContext, make_session -from ..session.prepare_session_data import prepare_session_data from ..utils import ( - cozo_query, - fix_uuid_if_present, - partialclass, - rewrap_exceptions, - verify_developer_id_query, - verify_developer_owns_resource_query, + pg_query, wrap_in_class, ) @@ -21,17 +13,107 @@ T = TypeVar("T") -# TODO: implement this part -# @rewrap_exceptions( -# { -# ValidationError: partialclass(HTTPException, status_code=400), -# TypeError: partialclass(HTTPException, status_code=400), -# } -# ) -@wrap_in_class( - ChatContext, - one=True, - transform=lambda d: { +query = """ +SELECT * FROM +( + SELECT jsonb_agg(u) AS users FROM ( + SELECT + session_lookup.participant_id, + users.user_id AS id, + users.developer_id, + users.name, + users.about, + users.created_at, + users.updated_at, + users.metadata + FROM session_lookup + INNER JOIN users ON session_lookup.participant_id = users.user_id + WHERE + session_lookup.developer_id = $1 AND + session_id = $2 AND + session_lookup.participant_type = 'user' + ) u +) AS users, +( + SELECT jsonb_agg(a) AS agents FROM ( + SELECT + session_lookup.participant_id, + agents.agent_id AS id, + agents.developer_id, + agents.canonical_name, + agents.name, + agents.about, + agents.instructions, + agents.model, + agents.created_at, + agents.updated_at, + agents.metadata, + agents.default_settings + FROM session_lookup + INNER JOIN agents ON session_lookup.participant_id = agents.agent_id + WHERE + session_lookup.developer_id = $1 AND + session_id = $2 AND + session_lookup.participant_type = 'agent' + ) a +) AS agents, +( + SELECT to_jsonb(s) AS session FROM ( + SELECT + sessions.session_id AS id, + sessions.developer_id, + sessions.situation, + sessions.system_template, + sessions.created_at, + sessions.metadata, + sessions.render_templates, + sessions.token_budget, + sessions.context_overflow, + sessions.forward_tool_calls, + sessions.recall_options + FROM sessions + WHERE + developer_id = $1 AND + session_id = $2 + LIMIT 1 + ) s +) AS session, +( + SELECT jsonb_agg(r) AS toolsets FROM ( + SELECT + session_lookup.participant_id, + tools.tool_id as id, + tools.developer_id, + tools.agent_id, + tools.task_id, + tools.task_version, + tools.type, + tools.name, + tools.description, + tools.spec, + tools.updated_at, + tools.created_at + FROM session_lookup + INNER JOIN tools ON session_lookup.participant_id = tools.agent_id + WHERE + session_lookup.developer_id = $1 AND + session_id = $2 AND + session_lookup.participant_type = 'agent' + ) r +) AS toolsets +""" + + +def _transform(d): + toolsets = {} + for tool in d["toolsets"]: + agent_id = tool["agent_id"] + if agent_id in toolsets: + toolsets[agent_id].append(tool) + else: + toolsets[agent_id] = [tool] + + return { **d, "session": make_session( agents=[a["id"] for a in d["agents"]], @@ -40,103 +122,44 @@ ), "toolsets": [ { - **ts, + "agent_id": agent_id, "tools": [ { tool["type"]: tool.pop("spec"), **tool, } - for tool in map(fix_uuid_if_present, ts["tools"]) + for tool in tools ], } - for ts in d["toolsets"] + for agent_id, tools in toolsets.items() ], - }, + } + + +# TODO: implement this part +# @rewrap_exceptions( +# { +# ValidationError: partialclass(HTTPException, status_code=400), +# TypeError: partialclass(HTTPException, status_code=400), +# } +# ) +@wrap_in_class( + ChatContext, + one=True, + transform=_transform, ) -@cozo_query +@pg_query @beartype -def prepare_chat_context( +async def prepare_chat_context( *, developer_id: UUID, session_id: UUID, -) -> tuple[list[str], dict]: +) -> tuple[list[str], list]: """ Executes a complex query to retrieve memory context based on session ID. """ - [*_, session_data_query], sd_vars = prepare_session_data.__wrapped__( - developer_id=developer_id, session_id=session_id - ) - - session_data_fields = ("session", "agents", "users") - - session_data_query += """ - :create _session_data_json { - agents: [Json], - users: [Json], - session: Json, - } - """ - - toolsets_query = """ - input[session_id] <- [[to_uuid($session_id)]] - - tools_by_agent[agent_id, collect(tool)] := - input[session_id], - *session_lookup{ - session_id, - participant_id: agent_id, - participant_type: "agent", - }, - - *tools { agent_id, tool_id, name, type, spec, description, updated_at, created_at }, - tool = { - "id": tool_id, - "name": name, - "type": type, - "spec": spec, - "description": description, - "updated_at": updated_at, - "created_at": created_at, - } - - agent_toolsets[collect(toolset)] := - tools_by_agent[agent_id, tools], - toolset = { - "agent_id": agent_id, - "tools": tools, - } - - ?[toolsets] := - agent_toolsets[toolsets] - - :create _toolsets_json { - toolsets: [Json], - } - """ - - combine_query = f""" - ?[{', '.join(session_data_fields)}, toolsets] := - *_session_data_json {{ {', '.join(session_data_fields)} }}, - *_toolsets_json {{ toolsets }} - - :limit 1 - """ - - queries = [ - verify_developer_id_query(developer_id), - verify_developer_owns_resource_query( - developer_id, "sessions", session_id=session_id - ), - session_data_query, - toolsets_query, - combine_query, - ] - return ( - queries, - { - "session_id": str(session_id), - **sd_vars, - }, + [query], + [developer_id, session_id], ) From ca12d656e2487ce107b5db10fab8427e6ac9ec3f Mon Sep 17 00:00:00 2001 From: whiterabbit1983 Date: Thu, 19 Dec 2024 12:38:09 +0000 Subject: [PATCH 03/16] refactor: Lint agents-api (CI) --- agents-api/agents_api/queries/chat/gather_messages.py | 1 + 1 file changed, 1 insertion(+) diff --git a/agents-api/agents_api/queries/chat/gather_messages.py b/agents-api/agents_api/queries/chat/gather_messages.py index 34a7c564f..4fd574368 100644 --- a/agents-api/agents_api/queries/chat/gather_messages.py +++ b/agents-api/agents_api/queries/chat/gather_messages.py @@ -9,6 +9,7 @@ from ...clients import litellm from ...common.protocol.developers import Developer from ...common.protocol.sessions import ChatContext + # from ..docs.search_docs_by_embedding import search_docs_by_embedding # from ..docs.search_docs_by_text import search_docs_by_text # from ..docs.search_docs_hybrid import search_docs_hybrid From 0aecd613642c3344520a102610f0bc1ddd3371f8 Mon Sep 17 00:00:00 2001 From: Dmitry Paramonov Date: Thu, 19 Dec 2024 15:49:10 +0300 Subject: [PATCH 04/16] feat: Add SQL validation --- agents-api/agents_api/exceptions.py | 9 ++ .../queries/chat/prepare_chat_context.py | 90 ++++++++++--------- 2 files changed, 56 insertions(+), 43 deletions(-) diff --git a/agents-api/agents_api/exceptions.py b/agents-api/agents_api/exceptions.py index 615958a87..f6fcc4741 100644 --- a/agents-api/agents_api/exceptions.py +++ b/agents-api/agents_api/exceptions.py @@ -49,3 +49,12 @@ class FailedEncodingSentinel: """Sentinel object returned when failed to encode payload.""" payload_data: bytes + + +class QueriesBaseException(AgentsBaseException): + pass + + +class InvalidSQLQuery(QueriesBaseException): + def __init__(self, query_name: str): + super().__init__(f"invalid query: {query_name}") diff --git a/agents-api/agents_api/queries/chat/prepare_chat_context.py b/agents-api/agents_api/queries/chat/prepare_chat_context.py index 23926ea4c..1d9bd52fb 100644 --- a/agents-api/agents_api/queries/chat/prepare_chat_context.py +++ b/agents-api/agents_api/queries/chat/prepare_chat_context.py @@ -1,9 +1,11 @@ from typing import Any, TypeVar from uuid import UUID +import sqlvalidator from beartype import beartype from ...common.protocol.sessions import ChatContext, make_session +from ...exceptions import InvalidSQLQuery from ..utils import ( pg_query, wrap_in_class, @@ -13,19 +15,19 @@ T = TypeVar("T") -query = """ -SELECT * FROM +sql_query = sqlvalidator.parse( + """SELECT * FROM ( SELECT jsonb_agg(u) AS users FROM ( SELECT session_lookup.participant_id, users.user_id AS id, - users.developer_id, - users.name, - users.about, - users.created_at, - users.updated_at, - users.metadata + users.developer_id, + users.name, + users.about, + users.created_at, + users.updated_at, + users.metadata FROM session_lookup INNER JOIN users ON session_lookup.participant_id = users.user_id WHERE @@ -39,16 +41,16 @@ SELECT session_lookup.participant_id, agents.agent_id AS id, - agents.developer_id, - agents.canonical_name, - agents.name, - agents.about, - agents.instructions, - agents.model, - agents.created_at, - agents.updated_at, - agents.metadata, - agents.default_settings + agents.developer_id, + agents.canonical_name, + agents.name, + agents.about, + agents.instructions, + agents.model, + agents.created_at, + agents.updated_at, + agents.metadata, + agents.default_settings FROM session_lookup INNER JOIN agents ON session_lookup.participant_id = agents.agent_id WHERE @@ -58,24 +60,24 @@ ) a ) AS agents, ( - SELECT to_jsonb(s) AS session FROM ( + SELECT to_jsonb(s) AS session FROM ( SELECT sessions.session_id AS id, - sessions.developer_id, - sessions.situation, - sessions.system_template, - sessions.created_at, - sessions.metadata, - sessions.render_templates, - sessions.token_budget, - sessions.context_overflow, - sessions.forward_tool_calls, - sessions.recall_options + sessions.developer_id, + sessions.situation, + sessions.system_template, + sessions.created_at, + sessions.metadata, + sessions.render_templates, + sessions.token_budget, + sessions.context_overflow, + sessions.forward_tool_calls, + sessions.recall_options FROM sessions WHERE developer_id = $1 AND session_id = $2 - LIMIT 1 + LIMIT 1 ) s ) AS session, ( @@ -83,16 +85,16 @@ SELECT session_lookup.participant_id, tools.tool_id as id, - tools.developer_id, - tools.agent_id, - tools.task_id, - tools.task_version, - tools.type, - tools.name, - tools.description, - tools.spec, - tools.updated_at, - tools.created_at + tools.developer_id, + tools.agent_id, + tools.task_id, + tools.task_version, + tools.type, + tools.name, + tools.description, + tools.spec, + tools.updated_at, + tools.created_at FROM session_lookup INNER JOIN tools ON session_lookup.participant_id = tools.agent_id WHERE @@ -100,8 +102,10 @@ session_id = $2 AND session_lookup.participant_type = 'agent' ) r -) AS toolsets -""" +) AS toolsets""" +) +if not sql_query.is_valid(): + raise InvalidSQLQuery("prepare_chat_context") def _transform(d): @@ -160,6 +164,6 @@ async def prepare_chat_context( """ return ( - [query], + [sql_query.format()], [developer_id, session_id], ) From 0d288c43ab50c5a855680ef02d41d1147853a310 Mon Sep 17 00:00:00 2001 From: Dmitry Paramonov Date: Fri, 20 Dec 2024 11:49:21 +0300 Subject: [PATCH 05/16] chore: Import other required queries --- agents-api/agents_api/queries/chat/gather_messages.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/agents-api/agents_api/queries/chat/gather_messages.py b/agents-api/agents_api/queries/chat/gather_messages.py index 4fd574368..94d5fe71a 100644 --- a/agents-api/agents_api/queries/chat/gather_messages.py +++ b/agents-api/agents_api/queries/chat/gather_messages.py @@ -10,10 +10,10 @@ from ...common.protocol.developers import Developer from ...common.protocol.sessions import ChatContext -# from ..docs.search_docs_by_embedding import search_docs_by_embedding -# from ..docs.search_docs_by_text import search_docs_by_text -# from ..docs.search_docs_hybrid import search_docs_hybrid -# from ..entry.get_history import get_history +from ..docs.search_docs_by_embedding import search_docs_by_embedding +from ..docs.search_docs_by_text import search_docs_by_text +from ..docs.search_docs_hybrid import search_docs_hybrid +from ..entries.get_history import get_history from ..sessions.get_session import get_session from ..utils import ( partialclass, From 473387bb2325cc6b0f0af96069732c3a2b46db7a Mon Sep 17 00:00:00 2001 From: whiterabbit1983 Date: Fri, 20 Dec 2024 08:50:13 +0000 Subject: [PATCH 06/16] refactor: Lint agents-api (CI) --- agents-api/agents_api/queries/chat/gather_messages.py | 1 - 1 file changed, 1 deletion(-) diff --git a/agents-api/agents_api/queries/chat/gather_messages.py b/agents-api/agents_api/queries/chat/gather_messages.py index 94d5fe71a..cbf3bf209 100644 --- a/agents-api/agents_api/queries/chat/gather_messages.py +++ b/agents-api/agents_api/queries/chat/gather_messages.py @@ -9,7 +9,6 @@ from ...clients import litellm from ...common.protocol.developers import Developer from ...common.protocol.sessions import ChatContext - from ..docs.search_docs_by_embedding import search_docs_by_embedding from ..docs.search_docs_by_text import search_docs_by_text from ..docs.search_docs_hybrid import search_docs_hybrid From 15659c57a0336ea9dff974b69f831ec5dddb5efc Mon Sep 17 00:00:00 2001 From: Dmitry Paramonov Date: Fri, 20 Dec 2024 12:37:19 +0300 Subject: [PATCH 07/16] chore: Move queries to another folder --- .../{models => queries}/tools/__init__.py | 0 .../{models => queries}/tools/create_tools.py | 21 +++++++++---------- .../{models => queries}/tools/delete_tool.py | 0 .../{models => queries}/tools/get_tool.py | 0 .../tools/get_tool_args_from_metadata.py | 0 .../{models => queries}/tools/list_tools.py | 0 .../{models => queries}/tools/patch_tool.py | 0 .../{models => queries}/tools/update_tool.py | 0 8 files changed, 10 insertions(+), 11 deletions(-) rename agents-api/agents_api/{models => queries}/tools/__init__.py (100%) rename agents-api/agents_api/{models => queries}/tools/create_tools.py (89%) rename agents-api/agents_api/{models => queries}/tools/delete_tool.py (100%) rename agents-api/agents_api/{models => queries}/tools/get_tool.py (100%) rename agents-api/agents_api/{models => queries}/tools/get_tool_args_from_metadata.py (100%) rename agents-api/agents_api/{models => queries}/tools/list_tools.py (100%) rename agents-api/agents_api/{models => queries}/tools/patch_tool.py (100%) rename agents-api/agents_api/{models => queries}/tools/update_tool.py (100%) diff --git a/agents-api/agents_api/models/tools/__init__.py b/agents-api/agents_api/queries/tools/__init__.py similarity index 100% rename from agents-api/agents_api/models/tools/__init__.py rename to agents-api/agents_api/queries/tools/__init__.py diff --git a/agents-api/agents_api/models/tools/create_tools.py b/agents-api/agents_api/queries/tools/create_tools.py similarity index 89% rename from agents-api/agents_api/models/tools/create_tools.py rename to agents-api/agents_api/queries/tools/create_tools.py index 578a1268d..0d2e0984c 100644 --- a/agents-api/agents_api/models/tools/create_tools.py +++ b/agents-api/agents_api/queries/tools/create_tools.py @@ -1,18 +1,18 @@ """This module contains functions for creating tools in the CozoDB database.""" +import sqlvalidator from typing import Any, TypeVar from uuid import UUID from beartype import beartype from fastapi import HTTPException -from pycozo.client import QueryException from pydantic import ValidationError from uuid_extensions import uuid7 from ...autogen.openapi_model import CreateToolRequest, Tool from ...metrics.counters import increase_counter from ..utils import ( - cozo_query, + pg_query, partialclass, rewrap_exceptions, verify_developer_id_query, @@ -24,14 +24,13 @@ T = TypeVar("T") -@rewrap_exceptions( - { - QueryException: partialclass(HTTPException, status_code=400), - ValidationError: partialclass(HTTPException, status_code=400), - TypeError: partialclass(HTTPException, status_code=400), - AssertionError: partialclass(HTTPException, status_code=400), - } -) +# @rewrap_exceptions( +# { +# ValidationError: partialclass(HTTPException, status_code=400), +# TypeError: partialclass(HTTPException, status_code=400), +# AssertionError: partialclass(HTTPException, status_code=400), +# } +# ) @wrap_in_class( Tool, transform=lambda d: { @@ -41,7 +40,7 @@ }, _kind="inserted", ) -@cozo_query +@pg_query @increase_counter("create_tools") @beartype def create_tools( diff --git a/agents-api/agents_api/models/tools/delete_tool.py b/agents-api/agents_api/queries/tools/delete_tool.py similarity index 100% rename from agents-api/agents_api/models/tools/delete_tool.py rename to agents-api/agents_api/queries/tools/delete_tool.py diff --git a/agents-api/agents_api/models/tools/get_tool.py b/agents-api/agents_api/queries/tools/get_tool.py similarity index 100% rename from agents-api/agents_api/models/tools/get_tool.py rename to agents-api/agents_api/queries/tools/get_tool.py diff --git a/agents-api/agents_api/models/tools/get_tool_args_from_metadata.py b/agents-api/agents_api/queries/tools/get_tool_args_from_metadata.py similarity index 100% rename from agents-api/agents_api/models/tools/get_tool_args_from_metadata.py rename to agents-api/agents_api/queries/tools/get_tool_args_from_metadata.py diff --git a/agents-api/agents_api/models/tools/list_tools.py b/agents-api/agents_api/queries/tools/list_tools.py similarity index 100% rename from agents-api/agents_api/models/tools/list_tools.py rename to agents-api/agents_api/queries/tools/list_tools.py diff --git a/agents-api/agents_api/models/tools/patch_tool.py b/agents-api/agents_api/queries/tools/patch_tool.py similarity index 100% rename from agents-api/agents_api/models/tools/patch_tool.py rename to agents-api/agents_api/queries/tools/patch_tool.py diff --git a/agents-api/agents_api/models/tools/update_tool.py b/agents-api/agents_api/queries/tools/update_tool.py similarity index 100% rename from agents-api/agents_api/models/tools/update_tool.py rename to agents-api/agents_api/queries/tools/update_tool.py From 8a44cdee8ad093f4fcde41445781c4e585a49893 Mon Sep 17 00:00:00 2001 From: whiterabbit1983 Date: Fri, 20 Dec 2024 09:38:56 +0000 Subject: [PATCH 08/16] refactor: Lint agents-api (CI) --- agents-api/agents_api/queries/tools/create_tools.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/agents-api/agents_api/queries/tools/create_tools.py b/agents-api/agents_api/queries/tools/create_tools.py index 0d2e0984c..a54fa6973 100644 --- a/agents-api/agents_api/queries/tools/create_tools.py +++ b/agents-api/agents_api/queries/tools/create_tools.py @@ -1,9 +1,9 @@ """This module contains functions for creating tools in the CozoDB database.""" -import sqlvalidator from typing import Any, TypeVar from uuid import UUID +import sqlvalidator from beartype import beartype from fastapi import HTTPException from pydantic import ValidationError @@ -12,8 +12,8 @@ from ...autogen.openapi_model import CreateToolRequest, Tool from ...metrics.counters import increase_counter from ..utils import ( - pg_query, partialclass, + pg_query, rewrap_exceptions, verify_developer_id_query, verify_developer_owns_resource_query, From 44122cad522f4fcbe00bc17d271ba9acfc373270 Mon Sep 17 00:00:00 2001 From: Dmitry Paramonov Date: Fri, 20 Dec 2024 14:10:53 +0300 Subject: [PATCH 09/16] feat: Add create tools query --- .../agents_api/queries/tools/create_tools.py | 103 +++++++----------- 1 file changed, 41 insertions(+), 62 deletions(-) diff --git a/agents-api/agents_api/queries/tools/create_tools.py b/agents-api/agents_api/queries/tools/create_tools.py index a54fa6973..d50e98e80 100644 --- a/agents-api/agents_api/queries/tools/create_tools.py +++ b/agents-api/agents_api/queries/tools/create_tools.py @@ -5,18 +5,14 @@ import sqlvalidator from beartype import beartype -from fastapi import HTTPException -from pydantic import ValidationError from uuid_extensions import uuid7 from ...autogen.openapi_model import CreateToolRequest, Tool +from ...exceptions import InvalidSQLQuery from ...metrics.counters import increase_counter from ..utils import ( - partialclass, pg_query, - rewrap_exceptions, - verify_developer_id_query, - verify_developer_owns_resource_query, + # rewrap_exceptions, wrap_in_class, ) @@ -24,6 +20,37 @@ T = TypeVar("T") +sql_query = sqlvalidator.parse( + """INSERT INTO tools +( + developer_id, + agent_id, + tool_id, + type, + name, + spec, + description +) +SELECT + $1, + $2, + $3, + $4, + $5, + $6, + $7 +WHERE NOT EXISTS ( + SELECT null FROM tools + WHERE (agent_id, name) = ($2, $5) +) +RETURNING * +""" +) + +if not sql_query.is_valid(): + raise InvalidSQLQuery("create_tools") + + # @rewrap_exceptions( # { # ValidationError: partialclass(HTTPException, status_code=400), @@ -48,8 +75,8 @@ def create_tools( developer_id: UUID, agent_id: UUID, data: list[CreateToolRequest], - ignore_existing: bool = False, -) -> tuple[list[str], dict]: + ignore_existing: bool = False, # TODO: what to do with this flag? +) -> tuple[list[str], list]: """ Constructs a datalog query for inserting tool records into the 'agent_functions' relation in the CozoDB. @@ -69,6 +96,7 @@ def create_tools( tools_data = [ [ + developer_id, str(agent_id), str(uuid7()), tool.type, @@ -79,57 +107,8 @@ def create_tools( for tool in data ] - ensure_tool_name_unique_query = """ - input[agent_id, tool_id, type, name, spec, description] <- $records - ?[tool_id] := - input[agent_id, _, type, name, _, _], - *tools{ - agent_id: to_uuid(agent_id), - tool_id, - type, - name, - spec, - description, - } - - :limit 1 - :assert none - """ - - # Datalog query for inserting new tool records into the 'tools' relation - create_query = """ - input[agent_id, tool_id, type, name, spec, description] <- $records - - # Do not add duplicate - ?[agent_id, tool_id, type, name, spec, description] := - input[agent_id, tool_id, type, name, spec, description], - not *tools{ - agent_id: to_uuid(agent_id), - type, - name, - } - - :insert tools { - agent_id, - tool_id, - type, - name, - spec, - description, - } - :returning - """ - - queries = [ - verify_developer_id_query(developer_id), - verify_developer_owns_resource_query(developer_id, "agents", agent_id=agent_id), - create_query, - ] - - if not ignore_existing: - queries.insert( - -1, - ensure_tool_name_unique_query, - ) - - return (queries, {"records": tools_data}) + return ( + sql_query.format(), + tools_data, + "fetchmany", + ) From b19a0010dd3276589ce829048700151cdbe402b4 Mon Sep 17 00:00:00 2001 From: Dmitry Paramonov Date: Fri, 20 Dec 2024 14:27:45 +0300 Subject: [PATCH 10/16] feat: Add delete tool query --- .../agents_api/queries/tools/delete_tool.py | 69 +++++++++---------- 1 file changed, 33 insertions(+), 36 deletions(-) diff --git a/agents-api/agents_api/queries/tools/delete_tool.py b/agents-api/agents_api/queries/tools/delete_tool.py index c79cdfd29..59f561cf1 100644 --- a/agents-api/agents_api/queries/tools/delete_tool.py +++ b/agents-api/agents_api/queries/tools/delete_tool.py @@ -1,19 +1,14 @@ from typing import Any, TypeVar from uuid import UUID +import sqlvalidator from beartype import beartype -from fastapi import HTTPException -from pycozo.client import QueryException -from pydantic import ValidationError from ...autogen.openapi_model import ResourceDeletedResponse from ...common.utils.datetime import utcnow +from ...exceptions import InvalidSQLQuery from ..utils import ( - cozo_query, - partialclass, - rewrap_exceptions, - verify_developer_id_query, - verify_developer_owns_resource_query, + pg_query, wrap_in_class, ) @@ -21,20 +16,34 @@ T = TypeVar("T") -@rewrap_exceptions( - { - QueryException: partialclass(HTTPException, status_code=400), - ValidationError: partialclass(HTTPException, status_code=400), - TypeError: partialclass(HTTPException, status_code=400), - } -) +sql_query = sqlvalidator.parse(""" +DELETE FROM + tools +WHERE + developer_id = $1 AND + agent_id = $2 AND + tool_id = $3 +RETURNING * +""") + +if not sql_query.is_valid(): + raise InvalidSQLQuery("delete_tool") + + +# @rewrap_exceptions( +# { +# QueryException: partialclass(HTTPException, status_code=400), +# ValidationError: partialclass(HTTPException, status_code=400), +# TypeError: partialclass(HTTPException, status_code=400), +# } +# ) @wrap_in_class( ResourceDeletedResponse, one=True, transform=lambda d: {"id": d["tool_id"], "deleted_at": utcnow(), "jobs": [], **d}, _kind="deleted", ) -@cozo_query +@pg_query @beartype def delete_tool( *, @@ -42,27 +51,15 @@ def delete_tool( agent_id: UUID, tool_id: UUID, ) -> tuple[list[str], dict]: + developer_id = str(developer_id) agent_id = str(agent_id) tool_id = str(tool_id) - delete_query = """ - # Delete function - ?[tool_id, agent_id] <- [[ - to_uuid($tool_id), - to_uuid($agent_id), - ]] - - :delete tools { - tool_id, + return ( + sql_query.format(), + [ + developer_id, agent_id, - } - :returning - """ - - queries = [ - verify_developer_id_query(developer_id), - verify_developer_owns_resource_query(developer_id, "agents", agent_id=agent_id), - delete_query, - ] - - return (queries, {"tool_id": tool_id, "agent_id": agent_id}) + tool_id, + ], + ) From e7d3079f380fa954c3e18c866bc120c8b16a9a50 Mon Sep 17 00:00:00 2001 From: Dmitry Paramonov Date: Fri, 20 Dec 2024 14:32:39 +0300 Subject: [PATCH 11/16] feat: Add get tool query --- .../agents_api/queries/tools/get_tool.py | 76 ++++++++----------- 1 file changed, 30 insertions(+), 46 deletions(-) diff --git a/agents-api/agents_api/queries/tools/get_tool.py b/agents-api/agents_api/queries/tools/get_tool.py index 465fd2efe..3662725b8 100644 --- a/agents-api/agents_api/queries/tools/get_tool.py +++ b/agents-api/agents_api/queries/tools/get_tool.py @@ -1,32 +1,39 @@ from typing import Any, TypeVar from uuid import UUID +import sqlvalidator from beartype import beartype -from fastapi import HTTPException -from pycozo.client import QueryException -from pydantic import ValidationError from ...autogen.openapi_model import Tool +from ...exceptions import InvalidSQLQuery from ..utils import ( - cozo_query, - partialclass, - rewrap_exceptions, - verify_developer_id_query, - verify_developer_owns_resource_query, + pg_query, wrap_in_class, ) ModelT = TypeVar("ModelT", bound=Any) T = TypeVar("T") +sql_query = sqlvalidator.parse(""" +SELECT * FROM tools +WHERE + developer_id = $1 AND + agent_id = $2 AND + tool_id = $3 +LIMIT 1 +""") -@rewrap_exceptions( - { - QueryException: partialclass(HTTPException, status_code=400), - ValidationError: partialclass(HTTPException, status_code=400), - TypeError: partialclass(HTTPException, status_code=400), - } -) +if not sql_query.is_valid(): + raise InvalidSQLQuery("get_tool") + + +# @rewrap_exceptions( +# { +# QueryException: partialclass(HTTPException, status_code=400), +# ValidationError: partialclass(HTTPException, status_code=400), +# TypeError: partialclass(HTTPException, status_code=400), +# } +# ) @wrap_in_class( Tool, transform=lambda d: { @@ -36,7 +43,7 @@ }, one=True, ) -@cozo_query +@pg_query @beartype def get_tool( *, @@ -44,38 +51,15 @@ def get_tool( agent_id: UUID, tool_id: UUID, ) -> tuple[list[str], dict]: + developer_id = str(developer_id) agent_id = str(agent_id) tool_id = str(tool_id) - get_query = """ - input[agent_id, tool_id] <- [[to_uuid($agent_id), to_uuid($tool_id)]] - - ?[ + return ( + sql_query.format(), + [ + developer_id, agent_id, tool_id, - type, - name, - spec, - updated_at, - created_at, - ] := input[agent_id, tool_id], - *tools { - agent_id, - tool_id, - name, - type, - spec, - updated_at, - created_at, - } - - :limit 1 - """ - - queries = [ - verify_developer_id_query(developer_id), - verify_developer_owns_resource_query(developer_id, "agents", agent_id=agent_id), - get_query, - ] - - return (queries, {"agent_id": agent_id, "tool_id": tool_id}) + ], + ) From 83f58aca92fc715cfbafc5f9f2f19f95cbf2da1e Mon Sep 17 00:00:00 2001 From: Dmitry Paramonov Date: Fri, 20 Dec 2024 14:45:28 +0300 Subject: [PATCH 12/16] feat: Add list tools query --- .../agents_api/queries/tools/list_tools.py | 92 ++++++++----------- 1 file changed, 37 insertions(+), 55 deletions(-) diff --git a/agents-api/agents_api/queries/tools/list_tools.py b/agents-api/agents_api/queries/tools/list_tools.py index 727bf8028..59fb1eff5 100644 --- a/agents-api/agents_api/queries/tools/list_tools.py +++ b/agents-api/agents_api/queries/tools/list_tools.py @@ -1,32 +1,43 @@ from typing import Any, Literal, TypeVar from uuid import UUID +import sqlvalidator from beartype import beartype -from fastapi import HTTPException -from pycozo.client import QueryException -from pydantic import ValidationError from ...autogen.openapi_model import Tool +from ...exceptions import InvalidSQLQuery from ..utils import ( - cozo_query, - partialclass, - rewrap_exceptions, - verify_developer_id_query, - verify_developer_owns_resource_query, + pg_query, wrap_in_class, ) ModelT = TypeVar("ModelT", bound=Any) T = TypeVar("T") +sql_query = sqlvalidator.parse(""" +SELECT * FROM tools +WHERE + developer_id = $1 AND + agent_id = $2 +ORDER BY + CASE WHEN $5 = 'created_at' AND $6 = 'desc' THEN s.created_at END DESC, + CASE WHEN $5 = 'created_at' AND $6 = 'asc' THEN s.created_at END ASC, + CASE WHEN $5 = 'updated_at' AND $6 = 'desc' THEN s.updated_at END DESC, + CASE WHEN $5 = 'updated_at' AND $6 = 'asc' THEN s.updated_at END ASC +LIMIT $3 OFFSET $4; +""") -@rewrap_exceptions( - { - QueryException: partialclass(HTTPException, status_code=400), - ValidationError: partialclass(HTTPException, status_code=400), - TypeError: partialclass(HTTPException, status_code=400), - } -) +if not sql_query.is_valid(): + raise InvalidSQLQuery("get_tool") + + +# @rewrap_exceptions( +# { +# QueryException: partialclass(HTTPException, status_code=400), +# ValidationError: partialclass(HTTPException, status_code=400), +# TypeError: partialclass(HTTPException, status_code=400), +# } +# ) @wrap_in_class( Tool, transform=lambda d: { @@ -38,7 +49,7 @@ **d, }, ) -@cozo_query +@pg_query @beartype def list_tools( *, @@ -49,46 +60,17 @@ def list_tools( sort_by: Literal["created_at", "updated_at"] = "created_at", direction: Literal["asc", "desc"] = "desc", ) -> tuple[list[str], dict]: + developer_id = str(developer_id) agent_id = str(agent_id) - sort = f"{'-' if direction == 'desc' else ''}{sort_by}" - - list_query = f""" - input[agent_id] <- [[to_uuid($agent_id)]] - - ?[ - agent_id, - id, - name, - type, - spec, - description, - updated_at, - created_at, - ] := input[agent_id], - *tools {{ - agent_id, - tool_id: id, - name, - type, - spec, - description, - updated_at, - created_at, - }} - - :limit $limit - :offset $offset - :sort {sort} - """ - - queries = [ - verify_developer_id_query(developer_id), - verify_developer_owns_resource_query(developer_id, "agents", agent_id=agent_id), - list_query, - ] - return ( - queries, - {"agent_id": agent_id, "limit": limit, "offset": offset}, + sql_query.format(), + [ + developer_id, + agent_id, + limit, + offset, + sort_by, + direction, + ], ) From 59b24ac9bf2031daff49c42ecce5e03c880b1ee9 Mon Sep 17 00:00:00 2001 From: Dmitry Paramonov Date: Fri, 20 Dec 2024 15:21:09 +0300 Subject: [PATCH 13/16] feat: Add patch tool query --- .../agents_api/queries/tools/patch_tool.py | 94 +++++++++---------- 1 file changed, 43 insertions(+), 51 deletions(-) diff --git a/agents-api/agents_api/queries/tools/patch_tool.py b/agents-api/agents_api/queries/tools/patch_tool.py index bc49b8121..aa663dec0 100644 --- a/agents-api/agents_api/queries/tools/patch_tool.py +++ b/agents-api/agents_api/queries/tools/patch_tool.py @@ -1,20 +1,14 @@ from typing import Any, TypeVar from uuid import UUID +import sqlvalidator from beartype import beartype -from fastapi import HTTPException -from pycozo.client import QueryException -from pydantic import ValidationError from ...autogen.openapi_model import PatchToolRequest, ResourceUpdatedResponse -from ...common.utils.cozo import cozo_process_mutate_data +from ...exceptions import InvalidSQLQuery from ...metrics.counters import increase_counter from ..utils import ( - cozo_query, - partialclass, - rewrap_exceptions, - verify_developer_id_query, - verify_developer_owns_resource_query, + pg_query, wrap_in_class, ) @@ -22,25 +16,46 @@ T = TypeVar("T") -@rewrap_exceptions( - { - QueryException: partialclass(HTTPException, status_code=400), - ValidationError: partialclass(HTTPException, status_code=400), - TypeError: partialclass(HTTPException, status_code=400), - } +sql_query = sqlvalidator.parse(""" +WITH updated_tools AS ( + UPDATE tools + SET + type = COALESCE($4, type), + name = COALESCE($5, name), + description = COALESCE($6, description), + spec = COALESCE($7, spec) + WHERE + developer_id = $1 AND + agent_id = $2 AND + tool_id = $3 + RETURNING * ) +SELECT * FROM updated_tools; +""") + +if not sql_query.is_valid(): + raise InvalidSQLQuery("patch_tool") + + +# @rewrap_exceptions( +# { +# QueryException: partialclass(HTTPException, status_code=400), +# ValidationError: partialclass(HTTPException, status_code=400), +# TypeError: partialclass(HTTPException, status_code=400), +# } +# ) @wrap_in_class( ResourceUpdatedResponse, one=True, transform=lambda d: {"id": d["tool_id"], "jobs": [], **d}, _kind="inserted", ) -@cozo_query +@pg_query @increase_counter("patch_tool") @beartype def patch_tool( *, developer_id: UUID, agent_id: UUID, tool_id: UUID, data: PatchToolRequest -) -> tuple[list[str], dict]: +) -> tuple[list[str], list]: """ Execute the datalog query and return the results as a DataFrame Updates the tool information for a given agent and tool ID in the 'cozodb' database. @@ -54,6 +69,7 @@ def patch_tool( ResourceUpdatedResponse: The updated tool data. """ + developer_id = str(developer_id) agent_id = str(agent_id) tool_id = str(tool_id) @@ -78,39 +94,15 @@ def patch_tool( if tool_spec: del patch_data[tool_type] - tool_cols, tool_vals = cozo_process_mutate_data( - { - **patch_data, - "agent_id": agent_id, - "tool_id": tool_id, - } - ) - - # Construct the datalog query for updating the tool information - patch_query = f""" - input[{tool_cols}] <- $input - - ?[{tool_cols}, spec, updated_at] := - *tools {{ - agent_id: to_uuid($agent_id), - tool_id: to_uuid($tool_id), - spec: old_spec, - }}, - input[{tool_cols}], - spec = concat(old_spec, $spec), - updated_at = now() - - :update tools {{ {tool_cols}, spec, updated_at }} - :returning - """ - - queries = [ - verify_developer_id_query(developer_id), - verify_developer_owns_resource_query(developer_id, "agents", agent_id=agent_id), - patch_query, - ] - return ( - queries, - dict(input=tool_vals, spec=tool_spec, agent_id=agent_id, tool_id=tool_id), + sql_query.format(), + [ + developer_id, + agent_id, + tool_id, + tool_type, + data.name, + data.description, + tool_spec, + ], ) From 32dbbbaac376757ddc535d40eef64d3d64259c3f Mon Sep 17 00:00:00 2001 From: Dmitry Paramonov Date: Fri, 20 Dec 2024 15:21:21 +0300 Subject: [PATCH 14/16] fix: Fix return types --- agents-api/agents_api/queries/tools/delete_tool.py | 2 +- agents-api/agents_api/queries/tools/get_tool.py | 2 +- agents-api/agents_api/queries/tools/list_tools.py | 4 ++-- 3 files changed, 4 insertions(+), 4 deletions(-) diff --git a/agents-api/agents_api/queries/tools/delete_tool.py b/agents-api/agents_api/queries/tools/delete_tool.py index 59f561cf1..17535e1e4 100644 --- a/agents-api/agents_api/queries/tools/delete_tool.py +++ b/agents-api/agents_api/queries/tools/delete_tool.py @@ -50,7 +50,7 @@ def delete_tool( developer_id: UUID, agent_id: UUID, tool_id: UUID, -) -> tuple[list[str], dict]: +) -> tuple[list[str], list]: developer_id = str(developer_id) agent_id = str(agent_id) tool_id = str(tool_id) diff --git a/agents-api/agents_api/queries/tools/get_tool.py b/agents-api/agents_api/queries/tools/get_tool.py index 3662725b8..af63be0c9 100644 --- a/agents-api/agents_api/queries/tools/get_tool.py +++ b/agents-api/agents_api/queries/tools/get_tool.py @@ -50,7 +50,7 @@ def get_tool( developer_id: UUID, agent_id: UUID, tool_id: UUID, -) -> tuple[list[str], dict]: +) -> tuple[list[str], list]: developer_id = str(developer_id) agent_id = str(agent_id) tool_id = str(tool_id) diff --git a/agents-api/agents_api/queries/tools/list_tools.py b/agents-api/agents_api/queries/tools/list_tools.py index 59fb1eff5..3dac84875 100644 --- a/agents-api/agents_api/queries/tools/list_tools.py +++ b/agents-api/agents_api/queries/tools/list_tools.py @@ -28,7 +28,7 @@ """) if not sql_query.is_valid(): - raise InvalidSQLQuery("get_tool") + raise InvalidSQLQuery("list_tools") # @rewrap_exceptions( @@ -59,7 +59,7 @@ def list_tools( offset: int = 0, sort_by: Literal["created_at", "updated_at"] = "created_at", direction: Literal["asc", "desc"] = "desc", -) -> tuple[list[str], dict]: +) -> tuple[list[str], list]: developer_id = str(developer_id) agent_id = str(agent_id) From 281e1a8f44c79cfd7081108a213b0a580446db26 Mon Sep 17 00:00:00 2001 From: Dmitry Paramonov Date: Fri, 20 Dec 2024 15:31:48 +0300 Subject: [PATCH 15/16] feat: Add update tool query --- .../agents_api/queries/tools/update_tool.py | 93 ++++++++----------- 1 file changed, 41 insertions(+), 52 deletions(-) diff --git a/agents-api/agents_api/queries/tools/update_tool.py b/agents-api/agents_api/queries/tools/update_tool.py index ef700a5f6..356e28bbf 100644 --- a/agents-api/agents_api/queries/tools/update_tool.py +++ b/agents-api/agents_api/queries/tools/update_tool.py @@ -1,44 +1,55 @@ from typing import Any, TypeVar from uuid import UUID +import sqlvalidator from beartype import beartype -from fastapi import HTTPException -from pycozo.client import QueryException -from pydantic import ValidationError from ...autogen.openapi_model import ( ResourceUpdatedResponse, UpdateToolRequest, ) -from ...common.utils.cozo import cozo_process_mutate_data +from ...exceptions import InvalidSQLQuery from ...metrics.counters import increase_counter from ..utils import ( - cozo_query, - partialclass, - rewrap_exceptions, - verify_developer_id_query, - verify_developer_owns_resource_query, + pg_query, wrap_in_class, ) ModelT = TypeVar("ModelT", bound=Any) T = TypeVar("T") - -@rewrap_exceptions( - { - QueryException: partialclass(HTTPException, status_code=400), - ValidationError: partialclass(HTTPException, status_code=400), - TypeError: partialclass(HTTPException, status_code=400), - } -) +sql_query = sqlvalidator.parse(""" +UPDATE tools +SET + type = $4, + name = $5, + description = $6, + spec = $7 +WHERE + developer_id = $1 AND + agent_id = $2 AND + tool_id = $3 +RETURNING *; +""") + +if not sql_query.is_valid(): + raise InvalidSQLQuery("update_tool") + + +# @rewrap_exceptions( +# { +# QueryException: partialclass(HTTPException, status_code=400), +# ValidationError: partialclass(HTTPException, status_code=400), +# TypeError: partialclass(HTTPException, status_code=400), +# } +# ) @wrap_in_class( ResourceUpdatedResponse, one=True, transform=lambda d: {"id": d["tool_id"], "jobs": [], **d}, _kind="inserted", ) -@cozo_query +@pg_query @increase_counter("update_tool") @beartype def update_tool( @@ -48,7 +59,8 @@ def update_tool( tool_id: UUID, data: UpdateToolRequest, **kwargs, -) -> tuple[list[str], dict]: +) -> tuple[list[str], list]: + developer_id = str(developer_id) agent_id = str(agent_id) tool_id = str(tool_id) @@ -72,38 +84,15 @@ def update_tool( update_data["spec"] = tool_spec del update_data[tool_type] - tool_cols, tool_vals = cozo_process_mutate_data( - { - **update_data, - "agent_id": agent_id, - "tool_id": tool_id, - } - ) - - # Construct the datalog query for updating the tool information - patch_query = f""" - input[{tool_cols}] <- $input - - ?[{tool_cols}, created_at, updated_at] := - *tools {{ - agent_id: to_uuid($agent_id), - tool_id: to_uuid($tool_id), - created_at - }}, - input[{tool_cols}], - updated_at = now() - - :put tools {{ {tool_cols}, created_at, updated_at }} - :returning - """ - - queries = [ - verify_developer_id_query(developer_id), - verify_developer_owns_resource_query(developer_id, "agents", agent_id=agent_id), - patch_query, - ] - return ( - queries, - dict(input=tool_vals, spec=tool_spec, agent_id=agent_id, tool_id=tool_id), + sql_query.format(), + [ + developer_id, + agent_id, + tool_id, + tool_type, + data.name, + data.description, + tool_spec, + ], ) From 6a52a4022ca8a52a70701f0f3878595759380f05 Mon Sep 17 00:00:00 2001 From: Dmitry Paramonov Date: Sat, 21 Dec 2024 21:05:17 +0300 Subject: [PATCH 16/16] WIP --- .../tools/get_tool_args_from_metadata.py | 33 +++++-------------- 1 file changed, 9 insertions(+), 24 deletions(-) diff --git a/agents-api/agents_api/queries/tools/get_tool_args_from_metadata.py b/agents-api/agents_api/queries/tools/get_tool_args_from_metadata.py index 2cdb92cb9..a8a9dba1a 100644 --- a/agents-api/agents_api/queries/tools/get_tool_args_from_metadata.py +++ b/agents-api/agents_api/queries/tools/get_tool_args_from_metadata.py @@ -2,16 +2,9 @@ from uuid import UUID from beartype import beartype -from fastapi import HTTPException -from pycozo.client import QueryException -from pydantic import ValidationError from ..utils import ( - cozo_query, - partialclass, - rewrap_exceptions, - verify_developer_id_query, - verify_developer_owns_resource_query, + pg_query, wrap_in_class, ) @@ -51,10 +44,6 @@ def tool_args_for_task( """ queries = [ - verify_developer_id_query(developer_id), - verify_developer_owns_resource_query( - developer_id, "tasks", task_id=task_id, parents=[("agents", "agent_id")] - ), get_query, ] @@ -95,25 +84,21 @@ def tool_args_for_session( """ queries = [ - verify_developer_id_query(developer_id), - verify_developer_owns_resource_query( - developer_id, "sessions", session_id=session_id - ), get_query, ] return (queries, {"agent_id": agent_id, "session_id": session_id}) -@rewrap_exceptions( - { - QueryException: partialclass(HTTPException, status_code=400), - ValidationError: partialclass(HTTPException, status_code=400), - TypeError: partialclass(HTTPException, status_code=400), - } -) +# @rewrap_exceptions( +# { +# QueryException: partialclass(HTTPException, status_code=400), +# ValidationError: partialclass(HTTPException, status_code=400), +# TypeError: partialclass(HTTPException, status_code=400), +# } +# ) @wrap_in_class(dict, transform=lambda x: x["values"], one=True) -@cozo_query +@pg_query @beartype def get_tool_args_from_metadata( *,