> ## Documentation Index
> Fetch the complete documentation index at: https://arize-ax.mintlify.site/docs/llms.txt
> Use this file to discover all available pages before exploring further.

# Prompt Learning via SDK

> Optimize your prompts in code

# SDK + Guides

The Prompt Learning SDK is fully open source. View the SDK [here](https://github.com/Arize-ai/prompt-learning).

Check out [Guides](/ax/cookbooks/improve/overview) for some example use cases of prompt optimization through Prompt Learning!

* [Optimizing Coding Agent Prompts for Execution](/ax/cookbooks/improve/optimizing-coding-agent-prompts-for-execution) - Optimizing Coding Agents
* [Optimizing your LLM-as-Judge evaluators](/ax/cookbooks/improve/optimizing-your-eval-prompts) - Optimizing your LLM-as-Judge evaluators
* [Optimizing LLMs for Structured Output Generation](/ax/cookbooks/improve/improving-structured-output-generation-with-prompt-learning) - Optimizing LLMs for Structured Output Generation

# SDK Components

The Prompt Learning SDK consists of several key components:

## Core Classes

**`PromptLearningOptimizer`**

The main class that orchestrates the prompt optimization process.

**`MetaPrompt`**

Handles the construction of meta-prompts used for optimization.

**`TiktokenSplitter`**

Manages token counting and batching for large datasets.

**`Annotator`**

Generates additional annotations to guide the optimization process.

## Key Features

* **Automatic batching** based on token limits
* **Template variable detection** and preservation
* **Multiple evaluation methods** support
* **Flexible input formats** (strings, message lists, PromptVersion objects)
* **OpenAI model integration** for optimization

# Setup

First clone the Prompt Learning repository.

```
git clone https://github.com/Arize-ai/prompt-learning.git
```

```bash theme={null}
# Set your OpenAI API key
export OPENAI_API_KEY="your-api-key-here"
```

# Basic Usage

## 1. Initialize the Optimizer

```python theme={null}
from optimizer_sdk.prompt_learning_optimizer import PromptLearningOptimizer

optimizer = PromptLearningOptimizer(
    prompt="You are a helpful assistant. Answer this question: {question}",
    model_choice="gpt-4o",
    openai_api_key="your-api-key"  # Optional if set in environment
)
```

## 2. Prepare Your Dataset

Your dataset should contain:

* **Input columns**: The data your prompt will use (e.g., `question`)
* **Output column**: The LLM's response (e.g., `answer`)
* **Feedback columns**: Evaluation results (e.g., `correctness`, `explanation`)

```python theme={null}
import pandas as pd

dataset = pd.DataFrame({
    "question": ["What is the capital of France?", "What is 2+2?"],
    "answer": ["Paris", "4"],
    "correctness": ["correct", "correct"],
    "explanation": ["Accurate answer", "Correct calculation"]
})
```

## 3. Run Evaluators (Optional)

If you don't have pre-existing feedback, you can run evaluators:

```python theme={null}
from phoenix.evals import OpenAIModel, llm_generate

def evaluate_output(dataset):
    """Custom evaluator function"""
    # Your evaluation logic here
    return dataset, ["correctness", "explanation"]

# Run evaluators
dataset, feedback_columns = optimizer.run_evaluators(
    dataset=dataset,
    evaluators=[evaluate_output],
    feedback_columns=[]
)
```

## 4. Optimize the Prompt

```python theme={null}
optimized_prompt = optimizer.optimize(
    dataset=dataset,
    output_column="answer",
    feedback_columns=["correctness", "explanation"],
    context_size_k=128000  # 128k token context window
)
```

# Advanced Usage

## Batch Processing and Context Management

The SDK automatically handles large datasets by splitting them into batches that fit within your specified context window.

**`context_size_k` Parameter**

* **Purpose**: Controls the maximum token limit for each optimization batch
* **Default**: 128,000 tokens
* **Impact**: Larger values allow more examples per batch but may increase memory usage
* **Recommendation**: Start with 128k and adjust based on your model's context window

```python theme={null}
# For models with smaller context windows
optimized_prompt = optimizer.optimize(
    dataset=dataset,
    output_column="answer",
    feedback_columns=["correctness"],
    context_size_k=8000  # 8k token limit
)

# For models with larger context windows
optimized_prompt = optimizer.optimize(
    dataset=dataset,
    output_column="answer",
    feedback_columns=["correctness"],
    context_size_k=128000  # 128k token limit
)
```

## Template Variable Handling

The SDK automatically detects and preserves template variables in your prompts:

```python theme={null}
# Your prompt with template variables
prompt = "You are a {role}. Answer this {question_type}: {question}"

# The SDK will preserve {role}, {question_type}, and {question}
# These variables must be present in your dataset columns
```

## Multiple Evaluation Criteria

You can use multiple evaluators and feedback columns for comprehensive optimization:

```python theme={null}
# Run multiple evaluators
dataset, feedback_columns = optimizer.run_evaluators(
    dataset=dataset,
    evaluators=[evaluate_accuracy, evaluate_style, evaluate_completeness],
    feedback_columns=[]
)

# Optimize using all feedback
optimized_prompt = optimizer.optimize(
    dataset=dataset,
    output_column="answer",
    feedback_columns=feedback_columns
)
```

## Custom Annotations

Use the annotator to generate additional guidance for optimization. This allows you to pass in all your outputs and evals into another LLM call for a final, comprehensive eval.

```python theme={null}
from optimizer_sdk.annotator import Annotator

# Create custom annotation prompts
annotation_prompts = [
    "Analyze the style and tone of responses",
    "Check for factual accuracy and completeness"
]

# Generate annotations
annotations = optimizer.create_annotation(
    prompt=prompt,
    template_variables=["question"],
    dataset=dataset,
    feedback_columns=["correctness"],
    annotator_prompts=annotation_prompts,
    output_column="answer"
)

# Use annotations in optimization
optimized_prompt = optimizer.optimize(
    dataset=dataset,
    output_column="answer",
    feedback_columns=["correctness"],
    annotations=annotations
)
```

# Complete Example

Here's a complete example showing the full workflow:

```python theme={null}
import pandas as pd
from optimizer_sdk.prompt_learning_optimizer import PromptLearningOptimizer
from phoenix.evals import OpenAIModel, llm_generate

# 1. Initialize optimizer
optimizer = PromptLearningOptimizer(
    prompt="You are a math tutor. Solve this problem: {problem}",
    model_choice="gpt-4o"
)

# 2. Prepare dataset
dataset = pd.DataFrame({
    "problem": ["2 + 2 = ?", "5 * 3 = ?", "10 / 2 = ?"],
    "answer": ["4", "15", "5"],
    "correctness": ["correct", "correct", "correct"],
    "explanation": ["Correct addition", "Correct multiplication", "Correct division"]
})

# 3. Optimize prompt
optimized_prompt = optimizer.optimize(
    dataset=dataset,
    output_column="answer",
    feedback_columns=["correctness", "explanation"],
    context_size_k=8000
)

print("Original prompt:", optimizer.prompt)
print("Optimized prompt:", optimized_prompt)
```

# Configuration Options

## Model Selection

The SDK supports various OpenAI models:

```python theme={null}
# Supported models
SUPPORTED_MODELS = [
    "o1",           # OpenAI o1
    "o3",           # OpenAI o3
    "gpt-4o",       # GPT-4 Omni
    "gpt-4",        # GPT-4
    "gpt-3.5-turbo", # GPT-3.5 Turbo
    "gpt-3.5",      # GPT-3.5
]

# Choose based on your needs
optimizer = PromptLearningOptimizer(
    prompt="Your prompt here",
    model_choice="gpt-4o"  # Best for complex optimization
)
```

## Input Format Flexibility

The SDK accepts multiple prompt formats:

```python theme={null}
# String format
optimizer = PromptLearningOptimizer(
    prompt="You are a helpful assistant: {input}",
    model_choice="gpt-4o"
)

# Message list format
optimizer = PromptLearningOptimizer(
    prompt=[
        {"role": "system", "content": "You are a helpful assistant"},
        {"role": "user", "content": "Answer this: {input}"}
    ],
    model_choice="gpt-4o"
)

# PromptVersion format (Phoenix integration)
from phoenix.client.types import PromptVersion
prompt_version = PromptVersion(...)
optimizer = PromptLearningOptimizer(
    prompt=prompt_version,
    model_choice="gpt-4o"
)
```

# Best Practices

## 1. Dataset Quality

* Ensure your dataset is representative of real-world usage
* Include diverse examples that cover edge cases
* Balance positive and negative feedback

## 2. Evaluation Criteria

* Define clear, measurable evaluation criteria
* Use multiple evaluators for comprehensive feedback
* Consider both objective (accuracy) and subjective (style) metrics

## 3. Context Window Management

* Start with smaller context windows for faster iteration
* Increase context size for more comprehensive optimization
* Monitor token usage to optimize costs

## 4. Iterative Improvement

* Run multiple optimization loops
* Monitor performance metrics across iterations
* Stop when performance plateaus or meets your criteria

## 5. Template Variable Preservation

* Always verify that template variables are preserved
* Test optimized prompts with new data
* Ensure backward compatibility

# Conclusion

The Prompt Learning SDK provides a powerful, automated approach to optimizing LLM prompts. By leveraging evaluation feedback and meta-prompt optimization, you can systematically improve prompt performance across various use cases.

Key benefits:

* **Automated optimization** reduces manual prompt engineering
* **Data-driven improvements** based on actual performance metrics
* **Scalable approach** for production systems
* **Flexible integration** with existing evaluation frameworks

Start with simple use cases and gradually incorporate more sophisticated evaluation criteria as you become familiar with the SDK's capabilities.