From e1275f17f5287f5acddf4adf872c89f0b6f42089 Mon Sep 17 00:00:00 2001 From: Dmitry Paramonov Date: Fri, 19 Apr 2024 13:29:06 +0300 Subject: [PATCH] fix: Fix queries calls --- agents-api/agents_api/activities/summarization.py | 15 +++++---------- agents-api/agents_api/activities/types.py | 6 ++++-- agents-api/agents_api/clients/worker/types.py | 6 ++++-- .../models/entry/entries_summarization.py | 4 ++-- 4 files changed, 15 insertions(+), 16 deletions(-) diff --git a/agents-api/agents_api/activities/summarization.py b/agents-api/agents_api/activities/summarization.py index 656b7bc7c..0e9159713 100644 --- 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 index 37fd8015d..4fa1d4d46 100644 --- 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 index 86bf564dd..16fa8404b 100644 --- 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 index 6350968c6..26e317342 100644 --- a/agents-api/agents_api/models/entry/entries_summarization.py +++ b/agents-api/agents_api/models/entry/entries_summarization.py @@ -89,8 +89,8 @@ def entries_summarization_query( entries = [ [ - new_entry.id, - session_id, + str(new_entry.id), + str(session_id), source, role, new_entry.name or "",