synth-ai 0.2.14__py3-none-any.whl → 0.4.4__py3-none-any.whl

synth_ai/sdk/api/train/graphgen.py

@@ -0,0 +1,1102 @@
+"""First-class SDK API for GraphGen (Graph Opt).
+
+**Status:** Alpha
+
+GraphGen is a simplified "Workflows API" for prompt optimization that:
+- Uses a simple JSON dataset format (GraphGenTaskSet) instead of TOML configs
+- Auto-generates task apps from the dataset (no user-managed task apps)
+- Has built-in verifier configurations (rubric, contrastive, gold_examples)
+- Wraps GEPA internally for the actual optimization
+
+Example CLI usage:
+    uvx synth-ai train --type graphgen --dataset my_tasks.json --poll
+
+Example SDK usage:
+    from synth_ai.sdk.api.train.graphgen import GraphGenJob
+    from synth_ai.sdk.api.train.graphgen_models import GraphGenTaskSet, GraphGenTask
+
+    # From a dataset file
+    job = GraphGenJob.from_dataset("my_tasks.json")
+    job.submit()
+    result = job.stream_until_complete()
+    print(f"Best score: {result.get('best_score')}")
+
+    # Or programmatically
+    dataset = GraphGenTaskSet(
+        metadata=GraphGenTaskSetMetadata(name="My Tasks"),
+        tasks=[GraphGenTask(id="t1", input={"question": "What is 2+2?"})],
+        gold_outputs=[GraphGenGoldOutput(output={"answer": "4"}, task_id="t1")],
+    )
+    job = GraphGenJob.from_dataset(dataset, policy_model="gpt-4o-mini", problem_spec="You are a helpful assistant.")
+    job.submit()
+"""
+
+from __future__ import annotations
+
+import asyncio
+import json
+import os
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import Any, Callable, Dict, List, Literal, Optional, Sequence, cast
+
+from synth_ai.core.telemetry import log_info
+
+from .graphgen_models import (
+    GraphGenJobConfig,
+    GraphGenTaskSet,
+    load_graphgen_taskset,
+    parse_graphgen_taskset,
+    SessionTraceInput,
+    GraphGenGraphVerifierResponse,
+)
+from .utils import ensure_api_base, http_get, http_post
+
+
+@dataclass
+class GraphGenJobResult:
+    """Result from a GraphGen job.
+
+    Contains the final status and results of a completed GraphGen workflow
+    optimization job, including the best score and snapshot ID for the
+    optimized graph.
+
+    Attributes:
+        graphgen_job_id: Unique identifier for the GraphGen job (e.g.,
+            "graphgen_abc123def456").
+        status: Current job status. One of: "pending", "running", "succeeded",
+            "failed", "cancelled".
+        best_score: Best evaluation score achieved during optimization. Higher
+            is better. None if job hasn't completed successfully.
+        best_snapshot_id: ID of the graph snapshot with the best score. Use this
+            to download or deploy the optimized graph.
+        error: Error message if the job failed, None otherwise.
+        dataset_name: Name of the dataset used for optimization.
+        task_count: Number of tasks in the dataset.
+        graph_evolve_job_id: ID of the underlying graph evolution job, if applicable.
+
+    Example:
+        >>> result = job.get_result()
+        >>> if result.status == "succeeded":
+        ...     print(f"Best score: {result.best_score}")
+        ...     print(f"Snapshot ID: {result.best_snapshot_id}")
+    """
+
+    graphgen_job_id: str
+    status: str
+    best_score: Optional[float] = None
+    best_snapshot_id: Optional[str] = None
+    error: Optional[str] = None
+    dataset_name: Optional[str] = None
+    task_count: Optional[int] = None
+    graph_evolve_job_id: Optional[str] = None
+
+
+@dataclass
+class GraphGenSubmitResult:
+    """Result from submitting a GraphGen job.
+
+    Returned immediately after job submission with initial job metadata
+    and configuration details.
+
+    Attributes:
+        graphgen_job_id: Unique identifier for the GraphGen job.
+        status: Initial job status (typically "pending" or "running").
+        dataset_name: Name of the dataset being used for optimization.
+        task_count: Number of tasks in the dataset.
+        rollout_budget: Total number of rollouts (evaluations) budgeted for
+            this optimization job.
+        policy_model: Name of the LLM model being used for the policy
+            (e.g., "gpt-4o-mini", "claude-3-5-sonnet").
+        verifier_mode: Evaluation mode being used. One of: "rubric", "contrastive",
+            "gold_examples", "verifier_graph".
+        graph_evolve_job_id: ID of the underlying graph evolution job, if applicable.
+
+    Example:
+        >>> submit_result = job.submit()
+        >>> print(f"Job {submit_result.graphgen_job_id} started")
+        >>> print(f"Optimizing {submit_result.task_count} tasks with {submit_result.rollout_budget} rollouts")
+    """
+
+    graphgen_job_id: str
+    status: str
+    dataset_name: str
+    task_count: int
+    rollout_budget: int
+    policy_model: str
+    verifier_mode: str
+    graph_evolve_job_id: Optional[str] = None
+
+
+class GraphGenJob:
+    """High-level SDK class for running GraphGen workflow optimization jobs.
+
+    GraphGen (Graph Opt) provides a simplified API for
+    graph/workflow optimization that doesn't require users to manage task apps.
+
+    Key differences from PromptLearningJob:
+    - Uses JSON dataset format (GraphGenTaskSet) instead of TOML configs
+    - No task app management required - GraphGen builds it internally
+    - Built-in verifier modes (rubric, contrastive, gold_examples)
+    - Graph-first: trains multi-node workflows by default (Graph-GEPA)
+    - Public graph downloads are redacted `.txt` exports only
+    - Simpler configuration with sensible defaults
+
+    Example:
+        >>> from synth_ai.sdk.api.train.graphgen import GraphGenJob
+        >>>
+        >>> # Create job from dataset file
+        >>> job = GraphGenJob.from_dataset(
+        ...     dataset="my_tasks.json",
+        ...     policy_model="gpt-4o-mini",
+        ...     rollout_budget=100,
+        ... )
+        >>>
+        >>> # Train a verifier graph
+        >>> verifier_job = GraphGenJob.from_dataset(
+        ...     dataset="verifier_dataset.json",
+        ...     graph_type="verifier",
+        ...     rollout_budget=200,
+        ... )
+        >>>
+        >>> # Train an RLM graph (massive context via tools)
+        >>> rlm_job = GraphGenJob.from_dataset(
+        ...     dataset="rlm_dataset.json",
+        ...     graph_type="rlm",
+        ...     configured_tools=[
+        ...         {"name": "materialize_context", "kind": "rlm_materialize", "stateful": True},
+        ...         {"name": "local_grep", "kind": "rlm_local_grep", "stateful": False},
+        ...     ],
+        ...     rollout_budget=100,
+        ... )
+        >>>
+        >>> # Submit and stream
+        >>> job.submit()
+        >>> result = job.stream_until_complete(timeout=3600.0)
+        >>> print(f"Best score: {result.get('best_score')}")
+        >>> 
+        >>> # Download public graph export
+        >>> export_txt = job.download_graph_txt()
+        >>> print(export_txt)
+        >>>
+        >>> # Run inference with optimized prompt
+        >>> output = job.run_inference({"question": "What is 2+2?"})
+        >>>
+        >>> # Run verifier with optimized verifier graph
+        >>> verification = verifier_job.run_verifier(trace_data)
+        >>> print(f"Outcome reward: {verification.outcome_reward}")
+    """
+
+    def __init__(
+        self,
+        *,
+        dataset: GraphGenTaskSet,
+        config: GraphGenJobConfig,
+        backend_url: str,
+        api_key: str,
+        auto_start: bool = True,
+        metadata: Optional[Dict[str, Any]] = None,
+    ) -> None:
+        """Initialize an GraphGen job.
+
+        Args:
+            dataset: The GraphGenTaskSet containing tasks and evaluation config
+            config: Job configuration (policy model, budget, etc.)
+            backend_url: Backend API URL
+            api_key: Synth API key
+            auto_start: Whether to start the job immediately after creation
+            metadata: Additional metadata for the job
+        """
+        self.dataset = dataset
+        self.config = config
+        self.backend_url = ensure_api_base(backend_url)
+        self.api_key = api_key
+        self.auto_start = auto_start
+        self.metadata = metadata or {}
+
+        self._graphgen_job_id: Optional[str] = None
+        self._graph_evolve_job_id: Optional[str] = None
+        self._submit_result: Optional[GraphGenSubmitResult] = None
+
+    @classmethod
+    def from_dataset(
+        cls,
+        dataset: str | Path | Dict[str, Any] | GraphGenTaskSet,
+        *,
+        graph_type: Literal["policy", "verifier", "rlm"] = "policy",
+        policy_model: str = "gpt-4o-mini",
+        rollout_budget: int = 100,
+        proposer_effort: Literal["low", "medium", "high"] = "medium",
+        verifier_model: Optional[str] = None,
+        verifier_provider: Optional[str] = None,
+        population_size: int = 4,
+        num_generations: Optional[int] = None,
+        problem_spec: Optional[str] = None,
+        target_llm_calls: Optional[int] = None,
+        configured_tools: Optional[List[Dict[str, Any]]] = None,
+        backend_url: Optional[str] = None,
+        api_key: Optional[str] = None,
+        auto_start: bool = True,
+        metadata: Optional[Dict[str, Any]] = None,
+    ) -> GraphGenJob:
+        """Create an GraphGen job from a dataset.
+
+        Args:
+            dataset: Dataset as file path, dict, or GraphGenTaskSet object
+            graph_type: Type of graph to train:
+                - "policy": Maps inputs to outputs (default).
+                - "verifier": Verifies/scores traces (requires verifier-compliant dataset).
+                - "rlm": Recursive Language Model - handles massive contexts via tool-based search
+                  and recursive LLM calls. Requires configured_tools parameter.
+            policy_model: Model to use for policy inference
+            rollout_budget: Total number of rollouts for optimization
+            proposer_effort: Proposer effort level ("medium" or "high").
+                "low" is not allowed as gpt-4.1-mini is too weak for graph generation.
+            verifier_model: Override verifier model from dataset
+            verifier_provider: Override verifier provider from dataset
+            population_size: Population size for GEPA
+            num_generations: Number of generations (auto-calculated if not specified)
+            problem_spec: Detailed problem specification for the graph proposer.
+                Include domain-specific info like valid output labels for classification.
+            target_llm_calls: Target number of LLM calls for the graph (1-10).
+                Controls how many LLM nodes the graph should use. Defaults to 5.
+            configured_tools: Optional list of tool bindings for RLM graphs.
+                Required for graph_type="rlm". Each tool should be a dict with 'name', 'kind', and 'stateful'.
+                Example: [{'name': 'materialize_context', 'kind': 'rlm_materialize', 'stateful': True}]
+            backend_url: Backend API URL (defaults to env or production)
+            api_key: API key (defaults to SYNTH_API_KEY env var)
+            auto_start: Whether to start the job immediately
+            metadata: Additional metadata for the job
+
+        Returns:
+            GraphGenJob instance
+
+        Example:
+            >>> # From file
+            >>> job = GraphGenJob.from_dataset("tasks.json")
+            >>>
+            >>> # From dict
+            >>> job = GraphGenJob.from_dataset({
+            ...     "metadata": {"name": "My Tasks"},
+            ...     "tasks": [{"id": "t1", "input": {"q": "Hi"}}],
+            ... }, problem_spec="You are helpful.")
+            >>>
+            >>> # From GraphGenTaskSet object
+            >>> job = GraphGenJob.from_dataset(my_taskset, policy_model="gpt-4o")
+        """
+        from synth_ai.core.env import get_backend_from_env
+
+        # Parse dataset
+        if isinstance(dataset, (str, Path)):
+            parsed_dataset = load_graphgen_taskset(dataset)
+        elif isinstance(dataset, dict):
+            parsed_dataset = parse_graphgen_taskset(dataset)
+        elif isinstance(dataset, GraphGenTaskSet):
+            parsed_dataset = dataset
+        else:
+            raise TypeError(
+                f"dataset must be a file path, dict, or GraphGenTaskSet, got {type(dataset)}"
+            )
+
+        # Resolve backend URL
+        if not backend_url:
+            backend_url = os.environ.get("BACKEND_BASE_URL", "").strip()
+            if not backend_url:
+                base, _ = get_backend_from_env()
+                backend_url = f"{base}/api" if not base.endswith("/api") else base
+
+        # Resolve API key
+        if not api_key:
+            api_key = os.environ.get("SYNTH_API_KEY")
+            if not api_key:
+                raise ValueError(
+                    "api_key is required (provide explicitly or set SYNTH_API_KEY env var)"
+                )
+
+        # Build config
+        config = GraphGenJobConfig(
+            graph_type=graph_type,
+            policy_model=policy_model,
+            rollout_budget=rollout_budget,
+            proposer_effort=proposer_effort,
+            verifier_model=verifier_model,
+            verifier_provider=verifier_provider,
+            population_size=population_size,
+            num_generations=num_generations,
+            problem_spec=problem_spec,
+            target_llm_calls=target_llm_calls,
+            configured_tools=configured_tools,
+        )
+
+        return cls(
+            dataset=parsed_dataset,
+            config=config,
+            backend_url=backend_url,
+            api_key=api_key,
+            auto_start=auto_start,
+            metadata=metadata,
+        )
+
+    @classmethod
+    def from_job_id(
+        cls,
+        job_id: str,
+        backend_url: Optional[str] = None,
+        api_key: Optional[str] = None,
+    ) -> GraphGenJob:
+        """Resume an existing GraphGen job by ID.
+
+        Args:
+            job_id: GraphGen job ID ("graphgen_*") or underlying GEPA job ID ("pl_*")
+            backend_url: Backend API URL (defaults to env or production)
+            api_key: API key (defaults to SYNTH_API_KEY env var)
+
+        Returns:
+            GraphGenJob instance for the existing job
+        """
+        from synth_ai.core.env import get_backend_from_env
+
+        # Resolve backend URL
+        if not backend_url:
+            backend_url = os.environ.get("BACKEND_BASE_URL", "").strip()
+            if not backend_url:
+                base, _ = get_backend_from_env()
+                backend_url = f"{base}/api" if not base.endswith("/api") else base
+
+        # Resolve API key
+        if not api_key:
+            api_key = os.environ.get("SYNTH_API_KEY")
+            if not api_key:
+                raise ValueError(
+                    "api_key is required (provide explicitly or set SYNTH_API_KEY env var)"
+                )
+
+        # Create minimal instance - dataset will be fetched from backend if needed
+        # For now, create a placeholder dataset
+        from .graphgen_models import GraphGenTaskSetMetadata, GraphGenTask
+        placeholder_dataset = GraphGenTaskSet(
+            metadata=GraphGenTaskSetMetadata(name="(resumed job)"),
+            tasks=[GraphGenTask(id="placeholder", input={})],
+        )
+
+        job = cls(
+            dataset=placeholder_dataset,
+            config=GraphGenJobConfig(),
+            backend_url=backend_url,
+            api_key=api_key,
+            auto_start=False,
+        )
+
+        # Accept GraphGen/GraphGen or graph_evolve/GEPA job IDs - backend handles resolution internally
+        valid_prefixes = ("graphgen_", "graphgen_", "graph_evolve_", "graph_evolve_", "pl_")
+        if not any(job_id.startswith(p) for p in valid_prefixes):
+            raise ValueError(
+                f"Unsupported job ID format: {job_id!r}. "
+                f"Expected one of: {valid_prefixes}"
+            )
+        job._graphgen_job_id = job_id
+        if job_id.startswith("pl_"):
+            job._graph_evolve_job_id = job_id
+        return job
+
+    @classmethod
+    def from_graph_evolve_job_id(
+        cls,
+        graph_evolve_job_id: str,
+        backend_url: Optional[str] = None,
+        api_key: Optional[str] = None,
+    ) -> GraphGenJob:
+        """Alias for resuming an GraphGen job from a GEPA job ID."""
+        return cls.from_job_id(graph_evolve_job_id, backend_url=backend_url, api_key=api_key)
+
+    @property
+    def job_id(self) -> Optional[str]:
+        """Get the GraphGen job ID (None if not yet submitted)."""
+        return self._graphgen_job_id
+
+    @property
+    def graph_evolve_job_id(self) -> Optional[str]:
+        """Get the underlying GEPA job ID if known."""
+        if self._graph_evolve_job_id:
+            return self._graph_evolve_job_id
+        if self._submit_result and self._submit_result.graph_evolve_job_id:
+            return self._submit_result.graph_evolve_job_id
+        return None
+
+    def _build_payload(self) -> Dict[str, Any]:
+        """Build the job creation payload."""
+        # Merge config num_generations into metadata if provided
+        metadata = dict(self.metadata) if self.metadata else {}
+        if self.config.num_generations is not None:
+            metadata["num_generations"] = self.config.num_generations
+        if self.config.population_size != 4:  # Only include if non-default
+            metadata["population_size"] = self.config.population_size
+        if self.config.num_parents != 2:
+            metadata["num_parents"] = self.config.num_parents
+        if self.config.evaluation_seeds is not None:
+            metadata["evaluation_seeds"] = self.config.evaluation_seeds
+
+        # Extract eval/feedback sample sizes from metadata as direct fields
+        eval_sample_size = metadata.pop("eval_sample_size", None)
+        feedback_sample_size = metadata.pop("feedback_sample_size", None)
+
+        # Build dataset dict and ensure it has an initial_prompt to satisfy legacy backend validation
+        # GraphGen is graph-first and doesn't really use this, so we use a placeholder or problem_spec
+        dataset_dict = self.dataset.model_dump()
+        if "initial_prompt" not in dataset_dict:
+            dataset_dict["initial_prompt"] = self.config.problem_spec or "Optimizing prompt graph..."
+
+        payload: Dict[str, Any] = {
+            "dataset": dataset_dict,
+            "initial_prompt": None,  # Top-level initial_prompt is ignored in favor of dataset.initial_prompt
+            "graph_type": self.config.graph_type,
+            "policy_model": self.config.policy_model,
+            "policy_provider": self.config.policy_provider,
+            "rollout_budget": self.config.rollout_budget,
+            "proposer_effort": self.config.proposer_effort,
+            "verifier_model": self.config.verifier_model,
+            "verifier_provider": self.config.verifier_provider,
+            "problem_spec": self.config.problem_spec,
+            "target_llm_calls": self.config.target_llm_calls,
+            "configured_tools": self.config.configured_tools,
+            "eval_sample_size": eval_sample_size,
+            "feedback_sample_size": feedback_sample_size,
+            "metadata": metadata,
+            "auto_start": self.auto_start,
+        }
+
+        # Strip unset optional fields so we don't send nulls to strict backends.
+        if payload.get("eval_sample_size") is None:
+            payload.pop("eval_sample_size", None)
+        if payload.get("feedback_sample_size") is None:
+            payload.pop("feedback_sample_size", None)
+        if payload.get("policy_provider") is None:
+            payload.pop("policy_provider", None)
+        if payload.get("verifier_model") is None:
+            payload.pop("verifier_model", None)
+        if payload.get("verifier_provider") is None:
+            payload.pop("verifier_provider", None)
+        if payload.get("problem_spec") is None:
+            payload.pop("problem_spec", None)
+        if payload.get("target_llm_calls") is None:
+            payload.pop("target_llm_calls", None)
+        if payload.get("configured_tools") is None:
+            payload.pop("configured_tools", None)
+
+        return payload
+
+    def submit(self) -> GraphGenSubmitResult:
+        """Submit the job to the backend.
+
+        Returns:
+            GraphGenSubmitResult with job IDs and initial status
+
+        Raises:
+            RuntimeError: If job submission fails
+        """
+        from .graphgen_validators import validate_graphgen_job_config
+
+        ctx: Dict[str, Any] = {"dataset_name": self.dataset.metadata.name}
+        log_info("GraphGenJob.submit invoked", ctx=ctx)
+
+        if self._graphgen_job_id:
+            raise RuntimeError(f"Job already submitted: {self._graphgen_job_id}")
+
+        # Validate config + dataset before expensive API call.
+        validate_graphgen_job_config(self.config, self.dataset)
+
+        payload = self._build_payload()
+
+        # Submit job - use /graphgen/jobs endpoint
+        create_url = f"{self.backend_url}/graphgen/jobs"
+        headers = {
+            "X-API-Key": self.api_key,
+            "Content-Type": "application/json",
+        }
+
+        import logging
+        logger = logging.getLogger(__name__)
+        logger.debug(f"Submitting GraphGen job to: {create_url}")
+
+        resp = http_post(create_url, headers=headers, json_body=payload, timeout=180.0)
+
+        if resp.status_code not in (200, 201):
+            error_msg = f"Job submission failed with status {resp.status_code}: {resp.text[:500]}"
+            if resp.status_code == 404:
+                error_msg += (
+                    f"\n\nPossible causes:"
+                    f"\n1. Backend route /api/graphgen/jobs not registered"
+                    f"\n2. GraphGen feature may not be enabled on this backend"
+                    f"\n3. Verify backend is running at: {self.backend_url}"
+                )
+            raise RuntimeError(error_msg)
+
+        try:
+            js = resp.json()
+        except Exception as e:
+            raise RuntimeError(f"Failed to parse response: {e}") from e
+
+        self._graphgen_job_id = js.get("graphgen_job_id")
+
+        if not self._graphgen_job_id:
+            raise RuntimeError("Response missing graphgen_job_id")
+
+        self._graph_evolve_job_id = js.get("graph_evolve_job_id")
+
+        self._submit_result = GraphGenSubmitResult(
+            graphgen_job_id=self._graphgen_job_id,
+            status=js.get("status", "queued"),
+            dataset_name=js.get("dataset_name", self.dataset.metadata.name),
+            task_count=js.get("task_count", len(self.dataset.tasks)),
+            rollout_budget=js.get("rollout_budget", self.config.rollout_budget),
+            policy_model=js.get("policy_model", self.config.policy_model),
+            verifier_mode=js.get("verifier_mode", self.dataset.verifier_config.mode),
+            graph_evolve_job_id=self._graph_evolve_job_id,
+        )
+
+        ctx["graphgen_job_id"] = self._graphgen_job_id
+        log_info("GraphGenJob.submit completed", ctx=ctx)
+
+        return self._submit_result
+
+    def get_status(self) -> Dict[str, Any]:
+        """Get current job status.
+
+        Returns:
+            Job status dictionary containing 'status', 'best_score', etc.
+
+        Raises:
+            RuntimeError: If job hasn't been submitted yet or API call fails.
+        """
+        if not self.job_id:
+            raise RuntimeError("Job not yet submitted. Call submit() first.")
+
+        url = f"{self.backend_url}/graphgen/jobs/{self.job_id}"
+        headers = {
+            "X-API-Key": self.api_key,
+        }
+
+        resp = http_get(url, headers=headers, timeout=30.0)
+
+        if resp.status_code != 200:
+            raise RuntimeError(
+                f"Failed to get job status: {resp.status_code} - {resp.text[:500]}"
+            )
+
+        data: Dict[str, Any] = resp.json()
+        gepa_id = data.get("graph_evolve_job_id")
+        if gepa_id:
+            self._graph_evolve_job_id = gepa_id
+        return data
+
+    def start(self) -> Dict[str, Any]:
+        """Start a queued GraphGen job.
+
+        This is only needed if the job was created with auto_start=False or ended up queued.
+
+        Returns:
+            Updated job status dictionary.
+        """
+        if not self.job_id:
+            raise RuntimeError("Job not yet submitted. Call submit() first.")
+
+        url = f"{self.backend_url}/graphgen/jobs/{self.job_id}/start"
+        headers = {
+            "X-API-Key": self.api_key,
+            "Content-Type": "application/json",
+        }
+
+        resp = http_post(url, headers=headers, json_body=None, timeout=60.0)
+        if resp.status_code != 200:
+            raise RuntimeError(
+                f"Failed to start job: {resp.status_code} - {resp.text[:500]}"
+            )
+        data: Dict[str, Any] = resp.json()
+        if self._submit_result and "status" in data:
+            self._submit_result.status = data.get("status", self._submit_result.status)
+        return data
+
+    def get_events(self, *, since_seq: int = 0, limit: int = 1000) -> Dict[str, Any]:
+        """Fetch events for this GraphGen job.
+
+        Args:
+            since_seq: Return events with sequence number greater than this.
+            limit: Maximum number of events to return.
+
+        Returns:
+            Backend envelope: {"events": [...], "has_more": bool, "next_seq": int}.
+        """
+        if not self.job_id:
+            raise RuntimeError("Job not yet submitted. Call submit() first.")
+
+        base = f"{self.backend_url}/graphgen/jobs/{self.job_id}/events"
+        url = f"{base}?since_seq={since_seq}&limit={limit}"
+        headers = {"X-API-Key": self.api_key}
+
+        resp = http_get(url, headers=headers, timeout=30.0)
+        if resp.status_code != 200:
+            raise RuntimeError(
+                f"Failed to get events: {resp.status_code} - {resp.text[:500]}"
+            )
+        return cast(Dict[str, Any], resp.json())
+
+    def get_metrics(
+        self,
+        *,
+        name: Optional[str] = None,
+        after_step: Optional[int] = None,
+        limit: int = 500,
+        run_id: Optional[str] = None,
+    ) -> Dict[str, Any]:
+        """Fetch metrics for this GraphGen job.
+
+        Args:
+            name: Optional metric name filter.
+            after_step: Optional step filter.
+            limit: Maximum number of metrics to return.
+            run_id: Optional run identifier filter.
+
+        Returns:
+            Dictionary containing 'metrics' list.
+        """
+        if not self.job_id:
+            raise RuntimeError("Job not yet submitted. Call submit() first.")
+
+        from urllib.parse import urlencode
+
+        params: Dict[str, Any] = {"limit": limit}
+        if name is not None:
+            params["name"] = name
+        if after_step is not None:
+            params["after_step"] = after_step
+        if run_id is not None:
+            params["run_id"] = run_id
+
+        qs = urlencode(params)
+        url = f"{self.backend_url}/graphgen/jobs/{self.job_id}/metrics?{qs}"
+        headers = {"X-API-Key": self.api_key}
+
+        resp = http_get(url, headers=headers, timeout=30.0)
+        if resp.status_code != 200:
+            raise RuntimeError(
+                f"Failed to get metrics: {resp.status_code} - {resp.text[:500]}"
+            )
+        return cast(Dict[str, Any], resp.json())
+
+    def stream_until_complete(
+        self,
+        *,
+        timeout: float = 3600.0,
+        interval: float = 5.0,
+        handlers: Optional[Sequence[Any]] = None,
+        on_event: Optional[Callable[[Dict[str, Any]], None]] = None,
+    ) -> Dict[str, Any]:
+        """Stream job events until completion using Server-Sent Events (SSE).
+
+        This method connects to the backend SSE stream and processes events in real-time
+        until the job reaches a terminal state (completed, failed, or cancelled).
+
+        Events include:
+        - job_started: Job execution began
+        - generation_started: New generation of candidates started
+        - candidate_evaluated: A candidate graph was evaluated
+        - generation_completed: Generation finished
+        - optimization_completed: Job finished successfully
+        - job_failed: Job encountered an error
+
+        Args:
+            timeout: Maximum seconds to wait for completion
+            interval: Seconds between status checks (for SSE reconnects)
+            handlers: Optional StreamHandler instances for custom event handling.
+                Defaults to GraphGenHandler which provides formatted CLI output.
+            on_event: Optional callback function called on each event.
+                Receives the event dict as argument.
+
+        Returns:
+            Final job status dictionary containing 'status', 'best_score', etc.
+
+        Raises:
+            RuntimeError: If job hasn't been submitted yet
+            TimeoutError: If timeout exceeded before job completion
+
+        Example:
+            >>> job.submit()
+            >>> result = job.stream_until_complete(timeout=1800.0)
+            >>> print(f"Best score: {result.get('best_score')}")
+        """
+        if not self.job_id:
+            raise RuntimeError("Job not yet submitted. Call submit() first.")
+
+        from synth_ai.sdk.streaming import (
+            GraphGenHandler,
+            JobStreamer,
+            StreamConfig,
+            StreamEndpoints,
+            StreamType,
+        )
+
+        # Build stream config
+        config = StreamConfig(
+            enabled_streams={StreamType.STATUS, StreamType.EVENTS, StreamType.METRICS},
+            max_events_per_poll=500,
+            deduplicate=True,
+        )
+
+        # Use provided handlers or default CLI handler
+        if handlers is None:
+            handlers = [GraphGenHandler()]
+
+        # Create streamer with GraphGen endpoints
+        # Backend handles GraphGen → GEPA resolution internally via job_relationships table
+        streamer = JobStreamer(
+            base_url=self.backend_url,
+            api_key=self.api_key,
+            job_id=self.job_id,  # Only GraphGen job ID - backend resolves to GEPA internally
+            endpoints=StreamEndpoints.graphgen(self.job_id),
+            config=config,
+            handlers=list(handlers),
+            interval_seconds=interval,
+            timeout_seconds=timeout,
+        )
+
+        # Run streaming
+        final_status = asyncio.run(streamer.stream_until_terminal())
+
+        return final_status
+
+    def poll_until_complete(
+        self,
+        *,
+        timeout: float = 3600.0,
+        interval: float = 5.0,
+        progress: bool = False,
+        on_status: Optional[Callable[[Dict[str, Any]], None]] = None,
+    ) -> Dict[str, Any]:
+        """Poll job until it reaches a terminal state.
+
+        Similar to PromptLearningJob.poll_until_complete(), this method polls
+        the backend periodically instead of using SSE streaming. Useful for
+        notebooks and environments where SSE may not work reliably.
+
+        Args:
+            timeout: Maximum seconds to wait (default: 3600 = 1 hour)
+            interval: Seconds between poll attempts (default: 5)
+            progress: If True, print status updates during polling (useful for notebooks)
+            on_status: Optional callback called on each status update
+
+        Returns:
+            Final job status dictionary containing 'status', 'best_score', etc.
+
+        Raises:
+            RuntimeError: If job hasn't been submitted yet
+            TimeoutError: If timeout is exceeded
+
+        Example:
+            >>> result = job.poll_until_complete(progress=True)
+            [00:15] running | score: 0.72
+            [00:30] running | score: 0.78
+            [00:45] succeeded | score: 0.85
+        """
+        if not self.job_id:
+            raise RuntimeError("Job not yet submitted. Call submit() first.")
+
+        import time
+
+        start_time = time.time()
+        elapsed = 0.0
+        last_data: Dict[str, Any] = {}
+
+        while elapsed <= timeout:
+            try:
+                status_data = self.get_status()
+                last_data = dict(status_data) if isinstance(status_data, dict) else {}
+
+                status = last_data.get("status", "unknown")
+                best_score = last_data.get("best_score")
+
+                # Progress output
+                if progress:
+                    mins, secs = divmod(int(elapsed), 60)
+                    score_str = f"score: {best_score:.2f}" if best_score is not None else "score: --"
+                    print(f"[{mins:02d}:{secs:02d}] {status} | {score_str}")
+
+                # Callback for custom handling
+                if on_status:
+                    on_status(last_data)
+
+                # Check terminal state
+                if status in ("succeeded", "completed", "failed", "error", "cancelled"):
+                    return last_data
+
+                # Sleep before next poll
+                time.sleep(interval)
+                elapsed = time.time() - start_time
+
+            except Exception as e:
+                # On error, continue polling (might be transient network issue)
+                import logging
+                logger = logging.getLogger(__name__)
+                logger.warning(f"Error polling job status: {e}")
+                time.sleep(interval)
+                elapsed = time.time() - start_time
+
+        # Timeout exceeded
+        raise TimeoutError(
+            f"Job {self.job_id} did not complete within {timeout}s timeout. "
+            f"Current status: {last_data.get('status', 'unknown')}"
+        )
+
+    def download_prompt(self) -> str:
+        """Download the optimized prompt from a completed job.
+
+        For graph-first jobs, prefer `download_graph_txt()`; this method is
+        mainly useful for legacy single-node prompt workflows.
+
+        Returns:
+            Optimized prompt text
+
+        Raises:
+            RuntimeError: If job hasn't been submitted or isn't complete
+        """
+        if not self.job_id:
+            raise RuntimeError("Job not yet submitted. Call submit() first.")
+
+        url = f"{self.backend_url}/graphgen/jobs/{self.job_id}/download"
+        headers = {
+            "X-API-Key": self.api_key,
+        }
+
+        resp = http_get(url, headers=headers, timeout=30.0)
+
+        if resp.status_code != 200:
+            raise RuntimeError(
+                f"Failed to download prompt: {resp.status_code} - {resp.text[:500]}"
+            )
+
+        data = resp.json()
+        return data.get("prompt", "")
+
+    def download_graph_txt(self) -> str:
+        """Download a PUBLIC (redacted) graph export for a completed job.
+
+        Graph-first GraphGen jobs produce multi-node graphs. The internal graph
+        YAML/spec is proprietary and never exposed. This helper downloads the
+        `.txt` export from:
+            GET /api/graphgen/jobs/{job_id}/graph.txt
+        """
+        if not self.job_id:
+            raise RuntimeError("Job not yet submitted. Call submit() first.")
+
+        url = f"{self.backend_url}/graphgen/jobs/{self.job_id}/graph.txt"
+        headers = {"X-API-Key": self.api_key}
+
+        resp = http_get(url, headers=headers, timeout=30.0)
+        if resp.status_code != 200:
+            raise RuntimeError(
+                f"Failed to download graph export: {resp.status_code} - {resp.text[:500]}"
+            )
+        return resp.text
+
+    def run_inference(
+        self,
+        input_data: Dict[str, Any],
+        *,
+        model: Optional[str] = None,
+        prompt_snapshot_id: Optional[str] = None,
+        graph_snapshot_id: Optional[str] = None,
+        timeout: float = 120.0,
+    ) -> Dict[str, Any]:
+        """Run inference with the optimized graph/workflow.
+
+        Args:
+            input_data: Input data matching the task format
+            model: Override model (default: use job's policy model)
+            prompt_snapshot_id: Legacy alias for selecting a specific snapshot.
+            graph_snapshot_id: Specific GraphSnapshot to use (default: best).
+                Preferred for graph-first jobs. If provided, it is sent as
+                `prompt_snapshot_id` for backward-compatible backend routing.
+            timeout: Request timeout in seconds (default: 120.0 = 2 minutes for image generation tasks)
+
+        Returns:
+            Output dictionary containing 'output', 'usage', etc.
+
+        Raises:
+            RuntimeError: If job hasn't been submitted or inference fails.
+            ValueError: If both prompt_snapshot_id and graph_snapshot_id are provided.
+        """
+        if not self.job_id:
+            raise RuntimeError("Job not yet submitted. Call submit() first.")
+
+        if prompt_snapshot_id and graph_snapshot_id:
+            raise ValueError("Provide only one of prompt_snapshot_id or graph_snapshot_id.")
+
+        url = f"{self.backend_url}/graphgen/graph/completions"
+        headers = {
+            "X-API-Key": self.api_key,
+            "Content-Type": "application/json",
+        }
+
+        payload: Dict[str, Any] = {
+            "job_id": self.job_id,
+            "input": input_data,
+        }
+        if model:
+            payload["model"] = model
+        snapshot_id = graph_snapshot_id or prompt_snapshot_id
+        if snapshot_id:
+            payload["prompt_snapshot_id"] = snapshot_id
+
+        # Use longer timeout for image generation tasks (can take 2-3 minutes)
+        resp = http_post(url, headers=headers, json_body=payload, timeout=timeout)
+
+        if resp.status_code != 200:
+            raise RuntimeError(
+                f"Inference failed: {resp.status_code} - {resp.text[:500]}"
+            )
+
+        return cast(Dict[str, Any], resp.json())
+
+    def run_inference_output(
+        self,
+        input_data: Dict[str, Any],
+        *,
+        model: Optional[str] = None,
+        prompt_snapshot_id: Optional[str] = None,
+        graph_snapshot_id: Optional[str] = None,
+    ) -> Any:
+        """Convenience wrapper returning only the model output."""
+        result = self.run_inference(
+            input_data,
+            model=model,
+            prompt_snapshot_id=prompt_snapshot_id,
+            graph_snapshot_id=graph_snapshot_id,
+        )
+        if isinstance(result, dict):
+            return result.get("output")
+        return None
+
+    def run_verifier(
+        self,
+        session_trace: Dict[str, Any] | SessionTraceInput,
+        *,
+        context: Optional[Dict[str, Any]] = None,
+        prompt_snapshot_id: Optional[str] = None,
+        graph_snapshot_id: Optional[str] = None,
+    ) -> GraphGenGraphVerifierResponse:
+        """Run a verifier graph on an execution trace.
+
+        This method is specifically for graphs trained with graph_type=\"verifier\".
+        It accepts a V3 trace and returns structured rewards.
+
+        Args:
+            session_trace: V3 session trace to evaluate. Can be a dict or SessionTraceInput.
+            context: Additional context for evaluation (e.g., rubric overrides, task description).
+            prompt_snapshot_id: Specific snapshot to use (default: best).
+            graph_snapshot_id: Specific GraphSnapshot to use (default: best).
+                Preferred for graph-first jobs.
+
+        Returns:
+            GraphGenGraphVerifierResponse containing structured rewards.
+
+        Raises:
+            RuntimeError: If job hasn't been submitted or inference fails.
+        """
+        if not self.job_id:
+            raise RuntimeError("Job not yet submitted. Call submit() first.")
+
+        if prompt_snapshot_id and graph_snapshot_id:
+            raise ValueError("Provide only one of prompt_snapshot_id or graph_snapshot_id.")
+
+        url = f"{self.backend_url}/graphgen/graph/verifier"
+        headers = {
+            "X-API-Key": self.api_key,
+            "Content-Type": "application/json",
+        }
+
+        # Convert trace to dict if it's a Pydantic model
+        if isinstance(session_trace, SessionTraceInput):
+            session_trace_data = session_trace.model_dump(mode="json")
+        else:
+            session_trace_data = session_trace
+
+        payload = {
+            "job_id": self.job_id,
+            "session_trace": session_trace_data,
+            "context": context,
+        }
+        
+        snapshot_id = graph_snapshot_id or prompt_snapshot_id
+        if snapshot_id:
+            payload["prompt_snapshot_id"] = snapshot_id
+
+        resp = http_post(url, headers=headers, json_body=payload, timeout=120.0)
+
+        if resp.status_code != 200:
+            raise RuntimeError(
+                f"Verifier inference failed: {resp.status_code} - {resp.text[:500]}"
+            )
+
+        return GraphGenGraphVerifierResponse.model_validate(resp.json())
+
+    def get_graph_record(
+        self,
+        *,
+        prompt_snapshot_id: Optional[str] = None,
+        graph_snapshot_id: Optional[str] = None,
+    ) -> Dict[str, Any]:
+        """Get the optimized graph record (snapshot) for a completed job.
+
+        Note: for graph-first jobs, this record is **redacted** and never
+        includes proprietary YAML/spec. Use `download_graph_txt()` for the
+        public export.
+
+        Args:
+            prompt_snapshot_id: Legacy alias for selecting a specific snapshot.
+            graph_snapshot_id: Specific GraphSnapshot to use (default: best).
+
+        Returns:
+            Graph record dictionary containing:
+            - job_id: The job ID
+            - snapshot_id: The snapshot ID used
+            - prompt: Extracted prompt text (legacy single-node only; may be empty)
+            - graph: Public graph record payload (e.g., export metadata)
+            - model: Model used for this graph (optional)
+
+        Raises:
+            RuntimeError: If job hasn't been submitted or API call fails.
+            ValueError: If both prompt_snapshot_id and graph_snapshot_id are provided.
+        """
+        if not self.job_id:
+            raise RuntimeError("Job not yet submitted. Call submit() first.")
+
+        if prompt_snapshot_id and graph_snapshot_id:
+            raise ValueError("Provide only one of prompt_snapshot_id or graph_snapshot_id.")
+
+        url = f"{self.backend_url}/graphgen/graph/record"
+        headers = {
+            "X-API-Key": self.api_key,
+            "Content-Type": "application/json",
+        }
+
+        payload: Dict[str, Any] = {
+            "job_id": self.job_id,
+        }
+        snapshot_id = graph_snapshot_id or prompt_snapshot_id
+        if snapshot_id:
+            payload["prompt_snapshot_id"] = snapshot_id
+
+        resp = http_post(url, headers=headers, json_body=payload, timeout=30.0)
+
+        if resp.status_code != 200:
+            raise RuntimeError(
+                f"Failed to get graph record: {resp.status_code} - {resp.text[:500]}"
+            )
+
+        return cast(Dict[str, Any], resp.json())
+
+
+__all__ = [
+    "GraphGenJob",
+    "GraphGenJobResult",
+    "GraphGenSubmitResult",
+]