gatewaysdk 0.3.2__py3-none-any.whl

gatewaysdk/algorithm/gepa/gepa.py

@@ -0,0 +1,570 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+from __future__ import annotations
+
+import asyncio
+import logging
+from typing import Any, ClassVar, Dict, Generic, List, Optional, TypeVar
+
+import openai
+
+from gatewaysdk.adapter import TraceAdapter
+from gatewaysdk.adapter.messages import TraceToMessages
+from gatewaysdk.algorithm.base import Algorithm
+from gatewaysdk.algorithm.gepa.adapter import (
+    GatewayDataInst,
+    GatewayGEPAAdapter,
+)
+from gatewaysdk.algorithm.gepa.lib.api import optimize
+from gatewaysdk.algorithm.gepa.tracing import GEPATracingContext
+from gatewaysdk.algorithm.utils import with_llm_proxy, with_store
+from gatewaysdk.types import Dataset, PromptTemplate, TracingConfig
+
+logger = logging.getLogger(__name__)
+
+T_task = TypeVar("T_task")
+
+
+TEMPLATE_AWARE_REFLECTION_PROMPT = """You are an expert at improving LLM prompts based on observed failures.
+
+## Current Prompt Template
+```
+<curr_instructions>
+```
+
+## Observed Failures
+The following examples show where the current prompt gave INCORRECT outputs. Study the pattern of failures carefully:
+```
+<inputs_outputs_feedback>
+```
+
+## Your Task
+Write an IMPROVED prompt template that fixes the observed failures.
+
+**CRITICAL REQUIREMENTS**:
+1. **Analyze the failure pattern**: Look at WHAT the prompt got wrong. Is it too strict? Too lenient? Missing context? Misunderstanding the task?
+2. **Make a CONCEPTUAL FIX**: Don't just tweak wording - fundamentally change the approach if the current strategy is flawed
+3. **Preserve placeholders**: Keep any {variable_name} placeholders exactly as they appear (e.g., {session}, {input}, {output}, or {{session}}, {{input}}, {{output}})
+4. **Be specific**: Add concrete criteria, examples, or decision rules based on what the failures reveal
+5. Read the inputs carefully and identify the input format and infer detailed task description about the task I wish to solve with the assistant
+6. Read all the assistant responses and the corresponding feedback. Identify all niche and domain specific factual information about the task and include it in the instruction, as a lot of it may not be available to the assistant in the future. The assistant may have utilized a generalizable strategy to solve the task, if so, include that in the instruction as well.
+
+**Think step by step**:
+- What is the PATTERN in the failures?
+- WHY is the current prompt failing on these cases?
+- What SPECIFIC change would fix this pattern?
+
+Output your improved template within ``` blocks."""
+
+
+class GEPA(Algorithm, Generic[T_task]):
+    """GEPA (Genetic-Pareto) algorithm for GatewaySDK.
+
+    This algorithm optimizes prompt templates (and potentially other text resources)
+    using an evolutionary approach with LLM-based reflection.
+
+    GEPA maintains a population of candidate prompts and evolves them by:
+    1. Evaluating candidates on training data
+    2. Using an LLM to reflect on failures and propose improvements
+    3. Selecting the best candidates based on validation performance
+    """
+
+    # Platform metadata
+    name: ClassVar[str] = "gepa"
+    display_name: ClassVar[str] = "GEPA"
+    description: ClassVar[str] = "Genetic-Pareto Prompt Adaptation"
+    requirements: ClassVar[Dict[str, Any]] = {"gpu": False, "min_runners": 1}
+    config_schema: ClassVar[Optional[Dict[str, Any]]] = {
+        "type": "object",
+        "properties": {
+            "reflection_model": {
+                "type": "string",
+                "default": "anthropic/claude-haiku-4-5",
+                "description": "LLM model used by the algorithm to analyze agent failures and generate improved prompts. Accepts any LiteLLM model string.",
+            },
+            "max_metric_calls": {
+                "type": "integer",
+                "default": 1000,
+                "minimum": 1,
+                "description": "Maximum number of evaluations (budget)",
+            },
+            "population_size": {
+                "type": "integer",
+                "default": 5,
+                "minimum": 1,
+                "maximum": 50,
+                "description": "Size of the population in evolutionary search",
+            },
+            "reflection_minibatch_size": {
+                "type": "integer",
+                "default": 3,
+                "minimum": 1,
+                "description": "Number of examples per reflection batch",
+            },
+            "rollout_batch_timeout": {
+                "type": "number",
+                "default": 300.0,
+                "minimum": 30.0,
+                "description": "Timeout in seconds for waiting for rollout results",
+            },
+            "candidate_selection_strategy": {
+                "type": "string",
+                "default": "pareto",
+                "enum": ["pareto", "current_best", "epsilon_greedy"],
+                "description": "Strategy for selecting the next candidate to optimize",
+            },
+            "perfect_score": {
+                "type": "number",
+                "default": 1.0,
+                "description": "Score at which an example is considered perfectly solved",
+            },
+            "seed": {
+                "type": "integer",
+                "default": 0,
+                "description": "Random seed for reproducibility",
+            },
+            "initial_prompt": {
+                "type": "string",
+                "description": "Initial system prompt template to optimize (uses {input} placeholder)",
+            },
+            "val_fraction": {
+                "type": "number",
+                "default": 0.2,
+                "minimum": 0.0,
+                "maximum": 0.5,
+                "description": "Fraction of dataset held out for validation",
+            },
+        },
+    }
+
+    # Algorithm-owned tracing configuration
+    TRACING_CONFIG = TracingConfig(
+        environment="gatewaysdk-gepa",
+        algorithm_name="gepa",
+    )
+
+    def _get_tracing_config_with_session(self, session_id: str) -> TracingConfig:
+        """Create a TracingConfig with a specific session_id for this run."""
+        return TracingConfig(
+            environment="gatewaysdk-gepa",
+            algorithm_name="gepa",
+            session_id=session_id,
+        )
+
+    def __init__(
+        self,
+        *,
+        max_metric_calls: int = 100,
+        reflection_minibatch_size: int = 5,
+        population_size: int = 4,
+        adapter: Optional[TraceAdapter] = None,
+        rollout_batch_timeout: float = 600.0,
+        reflection_model: str = "anthropic/claude-haiku-4-5",
+        reflection_prompt_template: Optional[str] = None,
+        reflection_metadata: Optional[Dict[str, Any]] = None,
+        target_model_config: Optional[Dict[str, Any]] = None,
+        **gepa_kwargs: Any,
+    ) -> None:
+        """Initialize the GEPA algorithm.
+
+        Args:
+            max_metric_calls: Maximum number of evaluations (budget).
+            reflection_minibatch_size: Batch size for reflection.
+            population_size: Size of the population in evolutionary search.
+            adapter: TraceAdapter to convert spans to messages. Defaults to TraceToMessages.
+            rollout_batch_timeout: Timeout for waiting for rollouts.
+            reflection_model: LiteLLM model string for the reflection LLM.
+            reflection_prompt_template: Custom prompt template for reflection.
+            reflection_metadata: Metadata for reflection LLM traces (Langfuse).
+            target_model_config: Direct model configuration (apiKey, baseUrl) to bypass proxy if needed.
+            **gepa_kwargs: Additional arguments passed to gepa.optimize.
+        """
+        super().__init__()
+        self.max_metric_calls = max_metric_calls
+        self.reflection_minibatch_size = reflection_minibatch_size
+        self.population_size = population_size
+        self.adapter = adapter or TraceToMessages()
+        self.rollout_batch_timeout = rollout_batch_timeout
+        self.reflection_model = reflection_model
+        self.reflection_metadata = reflection_metadata
+        self.target_model_config = target_model_config
+
+        # Set default reflection prompt template if not provided
+        if reflection_prompt_template is not None:
+            gepa_kwargs["reflection_prompt_template"] = reflection_prompt_template
+        elif "reflection_prompt_template" not in gepa_kwargs:
+            gepa_kwargs["reflection_prompt_template"] = TEMPLATE_AWARE_REFLECTION_PROMPT
+
+        self.gepa_kwargs = gepa_kwargs
+        self._best_candidate: Optional[Dict[str, str]] = None
+        self._best_score: float = 0.0
+        self._full_result: Optional[Any] = None  # Store full GEPAResult for history access
+        self._engine: Optional[Any] = None  # GEPAEngine reference for cancellation
+
+    def request_cancel(self) -> None:
+        """Signal GEPA to stop at the next iteration boundary."""
+        super().request_cancel()
+        if self._engine is not None:
+            self._engine.request_stop()
+            logger.info("GEPA engine stop requested via cancellation")
+
+    def get_best_prompt(self) -> Optional[PromptTemplate]:
+        """Get the best prompt found during optimization."""
+        if self._best_candidate:
+            # Return the first PromptTemplate value
+            for key, value in self._best_candidate.items():
+                return PromptTemplate(template=value, engine="f-string")
+        return None
+
+    @with_store
+    @with_llm_proxy(required=False)
+    async def run(
+        self,
+        llm_proxy,  # injected by decorator (may be None)
+        store,  # injected by decorator
+        train_dataset: Optional[Dataset[T_task]] = None,
+        val_dataset: Optional[Dataset[T_task]] = None,
+    ) -> None:
+        """Run the GEPA optimization loop.
+
+        Args:
+            train_dataset: Dataset used for optimization (training).
+            val_dataset: Dataset used for validation.
+        """
+        if train_dataset is None:
+            raise ValueError("train_dataset is required for GEPA optimization.")
+
+        store = self.get_store()
+        assert store is not None
+
+        loop = asyncio.get_running_loop()
+
+        # Get initial resources to find the target resource to optimize
+        initial_resources = self.get_initial_resources()
+        if not initial_resources:
+            raise ValueError("Initial resources must be set before running GEPA.")
+
+        # Find the target resource to optimize
+        target_resource_name = None
+        initial_candidate: Dict[str, str] = {}
+
+        for name, res in initial_resources.items():
+            if isinstance(res, PromptTemplate):
+                target_resource_name = name
+                initial_candidate[name] = res.template
+
+        if not target_resource_name:
+            raise ValueError("No PromptTemplate found in initial resources to optimize.")
+
+        logger.info(f"GEPA will optimize resource: {target_resource_name}")
+        logger.info(f"Initial prompt: {initial_candidate[target_resource_name][:100]}...")
+
+        # Setup Reflection LLM -- via proxy if available, otherwise direct
+        llm_resource = llm_proxy.as_resource() if llm_proxy is not None else None
+        reflection_model = self.reflection_model
+
+        if llm_resource is not None:
+            logger.info(f"GEPA reflection via LLM proxy: {llm_resource.model}")
+        else:
+            logger.info(f"GEPA reflection via direct litellm: {reflection_model}")
+
+        # Create tracing context for detailed execution tracking
+        # This generates a unique session_id for grouping all traces in this run
+        tracing_context = GEPATracingContext()
+
+        # Create a TracingConfig with this run's session_id
+        run_tracing_config = self._get_tracing_config_with_session(tracing_context.session_id)
+
+        logger.info(f"GEPA session started: {tracing_context.session_id}")
+
+        # Create the bridge adapter with tracing config and context
+        gepa_adapter = GatewayGEPAAdapter(
+            store=store,
+            loop=loop,
+            resource_name=target_resource_name,
+            adapter=self.adapter,
+            llm_proxy_resource=llm_resource,
+            rollout_batch_timeout=self.rollout_batch_timeout,
+            tracing_config=run_tracing_config,
+            tracing_context=tracing_context,
+        )
+
+        # Import here (inside run) to avoid circular imports:
+        # `gatewaysdk.algorithm.gepa.__init__` re-exports `GEPA` from this module.
+        from gatewaysdk.algorithm import gepa as gepa_module
+
+        reflection_failures = 0
+
+        @gepa_module.reflection
+        def reflection_lm(prompt: str) -> str:
+            """GEPA reflection LM callable -- string in, string out.
+
+            Two modes:
+            - With LLM proxy: routes through the proxy for centralized
+              key management and tracing.
+            - Without proxy: calls litellm directly using API keys from
+              the environment.  Tracing still works via tracing.init().
+            """
+            import litellm
+
+            tracing_context.set_phase("reflection")
+
+            model_name = llm_resource.model if llm_resource else reflection_model
+            logger.info(f"[reflection_lm] === REFLECTION (gen-{tracing_context.generation}) ===")
+            logger.info(f"[reflection_lm] Model: {model_name}, prompt: {len(prompt)} chars")
+
+            try:
+                from gatewaysdk.tracing.context import tracing_context as _tracing_ctx
+
+                gepa_tags = run_tracing_config.get_detailed_tags(
+                    phase="reflection",
+                    extra_tags=[f"gen-{tracing_context.generation}"],
+                )
+                reflection_session = f"{tracing_context.session_id}-reflection-gen{tracing_context.generation}"
+
+                with _tracing_ctx(
+                    tags=gepa_tags, environment=run_tracing_config.environment, session_id=reflection_session
+                ):
+                    if llm_resource is not None:
+                        # Route through LLM proxy
+                        proxy_endpoint = llm_resource.get_base_url(None, None)
+                        response = litellm.completion(
+                            model=llm_resource.model,
+                            messages=[{"role": "user", "content": prompt}],
+                            api_base=proxy_endpoint,
+                            api_key="dummy",
+                            temperature=0.7,
+                            extra_headers={"x-gateway-session-id": reflection_session},
+                        )
+                    else:
+                        # Direct litellm call -- API key from environment
+                        response = litellm.completion(
+                            model=reflection_model,
+                            messages=[{"role": "user", "content": prompt}],
+                            temperature=0.7,
+                        )
+
+                content = response.choices[0].message.content or ""
+                logger.info(f"[reflection_lm] SUCCESS - {len(content)} chars")
+
+                tracing_context.next_generation()
+                tracing_context.set_phase("train-eval")
+                return content
+
+            except Exception as e:
+                nonlocal reflection_failures
+                reflection_failures += 1
+                logger.error(f"[reflection_lm] FAILED ({reflection_failures} total): {e}", exc_info=True)
+                if reflection_failures > 3:
+                    raise RuntimeError(
+                        f"GEPA reflection failed {reflection_failures} times consecutively. "
+                        "Check LLM endpoint configuration."
+                    ) from e
+                logger.warning(f"[reflection_lm] Returning fallback (seed prompt)")
+                tracing_context.next_generation()
+                tracing_context.set_phase("train-eval")
+                return f"```\n{initial_candidate.get(target_resource_name, '')}\n```"
+
+        # Prepare Training Data
+        gepa_train_data: List[GatewayDataInst] = []
+        for i, item in enumerate(train_dataset):
+            item_id = str(i)
+            if isinstance(item, dict):
+                gepa_train_data.append({"input": item, "id": item.get("id", item_id)})
+            else:
+                gepa_train_data.append({"input": {"task": item}, "id": item_id})
+
+        # Prepare Validation Data (use train if not provided)
+        gepa_val_data: Optional[List[GatewayDataInst]] = None
+        if val_dataset is not None:
+            gepa_val_data = []
+            for i, item in enumerate(val_dataset):
+                item_id = f"val-{i}"
+                if isinstance(item, dict):
+                    gepa_val_data.append({"input": item, "id": item.get("id", item_id)})
+                else:
+                    gepa_val_data.append({"input": {"task": item}, "id": item_id})
+
+        logger.info(f"Starting GEPA optimization with {len(gepa_train_data)} training samples")
+        if gepa_val_data:
+            logger.info(f"Using {len(gepa_val_data)} validation samples")
+
+        # Create live state callback for real-time visualization
+        def on_step_callback(state) -> None:
+            """Update live state for real-time visualization after each GEPA iteration."""
+            logger.info(f"[on_step_callback] Called for iteration {state.i}")
+
+            if not hasattr(store, "log_state"):
+                logger.info(f"[on_step_callback] store has no log_state method")
+                return
+
+            logger.info(f"[on_step_callback] calling store.log_state")
+
+            try:
+                # Build candidates list with full data
+                candidates = []
+                for idx, cand in enumerate(state.program_candidates):
+                    # Get prompt text from candidate dict
+                    prompt_text = ""
+                    prompt_key = ""
+                    if cand:
+                        first_key = next(iter(cand.keys()), None)
+                        if first_key:
+                            prompt_key = first_key
+                            prompt_text = cand.get(first_key, "")
+
+                    # Get score from prog_candidate_val_subscores (dict of val_id -> score)
+                    val_scores = (
+                        state.prog_candidate_val_subscores[idx] if idx < len(state.prog_candidate_val_subscores) else {}
+                    )
+                    score = sum(val_scores.values()) / len(val_scores) if val_scores else None
+
+                    # Get parent indices from parent_program_for_candidate
+                    parent_list = (
+                        state.parent_program_for_candidate[idx] if idx < len(state.parent_program_for_candidate) else []
+                    )
+                    parent_idx = parent_list[0] if parent_list and parent_list[0] is not None else None
+
+                    candidates.append(
+                        {
+                            "idx": idx,
+                            "prompt": prompt_text[:2000],  # Truncate for sanity
+                            "promptKey": prompt_key,
+                            "score": score,
+                            "parent": parent_idx,
+                        }
+                    )
+
+                # Build pareto front from program_at_pareto_front_valset
+                pareto_front = {}
+                if hasattr(state, "program_at_pareto_front_valset"):
+                    for instance_id, best_idxs in state.program_at_pareto_front_valset.items():
+                        pareto_front[str(instance_id)] = list(best_idxs) if best_idxs else []
+
+                # Find best candidate by computing average validation score
+                best_idx = 0
+                best_score = 0.0
+                for idx, val_scores in enumerate(state.prog_candidate_val_subscores):
+                    if val_scores:
+                        avg = sum(val_scores.values()) / len(val_scores)
+                        if avg > best_score:
+                            best_score = avg
+                            best_idx = idx
+
+                # Log the live state via public API (F14)
+                store.log_state(
+                    {
+                        "candidates": candidates,
+                        "pareto_front": pareto_front,
+                        "best_idx": best_idx,
+                        "best_score": best_score,
+                        "generation": tracing_context.generation,
+                        "total_metric_calls": state.total_num_evals,
+                        "iteration": state.i,
+                    }
+                )
+
+                logger.info(
+                    f"[on_step_callback] Updated live state: {len(candidates)} candidates, best={best_score:.3f}"
+                )
+
+                # Log per-step metrics for the Metrics tab time series
+                if hasattr(store, "_run") and store._run:
+                    try:
+                        store._run.log(
+                            {
+                                "best_score": best_score,
+                                "num_candidates": len(candidates),
+                                "metric_calls": state.total_num_evals,
+                            },
+                            step=max(0, state.i),
+                        )
+                    except Exception:
+                        pass  # Metrics are best-effort; Evolution tab is primary
+            except Exception as e:
+                logger.warning(f"[on_step_callback] Failed to update live state: {e}")
+
+        # Create a cancellation stopper that checks Algorithm._cancel_event
+        cancel_event = self._cancel_event
+
+        class _CancelStopper:
+            """Bridges Algorithm.request_cancel() → GEPA engine stop."""
+
+            def __call__(self, gepa_state: Any) -> bool:
+                return cancel_event.is_set()
+
+        # Merge with any existing stop_callbacks from gepa_kwargs
+        extra_kwargs = dict(self.gepa_kwargs)
+        existing_stop_cbs = extra_kwargs.pop("stop_callbacks", None)
+        stop_cbs: list = []
+        if existing_stop_cbs is not None:
+            if isinstance(existing_stop_cbs, (list, tuple)):
+                stop_cbs.extend(existing_stop_cbs)
+            else:
+                stop_cbs.append(existing_stop_cbs)
+        stop_cbs.append(_CancelStopper())
+
+        # Run Optimization in Executor (GEPA is synchronous)
+        def run_optimization():
+            return optimize(
+                trainset=gepa_train_data,
+                valset=gepa_val_data,
+                adapter=gepa_adapter,
+                seed_candidate=initial_candidate,
+                reflection_lm=reflection_lm,
+                max_metric_calls=self.max_metric_calls,
+                reflection_minibatch_size=self.reflection_minibatch_size,
+                on_step_callback=on_step_callback,
+                stop_callbacks=stop_cbs,
+                **extra_kwargs,
+            )
+
+        try:
+            result = await loop.run_in_executor(None, run_optimization)
+
+            logger.info(f"GEPA optimization finished. Total metric calls: {result.total_metric_calls}")
+            logger.info(
+                f"Best candidate index: {result.best_idx}, score: {result.val_aggregate_scores[result.best_idx]}"
+            )
+            logger.info(f"Best candidate: {result.best_candidate}")
+
+            self._best_candidate = result.best_candidate
+            self._best_score = result.val_aggregate_scores[result.best_idx]
+            self._full_result = result  # Store full result for candidate history
+
+            # Log final summary metrics (these complement the per-step
+            # best_score series logged by on_step_callback above)
+            if hasattr(store, "_run") and store._run:
+                try:
+                    store._run.log(
+                        {
+                            "best_score": self._best_score,
+                            "num_candidates": len(result.candidates),
+                            "metric_calls": result.total_metric_calls,
+                        },
+                        step=result.total_metric_calls,
+                    )
+                    logger.info(f"Logged final GEPA metrics: best_score={self._best_score}")
+                except Exception as e:
+                    logger.warning(f"Failed to log final metrics: {e}")
+
+            # Update the store with the best candidate
+            final_resources = {}
+            for key, value in result.best_candidate.items():
+                original = initial_resources.get(key)
+                engine = original.engine if isinstance(original, PromptTemplate) else "f-string"
+                final_resources[key] = PromptTemplate(template=value, engine=engine)
+
+            await store.update_resources("gepa-final", final_resources)
+            logger.info("Updated store with best candidate resources (version: 'gepa-final').")
+
+            # Send full GEPA result to Gateway for experiment tracking
+            if hasattr(store, "complete_job"):
+                summary = result.to_dict()
+                # Add session_id for trace filtering in the UI
+                summary["session_id"] = tracing_context.session_id
+                store.complete_job(summary)
+
+        except Exception as e:
+            logger.error(f"GEPA optimization failed: {e}")
+            raise

gatewaysdk/algorithm/gepa/lib/__init__.py

@@ -0,0 +1,18 @@
+# Copyright (c) 2025 Lakshya A Agrawal and the GEPA contributors
+# https://github.com/gepa-ai/gepa
+
+from .adapters import default_adapter
+from .api import optimize
+from .core.adapter import EvaluationBatch, GEPAAdapter
+from .core.result import GEPAResult
+from .examples import aime
+from .utils.stop_condition import (
+    CompositeStopper,
+    FileStopper,
+    MaxMetricCallsStopper,
+    NoImprovementStopper,
+    ScoreThresholdStopper,
+    SignalStopper,
+    StopperProtocol,
+    TimeoutStopCondition,
+)

gatewaysdk/algorithm/gepa/lib/adapters/README.md

@@ -0,0 +1,12 @@
+# GEPA Adapters
+
+> GEPA 🤝 Any Framework
+
+This directory provides the interface to allow GEPA to plug into systems and frameworks of your choice! GEPA can interface with any system consisting of text components, by implementing `GEPAAdapter` in [../core/adapter.py](../core/adapter.py).
+
+Currently, GEPA has the following adapters:
+- [DSPy Adapter](./dspy_adapter/): This adapter integrates GEPA into [DSPy](https://dspy.ai/), to allow it to optimize any DSPy module's signature instructions.
+- [Default Adapter](./default_adapter/): This adapter integrates GEPA into a single-turn LLM environment, where the task is specified as a user message, and an answer string must be present in the assistant response. GEPA optimizes the system prompt.
+- [AnyMaths Adapter](./anymaths_adapter/): This adapter integrates GEPA with litellm and ollama to solve single-turn mathematical problems.
+
+If there are any frameworks you would like GEPA integrated into, please create an issue or PR!