teich 0.1.1a1__py3-none-any.whl

agentic_datagen/__init__.py

@@ -0,0 +1,18 @@
+"""Agentic Datagen v2 - Generate training data from Codex and Pi traces."""
+
+__version__ = "0.1.1a1"
+
+from .config import Config, load_config
+from .converter import TrainingExample, convert_trace_to_training_example, convert_traces_to_training_data
+from .formatter import format_and_mask
+from .loader import load_traces
+
+__all__ = [
+    "Config",
+    "TrainingExample",
+    "convert_trace_to_training_example",
+    "convert_traces_to_training_data",
+    "format_and_mask",
+    "load_traces",
+    "load_config",
+]

agentic_datagen/__main__.py

@@ -0,0 +1,4 @@
+from .cli import main
+
+if __name__ == "__main__":
+    main()

agentic_datagen/cli.py

@@ -0,0 +1,333 @@
+"""CLI for teich."""
+
+from __future__ import annotations
+
+from datetime import datetime, timezone
+from pathlib import Path
+from threading import RLock
+
+import typer
+from rich.console import Console, Group
+from rich.live import Live
+from rich.panel import Panel
+from rich.table import Table
+
+from .config import Config
+from .runner import CodexRunner, PiRunner, SessionProgressUpdate, TraceMetrics
+from .trace_readme import write_traces_readme
+
+console = Console()
+app = typer.Typer(
+    name="teich",
+    help="Generate agent training data using Codex or Pi",
+    no_args_is_help=True,
+)
+
+
+@app.command()
+def generate(
+    config: Path = typer.Option(
+        Path("config.yaml"),
+        "--config", "-c",
+        help="Path to configuration file",
+    ),
+    output: Path = typer.Option(
+        None,
+        "--output", "-o",
+        help="Override output directory",
+    ),
+    concurrency: int = typer.Option(
+        None,
+        "--concurrency", "-j",
+        min=1,
+        help="Number of prompts to run in parallel",
+    ),
+) -> None:
+    """Generate training traces from prompts."""
+    console.print(Panel.fit("Teich", style="bold blue"))
+
+    if not config.exists():
+        console.print(f"[red]Config file not found: {config}[/red]")
+        raise typer.Exit(1)
+
+    cfg = Config.from_yaml(config)
+    if output:
+        cfg.output.traces_dir = output
+    if concurrency is not None:
+        cfg.max_concurrency = concurrency
+
+    # Ensure output dir exists
+    cfg.output.traces_dir.mkdir(parents=True, exist_ok=True)
+    prompt_inputs = cfg.get_prompt_inputs()
+    effective_concurrency = (
+        max(1, min(cfg.max_concurrency, len(prompt_inputs))) if prompt_inputs else cfg.max_concurrency
+    )
+
+    console.print(f"[green]Loaded config: {config}[/green]")
+    console.print(
+        f"[blue]Processing {len(prompt_inputs)} prompts with concurrency {effective_concurrency}...[/blue]\n"
+    )
+
+    # Run generation
+    try:
+        agent_provider = cfg.get_agent_provider()
+        if agent_provider == "codex":
+            runner = CodexRunner(cfg)
+        elif agent_provider == "pi":
+            runner = PiRunner(cfg)
+        else:
+            console.print(
+                f"[red]Unsupported agent provider: {agent_provider}. Supported providers: codex, pi.[/red]"
+            )
+            raise typer.Exit(1)
+
+        with BatchProgressReporter(console) as reporter:
+            results = runner.run_all(
+                max_concurrency=cfg.max_concurrency,
+                progress_callback=reporter.update,
+            )
+        readme_path = write_traces_readme(
+            cfg.output.traces_dir,
+            pretty_name=cfg.output.pretty_name,
+            tags=cfg.output.tags,
+            model_id=cfg.model.model,
+            readme_file_name=cfg.output.readme_file_name,
+        )
+        totals = reporter.snapshot_totals()
+
+        console.print(f"\n[bold green]Success! Generated {len(results)} trace files:[/bold green]")
+        for r in results:
+            console.print(f"  - {r}")
+        console.print(
+            "\n[cyan]Usage:[/cyan] "
+            f"tokens={totals['total_tokens']} input={totals['input_tokens']} "
+            f"output={totals['output_tokens']} reasoning={totals['reasoning_tokens']} "
+            f"cache_read={totals['cache_read_tokens']}"
+        )
+        console.print(f"[cyan]API cost:[/cyan] ${totals['total_cost']:.6f}")
+        console.print(f"[green]Saved sandboxes: {cfg.output.sandbox_dir}[/green]")
+        console.print(f"\n[green]Wrote README: {readme_path}[/green]")
+
+    except Exception as e:
+        console.print(f"\n[red]Error: {e}[/red]")
+        raise typer.Exit(1)
+
+
+def _format_elapsed(update: SessionProgressUpdate) -> str:
+    if update.started_at is None:
+        return "--"
+    finished_at = update.finished_at or datetime.now(timezone.utc)
+    elapsed = max(0, int((finished_at - update.started_at).total_seconds()))
+    minutes, seconds = divmod(elapsed, 60)
+    hours, minutes = divmod(minutes, 60)
+    if hours:
+        return f"{hours:d}:{minutes:02d}:{seconds:02d}"
+    return f"{minutes:02d}:{seconds:02d}"
+
+
+def _format_api_tokens(metrics: TraceMetrics | None) -> str:
+    if not metrics or not metrics.total_tokens:
+        return "--"
+    return str(metrics.total_tokens)
+
+
+def _format_cost(metrics: TraceMetrics | None) -> str:
+    if not metrics:
+        return "--"
+    return f"${metrics.total_cost:.6f}"
+
+
+class BatchProgressReporter:
+    def __init__(self, progress_console: Console):
+        self.console = progress_console
+        self.enabled = progress_console.is_terminal
+        self._lock = RLock()
+        self._updates: dict[str, SessionProgressUpdate] = {}
+        self._live: Live | None = None
+
+    def __enter__(self) -> BatchProgressReporter:
+        if self.enabled:
+            self._live = Live(self._render(), console=self.console, refresh_per_second=4)
+            self._live.__enter__()
+        return self
+
+    def __exit__(self, exc_type, exc, exc_tb) -> None:
+        if self._live is not None:
+            self._live.__exit__(exc_type, exc, exc_tb)
+            self._live = None
+
+    def update(self, update: SessionProgressUpdate) -> None:
+        with self._lock:
+            previous = self._updates.get(update.prompt_id)
+            if previous and update.started_at is None:
+                update.started_at = previous.started_at
+            self._updates[update.prompt_id] = update
+            if self._live is not None:
+                self._live.update(self._render(), refresh=True)
+
+    def snapshot_totals(self) -> dict[str, float | int]:
+        with self._lock:
+            totals = {
+                "input_tokens": 0,
+                "output_tokens": 0,
+                "reasoning_tokens": 0,
+                "cache_read_tokens": 0,
+                "cache_write_tokens": 0,
+                "total_tokens": 0,
+                "est_total_tokens": 0,
+                "total_cost": 0.0,
+            }
+            for update in self._updates.values():
+                metrics = update.metrics
+                if not metrics:
+                    continue
+                totals["input_tokens"] += metrics.input_tokens
+                totals["output_tokens"] += metrics.output_tokens
+                totals["reasoning_tokens"] += metrics.reasoning_tokens
+                totals["cache_read_tokens"] += metrics.cache_read_tokens
+                totals["cache_write_tokens"] += metrics.cache_write_tokens
+                totals["total_tokens"] += metrics.total_tokens
+                totals["est_total_tokens"] += metrics.est_total_tokens
+                totals["total_cost"] += metrics.total_cost
+            return totals
+
+    def _render(self):
+        with self._lock:
+            summary = Table.grid(expand=True)
+            summary.add_column(justify="left")
+            summary.add_column(justify="left")
+            queued = sum(1 for update in self._updates.values() if update.status == "queued")
+            running = sum(1 for update in self._updates.values() if update.status == "running")
+            completed = sum(1 for update in self._updates.values() if update.status == "completed")
+            failed = sum(1 for update in self._updates.values() if update.status == "failed")
+            totals = self.snapshot_totals()
+            summary.add_row(
+                f"queued={queued} running={running} completed={completed} failed={failed}",
+                f"tokens={totals['total_tokens']} cost=${totals['total_cost']:.6f}",
+            )
+
+            table = Table(title="Generation Progress")
+            table.add_column("#", justify="right", no_wrap=True)
+            table.add_column("status", no_wrap=True)
+            table.add_column("elapsed", no_wrap=True)
+            table.add_column("tokens", justify="right", no_wrap=True)
+            table.add_column("cost", justify="right", no_wrap=True)
+            table.add_column("model", no_wrap=True)
+            table.add_column("prompt")
+            table.add_column("details")
+            active_updates = [
+                update
+                for update in self._updates.values()
+                if update.status in {"queued", "running"}
+            ]
+            for update in sorted(active_updates, key=lambda item: item.prompt_index):
+                metrics = update.metrics
+                model = metrics.model if metrics and metrics.model else "--"
+                details = "--"
+                if update.trace_path is not None:
+                    details = str(update.trace_path.name)
+                elif update.error:
+                    details = update.error
+                table.add_row(
+                    str(update.prompt_index),
+                    update.status,
+                    _format_elapsed(update),
+                    _format_api_tokens(metrics),
+                    _format_cost(metrics),
+                    model,
+                    update.prompt_preview,
+                    details,
+                )
+            return Group(Panel.fit(summary, title="Batch Summary"), table)
+
+
+@app.command()
+def init(
+    path: Path = typer.Argument(Path("."), help="Directory to initialize"),
+) -> None:
+    """Initialize a new teich project."""
+    console.print(Panel.fit("Initialize Project", style="bold blue"))
+
+    path.mkdir(parents=True, exist_ok=True)
+
+    # Create config.yaml
+    config_path = path / "config.yaml"
+    if not config_path.exists():
+        config_path.write_text(CONFIG_TEMPLATE, encoding="utf-8")
+        console.print(f"[green]Created: {config_path}[/green]")
+    else:
+        console.print(f"[yellow]Already exists: {config_path}[/yellow]")
+
+    # Create prompts.csv
+    prompts_path = path / "prompts.csv"
+    if not prompts_path.exists():
+        prompts_path.write_text(PROMPTS_TEMPLATE, encoding="utf-8")
+        console.print(f"[green]Created: {prompts_path}[/green]")
+    else:
+        console.print(f"[yellow]Already exists: {prompts_path}[/yellow]")
+
+    console.print(f"\n[bold green]Initialized in {path.absolute()}[/bold green]")
+    console.print("\n[yellow]Next:[/yellow]")
+    console.print("1. Set OPENAI_API_KEY in config.yaml or env")
+    console.print("2. Add prompt rows to prompts.csv")
+    console.print("3. Run: [cyan]uvx teich generate -c config.yaml[/cyan]")
+
+
+CONFIG_TEMPLATE = '''# Teich Configuration
+agent:
+  provider: codex
+
+model:
+  model: codex-mini-latest
+  approval_policy: never
+  sandbox: danger-full-access
+  reasoning_effort: null
+
+# Optional: API/model provider configuration
+# api:
+#   provider: openai
+#   base_url: null
+#   api_key: null
+
+# Optional: MCP servers
+# mcp_servers:
+#   - name: filesystem
+#     command: npx
+#     args: ["-y", "@modelcontextprotocol/server-filesystem", "/workspace"]
+
+# Load prompts from file (optional)
+prompts_file: prompts.csv
+
+# Or define prompts here
+prompts: []
+
+output:
+  traces_dir: ./output
+  sandbox_dir: ./sandbox
+  pretty_name: "My Agent Traces"
+  readme_file_name: README.md
+  tags:
+    - agent-traces
+    - codex
+
+max_concurrency: 1
+timeout_seconds: 600
+
+# Set via env: OPENAI_API_KEY or here
+openai_api_key: null
+developer_instructions: null
+'''
+
+PROMPTS_TEMPLATE = '''image,github_repo,prompt
+None,None,"Build a simple todo list app in React"
+None,None,"Create a Python script that fetches weather data from an API"
+None,armand0e/perplexica-mcp,"Add a small usability improvement and update the tests"
+'''
+
+
+def main():
+    app()
+
+
+if __name__ == "__main__":
+    main()

agentic_datagen/config.py

@@ -0,0 +1,289 @@
+"""Configuration management for teich."""
+
+from __future__ import annotations
+
+import csv
+import os
+from pathlib import Path
+import re
+
+import yaml
+from pydantic import BaseModel, Field, field_validator, model_validator
+
+
+GITHUB_REPO_ID_PATTERN = re.compile(r"^[A-Za-z0-9_.-]+/[A-Za-z0-9_.-]+$")
+
+
+def _get_env_alias(*names: str) -> str | None:
+    for name in names:
+        value = os.getenv(name)
+        if value:
+            return value
+    return None
+
+
+class MCPConfig(BaseModel):
+    """MCP server configuration."""
+    name: str
+    command: str | None = None
+    args: list[str] = Field(default_factory=list)
+    env: dict[str, str] = Field(default_factory=dict)
+    env_vars: list[str] = Field(default_factory=list)
+    cwd: str | None = None
+    url: str | None = None
+    bearer_token_env_var: str | None = None
+    http_headers: dict[str, str] = Field(default_factory=dict)
+    env_http_headers: dict[str, str] = Field(default_factory=dict)
+    startup_timeout_sec: int | None = None
+    tool_timeout_sec: int | None = None
+    enabled: bool = True
+    required: bool = False
+    enabled_tools: list[str] = Field(default_factory=list)
+    disabled_tools: list[str] = Field(default_factory=list)
+
+    @model_validator(mode="after")
+    def validate_transport(self) -> MCPConfig:
+        if not self.command and not self.url:
+            raise ValueError(
+                f"MCP server '{self.name}' must define either command or url"
+            )
+        if self.command and self.url:
+            raise ValueError(
+                f"MCP server '{self.name}' cannot define both command and url"
+            )
+        return self
+
+
+class APIConfig(BaseModel):
+    """API configuration for OpenAI-compatible endpoints."""
+    provider: str = "openai"  # openai, openrouter, azure, etc.
+    base_url: str | None = None  # e.g., https://openrouter.ai/api/v1
+    api_key: str | None = None  # Override global api_key for this provider
+    wire_api: str = "responses"
+
+
+class AgentConfig(BaseModel):
+    """Agent runtime selection."""
+    provider: str = "codex"
+
+
+class ModelConfig(BaseModel):
+    """Model configuration."""
+    model: str = "codex-mini-latest"
+    approval_policy: str = "never"
+    sandbox: str = "danger-full-access"
+    reasoning_effort: str | None = None
+    approval_mode: str | None = "none"
+    pi_model_overrides: dict[str, object] = Field(default_factory=dict)
+
+    @model_validator(mode="after")
+    def normalize_legacy_approval_mode(self) -> ModelConfig:
+        if self.approval_mode:
+            self.approval_policy = {
+                "none": "never",
+                "suggest": "on-request",
+                "auto-edit": "on-request",
+                "full-auto": "on-request",
+            }.get(self.approval_mode, self.approval_policy)
+        return self
+
+
+class OutputConfig(BaseModel):
+    """Output configuration."""
+    traces_dir: Path = Field(default=Path("./output"))
+    sandbox_dir: Path = Field(default=Path("./sandbox"))
+    pretty_name: str = "Agentic Training Traces"
+    tags: list[str] = Field(default_factory=lambda: ["agent-traces", "codex"])
+    readme_file_name: str = "README.md"
+
+
+class PromptInput(BaseModel):
+    """Structured prompt input row."""
+    image: str | None = None
+    github_repo: str | None = None
+    prompt: str
+
+    @staticmethod
+    def _normalize_optional_text(value: object) -> str | None:
+        if value is None:
+            return None
+        text = value if isinstance(value, str) else str(value)
+        normalized = text.strip()
+        if not normalized or normalized.lower() == "none":
+            return None
+        return normalized
+
+    @field_validator("image", "github_repo", mode="before")
+    @classmethod
+    def normalize_optional_fields(cls, value: object) -> str | None:
+        return cls._normalize_optional_text(value)
+
+    @field_validator("prompt", mode="before")
+    @classmethod
+    def normalize_prompt(cls, value: object) -> str:
+        if value is None:
+            return ""
+        return (value if isinstance(value, str) else str(value)).strip()
+
+    @field_validator("prompt")
+    @classmethod
+    def validate_prompt(cls, value: str) -> str:
+        if not value:
+            raise ValueError("Prompt cannot be empty")
+        return value
+
+    @field_validator("github_repo")
+    @classmethod
+    def validate_github_repo(cls, value: str | None) -> str | None:
+        if value is None:
+            return None
+        if not GITHUB_REPO_ID_PATTERN.fullmatch(value):
+            raise ValueError(
+                "github_repo must be in owner/repo format, e.g. armand0e/perplexica-mcp"
+            )
+        return value
+
+
+class Config(BaseModel):
+    """Main configuration."""
+    agent: AgentConfig = Field(default_factory=AgentConfig)
+    model: ModelConfig = Field(default_factory=ModelConfig)
+    api: APIConfig = Field(default_factory=APIConfig)
+    mcp_servers: list[MCPConfig] = Field(default_factory=list)
+    prompts: list[str] = Field(default_factory=list)
+    prompts_file: Path | None = None
+    output: OutputConfig = Field(default_factory=OutputConfig)
+    max_concurrency: int = Field(default=1, ge=1)
+    timeout_seconds: int = 600
+    openai_api_key: str | None = None
+    developer_instructions: str | None = None
+
+    @field_validator("prompts_file")
+    @classmethod
+    def validate_prompts_file(cls, v: Path | None) -> Path | None:
+        if v is not None and not v.exists():
+            raise ValueError(f"Prompts file not found: {v}")
+        return v
+
+    @classmethod
+    def from_yaml(cls, path: Path) -> Config:
+        """Load config from YAML file.
+
+        Environment variables override config file values (for testing):
+            TEICH_MODEL - Override model ID
+            TEICH_BASE_URL - Override API base URL
+            TEICH_API_KEY - Override API key
+            TEICH_PROVIDER - Override provider name
+        """
+        with open(path, "r", encoding="utf-8") as f:
+            data = yaml.safe_load(f) or {}
+
+        api_config = data.get("api")
+        if isinstance(api_config, dict) and "model" in api_config:
+            raise ValueError(
+                "Unsupported config key 'api.model'. Set the model under 'model.model' instead."
+            )
+        prompts_file = data.get("prompts_file")
+        if isinstance(prompts_file, str) and prompts_file.strip():
+            prompts_file_path = Path(prompts_file)
+            if not prompts_file_path.is_absolute():
+                data["prompts_file"] = (path.parent / prompts_file_path).resolve()
+
+        # Apply environment variable overrides
+        if model_env := _get_env_alias("TEICH_MODEL", "AGENTIC_DATAGEN_MODEL"):
+            data.setdefault("model", {})["model"] = model_env
+        if base_url_env := _get_env_alias("TEICH_BASE_URL", "AGENTIC_DATAGEN_BASE_URL"):
+            data.setdefault("api", {})["base_url"] = base_url_env
+        if api_key_env := _get_env_alias("TEICH_API_KEY", "AGENTIC_DATAGEN_API_KEY"):
+            data.setdefault("api", {})["api_key"] = api_key_env
+        if provider_env := _get_env_alias("TEICH_PROVIDER", "AGENTIC_DATAGEN_PROVIDER"):
+            data.setdefault("api", {})["provider"] = provider_env
+
+        return cls(**data)
+
+    def get_api_key(self) -> str | None:
+        """Get effective API key (from api config or global)."""
+        return self.api.api_key or self.openai_api_key
+
+    def get_effective_model(self) -> str:
+        """Get effective model identifier."""
+        return self.model.model
+
+    def get_base_url(self) -> str | None:
+        """Get effective base URL."""
+        return self.api.base_url
+
+    def get_agent_provider(self) -> str:
+        """Get active agent provider."""
+        return self.agent.provider.strip().lower() or "codex"
+
+    def get_prompts(self) -> list[str]:
+        """Get prompt text only for all configured prompt inputs."""
+        return [prompt_input.prompt for prompt_input in self.get_prompt_inputs()]
+
+    def get_prompt_inputs(self) -> list[PromptInput]:
+        """Get structured prompt inputs from config and prompts_file."""
+        prompt_inputs = [PromptInput(prompt=prompt) for prompt in self.prompts]
+        if self.prompts_file:
+            prompt_inputs.extend(self._load_prompt_inputs_from_file(self.prompts_file))
+        for prompt_input in prompt_inputs:
+            if prompt_input.image is not None:
+                raise ValueError(
+                    "Prompt image inputs are not supported yet. Leave the image column blank or set it to None."
+                )
+        return prompt_inputs
+
+    @staticmethod
+    def _load_prompt_inputs_from_file(path: Path) -> list[PromptInput]:
+        if path.suffix.lower() == ".csv":
+            return Config._load_prompt_inputs_from_csv(path)
+        return Config._load_prompt_inputs_from_text(path)
+
+    @staticmethod
+    def _load_prompt_inputs_from_text(path: Path) -> list[PromptInput]:
+        with path.open("r", encoding="utf-8") as handle:
+            return [
+                PromptInput(prompt=line.strip())
+                for line in handle
+                if line.strip() and not line.startswith("#")
+            ]
+
+    @staticmethod
+    def _load_prompt_inputs_from_csv(path: Path) -> list[PromptInput]:
+        with path.open("r", encoding="utf-8", newline="") as handle:
+            reader = csv.DictReader(handle)
+            fieldnames = [name.strip().lower() for name in reader.fieldnames or [] if isinstance(name, str)]
+            if "prompt" not in fieldnames:
+                raise ValueError("Prompt CSV must include a 'prompt' column")
+            prompt_inputs: list[PromptInput] = []
+            for row in reader:
+                normalized_row = {
+                    key.strip().lower(): value
+                    for key, value in row.items()
+                    if isinstance(key, str)
+                }
+                if not any(
+                    isinstance(value, str) and value.strip()
+                    for value in normalized_row.values()
+                ):
+                    continue
+                prompt_inputs.append(
+                    PromptInput(
+                        image=normalized_row.get("image"),
+                        github_repo=normalized_row.get("github_repo"),
+                        prompt=normalized_row.get("prompt") or "",
+                    )
+                )
+        return prompt_inputs
+
+
+def load_config(path: Path) -> Config:
+    """Load configuration from YAML file.
+
+    Environment variables (for testing):
+        TEICH_MODEL - Override model ID
+        TEICH_BASE_URL - Override API base URL
+        TEICH_API_KEY - Override API key
+        TEICH_PROVIDER - Override provider name
+    """
+    return Config.from_yaml(path)