From 93d41b66b14fa892ba1b5aab67dfffdd1a6c3f95 Mon Sep 17 00:00:00 2001 From: chopratejas Date: Fri, 13 Mar 2026 16:49:18 -0700 Subject: [PATCH] Commit message: Fix OpenAI streaming with backends and /v1 double-path bug Add stream_openai_message() to LiteLLM and any-llm backends so /v1/chat/completions with stream:true returns SSE events instead of a JSON blob. Clients (Kilo Code, Cursor, etc.) were hanging because the proxy ignored the stream flag when routing through a backend. Also strip trailing /v1 from OPENAI_TARGET_API_URL to prevent double-path URLs like /v1/v1/models. --- headroom/backends/anyllm.py | 55 +++++ headroom/backends/base.py | 24 ++ headroom/backends/litellm.py | 59 +++++ headroom/proxy/server.py | 190 +++++++++++---- headroom/transforms/kompress_compressor.py | 4 +- tests/test_backend_bugs.py | 59 +++++ tests/test_openai_streaming_backend.py | 262 +++++++++++++++++++++ 7 files changed, 609 insertions(+), 44 deletions(-) create mode 100644 tests/test_openai_streaming_backend.py diff --git a/headroom/backends/anyllm.py b/headroom/backends/anyllm.py index 592ed48bf..95d355c62 100644 --- a/headroom/backends/anyllm.py +++ b/headroom/backends/anyllm.py @@ -454,6 +454,61 @@ class AnyLLMBackend(Backend): return BackendResponse(body=body, status_code=status_code, error=str(e)) + async def stream_openai_message( + self, + body: dict[str, Any], + headers: dict[str, str], + ) -> AsyncIterator[str]: + """Stream OpenAI-format chat completion via any-llm. + + Yields SSE-formatted strings ready to send to the client. + """ + original_model = body.get("model", "gpt-4o") + + try: + kwargs: dict[str, Any] = { + "model": original_model, + "messages": body.get("messages", []), + "stream": True, + } + + for param in [ + "max_tokens", + "temperature", + "top_p", + "stop", + "tools", + "tool_choice", + "response_format", + "seed", + "n", + ]: + if param in body: + kwargs[param] = body[param] + + if "stream_options" in body: + kwargs["stream_options"] = body["stream_options"] + + stream_response = await self.llm.acompletion(**kwargs) + + async for chunk in cast(AsyncIterator[Any], stream_response): + chunk_dict = chunk.model_dump(exclude_none=True, exclude_unset=True) + yield f"data: {json.dumps(chunk_dict)}\n\n" + + yield "data: [DONE]\n\n" + + except Exception as e: + logger.error(f"any-llm OpenAI streaming error: {e}") + error_data = { + "error": { + "message": str(e), + "type": "api_error", + "code": "backend_error", + } + } + yield f"data: {json.dumps(error_data)}\n\n" + yield "data: [DONE]\n\n" + async def close(self) -> None: """Clean up (no-op for any-llm).""" pass diff --git a/headroom/backends/base.py b/headroom/backends/base.py index 9aa77204e..8231c4cce 100644 --- a/headroom/backends/base.py +++ b/headroom/backends/base.py @@ -140,6 +140,30 @@ class Backend(ABC): """ raise NotImplementedError(f"{self.name} backend does not support OpenAI format") + async def stream_openai_message( + self, + body: dict[str, Any], + headers: dict[str, str], + ) -> AsyncIterator[str]: + """Stream an OpenAI-format chat completion. + + Yields SSE-formatted strings: 'data: {...}\\n\\n' for each chunk, + ending with 'data: [DONE]\\n\\n'. + + Args: + body: Request body in OpenAI chat completion format (stream: true). + headers: Request headers. + + Yields: + SSE-formatted strings ready to send to client. + + Raises: + NotImplementedError: If backend doesn't support OpenAI streaming. + """ + raise NotImplementedError(f"{self.name} backend does not support OpenAI streaming") + # Make this an async generator (yield never reached but needed for type) + yield "" # type: ignore[misc] # pragma: no cover + async def close(self) -> None: # noqa: B027 """Clean up resources (e.g., close HTTP clients).""" pass diff --git a/headroom/backends/litellm.py b/headroom/backends/litellm.py index 8659bd04a..9c9deed15 100644 --- a/headroom/backends/litellm.py +++ b/headroom/backends/litellm.py @@ -819,3 +819,62 @@ class LiteLLMBackend(Backend): status_code=status_code, error=str(e), ) + + async def stream_openai_message( + self, + body: dict[str, Any], + headers: dict[str, str], + ) -> AsyncIterator[str]: + """Stream OpenAI-format chat completion via LiteLLM. + + Yields SSE-formatted strings ready to send to the client. + """ + original_model = body.get("model", "gpt-4") + litellm_model = self.map_model_id(original_model) + + try: + kwargs: dict[str, Any] = { + "model": litellm_model, + "messages": body.get("messages", []), + "stream": True, + } + + for param in [ + "max_tokens", + "temperature", + "top_p", + "stop", + "tools", + "tool_choice", + "response_format", + "seed", + "n", + ]: + if param in body: + kwargs[param] = body[param] + + if "stream_options" in body: + kwargs["stream_options"] = body["stream_options"] + + if self.provider == "bedrock" and self.region: + kwargs["aws_region_name"] = self.region + + response = await acompletion(**kwargs) + + async for chunk in response: + chunk_dict = chunk.model_dump(exclude_none=True, exclude_unset=True) + yield f"data: {json.dumps(chunk_dict)}\n\n" + + yield "data: [DONE]\n\n" + + except Exception as e: + logger.error(f"LiteLLM OpenAI streaming error: {e}") + error_data = { + "error": { + "message": str(e), + "type": "api_error", + "code": "backend_error", + } + } + yield f"data: {json.dumps(error_data)}\n\n" + yield "data: [DONE]\n\n" diff --git a/headroom/proxy/server.py b/headroom/proxy/server.py index 684e76fb8..f7364e25c 100644 --- a/headroom/proxy/server.py +++ b/headroom/proxy/server.py @@ -1313,12 +1313,17 @@ class HeadroomProxy: self.config = config # Override OPENAI_API_URL with config if set + # Strip trailing /v1 or /v1/ to avoid double-path (e.g., .../v1/v1/models) if config.openai_api_url: - HeadroomProxy.OPENAI_API_URL = config.openai_api_url + url = config.openai_api_url.rstrip("/") + if url.endswith("/v1"): + url = url[:-3] + HeadroomProxy.OPENAI_API_URL = url # Override GEMINI_API_URL with config if set if config.gemini_api_url: - HeadroomProxy.GEMINI_API_URL = config.gemini_api_url + gurl = config.gemini_api_url.rstrip("/") + HeadroomProxy.GEMINI_API_URL = gurl # Initialize providers self.anthropic_provider = AnthropicProvider() @@ -1651,26 +1656,45 @@ class HeadroomProxy: else: logger.info("Smart Routing: DISABLED (legacy sequential mode)") - # Eagerly load LLMLingua model at startup (avoids 5s delay on first request) - if self.config.llmlingua_enabled: + # Eagerly load ML compressors at startup (avoids download on first request) + # Kompress requires [ml] extra (torch + transformers). If not installed, skip. + self._kompress_status = "not installed" + from headroom.transforms.kompress_compressor import is_kompress_available + + if is_kompress_available() and self.config.optimize: + logger.info("Kompress: Downloading model (first-time only)...") for transform in self.anthropic_pipeline.transforms: if hasattr(transform, "eager_load_compressors"): transform.eager_load_compressors() - self._llmlingua_status = "enabled" + self._kompress_status = "enabled" + break + if self._kompress_status == "enabled": + logger.info("Kompress: ENABLED (ModernBERT token compressor)") + else: + if self.config.optimize: + logger.info( + "Kompress: not installed (pip install headroom-ai[ml] for ML compression)" + ) + + # LLMLingua fallback (only loads if Kompress is not available) + if self._kompress_status != "enabled" and self.config.llmlingua_enabled: + for transform in self.anthropic_pipeline.transforms: + if hasattr(transform, "_get_llmlingua"): + llmlingua = transform._get_llmlingua() + if llmlingua: + self._llmlingua_status = "enabled" break - # LLMLingua status with helpful hint + # LLMLingua status if self._llmlingua_status == "enabled": logger.info( f"LLMLingua: ENABLED (device={self.config.llmlingua_device}, " f"rate={self.config.llmlingua_target_rate})" ) + elif self._kompress_status == "enabled": + logger.info("LLMLingua: skipped (Kompress is active)") elif self._llmlingua_status == "lazy": logger.info("LLMLingua: LAZY (will load when prose content detected)") - elif self._llmlingua_status == "available": - logger.info("LLMLingua: available but disabled (use --llmlingua)") - elif self._llmlingua_status == "unavailable": - logger.info("LLMLingua: not installed (pip install headroom-ai[llmlingua])") elif self._llmlingua_status == "disabled": logger.info("LLMLingua: DISABLED") @@ -4340,6 +4364,67 @@ class HeadroomProxy: media_type="text/event-stream", ) + async def _stream_openai_via_backend( + self, + body: dict, + headers: dict, + model: str, + request_id: str, + start_time: float, + original_tokens: int, + optimized_tokens: int, + tokens_saved: int, + transforms_applied: list[str], + tags: dict[str, str], + optimization_latency: float, + pipeline_timing: dict[str, float] | None = None, + ) -> StreamingResponse: + """Stream OpenAI chat completion response from backend. + + Routes stream:true requests through the backend's stream_openai_message(), + yielding SSE events to the client. + """ + assert self.anthropic_backend is not None + + async def generate(): + try: + async for sse_chunk in self.anthropic_backend.stream_openai_message(body, headers): + yield sse_chunk.encode() if isinstance(sse_chunk, str) else sse_chunk + except Exception as e: + logger.error(f"[{request_id}] Backend streaming error: {e}") + error_data = { + "error": { + "message": str(e), + "type": "api_error", + "code": "backend_error", + } + } + yield f"data: {json.dumps(error_data)}\n\n".encode() + yield b"data: [DONE]\n\n" + finally: + total_latency = (time.time() - start_time) * 1000 + await self.metrics.record_request( + provider=self.anthropic_backend.name, + model=model, + input_tokens=optimized_tokens, + output_tokens=0, # Unknown in streaming + tokens_saved=tokens_saved, + latency_ms=total_latency, + cached=False, + overhead_ms=optimization_latency, + pipeline_timing=pipeline_timing, + ) + if tokens_saved > 0: + logger.info( + f"[{request_id}] {model}: {original_tokens:,} → {optimized_tokens:,} " + f"(saved {tokens_saved:,} tokens) via {self.anthropic_backend.name} [stream]" + ) + + return StreamingResponse( + generate(), + media_type="text/event-stream", + ) + async def handle_openai_chat( self, request: Request, @@ -4533,46 +4618,65 @@ class HeadroomProxy: if tools is not None: body["tools"] = tools - # Route through LiteLLM backend if configured (Databricks, Bedrock, etc.) + # Route through LiteLLM/any-llm backend if configured if self.anthropic_backend is not None: try: - # Use the backend's OpenAI-format method - backend_response = await self.anthropic_backend.send_openai_message(body, headers) + if stream: + # Streaming: use stream_openai_message() → SSE events + return await self._stream_openai_via_backend( + body, + headers, + model, + request_id, + start_time, + original_tokens, + optimized_tokens, + tokens_saved, + transforms_applied, + tags, + optimization_latency, + pipeline_timing=pipeline_timing, + ) + else: + # Non-streaming: use send_openai_message() → JSON + backend_response = await self.anthropic_backend.send_openai_message( + body, headers + ) + + if backend_response.error: + return JSONResponse( + status_code=backend_response.status_code, + content=backend_response.body, + ) + + # Track metrics + total_latency = (time.time() - start_time) * 1000 + usage = backend_response.body.get("usage", {}) + output_tokens = usage.get("completion_tokens", 0) + total_input_tokens = usage.get("prompt_tokens", optimized_tokens) + + await self.metrics.record_request( + provider=self.anthropic_backend.name, + model=model, + input_tokens=total_input_tokens, + output_tokens=output_tokens, + tokens_saved=tokens_saved, + latency_ms=total_latency, + cached=False, + overhead_ms=optimization_latency, + pipeline_timing=pipeline_timing, + ) + + if tokens_saved > 0: + logger.info( + f"[{request_id}] {model}: {original_tokens:,} → {optimized_tokens:,} " + f"(saved {tokens_saved:,} tokens) via {self.anthropic_backend.name}" + ) - if backend_response.error: return JSONResponse( status_code=backend_response.status_code, content=backend_response.body, ) - - # Track metrics - total_latency = (time.time() - start_time) * 1000 - usage = backend_response.body.get("usage", {}) - output_tokens = usage.get("completion_tokens", 0) - total_input_tokens = usage.get("prompt_tokens", optimized_tokens) - - await self.metrics.record_request( - provider=self.anthropic_backend.name, - model=model, - input_tokens=total_input_tokens, - output_tokens=output_tokens, - tokens_saved=tokens_saved, - latency_ms=total_latency, - cached=False, - overhead_ms=optimization_latency, - pipeline_timing=pipeline_timing, - ) - - if tokens_saved > 0: - logger.info( - f"[{request_id}] {model}: {original_tokens:,} → {optimized_tokens:,} " - f"(saved {tokens_saved:,} tokens) via {self.anthropic_backend.name}" - ) - - return JSONResponse( - status_code=backend_response.status_code, - content=backend_response.body, - ) except Exception as e: logger.error(f"[{request_id}] Backend error: {e}") return JSONResponse( diff --git a/headroom/transforms/kompress_compressor.py b/headroom/transforms/kompress_compressor.py index 6918d6e0b..49280c2e4 100644 --- a/headroom/transforms/kompress_compressor.py +++ b/headroom/transforms/kompress_compressor.py @@ -140,10 +140,12 @@ def _load_kompress(device: str = "auto") -> tuple[HeadroomCompressorModel, Any]: def is_kompress_available() -> bool: - """Check if Kompress dependencies are available.""" + """Check if Kompress dependencies are available (requires [ml] extra).""" try: import huggingface_hub # noqa: F401 import safetensors # noqa: F401 + import torch # noqa: F401 + import transformers # noqa: F401 return True except ImportError: diff --git a/tests/test_backend_bugs.py b/tests/test_backend_bugs.py index c7cd4b969..80a35cf0a 100644 --- a/tests/test_backend_bugs.py +++ b/tests/test_backend_bugs.py @@ -290,3 +290,62 @@ class TestVertexModelMap: def test_claude_3_legacy(self): assert "claude-3-haiku-20240307" in _VERTEX_MODEL_MAP + + +# ============================================================================= +# URL Normalization (trailing /v1 stripping) +# ============================================================================= + +pytest.importorskip("fastapi") + + +class TestOpenAIURLNormalization: + """Test that OPENAI_TARGET_API_URL with /v1 suffix is normalized.""" + + def test_v1_suffix_stripped(self): + from headroom.proxy.server import HeadroomProxy, ProxyConfig + + original = HeadroomProxy.OPENAI_API_URL + try: + config = ProxyConfig( + openai_api_url="http://localhost:4000/v1", + optimize=False, + cache_enabled=False, + rate_limit_enabled=False, + ) + proxy = HeadroomProxy(config) + assert proxy.OPENAI_API_URL == "http://localhost:4000" + finally: + HeadroomProxy.OPENAI_API_URL = original + + def test_v1_slash_suffix_stripped(self): + from headroom.proxy.server import HeadroomProxy, ProxyConfig + + original = HeadroomProxy.OPENAI_API_URL + try: + config = ProxyConfig( + openai_api_url="http://localhost:4000/v1/", + optimize=False, + cache_enabled=False, + rate_limit_enabled=False, + ) + proxy = HeadroomProxy(config) + assert proxy.OPENAI_API_URL == "http://localhost:4000" + finally: + HeadroomProxy.OPENAI_API_URL = original + + def test_no_v1_unchanged(self): + from headroom.proxy.server import HeadroomProxy, ProxyConfig + + original = HeadroomProxy.OPENAI_API_URL + try: + config = ProxyConfig( + openai_api_url="http://localhost:4000", + optimize=False, + cache_enabled=False, + rate_limit_enabled=False, + ) + proxy = HeadroomProxy(config) + assert proxy.OPENAI_API_URL == "http://localhost:4000" + finally: + HeadroomProxy.OPENAI_API_URL = original diff --git a/tests/test_openai_streaming_backend.py b/tests/test_openai_streaming_backend.py new file mode 100644 index 000000000..bf1d03ced --- /dev/null +++ b/tests/test_openai_streaming_backend.py @@ -0,0 +1,262 @@ +"""Test OpenAI /v1/chat/completions streaming through headroom proxy backends. + +Proves that streaming works end-to-end: client → headroom proxy → backend → OpenAI API. + +Two test modes: +1. Real API test (requires OPENAI_API_KEY): hits actual OpenAI with gpt-4o-mini +2. Mock test: proves the proxy returns SSE when stream:true with a backend configured + +Run with: + OPENAI_API_KEY=sk-... pytest tests/test_openai_streaming_backend.py -v +""" + +import os +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest + +fastapi = pytest.importorskip("fastapi") +httpx = pytest.importorskip("httpx") + +from fastapi.testclient import TestClient # noqa: E402 + +from headroom.backends.base import BackendResponse # noqa: E402 +from headroom.proxy.server import ProxyConfig, create_app # noqa: E402 + +# ============================================================================= +# Real API test (requires OPENAI_API_KEY) +# ============================================================================= + + +@pytest.mark.skipif(not os.environ.get("OPENAI_API_KEY"), reason="OPENAI_API_KEY not set") +class TestOpenAIStreamingRealAPI: + """Test streaming with real OpenAI API calls through the proxy.""" + + @pytest.fixture + def openai_api_key(self): + return os.environ["OPENAI_API_KEY"] + + @pytest.fixture + def direct_proxy_client(self): + """Proxy with NO backend — direct to OpenAI. This is the baseline.""" + config = ProxyConfig( + optimize=False, + cache_enabled=False, + rate_limit_enabled=False, + ) + app = create_app(config) + with TestClient(app) as client: + yield client + + @pytest.fixture + def litellm_backend_client(self): + """Proxy with litellm-openai backend — routes through LiteLLM.""" + config = ProxyConfig( + optimize=False, + cache_enabled=False, + rate_limit_enabled=False, + backend="litellm-openai", + ) + app = create_app(config) + with TestClient(app) as client: + yield client + + def test_baseline_streaming_works_direct(self, direct_proxy_client, openai_api_key): + """Baseline: streaming through proxy WITHOUT backend works (direct to OpenAI).""" + response = direct_proxy_client.post( + "/v1/chat/completions", + json={ + "model": "gpt-4o-mini", + "messages": [{"role": "user", "content": "Say 'hello' and nothing else."}], + "stream": True, + "max_tokens": 10, + }, + headers={"Authorization": f"Bearer {openai_api_key}"}, + ) + + assert response.status_code == 200, f"Got {response.status_code}: {response.text[:200]}" + + content_type = response.headers.get("content-type", "") + assert "text/event-stream" in content_type, ( + f"Direct proxy streaming broken: got content-type '{content_type}'" + ) + + # Verify we got actual SSE chunks + body = response.text + assert "data: " in body, "No SSE data chunks in response" + assert "data: [DONE]" in body, "Missing [DONE] terminator" + + def test_streaming_with_litellm_backend(self, litellm_backend_client, openai_api_key): + """CRITICAL: streaming through proxy WITH litellm backend must also stream. + + This test fails before the fix — the proxy returns a JSON blob + instead of SSE events, causing clients to hang. + """ + response = litellm_backend_client.post( + "/v1/chat/completions", + json={ + "model": "gpt-4o-mini", + "messages": [{"role": "user", "content": "Say 'hello' and nothing else."}], + "stream": True, + "max_tokens": 10, + }, + headers={"Authorization": f"Bearer {openai_api_key}"}, + ) + + assert response.status_code == 200, f"Got {response.status_code}: {response.text[:200]}" + + content_type = response.headers.get("content-type", "") + assert "text/event-stream" in content_type, ( + f"STREAMING BUG: litellm backend returned '{content_type}' instead of " + f"'text/event-stream'. Client sees a JSON blob, not SSE events.\n" + f"Response body (first 300 chars): {response.text[:300]}" + ) + + # Verify SSE format + body = response.text + assert "data: " in body, "No SSE data chunks in streaming response" + + def test_non_streaming_with_litellm_backend(self, litellm_backend_client, openai_api_key): + """Non-streaming with backend should return normal JSON (sanity check).""" + response = litellm_backend_client.post( + "/v1/chat/completions", + json={ + "model": "gpt-4o-mini", + "messages": [{"role": "user", "content": "Say 'hello' and nothing else."}], + "stream": False, + "max_tokens": 10, + }, + headers={"Authorization": f"Bearer {openai_api_key}"}, + ) + + assert response.status_code == 200, f"Got {response.status_code}: {response.text[:200]}" + + content_type = response.headers.get("content-type", "") + assert "application/json" in content_type + + data = response.json() + assert "choices" in data + assert data["choices"][0]["message"]["content"] + + +# ============================================================================= +# Mock test (no API key needed — proves the routing bug) +# ============================================================================= + + +class TestOpenAIStreamingMock: + """Prove the streaming bug with mocks — no API key needed.""" + + def test_streaming_request_returns_sse_not_json(self): + """When stream:true with a backend, content-type MUST be text/event-stream. + + This test FAILS before the fix: the proxy calls send_openai_message() + (non-streaming) and returns application/json even though stream:true. + """ + config = ProxyConfig( + optimize=False, + cache_enabled=False, + rate_limit_enabled=False, + backend="anyllm", + anyllm_provider="openai", + ) + + mock_backend = MagicMock() + mock_backend.name = "anyllm-openai" + mock_backend.send_openai_message = AsyncMock( + return_value=BackendResponse( + body={ + "id": "chatcmpl-123", + "object": "chat.completion", + "model": "test-model", + "choices": [ + { + "index": 0, + "message": {"role": "assistant", "content": "Hello!"}, + "finish_reason": "stop", + } + ], + "usage": {"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15}, + }, + status_code=200, + headers={"content-type": "application/json"}, + ) + ) + + with patch("headroom.proxy.server.AnyLLMBackend", return_value=mock_backend): + app = create_app(config) + + with TestClient(app) as client: + response = client.post( + "/v1/chat/completions", + json={ + "model": "test-model", + "messages": [{"role": "user", "content": "hello"}], + "stream": True, + }, + headers={"Authorization": "Bearer test-key"}, + ) + + assert response.status_code == 200, ( + f"Got {response.status_code}: {response.text[:200]}" + ) + + content_type = response.headers.get("content-type", "") + assert "text/event-stream" in content_type, ( + f"STREAMING BUG: stream:true with backend returned '{content_type}' " + f"instead of 'text/event-stream'. The proxy ignored the stream flag " + f"and returned a JSON blob. Clients expecting SSE will hang.\n" + f"Response: {response.text[:300]}" + ) + + def test_non_streaming_still_returns_json(self): + """Sanity: stream:false with backend should return JSON as before.""" + config = ProxyConfig( + optimize=False, + cache_enabled=False, + rate_limit_enabled=False, + backend="anyllm", + anyllm_provider="openai", + ) + + mock_backend = MagicMock() + mock_backend.name = "anyllm-openai" + mock_backend.send_openai_message = AsyncMock( + return_value=BackendResponse( + body={ + "id": "chatcmpl-123", + "object": "chat.completion", + "model": "test-model", + "choices": [ + { + "index": 0, + "message": {"role": "assistant", "content": "Hello!"}, + "finish_reason": "stop", + } + ], + "usage": {"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15}, + }, + status_code=200, + headers={"content-type": "application/json"}, + ) + ) + + with patch("headroom.proxy.server.AnyLLMBackend", return_value=mock_backend): + app = create_app(config) + + with TestClient(app) as client: + response = client.post( + "/v1/chat/completions", + json={ + "model": "test-model", + "messages": [{"role": "user", "content": "hello"}], + "stream": False, + }, + headers={"Authorization": "Bearer test-key"}, + ) + + assert response.status_code == 200 + content_type = response.headers.get("content-type", "") + assert "application/json" in content_type + data = response.json() + assert data["choices"][0]["message"]["content"] == "Hello!"