pumpsdk 0.3.0

Drop-in OpenAI-compatible Python SDK for the Pump LLM gateway

pump Python SDK

A drop-in replacement for the official OpenAI Python SDK for pump — a multi-tenant LLM gateway that exposes an OpenAI-compatible API.

Keep your existing OpenAI request code exactly as-is. The only things that change are the import, the base_url, and using your pump key (pk-...) as the API key.

Installation

pip install pumpsdk

The distribution is published as pumpsdk; the import package is pump.

Quickstart — OpenAI drop-in

# before:
# from openai import OpenAI
# client = OpenAI()

from pump.openai import OpenAI

client = OpenAI(
    base_url="https://api.pump.co/ai/v1",
    api_key="pk-...",  # your pump key (or set PUMP_API_KEY)
)

resp = client.chat.completions.create(
    model="gpt-4o-mini",
    messages=[{"role": "user", "content": "Hello!"}],
)

print(resp.choices[0].message.content)

That's the whole migration. Every standard parameter (temperature, tools, tool_choice, response_format, seed, stream, ...) is forwarded to the gateway unchanged.

Pump is the same client under a more obvious name — from pump import Pump and from pump.openai import OpenAI are interchangeable.

API key

Pass api_key="pk-..." explicitly, or omit it and the SDK reads PUMP_API_KEY from the environment:

export PUMP_API_KEY="pk-..."

from pump.openai import OpenAI

client = OpenAI()  # base_url defaults to https://api.pump.co/ai/v1

Responses & embeddings

client.responses.create(model="gpt-4o-mini", input="Write a haiku about pumps.")

client.embeddings.create(model="text-embedding-3-small", input="hello world")

Streaming

The streamed object is iterable directly — no context manager required:

stream = client.chat.completions.create(
    model="gpt-4o-mini",
    messages=[{"role": "user", "content": "Count to 5"}],
    stream=True,
)
for chunk in stream:
    delta = chunk.choices[0].delta
    if getattr(delta, "content", None):
        print(delta.content, end="", flush=True)

You can also use it as a context manager to guarantee the connection closes:

with client.chat.completions.create(..., stream=True) as stream:
    for chunk in stream:
        ...

Async

import asyncio
from pump.openai import AsyncOpenAI

async def main():
    client = AsyncOpenAI(api_key="pk-...")

    resp = await client.chat.completions.create(
        model="gpt-4o-mini",
        messages=[{"role": "user", "content": "Hello!"}],
    )
    print(resp.choices[0].message.content)

    # async streaming
    stream = await client.chat.completions.create(
        model="gpt-4o-mini",
        messages=[{"role": "user", "content": "Count to 5"}],
        stream=True,
    )
    async for chunk in stream:
        ...

    await client.aclose()

asyncio.run(main())

Anthropic drop-in

The same two-line migration works for the official Anthropic SDK:

# before:
# from anthropic import Anthropic

from pump.anthropic import Anthropic

client = Anthropic(api_key="pk-...")  # or set PUMP_API_KEY

msg = client.messages.create(
    model="claude-3-5-haiku-latest",
    max_tokens=256,
    messages=[{"role": "user", "content": "Hello!"}],
)
print(msg.content[0].text)

client.messages.count_tokens(...), client.models.list() and client.models.retrieve(id) work too, and Anthropic's native metadata body param ({"user_id": ...}) is forwarded upstream exactly like the official SDK. AsyncAnthropic is the async counterpart.

Any provider — the unified surface

For providers without a dedicated shim (or when you want one code path across all of them), use the pump-native client with a provider/model slug — the same convention OpenRouter, LiteLLM and Vercel AI Gateway use. The gateway splits the slug, uses a BYOK credential selected on your pump key when one covers that provider, otherwise falls back to Pump-managed provider credentials, and forwards the bare model name upstream:

from pump.pump import Pump

client = Pump(api_key="pk-...")

resp = client.chat.completions.create(
    model="cloudflare/@cf/meta/llama-3.1-8b-instruct",  # Cloudflare Workers AI
    messages=[{"role": "user", "content": "Hello!"}],
)

# bare model names keep working too — the gateway routes them by prefix/path:
#   model="gpt-4o-mini"             (openai)
#   model="claude-3-5-haiku-latest" (anthropic)

Reaching every endpoint

Resources cover the common surfaces, and every client also exposes raw verb methods for any other path the gateway passes through:

client.post("images/generations", model="gpt-image-1", prompt="a pump")
client.get("files", purpose="batch")   # kwargs become query params on GET
client.delete("files/file-abc")

Attaching metadata & tags (pump extension)

pump lets you attach observability metadata and tags to any request. These are not part of the provider request body — sending custom fields there would be rejected by the upstream provider. Instead, pass them as explicit kwargs and the SDK sends them as gateway headers (x-pump-metadata, x-pump-tags), which the gateway reads and strips before forwarding the request upstream.

resp = client.chat.completions.create(
    model="gpt-4o-mini",
    messages=[{"role": "user", "content": "Hello!"}],
    pump_metadata={"user_id": "u_123", "feature": "support-bot"},
    pump_tags=["prod", "support"],
)

On the OpenAI-flavored clients, metadata=... / tags=... are still accepted as backwards-compatible aliases for the headers. On the Anthropic client, metadata is never hijacked — it's Anthropic's own body field and is forwarded upstream; use pump_metadata for pump tracking.

Pump cache controls

Pump cache is explicit: pass pump_cache=True to opt into cache reads/writes. The cache is always isolated by company and provider. You can narrow it further with the Pump key or custom tags:

resp = client.chat.completions.create(
    model="gpt-4o-mini",
    messages=[{"role": "user", "content": "Answer from the FAQ."}],
    pump_cache=True,
    pump_cache_scope="key",      # "company" (default) or "key"
    pump_cache_tags=["faq"],     # isolate this call site's cache
)

pump_cache="read-only" reads but does not store misses; pump_cache="write-only" stores misses without serving hits. The same kwargs work on pump.anthropic.

You can also pass arbitrary headers via extra_headers:

client.chat.completions.create(
    model="gpt-4o-mini",
    messages=[...],
    extra_headers={"x-pump-trace-id": "abc123"},
)

Note: OpenAI's own body-level metadata field (for stored completions) is a different thing. If you need it, pass it through extra_body={"metadata": {...}}.

Response objects

Responses support both attribute access and dict access, and expose .model_dump() to get a plain dict:

resp = client.chat.completions.create(...)

resp.choices[0].message.content        # attribute access
resp["choices"][0]["message"]["content"]  # dict access
resp.model_dump()                       # plain dict

Error handling

The SDK raises a single error type, pump.PumpError. For HTTP failures it carries the status_code (and response) so you can branch on it:

from pump import PumpError

try:
    client.chat.completions.create(...)
except PumpError as e:
    if e.status_code == 429:
        ...  # rate limited
    elif e.status_code == 401:
        ...  # bad API key
    else:
        raise

Adding more providers

The SDK is built on a shared transport, so a new provider needs no SDK release at all if its API is OpenAI-compatible — it's reachable immediately through the unified surface with a provider/model slug once the gateway registers it. Providers with their own schema (like Anthropic) get a thin shim module that reuses the transport and only sets routing headers and the resource surface.

Development

pip install -e ".[dev]"
pytest

Tests use httpx.MockTransport and never call the real API.

License

MIT