From c05113fa0468fabe1da83d2d874553bb29b46398 Mon Sep 17 00:00:00 2001
From: Dmitry Paramonov <asmatic075@gmail.com>
Date: Wed, 22 May 2024 15:31:40 +0300
Subject: [PATCH] WIP

---
 .../agents_api/activities/summarization.py    | 39 +++++++++++++++++++
 1 file changed, 39 insertions(+)

diff --git a/agents-api/agents_api/activities/summarization.py b/agents-api/agents_api/activities/summarization.py
index a9843ab8e..f14f003d8 100644
--- a/agents-api/agents_api/activities/summarization.py
+++ b/agents-api/agents_api/activities/summarization.py
@@ -1,5 +1,6 @@
 #!/usr/bin/env python3
 
+import asyncio
 from uuid import UUID
 from typing import Callable
 from textwrap import dedent
@@ -10,6 +11,7 @@
     entries_summarization_query,
 )
 from agents_api.common.protocol.entries import Entry
+from agents_api.rec_sum import summarize_messages, get_entities, trim_messages
 from ..env import summarization_model_name
 
 
@@ -176,3 +178,40 @@ async def summarization(session_id: str) -> None:
         new_entry=new_entry,
         old_entry_ids=[e.id for e in entries],
     )
+
+
+# TODO: rename and remove the old one
+@activity.defn
+async def summarization_new(session_id: str) -> None:
+    session_id = UUID(session_id)
+    entries = [row for _, row in get_toplevel_entries_query(session_id=session_id).iterrows()]
+
+    assert len(entries) > 0, "no need to summarize on empty entries list"
+
+    trimmed_messages, entities = await asyncio.gather(
+        trim_messages(entries),
+        get_entities(entries),
+    )
+
+    summarized = await summarize_messages(trimmed_messages)
+
+    for msg in summarized:
+        # response = await run_prompt(
+        #     dialog=entries, previous_memories=[], model=summarization_model_name
+        # )
+
+        # TODO: set correct content
+        new_entry = Entry(
+            session_id=session_id,
+            source="summarizer",
+            role="system",
+            name="information",
+            content=msg,
+            timestamp=entries[-1].timestamp + 0.01,
+        )
+
+        entries_summarization_query(
+            session_id=session_id,
+            new_entry=new_entry,
+            old_entry_ids=[e.id for e in entries],
+        )