|
@@ -0,0 +1,358 @@
|
|
|
|
|
+"""Application registration, immutable binding, and AgentRunner assembly."""
|
|
|
|
|
+
|
|
|
|
|
+from __future__ import annotations
|
|
|
|
|
+
|
|
|
|
|
+from dataclasses import dataclass
|
|
|
|
|
+from types import MappingProxyType
|
|
|
|
|
+from typing import Any, Callable, Mapping
|
|
|
|
|
+
|
|
|
|
|
+from cyber_agent.application.models import (
|
|
|
|
|
+ AgentApplication,
|
|
|
|
|
+ AgentRole,
|
|
|
|
|
+ ApplicationRef,
|
|
|
|
|
+ RunLimits,
|
|
|
|
|
+ application_manifest,
|
|
|
|
|
+ role_manifest,
|
|
|
|
|
+ stable_hash,
|
|
|
|
|
+)
|
|
|
|
|
+from cyber_agent.application.ports import TaskContextProvider
|
|
|
|
|
+from cyber_agent.core.artifacts import ArtifactResolver
|
|
|
|
|
+from cyber_agent.core.validation import ValidationPolicy
|
|
|
|
|
+from cyber_agent.tools.registry import ToolRegistry
|
|
|
|
|
+
|
|
|
|
|
+
|
|
|
|
|
+@dataclass(frozen=True)
|
|
|
|
|
+class ApplicationServices:
|
|
|
|
|
+ """Non-serializable providers and private tool implementations."""
|
|
|
|
|
+
|
|
|
|
|
+ tool_registry: ToolRegistry
|
|
|
|
|
+ context_provider: TaskContextProvider | None = None
|
|
|
|
|
+ artifact_resolver: ArtifactResolver | None = None
|
|
|
|
|
+ candidate_repository: Any | None = None
|
|
|
|
|
+ quality_provider: Any | None = None
|
|
|
|
|
+ event_projector: Any | None = None
|
|
|
|
|
+ runtime_validator: Callable[[], None] | None = None
|
|
|
|
|
+
|
|
|
|
|
+
|
|
|
|
|
+@dataclass(frozen=True)
|
|
|
|
|
+class RoleBinding:
|
|
|
|
|
+ role: AgentRole
|
|
|
|
|
+ role_hash: str
|
|
|
|
|
+ tool_names: tuple[str, ...]
|
|
|
|
|
+ effective_limits: RunLimits
|
|
|
|
|
+ system_prompt: str
|
|
|
|
|
+
|
|
|
|
|
+
|
|
|
|
|
+@dataclass(frozen=True)
|
|
|
|
|
+class ApplicationBinding:
|
|
|
|
|
+ application: AgentApplication
|
|
|
|
|
+ application_ref: ApplicationRef
|
|
|
|
|
+ roles: Mapping[str, RoleBinding]
|
|
|
|
|
+ tool_registry: ToolRegistry
|
|
|
|
|
+ services: ApplicationServices
|
|
|
|
|
+ validation_policy: ValidationPolicy
|
|
|
|
|
+
|
|
|
|
|
+ def role(self, role_id: str) -> RoleBinding:
|
|
|
|
|
+ try:
|
|
|
|
|
+ return self.roles[role_id]
|
|
|
|
|
+ except KeyError as exc:
|
|
|
|
|
+ raise ValueError(f"Unknown application role: {role_id}") from exc
|
|
|
|
|
+
|
|
|
|
|
+ def configure_run_config(self, config: Any, role_id: str) -> RoleBinding:
|
|
|
|
|
+ """Apply an already-bound role to an existing RunConfig instance."""
|
|
|
|
|
+ binding = self.role(role_id)
|
|
|
|
|
+ role = binding.role
|
|
|
|
|
+ config.model = role.model
|
|
|
|
|
+ config.temperature = role.temperature
|
|
|
|
|
+ config.max_iterations = binding.effective_limits.max_iterations
|
|
|
|
|
+ config.extra_llm_params = dict(role.model_parameters)
|
|
|
|
|
+ config.tools = list(binding.tool_names)
|
|
|
|
|
+ config.tool_groups = []
|
|
|
|
|
+ config.exclude_tools = []
|
|
|
|
|
+ config.agent_type = role.role_id
|
|
|
|
|
+ config.system_prompt = binding.system_prompt
|
|
|
|
|
+ config.skills = []
|
|
|
|
|
+ config.max_parallel_children = binding.effective_limits.max_parallel_children
|
|
|
|
|
+ config.application_ref = self.application_ref
|
|
|
|
|
+ config.role_id = role.role_id
|
|
|
|
|
+ config.role_hash = binding.role_hash
|
|
|
|
|
+ config.effective_run_limits = binding.effective_limits.model_dump(mode="json")
|
|
|
|
|
+ return binding
|
|
|
|
|
+
|
|
|
|
|
+
|
|
|
|
|
+class ApplicationRegistry:
|
|
|
|
|
+ """Explicit, in-memory registry keyed by an exact application version."""
|
|
|
|
|
+
|
|
|
|
|
+ def __init__(self) -> None:
|
|
|
|
|
+ self._bindings: dict[tuple[str, str], ApplicationBinding] = {}
|
|
|
|
|
+
|
|
|
|
|
+ def register(
|
|
|
|
|
+ self,
|
|
|
|
|
+ application: AgentApplication,
|
|
|
|
|
+ services: ApplicationServices,
|
|
|
|
|
+ ) -> ApplicationBinding:
|
|
|
|
|
+ key = (application.application_id, application.application_version)
|
|
|
|
|
+ if key in self._bindings:
|
|
|
|
|
+ raise ValueError(
|
|
|
|
|
+ f"Application is already registered: {application.application_id}@"
|
|
|
|
|
+ f"{application.application_version}"
|
|
|
|
|
+ )
|
|
|
|
|
+ self._validate_services(application, services)
|
|
|
|
|
+ selected_specs = [
|
|
|
|
|
+ spec
|
|
|
|
|
+ for tool_set in application.tool_sets
|
|
|
|
|
+ for spec in tool_set.tools
|
|
|
|
|
+ ]
|
|
|
|
|
+ tool_names = [item.tool_name for item in selected_specs]
|
|
|
|
|
+ if len(tool_names) != len(set(tool_names)):
|
|
|
|
|
+ by_name: dict[str, Any] = {}
|
|
|
|
|
+ for spec in selected_specs:
|
|
|
|
|
+ prior = by_name.setdefault(spec.tool_name, spec)
|
|
|
|
|
+ if prior != spec:
|
|
|
|
|
+ raise ValueError(
|
|
|
|
|
+ f"Conflicting declarations for tool: {spec.tool_name}"
|
|
|
|
|
+ )
|
|
|
|
|
+ definitions = services.tool_registry.export_definitions(
|
|
|
|
|
+ sorted(set(tool_names))
|
|
|
|
|
+ )
|
|
|
|
|
+ private_registry = services.tool_registry.clone(
|
|
|
|
|
+ sorted(set(tool_names)),
|
|
|
|
|
+ freeze=True,
|
|
|
|
|
+ )
|
|
|
|
|
+ manifest = application_manifest(application, definitions)
|
|
|
|
|
+ application_ref = ApplicationRef(
|
|
|
|
|
+ application_id=application.application_id,
|
|
|
|
|
+ application_version=application.application_version,
|
|
|
|
|
+ config_hash=stable_hash(manifest),
|
|
|
|
|
+ )
|
|
|
|
|
+ definition_map = {item["tool_name"]: item for item in definitions}
|
|
|
|
|
+ role_bindings: dict[str, RoleBinding] = {}
|
|
|
|
|
+ for role in application.roles:
|
|
|
|
|
+ role_tools = tuple(dict.fromkeys(
|
|
|
|
|
+ spec.tool_name
|
|
|
|
|
+ for tool_set_id in role.tool_sets
|
|
|
|
|
+ for spec in application.tool_set(tool_set_id).tools
|
|
|
|
|
+ ))
|
|
|
|
|
+ role_limits = application.run_limits.tighten(role.run_limits)
|
|
|
|
|
+ role_bindings[role.role_id] = RoleBinding(
|
|
|
|
|
+ role=role,
|
|
|
|
|
+ role_hash=stable_hash(
|
|
|
|
|
+ role_manifest(application, role, definition_map)
|
|
|
|
|
+ ),
|
|
|
|
|
+ tool_names=role_tools,
|
|
|
|
|
+ effective_limits=role_limits,
|
|
|
|
|
+ system_prompt=self._compose_system_prompt(role),
|
|
|
|
|
+ )
|
|
|
|
|
+ validation_policy = (
|
|
|
|
|
+ ValidationPolicy.model_validate(application.validation_policy)
|
|
|
|
|
+ if application.validation_policy
|
|
|
|
|
+ else ValidationPolicy()
|
|
|
|
|
+ )
|
|
|
|
|
+ binding = ApplicationBinding(
|
|
|
|
|
+ application=application,
|
|
|
|
|
+ application_ref=application_ref,
|
|
|
|
|
+ roles=MappingProxyType(role_bindings),
|
|
|
|
|
+ tool_registry=private_registry,
|
|
|
|
|
+ services=services,
|
|
|
|
|
+ validation_policy=validation_policy,
|
|
|
|
|
+ )
|
|
|
|
|
+ self._bindings[key] = binding
|
|
|
|
|
+ return binding
|
|
|
|
|
+
|
|
|
|
|
+ def resolve(
|
|
|
|
|
+ self,
|
|
|
|
|
+ application_id: str,
|
|
|
|
|
+ application_version: str,
|
|
|
|
|
+ ) -> ApplicationBinding:
|
|
|
|
|
+ key = (application_id, application_version)
|
|
|
|
|
+ try:
|
|
|
|
|
+ return self._bindings[key]
|
|
|
|
|
+ except KeyError as exc:
|
|
|
|
|
+ raise ValueError(
|
|
|
|
|
+ f"Application is not registered: {application_id}@{application_version}"
|
|
|
|
|
+ ) from exc
|
|
|
|
|
+
|
|
|
|
|
+ def resolve_ref(
|
|
|
|
|
+ self,
|
|
|
|
|
+ application_ref: ApplicationRef | Mapping[str, Any],
|
|
|
|
|
+ ) -> ApplicationBinding:
|
|
|
|
|
+ if not isinstance(application_ref, ApplicationRef):
|
|
|
|
|
+ application_ref = ApplicationRef.model_validate(application_ref)
|
|
|
|
|
+ binding = self.resolve(
|
|
|
|
|
+ application_ref.application_id,
|
|
|
|
|
+ application_ref.application_version,
|
|
|
|
|
+ )
|
|
|
|
|
+ if binding.application_ref != application_ref:
|
|
|
|
|
+ raise ValueError("Application config hash does not match the registry")
|
|
|
|
|
+ return binding
|
|
|
|
|
+
|
|
|
|
|
+ @staticmethod
|
|
|
|
|
+ def _compose_system_prompt(role: AgentRole) -> str:
|
|
|
|
|
+ sections = [role.system_prompt.strip()]
|
|
|
|
|
+ if role.skills:
|
|
|
|
|
+ sections.append("# Application Skills")
|
|
|
|
|
+ sections.extend(
|
|
|
|
|
+ f"## {skill.name}@{skill.version}\n{skill.content.strip()}"
|
|
|
|
|
+ for skill in role.skills
|
|
|
|
|
+ )
|
|
|
|
|
+ return "\n\n".join(sections)
|
|
|
|
|
+
|
|
|
|
|
+ @staticmethod
|
|
|
|
|
+ def _validate_services(
|
|
|
|
|
+ application: AgentApplication,
|
|
|
|
|
+ services: ApplicationServices,
|
|
|
|
|
+ ) -> None:
|
|
|
|
|
+ requirements = (
|
|
|
|
|
+ (application.context_provider_ref, services.context_provider, "context_provider"),
|
|
|
|
|
+ (application.artifact_resolver_ref, services.artifact_resolver, "artifact_resolver"),
|
|
|
|
|
+ (application.candidate_repository_ref, services.candidate_repository, "candidate_repository"),
|
|
|
|
|
+ (application.quality_provider_ref, services.quality_provider, "quality_provider"),
|
|
|
|
|
+ (application.event_projector_ref, services.event_projector, "event_projector"),
|
|
|
|
|
+ )
|
|
|
|
|
+ missing = [name for ref, service, name in requirements if ref and service is None]
|
|
|
|
|
+ if missing:
|
|
|
|
|
+ raise ValueError(f"Application runtime services are missing: {missing}")
|
|
|
|
|
+
|
|
|
|
|
+
|
|
|
|
|
+class ApplicationRuntime:
|
|
|
|
|
+ """Assemble the existing AgentRunner for new and restored application runs."""
|
|
|
|
|
+
|
|
|
|
|
+ def __init__(
|
|
|
|
|
+ self,
|
|
|
|
|
+ *,
|
|
|
|
|
+ registry: ApplicationRegistry,
|
|
|
|
|
+ trace_store: Any,
|
|
|
|
|
+ llm_call: Callable[..., Any],
|
|
|
|
|
+ utility_llm_call: Callable[..., Any] | None = None,
|
|
|
|
|
+ ) -> None:
|
|
|
|
|
+ self.registry = registry
|
|
|
|
|
+ self.trace_store = trace_store
|
|
|
|
|
+ self.llm_call = llm_call
|
|
|
|
|
+ self.utility_llm_call = utility_llm_call
|
|
|
|
|
+
|
|
|
|
|
+ def build_runner(self, binding: ApplicationBinding):
|
|
|
|
|
+ from cyber_agent.core.runner import AgentRunner
|
|
|
|
|
+
|
|
|
|
|
+ return AgentRunner(
|
|
|
|
|
+ trace_store=self.trace_store,
|
|
|
|
|
+ tool_registry=binding.tool_registry,
|
|
|
|
|
+ llm_call=self.llm_call,
|
|
|
|
|
+ utility_llm_call=self.utility_llm_call,
|
|
|
|
|
+ validation_policy=binding.validation_policy,
|
|
|
|
|
+ artifact_resolver=binding.services.artifact_resolver,
|
|
|
|
|
+ application_binding=binding,
|
|
|
|
|
+ context_provider=binding.services.context_provider,
|
|
|
|
|
+ )
|
|
|
|
|
+
|
|
|
|
|
+ def new_run(
|
|
|
|
|
+ self,
|
|
|
|
|
+ application_id: str,
|
|
|
|
|
+ application_version: str,
|
|
|
|
|
+ *,
|
|
|
|
|
+ uid: str | None = None,
|
|
|
|
|
+ name: str | None = None,
|
|
|
|
|
+ root_task_anchor: Any = None,
|
|
|
|
|
+ max_iterations: int | None = None,
|
|
|
|
|
+ ):
|
|
|
|
|
+ from cyber_agent.core.runner import RunConfig
|
|
|
|
|
+
|
|
|
|
|
+ binding = self.registry.resolve(application_id, application_version)
|
|
|
|
|
+ self._validate_runtime_services(binding)
|
|
|
|
|
+ role = binding.role(binding.application.root_role)
|
|
|
|
|
+ if max_iterations is not None and max_iterations > role.effective_limits.max_iterations:
|
|
|
|
|
+ raise ValueError(
|
|
|
|
|
+ "max_iterations may only tighten the application run limit"
|
|
|
|
|
+ )
|
|
|
|
|
+ config = RunConfig(
|
|
|
|
|
+ uid=uid,
|
|
|
|
|
+ name=name,
|
|
|
|
|
+ root_task_anchor=root_task_anchor,
|
|
|
|
|
+ enable_research_flow=False,
|
|
|
|
|
+ )
|
|
|
|
|
+ binding.configure_run_config(config, role.role.role_id)
|
|
|
|
|
+ effective = self._compile_deployment_limits(role.effective_limits)
|
|
|
|
|
+ config.effective_run_limits = effective.model_dump(mode="json")
|
|
|
|
|
+ config.max_iterations = effective.max_iterations
|
|
|
|
|
+ config.max_parallel_children = effective.max_parallel_children
|
|
|
|
|
+ if max_iterations is not None:
|
|
|
|
|
+ config.max_iterations = max_iterations
|
|
|
|
|
+ config.effective_run_limits = {
|
|
|
|
|
+ **config.effective_run_limits,
|
|
|
|
|
+ "max_iterations": max_iterations,
|
|
|
|
|
+ }
|
|
|
|
|
+ return self.build_runner(binding), config
|
|
|
|
|
+
|
|
|
|
|
+ @staticmethod
|
|
|
|
|
+ def _compile_deployment_limits(limits: RunLimits) -> RunLimits:
|
|
|
|
|
+ from cyber_agent.core.resource_budget import ResourceBudget
|
|
|
|
|
+
|
|
|
|
|
+ deployment = ResourceBudget.from_environment()
|
|
|
|
|
+ return limits.model_copy(update={
|
|
|
|
|
+ "max_total_agents": min(
|
|
|
|
|
+ limits.max_total_agents,
|
|
|
|
|
+ deployment.max_total_agents,
|
|
|
|
|
+ ),
|
|
|
|
|
+ "max_llm_calls": min(
|
|
|
|
|
+ limits.max_llm_calls,
|
|
|
|
|
+ deployment.max_llm_calls,
|
|
|
|
|
+ ),
|
|
|
|
|
+ "max_total_tokens": min(
|
|
|
|
|
+ limits.max_total_tokens,
|
|
|
|
|
+ deployment.max_total_tokens,
|
|
|
|
|
+ ),
|
|
|
|
|
+ "max_total_cost_usd": min(
|
|
|
|
|
+ limits.max_total_cost_usd,
|
|
|
|
|
+ deployment.max_total_cost_usd,
|
|
|
|
|
+ ),
|
|
|
|
|
+ "max_duration_seconds": min(
|
|
|
|
|
+ limits.max_duration_seconds,
|
|
|
|
|
+ deployment.max_duration_seconds,
|
|
|
|
|
+ ),
|
|
|
|
|
+ "max_validation_tool_calls": min(
|
|
|
|
|
+ limits.max_validation_tool_calls,
|
|
|
|
|
+ deployment.max_validation_tool_calls,
|
|
|
|
|
+ ),
|
|
|
|
|
+ "max_validation_material_chars": min(
|
|
|
|
|
+ limits.max_validation_material_chars,
|
|
|
|
|
+ deployment.max_validation_material_chars,
|
|
|
|
|
+ ),
|
|
|
|
|
+ })
|
|
|
|
|
+
|
|
|
|
|
+ async def restore(self, trace_id: str):
|
|
|
|
|
+ from cyber_agent.core.run_snapshot import (
|
|
|
|
|
+ RunConfigSnapshotV2,
|
|
|
|
|
+ load_run_config_snapshot,
|
|
|
|
|
+ )
|
|
|
|
|
+ from cyber_agent.core.runner import RunConfig
|
|
|
|
|
+
|
|
|
|
|
+ trace = await self.trace_store.get_trace(trace_id)
|
|
|
|
|
+ if trace is None:
|
|
|
|
|
+ raise ValueError(f"Trace not found: {trace_id}")
|
|
|
|
|
+ snapshot = load_run_config_snapshot(trace.context)
|
|
|
|
|
+ if not isinstance(snapshot, RunConfigSnapshotV2):
|
|
|
|
|
+ raise ValueError("RunConfigSnapshotV1 is not an application run")
|
|
|
|
|
+ binding = self.registry.resolve_ref(snapshot.application_ref)
|
|
|
|
|
+ self._validate_runtime_services(binding)
|
|
|
|
|
+ role = binding.role(snapshot.role_id)
|
|
|
|
|
+ if role.role_hash != snapshot.role_hash:
|
|
|
|
|
+ raise ValueError("Application role hash does not match the registry")
|
|
|
|
|
+ config = RunConfig(trace_id=trace_id)
|
|
|
|
|
+ binding.configure_run_config(config, snapshot.role_id)
|
|
|
|
|
+ config.apply_snapshot(snapshot)
|
|
|
|
|
+ config.system_prompt = role.system_prompt
|
|
|
|
|
+ config.skills = []
|
|
|
|
|
+ runner = self.build_runner(binding)
|
|
|
|
|
+ if snapshot.uid != trace.uid or snapshot.agent_type != (
|
|
|
|
|
+ trace.agent_type or "default"
|
|
|
|
|
+ ):
|
|
|
|
|
+ raise ValueError(
|
|
|
|
|
+ "Application snapshot identity does not match Trace metadata"
|
|
|
|
|
+ )
|
|
|
|
|
+ runner._validate_application_snapshot(snapshot, trace)
|
|
|
|
|
+ return runner, config
|
|
|
|
|
+
|
|
|
|
|
+ @staticmethod
|
|
|
|
|
+ def _validate_runtime_services(binding: ApplicationBinding) -> None:
|
|
|
|
|
+ validator = binding.services.runtime_validator
|
|
|
|
|
+ if validator is not None:
|
|
|
|
|
+ validator()
|