fix(agents-api): fix create_agent and create_or_update_agent query

.env.example

@@ -6,6 +6,7 @@ COZO_AUTH_TOKEN=<your_cozo_auth_token>
 TEMPORAL_POSTGRES_PASSWORD=<your_temporal_postgres_password>
 LITELLM_POSTGRES_PASSWORD=<your_litellm_postgres_password>
 LITELLM_MASTER_KEY=<your_litellm_master_key>
+LITELLM_SALT_KEY=<your_litellm_salt_key>
 LITELLM_REDIS_PASSWORD=<your_litellm_redis_password>
 EMBEDDING_SERVICE_BASE=http://text-embeddings-inference-<gpu|cpu>  # Use the 'gpu' profile to run on GPU
 

agents-api/agents_api/activities/task_steps/raise_complete_async.py

@@ -13,7 +13,7 @@ async def raise_complete_async(context: StepContext, output: StepOutcome) -> Non
     # TODO: Create a transtition to "wait" and save the captured_token to the transition
 
     captured_token = activity.info().task_token
-    captured_token = captured_token.decode('latin-1')
+    captured_token = captured_token.decode("latin-1")
     transition_info = CreateTransitionRequest(
         current=context.cursor,
         type="wait",

agents-api/agents_api/activities/task_steps/transition_step.py

@@ -29,4 +29,4 @@ async def transition_step(
 
 transition_step = activity.defn(name="transition_step")(
     transition_step if not testing else mock_transition_step
-)
+)

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

@@ -80,6 +80,29 @@ def create_or_update_agent(
         }
     )
 
+    # TODO: remove this
+    ### # Create default agent settings
+    ### # Construct a query to insert default settings for the new agent
+    ### default_settings_query = f"""
+    ### %if {{
+    ###     len[count(agent_id)] :=
+    ###         *agent_default_settings{{agent_id}},
+    ###         agent_id = to_uuid($agent_id)
+
+    ###     ?[should_create] := len[count], count > 0
+    ### }}
+    ### %then {{
+    ###     ?[{settings_cols}] <- $settings_vals
+
+    ###     :put agent_default_settings {{
+    ###         {settings_cols}
+    ###     }}
+    ### }}
+    ### """
+
+    # FIXME: This create or update query will overwrite the settings
+    #        Need to find a way to only run the insert query if the agent_default_settings
+
     # Create default agent settings
     # Construct a query to insert default settings for the new agent
     default_settings_query = f"""
@@ -89,6 +112,7 @@ def create_or_update_agent(
             {settings_cols}
         }}
     """
+
     # create the agent
     # Construct a query to insert the new agent record into the agents table
     agent_query = """
@@ -127,7 +151,7 @@ def create_or_update_agent(
 
     queries = [
         verify_developer_id_query(developer_id),
-        default_settings_query if default_settings else None,
+        default_settings_query,
         agent_query,
     ]
 

agents-api/agents_api/models/session/create_or_update_session.py

@@ -26,6 +26,7 @@
 
 @rewrap_exceptions(
     {
+        AssertionError: partialclass(HTTPException, status_code=400),
         QueryException: partialclass(HTTPException, status_code=400),
         ValidationError: partialclass(HTTPException, status_code=400),
         TypeError: partialclass(HTTPException, status_code=400),

agents-api/agents_api/models/session/create_session.py

@@ -27,6 +27,7 @@
 
 @rewrap_exceptions(
     {
+        AssertionError: partialclass(HTTPException, status_code=400),
         QueryException: partialclass(HTTPException, status_code=400),
         ValidationError: partialclass(HTTPException, status_code=400),
         TypeError: partialclass(HTTPException, status_code=400),

agents-api/agents_api/models/session/prepare_session_data.py

@@ -127,6 +127,33 @@ def prepare_session_data(
             "instructions":      instructions,
         }
 
+    # Version where we don't have default settings
+    agent_data[collect(record)] :=
+        agents[agent_ids],
+        agent_id in agent_ids,
+        *agents{
+            agent_id,
+            model,
+            name,
+            about,
+            created_at,
+            updated_at,
+            metadata,
+            instructions,
+        },
+        not settings_data[agent_id, _],
+        record = {
+            "id":         agent_id,
+            "name":       name,
+            "model":      model,
+            "about":      about,
+            "created_at": created_at,
+            "updated_at": updated_at,
+            "metadata":   metadata,
+            "default_settings":  {},
+            "instructions":      instructions,
+        }
+
     user_data[collect(record)] :=
         users[user_ids],
         user_id in user_ids,

agents-api/agents_api/routers/agents/create_agent.py

@@ -19,6 +19,7 @@ async def create_agent(
     x_developer_id: Annotated[UUID4, Depends(get_developer_id)],
     data: CreateAgentRequest,
 ) -> ResourceCreatedResponse:
+    # TODO: Validate model name
     agent = models.agent.create_agent(
         developer_id=x_developer_id,
         data=data,

agents-api/agents_api/routers/agents/create_or_update_agent.py

@@ -21,6 +21,7 @@ async def create_or_update_agent(
     data: CreateOrUpdateAgentRequest,
     x_developer_id: Annotated[UUID4, Depends(get_developer_id)],
 ) -> ResourceCreatedResponse:
+    # TODO: Validate model name
     agent = models.agent.create_or_update_agent(
         developer_id=x_developer_id,
         agent_id=agent_id,

agents-api/agents_api/routers/sessions/create_or_update_session.py

@@ -7,7 +7,7 @@
 
 from ...autogen.openapi_model import (
     CreateOrUpdateSessionRequest,
-    ResourceCreatedResponse,
+    ResourceUpdatedResponse,
 )
 from ...dependencies.developer_id import get_developer_id
 from ...models.session.create_or_update_session import (
@@ -21,14 +21,11 @@ async def create_or_update_session(
     x_developer_id: Annotated[UUID4, Depends(get_developer_id)],
     session_id: UUID,
     data: CreateOrUpdateSessionRequest,
-) -> ResourceCreatedResponse:
-    session = create_session_query(
+) -> ResourceUpdatedResponse:
+    session_updated = create_session_query(
         developer_id=x_developer_id,
         session_id=session_id,
         data=data,
     )
 
-    return ResourceCreatedResponse(
-        id=session.id,
-        created_at=session.created_at,
-    )
+    return session_updated

agents-api/agents_api/workflows/task_execution/__init__.py

@@ -416,10 +416,10 @@ async def run(
                         task_steps.prompt_step,
                         context,
                         schedule_to_close_timeout=timedelta(
-                            seconds=30 if debug or testing else 600),
+                            seconds=30 if debug or testing else 600
+                        ),
                     )
-                    state = PartialTransition(
-                        output=new_response.output, type="resume")
+                    state = PartialTransition(output=new_response.output, type="resume")
 
             # case PromptStep(), StepOutcome(
             #     output=response