31 lines
1.1 KiB
Python
31 lines
1.1 KiB
Python
import tiktoken
|
|
|
|
from app.database.models import Message, Summary
|
|
|
|
|
|
class ContextBuilder:
|
|
def __init__(self, token_limit: int):
|
|
self.token_limit = token_limit
|
|
try:
|
|
self.encoder = tiktoken.get_encoding("cl100k_base")
|
|
except Exception:
|
|
self.encoder = None
|
|
|
|
def count_tokens(self, text: str) -> int:
|
|
return len(self.encoder.encode(text)) if self.encoder else max(1, len(text) // 4)
|
|
|
|
def build(self, summary: Summary | None, messages: list[Message]) -> str:
|
|
sections = [f"Current memory summary:\n{summary.text}" if summary else ""]
|
|
recent: list[str] = []
|
|
budget = self.token_limit - self.count_tokens(sections[0])
|
|
for item in reversed(messages):
|
|
line = f"[{item.sent_at.isoformat()}] {item.user_name}: {item.text}"
|
|
cost = self.count_tokens(line)
|
|
if cost > budget:
|
|
break
|
|
recent.append(line)
|
|
budget -= cost
|
|
if recent:
|
|
sections.append("Recent messages:\n" + "\n".join(reversed(recent)))
|
|
return "\n\n".join(part for part in sections if part)
|