39 lines
1.6 KiB
Python
39 lines
1.6 KiB
Python
from app.database.models import Summary
|
|
from app.database.repositories import MemoryRepository
|
|
from app.llm.prompts import SUMMARY_PROMPT
|
|
from app.memory.context_builder import ContextBuilder
|
|
|
|
|
|
class Summarizer:
|
|
def __init__(self, memory: MemoryRepository, builder: ContextBuilder, llm, trigger_tokens: int):
|
|
self.memory, self.builder, self.llm, self.trigger_tokens = (
|
|
memory,
|
|
builder,
|
|
llm,
|
|
trigger_tokens,
|
|
)
|
|
|
|
async def needs_summary(self, chat_id: int, thread_id: int | None) -> bool:
|
|
messages = await self.memory.context_messages(chat_id, thread_id, limit=300)
|
|
return self.builder.count_tokens("\n".join(m.text for m in messages)) >= self.trigger_tokens
|
|
|
|
async def summarize(self, chat_id: int, thread_id: int | None) -> Summary | None:
|
|
previous = await self.memory.latest_summary(chat_id, thread_id)
|
|
messages = await self.memory.context_messages(chat_id, thread_id, limit=300)
|
|
if not messages:
|
|
return None
|
|
context = self.builder.build(previous, messages)
|
|
text = await self.llm.complete(
|
|
[{"role": "user", "content": f"{SUMMARY_PROMPT}\n\n{context}"}]
|
|
)
|
|
summary = Summary(
|
|
chat_id=chat_id,
|
|
thread_id=thread_id,
|
|
version=(previous.version if previous else 0) + 1,
|
|
text=text,
|
|
from_message_id=messages[0].telegram_message_id,
|
|
to_message_id=messages[-1].telegram_message_id,
|
|
)
|
|
await self.memory.save_summary(summary)
|
|
return summary
|