f/patch methods (#22)

* feat: PATCH methods Signed-off-by: Diwank Singh Tomer <[email protected]> * fix: Use old instructions format * chore: Apply formatting --------- Signed-off-by: Diwank Singh Tomer <[email protected]> Co-authored-by: Dmitry Paramonov <[email protected]>
julep-ai · Apr 10, 2024 · 14d982d · 14d982d
1 parent 6a398b2
commit 14d982d
Show file tree

Hide file tree

Showing 36 changed files with 3,747 additions and 2,384 deletions.
diff --git a/agents-api/agents_api/autogen/openapi_model.py b/agents-api/agents_api/autogen/openapi_model.py
@@ -1,6 +1,6 @@
 # generated by datamodel-codegen:
-#   filename:  openapi.yaml
-#   timestamp: 2024-04-05T11:47:21+00:00
+#   filename:  openapi.json
+#   timestamp: 2024-04-10T05:34:23+00:00
 
 from __future__ import annotations
 
@@ -150,7 +150,7 @@ class UpdateSessionRequest(BaseModel):
     A valid request payload for updating a session
     """
 
-    situation: str | None = None
+    situation: str
     """
     Updated situation for this session
     """
@@ -165,11 +165,11 @@ class UpdateUserRequest(BaseModel):
     A valid request payload for updating a user
     """
 
-    about: str | None = None
+    about: str
     """
     About the user
     """
-    name: str | None = None
+    name: str
     """
     Name of the user
     """
@@ -514,7 +514,7 @@ class ChatSettings(BaseModel):
 
     """
     seed: Annotated[
-        int | None, Field(None, ge=-9223372036854775808, le=9223372036854775808)
+        int | None, Field(None, ge=-9223372036854776000, le=9223372036854776000)
     ]
     """
     This feature is in Beta.
@@ -945,15 +945,15 @@ class UpdateAgentRequest(BaseModel):
     A valid request payload for updating an agent
     """
 
-    about: str | None = None
+    about: str
     """
     About the agent
     """
     instructions: List[Instruction] | None = None
     """
     List of instructions for the agent
     """
-    name: str | None = None
+    name: str
     """
     Name of the agent
     """

diff --git a/agents-api/agents_api/models/agent/patch_agent.py b/agents-api/agents_api/models/agent/patch_agent.py
@@ -0,0 +1,77 @@
+from ...common.utils import json
+from uuid import UUID
+
+from ...common.utils.cozo import cozo_process_mutate_data
+from ...common.utils.datetime import utcnow
+from ...autogen.openapi_model import Instruction
+from ..instructions.create_instructions import create_instructions_query
+from ..instructions.delete_instructions import (
+    delete_instructions_by_agent_query,
+)
+
+
+def patch_agent_query(
+    agent_id: UUID,
+    developer_id: UUID,
+    default_settings: dict = {},
+    **update_data,
+) -> str:
+    instructions: list[Instruction] | None = update_data.pop("instructions")
+    del_instructions = delete_instructions_by_agent_query(agent_id=agent_id)
+    create_instructions = create_instructions_query(
+        agent_id=agent_id, instructions=instructions
+    )
+    # Agent update query
+    agent_update_cols, agent_update_vals = cozo_process_mutate_data(
+        {
+            **{k: v for k, v in update_data.items() if v is not None},
+            "agent_id": agent_id,
+            "developer_id": developer_id,
+            "updated_at": utcnow().timestamp(),
+        }
+    )
+
+    agent_update_query = f"""
+    {{
+        # update the agent
+        ?[{agent_update_cols}] <- {json.dumps(agent_update_vals)}
+
+        :update agents {{
+            {agent_update_cols}
+        }}
+        :returning
+    }}
+    """
+
+    # Settings update query
+    settings_cols, settings_vals = cozo_process_mutate_data(
+        {
+            **default_settings,
+            "agent_id": agent_id,
+        }
+    )
+
+    settings_update_query = f"""
+    {{
+        # update the agent settings
+        ?[{settings_cols}] <- {json.dumps(settings_vals)}
+
+        :update agent_default_settings {{
+            {settings_cols}
+        }}
+    }}
+    """
+
+    # Combine the queries
+    queries = [agent_update_query]
+
+    if len(default_settings) != 0:
+        queries.insert(0, settings_update_query)
+
+    if instructions:
+        queries.insert(0, create_instructions)
+        queries.insert(0, del_instructions)
+
+    combined_query = "\n".join(queries)
+
+    return combined_query
diff --git a/agents-api/agents_api/models/agent/patch_tool.py b/agents-api/agents_api/models/agent/patch_tool.py
@@ -0,0 +1,32 @@
+from ...common.utils import json
+from uuid import UUID
+
+from ...autogen.openapi_model import FunctionDef
+
+
+def patch_tool_by_id_query(
+    agent_id: UUID, tool_id: UUID, function: FunctionDef, embedding: list[float]
+) -> str:
+    # Agent update query
+    function = function.model_dump()
+
+    name = json.dumps(function["name"])
+    description = json.dumps(function["description"])
+    parameters = json.dumps(function.get("parameters", {}))
+
+    return f"""
+        ?[agent_id, tool_id, name, description, parameters, embedding, updated_at] <- [
+            [to_uuid("{agent_id}"), to_uuid("{tool_id}"), {name}, {description}, {parameters}, vec({embedding}), now()]
+        ]
+
+        :update agent_functions {{
+            agent_id,
+            tool_id,
+            name,
+            description,
+            parameters,
+            embedding,
+            updated_at,
+        }}
+        :returning
+    """
diff --git a/agents-api/agents_api/models/agent/update_agent.py b/agents-api/agents_api/models/agent/update_agent.py
@@ -36,7 +36,7 @@ def update_agent_query(
         # update the agent
         ?[{agent_update_cols}] <- {json.dumps(agent_update_vals)}
 
-        :update agents {{
+        :put agents {{
             {agent_update_cols}
         }}
         :returning
@@ -56,7 +56,7 @@ def update_agent_query(
         # update the agent settings
         ?[{settings_cols}] <- {json.dumps(settings_vals)}
 
-        :update agent_default_settings {{
+        :put agent_default_settings {{
             {settings_cols}
         }}
     }}

diff --git a/agents-api/agents_api/models/agent/update_tool.py b/agents-api/agents_api/models/agent/update_tool.py
@@ -19,7 +19,7 @@ def update_tool_by_id_query(
             [to_uuid("{agent_id}"), to_uuid("{tool_id}"), {name}, {description}, {parameters}, vec({embedding}), now()]
         ]
 
-        :update agent_functions {{
+        :put agent_functions {{
             agent_id,
             tool_id,
             name,

diff --git a/agents-api/agents_api/models/instructions/create_instructions.py b/agents-api/agents_api/models/instructions/create_instructions.py
@@ -7,9 +7,9 @@
 
 def create_instructions_query(
     agent_id: UUID,
-    instructions: list[Instruction] = [],
+    instructions: list[Instruction] | None = [],
 ):
-    instructions = [i.model_dump() for i in instructions]
+    instructions = [i.model_dump() for i in instructions or []]
 
     instruction_cols, instruction_rows = "", []
 

diff --git a/agents-api/agents_api/models/session/patch_session.py b/agents-api/agents_api/models/session/patch_session.py
@@ -0,0 +1,59 @@
+from ...common.utils import json
+from uuid import UUID
+
+from ...common.utils.cozo import cozo_process_mutate_data
+
+
+_fields = [
+    "situation",
+    "summary",
+    "metadata",
+    "created_at",
+    "session_id",
+    "developer_id",
+]
+
+
+def patch_session_query(
+    session_id: UUID,
+    developer_id: UUID,
+    **update_data,
+) -> str:
+    session_update_cols, session_update_vals = cozo_process_mutate_data(
+        {
+            **{k: v for k, v in update_data.items() if v is not None},
+        }
+    )
+
+    session_update_cols_lst = session_update_cols.split(",")
+    all_fields_lst = list(set(session_update_cols_lst).union(set(_fields)))
+    all_fields = ", ".join(all_fields_lst)
+    rest_fields = ", ".join(
+        list(
+            set(all_fields_lst)
+            - set([k for k, v in update_data.items() if v is not None])
+        )
+    )
+
+    session_update_query = f"""
+    {{
+        input[{session_update_cols}, session_id, developer_id] <- [
+            [{json.dumps(session_update_vals[0]).strip("[]")}, to_uuid("{session_id}"), to_uuid("{developer_id}")]
+        ]
+        
+        ?[{all_fields}, updated_at] :=
+            input[{session_update_cols}, session_id, developer_id],
+            *sessions{{
+                {rest_fields}, @ "NOW"
+            }},
+            updated_at = [floor(now()), true]
+
+        :update sessions {{
+            {all_fields}, updated_at
+        }}
+
+        :returning
+    }}
+    """
+
+    return session_update_query
diff --git a/agents-api/agents_api/models/user/patch_user.py b/agents-api/agents_api/models/user/patch_user.py
@@ -0,0 +1,26 @@
+from ...common.utils import json
+from uuid import UUID
+
+from ...common.utils.cozo import cozo_process_mutate_data
+from ...common.utils.datetime import utcnow
+
+
+def patch_user_query(developer_id: UUID, user_id: UUID, **update_data) -> str:
+    user_update_cols, user_update_vals = cozo_process_mutate_data(
+        {
+            **{k: v for k, v in update_data.items() if v is not None},
+            "user_id": user_id,
+            "developer_id": developer_id,
+            "updated_at": utcnow().timestamp(),
+        }
+    )
+
+    return f"""
+        # update the user
+        ?[{user_update_cols}] <- {json.dumps(user_update_vals)}
+
+        :update users {{
+            {user_update_cols}
+        }}
+        :returning
+    """
diff --git a/agents-api/agents_api/models/user/update_user.py b/agents-api/agents_api/models/user/update_user.py
@@ -19,7 +19,7 @@ def update_user_query(developer_id: UUID, user_id: UUID, **update_data) -> str:
         # update the user
         ?[{user_update_cols}] <- {json.dumps(user_update_vals)}
 
-        :update users {{
+        :put users {{
             {user_update_cols}
         }}
         :returning