Open source library

Polar Llama

A Python library for parallel LLM inference across providers, built on Polars DataFrames.

v0.5.2Released July 12, 2026
View on GitHubPyPI package
Version0.5.2latest0.5.10.5.00.3.00.2.20.2.10.2.00.1.7

Overview

Polar Llama is a Python library that enables parallel inference calls to multiple Large Language Model providers through Polars dataframes. It streamlines batch processing of AI queries without serial request delays, making it ideal for data-intensive AI applications. Since 0.2.2 it has grown well beyond inference: embeddings and vector search, taxonomy tagging, tool use / MCP, provider-native prompt caching, a DSPy-style prompt optimizer, and — as of 0.5.0 — a fully on-device MLX backend for Apple Silicon.

Concurrent Processing

Send multiple inference requests in parallel without waiting for individual completions

🔌Six Providers

OpenAI, Anthropic, Gemini, Groq, AWS Bedrock, and on-device MLX (Apple Silicon)

🧰Tool Use / MCP

Dataframe-native tool calling — emission, batch-parallel execution, and synthesis are all ordinary columns

🎛️Prompt Optimization

A DSPy-style engine (Signature, Predict, BootstrapFewShot, InstructionOptimizer) tunes prompts against labeled data

Installation

Using pip

bash
pip install polar-llama==0.5.2

On-device inference (Apple Silicon)

Pulls in mlx and mlx-lm; requires an Apple Silicon Mac and Python ≥ 3.10

bash
pip install "polar-llama[local]"

Development Installation

bash
maturin develop

Quick Start

Get started with a simple example

python
import polars as pl
from polar_llama import Provider
import dotenv

dotenv.load_dotenv()

# Example questions
questions = [
    'What is the capital of France?',
    'What is the difference between polars and pandas?'
]

df = pl.DataFrame({'Questions': questions})

# Using the fluent .llama namespace (recommended)
df = df.with_columns(
    answer=pl.col('Questions').llama.inference_async(
        provider=Provider.OPENAI,
        model='gpt-4o-mini'
    )
)

What’s New in 0.5.2

Latest

Fixed: in-process local inference on Gemma 3n

Two stacked bugs in the in-process MLX engine, both fixed at the model-load path

inference_local(engine="in_process") on Gemma 3n no longer crashes end-to-end. The mlx-lm #1384 batched shared-KV patch was previously applied only on the prompt-tuning bridge and benchmark paths, never on theinference_local load path — so MlxBatchEngine loaded Gemma 3n unpatched and batched generation raised ValueError: too many values to unpack (expected 2). Both this patch and a new guard are now applied automatically at model load, in polar_llama/local/engine.py::_apply_mlx_patches.

Fixing that unmasked a second, underlying bug: mlx_lm.generate.BatchGenerator.stats divided prompt_tokens / prompt_time with prompt_time == 0 in its teardown, raising a ZeroDivisionError during exception handling and silently replacing the real error. A new guarded, idempotent patch — apply_batchgen_stats_zerodiv_patch — wraps the context manager so a body exception is never masked and a zero-time exit yields tps = 0.0 instead of throwing.

Everything Since 0.2.2

0.5.2 is cumulative — every feature shipped in 0.2.2 is still here, plus everything added across 0.3.0, 0.5.0, and 0.5.1.

0.3.0 — Tool Use / MCP, Prompt Optimization, Caching

Released 2026-06-10

🧰Tool Use / MCP

tools_to_response_model, mcp_tools, execute_tool_calls, tool_results_to_message — batch-parallel tool calling over an MCP server or a Python executor

🗄️Provider-Native Prompt Caching

cache=True / CacheConfig shares a cached system prefix across rows via Anthropic cache_control, with 5-minute and 1-hour TTLs

🎛️Prompt Optimization Engine

Signature, Predict, evaluate, BootstrapFewShot, InstructionOptimizer — DSPy-style instruction and few-shot tuning

🌐Proxy / Gateway Support

OPENAI_BASE_URL and ANTHROPIC_BASE_URL overrides, plus POLAR_LLAMA_MAX_CONCURRENCY to bound in-flight requests (default 64)

Also updated the default model for every provider (previous defaults were retired or decommissioned — see the Provider Support table below), added native Gemini system_instruction support and JSON-schema structured outputs, fixed Gemini and Bedrock structured-output auth paths, and made Bedrock work from the synchronous inference expression. Minimum supported Python is now 3.9.

0.5.0 — Local MLX Inference Backend (Apple Silicon)

Released 2026-07-04

col(...).llama.inference_local(...) runs on-device batched generation on Apple Silicon — no API keys, no network — behind two engines:

EngineHow it runsRequires
server (default)Existing async fan-out talks HTTP to a local OpenAI-compatible endpoint you startedmlx_lm.server, vllm-mlx, or any server speaking /v1/chat/completions
in_processA map_batches UDF drives mlx-lm's BatchGenerator directly in-processpip install polar-llama[local] (Apple Silicon, Python ≥ 3.10)
Collapsed Prefix Prefill

Computes a shared prompt prefix once instead of re-prefilling it per row — 10.36× vs sequential on gemma-3n E4B (32 rows, 5 KB shared prompt) at 32/32 exact greedy parity

📉Batched Quantized KV Cache

BatchQuantizedKVCache (opt-in via POLAR_LLAMA_LOCAL_KV_BITS=4) cuts KV memory ~47% at fp16 output parity, roughly doubling batch/context that fits in 24 GB

🩹mlx-lm #1384 Fix

Runtime monkeypatch correcting a RoPE offset-aliasing bug that garbled batched generation on hybrid Gemma 3n / Gemma 4 models

🧪CPU-Safe Test Seam

A LocalEngine protocol with a FakeEngine implementation keeps batching/ordering/error-isolation logic testable on CI with no GPU

0.5.1 — Local Prompt-Tuning Bridge

Released 2026-07-05

polar_llama.local.make_local_inference_fn(model, ...) returns an inference_fn that drives the prompt optimizer’s Predict / BootstrapFewShot / InstructionOptimizer against on-device Gemma 3n via mlx-lm — no cloud or Rust path involved. It reuses the singleton-loaded weights and applies the mlx-lm #1384 batched fix automatically.

🚀Collapsed Prefill for inference_local

POLAR_LLAMA_LOCAL_COLLAPSE=1 shares the common prompt prefix across rows — ~2.8× faster on a full prompt-tuning schedule (3.4× on a demo-laden eval) at identical output

♻️Singleton Weight Reuse

MlxBatchEngine.get_model_and_tokenizer() reuses already-loaded weights instead of reloading per call

Fixed: InstructionOptimizer no longer crashes with TypeError: the truth value of a Series is ambiguous when a proposer model returns instructions as a JSON array instead of a newline-delimited string — list/Series values are now flattened to newline-delimited text. Note: POLAR_LLAMA_LOCAL_COLLAPSE is mutually exclusive with POLAR_LLAMA_LOCAL_KV_BITS (the quantized-KV path takes precedence when both are set).

Examples & Cookbooks

Tool Use: Emit, Execute, Synthesize

The agent loop unrolled into ordinary dataframe columns

python
from polar_llama import (
    mcp_tools, tools_to_response_model, execute_tool_calls,
    tool_results_to_message, combine_messages, inference_messages, Provider,
)

tools = mcp_tools("http://localhost:8811/mcp")     # tools/list introspection
ToolCalls = tools_to_response_model(tools)          # emission schema

df = (
    df
    # 1. Emit: the LLM parameterizes N calls per row (structured output)
    .with_columns(calls=pl.col("meal").llama.inference_async(
        provider=Provider.OPENAI, model="gpt-4o-mini", response_model=ToolCalls))
    # 2. Execute: all calls across all rows, in parallel, errors as data
    .with_columns(results=execute_tool_calls(
        pl.col("calls"), transport="http://localhost:8811/mcp", tools=tools))
    # 3. Synthesize: fold results back through a second inference pass
    .with_columns(answer=inference_messages(
        combine_messages(
            pl.col("meal").llama.to_message(role="user"),
            tool_results_to_message(pl.col("results")),
        ),
        provider=Provider.OPENAI, model="gpt-4o-mini"))
)

Prompt Optimization (DSPy-style)

Declare a task, then let an optimizer tune instructions or mine few-shot demos

python
import polars as pl
from polar_llama import Predict, Signature, BootstrapFewShot, InstructionOptimizer, evaluate

# 1. Declare the task
module = Predict(
    Signature("question -> answer", instructions="Answer concisely."),
    provider="openai",
    model="gpt-4o-mini",
)

# 2. Labeled training data
trainset = pl.DataFrame({
    "question": ["What is 2+2?", "Capital of France?", "Largest planet?"],
    "answer": ["4", "Paris", "Jupiter"],
})

# 3. A metric: (gold row, prediction) -> bool | float
def exact_match(example, prediction):
    return example["answer"].strip().lower() == (prediction["answer"] or "").strip().lower()

# 4a. Bootstrap few-shot demos from rows the model already gets right
compiled = BootstrapFewShot(metric=exact_match, max_demos=4).compile(module, trainset)

# 4b. Or search for better instructions (COPRO-style)
optimizer = InstructionOptimizer(metric=exact_match, n_candidates=4)
compiled = optimizer.compile(module, trainset)
print(optimizer.history)  # [(instructions, score), ...]

# 5. Run the optimized module on new data — outputs land in pred_* columns
result = compiled(pl.DataFrame({"question": ["What is 3+3?"]}))
print(result["pred_answer"])

# Score any module against a labeled set
print(evaluate(compiled, trainset, exact_match).score)

On-Device Inference (Apple Silicon, in-process engine)

Batched generation via mlx-lm, no API keys, no network

python
import polars as pl
import polar_llama  # registers the .llama namespace

df = pl.DataFrame({"ticket": [
    "My laptop won't turn on even when it's plugged in.",
    "I forgot my password and I'm locked out of my account.",
]})

tagged = df.with_columns(
    tag=pl.col("ticket").llama.inference_local(
        model="mlx-community/gemma-3n-E4B-it-lm-4bit",
        system="Classify the ticket as exactly one of: Hardware, Account, Billing.",
        engine="in_process",   # on-device mlx-lm, batched over the column
        max_tokens=16,
        temperature=0.0,
    )
)

Data Analysis Pipeline

Process customer feedback at scale

python
import polars as pl
from polar_llama import string_to_message, inference_async, Provider

# Load customer feedback data
feedback_df = pl.DataFrame({
    'customer_id': [101, 102, 103, 104, 105],
    'feedback': [
        'The product is amazing but shipping was slow',
        'Great quality, highly recommend!',
        'Disappointed with customer service',
        'Perfect for my needs, will buy again',
        'Product arrived damaged, requesting refund'
    ]
})

# Create sentiment analysis prompts
sentiment_prompt = """Analyze the sentiment of this customer feedback
and classify it as Positive, Negative, or Neutral.
Also provide a brief reason.

Feedback: {feedback}"""

df = feedback_df.with_columns(
    prompt=pl.format(sentiment_prompt, pl.col('feedback'))
)

# Convert to messages and run inference
df = df.with_columns(
    message=string_to_message("prompt", message_type='user')
)

df = df.with_columns(
    sentiment_analysis=inference_async('message',
                                      provider=Provider.OPENAI,
                                      model='gpt-4o-mini')
)

# Extract key insights
print(df.select(['customer_id', 'feedback', 'sentiment_analysis']))

Vector Similarity Search

Generate embeddings, then find nearest neighbors with HNSW

python
from polar_llama import knn_hnsw, embedding_async, Provider

# Create corpus of documents
corpus = pl.DataFrame({
    "doc": ["AI research", "cooking tips", "machine learning", "recipes"]
}).with_columns(
    embedding=embedding_async(pl.col("doc"), provider=Provider.OPENAI)
)

# Create query
query = pl.DataFrame({
    "query": ["artificial intelligence"]
}).with_columns(
    query_emb=embedding_async(pl.col("query"), provider=Provider.OPENAI),
    corpus_emb=pl.lit([corpus["embedding"].to_list()])
).with_columns(
    neighbors=knn_hnsw(
        pl.col("query_emb"),
        pl.col("corpus_emb").list.first(),
        k=2  # Find 2 nearest neighbors
    )
)

# Get nearest neighbor documents
indices = query["neighbors"][0]
print(corpus[indices]["doc"])  # ['AI research', 'machine learning']

Taxonomy-Based Tagging

Classify documents with reasoning, reflection, and confidence scores

python
import polars as pl
from polar_llama import tag_taxonomy, Provider

# Define your taxonomy
taxonomy = {
    "sentiment": {
        "description": "The emotional tone of the text",
        "values": {
            "positive": "Text expresses positive emotions or favorable opinions",
            "negative": "Text expresses negative emotions or unfavorable opinions",
            "neutral": "Text is factual and objective without clear emotional content"
        }
    },
    "urgency": {
        "description": "How urgent the content is",
        "values": {
            "high": "Requires immediate attention",
            "medium": "Should be addressed soon",
            "low": "Can be addressed at any time"
        }
    }
}

df = pl.DataFrame({
    "id": [1, 2],
    "message": [
        "URGENT: Server is down!",
        "Thanks for your help yesterday."
    ]
})

result = df.with_columns(
    tags=tag_taxonomy(
        pl.col("message"),
        taxonomy,
        provider=Provider.GROQ,
        model="llama-3.3-70b-versatile"
    )
)

# Extract specific values
result.select([
    "message",
    pl.col("tags").struct.field("sentiment").struct.field("value").alias("sentiment"),
    pl.col("tags").struct.field("sentiment").struct.field("confidence").alias("confidence"),
    pl.col("tags").struct.field("urgency").struct.field("value").alias("urgency")
])

Provider Support

Six inference targets — five hosted providers plus on-device MLX

OpenAI

Default model: gpt-4o-mini

python
df = df.with_columns(
    answer=inference_async('prompt',
                          provider=Provider.OPENAI,
                          model='gpt-4o-mini')
)

Anthropic (Claude)

Default model: claude-opus-4-8; supports cache=True for prompt caching

python
df = df.with_columns(
    answer=inference_async('prompt',
                          provider=Provider.ANTHROPIC,
                          model='claude-opus-4-8',
                          system_prompt='You are a helpful assistant.',
                          cache=True)
)

AWS Bedrock

Default model: us.anthropic.claude-haiku-4-5-20251001-v1:0; region resolved from AWS_REGION / AWS_DEFAULT_REGION

python
# Requires AWS credentials configured
df = df.with_columns(
    answer=inference_async('prompt',
                          provider='bedrock',
                          model='us.anthropic.claude-haiku-4-5-20251001-v1:0')
)

Google Gemini

Default model: gemini-2.5-flash; native system_instruction and JSON-schema structured outputs

python
df = df.with_columns(
    answer=inference_async('prompt',
                          provider=Provider.GEMINI,
                          model='gemini-2.5-flash')
)

Groq

Default model: llama-3.3-70b-versatile

python
df = df.with_columns(
    answer=inference_async('prompt',
                          provider=Provider.GROQ,
                          model='llama-3.3-70b-versatile')
)

Local (Apple Silicon / MLX)

No API key, no network — server engine points at a local OpenAI-compatible endpoint, in_process drives mlx-lm directly

python
# engine="server" (default): point at a local server you started yourself
df = df.with_columns(
    answer=pl.col('prompt').llama.inference_local(
        model='mlx-community/gemma-4-e2b-it-4bit',
        engine='server',
        base_url='http://localhost:8080',
    )
)

# engine="in_process": batched generation directly via mlx-lm
df = df.with_columns(
    answer=pl.col('prompt').llama.inference_local(
        model='mlx-community/gemma-4-e2b-it-4bit',
        engine='in_process',
        max_tokens=256,
    )
)

API Surface

Core expressions exported from polar_llama

FunctionPurpose
inference_async(expr, *, provider, model, response_model, cache, system_prompt)Parallel async inference; accepts cache=True/CacheConfig and system_prompt for provider-native prompt caching
inference(expr, *, provider, model, response_model)Synchronous inference (deprecated in favor of inference_async)
inference_messages(expr, *, provider, model, response_model, cache)Multi-turn conversation inference over JSON or List(Struct) message arrays
string_to_message(expr, *, message_type)Convert text to a {role, content} message
combine_messages(*exprs)Merge message columns/arrays into one ordered conversation
tag_taxonomy(expr, taxonomy, *, provider, model)Classify text against a taxonomy with reasoning, reflection, and confidence
embedding_async(expr, *, provider, model)Parallel embedding generation (OpenAI, Gemini, Bedrock)
cosine_similarity / dot_product / euclidean_distance(vec1, vec2)Rust-powered vector similarity metrics
knn_hnsw(query_expr, reference_expr, *, k)Approximate nearest-neighbor search via HNSW
mcp_tools(transport, *, timeout_s)Fetch tool definitions from an MCP server (tools/list)
tools_to_response_model(tools, *, model_name)Build a Pydantic emission schema so the LLM emits structured tool calls
execute_tool_calls(expr, *, transport, executor, tools, concurrency, timeout_s)Run every emitted call of every row in parallel; failures are data
tool_results_to_message(expr, *, role)Render tool results as a message for the synthesis inference pass
Signature / Predict / evaluate / BootstrapFewShot / InstructionOptimizerDSPy-style prompt optimization engine (polar_llama.optimize)
col(...).llama.inference_local(*, model, system, engine, base_url, max_tokens, temperature, top_p, stop)On-device inference on Apple Silicon via mlx-lm
polar_llama.local.make_local_inference_fn(model, *, engine, collapse, max_tokens, ...)Build an inference_fn that backs the optimizer with on-device Gemma 3n

Every expression is also available on the fluent .llama namespace (pl.col("text").llama.inference_async(...), .llama.to_message(...), .llama.embedding(...), .llama.execute_tool_calls(...), and so on).

Advanced Features

Environment Configuration

Set up your API keys and overrides in a .env file:

.envbash
OPENAI_API_KEY=your_openai_key
ANTHROPIC_API_KEY=your_anthropic_key
GEMINI_API_KEY=your_gemini_key
GROQ_API_KEY=your_groq_key
AWS_ACCESS_KEY_ID=your_aws_key
AWS_SECRET_ACCESS_KEY=your_aws_secret
AWS_REGION=us-east-1

# Optional proxy/gateway overrides (0.3.0+)
OPENAI_BASE_URL=https://your-proxy.example.com/v1
ANTHROPIC_BASE_URL=https://your-proxy.example.com

# Bound concurrent in-flight requests per batch (default: 64)
POLAR_LLAMA_MAX_CONCURRENCY=64

# Local MLX backend (Apple Silicon only)
POLAR_LLAMA_LOCAL_COLLAPSE=1        # collapsed prefix prefill for inference_local (0.5.1+)
POLAR_LLAMA_LOCAL_KV_BITS=4         # quantized batched KV cache, 0.5.0+ (mutually exclusive with COLLAPSE)

Prompt Caching

Share a cached system prefix across rows (Anthropic cache_control):

python
from polar_llama import CacheConfig, CacheStrategy

# Simple: enable automatic caching
df = df.with_columns(
    response=inference_async(
        pl.col("prompt"),
        provider=Provider.ANTHROPIC,
        model="claude-opus-4-8",
        system_prompt="You are a careful, concise research assistant.",
        cache=True,   # AUTO strategy, min_tokens=1024, ttl="5m"
    )
)

# Advanced: configure caching behavior explicitly
config = CacheConfig(strategy=CacheStrategy.SYSTEM_PROMPT, ttl="1h")
df = df.with_columns(
    response=inference_async(pl.col("prompt"), cache=config)
)

Testing

Run tests with configured providers:

bash
pip install -r tests/requirements.txt
pytest tests/ -v
cargo test --test model_client_tests -- --nocapture

# Local MLX backend: CPU-safe logic tests only (no GPU, no mlx import)
pytest tests/ -m "not local_gpu"

Changelog Highlights (0.3.0 → 0.5.2)

  • 0.5.2 — fixed in-process local inference on Gemma 3n (mlx-lm #1384 patch applied at load time; a masked ZeroDivisionError in BatchGenerator teardown no longer hides the real error)
  • 0.5.1 — local prompt-tuning bridge (make_local_inference_fn), collapsed-prefill opt-in for inference_local (~2.8× faster tuning schedules), InstructionOptimizer array/string flattening fix
  • 0.5.0 — local MLX inference backend (server + in_process engines), collapsed prefix prefill, batched quantized KV cache, mlx-lm #1384 fix
  • 0.3.0 — tool use / MCP integration, provider-native prompt caching, DSPy-style prompt optimizer, updated default models across all providers, TLS/rustls hardening, Windows wheel fix, ~400 lines of duplicated dispatch removed

Common Use Cases

📊Data Analysis

Process large datasets with AI insights — sentiment analysis, classification, entity extraction with validated structured outputs

🧰Tool-Augmented Pipelines

Let the LLM call databases, internal APIs, or MCP servers at scale — every call, result, and retry is an ordinary dataframe column

🎛️Prompt Tuning at Scale

Bootstrap few-shot demos or search for better instructions against a labeled dataset, entirely offline or on-device

🔒Private, On-Device Inference

Run classification, extraction, or tuning against Gemma models on Apple Silicon with no API keys and no data leaving the machine

Resources
GitHub RepositoryPyPI PackagePolars Documentation

Licensed under MIT.

Questions or issues? Open one on GitHub.