mantisdk 0.1.0__py3-none-any.whl

mantisdk/algorithm/gepa/gepa.py

@@ -0,0 +1,364 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+from __future__ import annotations
+
+import asyncio
+import logging
+from typing import Any, Dict, Generic, List, Optional, TypeVar
+
+import openai
+from mantisdk.algorithm.gepa.lib.api import optimize
+
+from mantisdk.adapter import TraceAdapter
+from mantisdk.adapter.messages import TraceToMessages
+from mantisdk.algorithm.base import Algorithm
+from mantisdk.algorithm.gepa.adapter import (
+    MantisdkDataInst,
+    MantisdkGEPAAdapter,
+)
+from mantisdk.algorithm.gepa.tracing import GEPATracingContext
+from mantisdk.algorithm.utils import with_llm_proxy, with_store
+from mantisdk.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 Mantisdk.
+
+    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
+    """
+
+    # Algorithm-owned tracing configuration.
+    # Defines environment and tags for all traces generated by GEPA.
+    # Note: session_id will be overridden per-run in the run() method
+    TRACING_CONFIG = TracingConfig(
+        environment="mantisdk-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="mantisdk-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_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_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_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
+
+    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=True)
+    async def run(
+        self,
+        llm_proxy,  # injected by decorator
+        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
+        assert llm_proxy 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
+        llm_resource = llm_proxy.as_resource()
+        
+        # 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 = MantisdkGEPAAdapter(
+            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:
+        # `mantisdk.algorithm.gepa.__init__` re-exports `GEPA` from this module.
+        from mantisdk.algorithm import gepa as gepa_module
+        
+        @gepa_module.reflection
+        def reflection_lm(prompt: str) -> str:
+            """GEPA reflection LM callable - takes a string prompt, returns string response.
+            
+            Routes through the LLMProxy for consistent tracing and OpenTelemetry export.
+            The proxy handles API keys and model routing - we just need to call it.
+            
+            IMPORTANT: The LLMProxy must be configured with callbacks: ["opentelemetry"]
+            (not "return_token_ids") when calling OpenAI providers.
+            
+            The reflection phase is tagged distinctly from validation-eval to enable
+            proper filtering in the Mantis UI. Tags include:
+            - "gepa" (algorithm)
+            - "reflection" (phase)
+            - "llm-reflection" (explicit marker)
+            - "gen-N" (generation number)
+            """
+            import litellm
+            
+            # Track phase transition: entering reflection, then starting new generation
+            tracing_context.set_phase("reflection")
+            
+            logger.info(f"[reflection_lm] === REFLECTION CALLED (gen-{tracing_context.generation}) ===")
+            logger.info(f"[reflection_lm] Session: {tracing_context.session_id}")
+            logger.info(f"[reflection_lm] Model: {llm_resource.model}")
+            logger.info(f"[reflection_lm] Proxy Endpoint: {llm_resource.endpoint}")
+            logger.info(f"[reflection_lm] Prompt length: {len(prompt)} chars")
+            logger.debug(f"[reflection_lm] Prompt preview: {prompt[:500]}...")
+            
+            try:
+                # Route through the LLMProxy for tracing
+                # The proxy already has the API key and model config from gepa_runner.py
+                logger.info(f"[reflection_lm] Calling via proxy: model={llm_resource.model} endpoint={llm_resource.endpoint}")
+                
+                # Include detailed tracing metadata for reflection traces
+                # Uses run_tracing_config which includes the session_id
+                # Note: call-type tagging ("reflection-call") is handled by @gepa_module.reflection decorator
+                import json
+                gepa_tags = [
+                    f"gen-{tracing_context.generation}",
+                ]
+                tracing_metadata = run_tracing_config.to_detailed_metadata(
+                    phase="reflection",
+                    extra_tags=gepa_tags,
+                )
+                
+                # Pass metadata via headers for proxy-side tracing context
+                # Note: call-type is handled automatically by the @gepa_module.reflection decorator
+                extra_headers = {
+                    "x-mantis-session-id": tracing_metadata.get("session_id", ""),
+                    "x-mantis-environment": tracing_metadata.get("environment", ""),
+                    "x-mantis-tags": json.dumps(tracing_metadata.get("tags", [])),
+                }
+                
+                response = litellm.completion(
+                    model=llm_resource.model,  # Model name as registered in proxy
+                    messages=[{"role": "user", "content": prompt}],
+                    api_base=llm_resource.endpoint,  # Proxy URL
+                    api_key="dummy",  # Proxy handles the real API key
+                    temperature=0.7,
+                    metadata=tracing_metadata,  # Pass environment/tags/session_id for OTEL
+                    extra_headers=extra_headers,  # Pass to proxy for server-side span attributes
+                )
+                content = response.choices[0].message.content or ""
+                logger.info(f"[reflection_lm] SUCCESS - Response length: {len(content)} chars")
+                logger.info(f"[reflection_lm] Response preview: {content[:300]}...")
+                
+                # After reflection completes, increment generation for next eval round
+                tracing_context.next_generation()
+                tracing_context.set_phase("train-eval")
+                
+                return content
+            except Exception as e:
+                logger.error(f"[reflection_lm] FAILED: {e}", exc_info=True)
+                # Return fallback instead of crashing - allows optimization to continue
+                logger.warning(f"[reflection_lm] Returning fallback (seed prompt)")
+                
+                # Still increment generation even on failure
+                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[MantisdkDataInst] = []
+        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[MantisdkDataInst]] = 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")
+        
+        # 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,
+                **self.gepa_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
+
+            # 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 Insight 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

mantisdk/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,
+)

mantisdk/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!