context_builder.py

"""Build role Agents and provenance-tagged context plans for team sessions."""

from __future__ import annotations

from dataclasses import dataclass, field
from typing import Any

from opencollab.adapters.llm.providers import RESPONSES
from opencollab.adapters.llm.types import model_capabilities
from opencollab.adapters.skills.null_skill_store import NullSkillStore
from opencollab.adapters.tools.base import Tool
from opencollab.adapters.trace import Tracer
from opencollab.application.event_bus import EventBus
from opencollab.application.ports import (
    AskUserPort,
    PermissionPort,
    SafetyPolicyFactory,
    SchedulerPort,
    SkillStorePort,
)
from opencollab.bootstrap.agent_profiles import (
    SingleAgentProfile,
    resolve_agent_profile,
)
from opencollab.bootstrap.config import (
    DEFAULT_TEMPERATURE,
    DEFAULT_THINKING,
    DEFAULT_THINKING_PARAMS,
    DEFAULT_TOP_P,
    resolve_model_route,
)
from opencollab.bootstrap.team_config import RoleConfig, TeamConfig
from opencollab.bootstrap.tool_registry import COORDINATION_TOOL_NAMES, build_tools_for_role
from opencollab.domain.agent import DEFAULT_MAX_TOKENS_PER_STEP, Agent
from opencollab.domain.context import (
    ContextLayer,
    ContextPlan,
    ContextPosition,
    ContextSource,
)
from opencollab.domain.scheduler import DelegationTask
from opencollab.domain.skill import SkillManifest


@dataclass
class SpawnConfig:
    """Shared LLM/runtime config inherited by every spawned agent session."""

    model: str
    provider: str
    api_key: str | None
    base_url: str | None
    llm_timeout: float
    tracer: Tracer | None
    event_bus: EventBus
    permission_policy: PermissionPort | None
    ask_policy: AskUserPort | None = None
    safety_policy_factory: SafetyPolicyFactory | None = None
    # Role-specific sampling settings override these inherited defaults.
    temperature: float = DEFAULT_TEMPERATURE
    top_p: float | None = DEFAULT_TOP_P
    max_output_tokens: int = DEFAULT_MAX_TOKENS_PER_STEP
    context_window: int | None = None
    # Roles may override the inherited thinking policy.
    thinking: bool = DEFAULT_THINKING
    thinking_params: dict | None = None
    wire_protocol: str = "chat_completions"
    reasoning_effort: str | None = None
    llm_max_retries: int = 3
    llm_connect_timeout: float = 30.0
    llm_first_event_timeout: float = 180.0
    llm_stream_idle_timeout: float = 180.0
    provider_error_time_budget: float = 0.0
    llm_stream_chat: bool = False
    thinking_params_explicit: bool | None = None
    auth_token: str | None = field(default=None, repr=False)
    model_routes: dict[str, Any] = field(default_factory=dict, repr=False)
    request_pool: str | None = None
    request_capacity: int | None = None
    logical_model: str | None = None
    workflow_cleanup_timeout: float = 2.0

    def model_settings(self, model: str | None = None) -> dict[str, Any]:
        return resolve_model_route(
            {key: getattr(self, key) for key in (
                "model", "provider", "api_key", "auth_token", "base_url", "wire_protocol",
                "context_window", "model_routes", "request_pool", "request_capacity", "logical_model",
            )},
            model,
        )

    def __post_init__(self) -> None:
        if self.thinking_params_explicit is None:
            self.thinking_params_explicit = self.thinking_params is not None
        self.thinking_params = dict(DEFAULT_THINKING_PARAMS if self.thinking_params is None else self.thinking_params)

    def reasoning_effort_for_role(self, role: RoleConfig) -> str | None:
        capabilities = model_capabilities(self.model_settings(role.model)["model"])
        if role.thinking is False and (
            not self.thinking or capabilities.honors_workflow_thinking_override
        ):
            return None
        return self.reasoning_effort


class ContextBuilder:
    """Assemble role context and resolve tool names for entry and spawned agents.

    System sources supply identity/team/skills/project text; a delegated task is
    user context. Tool schemas are passed separately rather than copied as prose.
    """

    def __init__(
        self,
        team_cfg: TeamConfig,
        cfg: SpawnConfig,
        *,
        skill_store: SkillStorePort | None = None,
        project_context: str | None = None,
        prebuilt_team: bool = False,
    ):
        self._team = team_cfg
        self._cfg = cfg
        # Catalog and dispatcher must share the same store; absence means no skills.
        self._skill_store: SkillStorePort = skill_store or NullSkillStore()
        # Empty project content contributes no startup source.
        self._project_context = project_context or None
        self._prebuilt_team = prebuilt_team

    def build_plan(
        self,
        role_name: str,
        *,
        task: str | None = None,
        context: str = "",
    ) -> ContextPlan:
        """Emit non-empty system sources and an optional delegated user task.

        Do not register empty memory/RAG placeholders without a content producer.
        """
        role = self._team.role_for(role_name)
        sources: list[ContextSource] = []
        # Match workflow composition: profile first, role second. Replacing the
        # profile prompt with a card would change the agent being compared.
        profile = self._profile_for(role)
        if profile is not None:
            sources.append(
                ContextSource(
                    name="profile",
                    layer=ContextLayer.IDENTITY,
                    position=ContextPosition.SYSTEM,
                    content=profile.system_prompt,
                )
            )
        sources.append(
            ContextSource(
                name="identity",
                layer=ContextLayer.IDENTITY,
                position=ContextPosition.SYSTEM,
                content=role.prompt,
            )
        )
        team_section = self._team_section(role_name, role)
        if team_section:
            sources.append(
                ContextSource(
                    name="team",
                    layer=ContextLayer.TEAM,
                    position=ContextPosition.SYSTEM,
                    content=team_section,
                )
            )
        # Expose the catalog only to roles allowed to invoke it.
        if "use_skill" in role.tools:
            catalog = _render_skill_catalog(self._skill_store.list_manifests())
            if catalog:
                sources.append(
                    ContextSource(
                        name="skills",
                        layer=ContextLayer.SKILL,
                        position=ContextPosition.SYSTEM,
                        content=catalog,
                    )
                )
        # SYSTEM position reaches both entry and spawned sessions.
        if self._project_context:
            sources.append(
                ContextSource(
                    name="project",
                    layer=ContextLayer.PROJECT,
                    position=ContextPosition.SYSTEM,
                    content=self._project_context,
                )
            )
        if task is not None:
            sources.append(
                ContextSource(
                    name="task",
                    layer=ContextLayer.TASK,
                    position=ContextPosition.USER_CONTEXT,
                    content=DelegationTask(
                        role=role_name, task=task, context=context
                    ).render(),
                )
            )
        return ContextPlan(sources=tuple(sources))

    def build_agent(
        self,
        role_name: str,
        *,
        scheduler: SchedulerPort | None = None,
        workflow_tool: Tool | None = None,
        authored_workflow_tool: Tool | None = None,
        ask_user_available: bool = False,
        allow_unisolated_shell: bool = False,
        plan: ContextPlan | None = None,
    ) -> Agent:
        """Build an Agent; human availability and host-shell permission are independent.

        See ``build_tools_for_role`` for permission resolution.
        """
        role = self._team.role_for(role_name)
        if plan is None:
            plan = self.build_plan(role_name)
        tools = build_tools_for_role(
            role.tools,
            scheduler=scheduler,
            workflow_tool=workflow_tool,
            authored_workflow_tool=authored_workflow_tool,
            skill_store=self._skill_store,
            ask_user_available=ask_user_available,
            allow_unisolated_shell=allow_unisolated_shell,
            tool_limits=self._tool_limits_for(role),
        )
        cfg = self._cfg
        model_settings = cfg.model_settings(role.model)
        model = model_settings["model"]
        if (
            model_settings["wire_protocol"] == RESPONSES
            and tools
            and not model_capabilities(model).supports_responses_tools
        ):
            raise ValueError(
                f"role {role_name!r} model {model!r} does not support "
                "Responses tools"
            )
        # A role override of 0.0 is meaningful (fully deterministic), so fall
        # back to the global default only when the role left it unset (None).
        temperature = (
            role.temperature if role.temperature is not None else cfg.temperature
        )
        thinking = role.thinking if role.thinking is not None else cfg.thinking
        thinking_params = (
            role.thinking_params
            if role.thinking_params is not None
            else cfg.thinking_params
        )
        return Agent(
            name=role_name,
            system_prompt=plan.system_prompt(),
            tools=tools,
            model=model,
            provider=model_settings["provider"],
            wire_protocol=model_settings["wire_protocol"],
            api_key=model_settings["api_key"],
            auth_token=model_settings["auth_token"],
            base_url=model_settings["base_url"],
            logical_model=model_settings.get("logical_model"),
            request_pool=model_settings.get("request_pool"),
            request_capacity=model_settings.get("request_capacity"),
            temperature=temperature,
            top_p=cfg.top_p,
            max_tokens_per_step=cfg.max_output_tokens,
            context_window=model_settings["context_window"],
            thinking=thinking,
            thinking_params=thinking_params,
            thinking_params_explicit=(
                role.thinking_params is not None or cfg.thinking_params_explicit
            ),
            reasoning_effort=cfg.reasoning_effort_for_role(role),
            llm_max_retries=cfg.llm_max_retries,
            llm_connect_timeout=cfg.llm_connect_timeout,
            llm_first_event_timeout=cfg.llm_first_event_timeout,
            llm_stream_idle_timeout=cfg.llm_stream_idle_timeout,
            llm_stream_chat=cfg.llm_stream_chat,
            provider_error_time_budget=cfg.provider_error_time_budget,
        )

    def _profile_for(self, role: RoleConfig) -> SingleAgentProfile | None:
        """The resolved profile this role is seated under, if it declared one."""
        return resolve_agent_profile(role.profile) if role.profile is not None else None

    def _tool_limits_for(self, role: RoleConfig) -> dict[str, dict[str, int]]:
        """Overlay explicit team limits on the profile's evaluated output caps.

        Taking a profile prompt without its tool caps would change that profile.
        """
        profile = self._profile_for(role)
        declared = self._team.tool_limits or {}
        if profile is None or not profile.tool_limits:
            return dict(declared)
        merged = {
            name: dict(values) for name, values in profile.tool_limits.items()
        }
        for name, values in declared.items():
            merged.setdefault(name, {}).update(values)
        return merged

    def _team_section(self, role_name: str, role: RoleConfig) -> str:
        topo = self._team.topology
        tools = set(role.tools)
        if not (tools & COORDINATION_TOOL_NAMES):
            return ""
        targets = sorted(self._team.roles if topo.allow_all
                         else topo.edges.get(role_name, frozenset()))
        lines = ["## Your team", "", "Configured roles you may address:"]
        lines += [f"- {target}" for target in targets]
        if not targets:
            lines.append("No other configured role is reachable from this role.")
        if self._prebuilt_team:
            lines.append("All configured agents are already seated. New agents cannot be spawned in this run.")
        elif "spawn_agent" in tools:
            lines.append("Use `spawn_agent` to delegate to a role; include the task and relevant context.")
            if topo.allow_all:
                lines.append("The open topology also permits custom role names, which use the generic role settings.")
        if "message_agent" in tools:
            lines.append("Use `message_agent` with to_role to contact an existing teammate asynchronously.")
        if "team_status" in tools:
            lines.append("Use `team_status` to inspect the live roster, agent ids and execution states.")
        return "\n".join(lines)


def _render_skill_catalog(manifests: tuple[SkillManifest, ...]) -> str:
    """Return the skill invocation catalog, or an empty string so assembly omits it."""
    if not manifests:
        return ""
    lines = [
        "## Skills",
        "",
        "Specialized instruction sets you can load on demand. When one matches "
        "the task, call `use_skill(name)` to pull its full instructions into "
        "your context.",
        "",
    ]
    lines += [f"- {m.name}: {m.description}" for m in manifests]
    return "\n".join(lines)


__all__ = ["ContextBuilder", "SpawnConfig"]