From 09d8dadf50f44e4db93df5de89d2832be8dd4aad Mon Sep 17 00:00:00 2001 From: Diwank Singh Tomer Date: Mon, 7 Oct 2024 01:15:28 -0400 Subject: [PATCH 1/2] fix: Add developer_id constraints in all queries when possible Signed-off-by: Diwank Singh Tomer --- agents-api/agents_api/env.py | 2 ++ .../models/agent/create_or_update_agent.py | 2 ++ agents-api/agents_api/models/agent/get_agent.py | 1 + .../models/session/prepare_session_data.py | 14 +++++++++++--- .../models/user/create_or_update_user.py | 2 ++ agents-api/agents_api/models/user/patch_user.py | 2 ++ agents-api/agents_api/models/utils.py | 7 +++++++ 7 files changed, 27 insertions(+), 3 deletions(-) diff --git a/agents-api/agents_api/env.py b/agents-api/agents_api/env.py index 42b2ce8d2..ef4755bd4 100644 --- a/agents-api/agents_api/env.py +++ b/agents-api/agents_api/env.py @@ -40,6 +40,8 @@ summarization_model_name: str = env.str( "SUMMARIZATION_MODEL_NAME", default="gpt-4-turbo" ) +do_verify_developer: bool = env.bool("DO_VERIFY_DEVELOPER", default=True) +do_verify_developer_owns_resource: bool = env.bool("DO_VERIFY_DEVELOPER_OWNS_RESOURCE", default=True) # Auth diff --git a/agents-api/agents_api/models/agent/create_or_update_agent.py b/agents-api/agents_api/models/agent/create_or_update_agent.py index 3902c1bb5..8a3975183 100644 --- a/agents-api/agents_api/models/agent/create_or_update_agent.py +++ b/agents-api/agents_api/models/agent/create_or_update_agent.py @@ -136,6 +136,7 @@ def create_or_update_agent( input[_agent_id, developer_id, model, name, about, metadata, instructions, updated_at], *agents{ agent_id, + developer_id, created_at, }, agent_id = to_uuid(_agent_id), @@ -144,6 +145,7 @@ def create_or_update_agent( input[_agent_id, developer_id, model, name, about, metadata, instructions, updated_at], not *agents{ agent_id, + developer_id, }, created_at = now(), agent_id = to_uuid(_agent_id), diff --git a/agents-api/agents_api/models/agent/get_agent.py b/agents-api/agents_api/models/agent/get_agent.py index bdae85fcb..008e39454 100644 --- a/agents-api/agents_api/models/agent/get_agent.py +++ b/agents-api/agents_api/models/agent/get_agent.py @@ -71,6 +71,7 @@ def get_agent(*, developer_id: UUID, agent_id: UUID) -> tuple[list[str], dict]: instructions, ] := input[id], *agents { + developer_id: to_uuid($developer_id), agent_id: id, model, name, diff --git a/agents-api/agents_api/models/session/prepare_session_data.py b/agents-api/agents_api/models/session/prepare_session_data.py index bbbd9c4cd..83ee0c219 100644 --- a/agents-api/agents_api/models/session/prepare_session_data.py +++ b/agents-api/agents_api/models/session/prepare_session_data.py @@ -56,12 +56,13 @@ def prepare_session_data( # This query retrieves session information by using `input` to pass parameters, get_query = """ - input[session_id] <- [[ + input[session_id, developer_id] <- [[ to_uuid($session_id), + to_uuid($developer_id), ]] participants[collect(participant_id), participant_type] := - input[session_id], + input[session_id, developer_id], *session_lookup{ session_id, participant_id, @@ -102,9 +103,11 @@ def prepare_session_data( } agent_data[collect(record)] := + input[session_id, developer_id], agents[agent_ids], agent_id in agent_ids, *agents{ + developer_id, agent_id, model, name, @@ -129,9 +132,11 @@ def prepare_session_data( # Version where we don't have default settings agent_data[collect(record)] := + input[session_id, developer_id], agents[agent_ids], agent_id in agent_ids, *agents{ + developer_id, agent_id, model, name, @@ -155,9 +160,11 @@ def prepare_session_data( } user_data[collect(record)] := + input[session_id, developer_id], users[user_ids], user_id in user_ids, *users{ + developer_id, user_id, name, about, @@ -175,8 +182,9 @@ def prepare_session_data( } session_data[record] := - input[session_id], + input[session_id, developer_id], *sessions{ + developer_id, session_id, situation, summary, diff --git a/agents-api/agents_api/models/user/create_or_update_user.py b/agents-api/agents_api/models/user/create_or_update_user.py index 97db913c5..d295d1d8a 100644 --- a/agents-api/agents_api/models/user/create_or_update_user.py +++ b/agents-api/agents_api/models/user/create_or_update_user.py @@ -70,6 +70,7 @@ def create_or_update_user( ?[user_id, developer_id, name, about, metadata, created_at, updated_at] := input[_user_id, developer_id, name, about, metadata, updated_at], *users{ + developer_id, user_id, created_at, }, @@ -78,6 +79,7 @@ def create_or_update_user( ?[user_id, developer_id, name, about, metadata, created_at, updated_at] := input[_user_id, developer_id, name, about, metadata, updated_at], not *users{ + developer_id, user_id, }, created_at = now(), user_id = to_uuid(_user_id), diff --git a/agents-api/agents_api/models/user/patch_user.py b/agents-api/agents_api/models/user/patch_user.py index 4498c6ded..265241d47 100644 --- a/agents-api/agents_api/models/user/patch_user.py +++ b/agents-api/agents_api/models/user/patch_user.py @@ -78,6 +78,7 @@ def patch_user( ?[{user_update_cols}, metadata] := input[{user_update_cols}], *users {{ + developer_id: to_uuid($developer_id), user_id: to_uuid($user_id), metadata: md, }}, @@ -101,5 +102,6 @@ def patch_user( "user_update_vals": user_update_vals, "metadata": metadata, "user_id": str(user_id), + "developer_id": str(developer_id), }, ) diff --git a/agents-api/agents_api/models/utils.py b/agents-api/agents_api/models/utils.py index c163642c0..f63646e1c 100644 --- a/agents-api/agents_api/models/utils.py +++ b/agents-api/agents_api/models/utils.py @@ -8,6 +8,7 @@ from pydantic import BaseModel from ..common.utils.cozo import uuid_int_list_to_uuid4 +from ..env import do_verify_developer, do_verify_developer_owns_resource P = ParamSpec("P") T = TypeVar("T") @@ -117,6 +118,9 @@ def mark_session_updated_query(developer_id: UUID | str, session_id: UUID | str) def verify_developer_id_query(developer_id: UUID | str) -> str: + if not do_verify_developer: + return "?[exists] := exists = true" + return f""" matched[count(developer_id)] := *developers{{ @@ -138,6 +142,9 @@ def verify_developer_owns_resource_query( parents: list[tuple[str, str]] | None = None, **resource_id, ) -> str: + if not do_verify_developer_owns_resource: + return "?[exists] := exists = true" + parents = parents or [] resource_id_key, resource_id_value = next(iter(resource_id.items())) From 2c8ebba55a75f8b9631bc2b072b0a54098bbe93c Mon Sep 17 00:00:00 2001 From: creatorrr Date: Mon, 7 Oct 2024 05:16:11 +0000 Subject: [PATCH 2/2] refactor: Lint agents-api (CI) --- agents-api/agents_api/env.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/agents-api/agents_api/env.py b/agents-api/agents_api/env.py index ef4755bd4..2b016fd40 100644 --- a/agents-api/agents_api/env.py +++ b/agents-api/agents_api/env.py @@ -41,7 +41,9 @@ "SUMMARIZATION_MODEL_NAME", default="gpt-4-turbo" ) do_verify_developer: bool = env.bool("DO_VERIFY_DEVELOPER", default=True) -do_verify_developer_owns_resource: bool = env.bool("DO_VERIFY_DEVELOPER_OWNS_RESOURCE", default=True) +do_verify_developer_owns_resource: bool = env.bool( + "DO_VERIFY_DEVELOPER_OWNS_RESOURCE", default=True +) # Auth