Customizing models used in the Optimizer

Opik Agent Optimizer leverages LiteLLM under the hood to provide comprehensive support for a wide array of Large Language Models (LLMs). This integration means you can use virtually any model provider or even locally hosted models with any of the Opik optimization algorithms.

The Agent Optimizer uses LLM’s to optimize prompts, as a result there are two different ways to configure the LLM calling behavior in the Optimizer:

1. ChatPrompt model: This is the LLM model used in your application to generate the response
2. Optimization model: This is the LLM model used by the Opik Optimizer to improve the ChatPrompt

Configuring ChatPrompt model

How it works

The ChatPrompt has two different parameters that you can use to configure the LLM response:

1. model: Specify a model to use
2. invoke: Specify a function to generate the response, this is more customizable than the model field

We recommend starting by using the model parameter that allows you to specify a provider or model to use to generate the LLM out. We recommend that you set the model string to the model that you use within your LLM application.

If you need to customize the LLM generation behavior further to match what you do in your LLM application, you can use the invoke parameter in the ChatPrompt. This parameter supports a function that takes the model string, the messages list and the tool list and returns the LLM response string.

Specifying a ChatPrompt model

You can configure the ChatPrompt model using the model parameter:

1
from opik_optimizer import ChatPrompt
2
3
prompt = ChatPrompt(
4
  name="demo_prompt",
5
  model="openai/gpt-4o-mini"
6
  messages=[
7
    {"role": "system", "content": "You are a helpful assistant."},
8
    {"role": "user", "content": "{question}"}
9
  ]
10
)

Customizing the LLM response

The invoke parameter of the ChatPrompt allows you to customize the LLM response. This allows you to fully customize the ChatPrompt behavior, the invoke function should follow the following parameters:

1
import litellm
2
from opik_optimizer import ChatPrompt
3
4
def invoke_llm(model, messages, tools):
5
    res = litellm.completion(
6
        model=model,
7
        messages=messages,
8
        tools=tools
9
    )
10
    return res.choices[0].delta.content
11
12
prompt = ChatPrompt(
13
  name="demo_prompt",
14
  messages=[
15
    {"role": "system", "content": "You are a helpful assistant."},
16
    {"role": "user", "content": "{question}"}
17
  ],
18
  invoke=invoke_llm
19
)

The invoke method can also be used to optimize prompts that utilize structured outputs:

1
from pydantic import BaseModel
2
from openai import OpenAI
3
4
client = OpenAI()
5
6
class CalendarEvent(BaseModel):
7
    name: str
8
    date: str
9
    participants: list[str]
10
11
def invoke_llm(model, messages, tools):
12
    completion = client.chat.completions.parse(
13
        model=model,
14
        messages=messages,
15
        response_format=CalendarEvent,
16
    )
17
18
    return completion.choices[0].message.parsed
19
20
prompt = ChatPrompt(
21
  name="demo_prompt",
22
  messages=[
23
    {"role": "system", "content": "You are a helpful assistant."},
24
    {"role": "user", "content": "{question}"}
25
  ],
26
  invoke=invoke_llm
27
)

Configuring Optimization model

How it Works

When you initialize an optimizer (e.g., MetaPromptOptimizer, FewShotBayesianOptimizer, MiproOptimizer, EvolutionaryOptimizer), the model parameter (and reasoning_model for MetaPromptOptimizer) accepts a LiteLLM model string. LiteLLM then handles the communication with the specified model provider or local endpoint.

This builds upon Opik’s existing tracing integration with LiteLLM, ensuring that not only are your optimization processes flexible in model choice, but calls made during optimization can also be seamlessly tracked if you have the OpikLogger callback configured in LiteLLM.

Specifying Models

You pass the LiteLLM model identifier string directly to the model parameter in the optimizer’s constructor. Here are some common examples:

OpenAI

1
from opik_optimizer import MetaPromptOptimizer
2
3
optimizer = MetaPromptOptimizer(
4
    model="openai/gpt-4o-mini", # Standard OpenAI model
5
    # ... other parameters
6
)

Azure OpenAI

1
from opik_optimizer import FewShotBayesianOptimizer
2
3
optimizer = FewShotBayesianOptimizer(
4
    model="azure/your-azure-deployment-name", # Your Azure deployment name
5
    # Ensure your environment variables for Azure (AZURE_API_KEY, AZURE_API_BASE, AZURE_API_VERSION) are set
6
    # ... other parameters
7
)

Anthropic

1
from opik_optimizer import FewShotBayesianOptimizer
2
3
optimizer = FewShotBayesianOptimizer(
4
    model="anthropic/claude-3-opus-20240229",
5
    # Ensure ANTHROPIC_API_KEY is set
6
    # ... other parameters
7
)

Google Gemini

1
from opik_optimizer import EvolutionaryOptimizer
2
3
optimizer = EvolutionaryOptimizer(
4
    model="gemini/gemini-1.5-pro-latest",
5
    # Ensure GEMINI_API_KEY is set
6
    # ... other parameters
7
)

Local Models (e.g., via Ollama)

LiteLLM allows you to connect to models served locally through tools like Ollama.

$
ollama serve
>
ollama pull llama3

1
from opik_optimizer import MetaPromptOptimizer
2
3
optimizer = MetaPromptOptimizer(
4
    model="ollama_chat/llama3", # Or "ollama/llama3" if not using the chat API specifically
5
    # LiteLLM defaults to http://localhost:11434 for Ollama.
6
    # If Ollama runs elsewhere, set an environment variable like OLLAMA_API_BASE_URL or pass api_base in model_kwargs
7
    # model_kwargs={"api_base": "http://my-ollama-host:11434"}
8
    # ... other parameters
9
)

Other Providers

LiteLLM supports numerous other providers like Cohere, Mistral AI, Bedrock, Vertex AI, Hugging Face Inference Endpoints, etc. Refer to the LiteLLM documentation on Providers for the correct model identifier strings and required environment variables for API keys.