headroom/examples/streaming_example.py
chopratejas 9c7d4512d6 Initial commit: Headroom SDK - LLM context optimization toolkit
A comprehensive SDK for optimizing LLM context windows, reducing token
usage while preserving critical information for AI agents.

Core Features:
- SmartCrusher: Statistical compression of tool outputs (70-85% reduction)
- CacheAligner: Prefix optimization for prompt cache hits
- RollingWindow: Intelligent context window management
- BM25/Hybrid relevance scoring for smart item selection

Integrations:
- OpenAI and Anthropic provider support
- LangChain integration (ChatModel, Callbacks, Runnable)
- MCP (Model Context Protocol) integration for tool compression

Test Coverage:
- 372 tests passing across all modules
- 35 performance benchmarks
- Real-world agent evaluations with 88% token savings

Key Components:
- headroom/transforms/: Core compression transforms
- headroom/providers/: OpenAI and Anthropic support
- headroom/integrations/: LangChain and MCP integrations
- headroom/relevance/: BM25 and hybrid scoring
- headroom/pricing/: Model pricing registry
- benchmarks/: Performance benchmark suite
- examples/: Usage examples and demos
2026-01-06 23:16:58 -08:00

71 lines
1.6 KiB
Python

#!/usr/bin/env python3
"""
Streaming example for Headroom SDK.
This example shows how to use Headroom with streaming responses.
"""
import os
import tempfile
from dotenv import load_dotenv
from openai import OpenAI
from headroom import HeadroomClient, OpenAIProvider
# Load API key from .env.local
load_dotenv(".env.local")
# Create base OpenAI client
base_client = OpenAI(api_key=os.environ.get("OPENAI_API_KEY", "sk-..."))
# Create provider for OpenAI models
provider = OpenAIProvider()
# Use temp directory for database
db_path = os.path.join(tempfile.gettempdir(), "headroom_streaming.db")
# Wrap with Headroom
client = HeadroomClient(
original_client=base_client,
provider=provider,
store_url=f"sqlite:///{db_path}",
default_mode="optimize",
)
def stream_example():
"""Example of streaming with Headroom."""
print("Streaming response:")
print("-" * 40)
messages = [
{
"role": "system",
"content": "You are a helpful assistant. Current Date: 2024-01-15. Be concise.",
},
{"role": "user", "content": "Count from 1 to 5 slowly."},
]
# Stream with optimization
stream = client.chat.completions.create(
model="gpt-4o-mini",
messages=messages,
stream=True,
headroom_mode="optimize",
max_tokens=100,
)
# Iterate over chunks
for chunk in stream:
if chunk.choices[0].delta.content:
print(chunk.choices[0].delta.content, end="", flush=True)
print()
print("-" * 40)
print("Stream complete!")
if __name__ == "__main__":
stream_example()
client.close()