Newer
Older
navi-1 / navi / core / context_builder.py
@Eugene Sukhodolskiy Eugene Sukhodolskiy on 29 Apr 5 KB Bootstrap test suite — Phase 1 unit tests
"""Build LLM context lists for agent turns.

Extracted from agent.py to reduce the Agent class surface area.
"""

from datetime import datetime, timezone
from typing import TYPE_CHECKING

from navi.llm.base import Message

import navi.config as _config

if TYPE_CHECKING:
    from navi.context_providers._loader import ContextProviderRegistry
    from navi.memory.store import MemoryStore
    from navi.profiles.base import AgentProfile


def render_todo_lines(session_id: str) -> list[str]:
    """Return a list of formatted todo lines for goal anchoring."""
    try:
        from navi.tools.todo import render_todo_lines as _rtl
        return _rtl(session_id)
    except Exception:
        return []


class ContextBuilder:
    """Handles construction of the full LLM context for each turn."""

    def __init__(
        self,
        profile_registry,
        memory_store: "MemoryStore | None" = None,
        cp_registry: "ContextProviderRegistry | None" = None,
    ) -> None:
        self._profiles = profile_registry
        self._memory = memory_store
        self._cp_registry = cp_registry
        self._system_prompt_cache: dict[str, str] = {}

    def build_system_prompt(self, profile: "AgentProfile") -> str:
        """Build the system prompt string for a profile (cached per profile)."""
        cached = self._system_prompt_cache.get(profile.id)
        if cached is not None:
            return cached

        parts: list[str] = []
        persona = _config.settings.navi_persona.strip()
        if persona:
            parts.append(persona)
        parts.append(profile.system_prompt)

        other = [p for p in self._profiles.all() if p.id != profile.id]
        if other:
            lines = [
                "## Available profiles",
                f"Current: **{profile.id}**",
            ]
            for p in other:
                desc = p.short_description or p.description
                lines.append(f"· {p.id}: {desc}")
            lines.append(
                "→ Switch profiles on your own judgment — do not ask for permission. "
                "When a task clearly fits another profile, call switch_profile immediately, "
                "then inform the user which profile is now active and why. "
                "Use list_profiles if you need details about a profile's capabilities."
            )
            parts.append("\n".join(lines))

        result = "\n\n---\n\n".join(parts)
        self._system_prompt_cache[profile.id] = result
        return result

    def invalidate_system_prompt_cache(self, profile_id: str | None = None) -> None:
        if profile_id is None:
            self._system_prompt_cache.clear()
        else:
            self._system_prompt_cache.pop(profile_id, None)

    async def _memory_msg(self) -> "Message | None":
        if self._memory is None:
            return None
        summary = await self._memory.get_summary()
        if not summary:
            return None
        return Message(
            role="system",
            content=f"## What I remember about the user\n\n{summary}",
        )

    async def _collect_context_injections(self, profile: "AgentProfile") -> list[Message]:
        if self._cp_registry is None:
            return []
        out: list[Message] = []
        for provider in self._cp_registry.all():
            if not provider.global_provider and provider.name not in profile.context_providers:
                continue
            try:
                text = await provider.get_context()
                if text:
                    out.append(Message(role="system", content=text))
            except Exception:
                pass
        return out

    def _build_goal_anchor(self, session_id: str, user_message: str) -> Message:
        lines = [
            "[Goal anchor]",
            f"Original request: {user_message}",
        ]
        todo_lines = render_todo_lines(session_id)
        if todo_lines:
            lines.append("Current todo:")
            lines.extend(todo_lines)
        lines.append("Stay on track — complete the remaining pending/in_progress steps.")
        lines.append("Use 1-based todo indexes. Mark completed steps done only after verification, with validation.")
        lines.append("Before final response, update todo for every completed step, including the final one.")
        return Message(role="system", content="\n".join(lines))

    def build(
        self,
        session_context: list[Message],
        profile: "AgentProfile",
        mem: "Message | None",
        iteration: int | None = None,
        max_iterations: int | None = None,
        extra_system: list[Message] | None = None,
    ) -> list[Message]:
        system_msg = Message(
            role="system",
            content=self.build_system_prompt(profile),
        )
        conv = [m for m in session_context if m.role != "system"]
        result: list[Message] = [system_msg]
        if mem:
            result.append(mem)
        if extra_system:
            result.extend(extra_system)
        result.extend(conv)

        if profile.iteration_budget_enabled and iteration is not None and max_iterations is not None:
            remaining = max_iterations - iteration
            if remaining <= 3:
                urgency = (
                    f" CRITICAL: only {remaining} iteration(s) left — finish or produce "
                    "a partial result now, do not start new subtasks."
                )
            elif remaining <= 7:
                urgency = " Start wrapping up: prioritize completing current work over starting new subtasks."
            else:
                urgency = ""
            result.append(Message(
                role="system",
                content=f"[Iteration {iteration + 1}/{max_iterations} — {remaining} remaining.{urgency}]",
            ))

        return result