Initial commit
This commit is contained in:
Jordan
551
src/context/engine.py
Normal file
551
src/context/engine.py
Normal file
@@ -0,0 +1,551 @@
|
||||
"""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.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
|
||||
# Debug history: last N context builds per session
|
||||
self._history: dict[str, list[dict[str, Any]]] = defaultdict(list)
|
||||
self._max_history = 20
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Public — build context for a model call
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
async def build_context(
|
||||
self,
|
||||
session: SessionState,
|
||||
agent: AgentProfile,
|
||||
artifacts: list[ArtifactSummary] | None = None,
|
||||
working_items: list[dict[str, Any]] | None = None,
|
||||
extra_instructions: str = "",
|
||||
) -> ContextPackage:
|
||||
"""Build a full ContextPackage for the given agent and session."""
|
||||
|
||||
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 state
|
||||
if "task_state" in allowed and session.current_task:
|
||||
sections.append(self._build_task_state(session.current_task))
|
||||
|
||||
# 5. Artifact memory — summarised, never raw
|
||||
if "artifact_memory" in allowed and artifacts:
|
||||
sections.append(self._build_artifact_memory(artifacts))
|
||||
|
||||
# 6. Working context — recent relevant items
|
||||
if "working_context" in allowed:
|
||||
sections.append(
|
||||
self._build_working_context(working_items or [], extra_instructions)
|
||||
)
|
||||
|
||||
# Compact to fit budget
|
||||
sections = self.compactor.compact_sections(sections)
|
||||
|
||||
# Assemble system prompt from sections
|
||||
system_prompt = self._assemble_system_prompt(sections)
|
||||
|
||||
# Build messages (just user message — no chat history)
|
||||
messages = self._build_messages(session)
|
||||
|
||||
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,
|
||||
)
|
||||
|
||||
# --- 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", " "),
|
||||
})
|
||||
|
||||
debug_entry = {
|
||||
"timestamp": time.time(),
|
||||
"agent": agent.role.value,
|
||||
"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,
|
||||
"working_items_count": len(working_items) if working_items else 0,
|
||||
}
|
||||
|
||||
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, working_items=%d",
|
||||
session.session_id[:8],
|
||||
agent.role.value,
|
||||
len(sections),
|
||||
total_tokens,
|
||||
len(artifacts) if artifacts else 0,
|
||||
len(working_items) if working_items else 0,
|
||||
)
|
||||
for s in section_summary:
|
||||
logger.debug(
|
||||
" Section [%s] prio=%d tokens=%d chars=%d",
|
||||
s["type"], s["priority"], s["tokens"], s["chars"],
|
||||
)
|
||||
|
||||
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 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",
|
||||
"- NUNCA recibirás salidas crudas de herramientas en tu contexto.",
|
||||
"- Los resultados de herramientas se resumen como artefactos.",
|
||||
"- Solicita rehidratación si necesitas el contenido completo.",
|
||||
"- Mantén las respuestas enfocadas en el paso actual.",
|
||||
"- 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 from the memory store.
|
||||
|
||||
Uses keyword matching against the task objective and step
|
||||
description to select only the most relevant docs.
|
||||
Max budget: ~15k tokens for knowledge.
|
||||
"""
|
||||
if not self.memory:
|
||||
return None
|
||||
|
||||
# Build search terms from current context
|
||||
search_terms = self._extract_search_terms(session)
|
||||
if not search_terms:
|
||||
# No task → load summaries of all docs (lightweight)
|
||||
return await self._build_knowledge_summaries_only()
|
||||
|
||||
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
|
||||
|
||||
# Score each doc by relevance
|
||||
scored = self._score_docs(all_docs, search_terms)
|
||||
|
||||
# Select top docs within token budget
|
||||
max_kb_tokens = 15_000
|
||||
selected: list[tuple[MemoryDocument, int]] = []
|
||||
token_budget = max_kb_tokens
|
||||
|
||||
for doc, score in scored:
|
||||
if score == 0:
|
||||
continue
|
||||
doc_tokens = estimate_tokens(doc.content)
|
||||
if doc_tokens > token_budget:
|
||||
# Include summary instead of full content
|
||||
summary_tokens = estimate_tokens(doc.summary or doc.title)
|
||||
if summary_tokens < token_budget:
|
||||
selected.append((doc, -1)) # -1 = summary only
|
||||
token_budget -= summary_tokens
|
||||
continue
|
||||
selected.append((doc, score))
|
||||
token_budget -= doc_tokens
|
||||
|
||||
if not selected:
|
||||
return await self._build_knowledge_summaries_only()
|
||||
|
||||
# Build section
|
||||
full_docs = [(d, s) for d, s in selected if s > 0]
|
||||
summary_docs = [(d, s) for d, s in selected if s == -1]
|
||||
|
||||
lines = [
|
||||
"# Knowledge Base",
|
||||
f"_{len(full_docs)} relevant doc(s) loaded, "
|
||||
f"{len(summary_docs)} summarized, "
|
||||
f"{len(all_docs) - len(selected)} filtered out_",
|
||||
"",
|
||||
]
|
||||
|
||||
for doc, _ in full_docs:
|
||||
lines.append(f"## {doc.title}")
|
||||
lines.append(doc.content)
|
||||
lines.append("")
|
||||
|
||||
if summary_docs:
|
||||
lines.append("## Other Available Docs (summaries)")
|
||||
for doc, _ in summary_docs:
|
||||
lines.append(f"- **{doc.title}**: {doc.summary[:200]}")
|
||||
lines.append("")
|
||||
|
||||
content = "\n".join(lines)
|
||||
return ContextSection(
|
||||
section_type=ContextSectionType.KNOWLEDGE_BASE,
|
||||
content=content,
|
||||
priority=60,
|
||||
token_estimate=estimate_tokens(content),
|
||||
)
|
||||
|
||||
async def _build_knowledge_summaries_only(self) -> ContextSection | None:
|
||||
"""Lightweight: only doc titles and summaries (no full content)."""
|
||||
if not self.memory:
|
||||
return None
|
||||
docs = await self.memory.list_documents(
|
||||
namespace="knowledge", memory_type=MemoryType.DOCUMENT
|
||||
)
|
||||
if not docs:
|
||||
return None
|
||||
lines = ["# Knowledge Base (summaries)", ""]
|
||||
for doc in docs:
|
||||
lines.append(f"- **{doc.title}**: {doc.summary[:150]}")
|
||||
content = "\n".join(lines)
|
||||
return ContextSection(
|
||||
section_type=ContextSectionType.KNOWLEDGE_BASE,
|
||||
content=content,
|
||||
priority=60,
|
||||
token_estimate=estimate_tokens(content),
|
||||
)
|
||||
|
||||
def _extract_search_terms(self, session: SessionState) -> set[str]:
|
||||
"""Extract keywords from the current task for doc matching."""
|
||||
terms: set[str] = set()
|
||||
if not session.current_task:
|
||||
return terms
|
||||
|
||||
text = session.current_task.objective.lower()
|
||||
step = session.current_task.current_step()
|
||||
if step:
|
||||
text += " " + step.description.lower()
|
||||
|
||||
# Split into words, filter short/common ones
|
||||
stopwords = {
|
||||
"de", "la", "el", "en", "un", "una", "los", "las", "del", "al",
|
||||
"por", "para", "con", "que", "como", "cómo", "qué", "es", "son",
|
||||
"se", "su", "más", "ya", "si", "no", "este", "esta", "esto",
|
||||
"the", "a", "an", "is", "are", "and", "or", "to", "in", "of",
|
||||
"for", "on", "with", "how", "what", "do", "does", "can",
|
||||
}
|
||||
for word in text.split():
|
||||
word = word.strip(".,;:!?¿¡()[]{}\"'`")
|
||||
if len(word) >= 3 and word not in stopwords:
|
||||
terms.add(word)
|
||||
|
||||
return terms
|
||||
|
||||
@staticmethod
|
||||
def _score_docs(
|
||||
docs: list[MemoryDocument], terms: set[str]
|
||||
) -> list[tuple[MemoryDocument, int]]:
|
||||
"""Score docs by keyword match against title, tags, and content."""
|
||||
scored: list[tuple[MemoryDocument, int]] = []
|
||||
|
||||
for doc in docs:
|
||||
score = 0
|
||||
title_lower = doc.title.lower()
|
||||
tags_lower = " ".join(doc.tags).lower()
|
||||
content_lower = doc.content[:2000].lower()
|
||||
|
||||
for term in terms:
|
||||
# Title match = high weight
|
||||
if term in title_lower:
|
||||
score += 10
|
||||
# Tag match = medium weight
|
||||
if term in tags_lower:
|
||||
score += 5
|
||||
# Content match = low weight
|
||||
if term in content_lower:
|
||||
score += 1
|
||||
|
||||
scored.append((doc, score))
|
||||
|
||||
scored.sort(key=lambda x: x[1], reverse=True)
|
||||
return scored
|
||||
|
||||
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 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
|
||||
lines.append(
|
||||
f" {marker} Step {i + 1} [{status_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),
|
||||
)
|
||||
|
||||
def _build_working_context(
|
||||
self,
|
||||
items: list[dict[str, Any]],
|
||||
extra_instructions: str,
|
||||
) -> ContextSection:
|
||||
lines = ["# Working Context"]
|
||||
if extra_instructions:
|
||||
lines.append(f"\n{extra_instructions}")
|
||||
for item in items[: settings.working_context_max_items]:
|
||||
role = item.get("role", "info")
|
||||
content_val = item.get("content", "")
|
||||
lines.append(f"[{role}] {content_val}")
|
||||
content = "\n".join(lines)
|
||||
return ContextSection(
|
||||
section_type=ContextSectionType.WORKING_CONTEXT,
|
||||
content=content,
|
||||
priority=30,
|
||||
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: rules → profile → task → artifacts → working
|
||||
order = [
|
||||
ContextSectionType.IMMUTABLE_RULES,
|
||||
ContextSectionType.PROJECT_PROFILE,
|
||||
ContextSectionType.KNOWLEDGE_BASE,
|
||||
ContextSectionType.TASK_STATE,
|
||||
ContextSectionType.ARTIFACT_MEMORY,
|
||||
ContextSectionType.WORKING_CONTEXT,
|
||||
]
|
||||
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) -> list[dict[str, Any]]:
|
||||
"""Build the messages array. We do NOT include chat history.
|
||||
|
||||
The user message is the current task objective (or a sentinel
|
||||
if no task is active).
|
||||
"""
|
||||
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."
|
||||
|
||||
return [{"role": "user", "content": user_content}]
|
||||
Reference in New Issue
Block a user