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

opik_optimizer/optimization_config/chat_prompt.py

@@ -1,4 +1,5 @@
-from typing import Any, Dict, List, Optional, Union, Callable
+from typing import Any
+from collections.abc import Callable
 
 import copy
 
@@ -10,7 +11,7 @@ from opik import track
 class Tool(BaseModel):
     name: str = Field(..., description="Name of the tool")
     description: str = Field(..., description="Description of the tool")
-    parameters: Dict[str, Any] = Field(
+    parameters: dict[str, Any] = Field(
         ..., description="JSON Schema defining the input parameters for the tool"
     )
 
@@ -33,14 +34,14 @@ class ChatPrompt:
     def __init__(
         self,
         name: str = "chat-prompt",
-        system: Optional[str] = None,
-        user: Optional[str] = None,
-        messages: Optional[List[Dict[str, str]]] = None,
-        tools: Optional[List[Dict[str, Any]]] = None,
-        function_map: Optional[Dict[str, Callable]] = None,
-        model: Optional[str] = None,
-        invoke: Optional[Callable] = None,
-        project_name: Optional[str] = "Default Project",
+        system: str | None = None,
+        user: str | None = None,
+        messages: list[dict[str, str]] | None = None,
+        tools: list[dict[str, Any]] | None = None,
+        function_map: dict[str, Callable] | None = None,
+        model: str | None = None,
+        invoke: Callable | None = None,
+        project_name: str | None = "Default Project",
         **model_kwargs: Any,
     ) -> None:
         if system is None and user is None and messages is None:
@@ -97,8 +98,8 @@ class ChatPrompt:
 
     def get_messages(
         self,
-        dataset_item: Optional[Dict[str, str]] = None,
-    ) -> List[Dict[str, str]]:
+        dataset_item: dict[str, str] | None = None,
+    ) -> list[dict[str, str]]:
         # This is a copy, so we can alter the messages:
         messages = self._standardize_prompts()
 
@@ -113,8 +114,8 @@ class ChatPrompt:
                         )
         return messages
 
-    def _standardize_prompts(self, **kwargs: Any) -> List[Dict[str, str]]:
-        standardize_messages: List[Dict[str, str]] = []
+    def _standardize_prompts(self, **kwargs: Any) -> list[dict[str, str]]:
+        standardize_messages: list[dict[str, str]] = []
 
         if self.system is not None:
             standardize_messages.append({"role": "system", "content": self.system})
@@ -128,13 +129,13 @@ class ChatPrompt:
 
         return copy.deepcopy(standardize_messages)
 
-    def to_dict(self) -> Dict[str, Union[str, List[Dict[str, str]]]]:
+    def to_dict(self) -> dict[str, str | list[dict[str, str]]]:
         """Convert ChatPrompt to a dictionary for JSON serialization.
 
         Returns:
             Dict containing the serializable representation of this ChatPrompt
         """
-        retval: Dict[str, Union[str, List[Dict[str, str]]]] = {}
+        retval: dict[str, str | list[dict[str, str]]] = {}
         if self.system is not None:
             retval["system"] = self.system
         if self.user is not None:
@@ -144,29 +145,49 @@ class ChatPrompt:
         return retval
 
     def copy(self) -> "ChatPrompt":
+        """Shallow clone preserving model configuration and tools."""
+
+        # TODO(opik-mcp): once we introduce a dedicated MCP prompt subclass,
+        # migrate callers away from generic copies so optimizer metadata stays typed.
+        model_kwargs = (
+            copy.deepcopy(self.model_kwargs) if self.model_kwargs is not None else {}
+        )
         return ChatPrompt(
+            name=self.name,
             system=self.system,
             user=self.user,
             messages=copy.deepcopy(self.messages),
-            tools=self.tools,
+            tools=copy.deepcopy(self.tools),
             function_map=self.function_map,
+            model=self.model,
+            invoke=self.invoke,
+            project_name=self.project_name,
+            **model_kwargs,
         )
 
-    def set_messages(self, messages: List[Dict[str, Any]]) -> None:
+    def set_messages(self, messages: list[dict[str, Any]]) -> None:
         self.system = None
         self.user = None
         self.messages = copy.deepcopy(messages)
 
+    # TODO(opik): remove this stop-gap once MetaPromptOptimizer supports MCP.
+    # Provides a second-pass flow so tool results can be appended before
+    # rerunning the model.
+    def with_messages(self, messages: list[dict[str, Any]]) -> "ChatPrompt":
+        cloned = self.copy()
+        cloned.set_messages(messages)
+        return cloned
+
     @classmethod
     def model_validate(
         cls,
         obj: Any,
         *,
-        strict: Optional[bool] = None,
-        from_attributes: Optional[bool] = None,
-        context: Optional[Any] = None,
-        by_alias: Optional[bool] = None,
-        by_name: Optional[bool] = None,
+        strict: bool | None = None,
+        from_attributes: bool | None = None,
+        context: Any | None = None,
+        by_alias: bool | None = None,
+        by_name: bool | None = None,
     ) -> "ChatPrompt":
         """Custom validation method to handle nested objects during deserialization."""
         return ChatPrompt(

opik_optimizer/optimization_config/configs.py

@@ -1,6 +1,6 @@
 """Module containing configuration classes for optimization."""
 
-from typing import Any, List
+from typing import Any
 
 import pydantic
 
@@ -12,6 +12,6 @@ class TaskConfig(pydantic.BaseModel):
 
     instruction_prompt: str
     use_chat_prompt: bool = False
-    input_dataset_fields: List[str]
+    input_dataset_fields: list[str]
     output_dataset_field: str
-    tools: List[Any] = []
+    tools: list[Any] = []

opik_optimizer/optimization_config/mappers.py

@@ -1,4 +1,5 @@
-from typing import Dict, Callable, Optional, Any, Union
+from typing import Any
+from collections.abc import Callable
 
 EVALUATED_LLM_TASK_OUTPUT = "llm_output"
 
@@ -8,8 +9,8 @@ class Mapper:
 
     def __init__(
         self,
-        name: Optional[str] = None,
-        transform: Optional[Callable[[Any], Any]] = None,
+        name: str | None = None,
+        transform: Callable[[Any], Any] | None = None,
     ):
         if name is not None and transform is not None:
             raise ValueError("Only one of name or transform can be provided")
@@ -27,9 +28,9 @@ class Mapper:
 
 def from_dataset_field(
     *,
-    name: Optional[str] = None,
-    transform: Optional[Callable[[Dict[str, Any]], Any]] = None,
-) -> Union[str, Callable[[Dict[str, Any]], Any]]:
+    name: str | None = None,
+    transform: Callable[[dict[str, Any]], Any] | None = None,
+) -> str | Callable[[dict[str, Any]], Any]:
     if name is not None and transform is not None:
         raise ValueError("Only one of name or transform can be provided")
 
@@ -47,8 +48,8 @@ def from_llm_response_text() -> str:
 
 
 def from_agent_output(
-    *, name: Optional[str] = None, transform: Optional[Callable[[Any], Any]] = None
-) -> Union[str, Callable[[Any], Any]]:
+    *, name: str | None = None, transform: Callable[[Any], Any] | None = None
+) -> str | Callable[[Any], Any]:
     if name is not None and transform is not None:
         raise ValueError("Only one of name or transform can be provided")
 

opik_optimizer/optimization_result.py

@@ -1,6 +1,6 @@
 """Module containing the OptimizationResult class."""
 
-from typing import Any, Dict, List, Optional
+from typing import Any
 
 import pydantic
 import rich
@@ -13,25 +13,25 @@ class OptimizationResult(pydantic.BaseModel):
 
     optimizer: str = "Optimizer"
 
-    prompt: List[Dict[str, str]]
+    prompt: list[dict[str, str]]
     score: float
     metric_name: str
 
-    optimization_id: Optional[str] = None
-    dataset_id: Optional[str] = None
+    optimization_id: str | None = None
+    dataset_id: str | None = None
 
     # Initial score
-    initial_prompt: Optional[List[Dict[str, str]]] = None
-    initial_score: Optional[float] = None
+    initial_prompt: list[dict[str, str]] | None = None
+    initial_score: float | None = None
 
-    details: Dict[str, Any] = pydantic.Field(default_factory=dict)
-    history: List[Dict[str, Any]] = []
-    llm_calls: Optional[int] = None
+    details: dict[str, Any] = pydantic.Field(default_factory=dict)
+    history: list[dict[str, Any]] = []
+    llm_calls: int | None = None
 
     # MIPRO specific
-    demonstrations: Optional[List[Dict[str, Any]]] = None
-    mipro_prompt: Optional[str] = None
-    tool_prompts: Optional[Dict[str, str]] = None
+    demonstrations: list[dict[str, Any]] | None = None
+    mipro_prompt: str | None = None
+    tool_prompts: dict[str, str] | None = None
 
     model_config = pydantic.ConfigDict(arbitrary_types_allowed=True)
 
@@ -40,7 +40,7 @@ class OptimizationResult(pydantic.BaseModel):
             optimization_id=self.optimization_id, dataset_id=self.dataset_id
         )
 
-    def model_dump(self, *kargs: Any, **kwargs: Any) -> Dict[str, Any]:
+    def model_dump(self, *kargs: Any, **kwargs: Any) -> dict[str, Any]:
         return super().model_dump(*kargs, **kwargs)
 
     def _calculate_improvement_str(self) -> str:
@@ -205,4 +205,11 @@ class OptimizationResult(pydantic.BaseModel):
         """
         console = get_console()
         console.print(self)
-        print("Optimization run link:", self.get_run_link())
+        # Gracefully handle cases where optimization tracking isn't available
+        if self.dataset_id and self.optimization_id:
+            try:
+                print("Optimization run link:", self.get_run_link())
+            except Exception:
+                print("Optimization run link: No optimization run link available")
+        else:
+            print("Optimization run link: No optimization run link available")

opik_optimizer/reporting_utils.py

@@ -1,6 +1,7 @@
+import json
 import logging
 from contextlib import contextmanager
-from typing import Any, Dict, List, Optional, Union
+from typing import Any
 
 from rich import box
 from rich.console import Console, Group
@@ -20,7 +21,7 @@ def get_console(*args: Any, **kwargs: Any) -> Console:
 
 
 @contextmanager
-def convert_tqdm_to_rich(description: Optional[str] = None, verbose: int = 1) -> Any:
+def convert_tqdm_to_rich(description: str | None = None, verbose: int = 1) -> Any:
     """Context manager to convert tqdm to rich."""
     import opik.evaluation.engine.evaluation_tasks_executor
 
@@ -66,7 +67,7 @@ def suppress_opik_logs() -> Any:
         opik_logger.setLevel(original_level)
 
 
-def display_messages(messages: List[Dict[str, str]], prefix: str = "") -> None:
+def display_messages(messages: list[dict[str, str]], prefix: str = "") -> None:
     for i, msg in enumerate(messages):
         panel = Panel(
             Text(msg.get("content", ""), overflow="fold"),
@@ -90,11 +91,53 @@ def display_messages(messages: List[Dict[str, str]], prefix: str = "") -> None:
             console.print(Text(prefix) + Text.from_ansi(line))
 
 
+def _format_tool_panel(tool: dict[str, Any]) -> Panel:
+    function_block = tool.get("function", {})
+    name = function_block.get("name") or tool.get("name", "unknown_tool")
+    description = function_block.get("description", "")
+    parameters = function_block.get("parameters", {})
+
+    body_lines: list[str] = []
+    if description:
+        body_lines.append(description)
+    if parameters:
+        formatted_schema = json.dumps(parameters, indent=2, sort_keys=True)
+        body_lines.append("\nSchema:\n" + formatted_schema)
+
+    content = Text(
+        "\n".join(body_lines) if body_lines else "(no metadata)", overflow="fold"
+    )
+    return Panel(
+        content,
+        title=f"tool: {name}",
+        title_align="left",
+        border_style="cyan",
+        width=PANEL_WIDTH,
+        padding=(1, 2),
+    )
+
+
+def _display_tools(tools: list[dict[str, Any]] | None) -> None:
+    if not tools:
+        return
+
+    console = get_console()
+    console.print(Text("\nTools registered:\n", style="bold"))
+    for tool in tools:
+        panel = _format_tool_panel(tool)
+        with console.capture() as capture:
+            console.print(panel)
+        rendered_panel = capture.get()
+        for line in rendered_panel.splitlines():
+            console.print(Text.from_ansi(line))
+    console.print("")
+
+
 def get_link_text(
     pre_text: str,
     link_text: str,
-    optimization_id: Optional[str] = None,
-    dataset_id: Optional[str] = None,
+    optimization_id: str | None = None,
+    dataset_id: str | None = None,
 ) -> Text:
     if optimization_id is not None and dataset_id is not None:
         optimization_url = get_optimization_run_url_by_id(
@@ -112,8 +155,8 @@ def get_link_text(
 
 def display_header(
     algorithm: str,
-    optimization_id: Optional[str] = None,
-    dataset_id: Optional[str] = None,
+    optimization_id: str | None = None,
+    dataset_id: str | None = None,
     verbose: int = 1,
 ) -> None:
     if verbose < 1:
@@ -140,8 +183,9 @@ def display_header(
 def display_result(
     initial_score: float,
     best_score: float,
-    best_prompt: List[Dict[str, str]],
+    best_prompt: list[dict[str, str]],
     verbose: int = 1,
+    tools: list[dict[str, Any]] | None = None,
 ) -> None:
     if verbose < 1:
         return
@@ -149,7 +193,7 @@ def display_result(
     console = get_console()
     console.print(Text("\n> Optimization complete\n"))
 
-    content: Union[Text, Panel] = []
+    content: Text | Panel = []
 
     if best_score > initial_score:
         if initial_score == 0:
@@ -199,9 +243,15 @@ def display_result(
         )
     )
 
+    if tools:
+        _display_tools(tools)
+
 
 def display_configuration(
-    messages: List[Dict[str, str]], optimizer_config: Dict[str, Any], verbose: int = 1
+    messages: list[dict[str, str]],
+    optimizer_config: dict[str, Any],
+    verbose: int = 1,
+    tools: list[dict[str, Any]] | None = None,
 ) -> None:
     """Displays the LLM messages and optimizer configuration using Rich panels."""
 
@@ -213,6 +263,7 @@ def display_configuration(
     console.print(Text("> Let's optimize the prompt:\n"))
 
     display_messages(messages)
+    _display_tools(tools)
 
     # Panel for configuration
     console.print(

opik_optimizer/task_evaluator.py

@@ -1,5 +1,6 @@
 import logging
-from typing import Any, Callable, Dict, List, Optional
+from typing import Any
+from collections.abc import Callable
 
 import opik
 from opik.evaluation import evaluator as opik_evaluator
@@ -38,14 +39,14 @@ def _create_metric_class(metric: Callable) -> base_metric.BaseMetric:
 
 def evaluate(
     dataset: opik.Dataset,
-    evaluated_task: Callable[[Dict[str, Any]], Dict[str, Any]],
+    evaluated_task: Callable[[dict[str, Any]], dict[str, Any]],
     metric: Callable,
     num_threads: int,
-    optimization_id: Optional[str] = None,
-    dataset_item_ids: Optional[List[str]] = None,
-    project_name: Optional[str] = None,
-    n_samples: Optional[int] = None,
-    experiment_config: Optional[Dict[str, Any]] = None,
+    optimization_id: str | None = None,
+    dataset_item_ids: list[str] | None = None,
+    project_name: str | None = None,
+    n_samples: int | None = None,
+    experiment_config: dict[str, Any] | None = None,
     verbose: int = 1,
 ) -> float:
     """
@@ -107,7 +108,7 @@ def evaluate(
         return 0.0
 
     # We may allow score aggregation customization.
-    score_results: List[score_result.ScoreResult] = [
+    score_results: list[score_result.ScoreResult] = [
         test_result.score_results[0] for test_result in result.test_results
     ]
     if not score_results:

opik_optimizer/utils/__init__.py

@@ -0,0 +1,15 @@
+"""Utility helpers exposed as part of the opik_optimizer package."""
+
+from .core import *  # noqa: F401,F403
+from .dataset_utils import *  # noqa: F401,F403
+from .prompt_segments import *  # noqa: F401,F403
+
+from . import core as _core
+from . import dataset_utils as _dataset_utils
+from . import prompt_segments as _prompt_segments
+
+__all__: list[str] = [
+    *getattr(_core, "__all__", []),
+    *getattr(_dataset_utils, "__all__", []),
+    *getattr(_prompt_segments, "__all__", []),
+]

opik_optimizer/{utils.py → utils/core.py}

@@ -2,18 +2,13 @@
 
 from typing import (
     Any,
-    Dict,
     Final,
     Literal,
-    Optional,
-    Type,
     TYPE_CHECKING,
-    List,
-    Callable,
 )
+from collections.abc import Callable
 
 import inspect
-import typing
 import base64
 import json
 import logging
@@ -22,6 +17,8 @@ import string
 import urllib.parse
 from types import TracebackType
 
+import requests
+
 import opik
 from opik.api_objects.opik_client import Opik
 from opik.api_objects.optimization import Optimization
@@ -30,8 +27,8 @@ ALLOWED_URL_CHARACTERS: Final[str] = ":/&?="
 logger = logging.getLogger(__name__)
 
 if TYPE_CHECKING:
-    from .optimizable_agent import OptimizableAgent
-    from .optimization_config.chat_prompt import ChatPrompt
+    from opik_optimizer.optimizable_agent import OptimizableAgent
+    from opik_optimizer.optimization_config.chat_prompt import ChatPrompt
 
 
 class OptimizationContextManager:
@@ -45,8 +42,8 @@ class OptimizationContextManager:
         client: Opik,
         dataset_name: str,
         objective_name: str,
-        name: Optional[str] = None,
-        metadata: Optional[Dict[str, Any]] = None,
+        name: str | None = None,
+        metadata: dict[str, Any] | None = None,
     ):
         """
         Initialize the optimization context.
@@ -63,9 +60,9 @@ class OptimizationContextManager:
         self.objective_name = objective_name
         self.name = name
         self.metadata = metadata
-        self.optimization: Optional[Optimization] = None
+        self.optimization: Optimization | None = None
 
-    def __enter__(self) -> Optional[Optimization]:
+    def __enter__(self) -> Optimization | None:
         """Create and return the optimization."""
         try:
             self.optimization = self.client.create_optimization(
@@ -88,9 +85,9 @@ class OptimizationContextManager:
 
     def __exit__(
         self,
-        exc_type: Optional[Type[BaseException]],
-        exc_val: Optional[BaseException],
-        exc_tb: Optional[TracebackType],
+        exc_type: type[BaseException] | None,
+        exc_val: BaseException | None,
+        exc_tb: TracebackType | None,
     ) -> Literal[False]:
         """Update optimization status based on context exit."""
         if self.optimization is None:
@@ -227,8 +224,8 @@ def optimization_context(
     client: Opik,
     dataset_name: str,
     objective_name: str,
-    name: Optional[str] = None,
-    metadata: Optional[Dict[str, Any]] = None,
+    name: str | None = None,
+    metadata: dict[str, Any] | None = None,
 ) -> OptimizationContextManager:
     """
     Create a context manager for handling optimization lifecycle.
@@ -258,7 +255,7 @@ def ensure_ending_slash(url: str) -> str:
 
 
 def get_optimization_run_url_by_id(
-    dataset_id: Optional[str], optimization_id: Optional[str]
+    dataset_id: str | None, optimization_id: str | None
 ) -> str:
     if dataset_id is None or optimization_id is None:
         raise ValueError(
@@ -276,11 +273,11 @@ def get_optimization_run_url_by_id(
     return urllib.parse.urljoin(ensure_ending_slash(url_override), run_path)
 
 
-def create_litellm_agent_class(prompt: "ChatPrompt") -> Type["OptimizableAgent"]:
+def create_litellm_agent_class(prompt: "ChatPrompt") -> type["OptimizableAgent"]:
     """
     Create a LiteLLMAgent from a chat prompt.
     """
-    from .optimizable_agent import OptimizableAgent
+    from opik_optimizer.optimizable_agent import OptimizableAgent
 
     if prompt.invoke is not None:
 
@@ -290,7 +287,7 @@ def create_litellm_agent_class(prompt: "ChatPrompt") -> Type["OptimizableAgent"]
             project_name = prompt.project_name
 
             def invoke(
-                self, messages: List[Dict[str, str]], seed: Optional[int] = None
+                self, messages: list[dict[str, str]], seed: int | None = None
             ) -> str:
                 return prompt.invoke(
                     self.model, messages, prompt.tools, **self.model_kwargs
@@ -307,13 +304,13 @@ def create_litellm_agent_class(prompt: "ChatPrompt") -> Type["OptimizableAgent"]
 
 
 def function_to_tool_definition(
-    func: Callable, description: Optional[str] = None
-) -> Dict[str, Any]:
+    func: Callable, description: str | None = None
+) -> dict[str, Any]:
     sig = inspect.signature(func)
     doc = description or func.__doc__ or ""
 
-    properties: Dict[str, Dict[str, str]] = {}
-    required: List[str] = []
+    properties: dict[str, dict[str, str]] = {}
+    required: list[str] = []
 
     for name, param in sig.parameters.items():
         param_type = (
@@ -350,7 +347,95 @@ def python_type_to_json_type(python_type: type) -> str:
         return "boolean"
     elif python_type in [dict]:
         return "object"
-    elif python_type in [list, typing.List]:
+    elif python_type in [list, list]:
         return "array"
     else:
         return "string"  # default fallback
+
+
+def search_wikipedia(query: str, use_api: bool = False) -> list[str]:
+    """
+    This agent is used to search wikipedia. It can retrieve additional details
+    about a topic.
+
+    Args:
+        query: The search query string
+        use_api: If True, directly use Wikipedia API instead of ColBERTv2.
+                If False (default), try ColBERTv2 first with API fallback.
+    """
+    if use_api:
+        # Directly use Wikipedia API when requested
+        try:
+            return _search_wikipedia_api(query)
+        except Exception as api_error:
+            print(f"Wikipedia API failed: {api_error}")
+            return [f"Wikipedia search unavailable. Query was: {query}"]
+
+    # Default behavior: Try ColBERTv2 first with API fallback
+    from .colbert import ColBERTv2
+
+    # Try ColBERTv2 first with a short timeout
+    try:
+        colbert = ColBERTv2(url="http://20.102.90.50:2017/wiki17_abstracts")
+        # Use a shorter timeout by modifying the max_retries parameter
+        results = colbert(query, k=3, max_retries=1)
+        return [str(item.text) for item in results if hasattr(item, "text")]
+    except Exception as e:
+        print(f"ColBERTv2 search failed: {e}")
+        # Fallback to Wikipedia API
+        try:
+            return _search_wikipedia_api(query)
+        except Exception as api_error:
+            print(f"Wikipedia API fallback also failed: {api_error}")
+            return [f"Wikipedia search unavailable. Query was: {query}"]
+
+
+def _search_wikipedia_api(query: str, max_results: int = 3) -> list[str]:
+    """
+    Fallback Wikipedia search using the Wikipedia API.
+    """
+    try:
+        # First, search for pages using the search API
+        search_params: dict[str, str | int] = {
+            "action": "query",
+            "format": "json",
+            "list": "search",
+            "srsearch": query,
+            "srlimit": max_results,
+            "srprop": "snippet",
+        }
+
+        headers = {
+            "User-Agent": "OpikOptimizer/1.0 (https://github.com/opik-ai/opik-optimizer)"
+        }
+        search_response = requests.get(
+            "https://en.wikipedia.org/w/api.php",
+            params=search_params,
+            headers=headers,
+            timeout=5,
+        )
+
+        if search_response.status_code != 200:
+            raise Exception(f"Search API returned status {search_response.status_code}")
+
+        search_data = search_response.json()
+
+        results = []
+        if "query" in search_data and "search" in search_data["query"]:
+            for item in search_data["query"]["search"][:max_results]:
+                page_title = item["title"]
+                snippet = item.get("snippet", "")
+
+                # Clean up the snippet (remove HTML tags)
+                import re
+
+                clean_snippet = re.sub(r"<[^>]+>", "", snippet)
+                clean_snippet = re.sub(r"&[^;]+;", " ", clean_snippet)
+
+                if clean_snippet.strip():
+                    results.append(f"{page_title}: {clean_snippet.strip()}")
+
+        return results if results else [f"No Wikipedia results found for: {query}"]
+
+    except Exception as e:
+        raise Exception(f"Wikipedia API request failed: {e}") from e