opik-optimizer 0.9.1__py3-none-any.whl → 1.0.0__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,22 +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[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
@@ -25,15 +32,15 @@ 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:
         """Helper to calculate improvement percentage string."""
-        initial_s = self.details.get("initial_score")
+        initial_s = self.initial_score
         final_s = self.score
 
         # Check if initial score exists and is a number
@@ -60,7 +67,7 @@ class OptimizationResult(pydantic.BaseModel):
         """Provides a clean, well-formatted plain-text summary."""
         separator = "=" * 80
         rounds_ran = len(self.details.get("rounds", []))
-        initial_score = self.details.get("initial_score")
+        initial_score = self.initial_score
         initial_score_str = (
             f"{initial_score:.4f}" if isinstance(initial_score, (int, float)) else "N/A"
         )
@@ -74,7 +81,6 @@ class OptimizationResult(pydantic.BaseModel):
             .replace("[dim]", "")
             .replace("[/dim]", "")
         )
-        stopped_early = self.details.get("stopped_early", "N/A")
 
         model_name = self.details.get("model", "N/A")
         temp = self.details.get("temperature")
@@ -101,7 +107,6 @@ class OptimizationResult(pydantic.BaseModel):
             f"Final Best Score: {final_score_str}",
             f"Total Improvement:{improvement_str.rjust(max(0, 18 - len('Total Improvement:')))}",
             f"Rounds Completed: {rounds_ran}",
-            f"Stopped Early:    {stopped_early}",
             "\nFINAL OPTIMIZED PROMPT / STRUCTURE:",
             "--------------------------------------------------------------------------------",
             f"{final_prompt_display}",
@@ -114,18 +119,15 @@ class OptimizationResult(pydantic.BaseModel):
         """Provides a rich, formatted output for terminals supporting Rich."""
         improvement_str = self._calculate_improvement_str()
         rounds_ran = len(self.details.get("rounds", []))
-        initial_score = self.details.get("initial_score")
+        initial_score = self.initial_score
         initial_score_str = (
             f"{initial_score:.4f}"
             if isinstance(initial_score, (int, float))
             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]")
-        temp = self.details.get("temperature")
-        temp_str = f"{temp:.1f}" if isinstance(temp, (int, float)) else "[dim]N/A[/dim]"
 
         table = rich.table.Table.grid(padding=(0, 1))
         table.add_column(style="dim")
@@ -135,13 +137,21 @@ class OptimizationResult(pydantic.BaseModel):
             "Optimizer:",
             f"[bold]{self.optimizer}[/bold]",
         )
-        table.add_row("Model Used:", f"{model_name} ([dim]Temp:[/dim] {temp_str})")
+        table.add_row("Model Used:", f"{model_name}")
         table.add_row("Metric Evaluated:", f"[bold]{self.metric_name}[/bold]")
         table.add_row("Initial Score:", initial_score_str)
         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("Stopped Early:", str(stopped_early))
+        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]"
@@ -168,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)