Improve agent trace generation with OpenSimula, dynamic initial context generation and monitoring

DESIGN.md

@@ -31,7 +31,7 @@ The library is designed around a few core concepts:
 - **Training Version Compatibility Guards**: Demo training requirements constrain `transformers` and `trl` to a compatible range (`transformers>=4.56.2,<5.0.0`, `trl>=0.29.1,<0.30.0`) to avoid runtime API mismatches during `SFTConfig` import and trainer startup. The same stack is declared as the `training` optional extra in `pyproject.toml` (aligned with `examples/demo_ui/training_scripts/requirements.txt`); install with `pip install -e ".[training]"` or `uv sync --extra training` so the Gradio-launched `train.py` subprocess can import `trl` and related packages.
 - **Environment Template Files**: Repository root includes a minimal `.env.example` and local `.env` template for demo runtime and training credentials (`GEMINI_API_KEY`, `DEEPSEEK_API_KEY`, `HF_TOKEN`, optional `HF_HUB_DISABLE_XET`); `.env` is gitignored to keep secrets out of version control.
 - **CLI Interface**: The `afterimage` command provides `generate`, `validate`, and `export` subcommands. Generation is driven by YAML config files that map to Pydantic models in `config.py`. The `export` command converts datasets to ShareGPT, Alpaca, or HuggingFace messages formats.
-- **Agent Trace Dataset Generation**: The `afterimage.agent_trace` subpackage provides environment-free synthetic agent-trace dataset generation combining ESAT methodology with a sub-millisecond local Declarative Tool Simulation Framework and dual-mode observation synthesis (`observation_mode: Literal["faker", "llm"]`, defaulting to `"llm"` as preferred production mode and `"faker"` as experimental). Key components include `SchemaArchitect` (LLM Pydantic response schema generator with static AST verification feedback loop), `SchemaVerifier` (6 structural invariant checks), `DeclarativeEngine` (4-tier fallback generator with parameter echoing annotations `param:<name>`/`echo` and semantic synonym matching), `LLMObservationSynthesizer` (LLM-driven structured observation generation as in original ESAT methodology), `DeclarativeEnvironment` & `DeclarativeTool` (dual-mode observation dispatching + stateful `SimulationContext` persistent entity & account balance deduction tables), `GridTaskSynthesizer` (360-bucket grid + `InverseFrequencySampler` + initial state context synthesizer), `ReActTrajectoryLoop` (multi-turn teacher execution against local tools), `TrajectoryJudge` (9-point LLM quality rubric), and `AsyncAgentTraceGenerator` facade. Explicit Pydantic response classes can be attached to `ToolActionSpec` via `response_model_cls` to bypass LLM schema generation. All subpackage components strictly align with the `afterimage.providers.llm_providers.LLMProvider` interface (`agenerate_content`, `agenerate_structured` with `schema`, and `astart_chat`). Model defaults follow `gemini-3.5-flash-lite` for execution/synthesis and `gemini-3.6-flash` for schema architecture and trajectory judging.
+- **Agent Trace Dataset Generation**: The `afterimage.agent_trace` subpackage provides environment-free synthetic agent-trace dataset generation combining ESAT and Simula methodology with a sub-millisecond local Declarative Tool Simulation Framework and dual-mode observation synthesis (`observation_mode: Literal["faker", "llm"]`, defaulting to `"llm"` as preferred production mode and `"faker"` as experimental). Key components include `BaseContextGenerator` hierarchy (`VirtualUserContextGenerator` identity/seed context generation using Faker, `PersonaContextGenerator` for `PersonaEntry` integration, `CallableContextGenerator`, `CompositeContextGenerator`), `SimulaTaskSynthesizer` (`task_synthesis_mode: Literal["grid", "simula"]` for deep OpenSimula factor taxonomy task synthesis), `SchemaArchitect` (LLM Pydantic response schema generator with static AST verification feedback loop), `SchemaVerifier` (6 structural invariant checks), `DeclarativeEngine` (4-tier fallback generator with parameter echoing annotations `param:<name>`/`echo` and semantic synonym matching), `LLMObservationSynthesizer` (LLM-driven structured observation generation as in original ESAT methodology), `DeclarativeEnvironment` & `DeclarativeTool` (dual-mode observation dispatching + stateful `SimulationContext` persistent entity & account balance deduction tables), `GridTaskSynthesizer` (360-bucket grid + `InverseFrequencySampler` + dynamic initial state context synthesizer), `ReActTrajectoryLoop` (multi-turn teacher execution against local tools), `TrajectoryJudge` (9-point LLM quality rubric), structured tool calling exporters in `exporters.py` (`to_openai_tools`, `to_anthropic_tools`, `to_hermes_tools`, `to_agent_dpo`), terminal progress bar indicators (`show_progress=True`, `progress_callback`), and `AsyncAgentTraceGenerator` facade. All public classes and methods feature Google-style Python docstrings for Sphinx documentation compatibility. Explicit Pydantic response classes can be attached to `ToolActionSpec` via `response_model_cls` to bypass LLM schema generation. All subpackage components strictly align with the `afterimage.providers.llm_providers.LLMProvider` interface (`agenerate_content`, `agenerate_structured` with `schema`, and `astart_chat`). Model defaults follow `gemini-3.5-flash-lite` for execution/synthesis and `gemini-3.6-flash` for schema architecture and trajectory judging.
 - **Local Model Support**: The `local` provider wraps the OpenAI-compatible API with local-friendly defaults: no API key required, no rate limiting via SmartKeyPool, extended timeouts (30s connect, 300s request), and clear connection error messages. Works with vLLM, Ollama, and llama.cpp servers.
 
 ## Directory Structure

afterimage/agent_trace/__init__.py

@@ -1,11 +1,17 @@
-"""
-AfterImage Agent Trace subpackage for environment-free synthetic agent-trace dataset generation.
-"""
+"""AfterImage Agent Trace subpackage for environment-free synthetic agent-trace dataset generation."""
 
+from .context import (
+    BaseContextGenerator,
+    CallableContextGenerator,
+    CompositeContextGenerator,
+    PersonaContextGenerator,
+    VirtualUserContextGenerator,
+)
 from .generator import AsyncAgentTraceGenerator
 from .llm_observation_synthesizer import LLMObservationSynthesizer
 from .schema_architect import SchemaArchitect
 from .simulation_engine import DeclarativeEngine, SimulationContext
+from .simula_task_synthesis import SimulaTaskSynthesizer
 from .task_synthesis import GridTaskSynthesizer, InverseFrequencySampler
 from .tool_environment import DeclarativeEnvironment, DeclarativeTool
 from .trajectory_generator import ReActTrajectoryLoop
@@ -27,6 +33,11 @@
 
 __all__ = [
     "AsyncAgentTraceGenerator",
+    "BaseContextGenerator",
+    "VirtualUserContextGenerator",
+    "PersonaContextGenerator",
+    "CallableContextGenerator",
+    "CompositeContextGenerator",
     "LLMObservationSynthesizer",
     "DeclarativeEngine",
     "SimulationContext",
@@ -36,6 +47,7 @@
     "SchemaVerifier",
     "VerificationReport",
     "GridTaskSynthesizer",
+    "SimulaTaskSynthesizer",
     "InverseFrequencySampler",
     "ReActTrajectoryLoop",
     "TrajectoryJudge",

afterimage/agent_trace/context.py

@@ -0,0 +1,255 @@
+"""Context generator framework for environment-free agent trace task synthesis.
+
+Provides extensible, composable abstractions for synthesizing initial environment
+context states (virtual user identities, persona attributes, seed database records)
+to ground synthetic AI agent task directives and tool interactions.
+"""
+
+from __future__ import annotations
+
+import abc
+import inspect
+import json
+import random
+from typing import Any, Callable, Dict, List, Optional, Union
+
+from faker import Faker
+
+from .types import AppDomainSpec, GridTaskBucket
+
+
+class BaseContextGenerator(abc.ABC):
+    """Abstract base class for all initial context generators.
+
+    Initial context generators synthesize structured state payloads (e.g., virtual
+    user profiles, database seed records, account IDs) that anchor synthetic agent
+    tasks into realistic, deterministic execution environments.
+    """
+
+    @abc.abstractmethod
+    async def generate_context(
+        self,
+        app_domains: Optional[Dict[str, AppDomainSpec]] = None,
+        bucket: Optional[GridTaskBucket] = None,
+    ) -> Dict[str, Any]:
+        """Synthesizes an initial context dictionary for task generation.
+
+        Args:
+            app_domains: Map of registered application domain specifications.
+            bucket: Optional grid task bucket specifying complexity constraints.
+
+        Returns:
+            Dict[str, Any]: Key-value pair payload representing initial context state.
+        """
+        pass
+
+    def render_prompt_snippet(self, context: Dict[str, Any]) -> str:
+        """Formats context dictionary into a clean markdown JSON snippet for LLM prompts.
+
+        Args:
+            context: Context state payload.
+
+        Returns:
+            str: Markdown-formatted JSON block for prompt insertion.
+        """
+        if not context:
+            return "{}"
+        try:
+            return json.dumps(context, indent=2, ensure_ascii=False)
+        except Exception:
+            return str(context)
+
+
+class VirtualUserContextGenerator(BaseContextGenerator):
+    """Generates realistic virtual user identity profiles using Faker.
+
+    Synthesizes localized virtual user identities including personal details,
+    account identifiers, contact info, physical location, device specs, and
+    domain-specific seed parameters.
+
+    Args:
+        locale: Locale string for Faker identity generation (e.g. ``"en_US"``).
+        seed: Optional integer seed for reproducible generation.
+        extra_fields_generator: Optional callable producing domain-specific extra fields.
+
+    Example:
+        >>> generator = VirtualUserContextGenerator(locale="en_US", seed=42)
+        >>> context = await generator.generate_context()
+        >>> print(context["user_name"])
+        'Alice Smith'
+    """
+
+    def __init__(
+        self,
+        locale: str = "en_US",
+        seed: Optional[int] = None,
+        extra_fields_generator: Optional[Callable[[], Dict[str, Any]]] = None,
+    ):
+        self.locale = locale
+        self.seed = seed
+        self.faker = Faker(locale)
+        if seed is not None:
+            self.faker.seed_instance(seed)
+        self.extra_fields_generator = extra_fields_generator
+
+    async def generate_context(
+        self,
+        app_domains: Optional[Dict[str, AppDomainSpec]] = None,
+        bucket: Optional[GridTaskBucket] = None,
+    ) -> Dict[str, Any]:
+        """Synthesizes a virtual user identity profile with realistic entity identifiers.
+
+        Args:
+            app_domains: Map of registered application domain specifications.
+            bucket: Optional grid task bucket.
+
+        Returns:
+            Dict[str, Any]: Dictionary containing virtual user attributes and IDs.
+        """
+        user_id = self.faker.random_int(min=101, max=999)
+        account_id = self.faker.random_int(min=1001, max=9999)
+        savings_account_id = account_id + 1
+
+        context: Dict[str, Any] = {
+            "user_id": user_id,
+            "user_name": self.faker.name(),
+            "user_email": self.faker.email(),
+            "user_phone": self.faker.phone_number(),
+            "account_id": account_id,
+            "savings_account_id": savings_account_id,
+            "checking_balance": round(random.uniform(500.0, 5500.0), 2),
+            "savings_balance": round(random.uniform(1000.0, 25000.0), 2),
+            "city": self.faker.city(),
+            "street_address": self.faker.street_address(),
+            "membership_tier": random.choice(["Standard", "Gold", "Platinum", "VIP"]),
+        }
+
+        if self.extra_fields_generator:
+            extra = self.extra_fields_generator()
+            if isinstance(extra, dict):
+                context.update(extra)
+
+        return context
+
+
+class PersonaContextGenerator(BaseContextGenerator):
+    """Integrates persona profiles into initial context payloads.
+
+    Wraps persona attributes (e.g., from ``afterimage.persona_generator.PersonaEntry``)
+    or persona dictionaries to inject user background, expertise level, and communication
+    preferences into task synthesis context.
+
+    Args:
+        persona: Persona entry object or dictionary containing persona details.
+
+    Example:
+        >>> persona_data = {"persona_name": "Tech Enthusiast", "expertise": "expert"}
+        >>> gen = PersonaContextGenerator(persona_data)
+        >>> ctx = await gen.generate_context()
+    """
+
+    def __init__(self, persona: Union[Dict[str, Any], Any]):
+        if hasattr(persona, "model_dump"):
+            self.persona_data = persona.model_dump(mode="json")
+        elif isinstance(persona, dict):
+            self.persona_data = dict(persona)
+        else:
+            self.persona_data = {"persona": str(persona)}
+
+    async def generate_context(
+        self,
+        app_domains: Optional[Dict[str, AppDomainSpec]] = None,
+        bucket: Optional[GridTaskBucket] = None,
+    ) -> Dict[str, Any]:
+        """Injects persona attributes into context state.
+
+        Args:
+            app_domains: Map of registered application domain specifications.
+            bucket: Optional grid task bucket.
+
+        Returns:
+            Dict[str, Any]: Persona context payload.
+        """
+        return {"persona_context": self.persona_data}
+
+
+class CallableContextGenerator(BaseContextGenerator):
+    """Wraps user-defined callables to produce initial context state.
+
+    Supports both synchronous and asynchronous functions returning dictionary state payloads.
+
+    Args:
+        func: Sync or async callable returning a dictionary.
+
+    Example:
+        >>> gen = CallableContextGenerator(lambda: {"custom_key": 42})
+        >>> ctx = await gen.generate_context()
+    """
+
+    def __init__(self, func: Callable[..., Any]):
+        self.func = func
+
+    async def generate_context(
+        self,
+        app_domains: Optional[Dict[str, AppDomainSpec]] = None,
+        bucket: Optional[GridTaskBucket] = None,
+    ) -> Dict[str, Any]:
+        """Invokes the wrapped callable to produce context state.
+
+        Args:
+            app_domains: Map of registered application domain specifications.
+            bucket: Optional grid task bucket.
+
+        Returns:
+            Dict[str, Any]: Result of callable execution.
+        """
+        if inspect.iscoroutinefunction(self.func):
+            res = await self.func()
+        else:
+            res = self.func()
+
+        if isinstance(res, dict):
+            return res
+        return {"data": res}
+
+
+class CompositeContextGenerator(BaseContextGenerator):
+    """Composes multiple initial context generators into a unified context provider.
+
+    Executes all child context generators in sequence, merging their produced state
+    dictionaries into a single context payload.
+
+    Args:
+        generators: List of child :class:`BaseContextGenerator` instances.
+
+    Example:
+        >>> gen = CompositeContextGenerator([
+        ...     VirtualUserContextGenerator(),
+        ...     CallableContextGenerator(lambda: {"order_id": 999})
+        ... ])
+        >>> ctx = await gen.generate_context()
+    """
+
+    def __init__(self, generators: List[BaseContextGenerator]):
+        self.generators = generators
+
+    async def generate_context(
+        self,
+        app_domains: Optional[Dict[str, AppDomainSpec]] = None,
+        bucket: Optional[GridTaskBucket] = None,
+    ) -> Dict[str, Any]:
+        """Generates context from all child generators and merges their dictionaries.
+
+        Args:
+            app_domains: Map of registered application domain specifications.
+            bucket: Optional grid task bucket.
+
+        Returns:
+            Dict[str, Any]: Merged context dictionary.
+        """
+        merged: Dict[str, Any] = {}
+        for gen in self.generators:
+            res = await gen.generate_context(app_domains=app_domains, bucket=bucket)
+            if isinstance(res, dict):
+                merged.update(res)
+        return merged