"""Few-shot example management for prompt engineering.
This module provides tools for managing, selecting, and formatting
few-shot examples to include in prompts.
"""
import random
from dataclasses import asdict, dataclass, field
from typing import Any, Callable, Dict, List, Optional, Union
from kerb.core.enums import SelectionStrategy
[docs]
@dataclass
class FewShotExample:
"""A few-shot example with input and output.
Attributes:
input (str): Example input
output (str): Example output
metadata (Dict[str, Any]): Additional metadata (category, difficulty, etc.)
"""
input: str
output: str
metadata: Dict[str, Any] = field(default_factory=dict)
[docs]
def to_dict(self) -> Dict[str, Any]:
"""Convert to dictionary representation."""
return asdict(self)
[docs]
class ExampleSelector:
"""Manages and selects few-shot examples for prompts.
Supports multiple selection strategies including random, semantic similarity,
and diversity-based selection.
"""
[docs]
def __init__(self, examples: Optional[List[FewShotExample]] = None):
"""Initialize the example selector.
Args:
examples (Optional[List[FewShotExample]]): Initial examples. Defaults to None.
"""
self.examples = examples or []
[docs]
def add(self, example: FewShotExample) -> None:
"""Add an example to the selector.
Args:
example (FewShotExample): Example to add
"""
self.examples.append(example)
[docs]
def select(
self,
k: int = 3,
strategy: Union[str, SelectionStrategy] = SelectionStrategy.RANDOM,
query: Optional[str] = None,
filter_fn: Optional[Callable[[FewShotExample], bool]] = None,
) -> List[FewShotExample]:
"""Select k examples using the specified strategy.
Args:
k (int): Number of examples to select. Defaults to 3.
strategy (Union[str, SelectionStrategy]): Selection strategy. Defaults to RANDOM.
query (Optional[str]): Query text for semantic selection. Defaults to None.
filter_fn (Optional[Callable]): Filter function to apply before selection.
Defaults to None.
Returns:
List[FewShotExample]: Selected examples
"""
# Apply filter if provided
candidates = self.examples
if filter_fn:
candidates = [ex for ex in candidates if filter_fn(ex)]
if not candidates:
return []
# Convert strategy to enum if needed
if isinstance(strategy, str):
strategy = SelectionStrategy(strategy.lower())
# Limit k to available examples
k = min(k, len(candidates))
if strategy == SelectionStrategy.RANDOM:
return random.sample(candidates, k)
elif strategy == SelectionStrategy.FIRST:
return candidates[:k]
elif strategy == SelectionStrategy.LAST:
return candidates[-k:]
elif strategy == SelectionStrategy.DIVERSE:
# Select diverse examples by maximizing dissimilarity
# Simple heuristic: space out selections evenly
if k >= len(candidates):
return candidates
step = len(candidates) / k
indices = [int(i * step) for i in range(k)]
return [candidates[i] for i in indices]
elif strategy == SelectionStrategy.SEMANTIC:
# Semantic selection requires embeddings
if query is None:
raise ValueError("Query text is required for semantic selection")
try:
from ..embedding import cosine_similarity, embed
# Embed query and examples
query_embedding = embed(query)
example_embeddings = [embed(ex.input) for ex in candidates]
# Calculate similarities
similarities = [
cosine_similarity(query_embedding, ex_emb)
for ex_emb in example_embeddings
]
# Select top k by similarity
sorted_indices = sorted(
range(len(similarities)),
key=lambda i: similarities[i],
reverse=True,
)
return [candidates[i] for i in sorted_indices[:k]]
except ImportError:
# Fall back to random if embedding is not available
return random.sample(candidates, k)
return candidates[:k]
[docs]
def create_example(
input_text: str, output_text: str, metadata: Optional[Dict[str, Any]] = None
) -> FewShotExample:
"""Create a few-shot example.
Args:
input_text (str): Example input
output_text (str): Example output
metadata (Optional[Dict[str, Any]]): Additional metadata. Defaults to None.
Returns:
FewShotExample: The created example
Examples:
>>> ex = create_example(
... input_text="What is 2+2?",
... output_text="4",
... metadata={"difficulty": "easy"}
... )
"""
return FewShotExample(input=input_text, output=output_text, metadata=metadata or {})
[docs]
def select_examples(
examples: List[FewShotExample],
k: int = 3,
strategy: Union[SelectionStrategy, str] = "random",
query: Optional[str] = None,
filter_fn: Optional[Callable[[FewShotExample], bool]] = None,
) -> List[FewShotExample]:
"""Select k examples from a list using the specified strategy.
Args:
examples: Available examples
k: Number of examples to select
strategy: Selection strategy (SelectionStrategy enum or string: "random", "similarity", "diverse", "recent", "fixed")
query: Query text for semantic selection
filter_fn: Filter function
Returns:
List[FewShotExample]: Selected examples
Examples:
>>> from kerb.core.enums import SelectionStrategy
>>> selected = select_examples(examples, k=3, strategy=SelectionStrategy.SIMILARITY, query="example query")
"""
from kerb.core.enums import validate_enum_or_string
# Validate and normalize strategy
strategy_val = validate_enum_or_string(strategy, SelectionStrategy, "strategy")
if isinstance(strategy_val, SelectionStrategy):
strategy_str = strategy_val.value
else:
strategy_str = strategy_val
selector = ExampleSelector(examples)
return selector.select(k=k, strategy=strategy_str, query=query, filter_fn=filter_fn)