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

opik_optimizer/meta_prompt_optimizer/meta_prompt_optimizer.py

@@ -1,19 +1,21 @@
+import copy
 import json
 import logging
 import os
-from typing import Any, Callable, Dict, List, Optional, Tuple, Type
+import textwrap
+import warnings
+from typing import Any, cast
+from collections.abc import Callable
 
 import litellm
 import opik
 from litellm.caching import Cache
 from litellm.types.caching import LiteLLMCacheType
 from opik import Dataset
-from opik.api_objects import opik_client
 from opik.environment import get_tqdm_for_current_environment
 from opik.evaluation.models.litellm import opik_monitor as opik_litellm_monitor
 
 from opik_optimizer import task_evaluator
-from opik_optimizer import utils
 
 from .. import _throttle
 from ..base_optimizer import BaseOptimizer, OptimizationRound
@@ -21,6 +23,15 @@ from ..optimization_config import chat_prompt, mappers
 from ..optimization_result import OptimizationResult
 from ..optimizable_agent import OptimizableAgent
 from . import reporting
+import re
+
+from ..mcp_utils.mcp import PROMPT_TOOL_FOOTER, PROMPT_TOOL_HEADER
+from ..mcp_utils.mcp_workflow import (
+    MCPExecutionConfig,
+    MCPSecondPassCoordinator,
+    extract_tool_arguments,
+)
+from ..utils.prompt_segments import apply_segment_updates, extract_prompt_segments
 
 tqdm = get_tqdm_for_current_environment()
 
@@ -34,6 +45,48 @@ logger = logging.getLogger(__name__)  # Gets logger configured by setup_logging
 _rate_limiter = _throttle.get_rate_limiter_for_current_opik_installation()
 
 
+def _sync_tool_description_in_system(prompt: chat_prompt.ChatPrompt) -> None:
+    if not prompt.system or not getattr(prompt, "tools", None):
+        return
+
+    description = (
+        prompt.tools[0].get("function", {}).get("description") if prompt.tools else None
+    )
+    if not description:
+        return
+
+    tool_name = (
+        prompt.tools[0].get("function", {}).get("name") if prompt.tools else None
+    )
+
+    system_text = cast(str, prompt.system)
+    if PROMPT_TOOL_HEADER not in system_text or PROMPT_TOOL_FOOTER not in system_text:
+        return
+
+    start = system_text.index(PROMPT_TOOL_HEADER) + len(PROMPT_TOOL_HEADER)
+    end = system_text.index(PROMPT_TOOL_FOOTER)
+    description_text = description.strip()
+    system_text = (
+        system_text[:start] + "\n" + description_text + "\n" + system_text[end:]
+    )
+    prompt.system = system_text
+
+    if tool_name:
+        pattern = rf"(-\s*{re.escape(tool_name)}:\s)(.*)"
+
+        def _tool_section_replacer(match: re.Match[str]) -> str:
+            return f"{match.group(1)}{description_text}"
+
+        system_text = re.sub(
+            pattern,
+            _tool_section_replacer,
+            system_text,
+            count=1,
+            flags=re.MULTILINE,
+        )
+        prompt.system = system_text
+
+
 class MetaPromptOptimizer(BaseOptimizer):
     """
     The Meta-Prompt Optimizer uses meta-prompting to improve prompts based on examples and performance.
@@ -82,13 +135,14 @@ class MetaPromptOptimizer(BaseOptimizer):
     def __init__(
         self,
         model: str,
-        reasoning_model: Optional[str] = None,
+        reasoning_model: str | None = None,
         rounds: int = DEFAULT_ROUNDS,
         num_prompts_per_round: int = DEFAULT_PROMPTS_PER_ROUND,
-        num_threads: Optional[int] = None,
+        num_threads: int | None = None,
         verbose: int = 1,
         enable_context: bool = True,
         n_threads: int = 12,
+        seed: int = 42,
         **model_kwargs: Any,
     ) -> None:
         """
@@ -103,22 +157,28 @@ class MetaPromptOptimizer(BaseOptimizer):
             **model_kwargs: Additional model parameters
         """
         if "project_name" in model_kwargs:
-            print(
-                "Removing `project_name` from constructor; it now belongs in the ChatPrompt()"
+            warnings.warn(
+                "The 'project_name' parameter in optimizer constructor is deprecated. "
+                "Set project_name in the ChatPrompt instead.",
+                DeprecationWarning,
+                stacklevel=2,
             )
             del model_kwargs["project_name"]
 
-        super().__init__(model=model, verbose=verbose, **model_kwargs)
+        super().__init__(model=model, verbose=verbose, seed=seed, **model_kwargs)
         self.reasoning_model = reasoning_model if reasoning_model is not None else model
         self.rounds = rounds
         self.num_prompts_per_round = num_prompts_per_round
         if num_threads is not None:
-            print("num_threads is deprecated; use n_threads instead")
+            warnings.warn(
+                "The 'num_threads' parameter is deprecated and will be removed in a future version. "
+                "Use 'n_threads' instead.",
+                DeprecationWarning,
+                stacklevel=2,
+            )
             n_threads = num_threads
         self.num_threads = n_threads
-        self.dataset: Optional[Dataset] = None
-        self._opik_client = opik_client.get_client_cached()
-        self.llm_call_counter = 0
+        self.dataset: Dataset | None = None
         self.enable_context = enable_context
         logger.debug(
             f"Initialized MetaPromptOptimizer with model={model}, reasoning_model={self.reasoning_model}"
@@ -127,16 +187,24 @@ class MetaPromptOptimizer(BaseOptimizer):
             f"Optimization rounds: {rounds}, Prompts/round: {num_prompts_per_round}"
         )
 
+    def get_optimizer_metadata(self) -> dict[str, Any]:
+        return {
+            "rounds": self.rounds,
+            "num_prompts_per_round": self.num_prompts_per_round,
+            "reasoning_model": self.reasoning_model,
+            "enable_context": self.enable_context,
+        }
+
     @_throttle.rate_limited(_rate_limiter)
     def _call_model(
         self,
         project_name: str,
-        messages: List[Dict[str, str]],
+        messages: list[dict[str, str]],
         is_reasoning: bool = False,
-        optimization_id: Optional[str] = None,
+        optimization_id: str | None = None,
     ) -> str:
         """Call the model with the given prompt and return the response."""
-        self.llm_call_counter += 1
+        self.increment_llm_counter()
         # Note: Basic retry logic could be added here using tenacity
         try:
             # Basic LLM parameters (e.g., temperature, max_tokens)
@@ -163,7 +231,7 @@ class MetaPromptOptimizer(BaseOptimizer):
             }
 
             # Prepare metadata that we want to be part of the LLM call context.
-            metadata_for_opik: Dict[str, Any] = {}
+            metadata_for_opik: dict[str, Any] = {}
             if project_name:
                 metadata_for_opik["project_name"] = (
                     project_name  # Top-level for general use
@@ -225,11 +293,12 @@ class MetaPromptOptimizer(BaseOptimizer):
         prompt: chat_prompt.ChatPrompt,
         dataset: opik.Dataset,
         metric: Callable,
-        n_samples: Optional[int] = None,
-        dataset_item_ids: Optional[List[str]] = None,
-        experiment_config: Optional[Dict] = None,
+        n_samples: int | None = None,
+        dataset_item_ids: list[str] | None = None,
+        experiment_config: dict | None = None,
         use_full_dataset: bool = True,
-        optimization_id: Optional[str] = None,
+        optimization_id: str | None = None,
+        mcp_config: MCPExecutionConfig | None = None,
         **kwargs: Any,
     ) -> float:
         """
@@ -266,58 +335,103 @@ class MetaPromptOptimizer(BaseOptimizer):
             subset_size = None  # Use all items for final checks
             logger.debug("Using full dataset for evaluation")
 
-        experiment_config = experiment_config or {}
-        experiment_config = {
-            **experiment_config,
-            **{
-                "optimizer": self.__class__.__name__,
-                "agent_class": self.agent_class.__name__,
-                "agent_config": prompt.to_dict(),
-                "metric": getattr(metric, "__name__", str(metric)),
-                "dataset": dataset.name,
-                "configuration": {
-                    "prompt": prompt.get_messages(),
-                    "n_samples": subset_size,
-                    "use_full_dataset": use_full_dataset,
-                },
-            },
-        }
-        if optimization_id:
-            experiment_config["optimization_id"] = optimization_id
-
-        def llm_task(dataset_item: Dict[str, Any]) -> Dict[str, str]:
-            # --- Step 1: Prepare the prompt for the LLM ---
-            # messages = [
-            #    {
-            #        "role": item["role"],
-            #        "content": item["content"].format(**dataset_item),
-            #    }
-            #    for item in prompt.get_messages()
-            # ]
-            # Step 1: create the agent
+        configuration_updates = self._drop_none(
+            {
+                "n_samples": subset_size,
+                "use_full_dataset": use_full_dataset,
+            }
+        )
+        meta_metadata = self._drop_none(
+            {
+                "optimization_id": optimization_id,
+                "stage": "trial_evaluation" if not use_full_dataset else "final_eval",
+            }
+        )
+        experiment_config = self._prepare_experiment_config(
+            prompt=prompt,
+            dataset=dataset,
+            metric=metric,
+            experiment_config=experiment_config,
+            configuration_updates=configuration_updates,
+            additional_metadata={"meta_prompt": meta_metadata}
+            if meta_metadata
+            else None,
+        )
+
+        def llm_task(dataset_item: dict[str, Any]) -> dict[str, str]:
             new_prompt = prompt.copy()
             messages = new_prompt.get_messages(dataset_item)
             new_prompt.set_messages(messages)
             agent = self.agent_class(new_prompt)
 
-            # --- Step 2: Call the model ---
-            try:
-                logger.debug(
-                    f"Calling LLM with prompt length: {sum(len(msg['content']) for msg in messages)}"
-                )
-                raw_model_output = agent.invoke(messages)
-                logger.debug(f"LLM raw response length: {len(raw_model_output)}")
-                logger.debug(f"LLM raw output: {raw_model_output}")
-            except Exception as e:
-                logger.error(f"Error calling model with prompt: {e}")
-                logger.error(f"Failed prompt: {messages}")
-                logger.error(
-                    f"Prompt length: {sum(len(msg['content']) for msg in messages)}"
+            if mcp_config is not None:
+                coordinator = mcp_config.coordinator
+                coordinator.reset()
+                try:
+                    logger.debug(
+                        "Calling MCP-enabled LLM with tool access; prompt length=%s",
+                        sum(len(msg["content"]) for msg in messages),
+                    )
+                    raw_model_output = agent.llm_invoke(
+                        messages=messages,
+                        seed=self.seed,
+                        allow_tool_use=True,
+                    )
+                except Exception as exc:
+                    logger.error("Error during MCP first pass: %s", exc)
+                    raise
+
+                second_pass_messages = coordinator.build_second_pass_messages(
+                    base_messages=messages,
+                    dataset_item=dataset_item,
                 )
-                raise
 
-            # --- Step 3: Clean the model's output before metric evaluation ---
-            cleaned_model_output = raw_model_output.strip()
+                if second_pass_messages is None and mcp_config.fallback_invoker:
+                    fallback_args = mcp_config.fallback_arguments(dataset_item)
+                    if fallback_args:
+                        logger.debug(
+                            "MCP fallback triggered for tool %s with args=%s",
+                            mcp_config.tool_name,
+                            fallback_args,
+                        )
+                        summary_override = mcp_config.fallback_invoker(fallback_args)
+                        second_pass_messages = coordinator.build_second_pass_messages(
+                            base_messages=messages,
+                            dataset_item=dataset_item,
+                            summary_override=summary_override,
+                        )
+
+                if second_pass_messages is not None:
+                    logger.debug(
+                        "Executing MCP second pass with %d messages",
+                        len(second_pass_messages),
+                    )
+                    final_response = agent.llm_invoke(
+                        messages=second_pass_messages,
+                        seed=self.seed,
+                        allow_tool_use=mcp_config.allow_tool_use_on_second_pass,
+                    )
+                else:
+                    final_response = raw_model_output
+
+                cleaned_model_output = final_response.strip()
+            else:
+                try:
+                    logger.debug(
+                        f"Calling LLM with prompt length: {sum(len(msg['content']) for msg in messages)}"
+                    )
+                    raw_model_output = agent.invoke(messages)
+                    logger.debug(f"LLM raw response length: {len(raw_model_output)}")
+                    logger.debug(f"LLM raw output: {raw_model_output}")
+                except Exception as e:
+                    logger.error(f"Error calling model with prompt: {e}")
+                    logger.error(f"Failed prompt: {messages}")
+                    logger.error(
+                        f"Prompt length: {sum(len(msg['content']) for msg in messages)}"
+                    )
+                    raise
+
+                cleaned_model_output = raw_model_output.strip()
 
             result = {
                 mappers.EVALUATED_LLM_TASK_OUTPUT: cleaned_model_output,
@@ -348,46 +462,39 @@ class MetaPromptOptimizer(BaseOptimizer):
         prompt: chat_prompt.ChatPrompt,
         dataset: Dataset,
         metric: Callable,
-        experiment_config: Optional[Dict] = None,
-        n_samples: Optional[int] = None,
+        experiment_config: dict | None = None,
+        n_samples: int | None = None,
         auto_continue: bool = False,
-        agent_class: Optional[Type[OptimizableAgent]] = None,
+        agent_class: type[OptimizableAgent] | None = None,
         **kwargs: Any,
     ) -> OptimizationResult:
+        mcp_config = kwargs.pop("mcp_config", None)
+        candidate_generator = kwargs.pop("candidate_generator", None)
+        candidate_generator_kwargs = kwargs.pop("candidate_generator_kwargs", None)
+
         """
         Optimize a prompt using meta-reasoning.
 
         Args:
+            prompt: The prompt to optimize
             dataset: The dataset to evaluate against
             metric: The metric to use for evaluation
             experiment_config: A dictionary to log with the experiments
             n_samples: The number of dataset items to use for evaluation
             auto_continue: If True, the algorithm may continue if goal not met
-            **kwargs: Additional arguments for evaluation
+            agent_class: Optional agent class to use
+            **kwargs: Additional arguments for evaluation, including:
+                mcp_config (MCPExecutionConfig | None): MCP tool calling configuration (default: None)
+                candidate_generator: Optional candidate generator
+                candidate_generator_kwargs: Optional kwargs for candidate generator
 
         Returns:
             OptimizationResult: Structured result containing optimization details
         """
-        if not isinstance(prompt, chat_prompt.ChatPrompt):
-            raise ValueError("Prompt must be a ChatPrompt object")
-
-        if not isinstance(dataset, Dataset):
-            raise ValueError("Dataset must be a Dataset object")
-
-        if not callable(metric):
-            raise ValueError(
-                "Metric must be a function that takes `dataset_item` and `llm_output` as arguments."
-            )
-
-        if prompt.model is None:
-            prompt.model = self.model
-        if prompt.model_kwargs is None:
-            prompt.model_kwargs = self.model_kwargs
-
-        if agent_class is None:
-            self.agent_class = utils.create_litellm_agent_class(prompt)
-        else:
-            self.agent_class = agent_class
+        # Use base class validation and setup methods
+        self.validate_optimization_inputs(prompt, dataset, metric)
+        self.configure_prompt_model(prompt)
+        self.agent_class = self.setup_agent_class(prompt, agent_class)
 
         total_items = len(dataset.get_items())
         if n_samples is not None and n_samples > total_items:
@@ -398,7 +505,7 @@ class MetaPromptOptimizer(BaseOptimizer):
 
         optimization = None
         try:
-            optimization = self._opik_client.create_optimization(
+            optimization = self.opik_client.create_optimization(
                 dataset_name=dataset.name,
                 objective_name=getattr(metric, "__name__", str(metric)),
                 metadata={"optimizer": self.__class__.__name__},
@@ -424,6 +531,7 @@ class MetaPromptOptimizer(BaseOptimizer):
                 "auto_continue": auto_continue,
             },
             verbose=self.verbose,
+            tools=getattr(prompt, "tools", None),
         )
 
         try:
@@ -436,6 +544,9 @@ class MetaPromptOptimizer(BaseOptimizer):
                 experiment_config=experiment_config,
                 n_samples=n_samples,
                 auto_continue=auto_continue,
+                mcp_config=mcp_config,
+                candidate_generator=candidate_generator,
+                candidate_generator_kwargs=candidate_generator_kwargs,
                 **kwargs,
             )
             if optimization:
@@ -449,40 +560,104 @@ class MetaPromptOptimizer(BaseOptimizer):
                 logger.debug("Optimization marked as cancelled")
             raise e
 
+    def optimize_mcp(
+        self,
+        prompt: chat_prompt.ChatPrompt,
+        dataset: Dataset,
+        metric: Callable,
+        *,
+        tool_name: str,
+        second_pass: MCPSecondPassCoordinator,
+        experiment_config: dict | None = None,
+        n_samples: int | None = None,
+        auto_continue: bool = False,
+        agent_class: type[OptimizableAgent] | None = None,
+        fallback_invoker: Callable[[dict[str, Any]], str] | None = None,
+        fallback_arguments: Callable[[Any], dict[str, Any]] | None = None,
+        allow_tool_use_on_second_pass: bool = False,
+        **kwargs: Any,
+    ) -> OptimizationResult:
+        panel_style = kwargs.pop("tool_panel_style", "bright_magenta")
+
+        if prompt.tools is None or not prompt.tools:
+            raise ValueError("Prompt must include tools for MCP optimization")
+
+        fallback_args_fn = fallback_arguments or extract_tool_arguments
+
+        if fallback_invoker is None:
+            function_map = prompt.function_map or {}
+            fallback_invoker = function_map.get(tool_name)
+
+        mcp_config = MCPExecutionConfig(
+            coordinator=second_pass,
+            tool_name=tool_name,
+            fallback_arguments=fallback_args_fn,
+            fallback_invoker=fallback_invoker,
+            allow_tool_use_on_second_pass=allow_tool_use_on_second_pass,
+        )
+
+        tool_segment_id = f"tool:{tool_name}"
+        segments = extract_prompt_segments(prompt)
+        if tool_segment_id not in {segment.segment_id for segment in segments}:
+            raise ValueError(f"Tool '{tool_name}' not present in prompt tools")
+
+        return self.optimize_prompt(
+            prompt=prompt,
+            dataset=dataset,
+            metric=metric,
+            experiment_config=experiment_config,
+            n_samples=n_samples,
+            auto_continue=auto_continue,
+            agent_class=agent_class,
+            mcp_config=mcp_config,
+            candidate_generator=self._generate_mcp_candidate_prompts,
+            candidate_generator_kwargs={
+                "tool_segment_id": tool_segment_id,
+                "tool_name": tool_name,
+                "panel_style": panel_style,
+            },
+            tool_panel_style=panel_style,
+            **kwargs,
+        )
+
     def _optimize_prompt(
         self,
-        optimization_id: Optional[str],
+        optimization_id: str | None,
         prompt: chat_prompt.ChatPrompt,
         dataset: Dataset,
         metric: Callable,
-        experiment_config: Optional[Dict],
-        n_samples: Optional[int],
+        experiment_config: dict | None,
+        n_samples: int | None,
         auto_continue: bool,
+        mcp_config: MCPExecutionConfig | None = None,
+        candidate_generator: None
+        | (Callable[..., list[chat_prompt.ChatPrompt]]) = None,
+        candidate_generator_kwargs: dict[str, Any] | None = None,
+        tool_panel_style: str = "bright_magenta",
         **kwargs: Any,
     ) -> OptimizationResult:
         self.auto_continue = auto_continue
         self.dataset = dataset
         self.prompt = prompt
-        self.llm_call_counter = 0  # Reset counter for run
+        self.reset_counters()  # Reset counters for run
         initial_prompt = prompt
 
         current_prompt = prompt
-        experiment_config = experiment_config or {}
-        experiment_config = {
-            **experiment_config,
-            **{
-                "optimizer": self.__class__.__name__,
-                "agent_class": self.agent_class.__name__,
-                "agent_config": prompt.to_dict(),
-                "metric": getattr(metric, "__name__", str(metric)),
-                "dataset": dataset.name,
-                "configuration": {
-                    "prompt": prompt.get_messages(),
-                    "rounds": self.rounds,
-                    "num_prompts_per_round": self.num_prompts_per_round,
-                },
-            },
-        }
+        configuration_updates = self._drop_none(
+            {
+                "rounds": self.rounds,
+                "num_prompts_per_round": self.num_prompts_per_round,
+            }
+        )
+        meta_metadata = {"stage": "initial"}
+        experiment_config = self._prepare_experiment_config(
+            prompt=prompt,
+            dataset=dataset,
+            metric=metric,
+            experiment_config=experiment_config,
+            configuration_updates=configuration_updates,
+            additional_metadata={"meta_prompt": meta_metadata},
+        )
 
         with reporting.display_evaluation(verbose=self.verbose) as baseline_reporter:
             initial_score = self._evaluate_prompt(
@@ -494,10 +669,11 @@ class MetaPromptOptimizer(BaseOptimizer):
                 experiment_config=experiment_config,
                 use_full_dataset=n_samples is None,
                 verbose=self.verbose,
+                mcp_config=mcp_config,
             )
             best_score = initial_score
             best_prompt = current_prompt
-            rounds: List[OptimizationRound] = []
+            rounds: list[OptimizationRound] = []
 
             baseline_reporter.set_score(initial_score)
 
@@ -510,8 +686,11 @@ class MetaPromptOptimizer(BaseOptimizer):
                 previous_best_score = best_score
 
                 # Step 1. Create a set of candidate prompts
+                generator = candidate_generator or self._generate_candidate_prompts
+                generator_kwargs = dict(candidate_generator_kwargs or {})
+
                 try:
-                    candidate_prompts = self._generate_candidate_prompts(
+                    candidate_prompts = generator(
                         project_name=self.agent_class.project_name,
                         current_prompt=best_prompt,
                         best_score=best_score,
@@ -519,25 +698,25 @@ class MetaPromptOptimizer(BaseOptimizer):
                         previous_rounds=rounds,
                         metric=metric,
                         optimization_id=optimization_id,
+                        **generator_kwargs,
                     )
                 except Exception as e:
                     round_reporter.failed_to_generate(self.num_prompts_per_round, e)
                     continue
 
                 # Step 2. Score each candidate prompt
-                prompt_scores: List[Tuple[chat_prompt.ChatPrompt, float]] = []
+                prompt_scores: list[tuple[chat_prompt.ChatPrompt, float]] = []
                 for candidate_count, prompt in enumerate(candidate_prompts):
                     with reporting.display_prompt_candidate_scoring_report(
                         verbose=self.verbose
                     ) as eval_report:
                         eval_report.set_generated_prompts(candidate_count, prompt)
 
-                        new_prompt = current_prompt.copy()
-                        new_prompt.set_messages(prompt.get_messages())
+                        candidate_prompt = prompt.copy()
 
                         try:
                             prompt_score = self._evaluate_prompt(
-                                prompt=new_prompt,
+                                prompt=candidate_prompt,
                                 optimization_id=optimization_id,
                                 dataset=dataset,
                                 metric=metric,
@@ -545,11 +724,12 @@ class MetaPromptOptimizer(BaseOptimizer):
                                 use_full_dataset=False,
                                 experiment_config=experiment_config,
                                 verbose=self.verbose,
+                                mcp_config=mcp_config,
                             )
 
                             eval_report.set_final_score(best_score, prompt_score)
                         except Exception:
-                            print("Failed evaluating agent; continuing...")
+                            logger.warning("Failed evaluating agent; continuing...")
                             prompt_score = 0
 
                     prompt_scores.append((prompt, prompt_score))
@@ -584,24 +764,39 @@ class MetaPromptOptimizer(BaseOptimizer):
                     best_score = best_cand_score_avg
                     best_prompt = best_candidate_this_round
 
+        if tool_panel_style and getattr(best_prompt, "tools", None):
+            description = (
+                best_prompt.tools[0].get("function", {}).get("description", "")
+                if best_prompt.tools
+                else ""
+            )
+            if description.strip():
+                reporting.display_tool_description(
+                    description.strip(),
+                    "Final tool description",
+                    tool_panel_style,
+                )
+
         reporting.display_result(
             initial_score,
             best_score,
             best_prompt.get_messages() if best_prompt is not None else [],
             verbose=self.verbose,
+            tools=getattr(best_prompt, "tools", None) if best_prompt else None,
         )
 
         return self._create_result(
             metric,
-            initial_prompt=initial_prompt.get_messages()
-            if initial_prompt is not None
-            else [],
+            initial_prompt=(
+                initial_prompt.get_messages() if initial_prompt is not None else []
+            ),
             best_prompt=best_prompt.get_messages() if best_prompt is not None else [],
             best_score=best_score,
             initial_score=initial_score,
             rounds=rounds,
             dataset_id=dataset.id,
             optimization_id=optimization_id,
+            best_tools=getattr(best_prompt, "tools", None) if best_prompt else None,
         )
 
     def _calculate_improvement(
@@ -620,19 +815,24 @@ class MetaPromptOptimizer(BaseOptimizer):
         current_best_prompt: chat_prompt.ChatPrompt,
         current_best_score: float,
         best_prompt_overall: chat_prompt.ChatPrompt,
-        evaluated_candidates: List[Tuple[chat_prompt.ChatPrompt, float]],
+        evaluated_candidates: list[tuple[chat_prompt.ChatPrompt, float]],
         previous_best_score: float,
         improvement_this_round: float,
     ) -> OptimizationRound:
         """Create an OptimizationRound object with the current round's data."""
-        generated_prompts_log = []
+        generated_prompts_log: list[dict[str, Any]] = []
         for prompt, score in evaluated_candidates:
             improvement_vs_prev = self._calculate_improvement(
                 score, previous_best_score
             )
+            tool_entries: list[Any] = []
+            if getattr(prompt, "tools", None):
+                tool_entries = copy.deepcopy(list(prompt.tools or []))
+
             generated_prompts_log.append(
                 {
                     "prompt": prompt.get_messages(),
+                    "tools": tool_entries,
                     "score": score,
                     "improvement": improvement_vs_prev,
                 }
@@ -651,13 +851,14 @@ class MetaPromptOptimizer(BaseOptimizer):
     def _create_result(
         self,
         metric: Callable,
-        initial_prompt: List[Dict[str, str]],
-        best_prompt: List[Dict[str, str]],
+        initial_prompt: list[dict[str, str]],
+        best_prompt: list[dict[str, str]],
         best_score: float,
         initial_score: float,
-        rounds: List[OptimizationRound],
-        dataset_id: Optional[str],
-        optimization_id: Optional[str],
+        rounds: list[OptimizationRound],
+        dataset_id: str | None,
+        optimization_id: str | None,
+        best_tools: list[dict[str, Any]] | None,
     ) -> OptimizationResult:
         """Create the final OptimizationResult object."""
         details = {
@@ -670,6 +871,18 @@ class MetaPromptOptimizer(BaseOptimizer):
             "temperature": self.model_kwargs.get("temperature"),
         }
 
+        if best_tools:
+            details["final_tools"] = best_tools
+
+        tool_prompts = None
+        if best_tools:
+            tool_prompts = {
+                (tool.get("function", {}).get("name") or f"tool_{idx}"): tool.get(
+                    "function", {}
+                ).get("description")
+                for idx, tool in enumerate(best_tools)
+            }
+
         return OptimizationResult(
             optimizer=self.__class__.__name__,
             prompt=best_prompt,
@@ -679,8 +892,10 @@ class MetaPromptOptimizer(BaseOptimizer):
             metric_name=getattr(metric, "__name__", str(metric)),
             details=details,
             llm_calls=self.llm_call_counter,
+            tool_calls=self.tool_call_counter,
             dataset_id=dataset_id,
             optimization_id=optimization_id,
+            tool_prompts=tool_prompts,
         )
 
     def _get_task_context(self, metric: Callable) -> str:
@@ -714,11 +929,11 @@ class MetaPromptOptimizer(BaseOptimizer):
         current_prompt: chat_prompt.ChatPrompt,
         best_score: float,
         round_num: int,
-        previous_rounds: List[OptimizationRound],
+        previous_rounds: list[OptimizationRound],
         metric: Callable,
-        optimization_id: Optional[str] = None,
-        project_name: Optional[str] = None,
-    ) -> List[chat_prompt.ChatPrompt]:
+        optimization_id: str | None = None,
+        project_name: str | None = None,
+    ) -> list[chat_prompt.ChatPrompt]:
         """Generate candidate prompts using meta-prompting."""
         with reporting.display_candidate_generation_report(
             self.num_prompts_per_round, verbose=self.verbose
@@ -819,7 +1034,7 @@ class MetaPromptOptimizer(BaseOptimizer):
                     )
 
                 # Extract and log valid prompts
-                valid_prompts: List[chat_prompt.ChatPrompt] = []
+                valid_prompts: list[chat_prompt.ChatPrompt] = []
                 for item in json_result["prompts"]:
                     if (
                         isinstance(item, dict)
@@ -870,7 +1085,128 @@ class MetaPromptOptimizer(BaseOptimizer):
                     f"Unexpected error during candidate prompt generation: {e}"
                 )
 
-    def _build_history_context(self, previous_rounds: List[OptimizationRound]) -> str:
+    def _generate_mcp_candidate_prompts(
+        self,
+        current_prompt: chat_prompt.ChatPrompt,
+        best_score: float,
+        round_num: int,
+        previous_rounds: list[OptimizationRound],
+        metric: Callable,
+        tool_segment_id: str,
+        tool_name: str,
+        optimization_id: str | None = None,
+        project_name: str | None = None,
+        panel_style: str = "bright_magenta",
+    ) -> list[chat_prompt.ChatPrompt]:
+        segments = {
+            segment.segment_id: segment
+            for segment in extract_prompt_segments(current_prompt)
+        }
+        if tool_segment_id not in segments:
+            raise ValueError(f"Tool segment '{tool_segment_id}' not found in prompt")
+
+        target_segment = segments[tool_segment_id]
+        current_description = target_segment.content
+        tool_metadata = target_segment.metadata.get("raw_tool", {})
+
+        history_context = self._build_history_context(previous_rounds)
+
+        instruction = textwrap.dedent(
+            f"""
+            Current tool name: {tool_name}
+            Current tool description:
+            ---
+            {current_description}
+            ---
+
+            Tool metadata (JSON):
+            {json.dumps(tool_metadata, indent=2)}
+
+            Current best score: {best_score:.4f}
+            {history_context}
+
+            Generate {self.num_prompts_per_round} improved descriptions for this tool.
+            Each description should clarify expected input arguments and set explicit expectations
+            for how the tool output must be used in the final response.
+            Avoid changing unrelated parts of the prompt. Focus only on the description text for `{tool_name}`.
+
+            Return a JSON object of the form:
+            {{
+              "prompts": [
+                {{
+                  "tool_description": "...",
+                  "improvement_focus": "...",
+                  "reasoning": "..."
+                }}
+              ]
+            }}
+            """
+        ).strip()
+
+        with reporting.display_candidate_generation_report(
+            self.num_prompts_per_round, verbose=self.verbose
+        ) as candidate_generation_report:
+            try:
+                content = self._call_model(
+                    project_name,
+                    messages=[
+                        {"role": "system", "content": self._REASONING_SYSTEM_PROMPT},
+                        {"role": "user", "content": instruction},
+                    ],
+                    is_reasoning=True,
+                    optimization_id=optimization_id,
+                )
+
+                try:
+                    json_result = json.loads(content)
+                except json.JSONDecodeError:
+                    import re
+
+                    json_match = re.search(r"\{.*\}", content, re.DOTALL)
+                    if not json_match:
+                        raise ValueError("No JSON object found in reasoning output")
+                    json_result = json.loads(json_match.group())
+
+                prompts_payload = json_result.get("prompts")
+                if not isinstance(prompts_payload, list):
+                    raise ValueError("Reasoning output missing 'prompts' list")
+
+                candidate_generation_report.set_generated_prompts()
+
+                candidates: list[chat_prompt.ChatPrompt] = []
+                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
+
+                    updated_prompt = apply_segment_updates(
+                        current_prompt,
+                        {tool_segment_id: description.strip()},
+                    )
+                    _sync_tool_description_in_system(updated_prompt)
+                    if (
+                        description.strip()
+                        and description.strip() != current_description.strip()
+                    ):
+                        reporting.display_tool_description(
+                            description.strip(),
+                            f"Round {round_num + 1} tool description",
+                            panel_style,
+                        )
+                    candidates.append(updated_prompt)
+
+                if not candidates:
+                    raise ValueError(
+                        "Reasoning output did not produce valid tool descriptions"
+                    )
+
+                return candidates
+            except Exception as exc:
+                raise ValueError(f"Error generating MCP prompt candidates: {exc}")
+
+    def _build_history_context(self, previous_rounds: list[OptimizationRound]) -> str:
         """Build context from previous optimization rounds."""
         if not previous_rounds:
             return ""
@@ -896,7 +1232,7 @@ class MetaPromptOptimizer(BaseOptimizer):
 
     def _get_evaluation_subset(
         self, dataset: opik.Dataset, min_size: int = 20, max_size: int = 100
-    ) -> List[Dict[str, Any]]:
+    ) -> list[dict[str, Any]]:
         """Get a random subset of the dataset for evaluation.
 
         Returns: