opik-optimizer 1.0.6__py3-none-any.whl → 2.0.0__py3-none-any.whl

opik_optimizer/evolutionary_optimizer/helpers.py

@@ -0,0 +1,10 @@
+from ..optimization_config import chat_prompt
+
+
+class Helpers:
+    def _get_task_description_for_llm(self, prompt: chat_prompt.ChatPrompt) -> str:
+        """Generates a concise task description for LLM prompts that need context."""
+        description = "Task: Given a list of AI messages with placeholder values, generate an effective prompt. "
+        description += f"The original high-level instruction being optimized is: '{prompt.get_messages()}'. "
+        description += "The goal is to create an effective prompt that guides a language model to perform this task well."
+        return description

opik_optimizer/evolutionary_optimizer/llm_support.py

@@ -0,0 +1,136 @@
+from typing import Any, TYPE_CHECKING
+
+import logging
+import os
+import time
+import random
+
+import litellm
+from litellm import exceptions as litellm_exceptions
+from litellm.caching import Cache
+from litellm.types.caching import LiteLLMCacheType
+from opik.evaluation.models.litellm import opik_monitor as opik_litellm_monitor
+
+from .. import _throttle
+
+
+logger = logging.getLogger(__name__)
+
+
+# Configure LiteLLM cache with safe fallback
+try:
+    # Prefer a disk cache in a user-writable location
+    cache_dir = os.path.join(os.path.expanduser("~"), ".cache", "litellm")
+    os.makedirs(cache_dir, exist_ok=True)
+    litellm.cache = Cache(type=LiteLLMCacheType.DISK, cache_dir=cache_dir)
+except (PermissionError, OSError, FileNotFoundError):
+    # Fall back to in-memory cache to avoid disk timeouts/locks
+    litellm.cache = Cache(type=LiteLLMCacheType.MEMORY)
+
+_rate_limiter = _throttle.get_rate_limiter_for_current_opik_installation()
+
+
+class LlmSupport:
+    if TYPE_CHECKING:
+        model: str
+        llm_call_counter: int
+        project_name: str | None
+        disable_litellm_monitoring: bool
+        temperature: float
+        max_tokens: int
+        top_p: float
+        frequency_penalty: float
+        presence_penalty: float
+
+        def increment_llm_counter(self) -> None: ...
+
+    @_throttle.rate_limited(_rate_limiter)
+    def _call_model(
+        self,
+        messages: list[dict[str, str]],
+        is_reasoning: bool = False,
+        optimization_id: str | None = None,
+    ) -> str:
+        """Call the model with the given prompt and return the response string."""
+        # Build base call params
+        llm_config_params: dict[str, Any] = {
+            "temperature": getattr(self, "temperature", 0.3),
+            "max_tokens": getattr(self, "max_tokens", 1000),
+            "top_p": getattr(self, "top_p", 1.0),
+            "frequency_penalty": getattr(self, "frequency_penalty", 0.0),
+            "presence_penalty": getattr(self, "presence_penalty", 0.0),
+        }
+
+        # Add Opik metadata unless disabled
+        try:
+            disable_monitoring_env = os.getenv(
+                "OPIK_OPTIMIZER_DISABLE_LITELLM_MONITORING", "0"
+            )
+            disable_monitoring = getattr(
+                self, "disable_litellm_monitoring", False
+            ) or disable_monitoring_env.lower() in ("1", "true", "yes")
+
+            if not disable_monitoring:
+                metadata_for_opik: dict[str, Any] = {}
+                pn = getattr(self, "project_name", None)
+                if pn:
+                    metadata_for_opik["project_name"] = pn
+                    metadata_for_opik["opik"] = {"project_name": pn}
+                if optimization_id and "opik" in metadata_for_opik:
+                    metadata_for_opik["opik"]["optimization_id"] = optimization_id
+                metadata_for_opik["optimizer_name"] = self.__class__.__name__
+                metadata_for_opik["opik_call_type"] = (
+                    "reasoning" if is_reasoning else "evaluation_llm_task_direct"
+                )
+                if metadata_for_opik:
+                    llm_config_params["metadata"] = metadata_for_opik
+
+                # Try to add Opik monitoring callbacks; fall back silently on failure
+                llm_config_params = (
+                    opik_litellm_monitor.try_add_opik_monitoring_to_params(  # type: ignore
+                        llm_config_params.copy()
+                    )
+                )
+        except Exception as e:
+            logger.debug(f"Skipping Opik-LiteLLM monitoring setup: {e}")
+
+        # Retry policy for transient errors
+        max_retries = int(os.getenv("OPIK_OPTIMIZER_LITELLM_MAX_RETRIES", "3"))
+        base_sleep = float(os.getenv("OPIK_OPTIMIZER_LITELLM_BACKOFF", "0.5"))
+
+        for attempt in range(max_retries + 1):
+            try:
+                logger.debug(
+                    f"Calling model '{self.model}' with messages: {messages}, params: {llm_config_params} (attempt {attempt + 1})"
+                )
+                response = litellm.completion(
+                    model=self.model, messages=messages, **llm_config_params
+                )
+                self.increment_llm_counter()
+                return response.choices[0].message.content
+            except (
+                litellm_exceptions.RateLimitError,
+                litellm_exceptions.APIConnectionError,
+                litellm_exceptions.InternalServerError,
+            ) as e:
+                if attempt < max_retries:
+                    sleep_s = min(10.0, base_sleep * (2**attempt)) + random.uniform(
+                        0, 0.25
+                    )
+                    logger.warning(
+                        f"LiteLLM transient error ({type(e).__name__}): {e}. Retrying in {sleep_s:.2f}s..."
+                    )
+                    time.sleep(sleep_s)
+                    continue
+                logger.error(f"LiteLLM error (final attempt): {e}")
+                raise
+            except litellm_exceptions.ContextWindowExceededError as e:
+                logger.error(f"LiteLLM Context Window Exceeded Error: {e}")
+                raise
+            except Exception as e:
+                logger.error(
+                    f"Error calling model '{self.model}': {type(e).__name__} - {e}"
+                )
+                raise
+        # Should never reach here
+        raise RuntimeError("LLM call did not return a response and did not raise")

opik_optimizer/evolutionary_optimizer/mcp.py

@@ -0,0 +1,249 @@
+"""MCP helper utilities for the Evolutionary Optimizer."""
+
+from __future__ import annotations
+
+import json
+import logging
+import re
+import textwrap
+from dataclasses import dataclass
+from typing import Any
+
+from opik_optimizer.optimization_config import chat_prompt
+from opik_optimizer.utils.prompt_segments import (
+    apply_segment_updates,
+    extract_prompt_segments,
+)
+
+from . import prompts as evo_prompts
+from . import reporting
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class EvolutionaryMCPContext:
+    tool_name: str
+    tool_segment_id: str
+    original_description: str
+    tool_metadata: dict[str, Any]
+    panel_style: str
+
+
+def _tool_metadata_json(metadata: dict[str, Any]) -> str:
+    try:
+        return json.dumps(metadata, indent=2)
+    except (
+        TypeError,
+        ValueError,
+    ):  # pragma: no cover - defensive, shouldn't happen under normal circumstances
+        return str(metadata)
+
+
+def generate_tool_description_variations(
+    optimizer: Any,
+    base_prompt: chat_prompt.ChatPrompt,
+    context: EvolutionaryMCPContext,
+    num_variations: int,
+) -> list[chat_prompt.ChatPrompt]:
+    if num_variations <= 0:
+        return []
+
+    instruction = textwrap.dedent(
+        evo_prompts.mcp_tool_rewrite_user_prompt(
+            tool_name=context.tool_name,
+            current_description=_current_tool_description(
+                base_prompt, context.tool_segment_id
+            )
+            or context.original_description,
+            tool_metadata_json=_tool_metadata_json(context.tool_metadata),
+            num_variations=num_variations,
+        )
+    ).strip()
+
+    try:
+        response = optimizer._call_model(  # type: ignore[attr-defined]
+            messages=[
+                {
+                    "role": "system",
+                    "content": evo_prompts.mcp_tool_rewrite_system_prompt(),
+                },
+                {"role": "user", "content": instruction},
+            ],
+            is_reasoning=True,
+            optimization_id=getattr(optimizer, "_current_optimization_id", None),
+        )
+
+        payload = _extract_json_payload(response)
+        prompts_payload = payload.get("prompts")
+        if not isinstance(prompts_payload, list):
+            raise ValueError("LLM response missing 'prompts' list")
+
+        candidates: list[chat_prompt.ChatPrompt] = []
+        seen: set[str] = set()
+        for item in prompts_payload:
+            if not isinstance(item, dict):
+                continue
+            description = item.get("tool_description")
+            if not isinstance(description, str) or not description.strip():
+                continue
+            normalized = description.strip()
+            if normalized in seen:
+                continue
+            seen.add(normalized)
+            updated_prompt = _apply_description(base_prompt, context, normalized)
+            reporting.display_tool_description(
+                normalized,
+                f"Candidate tool description ({context.tool_name})",
+                context.panel_style,
+            )
+            candidates.append(updated_prompt)
+            if len(candidates) >= num_variations:
+                break
+
+        return candidates
+    except Exception as exc:  # pragma: no cover - fallback path
+        logger.warning(f"Failed to generate MCP tool descriptions: {exc}")
+        return []
+
+
+def initialize_population_mcp(
+    optimizer: Any,
+    prompt: chat_prompt.ChatPrompt,
+    context: EvolutionaryMCPContext,
+) -> list[chat_prompt.ChatPrompt]:
+    population_size = getattr(optimizer, "population_size", 1)
+    with reporting.initializing_population(
+        verbose=getattr(optimizer, "verbose", 1)
+    ) as init_pop_report:
+        init_pop_report.start(population_size)
+
+        population = [prompt]
+        num_to_generate = max(0, population_size - 1)
+        if num_to_generate > 0:
+            candidates = generate_tool_description_variations(
+                optimizer,
+                prompt,
+                context,
+                num_to_generate,
+            )
+            population.extend(candidates[:num_to_generate])
+
+        seen: set[str] = set()
+        final_population: list[chat_prompt.ChatPrompt] = []
+        for candidate in population:
+            key = json.dumps(candidate.get_messages())
+            if key in seen:
+                continue
+            seen.add(key)
+            final_population.append(candidate)
+
+        while len(final_population) < population_size:
+            final_population.append(prompt)
+
+        init_pop_report.end(final_population)
+        return final_population[:population_size]
+
+
+def tool_description_mutation(
+    optimizer: Any,
+    prompt: chat_prompt.ChatPrompt,
+    context: EvolutionaryMCPContext,
+) -> chat_prompt.ChatPrompt | None:
+    candidates = generate_tool_description_variations(optimizer, prompt, context, 1)
+    if not candidates:
+        return None
+
+    description = _current_tool_description(candidates[0], context.tool_segment_id)
+    if description:
+        reporting.display_tool_description(
+            description,
+            f"Updated tool description ({context.tool_name})",
+            context.panel_style,
+        )
+    return candidates[0]
+
+
+def finalize_mcp_result(
+    result: Any,
+    context: EvolutionaryMCPContext,
+    panel_style: str,
+) -> None:
+    final_tools = (
+        result.details.get("final_tools") if isinstance(result.details, dict) else None
+    )
+    tool_prompts = {
+        (tool.get("function", {}).get("name") or tool.get("name")): tool.get(
+            "function", {}
+        ).get("description")
+        for tool in (final_tools or [])
+    }
+    if tool_prompts.get(context.tool_name):
+        reporting.display_tool_description(
+            tool_prompts[context.tool_name],
+            f"Final tool description ({context.tool_name})",
+            panel_style,
+        )
+
+    if not tool_prompts and context.original_description:
+        tool_prompts = {context.tool_name: context.original_description}
+
+    if tool_prompts:
+        result.tool_prompts = tool_prompts
+
+
+# ---------------------------------------------------------------------------
+# Internal helpers
+# ---------------------------------------------------------------------------
+
+
+def _current_tool_description(
+    prompt: chat_prompt.ChatPrompt,
+    tool_segment_id: str,
+) -> str:
+    segments = {
+        segment.segment_id: segment for segment in extract_prompt_segments(prompt)
+    }
+    target = segments.get(tool_segment_id)
+    return target.content if target else ""
+
+
+def _extract_json_payload(response: str) -> dict[str, Any]:
+    try:
+        return json.loads(response)
+    except json.JSONDecodeError:
+        match = re.search(r"\{.*\}", response, re.DOTALL)
+        if not match:
+            raise ValueError("No JSON object found in LLM response")
+        return json.loads(match.group())
+
+
+def _apply_description(
+    prompt: chat_prompt.ChatPrompt,
+    context: EvolutionaryMCPContext,
+    description: str,
+) -> chat_prompt.ChatPrompt:
+    updated_prompt = apply_segment_updates(
+        prompt,
+        {context.tool_segment_id: description},
+    )
+    _sync_system_description(updated_prompt, description)
+    return updated_prompt
+
+
+def _sync_system_description(prompt: chat_prompt.ChatPrompt, description: str) -> None:
+    if not prompt.system:
+        return
+
+    marker_start = "<<TOOL_DESCRIPTION>>"
+    marker_end = "<<END_TOOL_DESCRIPTION>>"
+
+    start = prompt.system.find(marker_start)
+    end = prompt.system.find(marker_end)
+    if start == -1 or end == -1 or end <= start:
+        return
+
+    prefix = prompt.system[: start + len(marker_start)]
+    suffix = prompt.system[end:]
+    formatted_description = f"\n{description.strip()}\n"
+    prompt.system = f"{prefix}{formatted_description}{suffix}"