opik-optimizer 0.9.1__py3-none-any.whl → 1.0.0__py3-none-any.whl

opik_optimizer/mipro_optimizer/mipro_optimizer.py

@@ -2,6 +2,7 @@ import os
 import random
 from datetime import datetime
 from typing import Callable, Dict, List, Literal, Optional, Union
+import logging
 
 import dspy
 import litellm
@@ -16,7 +17,6 @@ from ..optimization_result import OptimizationResult
 from ..utils import optimization_context
 from ..base_optimizer import BaseOptimizer
 from ..optimization_config.configs import TaskConfig
-from ..optimization_result import OptimizationResult
 from ._lm import LM
 from ._mipro_optimizer_v2 import MIPROv2
 from .utils import (
@@ -30,23 +30,26 @@ from .utils import (
 disk_cache_dir = os.path.expanduser("~/.litellm_cache")
 litellm.cache = Cache(type="disk", disk_cache_dir=disk_cache_dir)
 
-# Set up logging
-import logging
-
 logger = logging.getLogger(__name__)  # Inherits config from setup_logging
 
 
 class MiproOptimizer(BaseOptimizer):
-    def __init__(self, model, project_name: Optional[str] = None, verbose: int = 1, **model_kwargs):
-        super().__init__(model, project_name, verbose=verbose, **model_kwargs)
+    def __init__(
+        self,
+        model,
+        project_name: Optional[str] = None,
+        verbose: int = 1,
+        **model_kwargs,
+    ):
+        super().__init__(model=model, verbose=verbose, **model_kwargs)
         self.tools = []
+        self.project_name = project_name
         self.num_threads = self.model_kwargs.pop("num_threads", 6)
         self.model_kwargs["model"] = self.model
-        self.llm_call_counter = 0
         # FIXME: add mipro_optimizer=True - It does not count the LLM calls made internally by DSPy during MiproOptimizer.optimizer.compile().
-        lm = LM(**self.model_kwargs)
+        self.lm = LM(**self.model_kwargs)
         opik_callback = OpikCallback(project_name=self.project_name, log_graph=True)
-        dspy.configure(lm=lm, callbacks=[opik_callback])
+        dspy.configure(lm=self.lm, callbacks=[opik_callback])
         logger.debug(f"Initialized MiproOptimizer with model: {model}")
 
     def evaluate_prompt(
@@ -54,7 +57,7 @@ class MiproOptimizer(BaseOptimizer):
         dataset: Union[str, Dataset],
         metric: Callable,
         task_config: TaskConfig,
-        prompt: Union[str, dspy.Module, OptimizationResult] = None,
+        prompt: Optional[Union[str, dspy.Module, OptimizationResult]] = None,
         n_samples: int = 10,
         dataset_item_ids: Optional[List[str]] = None,
         experiment_config: Optional[Dict] = None,
@@ -85,7 +88,9 @@ class MiproOptimizer(BaseOptimizer):
         output_key = task_config.output_dataset_field
 
         # Kwargs might contain n_samples, passed from run_benchmark.py
-        n_samples = kwargs.pop("n_samples", None) # Get n_samples from kwargs if present
+        n_samples = kwargs.pop(
+            "n_samples", None
+        )  # Get n_samples from kwargs if present
 
         if isinstance(dataset, str):
             opik_client = opik.Opik(project_name=self.project_name)
@@ -154,28 +159,42 @@ class MiproOptimizer(BaseOptimizer):
         # Robust n_samples handling for selecting dataset_item_ids
         dataset_items_for_eval = dataset.get_items()
         num_total_items = len(dataset_items_for_eval)
-        dataset_item_ids_to_use = dataset_item_ids # Use provided IDs if any
+        dataset_item_ids_to_use = dataset_item_ids  # Use provided IDs if any
 
-        if n_samples is not None: # If n_samples is specified by the caller (run_benchmark.py)
+        if (
+            n_samples is not None
+        ):  # If n_samples is specified by the caller (run_benchmark.py)
             if dataset_item_ids is not None:
                 # This case should ideally be an error or a clear precedence rule.
                 # For now, let's assume if dataset_item_ids is provided, it takes precedence over n_samples.
-                logger.warning("MiproOptimizer.evaluate_prompt: Both n_samples and dataset_item_ids provided. Using provided dataset_item_ids.")
+                logger.warning(
+                    "MiproOptimizer.evaluate_prompt: Both n_samples and dataset_item_ids provided. Using provided dataset_item_ids."
+                )
                 # dataset_item_ids_to_use is already dataset_item_ids
             elif n_samples > num_total_items:
-                logger.warning(f"MiproOptimizer.evaluate_prompt: n_samples ({n_samples}) > total items ({num_total_items}). Using all {num_total_items} items.")
-                dataset_item_ids_to_use = None # opik.evaluation.evaluate handles None as all items
+                logger.warning(
+                    f"MiproOptimizer.evaluate_prompt: n_samples ({n_samples}) > total items ({num_total_items}). Using all {num_total_items} items."
+                )
+                dataset_item_ids_to_use = (
+                    None  # opik.evaluation.evaluate handles None as all items
+                )
             elif n_samples <= 0:
-                logger.warning(f"MiproOptimizer.evaluate_prompt: n_samples ({n_samples}) is <= 0. Using all {num_total_items} items.")
+                logger.warning(
+                    f"MiproOptimizer.evaluate_prompt: n_samples ({n_samples}) is <= 0. Using all {num_total_items} items."
+                )
                 dataset_item_ids_to_use = None
             else:
                 # n_samples is valid and dataset_item_ids was not provided, so sample now.
                 all_ids = [item["id"] for item in dataset_items_for_eval]
                 dataset_item_ids_to_use = random.sample(all_ids, n_samples)
-                logger.info(f"MiproOptimizer.evaluate_prompt: Sampled {n_samples} items for evaluation.")
-        else: # n_samples is None
+                logger.info(
+                    f"MiproOptimizer.evaluate_prompt: Sampled {n_samples} items for evaluation."
+                )
+        else:  # n_samples is None
             if dataset_item_ids is None:
-                logger.info(f"MiproOptimizer.evaluate_prompt: n_samples is None and dataset_item_ids is None. Using all {num_total_items} items.")
+                logger.info(
+                    f"MiproOptimizer.evaluate_prompt: n_samples is None and dataset_item_ids is None. Using all {num_total_items} items."
+                )
             # dataset_item_ids_to_use is already dataset_item_ids (which could be None)
 
         experiment_config = experiment_config or {}
@@ -231,10 +250,10 @@ class MiproOptimizer(BaseOptimizer):
     ) -> OptimizationResult:
         self._opik_client = opik.Opik()
         with optimization_context(
-                client=self._opik_client,
-                dataset_name=dataset.name,
-                objective_name=metric.__name__,
-                metadata={"optimizer": self.__class__.__name__},
+            client=self._opik_client,
+            dataset_name=dataset.name,
+            objective_name=metric.__name__,
+            metadata={"optimizer": self.__class__.__name__},
         ) as optimization:
             result = self._optimize_prompt(
                 dataset=dataset,
@@ -323,7 +342,9 @@ class MiproOptimizer(BaseOptimizer):
             if self.output_key not in row:
                 raise Exception("row does not contain output_key: %r" % self.output_key)
 
-        self.trainset = create_dspy_training_set(self.dataset, self.input_key, self.n_samples)
+        self.trainset = create_dspy_training_set(
+            self.dataset, self.input_key, self.n_samples
+        )
         self.data_signature = create_dspy_signature(
             self.input_key, self.output_key, self.prompt
         )
@@ -384,8 +405,10 @@ class MiproOptimizer(BaseOptimizer):
         """
         Continue to look for optimizations
         """
-        if not hasattr(self, 'optimizer') or not self.optimizer:
-            raise RuntimeError("MiproOptimizer not prepared. Call prepare_optimize_prompt first.")
+        if not hasattr(self, "optimizer") or not self.optimizer:
+            raise RuntimeError(
+                "MiproOptimizer not prepared. Call prepare_optimize_prompt first."
+            )
 
         self.results = self.optimizer.compile(
             student=self.module,
@@ -404,16 +427,30 @@ class MiproOptimizer(BaseOptimizer):
         # self.num_candidates is set in prepare_optimize_prompt, defaults to 10
         # If self.num_candidates is 0 or None, this logic might break or be odd.
         # Add a safeguard for num_candidates_per_round if self.num_candidates is not usable.
-        num_candidates_per_round = self.num_candidates if hasattr(self, 'num_candidates') and self.num_candidates and self.num_candidates > 0 else 1
+        num_candidates_per_round = (  # noqa
+            self.num_candidates
+            if hasattr(self, "num_candidates")
+            and self.num_candidates
+            and self.num_candidates > 0
+            else 1
+        )
 
         for i, candidate_data in enumerate(self.results.candidate_programs):
             program_module = candidate_data.get("program")
             instruction = "N/A"
-            if hasattr(program_module, 'signature') and hasattr(program_module.signature, 'instructions'):
+            if hasattr(program_module, "signature") and hasattr(
+                program_module.signature, "instructions"
+            ):
                 instruction = program_module.signature.instructions
-            elif hasattr(program_module, 'extended_signature') and hasattr(program_module.extended_signature, 'instructions'):
+            elif hasattr(program_module, "extended_signature") and hasattr(
+                program_module.extended_signature, "instructions"
+            ):
                 instruction = program_module.extended_signature.instructions
-            elif hasattr(program_module, 'predictor') and hasattr(program_module.predictor, 'signature') and hasattr(program_module.predictor.signature, 'instructions'):
+            elif (
+                hasattr(program_module, "predictor")
+                and hasattr(program_module.predictor, "signature")
+                and hasattr(program_module.predictor.signature, "instructions")
+            ):
                 instruction = program_module.predictor.signature.instructions
 
             # Remove R and C calculation for Mipro as its history is flat
@@ -426,13 +463,11 @@ class MiproOptimizer(BaseOptimizer):
                 # "candidate_in_round": current_candidate_in_round, # Remove candidate_in_round
                 "timestamp": datetime.now().isoformat(),
                 "prompt_candidate": instruction,
-                "parameters_used": {
-                    "program_summary": str(program_module)[:500]
-                },
-                "scores": [], # Initialize scores list
-                "tokens_used": None, # TODO: add tokens_used
-                "cost": None, # TODO: add cost
-                "duration_seconds": None, # TODO: add duration_seconds
+                "parameters_used": {"program_summary": str(program_module)[:500]},
+                "scores": [],  # Initialize scores list
+                "tokens_used": None,  # TODO: add tokens_used
+                "cost": None,  # TODO: add cost
+                "duration_seconds": None,  # TODO: add duration_seconds
             }
 
             current_score = candidate_data.get("score")
@@ -440,70 +475,103 @@ class MiproOptimizer(BaseOptimizer):
 
             # Unscale if it's a known 0-1 metric that MIPRO might scale to 0-100
             # For now, specifically targeting Levenshtein-like metrics
-            if isinstance(current_score, (float, int)) and \
-               ("levenshtein" in metric_name_for_history.lower() or "similarity" in metric_name_for_history.lower()):
+            if isinstance(current_score, (float, int)) and (
+                "levenshtein" in metric_name_for_history.lower()
+                or "similarity" in metric_name_for_history.lower()
+            ):
                 # Assuming scores like 32.4 are 0-1 scores scaled by 100
-                if abs(current_score) > 1.0: # A simple check to see if it looks scaled
-                    logger.debug(f"Mipro history: Unscaling score {current_score} for metric {metric_name_for_history} by dividing by 100.")
+                if abs(current_score) > 1.0:  # A simple check to see if it looks scaled
+                    logger.debug(
+                        f"Mipro history: Unscaling score {current_score} for metric {metric_name_for_history} by dividing by 100."
+                    )
                     current_score /= 100.0
-            
-            iter_detail["scores"].append({
-                "metric_name": metric_name_for_history,
-                "score": current_score,
-                "opik_evaluation_id": None # TODO: add opik_evaluation_id
-            })
+
+            iter_detail["scores"].append(
+                {
+                    "metric_name": metric_name_for_history,
+                    "score": current_score,
+                    "opik_evaluation_id": None,  # TODO: add opik_evaluation_id
+                }
+            )
             mipro_history_processed.append(iter_detail)
 
         if not self.best_programs:
             logger.warning("MIPRO compile returned no candidate programs.")
             return OptimizationResult(
                 optimizer="MiproOptimizer",
-                prompt=[{"role": "user", "content": getattr(self, 'prompt', "Error: Initial prompt not found")}], 
+                prompt=[
+                    {
+                        "role": "user",
+                        "content": getattr(
+                            self, "prompt", "Error: Initial prompt not found"
+                        ),
+                    }
+                ],
                 score=0.0,
-                metric_name=self.opik_metric.__name__ if hasattr(self, 'opik_metric') else "unknown_metric",
+                metric_name=self.opik_metric.__name__
+                if hasattr(self, "opik_metric")
+                else "unknown_metric",
                 details={"error": "No candidate programs generated by MIPRO"},
                 history=mipro_history_processed,
-                llm_calls=self.llm_call_counter
+                llm_calls=self.lm.llm_call_counter,
             )
 
         self.module = self.get_best().details["program"]
         best_program_details = self.get_best()
-        
+
         # Unscale the main score if necessary, similar to history scores
         final_best_score = best_program_details.score
         final_metric_name = best_program_details.metric_name
-        if isinstance(final_best_score, (float, int)) and \
-           final_metric_name and \
-           ("levenshtein" in final_metric_name.lower() or "similarity" in final_metric_name.lower()):
-            if abs(final_best_score) > 1.0: # A simple check to see if it looks scaled
-                logger.debug(f"Mipro main result: Unscaling score {final_best_score} for metric {final_metric_name} by dividing by 100.")
+        if (
+            isinstance(final_best_score, (float, int))
+            and final_metric_name
+            and (
+                "levenshtein" in final_metric_name.lower()
+                or "similarity" in final_metric_name.lower()
+            )
+        ):
+            if abs(final_best_score) > 1.0:  # A simple check to see if it looks scaled
+                logger.debug(
+                    f"Mipro main result: Unscaling score {final_best_score} for metric {final_metric_name} by dividing by 100."
+                )
                 final_best_score /= 100.0
 
         return OptimizationResult(
             optimizer="MiproOptimizer",
             prompt=best_program_details.prompt,
             tool_prompts=best_program_details.tool_prompts,
-            score=final_best_score, # Use the potentially unscaled score
+            score=final_best_score,  # Use the potentially unscaled score
             metric_name=final_metric_name,
             demonstrations=best_program_details.demonstrations,
             details=best_program_details.details,
             history=mipro_history_processed,
-            llm_calls=self.llm_call_counter
+            llm_calls=self.lm.llm_call_counter,
         )
 
     def get_best(self, position: int = 0) -> OptimizationResult:
-        if not hasattr(self, 'best_programs') or not self.best_programs:
-            logger.error("get_best() called but no best_programs found. MIPRO compile might have failed or yielded no results.")
+        if not hasattr(self, "best_programs") or not self.best_programs:
+            logger.error(
+                "get_best() called but no best_programs found. MIPRO compile might have failed or yielded no results."
+            )
             return OptimizationResult(
                 optimizer="MiproOptimizer",
-                prompt=[{"role": "user", "content": getattr(self, 'prompt', "Error: Initial prompt not found")}], 
-                score=0.0, 
-                metric_name=getattr(self, 'opik_metric', None).name if hasattr(self, 'opik_metric') and self.opik_metric else "unknown_metric",
+                prompt=[
+                    {
+                        "role": "user",
+                        "content": getattr(
+                            self, "prompt", "Error: Initial prompt not found"
+                        ),
+                    }
+                ],
+                score=0.0,
+                metric_name=getattr(self, "opik_metric", None).name
+                if hasattr(self, "opik_metric") and self.opik_metric
+                else "unknown_metric",
                 details={"error": "No programs generated or compile failed"},
                 history=[],
-                llm_calls=self.llm_call_counter
+                llm_calls=self.lm.llm_call_counter,
             )
-            
+
         score = self.best_programs[position]["score"]
         program_module = self.best_programs[position]["program"]
         state = program_module.dump_state()
@@ -528,5 +596,5 @@ class MiproOptimizer(BaseOptimizer):
             metric_name=self.opik_metric.__name__,
             demonstrations=demos,
             details={"program": program_module},
-            llm_calls=self.llm_call_counter
+            llm_calls=self.lm.llm_call_counter,
         )

opik_optimizer/mipro_optimizer/utils.py

@@ -1,4 +1,4 @@
-from typing import Any, Dict, List, Tuple, Union, Optional
+from typing import Dict, Optional
 
 import uuid
 import dspy
@@ -46,7 +46,10 @@ def opik_metric_to_dspy(metric, output):
     def opik_metric_score_wrapper(example, prediction, trace=None):
         try:
             # Calculate the score using the metric
-            score_result = metric(dataset_item=example.toDict(), llm_output=getattr(prediction, answer_field, ""))
+            score_result = metric(
+                dataset_item=example.toDict(),
+                llm_output=getattr(prediction, answer_field, ""),
+            )
             return (
                 score_result.value if hasattr(score_result, "value") else score_result
             )

opik_optimizer/optimizable_agent.py

@@ -0,0 +1,179 @@
+from typing import Dict, Any, List, Optional, TYPE_CHECKING
+import json
+import os
+
+
+from opik.opik_context import get_current_span_data
+
+import litellm
+from litellm.integrations.opik.opik import OpikLogger
+
+from . import _throttle
+
+_limiter = _throttle.get_rate_limiter_for_current_opik_installation()
+
+if TYPE_CHECKING:
+    from .optimization_config.chat_prompt import ChatPrompt
+
+
+def tools_to_dict(tools: Dict[str, Dict[str, Any]]) -> Dict[str, Any]:
+    retval = {}
+    for name in tools:
+        parts = {}
+        for part in tools[name]:
+            if isinstance(tools[name][part], (int, float, str)):
+                parts[part] = tools[name][part]
+        if parts:
+            retval[name] = parts
+    return retval
+
+
+class OptimizableAgent:
+    """
+    An agent class to subclass to make an Optimizable Agent.
+
+    Attributes:
+        model (Optional[str]): The model to use for the agent
+        model_kwargs (Dict[str, Any]): Additional keyword arguments for the model
+        project_name (Optional[str]): The project name for tracking
+    """
+
+    model: Optional[str] = None
+    model_kwargs: Dict[str, Any] = {}
+    project_name: Optional[str] = "Default Project"
+    input_dataset_field: Optional[str] = None
+    prompts: Dict[str, "ChatPrompt"]
+    prompt: "ChatPrompt"
+
+    def __init__(self, prompt: "ChatPrompt") -> None:
+        """
+        Initialize the OptimizableAgent.
+
+        Args:
+            prompt: a chat prompt
+        """
+        self.init_llm()
+        self.init_agent(prompt)
+
+    def init_llm(self) -> None:
+        """Initialize the LLM with the appropriate callbacks."""
+        # Litellm bug requires this (maybe problematic if multi-threaded)
+        os.environ["OPIK_PROJECT_NAME"] = str(self.project_name)
+        self.opik_logger = OpikLogger()
+        litellm.callbacks = [self.opik_logger]
+
+    def init_agent(self, prompt: "ChatPrompt") -> None:
+        """Initialize the agent with the provided configuration."""
+        # Register the tools, if any, for default LiteLLM Agent use:
+        self.prompt = prompt
+
+    @_throttle.rate_limited(_limiter)
+    def _llm_complete(
+        self,
+        messages: List[Dict[str, str]],
+        tools: Optional[List[Dict[str, str]]],
+        seed: int,
+    ) -> Any:
+        response = litellm.completion(
+            model=self.model,
+            messages=messages,
+            seed=seed,
+            tools=tools,
+            metadata={
+                "opik": {
+                    "current_span_data": get_current_span_data(),
+                },
+            },
+            **self.model_kwargs,
+        )
+        return response
+
+    def llm_invoke(
+        self,
+        query: Optional[str] = None,
+        messages: Optional[List[Dict[str, str]]] = None,
+        seed: Optional[int] = None,
+        allow_tool_use: Optional[bool] = False,
+    ) -> str:
+        """
+        NOTE: this is the default LiteLLM API. It is used
+        internally for the LiteLLM Agent.
+
+        Invoke the LLM with the provided query or messages.
+
+        Args:
+            query (Optional[str]): The query to send to the LLM
+            messages (Optional[List[Dict[str, str]]]): Messages to send to the LLM
+            seed (Optional[int]): Seed for reproducibility
+            allow_tool_use: If True, allow LLM to use tools
+
+        Returns:
+            str: The LLM's response
+        """
+        all_messages = []
+        if messages is not None:
+            all_messages.extend(messages)
+
+        if query is not None:
+            all_messages.append({"role": "user", "content": query})
+
+        if allow_tool_use and self.prompt.tools:
+            # Tool-calling loop
+            final_response = "I was unable to find the desired information."
+            count = 0
+            while count < 20:
+                count += 1
+                response = self._llm_complete(all_messages, self.prompt.tools, seed)
+                msg = response.choices[0].message
+                all_messages.append(msg.to_dict())
+                if msg.tool_calls:
+                    for tool_call in msg["tool_calls"]:
+                        tool_name = tool_call["function"]["name"]
+                        arguments = json.loads(tool_call["function"]["arguments"])
+                        tool_func = self.prompt.function_map.get(tool_name)
+                        try:
+                            tool_result = (
+                                tool_func(**arguments)
+                                if tool_func is not None
+                                else "Unknown tool"
+                            )
+                        except Exception:
+                            tool_result = f"Error in calling tool `{tool_name}`"
+                        all_messages.append(
+                            {
+                                "role": "tool",
+                                "tool_call_id": tool_call["id"],
+                                "content": str(tool_result),
+                            }
+                        )
+                else:
+                    final_response = msg["content"]
+                    break
+            result = final_response
+        else:
+            response = self._llm_complete(all_messages, None, seed)
+            result = response.choices[0].message.content
+        return result
+
+    def invoke_dataset_item(self, dataset_item: Dict[str, str]) -> str:
+        messages = self.prompt.get_messages(dataset_item)
+        return self.invoke(messages)
+
+    def invoke(
+        self,
+        messages: List[Dict[str, str]],
+        seed: Optional[int] = None,
+    ) -> str:
+        """
+        Invoke the agent with a dataset item.
+
+        Args:
+            dataset_item (Dict[str, Any]): The dataset item to process
+            seed (Optional[int]): Seed for reproducibility
+
+        Returns:
+            Dict[str, Any]: The agent's response
+        """
+        # Replace with agent invocation:
+        result = self.llm_invoke(messages=messages, seed=seed, allow_tool_use=True)
+        return result