231 lines
8.1 KiB
Python
231 lines
8.1 KiB
Python
"""Context manager: builds the LLM prompt context with guaranteed-recent + dynamic summarization."""
|
||
from __future__ import annotations
|
||
|
||
import json
|
||
import uuid
|
||
from typing import Any, Dict, List, Optional
|
||
|
||
from sqlalchemy import select
|
||
from sqlalchemy.ext.asyncio import AsyncSession
|
||
|
||
from app.core.llm import LlmClient
|
||
from app.core.settings_service import cast_setting, get_all_settings
|
||
from app.logging_setup import get_logger
|
||
from app.models import Message, World
|
||
from app.prompts.templates import get_prompt
|
||
|
||
log = get_logger("context")
|
||
|
||
|
||
async def build_orchestrator_messages(
|
||
db: AsyncSession,
|
||
world: World,
|
||
session_id: uuid.UUID,
|
||
action_text: str,
|
||
) -> tuple[List[Dict[str, Any]], Dict[str, Any]]:
|
||
"""Build the messages list for the orchestrator LLM call.
|
||
|
||
Returns (messages, settings_used).
|
||
"""
|
||
settings_map = await get_all_settings(db)
|
||
recent_n = int(cast_setting("context.recent_messages", settings_map.get("context.recent_messages", 10)))
|
||
threshold = int(cast_setting("context.compress_threshold", settings_map.get("context.compress_threshold", 20)))
|
||
summary_n = int(cast_setting("context.summary_messages", settings_map.get("context.summary_messages", 10)))
|
||
|
||
# Load all messages ordered by seq
|
||
result = await db.execute(
|
||
select(Message).where(Message.session_id == session_id).order_by(Message.seq)
|
||
)
|
||
all_msgs: List[Message] = list(result.scalars().all())
|
||
|
||
# Check if we need to compress
|
||
if len(all_msgs) >= threshold:
|
||
await _maybe_compress(db, session_id, all_msgs, summary_n, recent_n, world, settings_map)
|
||
# Reload after compression
|
||
result = await db.execute(
|
||
select(Message).where(Message.session_id == session_id).order_by(Message.seq)
|
||
)
|
||
all_msgs = list(result.scalars().all())
|
||
|
||
# Get summary message (the latest summary before the recent window)
|
||
summary_text = ""
|
||
visible_msgs = [m for m in all_msgs if not m.hidden]
|
||
if len(visible_msgs) > recent_n:
|
||
# Look for the latest summary
|
||
summaries = [m for m in all_msgs if m.kind == "summary"]
|
||
if summaries:
|
||
summary_text = summaries[-1].content
|
||
|
||
recent = visible_msgs[-recent_n:] if visible_msgs else []
|
||
|
||
# Build orchestrator system prompt with current state
|
||
defn = world.definition or {}
|
||
system_prompt_template = get_prompt("orchestrator", world.language)
|
||
player_state = world.state.get("player", {}) if world.state else {}
|
||
system_prompt = system_prompt_template.format(
|
||
world_name=world.name,
|
||
setting_description=defn.get("setting_description", "")[:800],
|
||
rules=json.dumps(defn.get("rules", {}), ensure_ascii=False)[:600],
|
||
current_time=world.current_time or "",
|
||
player_state=json.dumps(player_state, ensure_ascii=False)[:600],
|
||
plot_rails=json.dumps(defn.get("plot_rails", {}), ensure_ascii=False)[:400],
|
||
summary=summary_text or "(нет сводки)",
|
||
)
|
||
|
||
messages: List[Dict[str, Any]] = [{"role": "system", "content": system_prompt}]
|
||
|
||
# Add summary as a system note if present
|
||
if summary_text:
|
||
messages.append({"role": "system", "content": f"Сводка прошлого:\n{summary_text}"})
|
||
|
||
# Add recent visible messages
|
||
for m in recent:
|
||
if m.kind == "player_action":
|
||
messages.append({"role": "user", "content": m.content})
|
||
elif m.kind == "narrative_step":
|
||
messages.append({"role": "assistant", "content": m.content})
|
||
|
||
# The current action
|
||
messages.append({"role": "user", "content": f'Действие игрока: "{action_text}"'})
|
||
|
||
return messages, settings_map
|
||
|
||
|
||
async def _maybe_compress(
|
||
db: AsyncSession,
|
||
session_id: uuid.UUID,
|
||
all_msgs: List[Message],
|
||
summary_n: int,
|
||
recent_n: int,
|
||
world: World,
|
||
settings_map: Dict[str, Any],
|
||
) -> None:
|
||
"""If history exceeds threshold, summarize older messages into a single summary message."""
|
||
visible = [m for m in all_msgs if not m.hidden]
|
||
if len(visible) <= recent_n + summary_n:
|
||
return
|
||
|
||
# Take the messages that will be summarized (everything before the recent window)
|
||
to_summarize = visible[:-recent_n]
|
||
if not to_summarize:
|
||
return
|
||
|
||
# Build summarization input
|
||
summary_input_lines = []
|
||
for m in to_summarize:
|
||
prefix = {
|
||
"player_action": "Игрок",
|
||
"narrative_step": "Сцена",
|
||
"summary": "Сводка",
|
||
"orchestrator_plan": "GM",
|
||
"technical_offscreen": "За кадром",
|
||
}.get(m.kind, m.kind)
|
||
summary_input_lines.append(f"{prefix}: {m.content[:300]}")
|
||
summary_input = "\n\n".join(summary_input_lines)
|
||
|
||
llm = LlmClient(settings_map)
|
||
system_prompt = get_prompt("summarizer", world.language)
|
||
response = await llm.chat(
|
||
messages=[
|
||
{"role": "system", "content": system_prompt},
|
||
{"role": "user", "content": summary_input[:4000]},
|
||
],
|
||
temperature=float(cast_setting("llm.summary_temperature", settings_map.get("llm.summary_temperature", 0.3))),
|
||
max_tokens=300,
|
||
purpose="summary",
|
||
session_id=session_id,
|
||
db=db,
|
||
)
|
||
|
||
# Parse summary response
|
||
summary_text = response.text
|
||
facts: List[Dict[str, Any]] = []
|
||
import re as _re
|
||
json_match = _re.search(r"\{[\s\S]*\}", response.text)
|
||
if json_match:
|
||
try:
|
||
data = json.loads(json_match.group(0))
|
||
summary_text = data.get("summary", response.text)
|
||
facts = data.get("facts", [])
|
||
except json.JSONDecodeError:
|
||
pass
|
||
|
||
# Create summary message
|
||
next_seq = (max((m.seq for m in all_msgs), default=0)) + 1
|
||
summary_msg = Message(
|
||
session_id=session_id,
|
||
seq=next_seq,
|
||
role="system",
|
||
kind="summary",
|
||
content=summary_text,
|
||
payload={"summarized_count": len(to_summarize), "facts": facts},
|
||
is_pinned=True,
|
||
hidden=False,
|
||
)
|
||
db.add(summary_msg)
|
||
|
||
# Hide the summarized messages (but keep them in DB)
|
||
for m in to_summarize:
|
||
m.hidden = True
|
||
|
||
# Index facts into RAG glossary
|
||
if facts:
|
||
from app.core.rag import get_rag
|
||
from app.models import GlossaryEntry
|
||
rag = await get_rag(settings_map)
|
||
for f in facts:
|
||
if not isinstance(f, dict):
|
||
continue
|
||
entry = GlossaryEntry(
|
||
world_id=world.id,
|
||
session_id=session_id,
|
||
kind=f.get("kind", "lore"),
|
||
name=f.get("name", "unknown"),
|
||
description=f.get("description", ""),
|
||
payload={},
|
||
)
|
||
db.add(entry)
|
||
await db.flush()
|
||
await rag.upsert_glossary(
|
||
world_id=world.id,
|
||
entry_id=entry.id,
|
||
kind=entry.kind,
|
||
name=entry.name,
|
||
description=entry.description,
|
||
payload={},
|
||
settings_map=settings_map,
|
||
)
|
||
|
||
await db.commit()
|
||
log.info("context_compressed", session_id=str(session_id), summarized=len(to_summarize))
|
||
|
||
|
||
async def build_step_writer_messages(
|
||
db: AsyncSession,
|
||
world: World,
|
||
session_id: uuid.UUID,
|
||
outcome: str,
|
||
narrative_prompt: str,
|
||
) -> List[Dict[str, Any]]:
|
||
"""Build messages for the step writer LLM call."""
|
||
defn = world.definition or {}
|
||
player_state = world.state.get("player", {}) if world.state else {}
|
||
system_prompt = get_prompt("step_writer", world.language).format(
|
||
setting_description=defn.get("setting_description", "")[:600],
|
||
current_time=world.current_time or "",
|
||
player_state=json.dumps(player_state, ensure_ascii=False)[:400],
|
||
outcome=outcome,
|
||
narrative_prompt=narrative_prompt[:600],
|
||
)
|
||
return [{"role": "system", "content": system_prompt}]
|
||
|
||
|
||
async def build_subagent_messages(
|
||
world: World,
|
||
task: str,
|
||
context: str,
|
||
) -> List[Dict[str, Any]]:
|
||
"""Build messages for a clean-context sub-agent call."""
|
||
system_prompt = get_prompt("subagent", world.language).format(task=task, context=context[:600])
|
||
return [{"role": "system", "content": system_prompt}]
|