599 lines
22 KiB
Python
599 lines
22 KiB
Python
"""Context Engine — the central intelligence of the system.
|
|
|
|
Builds structured prompts from session state. Never includes raw tool
|
|
outputs. Handles compaction, artifact summarization, and selective
|
|
rehydration.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import logging
|
|
import time
|
|
from collections import defaultdict
|
|
from typing import Any
|
|
|
|
from ..config import settings
|
|
|
|
logger = logging.getLogger(__name__)
|
|
from ..models.agent import AgentProfile
|
|
from ..models.artifacts import ArtifactSummary
|
|
from ..memory.embeddings import EmbeddingService
|
|
from ..memory.store import MemoryStore
|
|
from ..models.context import (
|
|
ContextPackage,
|
|
ContextSection,
|
|
ContextSectionType,
|
|
MemoryDocument,
|
|
MemoryType,
|
|
)
|
|
from ..models.session import SessionState, TaskState
|
|
from .compactor import ContextCompactor, estimate_tokens
|
|
|
|
|
|
class ContextEngine:
|
|
"""Assembles the context package that gets sent to the model.
|
|
|
|
The engine enforces a strict contract:
|
|
- Raw tool outputs NEVER appear in the context
|
|
- Each section has a priority for compaction
|
|
- Immutable rules are always included in full
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
compactor: ContextCompactor | None = None,
|
|
memory_store: MemoryStore | None = None,
|
|
) -> None:
|
|
self.compactor = compactor or ContextCompactor(
|
|
max_tokens=settings.context_max_tokens
|
|
)
|
|
self.memory = memory_store
|
|
self._embed_service: EmbeddingService | None = None
|
|
# Debug history: last N context builds per session
|
|
self._history: dict[str, list[dict[str, Any]]] = defaultdict(list)
|
|
self._max_history = 20
|
|
# Full context of the LAST build per session (not accumulated)
|
|
self._last_full_context: dict[str, dict[str, Any]] = {}
|
|
|
|
# ------------------------------------------------------------------
|
|
# Public — build context for a model call
|
|
# ------------------------------------------------------------------
|
|
|
|
async def build_context(
|
|
self,
|
|
session: SessionState,
|
|
agent: AgentProfile,
|
|
artifacts: list[ArtifactSummary] | None = None,
|
|
conversation: list[dict[str, Any]] | None = None,
|
|
extra_instructions: str = "",
|
|
) -> ContextPackage:
|
|
"""Build a full ContextPackage for the given agent and session.
|
|
|
|
The conversation parameter contains real assistant/tool messages
|
|
with complete tool results. These go into the messages array,
|
|
not the system prompt — like professional agentic tools.
|
|
"""
|
|
|
|
sections: list[ContextSection] = []
|
|
allowed = set(agent.context_sections)
|
|
|
|
# 1. Immutable rules — highest priority, never trimmed
|
|
if "immutable_rules" in allowed:
|
|
sections.append(self._build_immutable_rules(session, agent))
|
|
|
|
# 2. Project profile
|
|
if "project_profile" in allowed:
|
|
sections.append(self._build_project_profile(session))
|
|
|
|
# 3. Knowledge base — loaded from memory store
|
|
if "knowledge_base" in allowed and self.memory:
|
|
kb_section = await self._build_knowledge_base(session)
|
|
if kb_section:
|
|
sections.append(kb_section)
|
|
|
|
# 4. Task history — compact summaries of past tasks in this session
|
|
if "task_state" in allowed and session.task_history:
|
|
sections.append(self._build_task_history(session))
|
|
|
|
# 5. Task state — current task (includes compacted previous steps)
|
|
if "task_state" in allowed and session.current_task:
|
|
sections.append(self._build_task_state(session.current_task))
|
|
|
|
# Compact to fit budget
|
|
sections = self.compactor.compact_sections(sections)
|
|
|
|
# Assemble system prompt from sections
|
|
system_prompt = self._assemble_system_prompt(sections)
|
|
|
|
# Build messages with real conversation history
|
|
messages = self._build_messages(session, conversation)
|
|
|
|
total_tokens = estimate_tokens(system_prompt) + sum(
|
|
estimate_tokens(m.get("content", "")) for m in messages
|
|
)
|
|
|
|
package = ContextPackage(
|
|
sections=sections,
|
|
system_prompt=system_prompt,
|
|
messages=messages,
|
|
total_token_estimate=total_tokens,
|
|
)
|
|
|
|
# Guardar contexto completo del último build (solo el último por sesión)
|
|
self._last_full_context[session.session_id] = {
|
|
"system_prompt": system_prompt,
|
|
"messages": messages,
|
|
"total_tokens": total_tokens,
|
|
"timestamp": time.time(),
|
|
}
|
|
|
|
# --- Debug: log and store context build ---
|
|
section_summary = []
|
|
for s in sections:
|
|
section_summary.append({
|
|
"type": s.section_type.value,
|
|
"priority": s.priority,
|
|
"tokens": s.token_estimate,
|
|
"chars": len(s.content),
|
|
"preview": s.content[:150].replace("\n", " "),
|
|
})
|
|
|
|
conv_len = len(conversation) if conversation else 0
|
|
debug_entry = {
|
|
"timestamp": time.time(),
|
|
"agent": agent.role,
|
|
"agent_name": agent.name,
|
|
"total_tokens": total_tokens,
|
|
"sections": section_summary,
|
|
"sections_count": len(sections),
|
|
"compacted": len(sections) < len(allowed),
|
|
"system_prompt_tokens": estimate_tokens(system_prompt),
|
|
"user_message_preview": messages[0]["content"][:200] if messages else "",
|
|
"artifacts_count": len(artifacts) if artifacts else 0,
|
|
"conversation_messages": conv_len,
|
|
}
|
|
|
|
history = self._history[session.session_id]
|
|
history.append(debug_entry)
|
|
if len(history) > self._max_history:
|
|
self._history[session.session_id] = history[-self._max_history:]
|
|
|
|
logger.info(
|
|
"Context built for [%s/%s] — %d sections, ~%d tokens, artifacts=%d, conversation=%d msgs",
|
|
session.session_id[:8],
|
|
agent.role,
|
|
len(sections),
|
|
total_tokens,
|
|
len(artifacts) if artifacts else 0,
|
|
conv_len,
|
|
)
|
|
|
|
return package
|
|
|
|
def get_debug_history(self, session_id: str) -> list[dict[str, Any]]:
|
|
"""Return the context build history for a session."""
|
|
return list(self._history.get(session_id, []))
|
|
|
|
def get_last_context_debug(self, session_id: str) -> dict[str, Any] | None:
|
|
"""Return the most recent context build for a session."""
|
|
history = self._history.get(session_id, [])
|
|
return history[-1] if history else None
|
|
|
|
def get_last_full_context(self, session_id: str) -> dict[str, Any] | None:
|
|
"""Return the full context (system_prompt + messages) of the last build."""
|
|
return self._last_full_context.get(session_id)
|
|
|
|
def rehydrate_artifact(
|
|
self,
|
|
artifact: ArtifactSummary,
|
|
full_content: str,
|
|
) -> ContextSection:
|
|
"""Selectively rehydrate an artifact into a full context section.
|
|
|
|
Used when the agent explicitly needs the full content of a
|
|
specific artifact (e.g. reviewing generated code).
|
|
"""
|
|
content = (
|
|
f"## Rehydrated Artifact: {artifact.title}\n"
|
|
f"Type: {artifact.artifact_type} | Source: {artifact.source_tool}\n"
|
|
f"---\n{full_content}\n---"
|
|
)
|
|
return ContextSection(
|
|
section_type=ContextSectionType.WORKING_CONTEXT,
|
|
content=content,
|
|
priority=5,
|
|
token_estimate=estimate_tokens(content),
|
|
)
|
|
|
|
def summarize_tool_output(
|
|
self,
|
|
tool_name: str,
|
|
raw_output: str,
|
|
session_id: str,
|
|
task_id: str,
|
|
) -> ArtifactSummary:
|
|
"""Delegate to compactor — raw output never enters context."""
|
|
return self.compactor.summarize_tool_output(
|
|
tool_name=tool_name,
|
|
raw_output=raw_output,
|
|
session_id=session_id,
|
|
task_id=task_id,
|
|
)
|
|
|
|
# ------------------------------------------------------------------
|
|
# Section builders
|
|
# ------------------------------------------------------------------
|
|
|
|
def _build_immutable_rules(
|
|
self, session: SessionState, agent: AgentProfile
|
|
) -> ContextSection:
|
|
parts = [
|
|
"# System Rules (Immutable)",
|
|
"",
|
|
agent.system_prompt,
|
|
"",
|
|
]
|
|
if session.immutable_rules:
|
|
parts.append("## Session Rules")
|
|
for rule in session.immutable_rules:
|
|
parts.append(f"- {rule}")
|
|
parts.extend(
|
|
[
|
|
"",
|
|
"## Contrato de Contexto",
|
|
"- Los resultados de herramientas se incluyen completos en la conversación.",
|
|
"- Los steps anteriores pueden estar compactados como resúmenes.",
|
|
"- Mantén las respuestas enfocadas en el paso actual.",
|
|
"- Si ya tienes la información necesaria, genera tu respuesta final.",
|
|
"- NO repitas llamadas a herramientas con los mismos argumentos.",
|
|
"- Responde SIEMPRE en español.",
|
|
]
|
|
)
|
|
content = "\n".join(parts)
|
|
return ContextSection(
|
|
section_type=ContextSectionType.IMMUTABLE_RULES,
|
|
content=content,
|
|
priority=100,
|
|
token_estimate=estimate_tokens(content),
|
|
)
|
|
|
|
def _build_project_profile(self, session: SessionState) -> ContextSection:
|
|
if not session.project_profile:
|
|
content = "# Project Profile\nNo project profile configured."
|
|
else:
|
|
lines = ["# Project Profile"]
|
|
for key, value in session.project_profile.items():
|
|
lines.append(f"- **{key}**: {value}")
|
|
content = "\n".join(lines)
|
|
return ContextSection(
|
|
section_type=ContextSectionType.PROJECT_PROFILE,
|
|
content=content,
|
|
priority=80,
|
|
token_estimate=estimate_tokens(content),
|
|
)
|
|
|
|
async def _build_knowledge_base(
|
|
self, session: SessionState
|
|
) -> ContextSection | None:
|
|
"""Load relevant knowledge documents via semantic search.
|
|
|
|
Uses embeddings to find the most relevant docs for the current
|
|
task. Always includes a title index of ALL docs so the agent
|
|
knows what exists and can request more.
|
|
"""
|
|
if not self.memory:
|
|
return None
|
|
|
|
all_docs: list[MemoryDocument] = []
|
|
all_docs.extend(await self.memory.list_documents(
|
|
namespace="knowledge",
|
|
memory_type=MemoryType.DOCUMENT,
|
|
))
|
|
all_docs.extend(await self.memory.list_documents(
|
|
namespace=f"knowledge:{session.session_id}",
|
|
memory_type=MemoryType.DOCUMENT,
|
|
))
|
|
|
|
if not all_docs:
|
|
return None
|
|
|
|
doc_map = {d.memory_id: d for d in all_docs}
|
|
|
|
# Rank docs by semantic similarity
|
|
query = self._build_search_query(session)
|
|
ranked_ids: list[str] = []
|
|
|
|
if query:
|
|
ranked_ids = await self._semantic_rank(query)
|
|
|
|
if not ranked_ids:
|
|
# No embeddings or no task — sort by size (smallest first)
|
|
ranked_ids = [
|
|
d.memory_id
|
|
for d in sorted(all_docs, key=lambda d: len(d.content))
|
|
]
|
|
|
|
# Include ALL docs — 42K tokens fits well within model context (128K)
|
|
max_kb_tokens = 50_000
|
|
token_budget = max_kb_tokens
|
|
full_docs: list[MemoryDocument] = []
|
|
|
|
for doc_id in ranked_ids:
|
|
doc = doc_map.get(doc_id)
|
|
if not doc:
|
|
continue
|
|
doc_tokens = estimate_tokens(doc.content)
|
|
if doc_tokens <= token_budget:
|
|
full_docs.append(doc)
|
|
token_budget -= doc_tokens
|
|
|
|
# Build section — ALWAYS include title index of ALL docs
|
|
included_ids = {d.memory_id for d in full_docs}
|
|
not_included = [d for d in all_docs if d.memory_id not in included_ids]
|
|
|
|
lines = [
|
|
"# Knowledge Base",
|
|
f"_{len(full_docs)} doc(s) loaded in full, "
|
|
f"{len(not_included)} available on request_",
|
|
"",
|
|
]
|
|
|
|
for doc in full_docs:
|
|
lines.append(f"## {doc.title}")
|
|
lines.append(doc.content)
|
|
lines.append("")
|
|
|
|
if not_included:
|
|
lines.append("## Other Available Docs")
|
|
lines.append("_Ask for any of these if you need the full content:_")
|
|
for doc in not_included:
|
|
lines.append(f"- **{doc.title}** ({doc.memory_id}): {doc.summary[:150]}")
|
|
lines.append("")
|
|
|
|
content = "\n".join(lines)
|
|
return ContextSection(
|
|
section_type=ContextSectionType.KNOWLEDGE_BASE,
|
|
content=content,
|
|
priority=60,
|
|
token_estimate=estimate_tokens(content),
|
|
)
|
|
|
|
async def _semantic_rank(self, query: str) -> list[str]:
|
|
"""Rank knowledge docs by cosine similarity to the query."""
|
|
try:
|
|
if not self._embed_service:
|
|
self._embed_service = EmbeddingService()
|
|
|
|
query_embedding = await self._embed_service.embed(query)
|
|
results = await self.memory.search_by_similarity(
|
|
query_embedding=query_embedding,
|
|
namespace="knowledge",
|
|
top_k=50,
|
|
)
|
|
return [doc_id for doc_id, _score in results]
|
|
|
|
except Exception as e:
|
|
logger.warning("Semantic search failed: %s — loading all docs", e)
|
|
return []
|
|
|
|
@staticmethod
|
|
def _build_search_query(session: SessionState) -> str:
|
|
"""Build a natural language query from the current task."""
|
|
if not session.current_task:
|
|
return ""
|
|
parts = [session.current_task.objective]
|
|
step = session.current_task.current_step()
|
|
if step:
|
|
parts.append(step.description)
|
|
if session.current_task.facts_extracted:
|
|
parts.extend(session.current_task.facts_extracted[-5:])
|
|
return " ".join(parts)
|
|
|
|
def _build_task_history(self, session: SessionState) -> ContextSection:
|
|
"""Build a compact summary of past tasks in this session.
|
|
|
|
Each completed task is ~50 tokens instead of hundreds.
|
|
The agent retains awareness of what was done before.
|
|
"""
|
|
lines = [
|
|
"# Session History",
|
|
f"_{len(session.task_history)} previous task(s) in this session_",
|
|
"",
|
|
]
|
|
|
|
for i, entry in enumerate(session.task_history):
|
|
status = entry.get("status", "?")
|
|
objective = entry.get("objective", "")[:100]
|
|
summary = entry.get("summary", "")[:150]
|
|
facts = entry.get("facts", [])
|
|
|
|
lines.append(f"**Task {i + 1}** [{status}]: {objective}")
|
|
if summary:
|
|
lines.append(f" Result: {summary}")
|
|
if facts:
|
|
lines.append(f" Facts: {'; '.join(facts[:5])}")
|
|
# Key structured data (recordNums, sectionIds, etc.)
|
|
key_data = entry.get("key_data", {})
|
|
if key_data:
|
|
kd_parts = []
|
|
for table, nums in key_data.get("tables", {}).items():
|
|
kd_parts.append(f"{table}: records {nums}")
|
|
for page, num in key_data.get("pages", {}).items():
|
|
kd_parts.append(f"page '{page}' = record {num}")
|
|
if key_data.get("sections"):
|
|
kd_parts.append(f"sections: {key_data['sections']}")
|
|
if key_data.get("modules"):
|
|
kd_parts.append(f"modules: {key_data['modules']}")
|
|
if kd_parts:
|
|
lines.append(f" Key data: {'; '.join(kd_parts)}")
|
|
review = entry.get("review", "")
|
|
if review:
|
|
lines.append(f" Review: {review[:100]}")
|
|
lines.append("")
|
|
|
|
content = "\n".join(lines)
|
|
return ContextSection(
|
|
section_type=ContextSectionType.TASK_STATE,
|
|
content=content,
|
|
priority=55, # Below knowledge (60), above artifacts (50)
|
|
token_estimate=estimate_tokens(content),
|
|
)
|
|
|
|
def _build_task_state(self, task: TaskState) -> ContextSection:
|
|
lines = [
|
|
"# Current Task",
|
|
f"**Objective**: {task.objective}",
|
|
f"**Status**: {task.status}",
|
|
f"**Step**: {task.current_step_index + 1}/{len(task.plan)}",
|
|
]
|
|
|
|
current = task.current_step()
|
|
if current:
|
|
lines.extend(
|
|
[
|
|
"",
|
|
"## Current Step",
|
|
f"- Description: {current.description}",
|
|
f"- Agent: {current.agent_role}",
|
|
f"- Status: {current.status}",
|
|
]
|
|
)
|
|
|
|
if task.facts_extracted:
|
|
lines.append("")
|
|
lines.append("## Established Facts")
|
|
for fact in task.facts_extracted[-10:]:
|
|
lines.append(f"- {fact}")
|
|
|
|
if task.constraints:
|
|
lines.append("")
|
|
lines.append("## Constraints")
|
|
for c in task.constraints:
|
|
lines.append(f"- {c}")
|
|
|
|
# Show compacted previous steps results
|
|
compacted_steps = [s for s in task.plan if s.compacted and s.result_summary]
|
|
if compacted_steps:
|
|
lines.append("")
|
|
lines.append("## Previous Steps (compacted)")
|
|
for step in compacted_steps:
|
|
lines.append(f"- [{step.agent_role}] {step.description}: {step.result_summary[:300]}")
|
|
|
|
# Show plan overview (compact)
|
|
if task.plan:
|
|
lines.append("")
|
|
lines.append("## Plan Overview")
|
|
for i, step in enumerate(task.plan):
|
|
marker = "→" if i == task.current_step_index else "·"
|
|
status_label = step.status.value
|
|
compacted_label = " (compacted)" if step.compacted else ""
|
|
lines.append(
|
|
f" {marker} Step {i + 1} [{status_label}{compacted_label}]: {step.description}"
|
|
)
|
|
|
|
content = "\n".join(lines)
|
|
return ContextSection(
|
|
section_type=ContextSectionType.TASK_STATE,
|
|
content=content,
|
|
priority=70,
|
|
token_estimate=estimate_tokens(content),
|
|
)
|
|
|
|
def _build_artifact_memory(
|
|
self, artifacts: list[ArtifactSummary]
|
|
) -> ContextSection:
|
|
content = self.compactor.compact_artifact_summaries(
|
|
artifacts, max_chars=settings.artifact_summary_max_chars
|
|
)
|
|
return ContextSection(
|
|
section_type=ContextSectionType.ARTIFACT_MEMORY,
|
|
content=content,
|
|
priority=50,
|
|
token_estimate=estimate_tokens(content),
|
|
)
|
|
|
|
# ------------------------------------------------------------------
|
|
# Assembly
|
|
# ------------------------------------------------------------------
|
|
|
|
def _assemble_system_prompt(self, sections: list[ContextSection]) -> str:
|
|
"""Combine sections into a single system prompt string."""
|
|
parts: list[str] = []
|
|
order = [
|
|
ContextSectionType.IMMUTABLE_RULES,
|
|
ContextSectionType.PROJECT_PROFILE,
|
|
ContextSectionType.KNOWLEDGE_BASE,
|
|
ContextSectionType.TASK_STATE,
|
|
]
|
|
section_map: dict[ContextSectionType, ContextSection] = {
|
|
s.section_type: s for s in sections
|
|
}
|
|
for st in order:
|
|
if st in section_map:
|
|
parts.append(section_map[st].content)
|
|
return "\n\n---\n\n".join(parts)
|
|
|
|
def _build_messages(
|
|
self,
|
|
session: SessionState,
|
|
conversation: list[dict[str, Any]] | None = None,
|
|
) -> list[dict[str, Any]]:
|
|
"""Build the messages array with real conversation history.
|
|
|
|
Includes the user objective message followed by the full
|
|
assistant/tool conversation — like professional agentic tools.
|
|
"""
|
|
if session.current_task:
|
|
step = session.current_task.current_step()
|
|
if step:
|
|
user_content = (
|
|
f"Execute this step: {step.description}\n"
|
|
f"Overall objective: {session.current_task.objective}"
|
|
)
|
|
else:
|
|
user_content = session.current_task.objective
|
|
else:
|
|
user_content = "Awaiting task assignment."
|
|
|
|
messages: list[dict[str, Any]] = []
|
|
|
|
# Include previous task exchanges as compact conversation history
|
|
if session.task_history:
|
|
history_lines = ["[HISTORIAL DE CONVERSACIÓN ANTERIOR — NO ejecutar de nuevo, solo contexto]"]
|
|
for entry in session.task_history[-10:]:
|
|
objective = entry.get("objective", "")[:200]
|
|
summary = entry.get("summary", "")
|
|
key_data = entry.get("key_data", {})
|
|
tools = entry.get("tools_used", [])
|
|
|
|
history_lines.append(f"Usuario pidió: {objective}")
|
|
if tools:
|
|
history_lines.append(f" Tools usadas: {', '.join(tools[:5])}")
|
|
if key_data:
|
|
kd_parts = []
|
|
for table, nums in key_data.get("tables", {}).items():
|
|
kd_parts.append(f"{table}: records {nums}")
|
|
if key_data.get("sections"):
|
|
kd_parts.append(f"sections: {key_data['sections'][:5]}")
|
|
if key_data.get("modules"):
|
|
kd_parts.append(f"modules: {key_data['modules'][:5]}")
|
|
if kd_parts:
|
|
history_lines.append(f" Datos clave: {'; '.join(kd_parts)}")
|
|
# Extract agent response from summary
|
|
if " → Agent: " in summary:
|
|
agent_part = summary.split(" → Agent: ", 1)[1][:200]
|
|
history_lines.append(f" Resultado: {agent_part}")
|
|
history_lines.append("")
|
|
|
|
messages.append({"role": "user", "content": "\n".join(history_lines)})
|
|
messages.append({"role": "assistant", "content": "Entendido, tengo el contexto del historial. ¿En qué puedo ayudarte ahora?"})
|
|
|
|
# Current user message
|
|
messages.append({"role": "user", "content": user_content})
|
|
|
|
# Append real conversation (assistant messages + tool results from current step)
|
|
if conversation:
|
|
messages.extend(conversation)
|
|
|
|
return messages
|