opik-optimizer 0.9.2__py3-none-any.whl → 1.0.1__py3-none-any.whl

opik_optimizer/optimization_config/chat_prompt.py

@@ -1,106 +1,176 @@
-from typing import Any, Dict, List, Literal, Optional
+from typing import Any, Dict, List, Optional, Union, Callable
+
+import copy
 
 from pydantic import BaseModel, Field
 
+from opik import track
+
 
 class Tool(BaseModel):
-    name: str =Field(
-        ...,
-        description="Name of the tool"
-    )
-    description: str = Field(
-        ...,
-        description="Description of the tool"
-    )
+    name: str = Field(..., description="Name of the tool")
+    description: str = Field(..., description="Description of the tool")
     parameters: Dict[str, Any] = Field(
-        ...,
-        description="JSON Schema defining the input parameters for the tool"
+        ..., description="JSON Schema defining the input parameters for the tool"
     )
 
+
 class ChatPrompt:
-    system: str
-    prompt: str
-    messages: List[Dict[Literal["role", "content"], str]]
+    """
+    The ChatPrompt lies at the core of Opik Optimizer. It is
+    either a series of messages, or a system and/or prompt.
+
+    The ChatPrompt must make reference to at least one field
+    in the associated database when used with optimizations.
+
+    Args:
+        system: the system prompt
+        prompt: contains {input-dataset-field}, if given
+        messages: a list of dictionaries with role/content, with
+            a content containing {input-dataset-field}
+    """
 
     def __init__(
         self,
+        name: str = "chat-prompt",
         system: Optional[str] = None,
-        prompt: Optional[str] = None,
-        messages: Optional[List[Dict[Literal["role", "content"], str]]] = None,
-        tools: Optional[List[Tool]] = 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",
+        **model_kwargs: Any,
+    ) -> None:
+        if system is None and user is None and messages is None:
+            raise ValueError(
+                "At least one of `system`, `user`, or `messages` must be provided"
+            )
+
+        if user is not None and messages is not None:
+            raise ValueError("`user` and `messages` cannot be provided together")
+
+        if system is not None and messages is not None:
+            raise ValueError("`system` and `messages` cannot be provided together")
+
+        if system is not None and not isinstance(system, str):
+            raise ValueError("`system` must be a string")
+
+        if user is not None and not isinstance(user, str):
+            raise ValueError("`user` must be a string")
+
+        if messages is not None:
+            if not isinstance(messages, list):
+                raise ValueError("`messages` must be a list")
+            else:
+                for message in messages:
+                    if not isinstance(message, dict):
+                        raise ValueError("`messages` must be a dictionary")
+                    elif "role" not in message or "content" not in message:
+                        raise ValueError(
+                            "`message` must have 'role' and 'content' keys."
+                        )
+        self.name = name
         self.system = system
-        self.prompt = prompt
+        self.user = user
         self.messages = messages
+        # ALl of the rest are just for the ChatPrompt LLM
+        # These are used from the prompt as controls:
+        self.tools = tools
+        if function_map:
+            self.function_map = {
+                key: (
+                    value
+                    if hasattr(value, "__wrapped__")
+                    else track(type="tool")(value)
+                )
+                for key, value in function_map.items()
+            }
+        else:
+            self.function_map = {}
+        # These are used for the LiteLLMAgent class:
+        self.model = model
+        self.model_kwargs = model_kwargs
+        self.invoke = invoke
+        self.project_name = project_name
 
-        self.formatted_messages = self._standardize_prompts()
-
-    def _standardize_prompts(
-        self, **kwargs: Any
-    ) -> List[Dict[Literal["role", "content"], str]]:
-        if (self.system is None and self.prompt is None and self.messages is None):
-            raise ValueError(
-                "At least one of `system`, `prompt` or `messages` must be provided"
-            )
+    def get_messages(
+        self,
+        dataset_item: Optional[Dict[str, str]] = None,
+    ) -> List[Dict[str, str]]:
+        # This is a copy, so we can alter the messages:
+        messages = self._standardize_prompts()
 
-        if (self.prompt is not None and self.messages is not None):
-            raise ValueError(
-                "`prompt` and `messages` cannot be provided together"
-            )
-        
-        if (self.system is not None and not isinstance(self.system, str)):
-            raise ValueError(
-                "`system` must be a string"
-            )
-        
-        if (self.prompt is not None and not isinstance(self.prompt, str)):
-            raise ValueError(
-                "`prompt` must be a string"
-            )
+        if dataset_item:
+            for key, value in dataset_item.items():
+                for message in messages:
+                    # Only replace user message content:
+                    label = "{" + key + "}"
+                    if label in message["content"]:
+                        message["content"] = message["content"].replace(
+                            label, str(value)
+                        )
+        return messages
 
-        if (self.messages is not None and not isinstance(self.messages, list)):
-            raise ValueError(
-                "`messages` must be a list"
-            )
+    def _standardize_prompts(self, **kwargs: Any) -> List[Dict[str, str]]:
+        standardize_messages: List[Dict[str, str]] = []
 
-        standardize_messages = []
-            
-        if (self.system is not None):
+        if self.system is not None:
             standardize_messages.append({"role": "system", "content": self.system})
-        
-        if (self.prompt is not None):
-            standardize_messages.append({"role": "user", "content": self.prompt})
-        
-        if (self.messages is not None):
+
+        if self.messages is not None:
             for message in self.messages:
                 standardize_messages.append(message)
-        
-        return standardize_messages
 
-    def format(self, **kwargs: Any) -> str:
-        return self.prompt.format(**kwargs)
+        if self.user is not None:
+            standardize_messages.append({"role": "user", "content": self.user})
 
-    def to_dict(self) -> Dict[str, Any]:
+        return copy.deepcopy(standardize_messages)
+
+    def to_dict(self) -> Dict[str, Union[str, List[Dict[str, str]]]]:
         """Convert ChatPrompt to a dictionary for JSON serialization.
-        
+
         Returns:
             Dict containing the serializable representation of this ChatPrompt
         """
-        return {
-            "system": self.system,
-            "prompt": self.prompt,
-            "messages": self.messages,
-            "formatted_messages": self.formatted_messages
-        }
+        retval: Dict[str, Union[str, List[Dict[str, str]]]] = {}
+        if self.system is not None:
+            retval["system"] = self.system
+        if self.user is not None:
+            retval["user"] = self.user
+        if self.messages is not None:
+            retval["messages"] = self.messages
+        return retval
+
+    def copy(self) -> "ChatPrompt":
+        return ChatPrompt(
+            system=self.system,
+            user=self.user,
+            messages=copy.deepcopy(self.messages),
+            tools=self.tools,
+            function_map=self.function_map,
+        )
+
+    def set_messages(self, messages: List[Dict[str, Any]]) -> None:
+        self.system = None
+        self.user = None
+        self.messages = copy.deepcopy(messages)
 
     @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) -> 'ChatPrompt':
+    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,
+    ) -> "ChatPrompt":
         """Custom validation method to handle nested objects during deserialization."""
         return ChatPrompt(
-            system=obj.get('system', None),
-            prompt=obj.get('prompt', None),
-            messages=obj.get('messages', None),
-            
+            system=obj.get("system", None),
+            prompt=obj.get("prompt", None),
+            messages=obj.get("messages", None),
         )
-        

opik_optimizer/optimization_config/configs.py

@@ -1,15 +1,16 @@
 """Module containing configuration classes for optimization."""
 
-from typing import Any, Dict, List, Literal, Union
+from typing import Any, List
 
 import pydantic
 
 
 class TaskConfig(pydantic.BaseModel):
     """Configuration for a prompt task."""
+
     model_config = pydantic.ConfigDict(arbitrary_types_allowed=True)
-    
-    instruction_prompt: Union[str, List[Dict[Literal["role", "content"], str]]]
+
+    instruction_prompt: str
     use_chat_prompt: bool = False
     input_dataset_fields: List[str]
     output_dataset_field: str

opik_optimizer/optimization_config/mappers.py

@@ -2,16 +2,21 @@ from typing import Dict, Callable, Optional, Any, Union
 
 EVALUATED_LLM_TASK_OUTPUT = "llm_output"
 
+
 class Mapper:
     """Base class for mapping functions that transform data between different formats."""
-    
-    def __init__(self, name: Optional[str] = None, transform: Optional[Callable[[Any], Any]] = None):
+
+    def __init__(
+        self,
+        name: Optional[str] = None,
+        transform: Optional[Callable[[Any], Any]] = None,
+    ):
         if name is not None and transform is not None:
             raise ValueError("Only one of name or transform can be provided")
-        
+
         self.name = name
         self.transform = transform
-    
+
     def __call__(self, data: Any) -> Any:
         if self.transform is not None:
             return self.transform(data)
@@ -19,7 +24,12 @@ class Mapper:
             return data[self.name]
         return data
 
-def from_dataset_field(*, name: str = None, transform: Optional[Callable[[Dict[str, Any]], Any]] = None) -> Union[str, Callable[[Dict[str, Any]], Any]]:
+
+def from_dataset_field(
+    *,
+    name: Optional[str] = None,
+    transform: Optional[Callable[[Dict[str, Any]], Any]] = None,
+) -> Union[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")
 
@@ -36,7 +46,9 @@ def from_llm_response_text() -> str:
     return EVALUATED_LLM_TASK_OUTPUT
 
 
-def from_agent_output(*, name: str = None, transform: Optional[Callable[[Any], Any]] = None) -> Union[str, Callable[[Any], Any]]:
+def from_agent_output(
+    *, name: Optional[str] = None, transform: Optional[Callable[[Any], Any]] = None
+) -> Union[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,26 +1,29 @@
 """Module containing the OptimizationResult class."""
 
-from typing import Any, Dict, List, Literal, Optional
+from typing import Any, Dict, List, Optional
 
 import pydantic
 import rich
 
-from .reporting_utils import get_console
+from .reporting_utils import get_console, get_link_text
 
 
 class OptimizationResult(pydantic.BaseModel):
     """Result oan optimization run."""
 
     optimizer: str = "Optimizer"
-    
-    prompt: List[Dict[Literal["role", "content"], str]]
+
+    prompt: List[Dict[str, str]]
     score: float
     metric_name: str
-    
+
+    optimization_id: Optional[str] = None
+    dataset_id: Optional[str] = None
+
     # Initial score
-    initial_prompt: Optional[List[Dict[Literal["role", "content"], str]]] = None
+    initial_prompt: Optional[List[Dict[str, str]]] = None
     initial_score: Optional[float] = None
-    
+
     details: Dict[str, Any] = pydantic.Field(default_factory=dict)
     history: List[Dict[str, Any]] = []
     llm_calls: Optional[int] = None
@@ -29,10 +32,10 @@ class OptimizationResult(pydantic.BaseModel):
     demonstrations: Optional[List[Dict[str, Any]]] = None
     mipro_prompt: Optional[str] = None
     tool_prompts: Optional[Dict[str, str]] = None
-    
+
     model_config = pydantic.ConfigDict(arbitrary_types_allowed=True)
 
-    def model_dump(self, *kargs, **kwargs) -> 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:
@@ -123,7 +126,6 @@ class OptimizationResult(pydantic.BaseModel):
             else "[dim]N/A[/dim]"
         )
         final_score_str = f"{self.score:.4f}"
-        stopped_early = self.details.get("stopped_early", "N/A")
 
         model_name = self.details.get("model", "[dim]N/A[/dim]")
 
@@ -141,6 +143,15 @@ class OptimizationResult(pydantic.BaseModel):
         table.add_row("Final Best Score:", f"[bold cyan]{final_score_str}[/bold cyan]")
         table.add_row("Total Improvement:", improvement_str)
         table.add_row("Rounds Completed:", str(rounds_ran))
+        table.add_row(
+            "Optimization run link:",
+            get_link_text(
+                pre_text="",
+                link_text="Open in Opik Dashboard",
+                dataset_id=self.dataset_id,
+                optimization_id=self.optimization_id,
+            ),
+        )
 
         # Display Chat Structure if available
         panel_title = "[bold]Final Optimized Prompt[/bold]"
@@ -167,9 +178,7 @@ class OptimizationResult(pydantic.BaseModel):
         except Exception:
             # Fallback to simple text prompt
             prompt_renderable = rich.text.Text(str(self.prompt or ""), overflow="fold")
-            panel_title = (
-                "[bold]Final Optimized Prompt (Instruction - fallback)[/bold]"
-            )
+            panel_title = "[bold]Final Optimized Prompt (Instruction - fallback)[/bold]"
 
         prompt_panel = rich.panel.Panel(
             prompt_renderable, title=panel_title, border_style="blue", padding=(1, 2)

opik_optimizer/reporting_utils.py

@@ -1,6 +1,6 @@
 import logging
 from contextlib import contextmanager
-from typing import Dict, List, Optional
+from typing import Any, Dict, List, Optional, Union
 
 from rich import box
 from rich.console import Console, Group
@@ -12,30 +12,29 @@ from .utils import get_optimization_run_url_by_id
 
 PANEL_WIDTH = 70
 
-def get_console(*args, **kwargs):
+
+def get_console(*args: Any, **kwargs: Any) -> Console:
     console = Console(*args, **kwargs)
     console.is_jupyter = False
     return console
 
+
 @contextmanager
-def convert_tqdm_to_rich(description: Optional[str] = None, verbose: int = 1):
+def convert_tqdm_to_rich(description: Optional[str] = None, verbose: int = 1) -> Any:
     """Context manager to convert tqdm to rich."""
     import opik.evaluation.engine.evaluation_tasks_executor
-    
-    def _tqdm_to_track(iterable, desc, disable, total):
+
+    def _tqdm_to_track(iterable: Any, desc: str, disable: bool, total: int) -> Any:
         disable = verbose == 0
         return track(
-            iterable,
-            description=description or desc,
-            disable=disable,
-            total=total
+            iterable, description=description or desc, disable=disable, total=total
         )
 
     original__tqdm = opik.evaluation.engine.evaluation_tasks_executor._tqdm
     opik.evaluation.engine.evaluation_tasks_executor._tqdm = _tqdm_to_track
 
-
     from opik.evaluation import report
+
     report.display_experiment_results = lambda *args, **kwargs: None
     report.display_experiment_link = lambda *args, **kwargs: None
 
@@ -45,32 +44,32 @@ def convert_tqdm_to_rich(description: Optional[str] = None, verbose: int = 1):
         opik.evaluation.engine.evaluation_tasks_executor._tqdm = original__tqdm
 
 
-
 @contextmanager
-def suppress_opik_logs():
+def suppress_opik_logs() -> Any:
     """Suppress Opik startup logs by temporarily increasing the log level."""
     # Optimizer log level
-    optimizer_logger = logging.getLogger('opik_optimizer')
-    
+    optimizer_logger = logging.getLogger("opik_optimizer")
+
     # Get the Opik logger
     opik_logger = logging.getLogger("opik.api_objects.opik_client")
-    
+
     # Store original log level
     original_level = opik_logger.level
-    
+
     # Set log level to ERROR to suppress INFO messages
     opik_logger.setLevel(optimizer_logger.level)
-    
+
     try:
         yield
     finally:
         # Restore original log level
         opik_logger.setLevel(original_level)
 
-def display_messages(messages: List[Dict[str, str]], prefix: str = ""):
+
+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"),
+            Text(msg.get("content", ""), overflow="fold"),
             title=f"{msg.get('role', 'message')}",
             title_align="left",
             border_style="dim",
@@ -90,68 +89,97 @@ def display_messages(messages: List[Dict[str, str]], prefix: str = ""):
         for line in rendered_panel.splitlines():
             console.print(Text(prefix) + Text.from_ansi(line))
 
-def display_header(
-    algorithm: str,
-    optimization_id: Optional[str]=None,
-    dataset_id: Optional[str]=None,
-    verbose: int = 1
-):
-    if verbose < 1:
-        return
 
-    if optimization_id is not None and dataset_id is not None:    
+def get_link_text(
+    pre_text: str,
+    link_text: str,
+    optimization_id: Optional[str] = None,
+    dataset_id: Optional[str] = None,
+) -> Text:
+    if optimization_id is not None and dataset_id is not None:
         optimization_url = get_optimization_run_url_by_id(
-            optimization_id=optimization_id,
-            dataset_id=dataset_id
+            optimization_id=optimization_id, dataset_id=dataset_id
         )
-        
+
         # Create a visually appealing panel with an icon and ensure link doesn't wrap
-        
-        link_text = Text("-> View optimization details in your Opik dashboard")
-        link_text.stylize(f"link {optimization_url}", 28, len(link_text))
+        link_text = Text(pre_text + link_text)
+        link_text.stylize(f"link {optimization_url}", len(pre_text), len(link_text))  # type: ignore
     else:
         link_text = Text("No optimization run link available", style="dim")
 
+    return link_text
+
+
+def display_header(
+    algorithm: str,
+    optimization_id: Optional[str] = None,
+    dataset_id: Optional[str] = None,
+    verbose: int = 1,
+) -> None:
+    if verbose < 1:
+        return
+
+    link_text = get_link_text(
+        pre_text="-> View optimization details ",
+        link_text="in your Opik dashboard",
+        optimization_id=optimization_id,
+        dataset_id=dataset_id,
+    )
+
     content = Text.assemble(
-        ("● ", "green"),  
-        "Running Opik Evaluation - ",
-        (algorithm, "blue"),
-        "\n\n"
+        ("● ", "green"), "Running Opik Evaluation - ", (algorithm, "blue"), "\n\n"
     ).append(link_text)
 
-    
-    panel = Panel(
-        content,
-        box=box.ROUNDED,
-        width=PANEL_WIDTH
-    )
+    panel = Panel(content, box=box.ROUNDED, width=PANEL_WIDTH)
 
     console = get_console()
     console.print(panel)
     console.print("\n")
 
 
-def display_result(initial_score, best_score, best_prompt, verbose: int = 1):
+def display_result(
+    initial_score: float,
+    best_score: float,
+    best_prompt: List[Dict[str, str]],
+    verbose: int = 1,
+) -> None:
     if verbose < 1:
         return
 
     console = get_console()
     console.print(Text("\n> Optimization complete\n"))
-    
+
+    content: Union[Text, Panel] = []
+
     if best_score > initial_score:
         if initial_score == 0:
-            content = [Text(f"Prompt was optimized and improved from {initial_score:.4f} to {best_score:.4f}", style="bold green")]
+            content += [
+                Text(
+                    f"Prompt was optimized and improved from {initial_score:.4f} to {best_score:.4f}",
+                    style="bold green",
+                )
+            ]
         else:
             perc_change = (best_score - initial_score) / initial_score
-            content = [Text(f"Prompt was optimized and improved from {initial_score:.4f} to {best_score:.4f} ({perc_change:.2%})", style="bold green")]
+            content += [
+                Text(
+                    f"Prompt was optimized and improved from {initial_score:.4f} to {best_score:.4f} ({perc_change:.2%})",
+                    style="bold green",
+                )
+            ]
     else:
-        content = [Text(f"Optimization run did not find a better prompt than the initial one.\nScore: {best_score:.4f}", style="dim bold red")]
-    
+        content += [
+            Text(
+                f"Optimization run did not find a better prompt than the initial one.\nScore: {best_score:.4f}",
+                style="dim bold red",
+            )
+        ]
+
     content.append(Text("\nOptimized prompt:"))
     for i, msg in enumerate(best_prompt):
         content.append(
             Panel(
-                Text(msg.get('content', ''), overflow="fold"),
+                Text(msg.get("content", ""), overflow="fold"),
                 title=f"{msg.get('role', 'message')}",
                 title_align="left",
                 border_style="dim",
@@ -167,12 +195,14 @@ def display_result(initial_score, best_score, best_prompt, verbose: int = 1):
             title_align="left",
             border_style="green",
             width=PANEL_WIDTH,
-            padding=(1, 2)
+            padding=(1, 2),
         )
     )
 
 
-def display_configuration(messages: List[Dict[str, str]], optimizer_config: Dict[str, str], verbose: int = 1):
+def display_configuration(
+    messages: List[Dict[str, str]], optimizer_config: Dict[str, Any], verbose: int = 1
+) -> None:
     """Displays the LLM messages and optimizer configuration using Rich panels."""
 
     if verbose < 1:
@@ -185,15 +215,16 @@ def display_configuration(messages: List[Dict[str, str]], optimizer_config: Dict
     display_messages(messages)
 
     # Panel for configuration
-    console.print(Text(f"\nUsing {optimizer_config['optimizer']} with the parameters: "))
-    
+    console.print(
+        Text(f"\nUsing {optimizer_config['optimizer']} with the parameters: ")
+    )
+
     for key, value in optimizer_config.items():
         if key == "optimizer":  # Already displayed in the introductory text
             continue
         parameter_text = Text.assemble(
-            Text(f"  - {key}: ", style="dim"), 
-            Text(str(value), style="cyan")      
+            Text(f"  - {key}: ", style="dim"), Text(str(value), style="cyan")
         )
         console.print(parameter_text)
-    
+
     console.print("\n")