feat: Add support for embeddings endpoint

This commit introduces a new OpenAI-compatible `/v1/embeddings` endpoint to the proxy, enabling the creation of text embeddings.

Key changes include:
- A new `aembedding` method in `RotatingClient` to handle API calls, key rotation, and rate limit errors for embedding requests.
- The `UsageManager` is updated to correctly track token usage from embedding responses.
- Request logging is enhanced to categorize logs into `completions` and `embeddings` subdirectories for better organization.
- A sanitizer for embedding requests is added to filter unsupported parameters before calling the OpenAI API.

src/proxy_app/main.py

@@ -10,9 +10,18 @@ import logging
 from pathlib import Path
 import sys
 import json
-from typing import AsyncGenerator, Any
+from typing import AsyncGenerator, Any, List, Optional
+from pydantic import BaseModel
 import argparse
 
+# --- Pydantic Models ---
+class EmbeddingRequest(BaseModel):
+    model: str
+    input: List[str]
+    input_type: Optional[str] = None
+    dimensions: Optional[int] = None
+    user: Optional[str] = None
+
 # --- Argument Parsing ---
 parser = argparse.ArgumentParser(description="API Key Proxy Server")
 parser.add_argument("--host", type=str, default="0.0.0.0", help="Host to bind the server to.")
@@ -188,7 +197,8 @@ async def streaming_response_wrapper(
             log_request_response(
                 request_data=request_data,
                 response_data=full_response,
-                is_streaming=True
+                is_streaming=True,
+                log_type="completion"
             )
 
 @app.post("/v1/chat/completions")
@@ -219,7 +229,8 @@ async def chat_completions(
                 log_request_response(
                     request_data=request_data,
                     response_data=response.dict(),
-                    is_streaming=False
+                    is_streaming=False,
+                    log_type="completion"
                 )
             return response
 
@@ -234,7 +245,53 @@ async def chat_completions(
             log_request_response(
                 request_data=request_data,
                 response_data={"error": str(e)},
-                is_streaming=request_data.get("stream", False)
+                is_streaming=request_data.get("stream", False),
+                log_type="completion"
+            )
+        raise HTTPException(status_code=500, detail=str(e))
+
+@app.post("/v1/embeddings")
+async def embeddings(
+    request: Request,
+    body: EmbeddingRequest,
+    client: RotatingClient = Depends(get_rotating_client),
+    _ = Depends(verify_api_key)
+):
+    """
+    OpenAI-compatible endpoint for creating embeddings.
+    """
+    try:
+        request_data = body.dict(exclude_none=True)
+        response = await client.aembedding(**request_data)
+        
+        if ENABLE_REQUEST_LOGGING:
+            response_summary = {
+                "model": response.model,
+                "object": response.object,
+                "usage": response.usage.dict(),
+                "data_count": len(response.data),
+                "embedding_dimensions": len(response.data[0].embedding) if response.data else 0
+            }
+            log_request_response(
+                request_data=request_data,
+                response_data=response_summary,
+                is_streaming=False,
+                log_type="embedding"
+            )
+        return response
+
+    except Exception as e:
+        logging.error(f"Embedding request failed: {e}")
+        if ENABLE_REQUEST_LOGGING:
+            try:
+                request_data = await request.json()
+            except json.JSONDecodeError:
+                request_data = {"error": "Could not parse request body"}
+            log_request_response(
+                request_data=request_data,
+                response_data={"error": str(e)},
+                is_streaming=False,
+                log_type="embedding"
             )
         raise HTTPException(status_code=500, detail=str(e))
 

src/proxy_app/request_logger.py

@@ -3,18 +3,39 @@ import os
 from datetime import datetime
 from pathlib import Path
 import uuid
+from typing import Literal
 
 LOGS_DIR = Path(__file__).resolve().parent.parent.parent / "logs"
+COMPLETIONS_LOGS_DIR = LOGS_DIR / "completions"
+EMBEDDINGS_LOGS_DIR = LOGS_DIR / "embeddings"
+
+# Create directories if they don't exist
 LOGS_DIR.mkdir(exist_ok=True)
+COMPLETIONS_LOGS_DIR.mkdir(exist_ok=True)
+EMBEDDINGS_LOGS_DIR.mkdir(exist_ok=True)
 
-def log_request_response(request_data: dict, response_data: dict, is_streaming: bool):
+def log_request_response(
+    request_data: dict,
+    response_data: dict,
+    is_streaming: bool,
+    log_type: Literal["completion", "embedding"]
+):
     """
-    Logs the request and response data to a single file in the logs directory.
+    Logs the request and response data to a file in the appropriate log directory.
     """
     try:
         timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
         unique_id = uuid.uuid4()
-        filename = LOGS_DIR / f"{timestamp}_{unique_id}.json"
+        
+        if log_type == "completion":
+            target_dir = COMPLETIONS_LOGS_DIR
+        elif log_type == "embedding":
+            target_dir = EMBEDDINGS_LOGS_DIR
+        else:
+            # Fallback to the main logs directory if log_type is invalid
+            target_dir = LOGS_DIR
+
+        filename = target_dir / f"{timestamp}_{unique_id}.json"
 
         log_content = {
             "request": request_data,

src/rotator_library/client.py

@@ -305,6 +305,87 @@ class RotatingClient:
         
         raise Exception("Failed to complete the request: No available API keys for the provider or all keys failed.")
 
+    async def aembedding(self, **kwargs) -> Any:
+        """
+        Performs an embedding call with smart key rotation and retry logic.
+        """
+        model = kwargs.get("model")
+        if not model:
+            raise ValueError("'model' is a required parameter.")
+
+        provider = model.split('/')[0]
+        if provider not in self.api_keys:
+            raise ValueError(f"No API keys configured for provider: {provider}")
+
+        keys_for_provider = self.api_keys[provider]
+        tried_keys = set()
+        last_exception = None
+
+        while len(tried_keys) < len(keys_for_provider):
+            current_key = None
+            key_acquired = False
+            try:
+                keys_to_try = [k for k in keys_for_provider if k not in tried_keys]
+                if not keys_to_try:
+                    break
+
+                current_key = await self.usage_manager.acquire_key(
+                    available_keys=keys_to_try,
+                    model=model
+                )
+                key_acquired = True
+                tried_keys.add(current_key)
+
+                litellm_kwargs = self.all_providers.get_provider_kwargs(**kwargs.copy())
+                litellm_kwargs = sanitize_request_payload(litellm_kwargs, model)
+
+                for attempt in range(self.max_retries):
+                    try:
+                        lib_logger.info(f"Attempting embedding call with key ...{current_key[-4:]} (Attempt {attempt + 1}/{self.max_retries})")
+                        
+                        response = await litellm.aembedding(api_key=current_key, **litellm_kwargs)
+                        
+                        await self.usage_manager.record_success(current_key, model, response)
+                        await self.usage_manager.release_key(current_key, model)
+                        key_acquired = False
+                        return response
+
+                    except Exception as e:
+                        last_exception = e
+                        log_failure(api_key=current_key, model=model, attempt=attempt + 1, error=e, request_data=kwargs)
+                        
+                        classified_error = classify_error(e)
+
+                        if classified_error.error_type in ['invalid_request', 'context_window_exceeded']:
+                            lib_logger.error(f"Unrecoverable error '{classified_error.error_type}' with key ...{current_key[-4:]}. Failing request.")
+                            raise last_exception
+
+                        if classified_error.error_type in ['server_error', 'api_connection']:
+                            await self.usage_manager.record_failure(current_key, model, classified_error)
+                            
+                            if attempt >= self.max_retries - 1:
+                                lib_logger.warning(f"Key ...{current_key[-4:]} failed on final retry for {classified_error.error_type}. Trying next key.")
+                                break
+
+                            base_wait = 5 if classified_error.error_type == 'api_connection' else 1
+                            wait_time = classified_error.retry_after or (base_wait * (2 ** attempt)) + random.uniform(0, 1)
+                            
+                            lib_logger.warning(f"Key ...{current_key[-4:]} encountered a {classified_error.error_type}. Retrying in {wait_time:.2f} seconds...")
+                            await asyncio.sleep(wait_time)
+                            continue
+
+                        await self.usage_manager.record_failure(current_key, model, classified_error)
+                        lib_logger.warning(f"Key ...{current_key[-4:]} encountered '{classified_error.error_type}'. Trying next key.")
+                        break
+            finally:
+                if key_acquired and current_key:
+                    await self.usage_manager.release_key(current_key, model)
+
+        if last_exception:
+            raise last_exception
+        
+        raise Exception("Failed to complete the request: No available API keys for the provider or all keys failed.")
+
     def token_count(self, model: str, text: str = None, messages: List[Dict[str, str]] = None) -> int:
         """Calculates the number of tokens for a given text or list of messages."""
         if not model:

src/rotator_library/request_sanitizer.py

@@ -4,6 +4,9 @@ def sanitize_request_payload(payload: Dict[str, Any], model: str) -> Dict[str, A
     """
     Removes unsupported parameters from the request payload based on the model.
     """
+    if "dimensions" in payload and not model.startswith("openai/text-embedding-3"):
+        del payload["dimensions"]
+        
     if payload.get("thinking") == {"type": "enabled", "budget_tokens": -1}:
         if model not in ["gemini/gemini-2.5-pro", "gemini/gemini-2.5-flash"]:
             del payload["thinking"]

src/rotator_library/usage_manager.py

@@ -244,11 +244,17 @@ class UsageManager:
             if completion_response and hasattr(completion_response, 'usage') and completion_response.usage:
                 usage = completion_response.usage
                 daily_model_data["prompt_tokens"] += usage.prompt_tokens
-                daily_model_data["completion_tokens"] += usage.completion_tokens
+                daily_model_data["completion_tokens"] += getattr(usage, 'completion_tokens', 0) # Not present in embedding responses
                 lib_logger.info(f"Recorded usage from final stream object for key ...{key[-4:]}")
                 try:
-                    cost = litellm.completion_cost(completion_response=completion_response)
-                    daily_model_data["approx_cost"] += cost
+                    # Differentiate cost calculation based on response type
+                    if isinstance(completion_response, litellm.EmbeddingResponse):
+                        cost = litellm.embedding_cost(embedding_response=completion_response)
+                    else:
+                        cost = litellm.completion_cost(completion_response=completion_response)
+                    
+                    if cost is not None:
+                        daily_model_data["approx_cost"] += cost
                 except Exception as e:
                     lib_logger.warning(f"Could not calculate cost for model {model}: {e}")
             else: