From 95fd6d86882238bd97bdf2ea23814e9271d23dec Mon Sep 17 00:00:00 2001 From: chopratejas Date: Sat, 24 Jan 2026 11:41:18 -0800 Subject: [PATCH] Add multi-provider batch API support with CCR post-processing This commit adds comprehensive batch API support for all three major LLM providers (Anthropic, OpenAI, Google/Gemini) with integrated CCR (Compress-Cache-Retrieve) functionality for asynchronous batch processing. ## Batch CCR Post-Processing Architecture When batch APIs are used, responses are processed asynchronously. If the model calls the CCR retrieval tool (`headroom_retrieve`) within a batch response, the system now handles this automatically: 1. **Batch Submit**: Request context (messages, tools, model) is stored in BatchContextStore keyed by batch_id 2. **Batch Results**: When results are retrieved, CCR tool calls are detected in the responses 3. **Continuation**: For each CCR tool call, the system executes local retrieval and makes a continuation API call to complete the response 4. **Result Update**: The batch result is updated with the complete response, transparent to the caller ## New Components - `headroom/ccr/batch_store.py`: TTL-based context storage for batch requests, enabling CCR retrieval during result processing - `headroom/ccr/batch_processor.py`: Processes batch results, detects CCR tool calls across all provider formats, executes continuations ## Provider Support ### Anthropic - POST /v1/messages/batches (create with compression) - GET /v1/messages/batches (list) - GET /v1/messages/batches/{id} (status) - GET /v1/messages/batches/{id}/results (with CCR post-processing) ### OpenAI - POST /v1/batches (create) - GET /v1/batches (list) - GET /v1/batches/{id} (status) - Batch file upload/download support ### Google/Gemini - Native API support: /v1beta/models/{model}:generateContent - Batch API: /v1beta/models/{model}:batchGenerateContent - Token counting: /v1beta/models/{model}:countTokens - OpenAI-compatible endpoint support ## CCR Enhancements - Added Google/Gemini format support to response_handler.py - Extended tool_injection.py with multiple marker patterns for different compressors (SmartCrusher, TextCompressor, LogCompressor, etc.) - Added Google functionCall/functionResponse handling ## Proxy Server Updates - Added Gemini native API handlers alongside OpenAI-compatible endpoints - Integrated batch context storage on submission - Added batch result processing with CCR continuation - Rate limiting and metrics tracking for all providers ## Test Coverage Added comprehensive integration tests (all skip gracefully without API keys): - test_proxy_batch_integration.py: Anthropic and OpenAI batch APIs - test_proxy_gemini_integration.py: Gemini via OpenAI-compatible endpoint - test_proxy_gemini_native_integration.py: Gemini native API - test_proxy_count_tokens_integration.py: Token counting endpoint - test_proxy_openai_responses_integration.py: OpenAI responses API - test_proxy_passthrough_integration.py: Passthrough endpoints ## Compression Results (Real API Testing) - Token savings: 83-98% on tool result content - CCR tool injection: Working across all providers - Model behavior: OpenAI gpt-4o-mini successfully called headroom_retrieve when presented with compressed data, proving end-to-end CCR functionality --- headroom/ccr/__init__.py | 31 +- headroom/ccr/batch_processor.py | 534 ++++ headroom/ccr/batch_store.py | 253 ++ headroom/ccr/response_handler.py | 57 +- headroom/ccr/tool_injection.py | 64 +- headroom/proxy/server.py | 2272 ++++++++++++++++- tests/test_ccr_tool_injection.py | 84 + tests/test_proxy_batch_integration.py | 522 ++++ tests/test_proxy_count_tokens_integration.py | 500 ++++ tests/test_proxy_gemini_integration.py | 254 ++ tests/test_proxy_gemini_native_integration.py | 374 +++ ...test_proxy_openai_responses_integration.py | 280 ++ tests/test_proxy_passthrough_integration.py | 486 ++++ 13 files changed, 5694 insertions(+), 17 deletions(-) create mode 100644 headroom/ccr/batch_processor.py create mode 100644 headroom/ccr/batch_store.py create mode 100644 tests/test_proxy_batch_integration.py create mode 100644 tests/test_proxy_count_tokens_integration.py create mode 100644 tests/test_proxy_gemini_integration.py create mode 100644 tests/test_proxy_gemini_native_integration.py create mode 100644 tests/test_proxy_openai_responses_integration.py create mode 100644 tests/test_proxy_passthrough_integration.py diff --git a/headroom/ccr/__init__.py b/headroom/ccr/__init__.py index fe612e1b0..84aee7f3a 100644 --- a/headroom/ccr/__init__.py +++ b/headroom/ccr/__init__.py @@ -3,18 +3,37 @@ This module provides tool injection and retrieval handling for the CCR architecture. When tool outputs are compressed, the LLM can retrieve more data if needed. -Three key components: +Four key components: 1. Tool Injection: Proxy injects headroom_retrieve tool into requests 2. Response Handler: Intercepts responses, handles CCR tool calls automatically 3. Context Tracker: Tracks compressed content across turns, enables proactive expansion +4. Batch Processing: Handles CCR tool calls in batch API results (async processing) Two distribution channels for the retrieval tool: 1. Tool Injection: Proxy injects tool into request when compression occurs 2. MCP Server: Standalone server exposes tool via MCP protocol When MCP is configured, tool injection is skipped to avoid duplicates. + +Batch API Support: +- On batch submit: Store request context (messages, tools) in BatchContextStore +- On batch results: Detect CCR tool calls, execute retrieval, make continuation calls +- Works with all providers: Anthropic, OpenAI, Google """ +from .batch_processor import ( + BatchResultProcessor, + BatchResultProcessorConfig, + ProcessedBatchResult, + process_batch_results, +) +from .batch_store import ( + BatchContext, + BatchContextStore, + BatchRequestContext, + get_batch_context_store, + reset_batch_context_store, +) from .context_tracker import ( CompressedContext, ContextTracker, @@ -70,6 +89,16 @@ __all__ = [ "ExpansionRecommendation", "get_context_tracker", "reset_context_tracker", + # Batch processing + "BatchContext", + "BatchContextStore", + "BatchRequestContext", + "BatchResultProcessor", + "BatchResultProcessorConfig", + "ProcessedBatchResult", + "get_batch_context_store", + "process_batch_results", + "reset_batch_context_store", # MCP server "CCRMCPServer", "create_ccr_mcp_server", diff --git a/headroom/ccr/batch_processor.py b/headroom/ccr/batch_processor.py new file mode 100644 index 000000000..584354805 --- /dev/null +++ b/headroom/ccr/batch_processor.py @@ -0,0 +1,534 @@ +"""Batch result post-processor for CCR tool call handling. + +When batch results are retrieved, this processor: +1. Detects CCR tool calls in each result +2. Executes the retrieval locally (from compression store) +3. Makes continuation API calls to get final responses +4. Returns the processed results with complete answers + +This module works with all three providers: +- Anthropic: Batch Message API +- OpenAI: Batch API +- Google/Gemini: Batch API + +Each provider has different result formats, but the logic is the same: +1. Parse result to detect CCR tool calls +2. Execute retrieval +3. Make continuation call with tool result +4. Replace partial result with complete result +""" + +from __future__ import annotations + +import logging +from dataclasses import dataclass +from typing import Any, Protocol + +import httpx + +from .batch_store import BatchContext, BatchRequestContext, get_batch_context_store +from .response_handler import CCRResponseHandler, ResponseHandlerConfig +from .tool_injection import CCR_TOOL_NAME + +logger = logging.getLogger(__name__) + + +class APIClient(Protocol): + """Protocol for making API calls.""" + + async def post( + self, + url: str, + headers: dict[str, str], + json: dict[str, Any], + ) -> httpx.Response: + """Make a POST request.""" + ... + + +@dataclass +class BatchResultProcessorConfig: + """Configuration for batch result processing.""" + + # Whether to process CCR tool calls automatically + enabled: bool = True + + # Timeout for continuation API calls (seconds) + continuation_timeout: int = 120 + + # Maximum continuation rounds per result + max_continuation_rounds: int = 3 + + +@dataclass +class ProcessedBatchResult: + """A processed batch result.""" + + custom_id: str + result: dict[str, Any] + was_processed: bool = False # True if CCR tool calls were handled + continuation_rounds: int = 0 + error: str | None = None + + +class BatchResultProcessor: + """Processes batch results to handle CCR tool calls. + + When a batch result contains a CCR tool call (headroom_retrieve), + this processor: + 1. Looks up the original request context + 2. Executes the retrieval from the compression store + 3. Makes a continuation API call with the tool result + 4. Returns the final (complete) response + + Usage: + processor = BatchResultProcessor(http_client) + + # Process results as they come in + processed = await processor.process_results( + batch_id="batch_123", + results=raw_results, + provider="anthropic" + ) + + # Results now have complete responses (CCR handled) + """ + + def __init__( + self, + http_client: httpx.AsyncClient, + config: BatchResultProcessorConfig | None = None, + ) -> None: + self.http_client = http_client + self.config = config or BatchResultProcessorConfig() + self.ccr_handler = CCRResponseHandler( + ResponseHandlerConfig( + enabled=True, + max_retrieval_rounds=self.config.max_continuation_rounds, + ) + ) + + # API base URLs + self.api_urls = { + "anthropic": "https://api.anthropic.com", + "openai": "https://api.openai.com", + "google": "https://generativelanguage.googleapis.com", + } + + async def process_results( + self, + batch_id: str, + results: list[dict[str, Any]], + provider: str, + ) -> list[ProcessedBatchResult]: + """Process batch results, handling CCR tool calls. + + Args: + batch_id: The batch ID (to look up context). + results: Raw batch results from the provider. + provider: The provider type. + + Returns: + List of processed results (with CCR handled). + """ + if not self.config.enabled: + return [ + ProcessedBatchResult( + custom_id=self._get_custom_id(r, provider), + result=r, + ) + for r in results + ] + + # Get batch context + store = get_batch_context_store() + batch_context = await store.get(batch_id) + + if batch_context is None: + logger.warning( + f"Batch context not found for {batch_id}, returning results without CCR processing" + ) + return [ + ProcessedBatchResult( + custom_id=self._get_custom_id(r, provider), + result=r, + ) + for r in results + ] + + # Process each result + processed = [] + for result in results: + custom_id = self._get_custom_id(result, provider) + request_context = batch_context.get_request(custom_id) + + if request_context is None: + logger.warning(f"Request context not found for {custom_id} in batch {batch_id}") + processed.append(ProcessedBatchResult(custom_id=custom_id, result=result)) + continue + + # Check if result contains CCR tool calls + response = self._extract_response(result, provider) + + if response and self.ccr_handler.has_ccr_tool_calls(response, provider): + # Process the CCR tool calls + try: + final_result = await self._process_single_result( + result, + response, + request_context, + batch_context, + provider, + ) + processed.append(final_result) + except Exception as e: + logger.error(f"Failed to process CCR for {custom_id}: {e}") + processed.append( + ProcessedBatchResult( + custom_id=custom_id, + result=result, + error=str(e), + ) + ) + else: + # No CCR tool calls, pass through + processed.append(ProcessedBatchResult(custom_id=custom_id, result=result)) + + return processed + + def _get_custom_id(self, result: dict[str, Any], provider: str) -> str: + """Extract the custom ID from a result.""" + if provider == "anthropic": + return result.get("custom_id", "") + elif provider == "openai": + return result.get("custom_id", "") + elif provider == "google": + # Google uses metadata.key + return result.get("metadata", {}).get("key", "") + return result.get("custom_id", result.get("id", "")) + + def _extract_response( + self, + result: dict[str, Any], + provider: str, + ) -> dict[str, Any] | None: + """Extract the actual response from a batch result.""" + if provider == "anthropic": + # Anthropic: result.result.message + return result.get("result", {}).get("message") + elif provider == "openai": + # OpenAI: response.body (the full chat completion) + return result.get("response", {}).get("body") + elif provider == "google": + # Google: response (the generateContent response) + return result.get("response") + return result.get("response") + + async def _process_single_result( + self, + original_result: dict[str, Any], + response: dict[str, Any], + request_context: BatchRequestContext, + batch_context: BatchContext, + provider: str, + ) -> ProcessedBatchResult: + """Process a single result with CCR tool calls. + + Args: + original_result: The original batch result. + response: The extracted response (with CCR tool calls). + request_context: The original request context. + batch_context: The batch context. + provider: The provider type. + + Returns: + Processed result with complete response. + """ + custom_id = request_context.custom_id + + # Create API call function for continuations + async def api_call_fn( + messages: list[dict[str, Any]], + tools: list[dict[str, Any]] | None, + ) -> dict[str, Any]: + return await self._make_continuation_call( + messages, + tools, + request_context, + batch_context, + provider, + ) + + # Use CCR handler to process the response + final_response = await self.ccr_handler.handle_response( + response, + request_context.messages, + request_context.tools, + api_call_fn, + provider, + ) + + # Update the result with the final response + updated_result = self._update_result( + original_result, + final_response, + provider, + ) + + return ProcessedBatchResult( + custom_id=custom_id, + result=updated_result, + was_processed=True, + continuation_rounds=self.ccr_handler._retrieval_count, + ) + + async def _make_continuation_call( + self, + messages: list[dict[str, Any]], + tools: list[dict[str, Any]] | None, + request_context: BatchRequestContext, + batch_context: BatchContext, + provider: str, + ) -> dict[str, Any]: + """Make a continuation API call. + + Args: + messages: The messages including tool results. + tools: The tools list. + request_context: The request context. + batch_context: The batch context. + provider: The provider type. + + Returns: + The API response. + """ + if provider == "anthropic": + return await self._anthropic_continuation( + messages, tools, request_context, batch_context + ) + elif provider == "openai": + return await self._openai_continuation(messages, tools, request_context, batch_context) + elif provider == "google": + return await self._google_continuation(messages, tools, request_context, batch_context) + else: + raise ValueError(f"Unknown provider: {provider}") + + async def _anthropic_continuation( + self, + messages: list[dict[str, Any]], + tools: list[dict[str, Any]] | None, + request_context: BatchRequestContext, + batch_context: BatchContext, + ) -> dict[str, Any]: + """Make Anthropic continuation call.""" + url = f"{self.api_urls['anthropic']}/v1/messages" + + headers = { + "Content-Type": "application/json", + "anthropic-version": "2023-06-01", + } + if batch_context.api_key: + headers["x-api-key"] = batch_context.api_key + + body = { + "model": request_context.model, + "messages": messages, + "max_tokens": request_context.extras.get("max_tokens", 4096), + } + if tools: + body["tools"] = tools + + response = await self.http_client.post( + url, + headers=headers, + json=body, + timeout=self.config.continuation_timeout, + ) + response.raise_for_status() + return response.json() + + async def _openai_continuation( + self, + messages: list[dict[str, Any]], + tools: list[dict[str, Any]] | None, + request_context: BatchRequestContext, + batch_context: BatchContext, + ) -> dict[str, Any]: + """Make OpenAI continuation call.""" + url = f"{self.api_urls['openai']}/v1/chat/completions" + + headers = { + "Content-Type": "application/json", + } + if batch_context.api_key: + headers["Authorization"] = f"Bearer {batch_context.api_key}" + + body = { + "model": request_context.model, + "messages": messages, + } + if tools: + body["tools"] = tools + + response = await self.http_client.post( + url, + headers=headers, + json=body, + timeout=self.config.continuation_timeout, + ) + response.raise_for_status() + return response.json() + + async def _google_continuation( + self, + messages: list[dict[str, Any]], + tools: list[dict[str, Any]] | None, + request_context: BatchRequestContext, + batch_context: BatchContext, + ) -> dict[str, Any]: + """Make Google/Gemini continuation call. + + Note: Google format uses 'contents' not 'messages', + and 'parts' format for messages. + """ + model = request_context.model + url = f"{self.api_urls['google']}/v1beta/models/{model}:generateContent" + + if batch_context.api_key: + url = f"{url}?key={batch_context.api_key}" + + headers = {"Content-Type": "application/json"} + + # Convert messages to Google format (contents with parts) + contents = self._messages_to_google_contents(messages) + + body: dict[str, Any] = {"contents": contents} + + # Add system instruction if present + if request_context.system_instruction: + body["systemInstruction"] = {"parts": [{"text": request_context.system_instruction}]} + + # Add tools + if tools: + body["tools"] = [{"functionDeclarations": tools}] + + response = await self.http_client.post( + url, + headers=headers, + json=body, + timeout=self.config.continuation_timeout, + ) + response.raise_for_status() + return response.json() + + def _messages_to_google_contents( + self, + messages: list[dict[str, Any]], + ) -> list[dict[str, Any]]: + """Convert standard messages to Google contents format.""" + contents = [] + + for msg in messages: + role = msg.get("role", "") + content = msg.get("content") + + # Handle Google format messages (already have parts) + if "parts" in msg: + google_role = "model" if role in ("assistant", "model") else "user" + contents.append({"role": google_role, "parts": msg["parts"]}) + continue + + # Map roles + if role == "system": + # Skip system messages (handled separately) + continue + elif role == "assistant": + google_role = "model" + else: + google_role = "user" + + # Convert content to parts + if isinstance(content, str): + contents.append({"role": google_role, "parts": [{"text": content}]}) + elif isinstance(content, list): + # Handle structured content (tool results, etc.) + parts = [] + for block in content: + if isinstance(block, dict): + if block.get("type") == "text": + parts.append({"text": block.get("text", "")}) + elif block.get("type") == "tool_result": + parts.append( + { + "functionResponse": { + "name": block.get("tool_use_id", CCR_TOOL_NAME), + "response": {"content": block.get("content", "")}, + } + } + ) + elif block.get("type") == "tool_use": + parts.append( + { + "functionCall": { + "name": block.get("name", ""), + "args": block.get("input", {}), + } + } + ) + if parts: + contents.append({"role": google_role, "parts": parts}) + + return contents + + def _update_result( + self, + original_result: dict[str, Any], + final_response: dict[str, Any], + provider: str, + ) -> dict[str, Any]: + """Update a batch result with the final processed response.""" + result = dict(original_result) + + if provider == "anthropic": + # Update result.result.message + if "result" not in result: + result["result"] = {} + result["result"]["message"] = final_response + # Update type if it was tool_use + result["result"]["type"] = "succeeded" + + elif provider == "openai": + # Update response.body + if "response" not in result: + result["response"] = {} + result["response"]["body"] = final_response + + elif provider == "google": + # Update response directly + result["response"] = final_response + + return result + + +# Convenience function +async def process_batch_results( + batch_id: str, + results: list[dict[str, Any]], + provider: str, + http_client: httpx.AsyncClient, +) -> list[ProcessedBatchResult]: + """Process batch results with CCR handling. + + This is a convenience function for one-off processing. + + Args: + batch_id: The batch ID. + results: Raw batch results. + provider: The provider type. + http_client: HTTP client for API calls. + + Returns: + Processed results. + """ + processor = BatchResultProcessor(http_client) + return await processor.process_results(batch_id, results, provider) diff --git a/headroom/ccr/batch_store.py b/headroom/ccr/batch_store.py new file mode 100644 index 000000000..f2ff3c73e --- /dev/null +++ b/headroom/ccr/batch_store.py @@ -0,0 +1,253 @@ +"""Batch context storage for CCR post-processing. + +When batches are submitted, we store the request context (messages, tools, model) +so that when results are retrieved, we can handle CCR tool calls and make +continuation API calls. + +This module provides: +1. BatchContext: Data class for stored batch context +2. BatchContextStore: TTL-based cache for batch contexts +""" + +from __future__ import annotations + +import asyncio +import logging +import time +from dataclasses import dataclass, field +from typing import Any + +logger = logging.getLogger(__name__) + +# Default TTL for batch contexts (24 hours - batches can take a while) +DEFAULT_BATCH_CONTEXT_TTL = 86400 + +# Maximum contexts to store (prevent memory issues) +MAX_BATCH_CONTEXTS = 10000 + + +@dataclass +class BatchRequestContext: + """Context for a single request within a batch.""" + + custom_id: ( + str # The request ID within the batch (custom_id for Anthropic, metadata.key for Google) + ) + messages: list[dict[str, Any]] + tools: list[dict[str, Any]] | None = None + model: str = "" + system_instruction: str | None = None # For Google format + + # Provider-specific extras + extras: dict[str, Any] = field(default_factory=dict) + + +@dataclass +class BatchContext: + """Context for an entire batch submission. + + Stores all request contexts so we can handle CCR tool calls + when results are retrieved. + """ + + batch_id: str + provider: str # "anthropic", "openai", "google" + created_at: float = field(default_factory=time.time) + expires_at: float = 0 + + # Map of custom_id -> BatchRequestContext + requests: dict[str, BatchRequestContext] = field(default_factory=dict) + + # API configuration for continuation calls + api_key: str | None = None + api_base_url: str | None = None + + def __post_init__(self) -> None: + if self.expires_at == 0: + self.expires_at = self.created_at + DEFAULT_BATCH_CONTEXT_TTL + + @property + def is_expired(self) -> bool: + """Check if this context has expired.""" + return time.time() > self.expires_at + + def add_request(self, request: BatchRequestContext) -> None: + """Add a request context to this batch.""" + self.requests[request.custom_id] = request + + def get_request(self, custom_id: str) -> BatchRequestContext | None: + """Get a request context by custom_id.""" + return self.requests.get(custom_id) + + +class BatchContextStore: + """Thread-safe store for batch contexts. + + Stores batch submission contexts with TTL so that when results + are retrieved, we can handle CCR tool calls. + + Features: + - TTL-based expiration + - Automatic cleanup of expired entries + - Thread-safe operations + - Memory limits + + Usage: + store = BatchContextStore() + + # On batch submit + context = BatchContext(batch_id="batch_123", provider="anthropic") + for req in batch_requests: + context.add_request(BatchRequestContext( + custom_id=req["custom_id"], + messages=req["params"]["messages"], + tools=req["params"].get("tools"), + model=req["params"]["model"], + )) + store.store(context) + + # On batch results retrieval + context = store.get("batch_123") + if context: + for result in results: + req_ctx = context.get_request(result["custom_id"]) + # ... handle CCR tool calls using req_ctx + """ + + def __init__( + self, + ttl: int = DEFAULT_BATCH_CONTEXT_TTL, + max_contexts: int = MAX_BATCH_CONTEXTS, + ) -> None: + self._contexts: dict[str, BatchContext] = {} + self._ttl = ttl + self._max_contexts = max_contexts + self._lock = asyncio.Lock() + self._cleanup_task: asyncio.Task | None = None + + async def store(self, context: BatchContext) -> None: + """Store a batch context. + + Args: + context: The batch context to store. + """ + async with self._lock: + # Enforce memory limit + if len(self._contexts) >= self._max_contexts: + # Remove oldest entries + await self._cleanup_oldest() + + # Set expiration + context.expires_at = time.time() + self._ttl + self._contexts[context.batch_id] = context + + logger.debug( + f"Stored batch context {context.batch_id} with " + f"{len(context.requests)} requests (provider={context.provider})" + ) + + async def get(self, batch_id: str) -> BatchContext | None: + """Get a batch context by ID. + + Args: + batch_id: The batch ID to look up. + + Returns: + The batch context, or None if not found or expired. + """ + async with self._lock: + context = self._contexts.get(batch_id) + + if context is None: + return None + + if context.is_expired: + del self._contexts[batch_id] + logger.debug(f"Batch context {batch_id} expired and removed") + return None + + return context + + async def remove(self, batch_id: str) -> bool: + """Remove a batch context. + + Args: + batch_id: The batch ID to remove. + + Returns: + True if removed, False if not found. + """ + async with self._lock: + if batch_id in self._contexts: + del self._contexts[batch_id] + return True + return False + + async def cleanup_expired(self) -> int: + """Remove all expired entries. + + Returns: + Number of entries removed. + """ + async with self._lock: + now = time.time() + expired = [batch_id for batch_id, ctx in self._contexts.items() if ctx.expires_at < now] + + for batch_id in expired: + del self._contexts[batch_id] + + if expired: + logger.debug(f"Cleaned up {len(expired)} expired batch contexts") + + return len(expired) + + async def _cleanup_oldest(self) -> None: + """Remove oldest entries to make room for new ones.""" + # Sort by creation time, remove oldest 10% + if not self._contexts: + return + + sorted_entries = sorted( + self._contexts.items(), + key=lambda x: x[1].created_at, + ) + + to_remove = max(1, len(sorted_entries) // 10) + for batch_id, _ in sorted_entries[:to_remove]: + del self._contexts[batch_id] + + logger.debug(f"Cleaned up {to_remove} oldest batch contexts") + + def stats(self) -> dict[str, Any]: + """Get store statistics.""" + return { + "total_contexts": len(self._contexts), + "max_contexts": self._max_contexts, + "ttl_seconds": self._ttl, + "providers": self._count_by_provider(), + } + + def _count_by_provider(self) -> dict[str, int]: + """Count contexts by provider.""" + counts: dict[str, int] = {} + for ctx in self._contexts.values(): + counts[ctx.provider] = counts.get(ctx.provider, 0) + 1 + return counts + + +# Global store instance +_batch_context_store: BatchContextStore | None = None + + +def get_batch_context_store() -> BatchContextStore: + """Get the global batch context store instance.""" + global _batch_context_store + if _batch_context_store is None: + _batch_context_store = BatchContextStore() + return _batch_context_store + + +def reset_batch_context_store() -> None: + """Reset the global batch context store (for testing).""" + global _batch_context_store + _batch_context_store = None diff --git a/headroom/ccr/response_handler.py b/headroom/ccr/response_handler.py index 9589e706d..6bfd8fe4a 100644 --- a/headroom/ccr/response_handler.py +++ b/headroom/ccr/response_handler.py @@ -113,7 +113,9 @@ class CCRResponseHandler: """ tool_calls = self._extract_tool_calls(response, provider) return any( - tc.get("name") == CCR_TOOL_NAME or tc.get("function", {}).get("name") == CCR_TOOL_NAME + tc.get("name") == CCR_TOOL_NAME + or tc.get("function", {}).get("name") == CCR_TOOL_NAME + or tc.get("functionCall", {}).get("name") == CCR_TOOL_NAME # Google format for tc in tool_calls ) @@ -136,6 +138,15 @@ class CCRResponseHandler: tool_calls = message.get("tool_calls", []) return list(tool_calls) if tool_calls else [] + elif provider == "google": + # Google/Gemini format: candidates[0].content.parts contains functionCall objects + # Each part with a functionCall has: {"functionCall": {"name": "...", "args": {...}}} + candidates = response.get("candidates", []) + if not candidates: + return [] + parts = candidates[0].get("content", {}).get("parts", []) + return [part for part in parts if "functionCall" in part] + return [] def _parse_ccr_tool_calls( @@ -157,8 +168,14 @@ class CCRResponseHandler: hash_key, query = parse_tool_call(tc, provider) if hash_key is not None: - # This is a CCR tool call - tool_call_id = tc.get("id", "") + # This is a CCR tool call - extract tool_call_id based on provider + if provider == "google": + # Google uses function name as identifier for matching responses + # The functionResponse.name must match the functionCall.name + tool_call_id = tc.get("functionCall", {}).get("name", CCR_TOOL_NAME) + else: + # Anthropic and OpenAI use explicit IDs + tool_call_id = tc.get("id", "") ccr_calls.append( CCRToolCall( tool_call_id=tool_call_id, @@ -295,6 +312,29 @@ class CCRResponseHandler: ] } + elif provider == "google": + # Google/Gemini: user message with functionResponse parts + # Format: {"role": "user", "parts": [{"functionResponse": {"name": "...", "response": {...}}}]} + parts = [] + for result in results: + # Parse the content JSON to include as response object + try: + response_data = json.loads(result.content) + except json.JSONDecodeError: + response_data = {"content": result.content} + parts.append( + { + "functionResponse": { + "name": result.tool_call_id, # tool_call_id contains the function name for Google + "response": response_data, + } + } + ) + return { + "role": "user", + "parts": parts, + } + else: # Generic format return { @@ -330,6 +370,17 @@ class CCRResponseHandler: "content": message.get("content"), "tool_calls": message.get("tool_calls"), } + elif provider == "google": + # Google/Gemini format: role is "model", content is in candidates[0].content.parts + candidates = response.get("candidates", []) + if candidates: + parts = candidates[0].get("content", {}).get("parts", []) + else: + parts = [] + return { + "role": "model", + "parts": parts, + } else: return { "role": "assistant", diff --git a/headroom/ccr/tool_injection.py b/headroom/ccr/tool_injection.py index cbf21ab1b..5a873e6bb 100644 --- a/headroom/ccr/tool_injection.py +++ b/headroom/ccr/tool_injection.py @@ -194,10 +194,23 @@ class CCRToolInjector: # Detected compression markers _detected_hashes: list[str] = field(default_factory=list) - _marker_pattern: re.Pattern = field( - default_factory=lambda: re.compile( - r"\[(\d+) items compressed to (\d+)\. Retrieve more: hash=([a-f0-9]+)\]" - ) + # Multiple marker patterns to match different compressors: + # - SmartCrusher: [100 items compressed to 10. Retrieve more: hash=abc123] + # - LLMLingua: [1000 items compressed to 300. Retrieve more: hash=abc123] + # - TextCompressor: [100 lines compressed to 10. Retrieve more: hash=abc123] + # - LogCompressor: [200 lines compressed to 20. Retrieve more: hash=abc123] + # - SearchCompressor: [50 matches compressed to 5. Retrieve more: hash=abc123] + # - Generic: any [... compressed ... hash=xxx] pattern + _marker_patterns: list[re.Pattern] = field( + default_factory=lambda: [ + # Standard format: [N compressed to M. Retrieve more: hash=xxx] + # Matches items, lines, matches, or any other type + re.compile(r"\[(\d+) \w+ compressed to (\d+)\. Retrieve more: hash=([a-f0-9]+)\]"), + # Legacy format without "to M" or "Retrieve more:" (old TextCompressor) + re.compile(r"\[(\d+) \w+ compressed\. hash=([a-f0-9]+)\]"), + # Generic fallback: any compression marker with hash (8+ chars) + re.compile(r"\[.*?compressed.*?hash=([a-f0-9]{8,})\]", re.IGNORECASE), + ] ) def __post_init__(self) -> None: @@ -249,14 +262,39 @@ class CCRToolInjector: if isinstance(item, dict) and item.get("type") == "text": self._scan_text(item.get("text", "")) + # Handle Google/Gemini format with parts + parts = message.get("parts", []) + if isinstance(parts, list): + for part in parts: + if isinstance(part, dict): + # Text parts + if "text" in part: + self._scan_text(part.get("text", "")) + # Function response parts (tool results) + elif "functionResponse" in part: + response = part.get("functionResponse", {}).get("response", {}) + if isinstance(response, str): + self._scan_text(response) + elif isinstance(response, dict): + # Scan string values in response + for value in response.values(): + if isinstance(value, str): + self._scan_text(value) + return self._detected_hashes def _scan_text(self, text: str) -> None: - """Scan text for compression markers.""" - matches = self._marker_pattern.findall(text) - for _original, _compressed, hash_key in matches: - if hash_key not in self._detected_hashes: - self._detected_hashes.append(hash_key) + """Scan text for compression markers from any compressor.""" + for pattern in self._marker_patterns: + matches = pattern.findall(text) + for match in matches: + # Extract hash_key from match (last group is always the hash) + if isinstance(match, tuple): + hash_key = match[-1] # Last capture group is the hash + else: + hash_key = match # Single capture group (generic pattern) + if hash_key and hash_key not in self._detected_hashes: + self._detected_hashes.append(hash_key) def inject_tool_definition( self, @@ -389,7 +427,7 @@ def parse_tool_call( Returns: Tuple of (hash, query) or (None, None) if not a CCR tool call. """ - # Get tool name + # Get tool name and input data based on provider format if provider == "anthropic": name = tool_call.get("name") input_data = tool_call.get("input", {}) @@ -402,7 +440,13 @@ def parse_tool_call( input_data = json.loads(args_str) except json.JSONDecodeError: input_data = {} + elif provider == "google": + # Google/Gemini format: {"functionCall": {"name": "...", "args": {...}}} + function_call = tool_call.get("functionCall", {}) + name = function_call.get("name") + input_data = function_call.get("args", {}) else: + # Generic fallback name = tool_call.get("name") input_data = tool_call.get("input", tool_call.get("args", {})) diff --git a/headroom/proxy/server.py b/headroom/proxy/server.py index b02c79290..4e689a320 100644 --- a/headroom/proxy/server.py +++ b/headroom/proxy/server.py @@ -57,11 +57,16 @@ from headroom.cache.compression_feedback import get_compression_feedback from headroom.cache.compression_store import get_compression_store from headroom.ccr import ( CCR_TOOL_NAME, + # Batch processing + BatchContext, + BatchRequestContext, + BatchResultProcessor, CCRResponseHandler, CCRToolInjector, ContextTracker, ContextTrackerConfig, ResponseHandlerConfig, + get_batch_context_store, parse_tool_call, ) from headroom.config import CacheAlignerConfig, CCRConfig, RollingWindowConfig, SmartCrusherConfig @@ -171,6 +176,7 @@ class ProxyConfig: host: str = "127.0.0.1" port: int = 8787 openai_api_url: str | None = None # Custom OpenAI API URL override + gemini_api_url: str | None = None # Custom Gemini API URL override # Optimization optimize: bool = True @@ -747,6 +753,7 @@ class HeadroomProxy: ANTHROPIC_API_URL = "https://api.anthropic.com" OPENAI_API_URL = "https://api.openai.com" + GEMINI_API_URL = "https://generativelanguage.googleapis.com" def __init__(self, config: ProxyConfig): self.config = config @@ -755,6 +762,10 @@ class HeadroomProxy: if config.openai_api_url: HeadroomProxy.OPENAI_API_URL = config.openai_api_url + # Override GEMINI_API_URL with config if set + if config.gemini_api_url: + HeadroomProxy.GEMINI_API_URL = config.gemini_api_url + # Initialize providers self.anthropic_provider = AnthropicProvider() self.openai_provider = OpenAIProvider() @@ -1579,6 +1590,994 @@ class HeadroomProxy: }, ) + async def handle_anthropic_batch_create( + self, + request: Request, + ) -> Response: + """Handle Anthropic POST /v1/messages/batches endpoint with compression. + + Anthropic batch format: + { + "requests": [ + { + "custom_id": "req-1", + "params": { + "model": "claude-sonnet-4-20250514", + "max_tokens": 1024, + "messages": [{"role": "user", "content": "Hello"}] + } + }, + ... + ] + } + + This method applies compression to each request's messages before forwarding. + """ + start_time = time.time() + request_id = await self._next_request_id() + + # Check request body size + content_length = request.headers.get("content-length") + if content_length and int(content_length) > MAX_REQUEST_BODY_SIZE: + return JSONResponse( + status_code=413, + content={ + "type": "error", + "error": { + "type": "request_too_large", + "message": f"Request body too large. Maximum size is {MAX_REQUEST_BODY_SIZE // (1024 * 1024)}MB", + }, + }, + ) + + # Parse request + try: + body = await request.json() + except json.JSONDecodeError as e: + return JSONResponse( + status_code=400, + content={ + "type": "error", + "error": { + "type": "invalid_request_error", + "message": f"Invalid JSON in request body: {e!s}", + }, + }, + ) + + requests_list = body.get("requests", []) + if not requests_list: + return JSONResponse( + status_code=400, + content={ + "type": "error", + "error": { + "type": "invalid_request_error", + "message": "Missing or empty 'requests' field in batch request", + }, + }, + ) + + # Extract headers + headers = dict(request.headers.items()) + headers.pop("host", None) + headers.pop("content-length", None) + + # Track compression stats across all batch requests + total_original_tokens = 0 + total_optimized_tokens = 0 + total_tokens_saved = 0 + compressed_requests = [] + + # Apply compression to each request in the batch + for batch_req in requests_list: + custom_id = batch_req.get("custom_id", "") + params = batch_req.get("params", {}) + messages = params.get("messages", []) + model = params.get("model", "unknown") + + if not messages or not self.config.optimize: + # No messages or optimization disabled - pass through unchanged + compressed_requests.append(batch_req) + continue + + # Count original tokens + tokenizer = get_tokenizer(model) + original_tokens = sum(tokenizer.count_text(str(m.get("content", ""))) for m in messages) + total_original_tokens += original_tokens + + # Apply optimization + try: + context_limit = self.anthropic_provider.get_context_limit(model) + result = self.anthropic_pipeline.apply( + messages=messages, + model=model, + model_limit=context_limit, + ) + + optimized_messages = result.messages + optimized_tokens = sum( + tokenizer.count_text(str(m.get("content", ""))) for m in optimized_messages + ) + total_optimized_tokens += optimized_tokens + tokens_saved = original_tokens - optimized_tokens + total_tokens_saved += tokens_saved + + # CCR Tool Injection: Inject retrieval tool if compression occurred + tools = params.get("tools") + if self.config.ccr_inject_tool and tokens_saved > 0: + injector = CCRToolInjector( + provider="anthropic", + inject_tool=True, + inject_system_instructions=self.config.ccr_inject_system_instructions, + ) + optimized_messages, tools, was_injected = injector.process_request( + optimized_messages, tools + ) + if was_injected: + logger.debug( + f"[{request_id}] CCR: Injected retrieval tool for batch request '{custom_id}'" + ) + + # Create compressed batch request + compressed_params = {**params, "messages": optimized_messages} + if tools is not None: + compressed_params["tools"] = tools + compressed_requests.append( + { + "custom_id": custom_id, + "params": compressed_params, + } + ) + + if tokens_saved > 0: + logger.debug( + f"[{request_id}] Batch request '{custom_id}': " + f"{original_tokens:,} -> {optimized_tokens:,} tokens " + f"(saved {tokens_saved:,})" + ) + + except Exception as e: + logger.warning( + f"[{request_id}] Optimization failed for batch request '{custom_id}': {e}" + ) + # Pass through unchanged on failure + compressed_requests.append(batch_req) + total_optimized_tokens += original_tokens + + # Update body with compressed requests + body["requests"] = compressed_requests + + optimization_latency = (time.time() - start_time) * 1000 + + # Forward request to Anthropic + url = f"{self.ANTHROPIC_API_URL}/v1/messages/batches" + + try: + response = await self._retry_request("POST", url, headers, body) + + # Record metrics + await self.metrics.record_request( + provider="anthropic", + model="batch", + input_tokens=total_optimized_tokens, + output_tokens=0, + tokens_saved=total_tokens_saved, + latency_ms=optimization_latency, + ) + + # Log compression stats + if total_tokens_saved > 0: + savings_percent = ( + (total_tokens_saved / total_original_tokens * 100) + if total_original_tokens > 0 + else 0 + ) + logger.info( + f"[{request_id}] Batch ({len(compressed_requests)} requests): " + f"{total_original_tokens:,} -> {total_optimized_tokens:,} tokens " + f"(saved {total_tokens_saved:,}, {savings_percent:.1f}%)" + ) + + # Store batch context for CCR result processing + if response.status_code == 200 and self.config.ccr_inject_tool: + try: + response_data = response.json() + batch_id = response_data.get("id") + if batch_id: + await self._store_anthropic_batch_context( + batch_id, + requests_list, + headers.get("x-api-key"), + ) + except Exception as e: + logger.warning(f"[{request_id}] Failed to store batch context: {e}") + + # Remove compression headers + response_headers = dict(response.headers) + response_headers.pop("content-encoding", None) + response_headers.pop("content-length", None) + + return Response( + content=response.content, + status_code=response.status_code, + headers=response_headers, + ) + + except Exception as e: + await self.metrics.record_failed() + logger.error(f"[{request_id}] Batch request failed: {type(e).__name__}: {e}") + return JSONResponse( + status_code=502, + content={ + "type": "error", + "error": { + "type": "api_error", + "message": "An error occurred while processing your batch request. Please try again.", + }, + }, + ) + + async def handle_anthropic_batch_passthrough( + self, + request: Request, + batch_id: str | None = None, + ) -> Response: + """Handle Anthropic batch passthrough endpoints. + + Used for: + - GET /v1/messages/batches - List batches + - GET /v1/messages/batches/{batch_id} - Get batch + - GET /v1/messages/batches/{batch_id}/results - Get batch results + - POST /v1/messages/batches/{batch_id}/cancel - Cancel batch + """ + start_time = time.time() + path = request.url.path + url = f"{self.ANTHROPIC_API_URL}{path}" + + # Preserve query string parameters (e.g., limit, after_id for list endpoint) + if request.url.query: + url = f"{url}?{request.url.query}" + + headers = dict(request.headers.items()) + headers.pop("host", None) + + body = await request.body() + + response = await self.http_client.request( + method=request.method, + url=url, + headers=headers, + content=body, + ) + + # Track metrics + latency_ms = (time.time() - start_time) * 1000 + await self.metrics.record_request( + provider="anthropic", + model="passthrough:batches", + input_tokens=0, + output_tokens=0, + tokens_saved=0, + latency_ms=latency_ms, + ) + + # Remove compression headers + response_headers = dict(response.headers) + response_headers.pop("content-encoding", None) + response_headers.pop("content-length", None) + + return Response( + content=response.content, + status_code=response.status_code, + headers=response_headers, + ) + + async def _store_anthropic_batch_context( + self, + batch_id: str, + requests_list: list[dict[str, Any]], + api_key: str | None, + ) -> None: + """Store batch context for CCR result processing. + + Args: + batch_id: The batch ID from the API response. + requests_list: The original batch requests. + api_key: The API key for continuation calls. + """ + store = get_batch_context_store() + context = BatchContext( + batch_id=batch_id, + provider="anthropic", + api_key=api_key, + api_base_url=self.ANTHROPIC_API_URL, + ) + + for batch_req in requests_list: + custom_id = batch_req.get("custom_id", "") + params = batch_req.get("params", {}) + context.add_request( + BatchRequestContext( + custom_id=custom_id, + messages=params.get("messages", []), + tools=params.get("tools"), + model=params.get("model", ""), + extras={ + "max_tokens": params.get("max_tokens", 4096), + "system": params.get("system"), + }, + ) + ) + + await store.store(context) + logger.debug(f"Stored batch context for {batch_id} with {len(requests_list)} requests") + + async def handle_anthropic_batch_results( + self, + request: Request, + batch_id: str, + ) -> Response: + """Handle Anthropic batch results with CCR post-processing. + + This endpoint: + 1. Fetches raw results from Anthropic + 2. Detects CCR tool calls in each result + 3. Executes retrieval and makes continuation calls + 4. Returns processed results with complete responses + """ + start_time = time.time() + + # Forward request to get raw results + url = f"{self.ANTHROPIC_API_URL}/v1/messages/batches/{batch_id}/results" + + if request.url.query: + url = f"{url}?{request.url.query}" + + headers = dict(request.headers.items()) + headers.pop("host", None) + + response = await self.http_client.get(url, headers=headers) + + if response.status_code != 200: + # Error - pass through + response_headers = dict(response.headers) + response_headers.pop("content-encoding", None) + response_headers.pop("content-length", None) + return Response( + content=response.content, + status_code=response.status_code, + headers=response_headers, + ) + + # Parse results - Anthropic batch results are JSONL format + raw_content = response.content.decode("utf-8") + results = [] + for line in raw_content.strip().split("\n"): + if line.strip(): + try: + results.append(json.loads(line)) + except json.JSONDecodeError: + continue + + if not results: + # No results to process + response_headers = dict(response.headers) + response_headers.pop("content-encoding", None) + response_headers.pop("content-length", None) + return Response( + content=response.content, + status_code=response.status_code, + headers=response_headers, + ) + + # Check if we have context and CCR processing is enabled + store = get_batch_context_store() + batch_context = await store.get(batch_id) + + if batch_context is None or not self.config.ccr_inject_tool: + # No context or CCR disabled - pass through + response_headers = dict(response.headers) + response_headers.pop("content-encoding", None) + response_headers.pop("content-length", None) + return Response( + content=response.content, + status_code=response.status_code, + headers=response_headers, + ) + + # Process results with CCR handler + processor = BatchResultProcessor(self.http_client) + processed = await processor.process_results(batch_id, results, "anthropic") + + # Convert back to JSONL format + processed_lines = [] + for p in processed: + processed_lines.append(json.dumps(p.result)) + if p.was_processed: + logger.info( + f"CCR: Processed batch result {p.custom_id} " + f"({p.continuation_rounds} continuation rounds)" + ) + + processed_content = "\n".join(processed_lines) + + # Track metrics + latency_ms = (time.time() - start_time) * 1000 + await self.metrics.record_request( + provider="anthropic", + model="batch:ccr-processed", + input_tokens=0, + output_tokens=0, + tokens_saved=0, + latency_ms=latency_ms, + ) + + return Response( + content=processed_content.encode("utf-8"), + status_code=200, + media_type="application/jsonl", + ) + + # ========================================================================= + # Google/Gemini Batch API Handlers + # ========================================================================= + + async def handle_google_batch_create( + self, + request: Request, + model: str, + ) -> Response: + """Handle Google POST /v1beta/models/{model}:batchGenerateContent endpoint. + + Google batch format: + { + "batch": { + "display_name": "my-batch", + "input_config": { + "requests": { + "requests": [ + { + "request": {"contents": [{"parts": [{"text": "..."}]}]}, + "metadata": {"key": "request-1"} + } + ] + } + } + } + } + + This method applies compression to each request's contents before forwarding. + """ + start_time = time.time() + request_id = await self._next_request_id() + + # Check request body size + content_length = request.headers.get("content-length") + if content_length and int(content_length) > MAX_REQUEST_BODY_SIZE: + return JSONResponse( + status_code=413, + content={ + "error": { + "code": 413, + "message": f"Request body too large. Maximum size is {MAX_REQUEST_BODY_SIZE // (1024 * 1024)}MB", + "status": "INVALID_ARGUMENT", + } + }, + ) + + # Parse request + try: + body = await request.json() + except json.JSONDecodeError as e: + return JSONResponse( + status_code=400, + content={ + "error": { + "code": 400, + "message": f"Invalid JSON in request body: {e!s}", + "status": "INVALID_ARGUMENT", + } + }, + ) + + # Extract batch config + batch_config = body.get("batch", {}) + input_config = batch_config.get("input_config", {}) + requests_wrapper = input_config.get("requests", {}) + requests_list = requests_wrapper.get("requests", []) + + if not requests_list: + # No inline requests - might be using file input, pass through + logger.debug(f"[{request_id}] Google batch: No inline requests, passing through") + return await self._google_batch_passthrough(request, model, body) + + # Extract headers + headers = dict(request.headers.items()) + headers.pop("host", None) + headers.pop("content-length", None) + + # Track compression stats + total_original_tokens = 0 + total_optimized_tokens = 0 + total_tokens_saved = 0 + compressed_requests = [] + + # Apply compression to each request in the batch + for idx, batch_req in enumerate(requests_list): + req_content = batch_req.get("request", {}) + metadata = batch_req.get("metadata", {}) + contents = req_content.get("contents", []) + + if not contents or not self.config.optimize: + # No contents or optimization disabled - pass through unchanged + compressed_requests.append(batch_req) + continue + + # Convert Google format to messages for compression + system_instruction = req_content.get("systemInstruction") + messages = self._gemini_contents_to_messages(contents, system_instruction) + + # Count original tokens + tokenizer = get_tokenizer(model) + original_tokens = sum(tokenizer.count_text(str(m.get("content", ""))) for m in messages) + total_original_tokens += original_tokens + + # Apply optimization + try: + # Default context limit for most models + context_limit = 128000 + + # Use OpenAI pipeline (similar message format after conversion) + result = self.openai_pipeline.apply( + messages=messages, + model=model, + model_limit=context_limit, + ) + + optimized_messages = result.messages + optimized_tokens = sum( + tokenizer.count_text(str(m.get("content", ""))) for m in optimized_messages + ) + total_optimized_tokens += optimized_tokens + tokens_saved = original_tokens - optimized_tokens + total_tokens_saved += tokens_saved + + # CCR Tool Injection: Inject retrieval tool if compression occurred + tools = req_content.get("tools") + # Extract existing function declarations if present + existing_funcs = None + if tools: + for tool in tools: + if "functionDeclarations" in tool: + existing_funcs = tool["functionDeclarations"] + break + + if self.config.ccr_inject_tool and tokens_saved > 0: + injector = CCRToolInjector( + provider="google", + inject_tool=True, + inject_system_instructions=self.config.ccr_inject_system_instructions, + ) + optimized_messages, injected_funcs, was_injected = injector.process_request( + optimized_messages, existing_funcs + ) + if was_injected: + logger.debug( + f"[{request_id}] CCR: Injected retrieval tool for Google batch request {idx}" + ) + existing_funcs = injected_funcs + + # Convert back to Google contents format + optimized_contents, optimized_sys_inst = self._messages_to_gemini_contents( + optimized_messages + ) + + # Create compressed batch request + compressed_req_content = {**req_content, "contents": optimized_contents} + if optimized_sys_inst: + compressed_req_content["systemInstruction"] = optimized_sys_inst + if existing_funcs is not None: + compressed_req_content["tools"] = [{"functionDeclarations": existing_funcs}] + + compressed_req = { + "request": compressed_req_content, + "metadata": metadata, + } + + compressed_requests.append(compressed_req) + + if tokens_saved > 0: + logger.debug( + f"[{request_id}] Google batch request {idx}: " + f"{original_tokens:,} -> {optimized_tokens:,} tokens " + f"(saved {tokens_saved:,})" + ) + + except Exception as e: + logger.warning( + f"[{request_id}] Optimization failed for Google batch request {idx}: {e}" + ) + # Pass through unchanged on failure + compressed_requests.append(batch_req) + total_optimized_tokens += original_tokens + + # Update body with compressed requests + body["batch"]["input_config"]["requests"]["requests"] = compressed_requests + + optimization_latency = (time.time() - start_time) * 1000 + + # Forward request to Google + url = f"{self.GEMINI_API_URL}/v1beta/models/{model}:batchGenerateContent" + + # Add API key to URL if present in headers + api_key = headers.pop("x-goog-api-key", None) + if api_key: + url = f"{url}?key={api_key}" + + try: + response = await self._retry_request("POST", url, headers, body) + + # Record metrics + await self.metrics.record_request( + provider="google", + model=f"batch:{model}", + input_tokens=total_optimized_tokens, + output_tokens=0, + tokens_saved=total_tokens_saved, + latency_ms=optimization_latency, + ) + + # Log compression stats + if total_tokens_saved > 0: + savings_percent = ( + (total_tokens_saved / total_original_tokens * 100) + if total_original_tokens > 0 + else 0 + ) + logger.info( + f"[{request_id}] Google batch compression: " + f"{total_original_tokens:,} -> {total_optimized_tokens:,} tokens " + f"({savings_percent:.1f}% saved across {len(requests_list)} requests)" + ) + + # Store batch context for CCR result processing + if response.status_code == 200 and self.config.ccr_inject_tool: + try: + response_data = response.json() + batch_name = response_data.get("name") + if batch_name: + await self._store_google_batch_context( + batch_name, + requests_list, + model, + api_key, + ) + except Exception as e: + logger.warning(f"[{request_id}] Failed to store Google batch context: {e}") + + # Remove compression headers + response_headers = dict(response.headers) + response_headers.pop("content-encoding", None) + response_headers.pop("content-length", None) + + return Response( + content=response.content, + status_code=response.status_code, + headers=response_headers, + ) + + except Exception as e: + logger.error(f"[{request_id}] Google batch request failed: {e}") + return JSONResponse( + status_code=500, + content={ + "error": { + "code": 500, + "message": f"Failed to forward batch request: {e!s}", + "status": "INTERNAL", + } + }, + ) + + async def _google_batch_passthrough( + self, + request: Request, + model: str, + body: dict | None = None, + ) -> Response: + """Pass through Google batch request without modification.""" + start_time = time.time() + + headers = dict(request.headers.items()) + headers.pop("host", None) + headers.pop("content-length", None) + + url = f"{self.GEMINI_API_URL}/v1beta/models/{model}:batchGenerateContent" + + # Add API key to URL if present in headers + api_key = headers.pop("x-goog-api-key", None) + if api_key: + url = f"{url}?key={api_key}" + + if body is None: + body_content = await request.body() + else: + body_content = json.dumps(body).encode() + + response = await self.http_client.post( + url, + headers=headers, + content=body_content, + ) + + # Track metrics + latency_ms = (time.time() - start_time) * 1000 + await self.metrics.record_request( + provider="google", + model=f"passthrough:batch:{model}", + input_tokens=0, + output_tokens=0, + tokens_saved=0, + latency_ms=latency_ms, + ) + + response_headers = dict(response.headers) + response_headers.pop("content-encoding", None) + response_headers.pop("content-length", None) + + return Response( + content=response.content, + status_code=response.status_code, + headers=response_headers, + ) + + async def handle_google_batch_passthrough( + self, + request: Request, + batch_name: str | None = None, + ) -> Response: + """Handle Google batch passthrough endpoints. + + Used for: + - GET /v1beta/batches/{batch_name} - Get batch status + - POST /v1beta/batches/{batch_name}:cancel - Cancel batch + - DELETE /v1beta/batches/{batch_name} - Delete batch + """ + start_time = time.time() + path = request.url.path + url = f"{self.GEMINI_API_URL}{path}" + + # Preserve query string parameters + if request.url.query: + url = f"{url}?{request.url.query}" + + headers = dict(request.headers.items()) + headers.pop("host", None) + + # Handle API key + api_key = headers.pop("x-goog-api-key", None) + if api_key: + if "?" in url: + url = f"{url}&key={api_key}" + else: + url = f"{url}?key={api_key}" + + body = await request.body() + + response = await self.http_client.request( + method=request.method, + url=url, + headers=headers, + content=body, + ) + + # Track metrics + latency_ms = (time.time() - start_time) * 1000 + await self.metrics.record_request( + provider="google", + model="passthrough:batches", + input_tokens=0, + output_tokens=0, + tokens_saved=0, + latency_ms=latency_ms, + ) + + response_headers = dict(response.headers) + response_headers.pop("content-encoding", None) + response_headers.pop("content-length", None) + + return Response( + content=response.content, + status_code=response.status_code, + headers=response_headers, + ) + + async def _store_google_batch_context( + self, + batch_name: str, + requests_list: list[dict[str, Any]], + model: str, + api_key: str | None, + ) -> None: + """Store Google batch context for CCR result processing. + + Args: + batch_name: The batch name from the API response. + requests_list: The original batch requests. + model: The model used for the batch. + api_key: The API key for continuation calls. + """ + store = get_batch_context_store() + context = BatchContext( + batch_id=batch_name, + provider="google", + api_key=api_key, + api_base_url=self.GEMINI_API_URL, + ) + + for batch_req in requests_list: + metadata = batch_req.get("metadata", {}) + custom_id = metadata.get("key", "") + req_content = batch_req.get("request", {}) + contents = req_content.get("contents", []) + system_instruction = req_content.get("systemInstruction") + + # Convert contents to messages format for CCR handler + messages = self._gemini_contents_to_messages(contents, system_instruction) + + # Extract system instruction text if present + sys_text = None + if system_instruction: + parts = system_instruction.get("parts", []) + if parts and isinstance(parts[0], dict): + sys_text = parts[0].get("text") + + context.add_request( + BatchRequestContext( + custom_id=custom_id, + messages=messages, + tools=req_content.get("tools"), + model=model, + system_instruction=sys_text, + ) + ) + + await store.store(context) + logger.debug( + f"Stored Google batch context for {batch_name} with {len(requests_list)} requests" + ) + + async def handle_google_batch_results( + self, + request: Request, + batch_name: str, + ) -> Response: + """Handle Google batch results with CCR post-processing. + + Google batch results endpoint returns the batch operation status. + When status is SUCCEEDED, results are embedded in the response. + This handler processes CCR tool calls in those results. + """ + start_time = time.time() + + # Forward request to get batch status/results + url = f"{self.GEMINI_API_URL}/v1beta/{batch_name}" + + if request.url.query: + url = f"{url}?{request.url.query}" + + headers = dict(request.headers.items()) + headers.pop("host", None) + + # Handle API key + api_key = headers.pop("x-goog-api-key", None) + if api_key: + if "?" in url: + url = f"{url}&key={api_key}" + else: + url = f"{url}?key={api_key}" + + response = await self.http_client.get(url, headers=headers) + + if response.status_code != 200: + # Error - pass through + response_headers = dict(response.headers) + response_headers.pop("content-encoding", None) + response_headers.pop("content-length", None) + return Response( + content=response.content, + status_code=response.status_code, + headers=response_headers, + ) + + # Parse response + try: + response_data = response.json() + except json.JSONDecodeError: + # Not JSON - pass through + response_headers = dict(response.headers) + response_headers.pop("content-encoding", None) + response_headers.pop("content-length", None) + return Response( + content=response.content, + status_code=response.status_code, + headers=response_headers, + ) + + # Check if batch has results (state must be SUCCEEDED) + metadata = response_data.get("metadata", {}) + state = metadata.get("state") + + if state != "SUCCEEDED": + # Batch not complete - pass through + response_headers = dict(response.headers) + response_headers.pop("content-encoding", None) + response_headers.pop("content-length", None) + return Response( + content=response.content, + status_code=response.status_code, + headers=response_headers, + ) + + # Extract results from response + # Google embeds results in the batch response + results = response_data.get("response", {}).get("responses", []) + + if not results: + # No results to process + response_headers = dict(response.headers) + response_headers.pop("content-encoding", None) + response_headers.pop("content-length", None) + return Response( + content=response.content, + status_code=response.status_code, + headers=response_headers, + ) + + # Check if we have context and CCR processing is enabled + store = get_batch_context_store() + batch_context = await store.get(batch_name) + + if batch_context is None or not self.config.ccr_inject_tool: + # No context or CCR disabled - pass through + response_headers = dict(response.headers) + response_headers.pop("content-encoding", None) + response_headers.pop("content-length", None) + return Response( + content=response.content, + status_code=response.status_code, + headers=response_headers, + ) + + # Process results with CCR handler + processor = BatchResultProcessor(self.http_client) + processed = await processor.process_results(batch_name, results, "google") + + # Update response with processed results + processed_results = [p.result for p in processed] + response_data["response"]["responses"] = processed_results + + for p in processed: + if p.was_processed: + logger.info( + f"CCR: Processed Google batch result {p.custom_id} " + f"({p.continuation_rounds} continuation rounds)" + ) + + # Track metrics + latency_ms = (time.time() - start_time) * 1000 + await self.metrics.record_request( + provider="google", + model="batch:ccr-processed", + input_tokens=0, + output_tokens=0, + tokens_saved=0, + latency_ms=latency_ms, + ) + + return JSONResponse(content=response_data, status_code=200) + async def _stream_response( self, url: str, @@ -1878,11 +2877,29 @@ class HeadroomProxy: }, ) - async def handle_passthrough(self, request: Request, base_url: str) -> Response: - """Pass through request unchanged.""" + async def handle_passthrough( + self, + request: Request, + base_url: str, + endpoint_name: str | None = None, + provider: str | None = None, + ) -> Response: + """Pass through request unchanged. + + Args: + request: The incoming request + base_url: The upstream API base URL + endpoint_name: Optional name for stats tracking (e.g., "models", "embeddings") + provider: Optional provider name for stats (e.g., "openai", "anthropic", "gemini") + """ + start_time = time.time() path = request.url.path url = f"{base_url}{path}" + # Preserve query string parameters + if request.url.query: + url = f"{url}?{request.url.query}" + headers = dict(request.headers.items()) headers.pop("host", None) @@ -1900,12 +2917,1043 @@ class HeadroomProxy: response_headers.pop("content-encoding", None) response_headers.pop("content-length", None) # Length changed after decompression + # Track stats for passthrough requests + if endpoint_name and provider: + latency_ms = (time.time() - start_time) * 1000 + await self.metrics.record_request( + provider=provider, + model=f"passthrough:{endpoint_name}", + input_tokens=0, + output_tokens=0, + tokens_saved=0, + latency_ms=latency_ms, + cached=False, + cost_usd=0, + savings_usd=0, + ) + return Response( content=response.content, status_code=response.status_code, headers=response_headers, ) + # ========================================================================= + # OpenAI Batch API with Compression + # ========================================================================= + + async def handle_batch_create(self, request: Request) -> Response: + """Handle POST /v1/batches - Create a batch with compression. + + Flow: + 1. Parse request to get input_file_id + 2. Download the JSONL file content from OpenAI + 3. Parse each line and compress the messages + 4. Create a new compressed JSONL file + 5. Upload compressed file to OpenAI + 6. Create batch with the new compressed file_id + 7. Return batch object with compression stats in metadata + """ + start_time = time.time() + request_id = await self._next_request_id() + + # Parse request + try: + body = await request.json() + except json.JSONDecodeError as e: + return JSONResponse( + status_code=400, + content={ + "error": { + "message": f"Invalid JSON in request body: {e!s}", + "type": "invalid_request_error", + "code": "invalid_json", + } + }, + ) + + input_file_id = body.get("input_file_id") + endpoint = body.get("endpoint") + completion_window = body.get("completion_window", "24h") + metadata = body.get("metadata", {}) + + if not input_file_id: + return JSONResponse( + status_code=400, + content={ + "error": { + "message": "input_file_id is required", + "type": "invalid_request_error", + "code": "missing_parameter", + } + }, + ) + + if not endpoint: + return JSONResponse( + status_code=400, + content={ + "error": { + "message": "endpoint is required", + "type": "invalid_request_error", + "code": "missing_parameter", + } + }, + ) + + # Only compress chat completions endpoint + if endpoint != "/v1/chat/completions": + # Pass through for other endpoints + return await self._batch_passthrough(request, body) + + headers = dict(request.headers.items()) + headers.pop("host", None) + headers.pop("content-length", None) + + try: + # Step 1: Download the input file from OpenAI + logger.info(f"[{request_id}] Batch: Downloading input file {input_file_id}") + file_content = await self._download_openai_file(input_file_id, headers) + + if file_content is None: + return JSONResponse( + status_code=404, + content={ + "error": { + "message": f"Failed to download file {input_file_id}", + "type": "invalid_request_error", + "code": "file_not_found", + } + }, + ) + + # Step 2: Parse and compress each line + logger.info(f"[{request_id}] Batch: Compressing JSONL content") + compressed_lines, stats = await self._compress_batch_jsonl(file_content, request_id) + + if stats["total_requests"] == 0: + return JSONResponse( + status_code=400, + content={ + "error": { + "message": "No valid requests found in input file", + "type": "invalid_request_error", + "code": "empty_file", + } + }, + ) + + # Step 3: Create compressed JSONL content + compressed_content = "\n".join(compressed_lines) + + # Step 4: Upload compressed file to OpenAI + logger.info(f"[{request_id}] Batch: Uploading compressed file") + new_file_id = await self._upload_openai_file( + compressed_content, f"compressed_{input_file_id}.jsonl", headers + ) + + if new_file_id is None: + return JSONResponse( + status_code=500, + content={ + "error": { + "message": "Failed to upload compressed file", + "type": "server_error", + "code": "upload_failed", + } + }, + ) + + # Step 5: Create batch with compressed file + logger.info(f"[{request_id}] Batch: Creating batch with compressed file {new_file_id}") + + # Add compression stats to metadata + compression_metadata = { + **metadata, + "headroom_compressed": "true", + "headroom_original_file_id": input_file_id, + "headroom_total_requests": str(stats["total_requests"]), + "headroom_tokens_saved": str(stats["total_tokens_saved"]), + "headroom_original_tokens": str(stats["total_original_tokens"]), + "headroom_compressed_tokens": str(stats["total_compressed_tokens"]), + "headroom_savings_percent": f"{stats['savings_percent']:.1f}", + } + + batch_body = { + "input_file_id": new_file_id, + "endpoint": endpoint, + "completion_window": completion_window, + "metadata": compression_metadata, + } + + url = f"{self.OPENAI_API_URL}/v1/batches" + response = await self.http_client.post(url, json=batch_body, headers=headers) # type: ignore[union-attr] + + total_latency = (time.time() - start_time) * 1000 + + # Log compression stats + logger.info( + f"[{request_id}] Batch created: {stats['total_requests']} requests, " + f"{stats['total_original_tokens']:,} -> {stats['total_compressed_tokens']:,} tokens " + f"(saved {stats['total_tokens_saved']:,} tokens, {stats['savings_percent']:.1f}%) " + f"in {total_latency:.0f}ms" + ) + + # Record metrics + await self.metrics.record_request( + provider="openai", + model="batch", + input_tokens=stats["total_compressed_tokens"], + output_tokens=0, + tokens_saved=stats["total_tokens_saved"], + latency_ms=total_latency, + ) + + # Return response with compression info in headers + response_headers = dict(response.headers) + response_headers.pop("content-encoding", None) + response_headers.pop("content-length", None) + response_headers["x-headroom-tokens-saved"] = str(stats["total_tokens_saved"]) + response_headers["x-headroom-savings-percent"] = f"{stats['savings_percent']:.1f}" + + return Response( + content=response.content, + status_code=response.status_code, + headers=response_headers, + ) + + except Exception as e: + logger.error(f"[{request_id}] Batch creation failed: {type(e).__name__}: {e}") + await self.metrics.record_failed() + return JSONResponse( + status_code=500, + content={ + "error": { + "message": "An error occurred while processing the batch request", + "type": "server_error", + "code": "batch_processing_error", + } + }, + ) + + async def _download_openai_file(self, file_id: str, headers: dict) -> str | None: + """Download file content from OpenAI.""" + url = f"{self.OPENAI_API_URL}/v1/files/{file_id}/content" + try: + response = await self.http_client.get(url, headers=headers) # type: ignore[union-attr] + if response.status_code == 200: + return response.text + logger.error(f"Failed to download file {file_id}: {response.status_code}") + return None + except Exception as e: + logger.error(f"Error downloading file {file_id}: {e}") + return None + + async def _upload_openai_file(self, content: str, filename: str, headers: dict) -> str | None: + """Upload a file to OpenAI for batch processing.""" + url = f"{self.OPENAI_API_URL}/v1/files" + + # Prepare multipart form data + # We need to use httpx's files parameter for multipart upload + files = { + "file": (filename, content.encode("utf-8"), "application/jsonl"), + } + data = { + "purpose": "batch", + } + + # Remove content-type from headers (httpx will set it for multipart) + upload_headers = {k: v for k, v in headers.items() if k.lower() != "content-type"} + + try: + response = await self.http_client.post( # type: ignore[union-attr] + url, files=files, data=data, headers=upload_headers + ) + if response.status_code == 200: + result = response.json() + return result.get("id") + logger.error(f"Failed to upload file: {response.status_code} - {response.text}") + return None + except Exception as e: + logger.error(f"Error uploading file: {e}") + return None + + async def _compress_batch_jsonl(self, content: str, request_id: str) -> tuple[list[str], dict]: + """Compress messages in each line of a batch JSONL file. + + Returns: + Tuple of (compressed_lines, stats_dict) + """ + lines = content.strip().split("\n") + compressed_lines = [] + total_original_tokens = 0 + total_compressed_tokens = 0 + total_requests = 0 + errors = 0 + + tokenizer = get_tokenizer("gpt-4") # Use gpt-4 tokenizer for batch + + for i, line in enumerate(lines): + if not line.strip(): + continue + + try: + request_obj = json.loads(line) + body = request_obj.get("body", {}) + messages = body.get("messages", []) + model = body.get("model", "gpt-4") + + if not messages: + # No messages to compress, pass through + compressed_lines.append(line) + total_requests += 1 + continue + + # Count original tokens + original_tokens = sum( + tokenizer.count_text(str(m.get("content", ""))) for m in messages + ) + total_original_tokens += original_tokens + + # Compress messages using the OpenAI pipeline + if self.config.optimize: + try: + context_limit = self.openai_provider.get_context_limit(model) + result = self.openai_pipeline.apply( + messages=messages, + model=model, + model_limit=context_limit, + ) + compressed_messages = result.messages + except Exception as e: + logger.warning(f"[{request_id}] Compression failed for line {i}: {e}") + compressed_messages = messages + else: + compressed_messages = messages + + # Count compressed tokens + compressed_tokens = sum( + tokenizer.count_text(str(m.get("content", ""))) for m in compressed_messages + ) + total_compressed_tokens += compressed_tokens + tokens_saved = original_tokens - compressed_tokens + + # CCR Tool Injection: Inject retrieval tool if compression occurred + tools = body.get("tools") + if self.config.ccr_inject_tool and tokens_saved > 0: + injector = CCRToolInjector( + provider="openai", + inject_tool=True, + inject_system_instructions=self.config.ccr_inject_system_instructions, + ) + compressed_messages, tools, was_injected = injector.process_request( + compressed_messages, tools + ) + if was_injected: + logger.debug( + f"[{request_id}] CCR: Injected retrieval tool for batch line {i}" + ) + + # Update body with compressed messages + body["messages"] = compressed_messages + if tools is not None: + body["tools"] = tools + request_obj["body"] = body + + compressed_lines.append(json.dumps(request_obj)) + total_requests += 1 + + except json.JSONDecodeError as e: + logger.warning(f"[{request_id}] Invalid JSON on line {i}: {e}") + errors += 1 + # Keep original line on error + compressed_lines.append(line) + total_requests += 1 + + total_tokens_saved = total_original_tokens - total_compressed_tokens + savings_percent = ( + (total_tokens_saved / total_original_tokens * 100) if total_original_tokens > 0 else 0 + ) + + stats = { + "total_requests": total_requests, + "total_original_tokens": total_original_tokens, + "total_compressed_tokens": total_compressed_tokens, + "total_tokens_saved": total_tokens_saved, + "savings_percent": savings_percent, + "errors": errors, + } + + return compressed_lines, stats + + async def _batch_passthrough(self, request: Request, body: dict) -> Response: + """Pass through batch request to OpenAI without compression.""" + headers = dict(request.headers.items()) + headers.pop("host", None) + headers.pop("content-length", None) + + url = f"{self.OPENAI_API_URL}/v1/batches" + response = await self.http_client.post(url, json=body, headers=headers) # type: ignore[union-attr] + + response_headers = dict(response.headers) + response_headers.pop("content-encoding", None) + response_headers.pop("content-length", None) + + return Response( + content=response.content, + status_code=response.status_code, + headers=response_headers, + ) + + async def handle_batch_list(self, request: Request) -> Response: + """Handle GET /v1/batches - List batches (passthrough).""" + return await self.handle_passthrough(request, self.OPENAI_API_URL) + + async def handle_batch_get(self, request: Request, batch_id: str) -> Response: + """Handle GET /v1/batches/{batch_id} - Get batch (passthrough).""" + return await self.handle_passthrough(request, self.OPENAI_API_URL) + + async def handle_batch_cancel(self, request: Request, batch_id: str) -> Response: + """Handle POST /v1/batches/{batch_id}/cancel - Cancel batch (passthrough).""" + return await self.handle_passthrough(request, self.OPENAI_API_URL) + + def _gemini_contents_to_messages( + self, contents: list[dict], system_instruction: dict | None = None + ) -> list[dict]: + """Convert Gemini contents[] format to OpenAI messages[] format for optimization. + + Gemini format: + contents: [{"role": "user", "parts": [{"text": "..."}]}] + systemInstruction: {"parts": [{"text": "..."}]} + + OpenAI format: + messages: [{"role": "user", "content": "..."}] + """ + messages = [] + + # Add system instruction as system message + if system_instruction: + parts = system_instruction.get("parts", []) + text_parts = [p.get("text", "") for p in parts if "text" in p] + if text_parts: + messages.append({"role": "system", "content": "\n".join(text_parts)}) + + # Convert contents to messages + for content in contents: + role = content.get("role", "user") + # Map Gemini roles to OpenAI roles + if role == "model": + role = "assistant" + + parts = content.get("parts", []) + text_parts = [p.get("text", "") for p in parts if "text" in p] + + if text_parts: + messages.append({"role": role, "content": "\n".join(text_parts)}) + + return messages + + def _messages_to_gemini_contents(self, messages: list[dict]) -> tuple[list[dict], dict | None]: + """Convert OpenAI messages[] format back to Gemini contents[] format. + + Returns: + (contents, system_instruction) tuple + """ + contents = [] + system_instruction = None + + for msg in messages: + role = msg.get("role", "user") + content = msg.get("content", "") + + if role == "system": + # Extract as systemInstruction + system_instruction = {"parts": [{"text": content}]} + else: + # Map OpenAI roles to Gemini roles + gemini_role = "model" if role == "assistant" else "user" + contents.append({"role": gemini_role, "parts": [{"text": content}]}) + + return contents, system_instruction + + async def handle_openai_responses( + self, + request: Request, + ) -> Response | StreamingResponse: + """Handle OpenAI /v1/responses endpoint (new Responses API). + + The Responses API differs from /v1/chat/completions: + - Input: `input` (string or array) instead of `messages` + - System: `instructions` instead of system message + - Output: `output[]` array instead of `choices[].message` + - State: `previous_response_id` for multi-turn + - Built-in tools: web_search, file_search, code_interpreter + """ + start_time = time.time() + request_id = await self._next_request_id() + + # Check request body size + content_length = request.headers.get("content-length") + if content_length and int(content_length) > MAX_REQUEST_BODY_SIZE: + return JSONResponse( + status_code=413, + content={ + "error": { + "message": f"Request body too large. Maximum size is {MAX_REQUEST_BODY_SIZE // (1024 * 1024)}MB", + "type": "invalid_request_error", + "code": "request_too_large", + } + }, + ) + + # Parse request + try: + body = await request.json() + except json.JSONDecodeError as e: + return JSONResponse( + status_code=400, + content={ + "error": { + "message": f"Invalid JSON in request body: {e!s}", + "type": "invalid_request_error", + "code": "invalid_json", + } + }, + ) + + model = body.get("model", "unknown") + stream = body.get("stream", False) + + # Convert Responses API input to messages format for optimization + # The Responses API accepts either a string or array of messages + input_data = body.get("input", "") + instructions = body.get("instructions") + + messages = [] + if instructions: + messages.append({"role": "system", "content": instructions}) + + if isinstance(input_data, str): + messages.append({"role": "user", "content": input_data}) + elif isinstance(input_data, list): + # Input is already an array of message objects + messages.extend(input_data) + + headers = dict(request.headers.items()) + headers.pop("host", None) + headers.pop("content-length", None) + tags = self._extract_tags(headers) + + # Rate limiting + if self.rate_limiter: + rate_key = headers.get("authorization", "default")[:20] + allowed, wait_seconds = await self.rate_limiter.check_request(rate_key) + if not allowed: + await self.metrics.record_rate_limited() + raise HTTPException( + status_code=429, + detail=f"Rate limited. Retry after {wait_seconds:.1f}s", + ) + + # Token counting on converted messages + tokenizer = get_tokenizer(model) + original_tokens = sum(tokenizer.count_text(str(m.get("content", ""))) for m in messages) + + # Note: We pass through to OpenAI without optimization for now + # The Responses API has different semantics that may not work well with compression + tokens_saved = 0 + transforms_applied: list[str] = [] + optimization_latency = (time.time() - start_time) * 1000 + + url = f"{self.OPENAI_API_URL}/v1/responses" + + try: + if stream: + # Streaming for Responses API uses semantic events + return await self._stream_response( + url, + headers, + body, + "openai", + model, + request_id, + original_tokens, + original_tokens, + tokens_saved, + transforms_applied, + tags, + optimization_latency, + ) + else: + response = await self._retry_request("POST", url, headers, body) + total_latency = (time.time() - start_time) * 1000 + + output_tokens = 0 + try: + resp_json = response.json() + usage = resp_json.get("usage", {}) + output_tokens = usage.get("output_tokens", 0) + except Exception: + pass + + # Cost tracking + cost_usd = savings_usd = None + if self.cost_tracker: + cost_usd = self.cost_tracker.estimate_cost( + model, original_tokens, output_tokens + ) + if cost_usd: + self.cost_tracker.record_cost(cost_usd) + + # Metrics + await self.metrics.record_request( + provider="openai", + model=model, + input_tokens=original_tokens, + output_tokens=output_tokens, + tokens_saved=tokens_saved, + latency_ms=total_latency, + cost_usd=cost_usd or 0, + savings_usd=savings_usd or 0, + ) + + logger.info(f"[{request_id}] /v1/responses {model}: {original_tokens:,} tokens") + + # Remove compression headers + response_headers = dict(response.headers) + response_headers.pop("content-encoding", None) + response_headers.pop("content-length", None) + + return Response( + content=response.content, + status_code=response.status_code, + headers=response_headers, + ) + except Exception as e: + await self.metrics.record_failed() + logger.error(f"[{request_id}] OpenAI responses request failed: {type(e).__name__}: {e}") + return JSONResponse( + status_code=502, + content={ + "error": { + "message": "An error occurred while processing your request. Please try again.", + "type": "server_error", + "code": "proxy_error", + } + }, + ) + + async def handle_gemini_generate_content( + self, + request: Request, + model: str, + ) -> Response | StreamingResponse: + """Handle Gemini native /v1beta/models/{model}:generateContent endpoint. + + Gemini's native API differs from OpenAI: + - Input: `contents[]` with `parts[]` instead of `messages` + - System: `systemInstruction` instead of system message + - Auth: `x-goog-api-key` header instead of `Authorization: Bearer` + - Output: `candidates[].content.parts[].text` + """ + start_time = time.time() + request_id = await self._next_request_id() + + # Check request body size + content_length = request.headers.get("content-length") + if content_length and int(content_length) > MAX_REQUEST_BODY_SIZE: + return JSONResponse( + status_code=413, + content={ + "error": { + "message": f"Request body too large. Maximum size is {MAX_REQUEST_BODY_SIZE // (1024 * 1024)}MB", + "code": 413, + } + }, + ) + + # Parse request + try: + body = await request.json() + except json.JSONDecodeError as e: + return JSONResponse( + status_code=400, + content={ + "error": { + "message": f"Invalid JSON in request body: {e!s}", + "code": 400, + } + }, + ) + + contents = body.get("contents", []) + + headers = dict(request.headers.items()) + headers.pop("host", None) + headers.pop("content-length", None) + tags = self._extract_tags(headers) + + # Rate limiting (use Gemini API key) + if self.rate_limiter: + rate_key = headers.get("x-goog-api-key", "default")[:20] + allowed, wait_seconds = await self.rate_limiter.check_request(rate_key) + if not allowed: + await self.metrics.record_rate_limited() + raise HTTPException( + status_code=429, + detail=f"Rate limited. Retry after {wait_seconds:.1f}s", + ) + + # Convert Gemini format to messages for optimization + system_instruction = body.get("systemInstruction") + messages = self._gemini_contents_to_messages(contents, system_instruction) + + # Token counting + tokenizer = get_tokenizer(model) + original_tokens = sum(tokenizer.count_text(str(m.get("content", ""))) for m in messages) + + # Optimization + transforms_applied: list[str] = [] + optimized_messages = messages + optimized_tokens = original_tokens + + if self.config.optimize and messages: + try: + # Use OpenAI pipeline (similar message format) + context_limit = self.openai_provider.get_context_limit(model) + result = self.openai_pipeline.apply( + messages=messages, + model=model, + model_limit=context_limit, + ) + if result.messages != messages: + optimized_messages = result.messages + transforms_applied = result.transforms_applied + optimized_tokens = sum( + tokenizer.count_text(str(m.get("content", ""))) for m in optimized_messages + ) + except Exception as e: + logger.warning(f"[{request_id}] Gemini optimization failed: {e}") + + tokens_saved = original_tokens - optimized_tokens + optimization_latency = (time.time() - start_time) * 1000 + + # Convert back to Gemini format if optimized + if optimized_messages != messages: + optimized_contents, optimized_system = self._messages_to_gemini_contents( + optimized_messages + ) + body["contents"] = optimized_contents + if optimized_system: + body["systemInstruction"] = optimized_system + elif "systemInstruction" in body: + del body["systemInstruction"] + + # Build URL - model is extracted from path + url = f"{self.GEMINI_API_URL}/v1beta/models/{model}:generateContent" + + # Check if streaming requested via query param + query_params = dict(request.query_params) + is_streaming = query_params.get("alt") == "sse" + + # Preserve API key in query params if present + if "key" in query_params: + url += f"?key={query_params['key']}" + + try: + if is_streaming: + # For streaming, use streamGenerateContent endpoint + stream_url = ( + f"{self.GEMINI_API_URL}/v1beta/models/{model}:streamGenerateContent?alt=sse" + ) + if "key" in query_params: + stream_url = f"{self.GEMINI_API_URL}/v1beta/models/{model}:streamGenerateContent?key={query_params['key']}&alt=sse" + + return await self._stream_response( + stream_url, + headers, + body, + "gemini", + model, + request_id, + original_tokens, + optimized_tokens, + tokens_saved, + transforms_applied, + tags, + optimization_latency, + ) + else: + response = await self._retry_request("POST", url, headers, body) + total_latency = (time.time() - start_time) * 1000 + + output_tokens = 0 + try: + resp_json = response.json() + usage = resp_json.get("usageMetadata", {}) + output_tokens = usage.get("candidatesTokenCount", 0) + except Exception: + pass + + # Cost tracking + cost_usd = savings_usd = None + if self.cost_tracker: + cost_usd = self.cost_tracker.estimate_cost( + model, optimized_tokens, output_tokens + ) + original_cost = self.cost_tracker.estimate_cost( + model, original_tokens, output_tokens + ) + if cost_usd and original_cost: + savings_usd = original_cost - cost_usd + self.cost_tracker.record_cost(cost_usd) + self.cost_tracker.record_savings(savings_usd) + + # Metrics + await self.metrics.record_request( + provider="gemini", + model=model, + input_tokens=optimized_tokens, + output_tokens=output_tokens, + tokens_saved=tokens_saved, + latency_ms=total_latency, + cost_usd=cost_usd or 0, + savings_usd=savings_usd or 0, + ) + + if tokens_saved > 0: + logger.info( + f"[{request_id}] Gemini {model}: {original_tokens:,} → {optimized_tokens:,} " + f"(saved {tokens_saved:,} tokens)" + ) + else: + logger.info(f"[{request_id}] Gemini {model}: {original_tokens:,} tokens") + + # Remove compression headers + response_headers = dict(response.headers) + response_headers.pop("content-encoding", None) + response_headers.pop("content-length", None) + + return Response( + content=response.content, + status_code=response.status_code, + headers=response_headers, + ) + except Exception as e: + await self.metrics.record_failed() + logger.error(f"[{request_id}] Gemini request failed: {type(e).__name__}: {e}") + return JSONResponse( + status_code=502, + content={ + "error": { + "message": "An error occurred while processing your request. Please try again.", + "code": 502, + } + }, + ) + + async def handle_gemini_stream_generate_content( + self, + request: Request, + model: str, + ) -> StreamingResponse: + """Handle Gemini streaming endpoint /v1beta/models/{model}:streamGenerateContent.""" + start_time = time.time() + request_id = await self._next_request_id() + + # Parse request + try: + body = await request.json() + except json.JSONDecodeError as e: + return JSONResponse( + status_code=400, + content={ + "error": { + "message": f"Invalid JSON in request body: {e!s}", + "code": 400, + } + }, + ) + + contents = body.get("contents", []) + + headers = dict(request.headers.items()) + headers.pop("host", None) + headers.pop("content-length", None) + tags = self._extract_tags(headers) + + # Token counting + tokenizer = get_tokenizer(model) + original_tokens = 0 + for content in contents: + parts = content.get("parts", []) + for part in parts: + if "text" in part: + original_tokens += tokenizer.count_text(part["text"]) + + optimization_latency = (time.time() - start_time) * 1000 + + # Build URL with SSE param + query_params = dict(request.query_params) + url = f"{self.GEMINI_API_URL}/v1beta/models/{model}:streamGenerateContent?alt=sse" + if "key" in query_params: + url = f"{self.GEMINI_API_URL}/v1beta/models/{model}:streamGenerateContent?key={query_params['key']}&alt=sse" + + return await self._stream_response( + url, + headers, + body, + "gemini", + model, + request_id, + original_tokens, + original_tokens, + 0, # tokens_saved + [], # transforms_applied + tags, + optimization_latency, + ) + + async def handle_gemini_count_tokens( + self, + request: Request, + model: str, + ) -> Response: + """Handle Gemini /v1beta/models/{model}:countTokens endpoint with compression. + + This endpoint counts tokens AFTER applying compression, so users can see + how many tokens they'll actually use after optimization. + + The request format is the same as generateContent: + {"contents": [...], "systemInstruction": {...}} + """ + start_time = time.time() + request_id = await self._next_request_id() + + # Parse request + try: + body = await request.json() + except json.JSONDecodeError as e: + return JSONResponse( + status_code=400, + content={ + "error": { + "message": f"Invalid JSON in request body: {e!s}", + "code": 400, + } + }, + ) + + contents = body.get("contents", []) + + headers = dict(request.headers.items()) + headers.pop("host", None) + headers.pop("content-length", None) + + # Convert Gemini format to messages for optimization + system_instruction = body.get("systemInstruction") + messages = self._gemini_contents_to_messages(contents, system_instruction) + + # Token counting (original) + tokenizer = get_tokenizer(model) + original_tokens = sum(tokenizer.count_text(str(m.get("content", ""))) for m in messages) + + # Apply compression using the same pipeline as generateContent + transforms_applied: list[str] = [] + optimized_messages = messages + + if self.config.optimize and messages: + try: + context_limit = self.openai_provider.get_context_limit(model) + result = self.openai_pipeline.apply( + messages=messages, + model=model, + model_limit=context_limit, + ) + if result.messages != messages: + optimized_messages = result.messages + transforms_applied = result.transforms_applied + except Exception as e: + logger.warning(f"[{request_id}] Gemini countTokens optimization failed: {e}") + + # Convert back to Gemini format for the API call + if optimized_messages != messages: + optimized_contents, optimized_system = self._messages_to_gemini_contents( + optimized_messages + ) + body["contents"] = optimized_contents + if optimized_system: + body["systemInstruction"] = optimized_system + elif "systemInstruction" in body: + del body["systemInstruction"] + + # Build URL + url = f"{self.GEMINI_API_URL}/v1beta/models/{model}:countTokens" + + # Preserve API key in query params if present + query_params = dict(request.query_params) + if "key" in query_params: + url += f"?key={query_params['key']}" + + try: + response = await self._retry_request("POST", url, headers, body) + total_latency = (time.time() - start_time) * 1000 + + # Parse response to get token count + compressed_tokens = 0 + try: + resp_json = response.json() + compressed_tokens = resp_json.get("totalTokens", 0) + except Exception: + pass + + # Track stats + tokens_saved = original_tokens - compressed_tokens if compressed_tokens > 0 else 0 + + await self.metrics.record_request( + provider="gemini", + model=model, + input_tokens=compressed_tokens, + output_tokens=0, + tokens_saved=tokens_saved, + latency_ms=total_latency, + cost_usd=0, + savings_usd=0, + ) + + if tokens_saved > 0: + logger.info( + f"[{request_id}] Gemini countTokens {model}: {original_tokens:,} → {compressed_tokens:,} " + f"(saved {tokens_saved:,} tokens, transforms: {transforms_applied})" + ) + else: + logger.info( + f"[{request_id}] Gemini countTokens {model}: {compressed_tokens:,} tokens" + ) + + # Remove compression headers + response_headers = dict(response.headers) + response_headers.pop("content-encoding", None) + response_headers.pop("content-length", None) + + return Response( + content=response.content, + status_code=response.status_code, + headers=response_headers, + ) + except Exception as e: + await self.metrics.record_failed() + logger.error(f"[{request_id}] Gemini countTokens failed: {type(e).__name__}: {e}") + return JSONResponse( + status_code=502, + content={ + "error": { + "message": "An error occurred while processing your request. Please try again.", + "code": 502, + } + }, + ) + # ============================================================================= # FastAPI App @@ -2562,13 +4610,231 @@ def create_app(config: ProxyConfig | None = None) -> FastAPI: @app.post("/v1/messages/count_tokens") async def anthropic_count_tokens(request: Request): - return await proxy.handle_passthrough(request, proxy.ANTHROPIC_API_URL) + return await proxy.handle_passthrough( + request, proxy.ANTHROPIC_API_URL, "count_tokens", "anthropic" + ) + + # Anthropic Message Batches API endpoints + @app.post("/v1/messages/batches") + async def anthropic_batch_create(request: Request): + """Create a message batch with compression applied to all requests.""" + return await proxy.handle_anthropic_batch_create(request) + + @app.get("/v1/messages/batches") + async def anthropic_batch_list(request: Request): + """List message batches (passthrough).""" + return await proxy.handle_anthropic_batch_passthrough(request) + + @app.get("/v1/messages/batches/{batch_id}") + async def anthropic_batch_get(request: Request, batch_id: str): + """Get a specific message batch (passthrough).""" + return await proxy.handle_anthropic_batch_passthrough(request, batch_id) + + @app.get("/v1/messages/batches/{batch_id}/results") + async def anthropic_batch_results(request: Request, batch_id: str): + """Get results for a message batch with CCR post-processing.""" + return await proxy.handle_anthropic_batch_results(request, batch_id) + + @app.post("/v1/messages/batches/{batch_id}/cancel") + async def anthropic_batch_cancel(request: Request, batch_id: str): + """Cancel a message batch (passthrough).""" + return await proxy.handle_anthropic_batch_passthrough(request, batch_id) # OpenAI endpoints @app.post("/v1/chat/completions") async def openai_chat(request: Request): return await proxy.handle_openai_chat(request) + @app.post("/v1/responses") + async def openai_responses(request: Request): + """OpenAI Responses API (new API introduced March 2025).""" + return await proxy.handle_openai_responses(request) + + # OpenAI Batch API endpoints (with compression!) + @app.post("/v1/batches") + async def create_batch(request: Request): + """Create a batch with automatic compression of messages.""" + return await proxy.handle_batch_create(request) + + @app.get("/v1/batches") + async def list_batches(request: Request): + """List batches (passthrough to OpenAI).""" + return await proxy.handle_batch_list(request) + + @app.get("/v1/batches/{batch_id}") + async def get_batch(request: Request, batch_id: str): + """Get batch details (passthrough to OpenAI).""" + return await proxy.handle_batch_get(request, batch_id) + + @app.post("/v1/batches/{batch_id}/cancel") + async def cancel_batch(request: Request, batch_id: str): + """Cancel a batch (passthrough to OpenAI).""" + return await proxy.handle_batch_cancel(request, batch_id) + + # Gemini native endpoints + @app.post("/v1beta/models/{model}:generateContent") + async def gemini_generate_content(request: Request, model: str): + """Gemini native generateContent API.""" + return await proxy.handle_gemini_generate_content(request, model) + + @app.post("/v1beta/models/{model}:streamGenerateContent") + async def gemini_stream_generate_content(request: Request, model: str): + """Gemini native streaming generateContent API.""" + return await proxy.handle_gemini_stream_generate_content(request, model) + + @app.post("/v1beta/models/{model}:countTokens") + async def gemini_count_tokens(request: Request, model: str): + """Gemini countTokens API with compression applied.""" + return await proxy.handle_gemini_count_tokens(request, model) + + # ========================================================================= + # Passthrough Endpoints (no compression needed) + # ========================================================================= + + # --- OpenAI Passthrough Endpoints --- + + @app.get("/v1/models") + async def list_models(request: Request): + """List models - route based on auth header. + + - x-api-key header present -> Anthropic + - Authorization: Bearer header -> OpenAI + """ + if request.headers.get("x-api-key"): + return await proxy.handle_passthrough( + request, proxy.ANTHROPIC_API_URL, "models", "anthropic" + ) + return await proxy.handle_passthrough(request, proxy.OPENAI_API_URL, "models", "openai") + + @app.get("/v1/models/{model_id}") + async def get_model(request: Request, model_id: str): + """Get model details - route based on auth header. + + - x-api-key header present -> Anthropic + - Authorization: Bearer header -> OpenAI + """ + if request.headers.get("x-api-key"): + return await proxy.handle_passthrough( + request, proxy.ANTHROPIC_API_URL, "models", "anthropic" + ) + return await proxy.handle_passthrough(request, proxy.OPENAI_API_URL, "models", "openai") + + @app.post("/v1/embeddings") + async def openai_embeddings(request: Request): + """OpenAI embeddings API - passthrough.""" + return await proxy.handle_passthrough(request, proxy.OPENAI_API_URL, "embeddings", "openai") + + @app.post("/v1/moderations") + async def openai_moderations(request: Request): + """OpenAI moderations API - passthrough.""" + return await proxy.handle_passthrough( + request, proxy.OPENAI_API_URL, "moderations", "openai" + ) + + @app.post("/v1/images/generations") + async def openai_images_generations(request: Request): + """OpenAI image generation API - passthrough.""" + return await proxy.handle_passthrough( + request, proxy.OPENAI_API_URL, "images/generations", "openai" + ) + + @app.post("/v1/audio/transcriptions") + async def openai_audio_transcriptions(request: Request): + """OpenAI audio transcription API (multipart/form-data) - passthrough.""" + return await proxy.handle_passthrough( + request, proxy.OPENAI_API_URL, "audio/transcriptions", "openai" + ) + + @app.post("/v1/audio/speech") + async def openai_audio_speech(request: Request): + """OpenAI text-to-speech API - passthrough.""" + return await proxy.handle_passthrough( + request, proxy.OPENAI_API_URL, "audio/speech", "openai" + ) + + # --- Gemini Passthrough Endpoints --- + + @app.get("/v1beta/models") + async def gemini_list_models(request: Request): + """Gemini list models API - passthrough.""" + return await proxy.handle_passthrough(request, proxy.GEMINI_API_URL, "models", "gemini") + + @app.get("/v1beta/models/{model_name}") + async def gemini_get_model(request: Request, model_name: str): + """Gemini get model API - passthrough. + + Note: This handles GET /v1beta/models/{model_name} but NOT :countTokens + which is handled by a separate POST route above. + """ + return await proxy.handle_passthrough(request, proxy.GEMINI_API_URL, "models", "gemini") + + @app.post("/v1beta/models/{model}:embedContent") + async def gemini_embed_content(request: Request, model: str): + """Gemini embedding API - passthrough.""" + return await proxy.handle_passthrough( + request, proxy.GEMINI_API_URL, "embedContent", "gemini" + ) + + @app.post("/v1beta/models/{model}:batchEmbedContents") + async def gemini_batch_embed_contents(request: Request, model: str): + """Gemini batch embeddings API - passthrough.""" + return await proxy.handle_passthrough( + request, proxy.GEMINI_API_URL, "batchEmbedContents", "gemini" + ) + + # Google/Gemini Batch API endpoints (with compression!) + @app.post("/v1beta/models/{model}:batchGenerateContent") + async def gemini_batch_create(request: Request, model: str): + """Create a Gemini batch with compression applied to all requests.""" + return await proxy.handle_google_batch_create(request, model) + + @app.get("/v1beta/batches/{batch_name}") + async def gemini_batch_get(request: Request, batch_name: str): + """Get a specific Gemini batch with CCR post-processing.""" + return await proxy.handle_google_batch_results(request, batch_name) + + @app.post("/v1beta/batches/{batch_name}:cancel") + async def gemini_batch_cancel(request: Request, batch_name: str): + """Cancel a Gemini batch (passthrough).""" + return await proxy.handle_google_batch_passthrough(request, batch_name) + + @app.delete("/v1beta/batches/{batch_name}") + async def gemini_batch_delete(request: Request, batch_name: str): + """Delete a Gemini batch (passthrough).""" + return await proxy.handle_google_batch_passthrough(request, batch_name) + + @app.post("/v1beta/cachedContents") + async def gemini_create_cached_content(request: Request): + """Gemini create cached content API - passthrough.""" + return await proxy.handle_passthrough( + request, proxy.GEMINI_API_URL, "cachedContents", "gemini" + ) + + @app.get("/v1beta/cachedContents") + async def gemini_list_cached_contents(request: Request): + """Gemini list cached contents API - passthrough.""" + return await proxy.handle_passthrough( + request, proxy.GEMINI_API_URL, "cachedContents", "gemini" + ) + + @app.get("/v1beta/cachedContents/{cache_id}") + async def gemini_get_cached_content(request: Request, cache_id: str): + """Gemini get cached content API - passthrough.""" + return await proxy.handle_passthrough( + request, proxy.GEMINI_API_URL, "cachedContents", "gemini" + ) + + @app.delete("/v1beta/cachedContents/{cache_id}") + async def gemini_delete_cached_content(request: Request, cache_id: str): + """Gemini delete cached content API - passthrough.""" + return await proxy.handle_passthrough( + request, proxy.GEMINI_API_URL, "cachedContents", "gemini" + ) + + # ========================================================================= + # Catch-all Passthrough + # ========================================================================= + # Passthrough - route to correct backend based on headers @app.api_route("/{path:path}", methods=["GET", "POST", "PUT", "DELETE"]) async def passthrough(request: Request, path: str): diff --git a/tests/test_ccr_tool_injection.py b/tests/test_ccr_tool_injection.py index 40c2558db..6d3f6a023 100644 --- a/tests/test_ccr_tool_injection.py +++ b/tests/test_ccr_tool_injection.py @@ -344,3 +344,87 @@ class TestSystemInstructions: assert "hash0" in instructions assert "hash4" in instructions assert "..." in instructions + + +class TestAlternativeMarkerFormats: + """Test CCR marker detection for different compressor formats. + + Different compressors use slightly different marker formats: + - SmartCrusher: [N items compressed to M. Retrieve more: hash=xxx] + - TextCompressor: [N lines compressed to M. Retrieve more: hash=xxx] + - LogCompressor: [N lines compressed to M. Retrieve more: hash=xxx] + - SearchCompressor: [N matches compressed to M. Retrieve more: hash=xxx] + - LLMLingua: [N items compressed to M. Retrieve more: hash=xxx] + + The CCRToolInjector should detect all these formats. + """ + + def test_textcompressor_format(self): + """Detects TextCompressor marker format (lines).""" + messages = [ + { + "role": "assistant", + "content": "Build output:\n[500 lines compressed to 50. Retrieve more: hash=aabbccddeeff00112233]", + }, + ] + + injector = CCRToolInjector() + hashes = injector.scan_for_markers(messages) + + assert len(hashes) == 1 + assert "aabbccddeeff00112233" in hashes + + def test_searchcompressor_format(self): + """Detects SearchCompressor marker format (matches).""" + messages = [ + { + "role": "assistant", + "content": "Search results:\n[100 matches compressed to 10. Retrieve more: hash=1122334455667788]", + }, + ] + + injector = CCRToolInjector() + hashes = injector.scan_for_markers(messages) + + assert len(hashes) == 1 + assert "1122334455667788" in hashes + + def test_mixed_compressor_formats(self): + """Detects multiple marker formats in same conversation.""" + messages = [ + { + "role": "assistant", + "content": "Search results:\n[50 matches compressed to 5. Retrieve more: hash=aaaa11111111]", + }, + { + "role": "assistant", + "content": "Build logs:\n[200 lines compressed to 20. Retrieve more: hash=bbbb22222222]", + }, + { + "role": "assistant", + "content": "Database:\n[1000 items compressed to 100. Retrieve more: hash=cccc33333333]", + }, + ] + + injector = CCRToolInjector() + hashes = injector.scan_for_markers(messages) + + assert len(hashes) == 3 + assert "aaaa11111111" in hashes + assert "bbbb22222222" in hashes + assert "cccc33333333" in hashes + + def test_generic_compressed_marker(self): + """Detects generic compression markers via fallback pattern.""" + messages = [ + { + "role": "assistant", + "content": "Data:\n[Content compressed for efficiency. hash=fedcba9876543210fedcba98]", + }, + ] + + injector = CCRToolInjector() + hashes = injector.scan_for_markers(messages) + + assert len(hashes) == 1 + assert "fedcba9876543210fedcba98" in hashes diff --git a/tests/test_proxy_batch_integration.py b/tests/test_proxy_batch_integration.py new file mode 100644 index 000000000..fa4bc33f0 --- /dev/null +++ b/tests/test_proxy_batch_integration.py @@ -0,0 +1,522 @@ +"""Integration tests for proxy batch APIs with compression. + +These tests verify that batch endpoints work correctly with real API calls +and compression enabled, testing token savings tracking. + +Required environment variables: +- OPENAI_API_KEY: For OpenAI /v1/batches endpoint +- ANTHROPIC_API_KEY: For Anthropic /v1/messages/batches endpoint + +IMPORTANT: Batch API tests create real batch jobs which may incur costs. +Use sparingly and clean up resources after testing. + +Run with: + OPENAI_API_KEY=... ANTHROPIC_API_KEY=... pytest tests/test_proxy_batch_integration.py -v +""" + +import json +import os + +import pytest + +pytest.importorskip("fastapi") +pytest.importorskip("httpx") + +from fastapi.testclient import TestClient + +from headroom.proxy.server import ProxyConfig, create_app + +# ============================================================================= +# Fixtures +# ============================================================================= + + +@pytest.fixture +def openai_batch_client(): + """Create test client for OpenAI batch API with compression enabled.""" + config = ProxyConfig( + optimize=True, # Enable compression for batch + cache_enabled=False, + rate_limit_enabled=False, + cost_tracking_enabled=False, + ) + app = create_app(config) + with TestClient(app) as client: + yield client + + +@pytest.fixture +def anthropic_batch_client(): + """Create test client for Anthropic batch API with compression enabled.""" + config = ProxyConfig( + optimize=True, # Enable compression for batch + cache_enabled=False, + rate_limit_enabled=False, + cost_tracking_enabled=False, + ) + app = create_app(config) + with TestClient(app) as client: + yield client + + +@pytest.fixture +def openai_api_key(): + """Get OpenAI API key from environment.""" + return os.environ.get("OPENAI_API_KEY") + + +@pytest.fixture +def anthropic_api_key(): + """Get Anthropic API key from environment.""" + return os.environ.get("ANTHROPIC_API_KEY") + + +def create_large_messages(num_items: int = 50) -> list[dict]: + """Create messages with large JSON data for compression testing.""" + # Create a list of items that will be compressible + items = [ + { + "id": i, + "name": f"Item number {i}", + "description": f"This is a detailed description for item {i}. It contains additional information.", + "status": "active" if i % 2 == 0 else "inactive", + "metadata": { + "created_at": f"2024-01-{(i % 28) + 1:02d}", + "updated_at": f"2024-06-{(i % 28) + 1:02d}", + "tags": [f"tag{i % 5}", f"category{i % 3}"], + }, + } + for i in range(num_items) + ] + large_json = json.dumps(items, indent=2) + + return [ + {"role": "system", "content": "You are a helpful data analyst assistant."}, + {"role": "user", "content": "I have some data I need you to analyze."}, + {"role": "assistant", "content": f"I've received your data:\n\n{large_json}"}, + {"role": "user", "content": "How many items have status 'active'?"}, + ] + + +# ============================================================================= +# OpenAI Batch API Tests +# ============================================================================= + + +@pytest.mark.skipif(not os.environ.get("OPENAI_API_KEY"), reason="OPENAI_API_KEY not set") +class TestOpenAIBatchCreate: + """Test OpenAI /v1/batches create endpoint with compression.""" + + def test_batch_create_validation_missing_input_file(self, openai_batch_client, openai_api_key): + """POST /v1/batches without input_file_id returns validation error.""" + response = openai_batch_client.post( + "/v1/batches", + headers={"Authorization": f"Bearer {openai_api_key}"}, + json={ + "endpoint": "/v1/chat/completions", + "completion_window": "24h", + }, + ) + assert response.status_code == 400 + data = response.json() + assert "error" in data + assert "input_file_id" in data["error"]["message"].lower() + + def test_batch_create_validation_missing_endpoint(self, openai_batch_client, openai_api_key): + """POST /v1/batches without endpoint returns validation error.""" + response = openai_batch_client.post( + "/v1/batches", + headers={"Authorization": f"Bearer {openai_api_key}"}, + json={ + "input_file_id": "file-abc123", + "completion_window": "24h", + }, + ) + assert response.status_code == 400 + data = response.json() + assert "error" in data + assert "endpoint" in data["error"]["message"].lower() + + def test_batch_create_with_compression(self, openai_batch_client, openai_api_key): + """Full batch creation flow with compression. + + This test: + 1. Creates a JSONL file with compressible content + 2. Uploads it to OpenAI + 3. Creates a batch with compression enabled + 4. Verifies compression stats are tracked + 5. Cancels the batch to avoid costs + """ + # Step 1: Create JSONL content with compressible messages + messages = create_large_messages(num_items=30) + jsonl_lines = [ + json.dumps( + { + "custom_id": f"request-{i}", + "method": "POST", + "url": "/v1/chat/completions", + "body": { + "model": "gpt-4o-mini", + "messages": messages, + "max_tokens": 100, + }, + } + ) + for i in range(3) # 3 requests in batch + ] + jsonl_content = "\n".join(jsonl_lines) + + # Step 2: Upload the JSONL file directly to OpenAI + import httpx + + upload_response = httpx.post( + "https://api.openai.com/v1/files", + headers={"Authorization": f"Bearer {openai_api_key}"}, + files={"file": ("batch_input.jsonl", jsonl_content.encode(), "application/jsonl")}, + data={"purpose": "batch"}, + ) + assert upload_response.status_code == 200, f"File upload failed: {upload_response.text}" + file_data = upload_response.json() + input_file_id = file_data["id"] + + try: + # Step 3: Create batch through proxy with compression + response = openai_batch_client.post( + "/v1/batches", + headers={"Authorization": f"Bearer {openai_api_key}"}, + json={ + "input_file_id": input_file_id, + "endpoint": "/v1/chat/completions", + "completion_window": "24h", + "metadata": {"test": "compression_integration"}, + }, + ) + assert response.status_code == 200, f"Batch creation failed: {response.text}" + batch_data = response.json() + + # Verify batch was created + assert "id" in batch_data + assert batch_data["object"] == "batch" + batch_id = batch_data["id"] + + # Verify compression stats in response headers + if "x-headroom-tokens-saved" in response.headers: + tokens_saved = int(response.headers["x-headroom-tokens-saved"]) + assert tokens_saved >= 0 + + if "x-headroom-savings-percent" in response.headers: + savings_percent = float(response.headers["x-headroom-savings-percent"]) + assert 0 <= savings_percent <= 100 + + # Verify compression metadata was added + metadata = batch_data.get("metadata", {}) + if metadata.get("headroom_compressed") == "true": + # Compression was applied + assert "headroom_tokens_saved" in metadata + assert "headroom_original_tokens" in metadata + assert "headroom_compressed_tokens" in metadata + tokens_saved = int(metadata["headroom_tokens_saved"]) + assert tokens_saved >= 0 + + # Step 4: Cancel the batch to avoid costs + cancel_response = openai_batch_client.post( + f"/v1/batches/{batch_id}/cancel", + headers={"Authorization": f"Bearer {openai_api_key}"}, + ) + # Cancel may succeed or fail if batch already completed/cancelled + assert cancel_response.status_code in [200, 400] + + finally: + # Cleanup: Delete the uploaded file + httpx.delete( + f"https://api.openai.com/v1/files/{input_file_id}", + headers={"Authorization": f"Bearer {openai_api_key}"}, + ) + + +@pytest.mark.skipif(not os.environ.get("OPENAI_API_KEY"), reason="OPENAI_API_KEY not set") +class TestOpenAIBatchList: + """Test OpenAI /v1/batches list endpoint passthrough.""" + + def test_list_batches(self, openai_batch_client, openai_api_key): + """GET /v1/batches returns list of batches.""" + response = openai_batch_client.get( + "/v1/batches", + headers={"Authorization": f"Bearer {openai_api_key}"}, + ) + assert response.status_code == 200 + data = response.json() + + # Verify list response format + assert "data" in data + assert "object" in data + assert data["object"] == "list" + + def test_list_batches_with_limit(self, openai_batch_client, openai_api_key): + """GET /v1/batches with limit parameter.""" + response = openai_batch_client.get( + "/v1/batches?limit=5", + headers={"Authorization": f"Bearer {openai_api_key}"}, + ) + assert response.status_code == 200 + data = response.json() + + assert len(data["data"]) <= 5 + + +# ============================================================================= +# Anthropic Batch API Tests +# ============================================================================= + + +@pytest.mark.skipif(not os.environ.get("ANTHROPIC_API_KEY"), reason="ANTHROPIC_API_KEY not set") +class TestAnthropicBatchCreate: + """Test Anthropic /v1/messages/batches create endpoint with compression.""" + + def test_batch_create_validation_missing_requests( + self, anthropic_batch_client, anthropic_api_key + ): + """POST /v1/messages/batches without requests returns validation error.""" + response = anthropic_batch_client.post( + "/v1/messages/batches", + headers={ + "x-api-key": anthropic_api_key, + "anthropic-version": "2023-06-01", + "anthropic-beta": "message-batches-2024-09-24", + }, + json={}, + ) + assert response.status_code == 400 + data = response.json() + assert "error" in data + + def test_batch_create_validation_empty_requests( + self, anthropic_batch_client, anthropic_api_key + ): + """POST /v1/messages/batches with empty requests list returns error.""" + response = anthropic_batch_client.post( + "/v1/messages/batches", + headers={ + "x-api-key": anthropic_api_key, + "anthropic-version": "2023-06-01", + "anthropic-beta": "message-batches-2024-09-24", + }, + json={"requests": []}, + ) + assert response.status_code == 400 + data = response.json() + assert "error" in data + + def test_batch_create_with_compression(self, anthropic_batch_client, anthropic_api_key): + """Create Anthropic batch with compression. + + This test: + 1. Creates a batch request with compressible messages + 2. Verifies the batch is created successfully + 3. Checks that compression stats are tracked + 4. Cancels the batch to avoid costs + """ + # Create messages with compressible content + messages = create_large_messages(num_items=25) + + # Create batch request in Anthropic format + batch_requests = [ + { + "custom_id": f"req-{i}", + "params": { + "model": "claude-3-5-haiku-20241022", + "max_tokens": 100, + "messages": messages, + }, + } + for i in range(2) # 2 requests in batch + ] + + response = anthropic_batch_client.post( + "/v1/messages/batches", + headers={ + "x-api-key": anthropic_api_key, + "anthropic-version": "2023-06-01", + "anthropic-beta": "message-batches-2024-09-24", + "content-type": "application/json", + }, + json={"requests": batch_requests}, + ) + assert response.status_code == 200, f"Batch creation failed: {response.text}" + batch_data = response.json() + + # Verify batch was created + assert "id" in batch_data + assert batch_data["type"] == "message_batch" + batch_id = batch_data["id"] + + # Verify processing status + assert "processing_status" in batch_data + assert batch_data["processing_status"] in ["in_progress", "ended", "canceling"] + + # Check proxy stats for compression + stats_response = anthropic_batch_client.get("/stats") + stats = stats_response.json() + # Batch requests should be tracked + assert stats["requests"]["total"] >= 1 + + # Cancel the batch to avoid costs + cancel_response = anthropic_batch_client.post( + f"/v1/messages/batches/{batch_id}/cancel", + headers={ + "x-api-key": anthropic_api_key, + "anthropic-version": "2023-06-01", + "anthropic-beta": "message-batches-2024-09-24", + }, + ) + # Cancel may succeed or return error if already processed + assert cancel_response.status_code in [200, 400, 409] + + +@pytest.mark.skipif(not os.environ.get("ANTHROPIC_API_KEY"), reason="ANTHROPIC_API_KEY not set") +class TestAnthropicBatchList: + """Test Anthropic /v1/messages/batches list endpoint passthrough.""" + + def test_list_batches(self, anthropic_batch_client, anthropic_api_key): + """GET /v1/messages/batches returns list of batches.""" + response = anthropic_batch_client.get( + "/v1/messages/batches", + headers={ + "x-api-key": anthropic_api_key, + "anthropic-version": "2023-06-01", + "anthropic-beta": "message-batches-2024-09-24", + }, + ) + assert response.status_code == 200 + data = response.json() + + # Verify list response format + assert "data" in data + + def test_list_batches_with_limit(self, anthropic_batch_client, anthropic_api_key): + """GET /v1/messages/batches with limit parameter.""" + response = anthropic_batch_client.get( + "/v1/messages/batches?limit=5", + headers={ + "x-api-key": anthropic_api_key, + "anthropic-version": "2023-06-01", + "anthropic-beta": "message-batches-2024-09-24", + }, + ) + assert response.status_code == 200 + data = response.json() + + assert len(data.get("data", [])) <= 5 + + +# ============================================================================= +# Compression Verification Tests +# ============================================================================= + + +@pytest.mark.skipif(not os.environ.get("OPENAI_API_KEY"), reason="OPENAI_API_KEY not set") +class TestBatchCompressionStats: + """Test that batch compression stats are properly tracked.""" + + def test_stats_track_batch_requests(self, openai_batch_client, openai_api_key): + """Verify batch requests update proxy stats correctly.""" + # Get initial stats + initial_stats = openai_batch_client.get("/stats").json() + initial_requests = initial_stats["requests"]["total"] + + # Make a batch list request (passthrough) + openai_batch_client.get( + "/v1/batches", + headers={"Authorization": f"Bearer {openai_api_key}"}, + ) + + # Verify stats updated + updated_stats = openai_batch_client.get("/stats").json() + assert updated_stats["requests"]["total"] >= initial_requests + + +@pytest.mark.skipif(not os.environ.get("ANTHROPIC_API_KEY"), reason="ANTHROPIC_API_KEY not set") +class TestAnthropicBatchCompressionStats: + """Test Anthropic batch compression stats tracking.""" + + def test_stats_track_anthropic_batch_requests(self, anthropic_batch_client, anthropic_api_key): + """Verify Anthropic batch requests update proxy stats.""" + # Get initial stats + initial_stats = anthropic_batch_client.get("/stats").json() + initial_requests = initial_stats["requests"]["total"] + + # Make a batch list request + anthropic_batch_client.get( + "/v1/messages/batches", + headers={ + "x-api-key": anthropic_api_key, + "anthropic-version": "2023-06-01", + "anthropic-beta": "message-batches-2024-09-24", + }, + ) + + # Verify stats updated + updated_stats = anthropic_batch_client.get("/stats").json() + assert updated_stats["requests"]["total"] >= initial_requests + + +# ============================================================================= +# Error Handling Tests +# ============================================================================= + + +class TestBatchErrorHandling: + """Test error handling for batch endpoints.""" + + @pytest.mark.skipif(not os.environ.get("OPENAI_API_KEY"), reason="OPENAI_API_KEY not set") + def test_openai_batch_invalid_file_id(self, openai_batch_client, openai_api_key): + """Invalid file ID returns appropriate error.""" + response = openai_batch_client.post( + "/v1/batches", + headers={"Authorization": f"Bearer {openai_api_key}"}, + json={ + "input_file_id": "file-nonexistent12345", + "endpoint": "/v1/chat/completions", + "completion_window": "24h", + }, + ) + # Should return error for non-existent file + assert response.status_code in [400, 404] + + def test_openai_batch_missing_auth(self, openai_batch_client): + """Missing authentication returns error (401 or 404 depending on routing).""" + response = openai_batch_client.post( + "/v1/batches", + json={ + "input_file_id": "file-abc123", + "endpoint": "/v1/chat/completions", + }, + ) + # Proxy may return 404 (no route match) or 401 (auth error) + assert response.status_code in [401, 404] + + def test_anthropic_batch_missing_auth(self, anthropic_batch_client): + """Missing authentication returns error (401 or 400 depending on validation).""" + response = anthropic_batch_client.post( + "/v1/messages/batches", + headers={ + "anthropic-version": "2023-06-01", + "anthropic-beta": "message-batches-2024-09-24", + }, + json={"requests": []}, + ) + # Proxy may return 400 (validation) or 401 (auth error) + assert response.status_code in [400, 401] + + @pytest.mark.skipif(not os.environ.get("OPENAI_API_KEY"), reason="OPENAI_API_KEY not set") + def test_openai_batch_invalid_json(self, openai_batch_client, openai_api_key): + """Invalid JSON body returns 400.""" + response = openai_batch_client.post( + "/v1/batches", + headers={ + "Authorization": f"Bearer {openai_api_key}", + "Content-Type": "application/json", + }, + content=b"not valid json", + ) + assert response.status_code == 400 diff --git a/tests/test_proxy_count_tokens_integration.py b/tests/test_proxy_count_tokens_integration.py new file mode 100644 index 000000000..40f53399d --- /dev/null +++ b/tests/test_proxy_count_tokens_integration.py @@ -0,0 +1,500 @@ +"""Integration tests for Gemini countTokens endpoint with compression. + +These tests verify that the Gemini /v1beta/models/{model}:countTokens endpoint +works correctly with compression enabled, properly counting tokens after +compression is applied. + +Required environment variables: +- GEMINI_API_KEY: For Gemini countTokens endpoint + +Run with: + GEMINI_API_KEY=... pytest tests/test_proxy_count_tokens_integration.py -v +""" + +import json +import os + +import pytest + +# Skip entire module if no API key +pytestmark = pytest.mark.skipif( + not os.environ.get("GEMINI_API_KEY"), reason="GEMINI_API_KEY not set" +) + +pytest.importorskip("fastapi") +pytest.importorskip("httpx") + +from fastapi.testclient import TestClient # noqa: E402 + +from headroom.proxy.server import ProxyConfig, create_app # noqa: E402 + +# ============================================================================= +# Fixtures +# ============================================================================= + + +@pytest.fixture +def gemini_client_optimized(): + """Create test client with optimization enabled for Gemini.""" + config = ProxyConfig( + optimize=True, # Enable compression + cache_enabled=False, + rate_limit_enabled=False, + cost_tracking_enabled=False, + ) + app = create_app(config) + with TestClient(app) as client: + yield client + + +@pytest.fixture +def gemini_client_passthrough(): + """Create test client with optimization disabled (passthrough mode).""" + config = ProxyConfig( + optimize=False, # Disable compression + cache_enabled=False, + rate_limit_enabled=False, + cost_tracking_enabled=False, + ) + app = create_app(config) + with TestClient(app) as client: + yield client + + +@pytest.fixture +def api_key(): + """Get Gemini API key from environment.""" + return os.environ.get("GEMINI_API_KEY") + + +def create_large_content(num_items: int = 50) -> list[dict]: + """Create Gemini-format contents with large compressible data.""" + # Create JSON data that can be compressed + items = [ + { + "id": i, + "name": f"Product Item {i}", + "description": f"This is a detailed description for product item {i}. " + f"It includes various specifications and features.", + "price": 99.99 + i * 0.5, + "category": f"category_{i % 5}", + "in_stock": i % 2 == 0, + "metadata": { + "sku": f"SKU-{i:05d}", + "weight": f"{i * 0.1:.2f}kg", + "dimensions": f"{10 + i}x{15 + i}x{5 + i}cm", + }, + } + for i in range(num_items) + ] + large_json = json.dumps(items, indent=2) + + return [ + { + "role": "user", + "parts": [{"text": "I have product data to analyze."}], + }, + { + "role": "model", + "parts": [{"text": f"Here is the product data:\n\n{large_json}"}], + }, + { + "role": "user", + "parts": [{"text": "How many products are in stock?"}], + }, + ] + + +def create_simple_content() -> list[dict]: + """Create simple Gemini-format contents for basic testing.""" + return [ + { + "role": "user", + "parts": [{"text": "What is 2 + 2?"}], + } + ] + + +# ============================================================================= +# Basic countTokens Tests +# ============================================================================= + + +class TestGeminiCountTokensBasic: + """Test basic Gemini countTokens functionality.""" + + def test_count_tokens_simple_content(self, gemini_client_optimized, api_key): + """Basic token counting works correctly.""" + response = gemini_client_optimized.post( + f"/v1beta/models/gemini-2.0-flash:countTokens?key={api_key}", + json={"contents": create_simple_content()}, + ) + assert response.status_code == 200 + data = response.json() + + # Verify response format + assert "totalTokens" in data + assert isinstance(data["totalTokens"], int) + assert data["totalTokens"] > 0 + + def test_count_tokens_with_system_instruction(self, gemini_client_optimized, api_key): + """Token counting includes system instruction.""" + response = gemini_client_optimized.post( + f"/v1beta/models/gemini-2.0-flash:countTokens?key={api_key}", + json={ + "contents": create_simple_content(), + "systemInstruction": {"parts": [{"text": "You are a helpful math assistant."}]}, + }, + ) + # Note: systemInstruction may not be supported by all models/versions + # Accept both success and 400 (if not supported) + assert response.status_code in [200, 400] + if response.status_code == 200: + data = response.json() + assert "totalTokens" in data + assert data["totalTokens"] > 0 + + def test_count_tokens_multi_turn(self, gemini_client_optimized, api_key): + """Token counting for multi-turn conversation.""" + contents = [ + {"role": "user", "parts": [{"text": "Hello, my name is Alice."}]}, + {"role": "model", "parts": [{"text": "Nice to meet you, Alice!"}]}, + {"role": "user", "parts": [{"text": "What is my name?"}]}, + ] + + response = gemini_client_optimized.post( + f"/v1beta/models/gemini-2.0-flash:countTokens?key={api_key}", + json={"contents": contents}, + ) + assert response.status_code == 200 + data = response.json() + + assert data["totalTokens"] > 0 + + +# ============================================================================= +# Compression Tests +# ============================================================================= + + +class TestGeminiCountTokensCompression: + """Test that compression reduces token count.""" + + def test_compression_reduces_token_count( + self, gemini_client_optimized, gemini_client_passthrough, api_key + ): + """Verify compression reduces token count for large content. + + This test compares token counts between: + - Passthrough mode (no compression) + - Optimized mode (compression enabled) + """ + large_contents = create_large_content(num_items=40) + + # Get token count without compression + passthrough_response = gemini_client_passthrough.post( + f"/v1beta/models/gemini-2.0-flash:countTokens?key={api_key}", + json={"contents": large_contents}, + ) + assert passthrough_response.status_code == 200 + passthrough_tokens = passthrough_response.json()["totalTokens"] + + # Get token count with compression + optimized_response = gemini_client_optimized.post( + f"/v1beta/models/gemini-2.0-flash:countTokens?key={api_key}", + json={"contents": large_contents}, + ) + assert optimized_response.status_code == 200 + optimized_tokens = optimized_response.json()["totalTokens"] + + # Compression should reduce token count (or at least not increase it) + # Note: compression effect depends on content and may vary + assert optimized_tokens <= passthrough_tokens * 1.1 # Allow 10% margin + + # For large content, we expect some savings + if passthrough_tokens > 1000: + assert optimized_tokens < passthrough_tokens, ( + f"Expected compression to reduce tokens from {passthrough_tokens} " + f"but got {optimized_tokens}" + ) + + def test_compression_stats_tracked(self, gemini_client_optimized, api_key): + """Verify compression stats are tracked in proxy stats.""" + large_contents = create_large_content(num_items=30) + + # Make countTokens request with large content + response = gemini_client_optimized.post( + f"/v1beta/models/gemini-2.0-flash:countTokens?key={api_key}", + json={"contents": large_contents}, + ) + assert response.status_code == 200 + + # Check proxy stats + stats_response = gemini_client_optimized.get("/stats") + assert stats_response.status_code == 200 + stats = stats_response.json() + + # Verify Gemini requests are tracked + assert stats["requests"]["total"] >= 1 + assert "gemini" in stats["requests"]["by_provider"] + + +class TestGeminiCountTokensLargeContent: + """Test countTokens with large content that benefits from compression.""" + + def test_very_large_json_content(self, gemini_client_optimized, api_key): + """Token counting handles very large JSON content.""" + large_contents = create_large_content(num_items=100) + + response = gemini_client_optimized.post( + f"/v1beta/models/gemini-2.0-flash:countTokens?key={api_key}", + json={"contents": large_contents}, + ) + assert response.status_code == 200 + data = response.json() + + assert "totalTokens" in data + assert data["totalTokens"] > 0 + + def test_repeated_data_compression(self, gemini_client_optimized, api_key): + """Content with repeated patterns compresses well.""" + # Create content with highly repetitive data + repeated_items = [{"id": i, "status": "active", "type": "item"} for i in range(200)] + repeated_json = json.dumps(repeated_items) + + contents = [ + {"role": "user", "parts": [{"text": "Analyze this data."}]}, + {"role": "model", "parts": [{"text": f"Data:\n{repeated_json}"}]}, + {"role": "user", "parts": [{"text": "Count the items."}]}, + ] + + response = gemini_client_optimized.post( + f"/v1beta/models/gemini-2.0-flash:countTokens?key={api_key}", + json={"contents": contents}, + ) + assert response.status_code == 200 + data = response.json() + + assert data["totalTokens"] > 0 + + def test_code_content_compression(self, gemini_client_optimized, api_key): + """Token counting handles code content.""" + code_sample = ''' +def calculate_statistics(data): + """Calculate statistics for the given data.""" + if not data: + return {"count": 0, "sum": 0, "average": 0} + + count = len(data) + total = sum(data) + average = total / count + + return { + "count": count, + "sum": total, + "average": average, + "min": min(data), + "max": max(data), + } + +# Example usage +numbers = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10] +result = calculate_statistics(numbers) +print(result) +''' + + contents = [ + {"role": "user", "parts": [{"text": "Can you explain this code?"}]}, + { + "role": "model", + "parts": [{"text": f"Here's the code:\n\n```python\n{code_sample}\n```"}], + }, + {"role": "user", "parts": [{"text": "What does calculate_statistics return?"}]}, + ] + + response = gemini_client_optimized.post( + f"/v1beta/models/gemini-2.0-flash:countTokens?key={api_key}", + json={"contents": contents}, + ) + assert response.status_code == 200 + data = response.json() + + assert data["totalTokens"] > 0 + + +# ============================================================================= +# Model Variant Tests +# ============================================================================= + + +class TestGeminiCountTokensModels: + """Test countTokens with different Gemini models.""" + + def test_gemini_flash_model(self, gemini_client_optimized, api_key): + """countTokens works with gemini-2.0-flash model.""" + response = gemini_client_optimized.post( + f"/v1beta/models/gemini-2.0-flash:countTokens?key={api_key}", + json={"contents": create_simple_content()}, + ) + assert response.status_code == 200 + assert "totalTokens" in response.json() + + def test_gemini_flash_lite_model(self, gemini_client_optimized, api_key): + """countTokens works with gemini-2.0-flash-lite model.""" + response = gemini_client_optimized.post( + f"/v1beta/models/gemini-2.0-flash-lite:countTokens?key={api_key}", + json={"contents": create_simple_content()}, + ) + # Model may or may not be available + assert response.status_code in [200, 404] + if response.status_code == 200: + assert "totalTokens" in response.json() + + +# ============================================================================= +# Error Handling Tests +# ============================================================================= + + +class TestGeminiCountTokensErrors: + """Test error handling for countTokens endpoint.""" + + def test_invalid_api_key(self, gemini_client_optimized): + """Invalid API key returns authentication error.""" + response = gemini_client_optimized.post( + "/v1beta/models/gemini-2.0-flash:countTokens?key=invalid-key-12345", + json={"contents": create_simple_content()}, + ) + assert response.status_code in [400, 401, 403] + + def test_invalid_model(self, gemini_client_optimized, api_key): + """Invalid model name returns error.""" + response = gemini_client_optimized.post( + f"/v1beta/models/nonexistent-model-xyz:countTokens?key={api_key}", + json={"contents": create_simple_content()}, + ) + assert response.status_code >= 400 + + def test_empty_contents(self, gemini_client_optimized, api_key): + """Empty contents may return error or zero tokens.""" + response = gemini_client_optimized.post( + f"/v1beta/models/gemini-2.0-flash:countTokens?key={api_key}", + json={"contents": []}, + ) + # May return error or success with 0 tokens + if response.status_code == 200: + data = response.json() + assert "totalTokens" in data + + def test_invalid_json_body(self, gemini_client_optimized, api_key): + """Invalid JSON body returns 400 error.""" + response = gemini_client_optimized.post( + f"/v1beta/models/gemini-2.0-flash:countTokens?key={api_key}", + headers={"Content-Type": "application/json"}, + content=b"not valid json", + ) + assert response.status_code == 400 + + def test_missing_contents_field(self, gemini_client_optimized, api_key): + """Missing contents field handled gracefully.""" + response = gemini_client_optimized.post( + f"/v1beta/models/gemini-2.0-flash:countTokens?key={api_key}", + json={}, + ) + # May return error or handle empty contents + assert response.status_code in [200, 400] + + +# ============================================================================= +# Stats Tracking Tests +# ============================================================================= + + +class TestGeminiCountTokensStats: + """Test proxy stats tracking for countTokens requests.""" + + def test_stats_track_gemini_provider(self, gemini_client_optimized, api_key): + """Stats correctly track Gemini provider.""" + # Clear stats by getting a fresh client + gemini_client_optimized.post( + f"/v1beta/models/gemini-2.0-flash:countTokens?key={api_key}", + json={"contents": create_simple_content()}, + ) + + stats = gemini_client_optimized.get("/stats").json() + assert "gemini" in stats["requests"]["by_provider"] + assert stats["requests"]["by_provider"]["gemini"] >= 1 + + def test_stats_track_model(self, gemini_client_optimized, api_key): + """Stats correctly track model used.""" + gemini_client_optimized.post( + f"/v1beta/models/gemini-2.0-flash:countTokens?key={api_key}", + json={"contents": create_simple_content()}, + ) + + stats = gemini_client_optimized.get("/stats").json() + # Model should be tracked in by_model + assert len(stats["requests"]["by_model"]) >= 1 + + def test_stats_track_tokens_saved(self, gemini_client_optimized, api_key): + """Stats track tokens saved from compression.""" + # Make request with large compressible content + large_contents = create_large_content(num_items=30) + gemini_client_optimized.post( + f"/v1beta/models/gemini-2.0-flash:countTokens?key={api_key}", + json={"contents": large_contents}, + ) + + stats = gemini_client_optimized.get("/stats").json() + # tokens.saved should be tracked (may be 0 if content wasn't compressed) + assert "tokens" in stats + assert "saved" in stats["tokens"] + + +# ============================================================================= +# Integration Tests +# ============================================================================= + + +class TestGeminiCountTokensIntegration: + """Integration tests combining multiple features.""" + + def test_full_workflow(self, gemini_client_optimized, api_key): + """Test complete workflow: count tokens, verify compression, check stats.""" + # Step 1: Count tokens with large content + large_contents = create_large_content(num_items=35) + + initial_stats = gemini_client_optimized.get("/stats").json() + initial_tokens_saved = initial_stats["tokens"]["saved"] + + # Step 2: Make countTokens request + response = gemini_client_optimized.post( + f"/v1beta/models/gemini-2.0-flash:countTokens?key={api_key}", + json={"contents": large_contents}, + ) + assert response.status_code == 200 + token_count = response.json()["totalTokens"] + assert token_count > 0 + + # Step 3: Verify stats updated + updated_stats = gemini_client_optimized.get("/stats").json() + assert updated_stats["requests"]["total"] > initial_stats["requests"]["total"] + + # Step 4: Verify tokens saved is tracked (may be negative for small overhead) + # Allow for some compression overhead + assert updated_stats["tokens"]["saved"] >= initial_tokens_saved - 100 + + def test_multiple_requests_accumulate_stats(self, gemini_client_optimized, api_key): + """Multiple requests correctly accumulate stats.""" + initial_stats = gemini_client_optimized.get("/stats").json() + initial_total = initial_stats["requests"]["total"] + + # Make several requests + for _ in range(3): + gemini_client_optimized.post( + f"/v1beta/models/gemini-2.0-flash:countTokens?key={api_key}", + json={"contents": create_simple_content()}, + ) + + updated_stats = gemini_client_optimized.get("/stats").json() + assert updated_stats["requests"]["total"] >= initial_total + 3 diff --git a/tests/test_proxy_gemini_integration.py b/tests/test_proxy_gemini_integration.py new file mode 100644 index 000000000..0aba8d09d --- /dev/null +++ b/tests/test_proxy_gemini_integration.py @@ -0,0 +1,254 @@ +"""Integration tests for the proxy with real Gemini API calls. + +These tests require a valid GEMINI_API_KEY environment variable. +They test the actual /v1/chat/completions endpoint with real API calls. + +Run with: + GEMINI_API_KEY=your-key pytest tests/test_proxy_gemini_integration.py -v +""" + +import json +import os + +import pytest + +# Skip entire module if no API key +pytestmark = pytest.mark.skipif( + not os.environ.get("GEMINI_API_KEY"), reason="GEMINI_API_KEY not set" +) + +pytest.importorskip("fastapi") +pytest.importorskip("httpx") + +from fastapi.testclient import TestClient # noqa: E402 + +from headroom.proxy.server import ProxyConfig, create_app # noqa: E402 + +GEMINI_BASE_URL = "https://generativelanguage.googleapis.com/v1beta/openai" + + +@pytest.fixture +def gemini_client(): + """Create test client configured to forward to Gemini.""" + config = ProxyConfig( + optimize=False, + cache_enabled=False, + rate_limit_enabled=False, + cost_tracking_enabled=False, + openai_api_url=GEMINI_BASE_URL, + ) + app = create_app(config) + with TestClient(app) as client: + yield client + + +@pytest.fixture +def api_key(): + """Get Gemini API key from environment.""" + return os.environ.get("GEMINI_API_KEY") + + +class TestGeminiChatCompletions: + """Test /v1/chat/completions with real Gemini API.""" + + def test_basic_completion(self, gemini_client, api_key): + """Basic chat completion works.""" + response = gemini_client.post( + "/v1/chat/completions", + headers={"Authorization": f"Bearer {api_key}"}, + json={ + "model": "gemini-2.0-flash", + "messages": [ + {"role": "user", "content": "What is 2+2? Reply with just the number."} + ], + }, + ) + assert response.status_code == 200 + data = response.json() + + # Verify OpenAI-compatible response format + assert "choices" in data + assert len(data["choices"]) > 0 + assert "message" in data["choices"][0] + assert "content" in data["choices"][0]["message"] + assert "4" in data["choices"][0]["message"]["content"] + + # Verify usage stats + assert "usage" in data + assert "prompt_tokens" in data["usage"] + assert "completion_tokens" in data["usage"] + + def test_multi_turn_conversation(self, gemini_client, api_key): + """Multi-turn conversations maintain context.""" + response = gemini_client.post( + "/v1/chat/completions", + headers={"Authorization": f"Bearer {api_key}"}, + json={ + "model": "gemini-2.0-flash", + "messages": [ + {"role": "system", "content": "You are a helpful assistant. Be concise."}, + {"role": "user", "content": "My name is TestUser123."}, + {"role": "assistant", "content": "Nice to meet you, TestUser123!"}, + {"role": "user", "content": "What is my name?"}, + ], + }, + ) + assert response.status_code == 200 + data = response.json() + + content = data["choices"][0]["message"]["content"].lower() + assert "testuser123" in content + + def test_streaming(self, gemini_client, api_key): + """Streaming responses work correctly.""" + response = gemini_client.post( + "/v1/chat/completions", + headers={"Authorization": f"Bearer {api_key}"}, + json={ + "model": "gemini-2.0-flash", + "stream": True, + "messages": [{"role": "user", "content": "Count from 1 to 3."}], + }, + ) + assert response.status_code == 200 + + # Parse SSE stream + chunks = [] + for line in response.text.strip().split("\n"): + if line.startswith("data: ") and line != "data: [DONE]": + chunk = json.loads(line[6:]) + chunks.append(chunk) + + assert len(chunks) > 0 + + # Verify chunk format + for chunk in chunks: + assert "choices" in chunk + assert "delta" in chunk["choices"][0] + assert chunk["object"] == "chat.completion.chunk" + + def test_function_calling(self, gemini_client, api_key): + """Function calling / tools work correctly.""" + response = gemini_client.post( + "/v1/chat/completions", + headers={"Authorization": f"Bearer {api_key}"}, + json={ + "model": "gemini-2.0-flash", + "messages": [{"role": "user", "content": "What is the weather in Paris?"}], + "tools": [ + { + "type": "function", + "function": { + "name": "get_weather", + "description": "Get the weather for a location", + "parameters": { + "type": "object", + "properties": { + "location": {"type": "string", "description": "City name"} + }, + "required": ["location"], + }, + }, + } + ], + "tool_choice": "auto", + }, + ) + assert response.status_code == 200 + data = response.json() + + # Verify tool call response + message = data["choices"][0]["message"] + assert "tool_calls" in message + assert len(message["tool_calls"]) > 0 + + tool_call = message["tool_calls"][0] + assert tool_call["function"]["name"] == "get_weather" + assert "paris" in tool_call["function"]["arguments"].lower() + + def test_json_mode(self, gemini_client, api_key): + """JSON response format works.""" + response = gemini_client.post( + "/v1/chat/completions", + headers={"Authorization": f"Bearer {api_key}"}, + json={ + "model": "gemini-2.0-flash", + "messages": [ + { + "role": "user", + "content": "Return a JSON object with keys 'name' and 'age' for a person named Alice who is 30 years old.", + } + ], + "response_format": {"type": "json_object"}, + }, + ) + assert response.status_code == 200 + data = response.json() + + content = data["choices"][0]["message"]["content"] + # Parse the response as JSON + parsed = json.loads(content) + assert "name" in parsed or "Name" in parsed + assert "age" in parsed or "Age" in parsed + + +class TestGeminiModels: + """Test /v1/models endpoint with Gemini.""" + + def test_list_models(self, gemini_client, api_key): + """Can list available models.""" + response = gemini_client.get("/v1/models", headers={"Authorization": f"Bearer {api_key}"}) + # This goes through passthrough handler + assert response.status_code == 200 + data = response.json() + + assert "data" in data or "object" in data + + +class TestProxyStats: + """Test that proxy stats track Gemini requests correctly.""" + + def test_stats_track_requests(self, gemini_client, api_key): + """Proxy stats track Gemini requests.""" + # Make a request + gemini_client.post( + "/v1/chat/completions", + headers={"Authorization": f"Bearer {api_key}"}, + json={"model": "gemini-2.0-flash", "messages": [{"role": "user", "content": "Hi"}]}, + ) + + # Check stats + stats_response = gemini_client.get("/stats") + assert stats_response.status_code == 200 + stats = stats_response.json() + + assert stats["requests"]["total"] >= 1 + assert stats["requests"]["by_provider"]["openai"] >= 1 + assert "gemini" in str(stats["requests"]["by_model"]).lower() + + +class TestErrorHandling: + """Test error handling with Gemini.""" + + def test_invalid_api_key(self, gemini_client): + """Invalid API key returns appropriate error.""" + response = gemini_client.post( + "/v1/chat/completions", + headers={"Authorization": "Bearer invalid-key-123"}, + json={"model": "gemini-2.0-flash", "messages": [{"role": "user", "content": "Hi"}]}, + ) + # Should return 4xx error + assert response.status_code >= 400 + + def test_invalid_model(self, gemini_client, api_key): + """Invalid model returns appropriate error.""" + response = gemini_client.post( + "/v1/chat/completions", + headers={"Authorization": f"Bearer {api_key}"}, + json={ + "model": "nonexistent-model-xyz", + "messages": [{"role": "user", "content": "Hi"}], + }, + ) + # Should return 4xx error + assert response.status_code >= 400 diff --git a/tests/test_proxy_gemini_native_integration.py b/tests/test_proxy_gemini_native_integration.py new file mode 100644 index 000000000..ca05a56af --- /dev/null +++ b/tests/test_proxy_gemini_native_integration.py @@ -0,0 +1,374 @@ +"""Integration tests for Gemini native API endpoint with real API calls. + +These tests require a valid GEMINI_API_KEY environment variable. +They test the /v1beta/models/{model}:generateContent endpoint with compression. + +Run with: + GEMINI_API_KEY=your-key pytest tests/test_proxy_gemini_native_integration.py -v +""" + +import json +import os + +import pytest + +# Skip entire module if no API key +pytestmark = pytest.mark.skipif( + not os.environ.get("GEMINI_API_KEY"), reason="GEMINI_API_KEY not set" +) + +pytest.importorskip("fastapi") +pytest.importorskip("httpx") + +from fastapi.testclient import TestClient # noqa: E402 + +from headroom.proxy.server import ProxyConfig, create_app # noqa: E402 + + +@pytest.fixture +def gemini_native_client(): + """Create test client for Gemini native API with optimization enabled.""" + config = ProxyConfig( + optimize=True, # Enable compression + cache_enabled=False, + rate_limit_enabled=False, + cost_tracking_enabled=False, + ) + app = create_app(config) + with TestClient(app) as client: + yield client + + +@pytest.fixture +def api_key(): + """Get Gemini API key from environment.""" + return os.environ.get("GEMINI_API_KEY") + + +class TestGeminiNativeGenerateContent: + """Test /v1beta/models/{model}:generateContent endpoint.""" + + def test_basic_generation(self, gemini_native_client, api_key): + """Basic text generation works.""" + response = gemini_native_client.post( + f"/v1beta/models/gemini-2.0-flash:generateContent?key={api_key}", + json={"contents": [{"parts": [{"text": "What is 2+2? Reply with just the number."}]}]}, + ) + assert response.status_code == 200 + data = response.json() + + # Verify Gemini native response format + assert "candidates" in data + assert len(data["candidates"]) > 0 + assert "content" in data["candidates"][0] + assert "parts" in data["candidates"][0]["content"] + text = data["candidates"][0]["content"]["parts"][0]["text"] + assert "4" in text + + # Verify usage metadata + assert "usageMetadata" in data + assert "promptTokenCount" in data["usageMetadata"] + + def test_with_system_instruction(self, gemini_native_client, api_key): + """System instruction works correctly.""" + response = gemini_native_client.post( + f"/v1beta/models/gemini-2.0-flash:generateContent?key={api_key}", + json={ + "contents": [{"parts": [{"text": "Hello"}]}], + "systemInstruction": {"parts": [{"text": "Always respond with exactly one word."}]}, + }, + ) + assert response.status_code == 200 + data = response.json() + text = data["candidates"][0]["content"]["parts"][0]["text"] + # Should be a short response due to system instruction + assert len(text.split()) <= 3 + + def test_multi_turn_conversation(self, gemini_native_client, api_key): + """Multi-turn conversations maintain context.""" + response = gemini_native_client.post( + f"/v1beta/models/gemini-2.0-flash:generateContent?key={api_key}", + json={ + "contents": [ + {"role": "user", "parts": [{"text": "My name is TestUser456."}]}, + {"role": "model", "parts": [{"text": "Nice to meet you, TestUser456!"}]}, + {"role": "user", "parts": [{"text": "What is my name?"}]}, + ] + }, + ) + assert response.status_code == 200 + data = response.json() + text = data["candidates"][0]["content"]["parts"][0]["text"].lower() + assert "testuser456" in text + + def test_function_calling(self, gemini_native_client, api_key): + """Function calling / tools work correctly.""" + response = gemini_native_client.post( + f"/v1beta/models/gemini-2.0-flash:generateContent?key={api_key}", + json={ + "contents": [{"parts": [{"text": "What is the weather in Tokyo?"}]}], + "tools": [ + { + "functionDeclarations": [ + { + "name": "get_weather", + "description": "Get current weather for a location", + "parameters": { + "type": "object", + "properties": { + "location": {"type": "string", "description": "City name"} + }, + "required": ["location"], + }, + } + ] + } + ], + }, + ) + assert response.status_code == 200 + data = response.json() + + # Verify function call response + parts = data["candidates"][0]["content"]["parts"] + function_call = None + for part in parts: + if "functionCall" in part: + function_call = part["functionCall"] + break + + assert function_call is not None + assert function_call["name"] == "get_weather" + assert "tokyo" in function_call["args"]["location"].lower() + + def test_generation_config(self, gemini_native_client, api_key): + """Generation config parameters are respected.""" + response = gemini_native_client.post( + f"/v1beta/models/gemini-2.0-flash:generateContent?key={api_key}", + json={ + "contents": [{"parts": [{"text": "Write a very short poem about AI."}]}], + "generationConfig": {"maxOutputTokens": 50, "temperature": 0.1}, + }, + ) + assert response.status_code == 200 + data = response.json() + # Response should be limited by maxOutputTokens + assert data["usageMetadata"]["candidatesTokenCount"] <= 60 # Some buffer + + +class TestGeminiNativeCompression: + """Test that compression works with Gemini native API.""" + + def test_compression_on_model_message(self, gemini_native_client, api_key): + """Large data in model message gets compressed.""" + # Create large JSON data (simulating tool output) + items = [ + {"id": i, "name": f"Item {i}", "desc": f"Description for item {i}"} for i in range(100) + ] + tool_output = json.dumps(items) + + # Send as model message (like tool returning data) + response = gemini_native_client.post( + f"/v1beta/models/gemini-2.0-flash:generateContent?key={api_key}", + json={ + "contents": [ + {"role": "user", "parts": [{"text": "Get items from database"}]}, + {"role": "model", "parts": [{"text": f"Here are the results:\n{tool_output}"}]}, + {"role": "user", "parts": [{"text": "How many items are there?"}]}, + ] + }, + ) + assert response.status_code == 200 + data = response.json() + text = data["candidates"][0]["content"]["parts"][0]["text"] + # Model should correctly count the items + assert "100" in text + + # Check that compression happened via stats + stats = gemini_native_client.get("/stats").json() + # At least some tokens should have been saved + assert stats["tokens"]["saved"] >= 0 # May or may not compress depending on size + + def test_user_messages_protected(self, gemini_native_client, api_key): + """User messages are not compressed (by design).""" + # Large data in user message + items = [{"id": i} for i in range(50)] + user_data = json.dumps(items) + + # First request with data in user message + response = gemini_native_client.post( + f"/v1beta/models/gemini-2.0-flash:generateContent?key={api_key}", + json={ + "contents": [ + {"role": "user", "parts": [{"text": f"Analyze this data: {user_data}"}]} + ] + }, + ) + assert response.status_code == 200 + # The request should succeed - user messages are protected from compression + + +class TestGeminiNativeStats: + """Test that proxy stats track Gemini native requests correctly.""" + + def test_stats_track_gemini_provider(self, gemini_native_client, api_key): + """Stats show requests under 'gemini' provider.""" + # Make a request + gemini_native_client.post( + f"/v1beta/models/gemini-2.0-flash:generateContent?key={api_key}", + json={"contents": [{"parts": [{"text": "Hi"}]}]}, + ) + + stats = gemini_native_client.get("/stats").json() + assert "gemini" in stats["requests"]["by_provider"] + assert stats["requests"]["by_provider"]["gemini"] >= 1 + + def test_stats_track_model(self, gemini_native_client, api_key): + """Stats track the specific model used.""" + gemini_native_client.post( + f"/v1beta/models/gemini-2.0-flash:generateContent?key={api_key}", + json={"contents": [{"parts": [{"text": "Hi"}]}]}, + ) + + stats = gemini_native_client.get("/stats").json() + assert "gemini-2.0-flash" in stats["requests"]["by_model"] + + +class TestGeminiNativeErrorHandling: + """Test error handling for Gemini native API.""" + + def test_invalid_api_key(self, gemini_native_client): + """Invalid API key returns appropriate error.""" + response = gemini_native_client.post( + "/v1beta/models/gemini-2.0-flash:generateContent?key=invalid-key-123", + json={"contents": [{"parts": [{"text": "Hi"}]}]}, + ) + assert response.status_code >= 400 + + def test_invalid_model(self, gemini_native_client, api_key): + """Invalid model returns appropriate error.""" + response = gemini_native_client.post( + f"/v1beta/models/nonexistent-model-xyz:generateContent?key={api_key}", + json={"contents": [{"parts": [{"text": "Hi"}]}]}, + ) + assert response.status_code >= 400 + + def test_empty_contents(self, gemini_native_client, api_key): + """Empty contents handled gracefully.""" + response = gemini_native_client.post( + f"/v1beta/models/gemini-2.0-flash:generateContent?key={api_key}", json={"contents": []} + ) + # Should either return error or handle gracefully + assert response.status_code in [200, 400] + + +class TestGeminiNativeHeaderAuth: + """Test authentication via x-goog-api-key header.""" + + def test_header_auth(self, gemini_native_client, api_key): + """API key in header works.""" + response = gemini_native_client.post( + "/v1beta/models/gemini-2.0-flash:generateContent", + headers={"x-goog-api-key": api_key}, + json={"contents": [{"parts": [{"text": "Hi"}]}]}, + ) + assert response.status_code == 200 + + +class TestGeminiNativeCountTokens: + """Test /v1beta/models/{model}:countTokens endpoint with compression.""" + + def test_count_tokens_basic(self, gemini_native_client, api_key): + """Basic token counting works.""" + response = gemini_native_client.post( + f"/v1beta/models/gemini-2.0-flash:countTokens?key={api_key}", + json={"contents": [{"parts": [{"text": "Hello, world!"}]}]}, + ) + assert response.status_code == 200 + data = response.json() + + # Verify response format + assert "totalTokens" in data + assert isinstance(data["totalTokens"], int) + assert data["totalTokens"] > 0 + + def test_count_tokens_with_system_instruction(self, gemini_native_client, api_key): + """Token counting includes system instruction.""" + response = gemini_native_client.post( + f"/v1beta/models/gemini-2.0-flash:countTokens?key={api_key}", + json={ + "contents": [{"parts": [{"text": "Hello"}]}], + "systemInstruction": {"parts": [{"text": "You are a helpful assistant."}]}, + }, + ) + # Note: systemInstruction may not be supported by countTokens in all versions + assert response.status_code in [200, 400] + if response.status_code == 200: + data = response.json() + assert "totalTokens" in data + assert data["totalTokens"] > 0 + + def test_count_tokens_reflects_compression(self, gemini_native_client, api_key): + """Token count reflects compressed content size.""" + # Create large repetitive JSON data that should compress + items = [ + { + "id": i, + "name": f"Item {i}", + "description": f"This is the description for item number {i}", + } + for i in range(100) + ] + tool_output = json.dumps(items) + + # Count tokens with large data in model message (which gets compressed) + response = gemini_native_client.post( + f"/v1beta/models/gemini-2.0-flash:countTokens?key={api_key}", + json={ + "contents": [ + {"role": "user", "parts": [{"text": "Get items from database"}]}, + {"role": "model", "parts": [{"text": f"Here are the results:\n{tool_output}"}]}, + {"role": "user", "parts": [{"text": "Summarize these items"}]}, + ] + }, + ) + assert response.status_code == 200 + data = response.json() + + # Verify we got a token count + assert "totalTokens" in data + compressed_tokens = data["totalTokens"] + assert compressed_tokens > 0 + + # Check stats to verify compression was applied + stats = gemini_native_client.get("/stats").json() + # The request should have been tracked + assert stats["requests"]["by_provider"].get("gemini", 0) >= 1 + + def test_count_tokens_multi_turn(self, gemini_native_client, api_key): + """Token counting works for multi-turn conversations.""" + response = gemini_native_client.post( + f"/v1beta/models/gemini-2.0-flash:countTokens?key={api_key}", + json={ + "contents": [ + {"role": "user", "parts": [{"text": "My name is Alice."}]}, + {"role": "model", "parts": [{"text": "Nice to meet you, Alice!"}]}, + {"role": "user", "parts": [{"text": "What is my name?"}]}, + ] + }, + ) + assert response.status_code == 200 + data = response.json() + assert "totalTokens" in data + assert data["totalTokens"] > 0 + + def test_count_tokens_header_auth(self, gemini_native_client, api_key): + """API key in header works for countTokens.""" + response = gemini_native_client.post( + "/v1beta/models/gemini-2.0-flash:countTokens", + headers={"x-goog-api-key": api_key}, + json={"contents": [{"parts": [{"text": "Hello"}]}]}, + ) + assert response.status_code == 200 + data = response.json() + assert "totalTokens" in data diff --git a/tests/test_proxy_openai_responses_integration.py b/tests/test_proxy_openai_responses_integration.py new file mode 100644 index 000000000..74f8a8504 --- /dev/null +++ b/tests/test_proxy_openai_responses_integration.py @@ -0,0 +1,280 @@ +"""Integration tests for OpenAI /v1/responses endpoint with real API calls. + +These tests require a valid OPENAI_API_KEY environment variable. +They test the /v1/responses endpoint (introduced March 2025) with compression. + +Run with: + OPENAI_API_KEY=your-key pytest tests/test_proxy_openai_responses_integration.py -v +""" + +import json +import os + +import pytest + +# Skip entire module if no API key +pytestmark = pytest.mark.skipif( + not os.environ.get("OPENAI_API_KEY"), reason="OPENAI_API_KEY not set" +) + +pytest.importorskip("fastapi") +pytest.importorskip("httpx") + +from fastapi.testclient import TestClient # noqa: E402 + +from headroom.proxy.server import ProxyConfig, create_app # noqa: E402 + + +@pytest.fixture +def openai_responses_client(): + """Create test client for OpenAI responses API with optimization enabled.""" + config = ProxyConfig( + optimize=True, # Enable compression + cache_enabled=False, + rate_limit_enabled=False, + cost_tracking_enabled=False, + ) + app = create_app(config) + with TestClient(app) as client: + yield client + + +@pytest.fixture +def api_key(): + """Get OpenAI API key from environment.""" + return os.environ.get("OPENAI_API_KEY") + + +class TestOpenAIResponsesBasic: + """Test /v1/responses endpoint basic functionality.""" + + def test_basic_generation(self, openai_responses_client, api_key): + """Basic text generation works.""" + response = openai_responses_client.post( + "/v1/responses", + headers={"Authorization": f"Bearer {api_key}"}, + json={"model": "gpt-4o-mini", "input": "What is 2+2? Reply with just the number."}, + ) + assert response.status_code == 200 + data = response.json() + + # Verify responses API format + assert "id" in data + assert "output" in data + assert len(data["output"]) > 0 + assert data["output"][0]["type"] == "message" + assert data["output"][0]["role"] == "assistant" + + # Get the text content + content = data["output"][0]["content"] + assert len(content) > 0 + text = content[0].get("text", "") + assert "4" in text + + # Verify usage metadata + assert "usage" in data + + def test_with_instructions(self, openai_responses_client, api_key): + """System instructions work correctly.""" + response = openai_responses_client.post( + "/v1/responses", + headers={"Authorization": f"Bearer {api_key}"}, + json={ + "model": "gpt-4o-mini", + "input": "Hello", + "instructions": "Always respond with exactly one word.", + }, + ) + assert response.status_code == 200 + data = response.json() + + content = data["output"][0]["content"] + text = content[0].get("text", "") + # Should be a short response due to instructions + assert len(text.split()) <= 3 + + def test_input_as_array(self, openai_responses_client, api_key): + """Input can be an array of messages.""" + response = openai_responses_client.post( + "/v1/responses", + headers={"Authorization": f"Bearer {api_key}"}, + json={ + "model": "gpt-4o-mini", + "input": [ + {"role": "user", "content": "My name is TestUser789."}, + {"role": "assistant", "content": "Nice to meet you, TestUser789!"}, + {"role": "user", "content": "What is my name?"}, + ], + }, + ) + assert response.status_code == 200 + data = response.json() + + content = data["output"][0]["content"] + text = content[0].get("text", "").lower() + assert "testuser789" in text + + def test_generation_parameters(self, openai_responses_client, api_key): + """Generation parameters are respected.""" + response = openai_responses_client.post( + "/v1/responses", + headers={"Authorization": f"Bearer {api_key}"}, + json={ + "model": "gpt-4o-mini", + "input": "Write a very short poem about AI.", + "max_output_tokens": 50, + "temperature": 0.1, + }, + ) + assert response.status_code == 200 + data = response.json() + # Response should be limited by max_output_tokens + assert data["usage"]["output_tokens"] <= 60 # Some buffer + + +class TestOpenAIResponsesTools: + """Test function calling / tools with /v1/responses endpoint.""" + + def test_function_calling(self, openai_responses_client, api_key): + """Function calling works correctly.""" + # Note: /v1/responses uses a different tools format than /v1/chat/completions + # - name, description, parameters are at top level, not nested under "function" + response = openai_responses_client.post( + "/v1/responses", + headers={"Authorization": f"Bearer {api_key}"}, + json={ + "model": "gpt-4o-mini", + "input": "What is the weather in Tokyo?", + "tools": [ + { + "type": "function", + "name": "get_weather", + "description": "Get current weather for a location", + "parameters": { + "type": "object", + "properties": { + "location": {"type": "string", "description": "City name"} + }, + "required": ["location"], + }, + } + ], + }, + ) + assert response.status_code == 200 + data = response.json() + + # Find tool call in output + output = data["output"] + tool_call_found = False + for item in output: + if item.get("type") == "function_call": + tool_call_found = True + assert item["name"] == "get_weather" + args = ( + json.loads(item["arguments"]) + if isinstance(item["arguments"], str) + else item["arguments"] + ) + assert "tokyo" in args.get("location", "").lower() + break + + assert tool_call_found, "Expected function_call in output" + + +class TestOpenAIResponsesCompression: + """Test that compression works with /v1/responses endpoint.""" + + def test_compression_on_assistant_message(self, openai_responses_client, api_key): + """Large data in assistant message gets compressed.""" + # Create large JSON data (simulating tool output) + items = [ + {"id": i, "name": f"Item {i}", "desc": f"Description for item {i}"} for i in range(100) + ] + tool_output = json.dumps(items) + + # Send as multi-turn with assistant message containing data + response = openai_responses_client.post( + "/v1/responses", + headers={"Authorization": f"Bearer {api_key}"}, + json={ + "model": "gpt-4o-mini", + "input": [ + {"role": "user", "content": "Get items from database"}, + {"role": "assistant", "content": f"Here are the results:\n{tool_output}"}, + {"role": "user", "content": "How many items are there?"}, + ], + }, + ) + assert response.status_code == 200 + data = response.json() + + content = data["output"][0]["content"] + text = content[0].get("text", "") + # Model should correctly count the items + assert "100" in text + + # Check that compression happened via stats + stats = openai_responses_client.get("/stats").json() + # At least some tokens should have been saved + assert stats["tokens"]["saved"] >= 0 # May or may not compress depending on size + + +class TestOpenAIResponsesStats: + """Test that proxy stats track /v1/responses requests correctly.""" + + def test_stats_track_openai_provider(self, openai_responses_client, api_key): + """Stats show requests under 'openai' provider.""" + # Make a request + openai_responses_client.post( + "/v1/responses", + headers={"Authorization": f"Bearer {api_key}"}, + json={"model": "gpt-4o-mini", "input": "Hi"}, + ) + + stats = openai_responses_client.get("/stats").json() + assert "openai" in stats["requests"]["by_provider"] + assert stats["requests"]["by_provider"]["openai"] >= 1 + + def test_stats_track_model(self, openai_responses_client, api_key): + """Stats track the specific model used.""" + openai_responses_client.post( + "/v1/responses", + headers={"Authorization": f"Bearer {api_key}"}, + json={"model": "gpt-4o-mini", "input": "Hi"}, + ) + + stats = openai_responses_client.get("/stats").json() + assert "gpt-4o-mini" in stats["requests"]["by_model"] + + +class TestOpenAIResponsesErrorHandling: + """Test error handling for /v1/responses endpoint.""" + + def test_invalid_api_key(self, openai_responses_client): + """Invalid API key returns appropriate error.""" + response = openai_responses_client.post( + "/v1/responses", + headers={"Authorization": "Bearer invalid-key-123"}, + json={"model": "gpt-4o-mini", "input": "Hi"}, + ) + assert response.status_code >= 400 + + def test_invalid_model(self, openai_responses_client, api_key): + """Invalid model returns appropriate error.""" + response = openai_responses_client.post( + "/v1/responses", + headers={"Authorization": f"Bearer {api_key}"}, + json={"model": "nonexistent-model-xyz", "input": "Hi"}, + ) + assert response.status_code >= 400 + + def test_missing_input(self, openai_responses_client, api_key): + """Missing input handled gracefully.""" + response = openai_responses_client.post( + "/v1/responses", + headers={"Authorization": f"Bearer {api_key}"}, + json={"model": "gpt-4o-mini"}, + ) + # Should either return error or handle gracefully + assert response.status_code in [200, 400, 422] diff --git a/tests/test_proxy_passthrough_integration.py b/tests/test_proxy_passthrough_integration.py new file mode 100644 index 000000000..872d4e87a --- /dev/null +++ b/tests/test_proxy_passthrough_integration.py @@ -0,0 +1,486 @@ +"""Integration tests for proxy passthrough endpoints with real API calls. + +These tests verify that passthrough endpoints work correctly with real API calls +to OpenAI, Gemini, and Anthropic APIs. + +Required environment variables: +- OPENAI_API_KEY: For OpenAI /v1/models, /v1/embeddings, /v1/moderations +- GEMINI_API_KEY: For Gemini /v1beta/models, :embedContent +- ANTHROPIC_API_KEY: For Anthropic /v1/models + +Run with: + OPENAI_API_KEY=... GEMINI_API_KEY=... ANTHROPIC_API_KEY=... pytest tests/test_proxy_passthrough_integration.py -v +""" + +import os + +import pytest + +pytest.importorskip("fastapi") +pytest.importorskip("httpx") + +from fastapi.testclient import TestClient + +from headroom.proxy.server import ProxyConfig, create_app + +# ============================================================================= +# Fixtures +# ============================================================================= + + +@pytest.fixture +def openai_client(): + """Create test client configured for OpenAI passthrough.""" + config = ProxyConfig( + optimize=True, + cache_enabled=False, + rate_limit_enabled=False, + cost_tracking_enabled=False, + ) + app = create_app(config) + with TestClient(app) as client: + yield client + + +@pytest.fixture +def gemini_client(): + """Create test client configured for Gemini passthrough.""" + config = ProxyConfig( + optimize=True, + cache_enabled=False, + rate_limit_enabled=False, + cost_tracking_enabled=False, + ) + app = create_app(config) + with TestClient(app) as client: + yield client + + +@pytest.fixture +def anthropic_client(): + """Create test client configured for Anthropic passthrough.""" + config = ProxyConfig( + optimize=True, + cache_enabled=False, + rate_limit_enabled=False, + cost_tracking_enabled=False, + ) + app = create_app(config) + with TestClient(app) as client: + yield client + + +@pytest.fixture +def openai_api_key(): + """Get OpenAI API key from environment.""" + return os.environ.get("OPENAI_API_KEY") + + +@pytest.fixture +def gemini_api_key(): + """Get Gemini API key from environment.""" + return os.environ.get("GEMINI_API_KEY") + + +@pytest.fixture +def anthropic_api_key(): + """Get Anthropic API key from environment.""" + return os.environ.get("ANTHROPIC_API_KEY") + + +# ============================================================================= +# OpenAI Passthrough Tests +# ============================================================================= + + +@pytest.mark.skipif(not os.environ.get("OPENAI_API_KEY"), reason="OPENAI_API_KEY not set") +class TestOpenAIModels: + """Test OpenAI /v1/models endpoint passthrough.""" + + def test_list_models(self, openai_client, openai_api_key): + """GET /v1/models returns list of available models.""" + response = openai_client.get( + "/v1/models", headers={"Authorization": f"Bearer {openai_api_key}"} + ) + assert response.status_code == 200 + data = response.json() + + # Verify OpenAI models list format + assert "data" in data + assert "object" in data + assert data["object"] == "list" + assert len(data["data"]) > 0 + + # Verify model object structure + model = data["data"][0] + assert "id" in model + assert "object" in model + assert model["object"] == "model" + + def test_get_specific_model(self, openai_client, openai_api_key): + """GET /v1/models/{model_id} returns model details.""" + response = openai_client.get( + "/v1/models/gpt-4o-mini", headers={"Authorization": f"Bearer {openai_api_key}"} + ) + assert response.status_code == 200 + data = response.json() + + assert data["id"] == "gpt-4o-mini" + assert data["object"] == "model" + + def test_invalid_api_key(self, openai_client): + """Invalid API key returns authentication error.""" + response = openai_client.get( + "/v1/models", headers={"Authorization": "Bearer invalid-key-12345"} + ) + assert response.status_code == 401 + + +@pytest.mark.skipif(not os.environ.get("OPENAI_API_KEY"), reason="OPENAI_API_KEY not set") +class TestOpenAIEmbeddings: + """Test OpenAI /v1/embeddings endpoint passthrough.""" + + def test_create_embedding(self, openai_client, openai_api_key): + """POST /v1/embeddings creates embeddings successfully.""" + response = openai_client.post( + "/v1/embeddings", + headers={"Authorization": f"Bearer {openai_api_key}"}, + json={ + "model": "text-embedding-3-small", + "input": "The quick brown fox jumps over the lazy dog.", + }, + ) + assert response.status_code == 200 + data = response.json() + + # Verify embedding response format + assert "data" in data + assert "model" in data + assert "usage" in data + assert data["object"] == "list" + + # Verify embedding data + embedding = data["data"][0] + assert "embedding" in embedding + assert "index" in embedding + assert embedding["object"] == "embedding" + assert isinstance(embedding["embedding"], list) + assert len(embedding["embedding"]) > 0 + + # Verify usage + assert "prompt_tokens" in data["usage"] + assert "total_tokens" in data["usage"] + + def test_create_embedding_batch(self, openai_client, openai_api_key): + """POST /v1/embeddings with multiple inputs creates batch embeddings.""" + response = openai_client.post( + "/v1/embeddings", + headers={"Authorization": f"Bearer {openai_api_key}"}, + json={ + "model": "text-embedding-3-small", + "input": ["First text to embed", "Second text to embed", "Third text to embed"], + }, + ) + assert response.status_code == 200 + data = response.json() + + # Should return 3 embeddings + assert len(data["data"]) == 3 + for i, embedding in enumerate(data["data"]): + assert embedding["index"] == i + assert isinstance(embedding["embedding"], list) + + def test_embedding_invalid_model(self, openai_client, openai_api_key): + """Invalid model returns appropriate error.""" + response = openai_client.post( + "/v1/embeddings", + headers={"Authorization": f"Bearer {openai_api_key}"}, + json={"model": "nonexistent-embedding-model", "input": "Test text"}, + ) + assert response.status_code >= 400 + + +@pytest.mark.skipif(not os.environ.get("OPENAI_API_KEY"), reason="OPENAI_API_KEY not set") +class TestOpenAIModerations: + """Test OpenAI /v1/moderations endpoint passthrough.""" + + def test_moderation_safe_content(self, openai_client, openai_api_key): + """POST /v1/moderations on safe content returns no flags.""" + response = openai_client.post( + "/v1/moderations", + headers={"Authorization": f"Bearer {openai_api_key}"}, + json={"input": "I love sunny days and playing with my dog in the park."}, + ) + assert response.status_code == 200 + data = response.json() + + # Verify moderation response format + assert "id" in data + assert "model" in data + assert "results" in data + + # Safe content should not be flagged + result = data["results"][0] + assert "flagged" in result + assert "categories" in result + assert "category_scores" in result + # Safe content should generally not be flagged + # (though model may have false positives occasionally) + + def test_moderation_batch(self, openai_client, openai_api_key): + """POST /v1/moderations with multiple inputs.""" + response = openai_client.post( + "/v1/moderations", + headers={"Authorization": f"Bearer {openai_api_key}"}, + json={ + "input": [ + "Hello, how are you today?", + "What a beautiful sunset!", + "I enjoy reading books.", + ] + }, + ) + assert response.status_code == 200 + data = response.json() + + # Should return 3 moderation results + assert len(data["results"]) == 3 + + +# ============================================================================= +# Gemini Passthrough Tests +# ============================================================================= + + +@pytest.mark.skipif(not os.environ.get("GEMINI_API_KEY"), reason="GEMINI_API_KEY not set") +class TestGeminiModels: + """Test Gemini /v1beta/models endpoint passthrough.""" + + def test_list_models(self, gemini_client, gemini_api_key): + """GET /v1beta/models returns list of available models.""" + response = gemini_client.get(f"/v1beta/models?key={gemini_api_key}") + assert response.status_code == 200 + data = response.json() + + # Verify Gemini models list format + assert "models" in data + assert len(data["models"]) > 0 + + # Verify model object structure + model = data["models"][0] + assert "name" in model + assert "displayName" in model or "description" in model + + def test_get_specific_model(self, gemini_client, gemini_api_key): + """GET /v1beta/models/{model} returns model details.""" + response = gemini_client.get(f"/v1beta/models/gemini-2.0-flash?key={gemini_api_key}") + assert response.status_code == 200 + data = response.json() + + assert "name" in data + assert "gemini" in data["name"].lower() + + +@pytest.mark.skipif(not os.environ.get("GEMINI_API_KEY"), reason="GEMINI_API_KEY not set") +class TestGeminiEmbedContent: + """Test Gemini /v1beta/models/{model}:embedContent endpoint passthrough.""" + + def test_embed_content(self, gemini_client, gemini_api_key): + """POST :embedContent creates embeddings successfully.""" + response = gemini_client.post( + f"/v1beta/models/text-embedding-004:embedContent?key={gemini_api_key}", + json={"content": {"parts": [{"text": "The quick brown fox jumps over the lazy dog."}]}}, + ) + assert response.status_code == 200 + data = response.json() + + # Verify embedding response format + assert "embedding" in data + assert "values" in data["embedding"] + assert isinstance(data["embedding"]["values"], list) + assert len(data["embedding"]["values"]) > 0 + + def test_embed_content_with_task_type(self, gemini_client, gemini_api_key): + """POST :embedContent with task type specified.""" + response = gemini_client.post( + f"/v1beta/models/text-embedding-004:embedContent?key={gemini_api_key}", + json={ + "content": {"parts": [{"text": "What is the capital of France?"}]}, + "taskType": "RETRIEVAL_QUERY", + }, + ) + assert response.status_code == 200 + data = response.json() + + assert "embedding" in data + assert "values" in data["embedding"] + + +@pytest.mark.skipif(not os.environ.get("GEMINI_API_KEY"), reason="GEMINI_API_KEY not set") +class TestGeminiBatchEmbedContents: + """Test Gemini /v1beta/models/{model}:batchEmbedContents endpoint passthrough.""" + + def test_batch_embed_contents(self, gemini_client, gemini_api_key): + """POST :batchEmbedContents creates batch embeddings.""" + # Note: batchEmbedContents requires model field in each request + response = gemini_client.post( + f"/v1beta/models/text-embedding-004:batchEmbedContents?key={gemini_api_key}", + json={ + "requests": [ + { + "model": "models/text-embedding-004", + "content": {"parts": [{"text": "First document to embed"}]}, + }, + { + "model": "models/text-embedding-004", + "content": {"parts": [{"text": "Second document to embed"}]}, + }, + { + "model": "models/text-embedding-004", + "content": {"parts": [{"text": "Third document to embed"}]}, + }, + ] + }, + ) + # May return 400 if format changed, or 200 on success + assert response.status_code in [200, 400] + if response.status_code == 200: + data = response.json() + # Verify batch embedding response format + assert "embeddings" in data + assert len(data["embeddings"]) == 3 + + for embedding in data["embeddings"]: + assert "values" in embedding + assert isinstance(embedding["values"], list) + + +# ============================================================================= +# Anthropic Passthrough Tests +# ============================================================================= + + +@pytest.mark.skipif(not os.environ.get("ANTHROPIC_API_KEY"), reason="ANTHROPIC_API_KEY not set") +class TestAnthropicModels: + """Test Anthropic /v1/models endpoint passthrough.""" + + def test_list_models(self, anthropic_client, anthropic_api_key): + """GET /v1/models returns list of available models with x-api-key header.""" + response = anthropic_client.get( + "/v1/models", + headers={"x-api-key": anthropic_api_key, "anthropic-version": "2023-06-01"}, + ) + assert response.status_code == 200 + data = response.json() + + # Verify Anthropic models list format + assert "data" in data + assert len(data["data"]) > 0 + + # Verify model object structure + model = data["data"][0] + assert "id" in model + assert "type" in model + + def test_get_specific_model(self, anthropic_client, anthropic_api_key): + """GET /v1/models/{model_id} returns model details.""" + # First get the list to find a valid model ID + list_response = anthropic_client.get( + "/v1/models", + headers={"x-api-key": anthropic_api_key, "anthropic-version": "2023-06-01"}, + ) + assert list_response.status_code == 200 + models = list_response.json().get("data", []) + + if not models: + pytest.skip("No models available") + + # Use the first available model + model_id = models[0]["id"] + + response = anthropic_client.get( + f"/v1/models/{model_id}", + headers={"x-api-key": anthropic_api_key, "anthropic-version": "2023-06-01"}, + ) + assert response.status_code == 200 + data = response.json() + + assert "id" in data + assert data["id"] == model_id + + def test_invalid_api_key(self, anthropic_client): + """Invalid API key returns authentication error.""" + response = anthropic_client.get( + "/v1/models", + headers={"x-api-key": "invalid-key-12345", "anthropic-version": "2023-06-01"}, + ) + assert response.status_code == 401 + + +# ============================================================================= +# Proxy Stats Tests +# ============================================================================= + + +@pytest.mark.skipif(not os.environ.get("OPENAI_API_KEY"), reason="OPENAI_API_KEY not set") +class TestPassthroughStats: + """Test that passthrough requests are tracked in proxy stats.""" + + def test_stats_track_passthrough_requests(self, openai_client, openai_api_key): + """Verify passthrough requests are tracked in stats.""" + # Make a passthrough request + openai_client.get("/v1/models", headers={"Authorization": f"Bearer {openai_api_key}"}) + + # Check stats + stats_response = openai_client.get("/stats") + assert stats_response.status_code == 200 + stats = stats_response.json() + + # Verify stats structure + assert "requests" in stats + assert "total" in stats["requests"] + assert stats["requests"]["total"] >= 1 + + def test_stats_track_embeddings_requests(self, openai_client, openai_api_key): + """Verify embeddings passthrough requests are tracked.""" + # Make an embeddings request + openai_client.post( + "/v1/embeddings", + headers={"Authorization": f"Bearer {openai_api_key}"}, + json={"model": "text-embedding-3-small", "input": "Test embedding"}, + ) + + # Check stats + stats_response = openai_client.get("/stats") + stats = stats_response.json() + + # Should track embeddings under openai provider + assert "openai" in stats["requests"]["by_provider"] + + +# ============================================================================= +# Error Handling Tests +# ============================================================================= + + +class TestPassthroughErrorHandling: + """Test error handling for passthrough endpoints.""" + + def test_missing_auth_header_openai(self, openai_client): + """Missing auth header returns appropriate error.""" + response = openai_client.get("/v1/models") + # OpenAI requires authentication + assert response.status_code >= 400 + + @pytest.mark.skipif(not os.environ.get("OPENAI_API_KEY"), reason="OPENAI_API_KEY not set") + def test_invalid_json_body(self, openai_client, openai_api_key): + """Invalid JSON body returns 400 error.""" + response = openai_client.post( + "/v1/embeddings", + headers={ + "Authorization": f"Bearer {openai_api_key}", + "Content-Type": "application/json", + }, + content=b"not valid json", + ) + assert response.status_code >= 400