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

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

@@ -2,18 +2,14 @@
 
 from typing import (
     Any,
-    Dict,
     Final,
     Literal,
-    Optional,
-    Type,
     TYPE_CHECKING,
-    List,
-    Callable,
 )
+from collections.abc import Callable
 
+import ast
 import inspect
-import typing
 import base64
 import json
 import logging
@@ -22,16 +18,20 @@ 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
 
-ALLOWED_URL_CHARACTERS: Final[str] = ":/&?="
-logger = logging.getLogger(__name__)
+from .colbert import ColBERTv2
 
 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
+
+ALLOWED_URL_CHARACTERS: Final[str] = ":/&?="
+logger = logging.getLogger(__name__)
 
 
 class OptimizationContextManager:
@@ -45,8 +45,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 +63,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 +88,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:
@@ -205,7 +205,7 @@ def json_to_dict(json_str: str) -> Any:
 
     try:
         return json.loads(cleaned_json_string)
-    except json.JSONDecodeError:
+    except json.JSONDecodeError as json_error:
         if cleaned_json_string.startswith("```json"):
             cleaned_json_string = cleaned_json_string[7:]
             if cleaned_json_string.endswith("```"):
@@ -217,18 +217,52 @@ def json_to_dict(json_str: str) -> Any:
 
         try:
             return json.loads(cleaned_json_string)
-        except json.JSONDecodeError as e:
-            print(f"Failed to parse JSON string: {json_str}")
-            logger.debug(f"Failed to parse JSON string: {json_str}")
-            raise e
+        except json.JSONDecodeError:
+            try:
+                literal_result = ast.literal_eval(cleaned_json_string)
+            except (ValueError, SyntaxError):
+                logger.debug("Failed to parse JSON string: %s", json_str)
+                raise json_error
+
+            normalized = _convert_literals_to_json_compatible(literal_result)
+
+            try:
+                return json.loads(json.dumps(normalized))
+            except (TypeError, ValueError) as serialization_error:
+                logger.debug(
+                    "Failed to serialise literal-evaluated payload %r: %s",
+                    literal_result,
+                    serialization_error,
+                )
+                raise json_error
+
+
+def _convert_literals_to_json_compatible(value: Any) -> Any:
+    """Convert Python literals to JSON-compatible structures."""
+    if isinstance(value, dict):
+        return {
+            key: _convert_literals_to_json_compatible(val) for key, val in value.items()
+        }
+    if isinstance(value, list):
+        return [_convert_literals_to_json_compatible(item) for item in value]
+    if isinstance(value, tuple):
+        return [_convert_literals_to_json_compatible(item) for item in value]
+    if isinstance(value, set):
+        return [
+            _convert_literals_to_json_compatible(item)
+            for item in sorted(value, key=repr)
+        ]
+    if isinstance(value, (str, int, float, bool)) or value is None:
+        return value
+    return str(value)
 
 
 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 +292,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 +310,17 @@ 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", optimizer_ref: Any = None
+) -> type["OptimizableAgent"]:
     """
     Create a LiteLLMAgent from a chat prompt.
+
+    Args:
+        prompt: The chat prompt to use
+        optimizer_ref: Optional optimizer instance to attach to the agent
     """
-    from .optimizable_agent import OptimizableAgent
+    from opik_optimizer.optimizable_agent import OptimizableAgent
 
     if prompt.invoke is not None:
 
@@ -288,9 +328,10 @@ def create_litellm_agent_class(prompt: "ChatPrompt") -> Type["OptimizableAgent"]
             model = prompt.model
             model_kwargs = prompt.model_kwargs
             project_name = prompt.project_name
+            optimizer = optimizer_ref
 
             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
@@ -302,18 +343,19 @@ def create_litellm_agent_class(prompt: "ChatPrompt") -> Type["OptimizableAgent"]
             model = prompt.model
             model_kwargs = prompt.model_kwargs
             project_name = prompt.project_name
+            optimizer = optimizer_ref
 
     return LiteLLMAgent
 
 
 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 +392,92 @@ 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 | None = 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: (Optional) 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
+    # 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:
+        # 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

opik_optimizer/utils/dataset_utils.py

@@ -0,0 +1,49 @@
+from __future__ import annotations
+
+import hashlib
+import secrets
+import time
+from functools import lru_cache
+from importlib import resources
+from typing import Any
+from collections.abc import Iterable
+
+
+@lru_cache(maxsize=None)
+def dataset_suffix(package: str, filename: str) -> str:
+    """Return a stable checksum-based suffix for a JSONL dataset file."""
+    text = resources.files(package).joinpath(filename).read_text(encoding="utf-8")
+    return hashlib.md5(text.encode("utf-8")).hexdigest()[:8]
+
+
+def generate_uuid7_str() -> str:
+    """Generate a UUIDv7-compatible string, emulating the layout if unavailable."""
+    import uuid
+
+    if hasattr(uuid, "uuid7"):
+        return str(uuid.uuid7())  # type: ignore[attr-defined]
+
+    unix_ts_ms = int(time.time() * 1000) & ((1 << 48) - 1)
+    rand_a = secrets.randbits(12)
+    rand_b = secrets.randbits(62)
+
+    uuid_int = unix_ts_ms << 80
+    uuid_int |= 0x7 << 76  # version 7
+    uuid_int |= rand_a << 64
+    uuid_int |= 0b10 << 62  # RFC4122 variant
+    uuid_int |= rand_b
+
+    return str(uuid.UUID(int=uuid_int))
+
+
+def attach_uuids(records: Iterable[dict[str, Any]]) -> list[dict[str, Any]]:
+    """Copy records and assign a fresh UUIDv7 `id` to each."""
+    payload: list[dict[str, Any]] = []
+    for record in records:
+        rec = dict(record)
+        rec["id"] = generate_uuid7_str()
+        payload.append(rec)
+    return payload
+
+
+__all__ = ["dataset_suffix", "generate_uuid7_str", "attach_uuids"]

opik_optimizer/utils/prompt_segments.py

@@ -0,0 +1,186 @@
+"""Prompt segmentation helpers for targeted prompt updates.
+
+These utilities operate on existing ``ChatPrompt`` instances without
+changing their constructor, allowing callers to identify and update
+specific sections (system message, individual chat messages, or tool
+descriptions) while preserving backwards compatibility for the rest of
+the optimizer stack.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Any
+from collections.abc import Iterable
+
+import copy
+
+from ..optimization_config.chat_prompt import ChatPrompt
+
+
+PROMPT_SEGMENT_PREFIX_TOOL = "tool:"
+PROMPT_SEGMENT_PREFIX_MESSAGE = "message:"
+
+
+@dataclass
+class PromptSegment:
+    """Lightweight view over a prompt component that may be edited."""
+
+    segment_id: str
+    kind: str
+    role: str | None
+    content: str
+    metadata: dict[str, Any]
+
+    def is_tool(self) -> bool:
+        return self.segment_id.startswith(PROMPT_SEGMENT_PREFIX_TOOL)
+
+
+def _normalise_tool(tool: dict[str, Any]) -> dict[str, Any]:
+    """Return tools in the ``{"function": {...}}`` structure for consistency."""
+
+    if "function" in tool:
+        return copy.deepcopy(tool)
+
+    normalised = copy.deepcopy(tool)
+    function_block = {
+        "name": normalised.pop("name", None),
+        "description": normalised.pop("description", ""),
+        "parameters": normalised.pop("parameters", None),
+    }
+    normalised = {"function": function_block, **normalised}
+    return normalised
+
+
+def extract_prompt_segments(prompt: ChatPrompt) -> list[PromptSegment]:
+    """Extract individual editable segments from ``prompt``.
+
+    The extraction preserves order for chat messages while assigning
+    stable segment identifiers:
+
+    * ``system`` for the system field (if present)
+    * ``user`` for the top-level user field (if present)
+    * ``message:<index>`` for entries in ``messages``
+    * ``tool:<name>`` for tool descriptions
+    """
+
+    segments: list[PromptSegment] = []
+
+    if prompt.system is not None:
+        segments.append(
+            PromptSegment(
+                segment_id="system",
+                kind="system",
+                role="system",
+                content=prompt.system,
+                metadata={},
+            )
+        )
+
+    if prompt.messages is not None:
+        for idx, message in enumerate(prompt.messages):
+            segments.append(
+                PromptSegment(
+                    segment_id=f"{PROMPT_SEGMENT_PREFIX_MESSAGE}{idx}",
+                    kind="message",
+                    role=message.get("role"),
+                    content=message.get("content", ""),
+                    metadata={
+                        key: value for key, value in message.items() if key != "content"
+                    },
+                )
+            )
+
+    if prompt.user is not None:
+        segments.append(
+            PromptSegment(
+                segment_id="user",
+                kind="user",
+                role="user",
+                content=prompt.user,
+                metadata={},
+            )
+        )
+
+    if prompt.tools:
+        for tool in prompt.tools:
+            normalised = _normalise_tool(tool)
+            function_block = normalised.get("function", {})
+            tool_name = function_block.get("name")
+            if not tool_name:
+                continue
+            segments.append(
+                PromptSegment(
+                    segment_id=f"{PROMPT_SEGMENT_PREFIX_TOOL}{tool_name}",
+                    kind="tool",
+                    role="tool",
+                    content=function_block.get("description", ""),
+                    metadata={
+                        "parameters": function_block.get("parameters"),
+                        "raw_tool": normalised,
+                    },
+                )
+            )
+
+    return segments
+
+
+def apply_segment_updates(
+    prompt: ChatPrompt,
+    updates: dict[str, str],
+) -> ChatPrompt:
+    """Return a new ``ChatPrompt`` with selected segments replaced.
+
+    ``updates`` maps segment identifiers (as produced by
+    ``extract_prompt_segments``) to replacement strings.
+    """
+
+    system = updates.get("system", prompt.system)
+    user = updates.get("user", prompt.user)
+
+    messages: list[dict[str, Any]] | None = None
+    if prompt.messages is not None:
+        new_messages: list[dict[str, Any]] = []
+        for idx, message in enumerate(prompt.messages):
+            segment_id = f"{PROMPT_SEGMENT_PREFIX_MESSAGE}{idx}"
+            replacement = updates.get(segment_id)
+            if replacement is not None:
+                updated_message = copy.deepcopy(message)
+                updated_message["content"] = replacement
+                new_messages.append(updated_message)
+            else:
+                new_messages.append(copy.deepcopy(message))
+        messages = new_messages
+
+    tools = copy.deepcopy(prompt.tools) if prompt.tools else None
+    if tools:
+        for tool in tools:
+            normalised = _normalise_tool(tool)
+            function_block = normalised.get("function", {})
+            tool_name = function_block.get("name")
+            if not tool_name:
+                continue
+            segment_id = f"{PROMPT_SEGMENT_PREFIX_TOOL}{tool_name}"
+            replacement = updates.get(segment_id)
+            if replacement is not None:
+                function_block["description"] = replacement
+            tool.update(normalised)
+
+    return ChatPrompt(
+        name=prompt.name,
+        system=system,
+        user=user,
+        messages=messages,
+        tools=tools,
+        function_map=prompt.function_map,
+        model=prompt.model,
+        invoke=prompt.invoke,
+        project_name=prompt.project_name,
+        **prompt.model_kwargs,
+    )
+
+
+def segment_ids_for_tools(segments: Iterable[PromptSegment]) -> list[str]:
+    """Convenience helper returning IDs of tool segments."""
+
+    return [segment.segment_id for segment in segments if segment.is_tool()]