Agentic Optimizer

[dsblank]

Details

Previously, we could only optimize prompts for a generic LiteLLM model. This PR allows optimizing chat-prompts on any agent framework.

Here is an example of the changes in the simplest case:

BEFORE:

from typing import Any, Dict
from opik.evaluation.metrics import LevenshteinRatio
from opik.evaluation.metrics.score_result import ScoreResult
from opik_optimizer.datasets import hotpot_300
from opik_optimizer import (
    ChatPrompt, 
    FewShotBayesianOptimizer,
)
dataset = hotpot_300()

def levenshtein_ratio(dataset_item: Dict[str, Any], llm_output: str) -> ScoreResult:
    metric = LevenshteinRatio()
    return metric.score(reference=dataset_item["answer"], output=llm_output)

project_name = "optimize-few-shot-bayesian-hotpot"

prompt = ChatPrompt(system="Answer the question.", user="{question}")

optimizer = FewShotBayesianOptimizer(
    model="openai/gpt-4o-mini",
    project_name=project_name,
    min_examples=3,
    max_examples=8,
    n_threads=16,
    seed=42,
)

optimization_result = optimizer.optimize_prompt(
    prompt=prompt,
    dataset=dataset,
    metric=levenshtein_ratio,
    n_trials=10,
    n_samples=150,
)

optimization_result.display()

AFTER:

from typing import Dict, Any
from opik.evaluation.metrics import LevenshteinRatio
from opik.evaluation.metrics.score_result import ScoreResult
from opik_optimizer.datasets import hotpot_300
from opik_optimizer import (
    OptimizableAgent,
    ChatPrompt,
    FewShotBayesianOptimizer,
    AgentConfig,
)

dataset = hotpot_300()

def levenshtein_ratio(dataset_item: Dict[str, Any], llm_output: str) -> ScoreResult:
    metric = LevenshteinRatio()
    return metric.score(reference=dataset_item["answer"], output=llm_output)

class LiteLLMAgent(OptimizableAgent):
    model = "openai/gpt-4o-mini"
    project_name = "optimize-few-shot-bayesian-hotpot"

agent_config = AgentConfig(
    chat_prompt=ChatPrompt(system="Answer the question.", user="{question}")
)

optimizer = FewShotBayesianOptimizer(
    model="openai/gpt-4o-mini",
    min_examples=3,
    max_examples=8,
    n_threads=16,
    seed=42,
)

optimization_result = optimizer.optimize_agent(
    agent_class=LiteLLMAgent,
    agent_config=agent_config,
    dataset=dataset,
    metric=levenshtein_ratio,
    n_trials=10,
    n_samples=50,
)
optimization_result.display()

The main difference is replacing:

project_name = "optimize-few-shot-bayesian-hotpot"

prompt = ChatPrompt(system="Answer the question.", user="{question}")

with:

class LiteLLMAgent(OptimizableAgent):
    model = "openai/gpt-4o-mini"
    project_name = "optimize-few-shot-bayesian-hotpot"

agent_config = AgentConfig(
    chat_prompt=ChatPrompt(system="Answer the question.", user="{question}")
)

And passing agent_class and agent_config to the optimize_agent() method.