Wurun

Async OpenAI API wrapper optimized for Jupyter notebooks with connection pooling, retry logic, and batch processing.

Features

• HTTP/2 Connection Pooling - Shared client for efficient API calls
• Robust Retry Logic - Exponential backoff, and fail-fast on errors that retrying cannot fix
• Batch Processing - Concurrent API calls with semaphore control
• Jupyter Optimized - Clean notebook output and error handling
• Zero Configuration - Simple setup, works with OpenAI and Azure OpenAI

Installation

# From PyPI
pip install wurun

# With pandas DataFrame support
pip install "wurun[dataframe]"

# Development (testing and linting tools)
pip install -e ".[dev,dataframe]"

Quick Start

from wurun import Wurun

# Setup once per kernel
await Wurun.setup(
    endpoint="https://api.openai.com/v1",
    api_key="your-api-key",
    deployment_name="gpt-3.5-turbo"
)

# Single question
messages = [{"role": "user", "content": "Explain asyncio"}]
answer = await Wurun.ask(messages)
print(answer)

# Custom parameters
answer = await Wurun.ask(messages, max_tokens=512, temperature=0.7)

# Batch processing
questions = [
    [{"role": "user", "content": "What is Python?"}],
    [{"role": "user", "content": "What is JavaScript?"}]
]
answers = await Wurun.run_gather(questions, concurrency=2)

# Cleanup
await Wurun.close()

DataFrame Processing

import pandas as pd
df = pd.DataFrame({
    "id": [1, 2, 3],
    "messages": [
        [{"role": "user", "content": "What is Python?"}],
        [{"role": "user", "content": "What is JavaScript?"}],
        [{"role": "user", "content": "What is Go?"}]
    ]
})
df["answer"] = await Wurun.run_dataframe(df, "messages", concurrency=2)

Results are returned in the same order as the DataFrame rows, so they can be assigned straight back to a column.

🧪 Practical Demo on Kaggle

API Reference

Setup

• Wurun.setup() - Initialize client (call once per kernel)
• Wurun.close() - Clean up resources

Single Calls

• Wurun.ask() - Single API call with retry logic
• return_meta=True - Return (answer, meta) instead of just the answer
• max_tokens=1024 - Maximum tokens in response (default: 1024)
• temperature=0 - Response randomness (default: 0)
• timeout=None - Per-request timeout; None inherits the setup() timeout

Batch Processing

• Wurun.run_gather() - Preserve input order
• Wurun.run_as_completed() - Yield (index, answer) in completion order
• concurrency parameter controls parallel requests
• max_tokens and temperature available on all batch methods

DataFrame Processing

• Wurun.run_dataframe() - Process messages from DataFrame column

Notebook Helpers

• Wurun.print_qna_ordered() - Pretty print Q&A format
• Wurun.print_as_ready() - Print results as they complete

Configuration

await Wurun.setup(
    endpoint="https://api.openai.com/v1",
    api_key="your-key",
    deployment_name="gpt-3.5-turbo",
    timeout=60.0,        # per-request cap in seconds
    max_connections=32,
    max_keepalive=16,
    http2=True,
    max_retries=2        # SDK-level retries, see the note below
)

setup() is safe to call again from the same kernel. The connection pool is reused when timeout, max_connections, max_keepalive and http2 are unchanged, and rebuilt when any of them differ - so re-running your setup cell with a new timeout actually takes effect.

Retry budget: setup(max_retries=...) is applied by the OpenAI SDK underneath ask(attempts=...). The worst case is attempts * (max_retries + 1) HTTP requests (by default 5 * 3 = 15). Lower one of the two if you want a tighter bound.

Error Handling

ask() never raises on API failures - it returns an "[ERROR] ..." string so one bad row cannot abort a whole batch. Use return_meta=True to detect failures reliably instead of matching on the string:

answer, meta = await Wurun.ask(messages, return_meta=True)
if meta["error"]:
    print(f"failed with {meta['error_type']} after {meta['retries']} retries")
else:
    print(f"Latency: {meta['latency']:.2f}s, Retries: {meta['retries']}")

The meta dict contains:

Only transient failures are retried: rate limits, connection/timeout errors, HTTP 408/429 and any 5xx. Non-retryable statuses such as 400, 401, 403 and 404 return immediately rather than burning the retry budget on a request that cannot succeed.

# Custom retry settings
result = await Wurun.ask(
    messages,
    attempts=3,
    initial_backoff=1.0,
    max_backoff=10.0,
    max_tokens=512,
    temperature=0.7
)

Development

# Install dev dependencies
pip install -e ".[dev,dataframe]"

# Run everything CI runs (format, lint, types, tests)
make check

# Or just the tests
pytest test_wurun.py -v

Release Process

1. Create PR: Make changes and create pull request to main
2. Auto Draft: Release Drafter automatically creates/updates draft release
3. Publish Release: Go to GitHub Releases, edit draft, and publish
4. Auto Deploy: Publishing triggers automatic PyPI deployment

Manual Version Update

# Updates pyproject.toml, wurun.py and CITATION.cff together
python scripts/update_version.py 1.2.3

License

Apache License 2.0