fix(agents-api): Fixed tests

Signed-off-by: Diwank Tomer <diwank@julep.ai>

agents-api/agents_api/models/agent/create_agent.py

@@ -23,6 +23,11 @@
 
 @rewrap_exceptions(
     {
+        lambda e: isinstance(e, QueryException)
+        and "asserted to return some results, but returned none"
+        in str(e): lambda *_: HTTPException(
+            detail="developer not found", status_code=403
+        ),
         QueryException: partialclass(HTTPException, status_code=400),
         ValidationError: partialclass(HTTPException, status_code=400),
         TypeError: partialclass(HTTPException, status_code=400),
@@ -62,7 +67,7 @@ def create_agent(
     )
     data.default_settings = data.default_settings or {}
 
-    agent_data = data.model_dump()
+    agent_data = data.model_dump(exclude_unset=True)
     default_settings = agent_data.pop("default_settings")
 
     settings_cols, settings_vals = cozo_process_mutate_data(

agents-api/agents_api/models/agent/delete_agent.py

@@ -23,6 +23,11 @@
 
 @rewrap_exceptions(
     {
+        lambda e: isinstance(e, QueryException)
+        and "asserted to return some results, but returned none"
+        in str(e): lambda *_: HTTPException(
+            detail="developer not found or doesnt own resource", status_code=404
+        ),
         QueryException: partialclass(HTTPException, status_code=400),
         ValidationError: partialclass(HTTPException, status_code=400),
         TypeError: partialclass(HTTPException, status_code=400),
@@ -57,14 +62,18 @@ def delete_agent(*, developer_id: UUID, agent_id: UUID) -> tuple[list[str], dict
         verify_developer_owns_resource_query(developer_id, "agents", agent_id=agent_id),
         """
         # Delete docs
-        ?[agent_id, doc_id] :=
-            *agent_docs{
-                agent_id,
+        ?[owner_id, owner_type, doc_id] :=
+            *docs{
+                owner_type,
+                owner_id,
                 doc_id,
-            }, agent_id = to_uuid($agent_id)
+            },
+            owner_id = to_uuid($agent_id),
+            owner_type = "agent"
 
-        :delete agent_docs {
-            agent_id,
+        :delete docs {
+            owner_type,
+            owner_id,
             doc_id
         }
         :returning

agents-api/agents_api/models/agent/get_agent.py

@@ -18,6 +18,11 @@
 
 @rewrap_exceptions(
     {
+        lambda e: isinstance(e, QueryException)
+        and "asserted to return some results, but returned none"
+        in str(e): lambda *_: HTTPException(
+            detail="developer not found or doesnt own resource", status_code=404
+        ),
         QueryException: partialclass(HTTPException, status_code=400),
         ValidationError: partialclass(HTTPException, status_code=400),
         TypeError: partialclass(HTTPException, status_code=400),

agents-api/agents_api/models/docs/delete_doc.py

@@ -12,6 +12,7 @@
     partialclass,
     rewrap_exceptions,
     verify_developer_id_query,
+    verify_developer_owns_resource_query,
     wrap_in_class,
 )
 
@@ -37,6 +38,8 @@
 def delete_doc(
     *,
     developer_id: UUID,
+    owner_id: UUID,
+    owner_type: str,
     doc_id: UUID,
 ) -> tuple[list[str], dict]:
     """Constructs and returns a datalog query for deleting documents and associated information snippets.
@@ -52,6 +55,7 @@ def delete_doc(
     """
     # Convert UUID parameters to string format for use in the datalog query
     doc_id = str(doc_id)
+    owner_id = str(owner_id)
 
     # The following query is divided into two main parts:
     # 1. Deleting information snippets associated with the document
@@ -75,16 +79,19 @@ def delete_doc(
 
     delete_doc_query = """
         # Delete the docs
-        ?[doc_id] <- [[ to_uuid($doc_id) ]]
+        ?[doc_id, owner_type, owner_id] <- [[ to_uuid($doc_id), $owner_type, to_uuid($owner_id) ]]
 
-        :delete docs { doc_id }
+        :delete docs { doc_id, owner_type, owner_id }
         :returning
     """
 
     queries = [
         verify_developer_id_query(developer_id),
+        verify_developer_owns_resource_query(
+            developer_id, f"{owner_type}s", **{f"{owner_type}_id": owner_id}
+        ),
         delete_snippets_query,
         delete_doc_query,
     ]
 
-    return (queries, {"doc_id": doc_id})
+    return (queries, {"doc_id": doc_id, "owner_type": owner_type, "owner_id": owner_id})

agents-api/agents_api/models/entry/create_entries.py

@@ -1,3 +1,4 @@
+import json
 from uuid import UUID, uuid4
 
 from beartype import beartype
@@ -44,14 +45,18 @@ def create_entries(
     developer_id = str(developer_id)
     session_id = str(session_id)
 
-    data_dicts = [item.model_dump() for item in data]
+    data_dicts = [item.model_dump(exclude_unset=True) for item in data]
 
     for item in data_dicts:
         item["content"] = content_to_json(item["content"])
         item["session_id"] = session_id
         item["entry_id"] = item.pop("id", None) or str(uuid4())
         item["created_at"] = (item.get("created_at") or utcnow()).timestamp()
 
+        if not item.get("token_count"):
+            item["token_count"] = len(json.dumps(item)) // 3.5
+            item["tokenizer"] = "character_count"
+
     cols, rows = cozo_process_mutate_data(data_dicts)
 
     # Construct a datalog query to insert the processed entries into the 'cozodb' database.

agents-api/agents_api/models/execution/create_execution.py

@@ -40,7 +40,7 @@ def create_execution(
     task_id: UUID,
     execution_id: UUID | None = None,
     data: Annotated[CreateExecutionRequest | dict, dict_like(CreateExecutionRequest)],
-    workflow_hande: WorkflowHandle,
+    workflow_handle: WorkflowHandle,
 ) -> tuple[list[str], dict]:
     execution_id = execution_id or uuid4()
 
@@ -58,10 +58,10 @@ def create_execution(
     temporal_columns, temporal_values = cozo_process_mutate_data(
         {
             "execution_id": execution_id,
-            "id": workflow_hande.id,
-            "run_id": workflow_hande.run_id,
-            "first_execution_run_id": workflow_hande.first_execution_run_id,
-            "result_run_id": workflow_hande.result_run_id,
+            "id": workflow_handle.id,
+            "run_id": workflow_handle.run_id,
+            "first_execution_run_id": workflow_handle.first_execution_run_id,
+            "result_run_id": workflow_handle.result_run_id,
         }
     )
 

agents-api/agents_api/models/user/delete_user.py

@@ -58,41 +58,19 @@ def delete_user(*, developer_id: UUID, user_id: UUID) -> tuple[list[str], dict]:
         """
         # Delete docs
         ?[user_id, doc_id] :=
-            *user_docs{
-                user_id,
+            *docs{
+                owner_id: user_id,
+                owner_type: "user",
                 doc_id,
             }, user_id = to_uuid($user_id)
 
-        :delete user_docs {
+        :delete docs {
             user_id,
             doc_id
         }
         :returning
         """,
         """
-        # Delete tools
-        ?[user_id, tool_id] :=
-            *tools{
-                user_id,
-                tool_id,
-            }, user_id = to_uuid($user_id)
-
-        :delete tools {
-            user_id,
-            tool_id
-        }
-        :returning
-        """,
-        """
-        # Delete default user settings
-        ?[user_id] <- [[$user_id]]
-
-        :delete user_default_settings {
-            user_id
-        }
-        :returning
-        """,
-        """
         # Delete the user
         ?[user_id, developer_id] <- [[$user_id, $developer_id]]
 

agents-api/agents_api/models/user/update_user.py

@@ -32,7 +32,7 @@
 @cozo_query
 @beartype
 def update_user(
-    *, developer_id: UUID, user_id: UUID, update_user: UpdateUserRequest
+    *, developer_id: UUID, user_id: UUID, data: UpdateUserRequest
 ) -> tuple[list[str], dict]:
     """Updates user information in the 'cozodb' database.
 
@@ -47,7 +47,7 @@ def update_user(
     """
     user_id = str(user_id)
     developer_id = str(developer_id)
-    update_data = update_user.model_dump()
+    update_data = data.model_dump()
 
     # Prepares the update data by filtering out None values and adding user_id and developer_id.
     user_update_cols, user_update_vals = cozo_process_mutate_data(

agents-api/agents_api/models/utils.py

@@ -188,7 +188,7 @@ def wrapper(*args: P.args, **kwargs: P.kwargs):
             transform = transform or (lambda x: x)
 
             if one:
-                assert len(data) >= 1
+                assert len(data) >= 1, "Expected one result, got none"
                 return cls(**transform(data[0]))
 
             return [cls(**item) for item in map(transform, data)]
@@ -232,7 +232,9 @@ def wrapper(*args: P.args, **kwargs: P.kwargs) -> T:
                             else transform(error)
                         )
 
-                        raise transform(new_error) from error
+                        setattr(new_error, "__cause__", error)
+
+                        raise new_error from error
 
                 raise