Skip to content

Commit

Permalink
feat: Add get tool query
Browse files Browse the repository at this point in the history
  • Loading branch information
whiterabbit1983 committed Dec 20, 2024
1 parent b19a001 commit 3b54467
Showing 1 changed file with 29 additions and 46 deletions.
75 changes: 29 additions & 46 deletions agents-api/agents_api/queries/tools/get_tool.py
Original file line number Diff line number Diff line change
@@ -1,32 +1,38 @@
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
""")

@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: {
Expand All @@ -36,46 +42,23 @@
},
one=True,
)
@cozo_query
@pg_query
@beartype
def get_tool(
*,
developer_id: UUID,
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})
],
)

0 comments on commit 3b54467

Please sign in to comment.