Source code for design_research_agents.integration

"""Study-facing integration helpers for executable agents."""

from __future__ import annotations

import hashlib
import inspect
import json
from collections.abc import Callable, Mapping, Sequence
from dataclasses import dataclass, field
from datetime import UTC, datetime
from importlib import import_module
from typing import Any

type AgentBinding = Any | Callable[[Any], Any]

_AGENT_EXECUTION_PARAMETER_NAMES = frozenset(
    {
        "prompt",
        "input",
        "request_id",
        "dependencies",
    }
)


[docs] @dataclass(slots=True) class AgentExecutionEnvelope: """Normalized execution envelope used by study-orchestration consumers.""" output: dict[str, Any] = field(default_factory=dict) metrics: dict[str, Any] = field(default_factory=dict) events: list[dict[str, Any]] = field(default_factory=list) trace_refs: list[str] = field(default_factory=list) metadata: dict[str, Any] = field(default_factory=dict)
[docs] @dataclass(slots=True, frozen=True, kw_only=True) class StudyCondition: """Minimal public condition descriptor for study orchestration. Attributes: condition_id: Stable condition identifier. label: Optional human-readable condition label. metadata: Optional condition metadata for downstream analysis. """ condition_id: str """Stable condition identifier.""" label: str | None = None """Optional human-readable condition label.""" metadata: Mapping[str, object] = field(default_factory=dict) """Optional condition metadata.""" def __post_init__(self) -> None: """Normalize condition fields for stable downstream use.""" normalized_condition_id = self.condition_id.strip() if not normalized_condition_id: raise ValueError("condition_id must be non-empty.") normalized_label = self.label.strip() if isinstance(self.label, str) else None object.__setattr__(self, "condition_id", normalized_condition_id) object.__setattr__(self, "label", normalized_label or None) object.__setattr__(self, "metadata", dict(self.metadata))
[docs] @dataclass(slots=True, frozen=True, kw_only=True) class AgentRunRequest: """Typed request object for running one agent inside a study. Attributes: agent_ref: Public agent reference, executable object, or binding key. prompt: Prompt or problem-like payload for the run. request_id: Optional run identifier. dependencies: Optional dependencies exposed to the agent. agent_bindings: Optional mapping of binding keys to executable agents or factories. """ agent_ref: Any """Public agent reference, executable object, or binding key.""" prompt: Any """Prompt or problem-like payload for the run.""" request_id: str | None = None """Optional run identifier.""" dependencies: Mapping[str, object] = field(default_factory=dict) """Optional dependencies exposed to the agent.""" agent_bindings: Mapping[str, AgentBinding] = field(default_factory=dict) """Optional mapping of binding keys to executable agents or factories.""" def __post_init__(self) -> None: """Normalize request metadata mappings.""" normalized_request_id = self.request_id.strip() if isinstance(self.request_id, str) else None object.__setattr__(self, "request_id", normalized_request_id or None) object.__setattr__(self, "dependencies", dict(self.dependencies)) object.__setattr__(self, "agent_bindings", dict(self.agent_bindings))
[docs] def execute_agent_request(request: AgentRunRequest) -> AgentExecutionEnvelope: """Execute one typed agent-run request. Args: request: Agent-run request to execute. Returns: Normalized execution envelope. """ return execute_agent_run( request.agent_ref, prompt=request.prompt, request_id=request.request_id, dependencies=request.dependencies, agent_bindings=request.agent_bindings, )
[docs] def execute_agent_run( agent_ref: Any, *, prompt: Any, request_id: str | None, dependencies: Mapping[str, object] | None, agent_bindings: Mapping[str, AgentBinding] | None = None, ) -> AgentExecutionEnvelope: """Execute one public agent reference through the stable prompt/dependencies contract.""" resolved_dependencies = dict(dependencies or {}) condition = resolved_dependencies.get("condition") executable = _resolve_agent_ref( agent_ref, condition=condition, agent_bindings=agent_bindings, ) raw = _invoke_agent( executable=executable, prompt=prompt, request_id=request_id, dependencies=resolved_dependencies, ) return _normalize_agent_execution(raw=raw, request_id=request_id)
[docs] def normalize_agent_execution( raw: Any, *, request_id: str | None = None, ) -> AgentExecutionEnvelope: """Normalize raw agent output into the study-facing execution envelope.""" return _normalize_agent_execution(raw=raw, request_id=request_id)
def _resolve_agent_ref( agent_ref: Any, *, condition: Any, agent_bindings: Mapping[str, AgentBinding] | None, ) -> Any: """Resolve one agent reference into an executable object.""" if not isinstance(agent_ref, str): return agent_ref if agent_bindings and agent_ref in agent_bindings: return _materialize_agent_binding(agent_bindings[agent_ref], condition) exported = getattr(import_module("design_research_agents"), agent_ref, None) if exported is None: raise ValueError(f"Unknown agent reference '{agent_ref}'. Register an agent binding or use a public export.") if inspect.isclass(exported): return exported() return exported def _materialize_agent_binding(binding: AgentBinding, condition: Any) -> Any: """Resolve one binding into a concrete executable for the current condition.""" if _is_condition_scoped_binding(binding): return binding(condition) return binding def _is_condition_scoped_binding(binding: AgentBinding) -> bool: """Return whether one binding is a condition-to-executable builder.""" if not callable(binding): return False if hasattr(binding, "run") and callable(binding.run): return False try: parameters = inspect.signature(binding).parameters except (TypeError, ValueError): return False if any(name in _AGENT_EXECUTION_PARAMETER_NAMES for name in parameters): return False if any( parameter.kind in (inspect.Parameter.VAR_POSITIONAL, inspect.Parameter.VAR_KEYWORD) for parameter in parameters.values() ): return False positional_parameters = [ parameter for parameter in parameters.values() if parameter.kind in ( inspect.Parameter.POSITIONAL_ONLY, inspect.Parameter.POSITIONAL_OR_KEYWORD, ) ] return len(parameters) == 1 and len(positional_parameters) == 1 def _invoke_agent( *, executable: Any, prompt: Any, request_id: str | None, dependencies: Mapping[str, object], ) -> Any: """Invoke one agent object or callable through the public prompt/dependencies contract.""" if hasattr(executable, "run") and callable(executable.run): return _invoke_callable( callable_obj=executable.run, prompt=prompt, request_id=request_id, dependencies=dependencies, ) if callable(executable): return _invoke_callable( callable_obj=executable, prompt=prompt, request_id=request_id, dependencies=dependencies, ) raise ValueError("Resolved agent object is not executable.") def _invoke_callable( *, callable_obj: Callable[..., Any], prompt: Any, request_id: str | None, dependencies: Mapping[str, object], ) -> Any: """Invoke one callable using the stable agent-entrypoint contract.""" parameters = inspect.signature(callable_obj).parameters kwargs: dict[str, Any] = {} if "prompt" in parameters and parameters["prompt"].kind is not inspect.Parameter.POSITIONAL_ONLY: kwargs["prompt"] = prompt if "input" in parameters and parameters["input"].kind is not inspect.Parameter.POSITIONAL_ONLY: kwargs["input"] = prompt if "request_id" in parameters: kwargs["request_id"] = request_id if "dependencies" in parameters: kwargs["dependencies"] = dependencies if kwargs: return callable_obj(**kwargs) if not parameters: return callable_obj() positional_parameters = [ parameter for parameter in parameters.values() if parameter.kind in ( inspect.Parameter.POSITIONAL_ONLY, inspect.Parameter.POSITIONAL_OR_KEYWORD, ) ] if len(positional_parameters) == 1: return callable_obj(prompt) raise ValueError( "Agent execution target must accept the public prompt/request_id/dependencies " "contract or a single prompt argument." ) def _normalize_agent_execution( *, raw: Any, request_id: str | None, ) -> AgentExecutionEnvelope: """Normalize raw execution output to the canonical study-facing envelope.""" if _is_execution_result(raw): return _normalize_execution_result(raw=raw, request_id=request_id) if isinstance(raw, Mapping): output = dict(raw.get("output", raw.get("outputs", {}))) if not output and "text" in raw: output = {"text": raw["text"]} metrics = dict(raw.get("metrics", {})) trace_refs = [str(value) for value in raw.get("trace_refs", [])] metadata = dict(raw.get("metadata", {})) events = _normalize_events( raw_events=raw.get("events", []), request_id=request_id, output=output, ) return AgentExecutionEnvelope( output=output, metrics=metrics, events=events, trace_refs=trace_refs, metadata=metadata, ) output = {"text": str(raw)} return AgentExecutionEnvelope( output=output, metrics={}, events=_normalize_events(raw_events=[], request_id=request_id, output=output), trace_refs=[], metadata={}, ) def _is_execution_result(raw: Any) -> bool: """Return whether one object looks like a design-research-agents `ExecutionResult`.""" if raw is None or isinstance(raw, Mapping): return False return all(hasattr(raw, attribute) for attribute in ("success", "output", "metadata")) def _normalize_execution_result( *, raw: Any, request_id: str | None, ) -> AgentExecutionEnvelope: """Normalize a design-research-agents `ExecutionResult` into the study-facing envelope.""" output_mapping = _extract_output_mapping(getattr(raw, "output", {})) output = _extract_execution_output(output_mapping) raw_metadata = getattr(raw, "metadata", {}) metadata = dict(raw_metadata) if isinstance(raw_metadata, Mapping) else {} model_response = getattr(raw, "model_response", None) metadata.update(_extract_model_metadata(model_response)) metrics: dict[str, Any] = {} raw_metrics = output_mapping.get("metrics") if isinstance(raw_metrics, Mapping): metrics.update(raw_metrics) _merge_usage_metrics(metrics, model_response) trace_refs = _extract_trace_refs(metadata) raw_events = output_mapping.get("events", []) normalized_raw_events = ( raw_events if isinstance(raw_events, Sequence) and not isinstance(raw_events, (str, bytes)) else [] ) events = _normalize_events( raw_events=normalized_raw_events, request_id=request_id, output=output, ) return AgentExecutionEnvelope( output=output, metrics=metrics, events=events, trace_refs=trace_refs, metadata=metadata, ) def _extract_output_mapping(raw_output: Any) -> dict[str, Any]: """Normalize the raw execution output envelope to a plain mapping.""" if isinstance(raw_output, Mapping): return dict(raw_output) return {} def _extract_execution_output(output_mapping: Mapping[str, Any]) -> dict[str, Any]: """Resolve the canonical output payload from a workflow-style envelope.""" final_output = output_mapping.get("final_output") if isinstance(final_output, Mapping): return dict(final_output) if isinstance(final_output, str): return {"text": final_output} if final_output is not None: return {"final_output": final_output} if "text" in output_mapping: return {"text": str(output_mapping["text"])} if "model_text" in output_mapping: return {"text": str(output_mapping["model_text"])} return dict(output_mapping) def _merge_usage_metrics(metrics: dict[str, Any], model_response: Any) -> None: """Merge token-usage metadata from an optional model response.""" usage = getattr(model_response, "usage", None) if isinstance(usage, Mapping): prompt_tokens = usage.get("prompt_tokens") completion_tokens = usage.get("completion_tokens") total_tokens = usage.get("total_tokens") else: prompt_tokens = getattr(usage, "prompt_tokens", None) completion_tokens = getattr(usage, "completion_tokens", None) total_tokens = getattr(usage, "total_tokens", None) if isinstance(prompt_tokens, int): metrics.setdefault("input_tokens", prompt_tokens) if isinstance(completion_tokens, int): metrics.setdefault("output_tokens", completion_tokens) if isinstance(total_tokens, int): metrics.setdefault("total_tokens", total_tokens) def _extract_model_metadata(model_response: Any) -> dict[str, Any]: """Extract stable model metadata from an optional LLM response.""" metadata: dict[str, Any] = {} model_name = getattr(model_response, "model", None) if isinstance(model_name, str) and model_name.strip(): metadata["model_name"] = model_name model_provider = getattr(model_response, "provider", None) if isinstance(model_provider, str) and model_provider.strip(): metadata["model_provider"] = model_provider return metadata def _extract_trace_refs(metadata: Mapping[str, Any]) -> list[str]: """Extract canonical trace references from execution metadata.""" trace_refs: list[str] = [] trace_path = metadata.get("trace_path") if isinstance(trace_path, str) and trace_path.strip(): trace_refs.append(trace_path) raw_trace_refs = metadata.get("trace_refs") if isinstance(raw_trace_refs, Sequence) and not isinstance(raw_trace_refs, (str, bytes)): for value in raw_trace_refs: if isinstance(value, str) and value.strip() and value not in trace_refs: trace_refs.append(value) return trace_refs def _normalize_events( *, raw_events: Sequence[Any], request_id: str | None, output: Mapping[str, Any], ) -> list[dict[str, Any]]: """Normalize raw event payloads into canonical event dictionaries.""" events: list[dict[str, Any]] = [] for index, raw_event in enumerate(raw_events): if not isinstance(raw_event, Mapping): continue meta_json = raw_event.get("meta_json", {}) if not isinstance(meta_json, Mapping): meta_json = {"value": meta_json} event_payload: dict[str, Any] = { "timestamp": str(raw_event.get("timestamp", _utc_now_iso())), "record_id": str( raw_event.get( "record_id", _hash_identifier( "evt", { "request_id": request_id, "index": index, "event_type": raw_event.get("event_type", "event"), }, ), ) ), "text": str(raw_event.get("text", "")), "session_id": str(raw_event.get("session_id", request_id or "")), "actor_id": str(raw_event.get("actor_id", "agent")), "event_type": str(raw_event.get("event_type", "event")), "meta_json": dict(meta_json), } for key in ("level", "trial_id", "step_id", "tool_name", "evaluation_id", "run_id"): if key in raw_event: event_payload[key] = raw_event[key] events.append(event_payload) if events: return events return [ { "timestamp": _utc_now_iso(), "record_id": _hash_identifier("evt", {"request_id": request_id, "index": 0}), "text": str(output.get("text", "")), "session_id": str(request_id or ""), "actor_id": "agent", "event_type": "assistant_output", "meta_json": {"auto_generated": True}, "level": "step", } ] def _utc_now_iso() -> str: """Return the current UTC timestamp in ISO-8601 format.""" return datetime.now(UTC).isoformat() def _hash_identifier(prefix: str, payload: Mapping[str, Any]) -> str: """Build one deterministic identifier for normalized event records.""" encoded = json.dumps(payload, sort_keys=True, separators=(",", ":"), default=str).encode("utf-8") digest = hashlib.sha1(encoded).hexdigest()[:12] return f"{prefix}-{digest}" __all__ = [ "AgentBinding", "AgentExecutionEnvelope", "AgentRunRequest", "StudyCondition", "execute_agent_request", "execute_agent_run", "normalize_agent_execution", ]