Merge pull request #1012 from julep-ai/x/sessions-updated-at-precision

fix(agents-api): Fix sessions `updated_at` precision issue

Author: HamadaSalhab

agents-api/agents_api/activities/execute_system.py

@@ -16,9 +16,10 @@
     HybridDocSearchRequest,
     SystemDef,
     TextOnlyDocSearchRequest,
+    UpdateSessionRequest,
     VectorDocSearchRequest,
 )
-from ..common.protocol.tasks import StepContext
+from ..common.protocol.tasks import ExecutionInput, StepContext
 from ..common.storage_handler import auto_blob_store, load_from_blob_store_if_remote
 from ..env import testing
 from ..models.developer import get_developer
@@ -40,6 +41,9 @@ async def execute_system(
     if set(arguments.keys()) == {"bucket", "key"}:
         arguments = await load_from_blob_store_if_remote(arguments)
 
+    if not isinstance(context.execution_input, ExecutionInput):
+        raise TypeError("Expected ExecutionInput type for context.execution_input")
+
     arguments["developer_id"] = context.execution_input.developer_id
 
     # Unbox all the arguments
@@ -106,10 +110,11 @@ async def execute_system(
             await bg_runner()
             return res
 
+        # Handle create operations
         if system.operation == "create" and system.resource == "session":
             developer_id = arguments.pop("developer_id")
             session_id = arguments.pop("session_id", None)
-            data = CreateSessionRequest(**arguments)
+            create_session_request = CreateSessionRequest(**arguments)
 
             # In case sessions.create becomes asynchronous in the future
             if asyncio.iscoroutinefunction(handler):
@@ -118,7 +123,35 @@ async def execute_system(
             # Run the synchronous function in another process
             loop = asyncio.get_running_loop()
             return await loop.run_in_executor(
-                process_pool_executor, partial(handler, developer_id, session_id, data)
+                process_pool_executor,
+                partial(
+                    handler,
+                    developer_id=developer_id,
+                    session_id=session_id,
+                    data=create_session_request,
+                ),
+            )
+
+        # Handle update operations
+        if system.operation == "update" and system.resource == "session":
+            developer_id = arguments.pop("developer_id")
+            session_id = arguments.pop("session_id")
+            update_session_request = UpdateSessionRequest(**arguments)
+
+            # In case sessions.update becomes asynchronous in the future
+            if asyncio.iscoroutinefunction(handler):
+                return await handler()
+
+            # Run the synchronous function in another process
+            loop = asyncio.get_running_loop()
+            return await loop.run_in_executor(
+                process_pool_executor,
+                partial(
+                    handler,
+                    developer_id=developer_id,
+                    session_id=session_id,
+                    data=update_session_request,
+                ),
             )
 
         # Handle regular operations

agents-api/agents_api/models/chat/prepare_chat_context.py

@@ -37,6 +37,7 @@
         "session": make_session(
             agents=[a["id"] for a in d["agents"]],
             users=[u["id"] for u in d["users"]],
+            updated_at=d["session"].pop("updated_at") / (1000000.0),
             **d["session"],
         ),
         "toolsets": [

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

@@ -27,7 +27,14 @@
         TypeError: partialclass(HTTPException, status_code=400),
     }
 )
-@wrap_in_class(make_session, one=True)
+@wrap_in_class(
+    make_session,
+    one=True,
+    transform=lambda d: {
+        "updated_at": d.pop("updated_at") / (1000000.0),
+        **d,
+    },
+)
 @cozo_query
 @beartype
 def get_session(

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

@@ -32,6 +32,7 @@
     one=True,
     transform=lambda d: {
         "session": make_session(
+            updated_at=d["session"].pop("updated_at") / (1000000.0),
             **d["session"],
             agents=[a["id"] for a in d["agents"]],
             users=[u["id"] for u in d["users"]],

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

@@ -45,7 +45,7 @@
     one=True,
     transform=lambda d: {
         "id": d["session_id"],
-        "updated_at": d.pop("updated_at")[0],
+        "updated_at": d.pop("updated_at")[0] / (1000000.0),
         "jobs": [],
         **d,
     },