fix: Fix queries calls

julep-ai · Apr 19, 2024 · 44ae715 · 44ae715
1 parent 7000518
commit 44ae715
Show file tree

Hide file tree

Showing 4 changed files with 18 additions and 22 deletions.
diff --git a/agents-api/agents_api/activities/summarization.py b/agents-api/agents_api/activities/summarization.py
@@ -4,7 +4,6 @@
 from typing import Callable
 from textwrap import dedent
 from temporalio import activity
-from agents_api.clients.cozo import client
 from agents_api.models.entry.entries_summarization import (
     get_toplevel_entries_query,
     entries_summarization_query,
@@ -154,9 +153,7 @@ async def summarization(session_id: str) -> None:
     session_id = UUID(session_id)
     entries = [
         Entry(**row)
-        for _, row in client.run(
-            get_toplevel_entries_query(session_id=session_id)
-        ).iterrows()
+        for _, row in get_toplevel_entries_query(session_id=session_id).iterrows()
     ]
 
     assert len(entries) > 0, "no need to summarize on empty entries list"
@@ -172,10 +169,8 @@ async def summarization(session_id: str) -> None:
         timestamp=entries[-1].timestamp + 0.01,
     )
 
-    client.run(
-        entries_summarization_query(
-            session_id=session_id,
-            new_entry=new_entry,
-            old_entry_ids=[e.id for e in entries],
-        )
+    entries_summarization_query(
+        session_id=session_id,
+        new_entry=new_entry,
+        old_entry_ids=[e.id for e in entries],
     )
diff --git a/agents-api/agents_api/activities/types.py b/agents-api/agents_api/activities/types.py
@@ -25,10 +25,12 @@ class ChatML(BaseModel):
     token_count: Optional[int] = None
 
 
-class BaseTask(BaseModel): ...
+class BaseTask(BaseModel):
+    ...
 
 
-class BaseTaskArgs(BaseModel): ...
+class BaseTaskArgs(BaseModel):
+    ...
 
 
 class AddPrinciplesTaskArgs(BaseTaskArgs):

diff --git a/agents-api/agents_api/clients/worker/types.py b/agents-api/agents_api/clients/worker/types.py
@@ -25,10 +25,12 @@ class ChatML(BaseModel):
     token_count: Optional[int] = None
 
 
-class BaseTask(BaseModel): ...
+class BaseTask(BaseModel):
+    ...
 
 
-class BaseTaskArgs(BaseModel): ...
+class BaseTaskArgs(BaseModel):
+    ...
 
 
 class MemoryManagementTaskArgs(BaseTaskArgs):

diff --git a/agents-api/agents_api/models/entry/entries_summarization.py b/agents-api/agents_api/models/entry/entries_summarization.py
@@ -83,18 +83,15 @@ def entries_summarization_query(
     # Create a list of relations indicating which entries the new entry summarizes.
 
     # Convert the new entry's source information into JSON format for storage.
-    source = json.dumps(new_entry.source)
-    role = json.dumps(new_entry.role)
-    content = json.dumps(new_entry.content)
 
     entries = [
         [
-            new_entry.id,
-            session_id,
-            source,
-            role,
+            str(new_entry.id),
+            str(session_id),
+            new_entry.source,
+            new_entry.role,
             new_entry.name or "",
-            content,
+            new_entry.content,
             new_entry.token_count,
             new_entry.tokenizer,
             new_entry.created_at,