Agent/examples/application_reference/application.py

from __future__ import annotations

from dataclasses import dataclass
from pathlib import Path

from cyber_agent.application import (
    AgentApplication,
    AgentRole,
    ApplicationServices,
    ProviderRef,
    QualityRuleSpec,
    ToolSet,
    ToolSpec,
)
from cyber_agent.tools import get_tool_registry

from .providers import (
    PlaceholderQualityProvider,
    ReferenceArtifactStore,
    ReferenceCandidateRepository,
    ReferenceContextProvider,
    ReferenceProjector,
)
from .tools import FactVerifier


_ROLES = Path(__file__).with_name("roles")


def _prompt(role_id: str) -> str:
    return (_ROLES / f"{role_id}.prompt").read_text(encoding="utf-8").strip()


@dataclass(frozen=True)
class ReferenceComponents:
    application: AgentApplication
    services: ApplicationServices
    artifacts: ReferenceArtifactStore
    candidates: ReferenceCandidateRepository
    context: ReferenceContextProvider
    facts: FactVerifier
    projector: ReferenceProjector


def build_reference_components(
    *,
    artifacts: ReferenceArtifactStore | None = None,
    candidates: ReferenceCandidateRepository | None = None,
    projector: ReferenceProjector | None = None,
) -> ReferenceComponents:
    """Build a fresh process-local binding around optionally persistent backends."""
    artifacts = artifacts or ReferenceArtifactStore()
    candidates = candidates or ReferenceCandidateRepository(artifacts)
    context = ReferenceContextProvider()
    facts = FactVerifier(artifacts)
    projector = projector or ReferenceProjector()
    registry = get_tool_registry().clone()

    async def verify_fact(
        question: str,
        context: dict,
        uid: str = "",
    ) -> dict:
        trace = await context["store"].get_trace(context["trace_id"])
        if trace is None:
            raise ValueError("Fact tool Trace not found")
        root_trace_id = trace.context.get("root_trace_id") or trace.trace_id
        ref = await facts.verify(
            question=question,
            root_trace_id=root_trace_id,
            uid=uid or None,
        )
        return {"artifact_ref": ref.model_dump(mode="json")}

    registry.register(verify_fact, hidden_params=["context", "uid"])
    application = AgentApplication(
        application_id="application_reference",
        application_version="1",
        root_role="editor",
        roles=(
            AgentRole(
                role_id="editor",
                model="reference-model",
                system_prompt=_prompt("editor"),
                tool_sets=("editor-tools",),
                allowed_child_roles=("fact_checker", "writer"),
            ),
            AgentRole(
                role_id="fact_checker",
                model="reference-model",
                system_prompt=_prompt("fact_checker"),
                tool_sets=("fact-tools",),
            ),
            AgentRole(
                role_id="writer",
                model="reference-model",
                system_prompt=_prompt("writer"),
                tool_sets=("writer-tools",),
            ),
        ),
        tool_sets=(
            ToolSet(
                tool_set_id="editor-tools",
                tools=tuple(ToolSpec(
                    tool_name=name,
                    implementation_version="framework-1",
                ) for name in (
                    "agent",
                    "update_task_progress",
                    "review_task_result",
                )),
            ),
            ToolSet(
                tool_set_id="fact-tools",
                tools=(
                    ToolSpec(
                        tool_name="verify_fact",
                        implementation_version="1",
                        side_effect_class="read",
                    ),
                    ToolSpec(
                        tool_name="update_task_progress",
                        implementation_version="framework-1",
                    ),
                    ToolSpec(
                        tool_name="submit_task_report",
                        implementation_version="framework-1",
                    ),
                ),
            ),
            ToolSet(
                tool_set_id="writer-tools",
                tools=tuple(ToolSpec(
                    tool_name=name,
                    implementation_version="framework-1",
                    side_effect_class=("write" if name == "manage_candidate" else "none"),
                ) for name in (
                    "manage_candidate",
                    "update_task_progress",
                    "submit_task_report",
                )),
            ),
        ),
        artifact_resolver_ref=ProviderRef(
            provider_id="reference-artifacts",
            provider_version="1",
        ),
        context_provider_ref=ProviderRef(
            provider_id="reference-context",
            provider_version="1",
        ),
        candidate_repository_ref=ProviderRef(
            provider_id="reference-candidates",
            provider_version="1",
        ),
        quality_provider_ref=ProviderRef(
            provider_id="reference-quality",
            provider_version="1",
        ),
        event_projector_ref=ProviderRef(
            provider_id="reference-projector",
            provider_version="1",
        ),
        quality_rules=(QualityRuleSpec(
            rule_id="no_placeholder",
            version="1",
            criterion="Candidate output contains no placeholder description.",
        ),),
    )
    services = ApplicationServices(
        tool_registry=registry,
        context_provider=context,
        artifact_resolver=artifacts,
        candidate_repository=candidates,
        quality_provider=PlaceholderQualityProvider(),
        event_projector=projector,
    )
    return ReferenceComponents(
        application=application,
        services=services,
        artifacts=artifacts,
        candidates=candidates,
        context=context,
        facts=facts,
        projector=projector,
    )