Python SDK

# pip
pip install bitfab-py

# Poetry
poetry add bitfab-py

# uv
uv add bitfab-py

import os
from bitfab import Bitfab

bitfab = Bitfab(api_key=os.environ["BITFAB_API_KEY"])

Bitfab(api_key="...")

# Omit api_key entirely: the SDK reads BITFAB_API_KEY from the environment
Bitfab()

# Disable tracing (functions still execute, but no spans are sent)
Bitfab(api_key="...", enabled=False)

# Pass a callable to defer resolution explicitly (resolved at first use):
Bitfab(api_key=lambda: os.environ.get("BITFAB_API_KEY"))

# Raises on the first traced call if no key resolves, instead of disabling quietly
Bitfab(api_key=os.environ.get("BITFAB_API_KEY"), strict=True)

order_service = bitfab.get_function("order-processing")

@order_service.span()
def process_order(order_id: str) -> dict:
    return {"order_id": order_id}

@order_service.span()
def validate_order(order_id: str) -> dict:
    return {"valid": True}

# lib/bitfab_client.py -- single source of truth
import os
from bitfab import Bitfab
bitfab = Bitfab(api_key=os.environ["BITFAB_API_KEY"])
order_service = bitfab.get_function("order-processing")

# services/process_order.py
from lib.bitfab_client import order_service

@order_service.span()
def process_order(order_id: str) -> dict:
    return {"order_id": order_id}

# services/validate_order.py
from lib.bitfab_client import order_service

@order_service.span()
def validate_order(order_id: str) -> dict:
    return {"valid": True}

@bitfab.span("one-off-operation")
def standalone_task() -> str:
    return "done"

@bitfab.span("outer", type="agent")
def outer():
    inner()  # Becomes a child of "outer"

@bitfab.span("inner", type="function")
def inner():
    pass

SpanType = Literal[
    "llm",        # LLM calls
    "agent",      # Agent workflows
    "function",   # Function calls
    "guardrail",  # Safety checks
    "handoff",    # Human handoffs
    "custom"      # Default
]

# Function name is automatically captured as span name
@bitfab.span("order-processing")
def process_order(order_id: str) -> dict:
    return {"order_id": order_id}
# Span name: "process_order"

# Override with name option
@bitfab.span("order-processing", name="OrderProcessor")
def process_order(order_id: str) -> dict:
    return {"order_id": order_id}
# Span name: "OrderProcessor"

# Set span type
@bitfab.span("safety-check", type="guardrail")
def check_content(content: str) -> dict:
    return {"safe": True}

from bitfab import finalizers

@bitfab.span("chat", type="llm", finalize=finalizers.openai_chunks)
async def chat(messages):
    stream = await client.chat.completions.create(
        model="gpt-4o", messages=messages, stream=True
    )
    async for chunk in stream:
        yield chunk  # caller still receives every chunk

# The span records { text, finish_reason, usage, tool_calls } in the background.

@bitfab.span("chat", type="llm", finalize=lambda chunks: {"text": "".join(
    c.choices[0].delta.content or "" for c in chunks
)})
async def chat(messages):
    ...
    async for chunk in stream:
        yield chunk

from bitfab import get_current_span

@bitfab.span("order-processing", type="function")
def process_order(order_id: str) -> dict:
    user_id = get_current_user()
    get_current_span().add_context({"user_id": user_id, "order_id": order_id})
    return {"order_id": order_id, "status": "completed"}

get_current_span().add_context({"user_id": "u-123"})
get_current_span().add_context({"request_id": "req-789"})
# Result: contexts: [{"user_id": "u-123"}, {"request_id": "req-789"}]

from bitfab import get_current_span

@bitfab.span("classification", type="llm")
def classify_text(text: str) -> str:
    prompt = f"Classify the following text: {text}"
    get_current_span().set_prompt(prompt)
    result = llm.complete(prompt)
    return result

from bitfab import get_current_trace

@bitfab.span("order-processing", type="function")
def process_order(order_id: str) -> dict:
    trace = get_current_trace()

    # Set session ID (stored as database column, filterable in dashboard)
    trace.set_session_id("session-123")

    # Set trace metadata (stored in raw trace data)
    trace.set_metadata({"region": "us-west-2", "environment": "production"})

    # Add context entries (stored as key-value pairs, accumulates across calls)
    trace.add_context({"workflow": "checkout-flow", "batch_id": "batch-2024-01"})

    return {"order_id": order_id, "status": "completed"}

from bitfab import get_current_trace

@bitfab.span("order-processing", type="function")
def process_order(order_id: str) -> dict:
    if is_health_check(order_id):
        get_current_trace().drop()

    return {"order_id": order_id, "status": "completed"}

trace = client.get_trace(trace_id)
trace.add_context({"refund_status": "approved"})
trace.set_metadata({"region": "us-west"})
trace.set_session_id("session_xyz")

# Optional: wait for confirmation
thread = trace.set_metadata({"status": "complete"})
if thread:
    thread.join(timeout=5.0)

span = client.get_trace_span(trace_id, name="GenerateAnswer")
first = client.get_trace_span(
    trace_id, name="GenerateAnswer", occurrence="first"
)
exact = client.get_trace_span(trace_id, id=span_id)

@bitfab.span("risky-service")
def risky():
    raise ValueError("error")

try:
    risky()
except ValueError:
    pass
# Span records error and timing

from bitfab import flush_traces

flush_traces(timeout=30.0)  # Default: 30s

@bitfab.span("my-function-key")
def my_function(text: str) -> dict:
    return {"processed": text.upper()}

result = bitfab.replay(my_function, limit=5)

# Or replay specific traces by ID
result = bitfab.replay(my_function, trace_ids=["trace-abc", "trace-def"])

print(f"Test Run: {result['test_run_url']}")
for item in result["items"]:
    print(f"  Input: {item['input']}")
    print(f"  Result: {item['result']}")
    print(f"  Original: {item['original_output']}")
    print(f"  Duration (ms): {item['duration_ms']}")
    print(f"  Tokens: {item['tokens']}")       # {"input", "output", "cached", "total"} or None
    print(f"  Model: {item['model']}")

{
    "items": [
        {
            "input": [...],             # The inputs passed to fn
            "result": ...,              # What fn returned
            "original_output": ...,     # What the original trace produced
            "error": None | str,        # Error message if fn raised
            "duration_ms": int | None,  # Original trace duration in ms
            "tokens": {                 # Original trace token usage, or None
                "input": int | None,
                "output": int | None,
                "cached": int | None,
                "total": int | None,
            } | None,
            "model": str | None,        # Original model name, or None
            "trace_id": str | None,     # Server trace ID for the replayed execution
            "db_snapshot_ref": dict | None,  # The source trace's snapshot pin, if any
        }
    ],
    "test_run_id": "...",
    "test_run_url": "..."
}

# scripts/replay.py
from my_app.agent import graph              # the compiled LangGraph graph
from my_app.bitfab_client import bitfab     # same client as instrumentation

handler = bitfab.get_langgraph_callback_handler("my-agent")  # same key


def replay_my_agent(state):
    config = {"callbacks": [handler], "configurable": build_replay_config()}
    return graph.invoke(state, config=config)


result = bitfab.replay("my-agent", replay_my_agent, limit=10)

@bitfab.span("fetch-article-from-db", mock_on_replay=True)
def fetch_article_from_db(article_id: str) -> Article:
    return db.articles.find_by_id(article_id)


@bitfab.span("summarize-article")
def summarize_article(article: Article) -> Summary:
    # Real summarization, no flag -- this is what we're iterating on.
    return Summary(...)


@bitfab.span("process-article")
def process_article(article_id: str) -> Summary:
    return summarize_article(fetch_article_from_db(article_id))


# During replay, fetch-article-from-db returns its recorded output;
# summarize-article runs real so you can iterate on it.
result = bitfab.replay(process_article, limit=10)

from bitfab import MockOverride

result = client.replay(
    process_article,
    mock="none",  # run everything real...
    mock_override=MockOverride(
        # ...except this span, which gets the value you supply
        match=lambda node: node.trace_function_key == "fetch-article-from-db",
        value={"id": "fixed", "title": "Fixed title"},  # flat value
    ),
)

value=lambda ctx: {**ctx.get_original_output(), "score": 1}

client.register_mock_override(
    MockOverride(match=lambda node: node.type == "llm", value={"label": "refund"})
)
# Ordered form (equivalent): client.register_mock_override(match, value)
client.clear_mock_overrides()

# Recorded as (user_id, limit); current signature is (opts: dict).
def adapt(args, kwargs, ctx):
    user_id, limit = args
    return [{"user_id": user_id, "limit": limit}], {}


result = bitfab.replay(my_function, adapt_inputs=adapt)

# scripts/replay_adapters/extraction.py
def adapt_inputs(args, kwargs, ctx):
    user_id, limit = args
    return [{"user_id": user_id, "limit": limit}], {}

# scripts/replay.py
from replay_adapters.extraction import adapt_inputs

bitfab.replay(my_function, limit=limit, adapt_inputs=adapt_inputs)

with open("src/foo.py") as f:
    before = f.read()

# ...edit src/foo.py...

with open("src/foo.py") as f:
    after = f.read()

result = bitfab.replay(
    my_function,
    code_change_description="fix off-by-one in retry logic",
    code_change_files=[{"path": "src/foo.py", "before": before, "after": after}],
)

result = bitfab.replay(my_function, limit=limit)

# Human-readable summary goes to stderr, so stdout stays pure JSON.
print(f"Test run: {result['test_run_url']}", file=sys.stderr)
print(f"Items:    {len(result['items'])}", file=sys.stderr)

# Then: full structured dump, ready for json.loads.
print(json.dumps(result, indent=2, default=str))

#!/usr/bin/env python3
"""
Replay production traces through instrumented functions.

Uses bitfab.replay() to fetch real traces and re-run them through
the current code, creating a test run for side-by-side comparison.

Usage:
    python scripts/replay.py <pipeline>
    python scripts/replay.py <pipeline> --limit 20
    python scripts/replay.py <pipeline> --trace-ids id1,id2
"""
import argparse
import json
import sys
from dotenv import load_dotenv
from bitfab import report_replay_progress
from lib.bitfab_client import bitfab
from services.extraction import extract_memories
from services.search import search_documents

load_dotenv()

FUNCTIONS = {
    "extraction": "my-extraction-pipeline",
    "search": "my-search-pipeline",
}


# Each pipeline gets its own replay function. If a signature drifts after traces
# were captured, add an input adapter (see the "Adapting inputs" section): write
# scripts/replay_adapters/<name>.py with an `adapt_inputs(args, kwargs, ctx)`
# function, import it, and pass it here, e.g.
#   from replay_adapters.extraction import adapt_inputs
#   return bitfab.replay(fn, limit=limit, trace_ids=trace_ids, adapt_inputs=adapt_inputs)

# limit is ignored when trace_ids is passed -- an explicit ID list already
# determines how many traces replay.

def replay_extraction(limit: int, trace_ids: list[str] | None):
    def fn(conversation: str, existing_items: list):
        return extract_memories(conversation, existing_items)
    if trace_ids:
        return bitfab.replay(fn, trace_ids=trace_ids, on_progress=report_replay_progress)
    return bitfab.replay(fn, limit=limit, on_progress=report_replay_progress)


def replay_search(limit: int, trace_ids: list[str] | None):
    def fn(query: str, opts: dict):
        return search_documents(query, user_id=opts["user_id"], limit=opts.get("limit", 10))
    if trace_ids:
        return bitfab.replay(fn, trace_ids=trace_ids, on_progress=report_replay_progress)
    return bitfab.replay(fn, limit=limit, on_progress=report_replay_progress)


REPLAY_FNS = {
    "extraction": replay_extraction,
    "search": replay_search,
}


def main():
    parser = argparse.ArgumentParser(description="Replay production traces")
    parser.add_argument("pipeline", choices=FUNCTIONS.keys())
    parser.add_argument("--limit", type=int, default=10)
    parser.add_argument("--trace-ids", type=str)
    args = parser.parse_args()

    trace_ids = [tid.strip() for tid in args.trace_ids.split(",")] if args.trace_ids else None
    function_key = FUNCTIONS[args.pipeline]

    # Human-readable output goes to stderr; stdout carries only the ReplayResult JSON.
    print(f"[replay] Replaying {len(trace_ids) if trace_ids else args.limit} traces from \"{function_key}\"...\n", file=sys.stderr)

    result = REPLAY_FNS[args.pipeline](args.limit, trace_ids)
    print(f"Test run: {result['test_run_url']}\n", file=sys.stderr)

    changed = same = errors = 0
    for item in result["items"]:
        raw_input = item.get("input") or []
        label = str(raw_input[0])[:80] if raw_input else "unknown"

        if item["error"]:
            print(f'  ✗ "{label}"', file=sys.stderr)
            print(f"    Error: {item['error']}", file=sys.stderr)
            errors += 1
        else:
            orig = item["original_output"]
            new = item["result"]
            orig_str = orig if isinstance(orig, str) else json.dumps(orig, default=str)
            new_str = new if isinstance(new, str) else json.dumps(new, default=str)
            is_same = orig_str == new_str
            marker = "=" if is_same else "Δ"

            print(f'  {marker} "{label}"', file=sys.stderr)
            print(f"    Original: {orig_str}", file=sys.stderr)
            print(f"    New:      {new_str}", file=sys.stderr)

            if is_same:
                same += 1
            else:
                changed += 1

    print("\n─── Summary ───", file=sys.stderr)
    print(f"  Pipeline: {args.pipeline}", file=sys.stderr)
    print(f"  Replayed: {len(result['items'])}", file=sys.stderr)
    print(f"  Same:     {same}", file=sys.stderr)
    print(f"  Changed:  {changed}", file=sys.stderr)
    if errors > 0:
        print(f"  Errors:   {errors}", file=sys.stderr)
    print(f"\n  {result['test_run_url']}", file=sys.stderr)

    # Full ReplayResult as JSON to stdout (the Replay Output Contract).
    print(json.dumps(result, indent=2, default=str))


if __name__ == "__main__":
    main()

Bitfab(
    api_key: str,                    # Required
    service_url: str | None = None,  # Default: https://bitfab.ai
    env_vars: dict[str, str] | None = None,  # For local function execution
    enabled: bool = True,            # Enable/disable tracing
    baml_client: Any = None          # Generated BAML client (for wrap_baml)
)