Skip to content

Commit b213390

Browse files
jhayniejhaynie
authored
feat: add SSE streaming support with billing extraction (#7)
* feat: add SSE streaming support with billing extraction - Add streaming SSE parser for OpenAI and Anthropic responses - Implement StreamingResponseExtractor interface for all providers - Use http.ResponseController for efficient flushing - Add simdjson-go for fast JSON parsing of usage data - Auto-inject stream_options.include_usage for billing - Extract token usage from streaming responses for cost calculation - Add comprehensive unit tests for streaming functionality * docs: add streaming support documentation * fix: use exclusion list for stream_options injection, add Anthropic streaming tests Replace whitelist-based provider switch with exclusion list approach for stream_options injection. Providers with native streaming usage reporting (Anthropic, Bedrock, Google AI) are excluded; all others automatically get stream_options injected. This is more extensible for new providers. Also adds comprehensive Anthropic streaming tests: - TestParseAnthropicSSEEvent - TestExtractUsageFromAnthropicEvent - TestAnthropicSSEParser (full event stream) - TestAutoRouter_AnthropicStreamingNoStreamOptions * fix: improve test assertions in streaming_test.go - Add error checking for ParseAnthropicSSEEvent calls to prevent panics - Require strict nil result for expected-nil test cases - Replace no-op _ = chunk.ID with real assertions for ID and content * fix: address PR review findings Critical fixes: - ForwardStreaming now wraps upstream call with interceptor chain - Use HTTP trailers for billing headers in streaming (fixes timing) - SSE events now forwarded even when parsing fails - stream_options merged instead of overwritten Robustness improvements: - Add nil guards in BillingCalculator.Calculate - Increase SSEParser buffer to 1MB for large payloads - FormatSSEEvent splits multi-line data correctly - Add TotalTokens fallback in extractOpenAIStd - Support Anthropic ephemeral cache token fields Test improvements: - Remove unused mockExtractor embedding - Add cacheUsage nil assertion in tests - Add error checking for ParseAnthropicSSEEvent calls Dependency fix: - Run go mod tidy to mark simdjson-go as direct dependency * feat: implement true streaming for Bedrock and Google AI Previously both providers buffered entire responses before forwarding to client, defeating SSE streaming. Now: Resolvers: - Google AI: Use :streamGenerateContent?alt=sse for streaming - Bedrock: Use /converse-stream for streaming Streaming extractors: - Google AI: Incremental SSE parsing, forward each event immediately - Bedrock: Binary AWS event stream parsing, forward each message immediately Tests: - Verify incremental streaming (data arrives before stream completes) - Test metadata extraction from streams - Test resolver endpoint selection based on stream flag All 100+ tests pass. * fix: resolve data race in streaming tests Use thread-safe response writer with channel-based signaling instead of polling shared buffer. This eliminates the race condition where one goroutine writes to the buffer while another reads its length. - Added threadSafeResponseWriter with mutex-protected buffer - Use channel to signal first write time (eliminates race on time.Time) - Added Flush() method to satisfy http.Flusher interface - Both bedrock and googleai streaming tests now pass with -race
1 parent b1ae20e commit b213390

29 files changed

Lines changed: 4274 additions & 89 deletions

‎DESIGN.md‎

Lines changed: 143 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -394,6 +394,138 @@ curl -X POST http://localhost:8080/v1/chat/completions \
394394

395395
---
396396

397+
## Streaming
398+
399+
The proxy fully supports SSE (Server-Sent Events) streaming with efficient token usage extraction for billing.
400+
401+
### StreamingResponseExtractor
402+
403+
Extends `ResponseExtractor` for streaming responses:
404+
405+
```
406+
type StreamingResponseExtractor interface {
407+
ResponseExtractor
408+
ExtractStreamingWithController(resp, w, rc) -> (ResponseMetadata, error)
409+
IsStreamingResponse(resp) -> bool
410+
}
411+
```
412+
413+
All built-in providers implement this interface for streaming support.
414+
415+
### Streaming Flow
416+
417+
```
418+
+------------------+
419+
| Incoming Request |
420+
| stream: true |
421+
+--------+---------+
422+
|
423+
Parse body, detect
424+
stream: true flag
425+
|
426+
+--------v---------+
427+
| AutoRouter |
428+
| ForwardStreaming |
429+
+--------+---------+
430+
|
431+
Auto-inject
432+
stream_options:
433+
{include_usage:true}
434+
|
435+
+--------v---------+
436+
| HTTP Request |
437+
| to upstream |
438+
+--------+---------+
439+
|
440+
+--------v---------+
441+
| Upstream Response|
442+
| text/event-stream|
443+
+--------+---------+
444+
|
445+
+--------v---------+
446+
| StreamingExtractor|
447+
| Parse SSE events |
448+
| Extract usage |
449+
| Flush each chunk |
450+
+--------+---------+
451+
|
452+
+--------v---------+
453+
| BillingCalculator|
454+
| Calculate cost |
455+
+--------+---------+
456+
|
457+
+--------v---------+
458+
| HTTP Response |
459+
| to client |
460+
+------------------+
461+
```
462+
463+
### Usage Extraction
464+
465+
**OpenAI**: Usage is sent in the final chunk before `[DONE]`:
466+
467+
```json
468+
data: {"id":"...","usage":{"prompt_tokens":100,"completion_tokens":50,"total_tokens":150}}
469+
data: [DONE]
470+
```
471+
472+
**Anthropic**: Usage is sent in `message_start` and `message_delta` events:
473+
474+
```json
475+
data: {"type":"message_start","message":{"usage":{"input_tokens":100}}}
476+
...
477+
data: {"type":"message_delta","usage":{"output_tokens":50}}
478+
data: {"type":"message_stop"}
479+
```
480+
481+
### Auto stream_options Injection
482+
483+
When `BillingCalculator` is configured and the request has `stream: true`, the proxy automatically injects:
484+
485+
```json
486+
{
487+
"stream": true,
488+
"stream_options": { "include_usage": true }
489+
}
490+
```
491+
492+
This ensures OpenAI returns token usage in the streaming response for billing calculation.
493+
494+
### Efficient Flushing
495+
496+
Uses `http.ResponseController` for optimal streaming:
497+
498+
```go
499+
rc := http.NewResponseController(w)
500+
501+
for each SSE event {
502+
w.Write(event)
503+
rc.Flush() // Immediate flush after each chunk
504+
}
505+
```
506+
507+
Non-streaming responses also use chunked read/write/flush with a 512KB buffer for better performance.
508+
509+
### Billing with Streaming
510+
511+
```go
512+
adapter, _ := modelsdev.LoadFromURL()
513+
514+
billingCallback := func(r llmproxy.BillingResult) {
515+
log.Printf("Cost: $%.6f", r.TotalCost)
516+
}
517+
518+
router := llmproxy.NewAutoRouter(
519+
llmproxy.WithAutoRouterBillingCalculator(
520+
llmproxy.NewBillingCalculator(adapter.GetCostLookup(), billingCallback),
521+
),
522+
)
523+
```
524+
525+
After the stream completes, the billing callback is invoked with the extracted token usage.
526+
527+
---
528+
397529
## Providers
398530

399531
Nine providers are included. Six share the OpenAI-compatible base; three have fully custom implementations.
@@ -884,19 +1016,24 @@ Matches the signature of `github.com/agentuity/go-common/logger` without requiri
8841016
```
8851017
llmproxy/
8861018
├── apitype.go # API type detection and constants
887-
├── autorouter.go # AutoRouter, provider/API auto-detection
1019+
├── autorouter.go # AutoRouter, provider/API auto-detection, streaming
8881020
├── billing.go # CostInfo, CostLookup, BillingResult, CalculateCost
1021+
├── billing_calculator.go # BillingCalculator for streaming/non-streaming
8891022
├── detection.go # Provider detection from model/header
8901023
├── enricher.go # RequestEnricher interface
891-
├── extractor.go # ResponseExtractor interface
1024+
├── extractor.go # ResponseExtractor, StreamingResponseExtractor interface
8921025
├── interceptor.go # Interceptor, InterceptorChain, RoundTripFunc
1026+
├── internal/
1027+
│ └── fastjson/
1028+
│ └── extractor.go # Fast JSON parsing with simdjson-go
8931029
├── logger.go # Logger interface, LoggerFunc adapter
8941030
├── metadata.go # BodyMetadata, ResponseMetadata, Message, Usage, Choice
8951031
├── parser.go # BodyParser interface
8961032
├── provider.go # Provider interface, BaseProvider
8971033
├── proxy.go # Proxy struct, Forward method
8981034
├── registry.go # Registry interface, MapRegistry
8991035
├── resolver.go # URLResolver interface
1036+
├── streaming.go # SSE parser, streaming types, usage extraction
9001037
├── interceptors/
9011038
│ ├── addheader.go # AddHeaderInterceptor
9021039
│ ├── billing.go # BillingInterceptor
@@ -911,14 +1048,18 @@ llmproxy/
9111048
│ └── adapter.go # models.dev pricing adapter
9121049
├── providers/
9131050
│ ├── anthropic/ # Anthropic Messages API
1051+
│ │ └── streaming_extractor.go # Anthropic SSE streaming
9141052
│ ├── azure/ # Azure OpenAI
9151053
│ ├── bedrock/ # AWS Bedrock Converse API
1054+
│ │ └── streaming_extractor.go # Bedrock streaming
9161055
│ ├── fireworks/ # Fireworks (OpenAI-compatible)
9171056
│ ├── googleai/ # Google AI Gemini
1057+
│ │ └── streaming_extractor.go # Google AI streaming
9181058
│ ├── groq/ # Groq (OpenAI-compatible)
9191059
│ ├── openai/ # OpenAI (Chat Completions + Responses)
9201060
│ ├── openai_compatible/ # Base for OpenAI-compatible providers
9211061
│ │ ├── multiapi.go # Multi-API parser/extractor
1062+
│ │ ├── streaming_extractor.go # SSE streaming with usage extraction
9221063
│ │ ├── responses_parser.go # Responses API parser
9231064
│ │ └── responses_extractor.go # Responses API extractor
9241065
│ └── xai/ # x.AI (OpenAI-compatible)

‎README.md‎

Lines changed: 32 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -108,6 +108,7 @@ curl -X POST http://localhost:8080/ \
108108
- **9 Provider Implementations**: OpenAI, Anthropic, Groq, Fireworks, x.AI, Google AI, AWS Bedrock, Azure OpenAI, OpenAI-compatible base
109109
- **AutoRouter**: Single endpoint with automatic provider/API detection
110110
- **Responses API**: Full support for OpenAI's new Responses API
111+
- **SSE Streaming**: Full streaming support with efficient token usage extraction
111112
- **8 Built-in Interceptors**: Logging, Metrics, Retry, Billing, Tracing (OTel), HeaderBan, AddHeader, PromptCaching
112113
- **Pricing Integration**: models.dev adapter with markup support
113114
- **Prompt Caching**: prompt caching support for Anthropic, OpenAI, xAI, Fireworks, and Bedrock
@@ -153,6 +154,37 @@ curl -X POST http://localhost:8080/v1/chat/completions \
153154
-d '{"model":"gpt-4","messages":[{"role":"user","content":"Hello"}]}'
154155
```
155156

157+
## Streaming
158+
159+
SSE streaming is fully supported with automatic token usage extraction for billing:
160+
161+
```bash
162+
# Streaming with automatic usage extraction
163+
curl -X POST http://localhost:8080/ \
164+
-H 'Content-Type: application/json' \
165+
-d '{"model":"gpt-4","stream":true,"messages":[{"role":"user","content":"Hello"}]}'
166+
```
167+
168+
**Key Features:**
169+
170+
- **Efficient flushing**: Uses `http.ResponseController` for immediate SSE delivery
171+
- **Token extraction**: Extracts usage from streaming responses for billing
172+
- **Auto stream_options**: Automatically injects `stream_options.include_usage` when billing is configured
173+
- **Works with billing**: Billing is calculated after stream completes
174+
175+
**Example with billing:**
176+
177+
```go
178+
adapter, _ := modelsdev.LoadFromURL()
179+
billingCallback := func(r llmproxy.BillingResult) {
180+
log.Printf("Cost: $%.6f (tokens: %d/%d)", r.TotalCost, r.PromptTokens, r.CompletionTokens)
181+
}
182+
183+
router := llmproxy.NewAutoRouter(
184+
llmproxy.WithAutoRouterBillingCalculator(llmproxy.NewBillingCalculator(adapter.GetCostLookup(), billingCallback)),
185+
)
186+
```
187+
156188
## Providers
157189

158190
| Provider | Auth | API Format | Notes |

0 commit comments

Comments
 (0)