Prompt Learning SDK and CLI

repository·main·Indexed 43 days ago

https://github.com/arize-ai/prompt-learning

An SDK and CLI tool designed to optimize LLM prompts using natural language feedback instead of scalar scores. It utilizes an Agent-Evaluator-Meta Prompt loop to iteratively refine prompts. The package includes the PromptLearningOptimizer for improving prompts via meta-prompting, support for OpenAI and Google AI providers, and integration with Arize Phoenix for experiment tracking. Version 0.1.0 includes features for pricing controls and production-ready prompt optimization.

Tokens
28.2K
Snippets
85
Records
120
Agent score
55%

What's inside Prompt Learning

  1. Project file organization and structure

    main

    The repository is organized into shared utilities and agent-specific implementations:

    Shared (root level)

    • constants.py: Paths, prompts, and model configurations.
    • container_helpers.py: Docker container lifecycle management.
    • evals.py: Shared LLM-as-judge evaluation logic.
    • phoenix_experiments.py: Logging via Phoenix REST API.

    Cline (cline/)

    • cline_helpers.py: Low-level Docker and execution utilities.
    • act_mode/: Orchestration for Act Mode (includes run_act.py and optimization notebooks).
    • plan_mode/: Orchestration for Plan Mode.

    Claude Code (claude_code/)

    • claude_code_helpers.py: Helpers for running the Claude Code CLI.
    • run_claude.py: Parallel execution logic for Claude Code.
    • optimize_claude_code.py: Automated optimization loop for Claude Code.
  2. Overview of the SWE-bench harness functionality

    main

    The harness in this directory is designed to run coding agents (such as Cline or Claude Code) against SWE-bench instances within isolated Docker containers.

    Workflow:

    1. Materialization: It extracts each SWE-bench repository from a prebuilt Docker image into a host workspace directory.
    2. Isolation: It starts a container with that workspace bind-mounted to /testbed.
    3. Execution: It runs the specified coding agent on the provided problem statement.
    4. Evaluation: It collects the agent's outputs (either patches or plans) and evaluates their correctness.
  3. Extend the library with custom AI providers

    main
    The library uses a ModelProvider abstraction to decouple the core optimization logic from specific AI implementations. To add a new provider (e.g., a custom LLM or a new service), you should implement the abstract interface defined in providers/base_provider.py. This allows you to use the PromptLearningOptimizer with your custom provider via dependency injection.
  4. How Prompt Learning optimization works

    main

    The optimization process follows these steps:

    1. Input: Takes a baseline prompt and an evaluation dataset.
    2. Batching: Splits the data into batches that fit within the LLM's context window (using tiktoken_splitter.py).
    3. Analysis: Uses a meta-prompt to analyze failures in the current prompt and suggest improvements.
    4. Iteration: Iteratively refines the prompt across the different batches.
    5. Output: Returns the optimized prompt in the same format as the original input.
  5. How feedback columns drive prompt optimization

    main

    Prompt optimization relies on feedback columns in your dataset. These columns must contain natural language descriptions (critiques, annotations, or explanations) of what went wrong or how an output could be improved.

    Dataset Requirements:

    1. Input columns: Variables used in your prompt template (e.g., {text}).
    2. Output column: The LLM's response for each input.
    3. Feedback column(s): Textual critique of the output.

    Example CSV Structure:

    input,output,feedback
    "Generate a tech page","{ ""sections"": [...]}","Missing 'updatedAt' field; top-level key should be 'page'"
    "Generate a menu","{ ""menu"": [...]}","Good structure but missing required 'metadata' section"
  6. Project Architecture of Prompt Learning

    main

    The prompt-learning repository is organized into several functional modules:

    • cli/: The command-line interface, including the entry point (main.py) and command implementations like optimize and image.
    • core/: Business logic for cost tracking (pricing.py), token-aware batch splitting (dataset_splitter.py), and error handling (exceptions.py).
    • interfaces/: Abstractions such as token_counter.py.
    • providers/: Implementations for different AI providers (e.g., google_provider.py) based on a base_provider.py interface.
    • optimizer_sdk/: The core SDK containing the main optimizer (prompt_learning_optimizer.py), meta-prompt templates (meta_prompt.py), and feedback annotation (annotator.py).
    • evaluators/: Built-in assessment tools like image_evaluator.py for image quality.
  7. Manage costs and budgets

    main

    The project includes built-in budget enforcement to prevent runaway costs.

    • CLI Budgeting: Use the --budget flag with prompt-learn optimize to set a hard limit.
    • Real-time Tracking: Use the --verbose flag to see cost tracking during execution.
    • Automatic Enforcement: Optimization will automatically stop once the budget limit is reached.
  8. Understand the Prompt Learning optimization loop

    main

    The project follows a cyclical optimization pattern to improve agent performance:

    1. Evaluation Phase: Run run_act() to evaluate an agent (like Cline) on SWE-bench instances to generate a dataset of attempts and outcomes.
    2. Qualitative Assessment: Call evaluate_results() to add LLM-based reasoning and correctness scores to that dataset.
    3. Optimization Phase: Pass the feedback-enriched dataset to PromptLearningOptimizer.optimize() to generate an improved ruleset or prompt.
    4. Iteration: Use the newly optimized prompt in the next evaluation round. This loop is typically repeated 3-5 times.
  9. Install and build the SWE-bench dataset

    main

    To run coding agents against SWE-bench, you must first clone the official SWE-bench repository, install it in editable mode, and build the necessary Docker images.

    Prerequisites:

    • Python 3.10+
    • pip install swebench pandas datasets
    • An OPENAI_API_KEY environment variable set for the optimizer_sdk and evaluation.
    • Docker Desktop running with at least 12-16 GB of memory and ~120 GB of free disk space.

    Steps:

    1. Clone and install SWE-bench:
      git clone https://github.com/princeton-nlp/SWE-bench.git
      cd SWE-bench
      pip install -e .
    2. Build images for SWE-bench Lite (300 rows from the test split):
      python -m swebench.harness.prepare_images \
        --dataset_name SWE-bench/SWE-bench_Lite \
        --split test \
        --max_workers 4 \
        --namespace swebench-images \
        --tag latest \
        --env_image_tag latest
      Note: On Apple Silicon/M-series, use --namespace '' to build locally for your architecture. Adjust --max_workers based on your system resources.
    3. Verify the images exist:
      docker images | grep '^sweb.env'
      docker images | grep '^sweb.eval'
    git clone https://github.com/princeton-nlp/SWE-bench.git
    cd SWE-bench
    pip install -e .
    
    python -m swebench.harness.prepare_images \
      --dataset_name SWE-bench/SWE-bench_Lite \
      --split test \
      --max_workers 4 \
      --namespace swebench-images \
      --tag latest \
      --env_image_tag latest
  10. Configure environment variables for OpenAI or Google AI

    main

    Set your API keys as environment variables before running the CLI or SDK.

    For OpenAI (the default provider), use OPENAI_API_KEY. For Google AI / Gemini, use either GOOGLE_API_KEY or GEMINI_API_KEY.

    # For OpenAI (default provider)
    export OPENAI_API_KEY="your-openai-key"
    
    # For Google AI / Gemini
    export GOOGLE_API_KEY="your-google-key"
    # or
    export GEMINI_API_KEY="your-google-key"
  11. Prerequisites for the Cline SWE-bench Harness

    main

    Before using the harness, ensure you have completed the shared SWE-bench setup described in ../SETUP_SWEBENCH.md. Additionally, the following tools must be installed and available on your system:

    • Node.js
    • npx + tsx (accessible via npx tsx)
    • grpcurl (on macOS: brew install grpcurl)
    • ripgrep (on macOS: brew install ripgrep).

    Important: You must symlink ripgrep for the Cline core to function. From the Cline repository root, run:

    ln -sf "$(command -v rg)" dist-standalone/rg
    ln -sf "$(command -v rg)" dist-standalone/rg
  12. Configure Environment Variables for Providers

    main

    Before running optimizations, set your API keys in your environment. The system supports OpenAI and Google AI (Gemini).

    • For OpenAI: OPENAI_API_KEY
    • For Google AI: GOOGLE_API_KEY or GEMINI_API_KEY
    export OPENAI_API_KEY="your-openai-key"
    export GOOGLE_API_KEY="your-google-ai-key"