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

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,26 @@ 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
+    tool_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 +41,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 +206,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/colbert.py

@@ -0,0 +1,236 @@
+"""
+Minimal ColBERTv2 implementation extracted from dspy (MIT license).
+
+This module provides a lightweight implementation of ColBERTv2 search functionality
+without requiring the full dspy dependency.
+"""
+
+import copy
+import time
+from typing import Any
+import requests  # type: ignore[import-untyped]
+from requests.adapters import HTTPAdapter  # type: ignore[import-untyped]
+from urllib3.util.retry import Retry
+
+
+def _create_session_with_retries(max_retries: int = 4) -> requests.Session:
+    """
+    Create a requests session with retry configuration.
+
+    Args:
+        max_retries: Maximum number of retry attempts
+
+    Returns:
+        Configured requests session
+    """
+    session = requests.Session()
+
+    retry_strategy = Retry(
+        total=max_retries,
+        backoff_factor=1,  # Wait 1, 2, 4, 8 seconds between retries
+        status_forcelist=[429, 500, 502, 503, 504],  # HTTP status codes to retry on
+        allowed_methods=["HEAD", "GET", "POST", "PUT", "DELETE", "OPTIONS", "TRACE"],
+    )
+
+    adapter = HTTPAdapter(max_retries=retry_strategy)
+    session.mount("http://", adapter)
+    session.mount("https://", adapter)
+
+    return session
+
+
+class dotdict(dict):
+    """Dictionary with attribute access (extracted from dspy)."""
+
+    def __getattr__(self, key: str) -> Any:
+        if key.startswith("__") and key.endswith("__"):
+            return super().__getattribute__(key)
+        try:
+            return self[key]
+        except KeyError:
+            raise AttributeError(
+                f"'{type(self).__name__}' object has no attribute '{key}'"
+            )
+
+    def __setattr__(self, key: str, value: Any) -> None:
+        if key.startswith("__") and key.endswith("__"):
+            super().__setattr__(key, value)
+        else:
+            self[key] = value
+
+    def __delattr__(self, key: str) -> None:
+        if key.startswith("__") and key.endswith("__"):
+            super().__delattr__(key)
+        else:
+            del self[key]
+
+    def __deepcopy__(self, memo: dict[Any, Any]) -> "dotdict":
+        # Use the default dict copying method to avoid infinite recursion.
+        return dotdict(copy.deepcopy(dict(self), memo))
+
+
+def colbertv2_get_request(
+    url: str, query: str, k: int, max_retries: int = 4
+) -> list[dict[str, Any]]:
+    """
+    Make a GET request to ColBERTv2 server with retry logic.
+
+    Args:
+        url: The ColBERTv2 server URL
+        query: The search query
+        k: Number of results to return
+        max_retries: Maximum number of retry attempts
+
+    Returns:
+        List of search results
+    """
+    assert k <= 100, (
+        "Only k <= 100 is supported for the hosted ColBERTv2 server at the moment."
+    )
+
+    session = _create_session_with_retries(max_retries)
+    payload: dict[str, str | int] = {"query": query, "k": k}
+
+    # Application-level retry for server connection errors
+    for attempt in range(max_retries):
+        try:
+            res = session.get(url, params=payload, timeout=5)
+            response_data = res.json()
+
+            # Check for application-level errors (server connection issues, etc.)
+            if "error" in response_data and response_data["error"]:
+                error_msg = response_data.get("message", "Unknown error")
+                # If it's a connection error, retry; otherwise, fail immediately
+                if (
+                    "Cannot connect to host" in error_msg
+                    or "Connection refused" in error_msg
+                ):
+                    if attempt == max_retries - 1:
+                        raise Exception(f"ColBERTv2 server error: {error_msg}")
+                    time.sleep(1)  # Wait 1 second before retrying
+                    continue
+                else:
+                    raise Exception(f"ColBERTv2 server error: {error_msg}")
+
+            if "topk" not in response_data:
+                raise Exception(
+                    f"Unexpected response format from ColBERTv2 server: {list(response_data.keys())}"
+                )
+
+            topk = response_data["topk"][:k]
+            topk = [{**d, "long_text": d["text"]} for d in topk]
+            return topk[:k]
+
+        except requests.RequestException as e:
+            if attempt == max_retries - 1:
+                raise Exception(f"ColBERTv2 request failed: {str(e)}")
+            time.sleep(1)  # Wait 1 second before retrying
+
+    # This should never be reached, but mypy requires a return statement
+    raise Exception("Unexpected end of retry loop")
+
+
+def colbertv2_post_request(
+    url: str, query: str, k: int, max_retries: int = 4
+) -> list[dict[str, Any]]:
+    """
+    Make a POST request to ColBERTv2 server with retry logic.
+
+    Args:
+        url: The ColBERTv2 server URL
+        query: The search query
+        k: Number of results to return
+        max_retries: Maximum number of retry attempts
+
+    Returns:
+        List of search results
+    """
+    session = _create_session_with_retries(max_retries)
+    headers = {"Content-Type": "application/json; charset=utf-8"}
+    payload = {"query": query, "k": k}
+
+    # Application-level retry for server connection errors
+    for attempt in range(max_retries):
+        try:
+            res = session.post(url, json=payload, headers=headers, timeout=5)
+            response_data = res.json()
+
+            # Check for application-level errors (server connection issues, etc.)
+            if "error" in response_data and response_data["error"]:
+                error_msg = response_data.get("message", "Unknown error")
+                # If it's a connection error, retry; otherwise, fail immediately
+                if (
+                    "Cannot connect to host" in error_msg
+                    or "Connection refused" in error_msg
+                ):
+                    if attempt == max_retries - 1:
+                        raise Exception(f"ColBERTv2 server error: {error_msg}")
+                    time.sleep(1)  # Wait 1 second before retrying
+                    continue
+                else:
+                    raise Exception(f"ColBERTv2 server error: {error_msg}")
+
+            if "topk" not in response_data:
+                raise Exception(
+                    f"Unexpected response format from ColBERTv2 server: {list(response_data.keys())}"
+                )
+
+            return response_data["topk"][:k]
+
+        except requests.RequestException as e:
+            if attempt == max_retries - 1:
+                raise Exception(f"ColBERTv2 request failed: {str(e)}")
+            time.sleep(1)  # Wait 1 second before retrying
+
+    # This should never be reached, but mypy requires a return statement
+    raise Exception("Unexpected end of retry loop")
+
+
+class ColBERTv2:
+    """Wrapper for the ColBERTv2 Retrieval (extracted from dspy)."""
+
+    def __init__(
+        self,
+        url: str = "http://0.0.0.0",
+        port: str | int | None = None,
+        post_requests: bool = False,
+    ):
+        """
+        Initialize ColBERTv2 client.
+
+        Args:
+            url: Base URL for the ColBERTv2 server
+            port: Optional port number
+            post_requests: Whether to use POST requests instead of GET
+        """
+        self.post_requests = post_requests
+        self.url = f"{url}:{port}" if port else url
+
+    def __call__(
+        self,
+        query: str,
+        k: int = 10,
+        simplify: bool = False,
+        max_retries: int = 4,
+    ) -> list[str] | list[dotdict]:
+        """
+        Search using ColBERTv2.
+
+        Args:
+            query: The search query
+            k: Number of results to return
+            simplify: If True, return only text strings; if False, return dotdict objects
+            max_retries: Maximum number of retry attempts
+
+        Returns:
+            List of search results (either strings or dotdict objects)
+        """
+        if self.post_requests:
+            topk_results = colbertv2_post_request(self.url, query, k, max_retries)
+        else:
+            topk_results = colbertv2_get_request(self.url, query, k, max_retries)
+
+        if simplify:
+            return [psg["long_text"] for psg in topk_results]
+
+        return [dotdict(psg) for psg in topk_results]