TaskLLM

A library for optimizing LLM tasks, including prompt engineering and bandit-based training.

Project Overview

TaskLLM is a Python library designed to help developers optimize their interactions with Large Language Models (LLMs). It provides tools for:

• Instrumenting LLM tasks to track inputs, outputs, and quality metrics
• Making LLM calls with both simple text responses and structured outputs
• Optimizing prompts through bandit-based training algorithms

The library is particularly useful for developers who want to:

• Systematically improve prompt performance
• Track and analyze LLM interactions
• Convert unstructured LLM outputs into structured data
• Implement quality assessment for LLM outputs

Installation

TaskLLM requires Python 3.11 or higher.

# Install from PyPI, not yet available.
pip install taskllm

# Or install from source
git clone https://github.com/your-username/taskllm.git
cd taskllm
pip install -e .

Quick Start

Please see one of the examples for a quick start. It will show you how to train a prompt for a specific task.

Video to come out shortly!

Core Components

Optimization

The BanditTrainer class helps you find the best prompts for your tasks:

from taskllm.optimizer.methods import BanditTrainer
from taskllm.optimizer.data import DataSet, Row

trainer = BanditTrainer(
    all_rows=your_dataset,
    task_guidance="your task description",
    keys=["input_field1", "input_field2"],
    expected_output_type=YourOutputModel,
    scoring_function=your_scoring_function
)

await trainer.train()
best_prompt = await trainer.get_best_prompt()

Examples

The repository includes several examples demonstrating how to use TaskLLM in different scenarios:

Jokes Example

Determines whether jokes are funny using a bandit-based prompt optimizer:

# From examples/jokes/run.py
trainer = BanditTrainer(
    all_rows=dataset,
    task_guidance="write a prompt that determines whether a joke is funny based on the category of joke",
    keys=["joke"],
    expected_output_type=IsJokeFunny,
    scoring_function=funny_scoring_function
)

Tweet Sentiment Analysis

Analyzes the sentiment of tweets (positive, negative, or neutral):

# From examples/tweet_sentiment/run.py
trainer = BanditTrainer(
    all_rows=dataset,
    task_guidance="what is the sentiment of this tweet?",
    keys=["tweet"],
    expected_output_type=TweetSentiment,
    scoring_function=sentiment_scoring_function
)

Starbucks Reviews

Rates Starbucks reviews on a scale of 1-5:

# From examples/starbucks/run.py
trainer = BanditTrainer(
    all_rows=dataset,
    task_guidance="determine the rating of this review",
    keys=["review", "name", "location", "date"],
    expected_output_type=StarbucksReviewRating,
    scoring_function=sentiment_scoring_function,
    prompt_mode=PromptMode.ADVANCED
)

To run any of these examples:

cd examples/[example_directory]
python run.py

Advanced Usage

Configuration Options

You can customize the LLM configuration:

from taskllm.ai import LLMConfig

custom_config = LLMConfig(
    model="gpt-4o",
    temperature=0.7,
    max_tokens=2000,
    top_p=0.95,
    frequency_penalty=0.5,
    presence_penalty=0.5
)

Quality Labeling

Enable interactive quality assessment for your tasks:

@instrument_task("your_task", enable_quality_labeling=True)
def your_function():
    # After execution, you'll be prompted to rate the quality

Caching Strategies

TaskLLM automatically caches LLM responses to save time and costs:

# Disable caching for specific calls
response = await simple_llm_call(
    messages=[...],
    config=config,
    use_cache=False
)

Custom Prompt Modes

Use advanced prompt modes for more sophisticated optimization:

from taskllm.optimizer.prompt.meta import PromptMode

trainer = BanditTrainer(
    # ...other parameters
    prompt_mode=PromptMode.ADVANCED
)

API Reference

Key Modules

• taskllm.instrument: Functions for tracking and logging LLM tasks
• taskllm.ai: Interface for making LLM calls
• taskllm.optimizer: Tools for optimizing prompts
  • taskllm.optimizer.data: Data structures for optimization
  • taskllm.optimizer.methods: Optimization algorithms
  • taskllm.optimizer.prompt: Prompt management and templates

Contributing

Development Setup

1. Clone the repository

   git clone https://github.com/your-username/taskllm.git
   cd taskllm
   

2. Create a virtual environment

   python -m venv .venv
   source .venv/bin/activate  # On Windows: .venv\Scripts\activate
   

3. Install development dependencies

   pip install -e ".[dev]"
   

Testing

Run tests using:

python -m pytest

Contribution Guidelines

1. Fork the repository
2. Create a feature branch
3. Add your changes
4. Run tests
5. Submit a pull request

License

MIT License

---

For more information, check out the examples directory or open an issue on GitHub.