langfuse-instrumented-dspy-lm 0.1.1

Custom DSPy LM class for token usage tracking with Langfuse

langfuse-instrumented-dspy-lm

Track token usage and costs in your DSPy + Langfuse applications.

Using DSPy with Langfuse but not seeing token counts or costs in your traces? You're not alone. The standard integration doesn't capture this data. This library fixes that with a drop-in replacement for dspy.LM.

Quick Start

Follow the Langfuse DSPy integration guide, but use LangfuseInstrumentedLm instead of dspy.LM in the final step.

Step 1: Install Packages

pip install langfuse dspy openinference-instrumentation-dspy langfuse-instrumented-dspy-lm

Step 2: Configure Environment

import os

os.environ["LANGFUSE_PUBLIC_KEY"] = "pk-lf-..."
os.environ["LANGFUSE_SECRET_KEY"] = "sk-lf-..."
os.environ["LANGFUSE_HOST"] = "https://cloud.langfuse.com"  # or your self-hosted URL

# Your LLM provider API key
os.environ["OPENAI_API_KEY"] = "sk-..."  # or ANTHROPIC_API_KEY, etc.

Step 3: Initialize Langfuse and Enable DSPy Tracing

Initialize the Langfuse client before enabling the DSPy instrumentor. This ensures Langfuse is ready to receive traces.

from langfuse import get_client
from openinference.instrumentation.dspy import DSPyInstrumentor

langfuse = get_client()
langfuse.auth_check()  # verify connection

DSPyInstrumentor().instrument()

Step 4: Configure DSPy with Instrumented LM

import dspy
from langfuse_instrumented_dspy_lm import LangfuseInstrumentedLm

# Use LangfuseInstrumentedLm instead of dspy.LM
lm = LangfuseInstrumentedLm("openai/gpt-4o-mini", max_tokens=1000)
dspy.configure(lm=lm)

Step 5: Use DSPy as Normal

qa = dspy.Predict("question -> answer")
result = qa(question="What is the capital of France?")
print(result.answer)

Step 6: Flush Traces Before Exit

In scripts and short-lived processes, flush traces to ensure they are delivered before the process exits:

from opentelemetry import trace as otel_trace

# Flush OpenTelemetry spans
provider = otel_trace.get_tracer_provider()
if hasattr(provider, "force_flush"):
    provider.force_flush()

# Flush and shutdown Langfuse
langfuse.flush()
langfuse.shutdown()

Note: In long-running applications (e.g., web servers), traces are sent automatically in the background and you typically don't need to flush manually.

That's it! Your Langfuse traces will now include token usage and cost data.

Complete Example

import os
import dspy
from langfuse import get_client
from openinference.instrumentation.dspy import DSPyInstrumentor
from opentelemetry import trace as otel_trace
from langfuse_instrumented_dspy_lm import LangfuseInstrumentedLm

# Configure environment
os.environ["LANGFUSE_PUBLIC_KEY"] = "pk-lf-..."
os.environ["LANGFUSE_SECRET_KEY"] = "sk-lf-..."
os.environ["LANGFUSE_HOST"] = "https://cloud.langfuse.com"
os.environ["OPENAI_API_KEY"] = "sk-..."

# Initialize Langfuse client before instrumenting
langfuse = get_client()
langfuse.auth_check()

# Enable tracing
DSPyInstrumentor().instrument()

# Configure DSPy with instrumented LM
dspy.configure(
    lm=LangfuseInstrumentedLm(
        "openai/gpt-4o-mini",
        max_tokens=1000,
        temperature=0.7,
    )
)

# Use DSPy - token usage is now tracked!
qa = dspy.Predict("question -> answer")
result = qa(question="What is the capital of France?")
print(result.answer)

# Flush traces before exit
provider = otel_trace.get_tracer_provider()
if hasattr(provider, "force_flush"):
    provider.force_flush()
langfuse.flush()
langfuse.shutdown()

Using with @observe() Decorator

Combine with Langfuse's decorator for hierarchical tracing:

from langfuse.decorators import langfuse_context, observe

@observe()
def answer_question(question: str) -> str:
    qa = dspy.Predict("question -> answer")
    result = qa(question=question)
    return result.answer

answer = answer_question("What is the capital of France?")

# Flush to ensure traces are sent
langfuse_context.flush()

Supported Models

Any model supported by LiteLLM works:

# OpenAI
LangfuseInstrumentedLm("openai/gpt-4o-mini")

# Anthropic
LangfuseInstrumentedLm("anthropic/claude-3-5-sonnet-20241022")

# Google
LangfuseInstrumentedLm("gemini/gemini-2.0-flash")

# Azure OpenAI
LangfuseInstrumentedLm("azure/gpt-4o-mini")

Streaming Support

Token usage is automatically tracked when using dspy.streamify() for streaming:

import asyncio
import dspy
from langfuse_instrumented_dspy_lm import LangfuseInstrumentedLm

lm = LangfuseInstrumentedLm("openai/gpt-4o-mini", max_tokens=1000)
dspy.configure(lm=lm)

predict = dspy.Predict("question -> answer")
stream_predict = dspy.streamify(
    predict,
    stream_listeners=[dspy.streaming.StreamListener(signature_field_name="answer")],
)

async def main():
    output = stream_predict(question="What is the capital of France?")
    async for chunk in output:
        print(chunk)

asyncio.run(main())

Captured Metrics

Metric	Description
Prompt tokens	Input token count
Completion tokens	Output token count
Total tokens	Combined token count
Cost	Total cost (when available)
Reasoning tokens	Thinking/reasoning tokens (o1, etc.)
Cached tokens	Prompt cache hits

---

For Developers

Why This Library Exists

When using Langfuse with DSPy via the standard integration, token usage and cost per LLM call are not captured. This happens because:

1. Langfuse relies on openinference-instrumentation-dspy for auto-instrumenting DSPy
2. The instrumentation wraps dspy.LM.__call__, which only returns parsed output (not token usage)
3. Token usage is available in dspy.LM.forward() response, but the wrapper never sees it

LangfuseInstrumentedLm solves this by overriding forward() and aforward() to:

1. Call the parent method to get the full LiteLLM response
2. Extract token usage from the response
3. Set OpenInference span attributes on the current span
4. Return the response unchanged

DSPy Application
       |
       v
DSPyInstrumentor wrapper (creates span)
       |
       v
LangfuseInstrumentedLm.__call__() [inherited]
       |
       v
LangfuseInstrumentedLm.forward() <-- captures tokens here
       |
       +-- super().forward() -> LiteLLM ModelResponse
       +-- _set_span_attributes(response)
       +-- return response

API Reference

class LangfuseInstrumentedLm(dspy.LM):
    """Drop-in replacement for dspy.LM with token tracking."""

Constructor accepts all dspy.LM / LiteLLM parameters:

lm = LangfuseInstrumentedLm(
    model="openai/gpt-4o-mini",
    max_tokens=1000,
    temperature=0.7,
    cache=False,  # Disable to get fresh token counts each call
)

Development Setup

git clone https://github.com/unravel-team/langfuse-instrumented-dspy-lm.git
cd langfuse-instrumented-dspy-lm
uv sync --group dev

Commands

# Lint
uv run ruff check src/
uv run ruff format src/

# Build
uv build

# Test import
uv run python -c "from langfuse_instrumented_dspy_lm import LangfuseInstrumentedLm; print('OK')"

Publishing

# Test with TestPyPI
uv build
uv publish --index-url https://test.pypi.org/legacy/ --token <TOKEN>

# Publish to PyPI
uv build
uv publish --token <PYPI_TOKEN>

Contributing

Contributions welcome! Please open an issue or pull request.

License

MIT - see LICENSE