new: implement environment-free agent-trace generation

DESIGN.md

@@ -31,6 +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). 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 `SimulationContext` entity lookup pools), `GridTaskSynthesizer` (360-bucket grid + `InverseFrequencySampler` + procedural task rewriter), `ReActTrajectoryLoop` (multi-turn teacher execution against local tools), `TrajectoryJudge` (9-point LLM quality rubric), and `AsyncAgentTraceGenerator` facade. 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/__init__.py

@@ -1,5 +1,6 @@
 import importlib.metadata
 
+from afterimage.agent_trace import AsyncAgentTraceGenerator  # noqa
 from afterimage.async_conversation_generator import AsyncConversationGenerator  # noqa
 from afterimage.callbacks import (
     AndStoppingCallback,  # noqa

afterimage/agent_trace/__init__.py

@@ -0,0 +1,49 @@
+"""
+AfterImage Agent Trace subpackage for environment-free synthetic agent-trace dataset generation.
+"""
+
+from .generator import AsyncAgentTraceGenerator
+from .schema_architect import SchemaArchitect
+from .simulation_engine import DeclarativeEngine, SimulationContext
+from .task_synthesis import GridTaskSynthesizer, InverseFrequencySampler
+from .tool_environment import DeclarativeEnvironment, DeclarativeTool
+from .trajectory_generator import ReActTrajectoryLoop
+from .trajectory_judge import TrajectoryJudge
+from .types import (
+    AgentTrajectory,
+    AppDomainSpec,
+    GridTaskBucket,
+    JudgeVerdict,
+    RubricScores,
+    ToolActionSpec,
+    ToolCall,
+    ToolObservation,
+    ToolParameterSpec,
+    TrajectoryTurn,
+)
+from .verifier import SchemaVerifier, VerificationReport
+
+__all__ = [
+    "AsyncAgentTraceGenerator",
+    "DeclarativeEngine",
+    "SimulationContext",
+    "DeclarativeEnvironment",
+    "DeclarativeTool",
+    "SchemaArchitect",
+    "SchemaVerifier",
+    "VerificationReport",
+    "GridTaskSynthesizer",
+    "InverseFrequencySampler",
+    "ReActTrajectoryLoop",
+    "TrajectoryJudge",
+    "AgentTrajectory",
+    "AppDomainSpec",
+    "GridTaskBucket",
+    "JudgeVerdict",
+    "RubricScores",
+    "ToolActionSpec",
+    "ToolCall",
+    "ToolObservation",
+    "ToolParameterSpec",
+    "TrajectoryTurn",
+]

afterimage/agent_trace/generator.py

@@ -0,0 +1,161 @@
+import asyncio
+import logging
+from typing import Any, Dict, List, Optional, Union
+from pydantic import BaseModel
+
+from ..key_management import SmartKeyPool
+from ..providers.llm_providers import LLMFactory, LLMProvider
+from ..storage import BaseStorage, JSONLStorage
+from ..types import Conversation, ConversationEntry, Role
+from .schema_architect import SchemaArchitect
+from .task_synthesis import GridTaskSynthesizer
+from .tool_environment import DeclarativeEnvironment
+from .trajectory_generator import ReActTrajectoryLoop
+from .trajectory_judge import TrajectoryJudge
+from .types import AgentTrajectory, AppDomainSpec, ToolActionSpec
+
+logger = logging.getLogger(__name__)
+
+
+class AsyncAgentTraceGenerator:
+    """Async Environment-Free Synthetic Agent-Trace Dataset Generator Facade."""
+
+    def __init__(
+        self,
+        api_key: Optional[Union[str, List[str], SmartKeyPool]] = None,
+        llm_provider: Optional[LLMProvider] = None,
+        provider: str = "gemini",
+        architect_model: str = "gemini-3.6-flash",
+        teacher_model: str = "gemini-3.5-flash-lite",
+        judge_model: str = "gemini-3.6-flash",
+        storage: Optional[BaseStorage] = None,
+    ):
+        if llm_provider:
+            self.llm_provider = llm_provider
+        else:
+            self.llm_provider = LLMFactory.create(
+                provider=provider,
+                api_key=api_key,
+                model_name=architect_model,
+            )
+
+        self.architect = SchemaArchitect(
+            llm_provider=self.llm_provider,
+            model_name=architect_model,
+        )
+        self.synthesizer = GridTaskSynthesizer(
+            llm_provider=self.llm_provider,
+            model_name=teacher_model,
+        )
+        self.teacher_loop = ReActTrajectoryLoop(
+            llm_provider=self.llm_provider,
+            model_name=teacher_model,
+        )
+        self.judge = TrajectoryJudge(
+            llm_provider=self.llm_provider,
+            model_name=judge_model,
+        )
+
+        self.environment = DeclarativeEnvironment()
+        self.storage = storage or JSONLStorage(conversations_path="outputs/agent_trajectories.jsonl")
+
+    async def register_app_domain(
+        self, app_name: str, app_description: str, actions: List[ToolActionSpec]
+    ) -> AppDomainSpec:
+        """Runs SchemaArchitect to generate and register Pydantic response models for an app domain."""
+        app_spec, model_classes = await self.architect.generate_app_domain_schema(
+            app_name=app_name,
+            app_description=app_description,
+            actions=actions,
+        )
+        self.environment.register_app_domain(app_spec, model_classes=model_classes)
+        return app_spec
+
+    async def generate_single(self, max_turns: int = 6) -> Optional[AgentTrajectory]:
+        """Synthesizes a single agent trajectory (task -> ReAct loop -> judge)."""
+        if not self.environment.app_domains:
+            raise ValueError("No app domains registered. Call register_app_domain() first.")
+
+        # 1. Task synthesis via 360-bucket grid & task rewriter
+        task, selected_apps, bucket = await self.synthesizer.synthesize_task(
+            app_domains=self.environment.app_domains
+        )
+
+        # 2. ReAct teacher trajectory loop against DeclarativeEnvironment (< 1ms tool calls)
+        trajectory = await self.teacher_loop.run_trajectory(
+            task=task,
+            environment=self.environment,
+            domain_apps=selected_apps,
+        )
+        trajectory.metadata["grid_bucket"] = bucket.model_dump()
+
+        # 3. Trajectory Judge Quality Filtering
+        verdict = await self.judge.evaluate_trajectory(trajectory)
+        trajectory.judge_verdict = verdict
+
+        if verdict.is_valid:
+            return trajectory
+        return None
+
+    async def generate(
+        self,
+        num_trajectories: int = 10,
+        max_turns: int = 6,
+        max_concurrency: int = 4,
+    ) -> List[AgentTrajectory]:
+        """Generates multiple synthetic agent trajectories concurrently."""
+        sem = asyncio.Semaphore(max_concurrency)
+        accepted_trajectories: List[AgentTrajectory] = []
+
+        async def _worker() -> Optional[AgentTrajectory]:
+            async with sem:
+                try:
+                    return await self.generate_single(max_turns=max_turns)
+                except Exception as e:
+                    logger.warning(f"Error during trajectory generation worker: {e}")
+                    return None
+
+        tasks = [_worker() for _ in range(num_trajectories)]
+        results = await asyncio.gather(*tasks)
+
+        conversations = []
+        for res in results:
+            if res:
+                accepted_trajectories.append(res)
+                conv = self._trajectory_to_conversation(res)
+                conversations.append(conv)
+
+        if conversations:
+            self.storage.save_conversations(conversations)
+
+        return accepted_trajectories
+
+    def _trajectory_to_conversation(self, traj: AgentTrajectory) -> Conversation:
+        """Converts an AgentTrajectory into AfterImage's base Conversation schema."""
+        entries: List[ConversationEntry] = [
+            ConversationEntry(role=Role.USER, content=traj.task)
+        ]
+        for t in traj.turns:
+            entry_text = f"Thought: {t.agent_thought}"
+            if t.tool_call:
+                entry_text += f"\nAction: {t.tool_call.app}.{t.tool_call.action}\nAction Input: {t.tool_call.parameters}"
+            entries.append(ConversationEntry(role=Role.ASSISTANT, content=entry_text))
+
+            if t.observation:
+                entries.append(
+                    ConversationEntry(
+                        role=Role.USER, content=f"Observation: {t.observation.observation}"
+                    )
+                )
+
+        if traj.final_answer:
+            entries.append(
+                ConversationEntry(role=Role.ASSISTANT, content=f"Final Answer: {traj.final_answer}")
+            )
+
+        metadata = traj.metadata
+        if traj.judge_verdict:
+            metadata["judge_verdict"] = traj.judge_verdict.model_dump()
+        metadata["trajectory_id"] = traj.trajectory_id
+
+        return Conversation(conversations=entries, metadata=metadata)

afterimage/agent_trace/schema_architect.py

@@ -0,0 +1,139 @@
+import asyncio
+from typing import Dict, List, Optional, Tuple, Type
+from pydantic import BaseModel
+
+from ..providers.llm_providers import LLMProvider
+from .types import AppDomainSpec, ToolActionSpec, ToolParameterSpec
+from .verifier import SchemaVerifier, VerificationReport
+
+
+SCHEMA_ARCHITECT_PROMPT_TEMPLATE = """You are an expert LLM Schema Architect for synthetic tool environments.
+Your task is to generate valid Python source code containing Pydantic V2 response models for an application domain.
+
+App Name: {app_name}
+App Description: {app_description}
+Target Endpoints: {endpoints_desc}
+
+### ANNOTATION PROTOCOL CONSTRAINTS:
+You MUST annotate fields using Pydantic's `Field(json_schema_extra={{...}})` parameter following these exact keys:
+1. `json_schema_extra={{"generator": "id"}}` for primary keys (e.g. user_id, order_id).
+2. `json_schema_extra={{"generator": "fk:<entity>.<field>"}}` for foreign keys (e.g. fk:user.user_id, fk:order.order_id).
+3. `json_schema_extra={{"generator": "money"}}` for monetary float amounts.
+4. `json_schema_extra={{"generator": "faker:<method>"}}` for fake string data (e.g. faker:name, faker:email, faker:company, faker:street_address).
+5. `json_schema_extra={{"generator": "enum", "values": ["v1", "v2"]}}` for categorical strings.
+
+### CODE FORMATTING REQUIREMENTS:
+- Output ONLY valid, executable Python code inside a ```python ``` code block.
+- Import `from pydantic import BaseModel, Field, EmailStr`.
+- Inherit all response models from `BaseModel`.
+- Ensure every model primary key has `generator: "id"` so foreign key references in other models can resolve safely.
+
+{feedback_section}
+"""
+
+
+class SchemaArchitect:
+    """LLM Schema Architect that generates Pydantic response model specifications with static verification feedback loop."""
+
+    def __init__(
+        self,
+        llm_provider: LLMProvider,
+        model_name: str = "gemini-3.6-flash",
+        max_retries: int = 3,
+    ):
+        self.llm_provider = llm_provider
+        self.model_name = model_name
+        self.max_retries = max_retries
+        self.verifier = SchemaVerifier()
+
+    async def generate_app_domain_schema(
+        self,
+        app_name: str,
+        app_description: str,
+        actions: List[ToolActionSpec],
+        existing_primary_ids: Optional[set[str]] = None,
+    ) -> Tuple[AppDomainSpec, Dict[str, Type[BaseModel]]]:
+        """Generates Pydantic response models code for an app domain with static verification retries."""
+        endpoints_desc = "\n".join(
+            [f"- {a.action_name}: {a.description} -> ResponseModel: {a.response_model_name}" for a in actions]
+        )
+
+        feedback_section = ""
+        last_code = ""
+        report = VerificationReport(is_valid=False)
+
+        for attempt in range(1, self.max_retries + 1):
+            prompt = SCHEMA_ARCHITECT_PROMPT_TEMPLATE.format(
+                app_name=app_name,
+                app_description=app_description,
+                endpoints_desc=endpoints_desc,
+                feedback_section=feedback_section,
+            )
+
+            response = await self.llm_provider.agenerate(
+                prompt=prompt,
+                model_name=self.model_name,
+                temperature=0.2,
+            )
+            raw_text = response.text
+            code_str = self._extract_python_code(raw_text)
+
+            report = self.verifier.verify_code(code_str, existing_declared_ids=existing_primary_ids)
+            if report.is_valid:
+                last_code = code_str
+                break
+
+            feedback_section = f"\n### PREVIOUS ATTEMPT VERIFICATION ERRORS (ATTEMPT {attempt}/{self.max_retries}):\n{report.get_prompt_feedback()}\nPlease fix all listed errors and output the corrected Python code."
+
+        if not report.is_valid and not last_code:
+            # Fallback minimal schema if LLM retries failed
+            last_code = self._generate_fallback_code(actions)
+
+        # Exec verified code into runtime model classes dictionary
+        model_classes = self._compile_model_classes(last_code)
+
+        app_spec = AppDomainSpec(
+            app_name=app_name,
+            description=app_description,
+            actions=actions,
+            response_models_code=last_code,
+        )
+        return app_spec, model_classes
+
+    def _extract_python_code(self, text: str) -> str:
+        """Extracts python code block from LLM response."""
+        if "```python" in text:
+            return text.split("```python", 1)[1].split("```", 1)[0].strip()
+        elif "```" in text:
+            return text.split("```", 1)[1].split("```", 1)[0].strip()
+        return text.strip()
+
+    def _compile_model_classes(self, code_str: str) -> Dict[str, Type[BaseModel]]:
+        """Compiles Python source code string into live Pydantic BaseModel classes."""
+        local_scope: Dict[str, Any] = {}
+        try:
+            exec(code_str, local_scope)
+        except Exception:
+            return {}
+
+        model_classes: Dict[str, Type[BaseModel]] = {}
+        for k, v in local_scope.items():
+            if isinstance(v, type) and issubclass(v, BaseModel) and v is not BaseModel:
+                model_classes[k] = v
+        return model_classes
+
+    def _generate_fallback_code(self, actions: List[ToolActionSpec]) -> str:
+        """Generates fallback basic Pydantic code if LLM generation exhausts retries."""
+        lines = [
+            "from pydantic import BaseModel, Field",
+            "from typing import List, Optional\n",
+        ]
+        for a in actions:
+            model_name = a.response_model_name
+            lines.extend([
+                f"class {model_name}(BaseModel):",
+                f"    id: int = Field(json_schema_extra={{'generator': 'id'}})",
+                f"    status: str = Field(default='success')",
+                f"    message: str = Field(default='Action completed')\n",
+            ])
+        return "\n".join(lines)