fix: Add developer_id constraints in all queries when possible

Signed-off-by: Diwank Singh Tomer <[email protected]>

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

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),
 

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,

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,

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),

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),
         },
     )

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()))