fix(agents-api): Minor fix to tests

Signed-off-by: Diwank Tomer <[email protected]>
julep-ai · Aug 14, 2024 · 82c8af2 · 82c8af2
1 parent 3a38e70
commit 82c8af2
Show file tree

Hide file tree

Showing 3 changed files with 23 additions and 32 deletions.
diff --git a/agents-api/agents_api/models/entry/create_entries.py b/agents-api/agents_api/models/entry/create_entries.py
@@ -78,8 +78,7 @@ def create_entries(
         verify_developer_owns_resource_query(
             developer_id, "sessions", session_id=session_id
         ),
-        mark_session_as_updated
-        and mark_session_updated_query(developer_id, session_id),
+        mark_session_updated_query(developer_id, session_id) if mark_session_as_updated else "",
         create_query,
     ]
 

diff --git a/agents-api/agents_api/routers/tasks/create_task.py b/agents-api/agents_api/routers/tasks/create_task.py
@@ -1,14 +1,13 @@
 from typing import Annotated
-from uuid import uuid4
 
-import pandas as pd
 from fastapi import Depends
 from pydantic import UUID4
 from starlette.status import HTTP_201_CREATED
 
 from agents_api.autogen.openapi_model import (
     CreateTaskRequest,
     ResourceCreatedResponse,
+    Task,
 )
 from agents_api.dependencies.developer_id import get_developer_id
 from agents_api.models.task.create_task import create_task as create_task_query
@@ -18,29 +17,18 @@
 
 @router.post("/agents/{agent_id}/tasks", status_code=HTTP_201_CREATED, tags=["tasks"])
 async def create_task(
-    request: CreateTaskRequest,
+    data: CreateTaskRequest,
     agent_id: UUID4,
     x_developer_id: Annotated[UUID4, Depends(get_developer_id)],
 ) -> ResourceCreatedResponse:
-    task_id = uuid4()
-
     # TODO: Do thorough validation of the task spec
 
-    workflows = [
-        {"name": "main", "steps": [w.model_dump() for w in request.main]},
-    ] + [{"name": name, "steps": steps} for name, steps in request.model_extra.items()]
-
-    resp: pd.DataFrame = create_task_query(
-        agent_id=agent_id,
-        task_id=task_id,
+    task: Task = create_task_query(
         developer_id=x_developer_id,
-        name=request.name,
-        description=request.description,
-        input_schema=request.input_schema or {},
-        tools_available=request.tools or [],
-        workflows=workflows,
+        agent_id=agent_id,
+        data=data,
     )
 
     return ResourceCreatedResponse(
-        id=resp["task_id"][0], created_at=resp["created_at"][0]
+        id=task.id, created_at=task.created_at, jobs=[]
     )
diff --git a/agents-api/tests/test_task_routes.py b/agents-api/tests/test_task_routes.py
@@ -10,12 +10,14 @@
 def _(client=client, agent=test_agent):
     data = dict(
         name="test user",
-        main={
-            "kind_": "evaluate",
-            "evaluate": {
-                "additionalProp1": "value1",
-            },
-        },
+        main=[
+            {
+                "kind_": "evaluate",
+                "evaluate": {
+                    "additionalProp1": "value1",
+                },
+            }
+        ],
     )
 
     response = client.request(
@@ -31,12 +33,14 @@ def _(client=client, agent=test_agent):
 def _(make_request=make_request, agent=test_agent):
     data = dict(
         name="test user",
-        main={
-            "kind_": "evaluate",
-            "evaluate": {
-                "additionalProp1": "value1",
-            },
-        },
+        main=[
+            {
+                "kind_": "evaluate",
+                "evaluate": {
+                    "additionalProp1": "value1",
+                },
+            }
+        ],
     )
 
     response = make_request(