Spaces:
Running
Running
| """OpenAI SDK-compatible embeddings API, intentionally separate from Voyage routes. | |
| Configure OpenAI(base_url="https://HOST/openai/v1", api_key="..."). | |
| The client is responsible for model-specific query/document prefixes. | |
| """ | |
| from __future__ import annotations | |
| import base64 | |
| import logging | |
| import struct | |
| from typing import Literal | |
| from fastapi import APIRouter, Depends | |
| from fastapi.responses import JSONResponse | |
| from pydantic import BaseModel, ConfigDict, Field, field_validator | |
| from app.embedding import MODEL_ID as E5_ID, TooManyTokens | |
| from app.bge_m3 import BGE_M3_ID | |
| from app.models import ModelRegistry, QWEN_ID, ServiceBusy | |
| from app.voyage_api import authenticate, resident_models | |
| router = APIRouter(prefix="/openai/v1", tags=["OpenAI-compatible"]) | |
| log = logging.getLogger(__name__) | |
| DenseModelName = Literal[ | |
| "e5", "qwen", "bge-m3", "intfloat/multilingual-e5-small", | |
| "Qwen/Qwen3-Embedding-0.6B", "BAAI/bge-m3" | |
| ] | |
| class OpenAIEmbeddingRequest(BaseModel): | |
| model_config = ConfigDict(extra="forbid") | |
| model: DenseModelName | |
| input: str | list[str] = Field( | |
| description="Raw input; client must add model-specific query/document prefixes" | |
| ) | |
| dimensions: int | None = Field(default=None, ge=1) | |
| encoding_format: Literal["float", "base64"] = "float" | |
| user: str | None = None | |
| # Non-standard client probe flag (e.g. OpenAI-compatible integrations). | |
| # It does not change server-side processing; accept it for compatibility. | |
| drop_params: bool | None = None | |
| def validate_input(cls, value): | |
| texts = [value] if isinstance(value, str) else value | |
| if not texts or len(texts) > 1000: | |
| raise ValueError("input must contain between 1 and 1000 texts") | |
| if any(not text.strip() for text in texts): | |
| raise ValueError("input cannot contain blank strings") | |
| return value | |
| def api_error(status: int, message: str, error_type="invalid_request_error", code=None): | |
| return JSONResponse(status_code=status, content={ | |
| "error": {"message": message, "type": error_type, "param": None, "code": code} | |
| }) | |
| def create_embedding( | |
| request: OpenAIEmbeddingRequest, | |
| registry: ModelRegistry = Depends(resident_models), | |
| ): | |
| model = request.model | |
| if model in ("bge-m3", BGE_M3_ID): | |
| canonical, engine = BGE_M3_ID, registry.bge_m3 | |
| elif model in ("qwen", QWEN_ID): | |
| canonical, engine = QWEN_ID, registry.qwen | |
| else: | |
| canonical, engine = E5_ID, registry.e5 | |
| if engine is None: | |
| return api_error(503, f"Model {model} is not enabled", "server_error") | |
| if request.dimensions is not None and request.dimensions != engine.dimensions: | |
| log.warning( | |
| "OpenAI embeddings rejected: reason=unsupported_dimensions model=%s " | |
| "requested_dimensions=%d expected_dimensions=%d", | |
| model, request.dimensions, engine.dimensions, | |
| ) | |
| return api_error( | |
| 400, | |
| f"Model {model} supports exactly {engine.dimensions} dimensions", | |
| code="unsupported_dimensions", | |
| ) | |
| texts = [request.input] if isinstance(request.input, str) else request.input | |
| try: | |
| results = registry.openai_encode(texts, canonical) | |
| except TooManyTokens as exc: | |
| log.warning( | |
| "OpenAI embeddings rejected: reason=context_length_exceeded model=%s " | |
| "batch_size=%d input_index=%s input_tokens=%d max_tokens=%d", | |
| model, len(texts), exc.input_index, exc.count, exc.limit, | |
| ) | |
| return api_error(400, str(exc), code="context_length_exceeded") | |
| except ValueError as exc: | |
| # Never log input texts or untrusted tokenizer exception contents. | |
| reason = ("batch_token_budget_exceeded" | |
| if str(exc) == "Batch exceeds 65536 total input tokens" | |
| else "invalid_batch") | |
| log.warning( | |
| "OpenAI embeddings rejected: reason=%s model=%s batch_size=%d", | |
| reason, model, len(texts), | |
| ) | |
| return api_error(400, str(exc)) | |
| except ServiceBusy as exc: | |
| response = api_error(429, str(exc), "rate_limit_error", "rate_limit_exceeded") | |
| response.headers["Retry-After"] = "2" | |
| return response | |
| except Exception: | |
| log.exception("OpenAI-compatible embedding failed") | |
| return api_error(500, "Embedding generation failed", "server_error") | |
| def encode(vector): | |
| if request.encoding_format == "float": | |
| return vector | |
| raw = struct.pack(f"<{len(vector)}f", *vector) | |
| return base64.b64encode(raw).decode("ascii") | |
| total_tokens = sum(result.tokens_count for result in results) | |
| # Compact preview of successful inputs for source attribution. | |
| previews = [(i, value[:240]) for i, value in enumerate(texts[:3])] | |
| log.warning("Embeddings input preview: model=%s chunks=%d previews=%r", | |
| model, len(texts), previews) | |
| return { | |
| "object": "list", | |
| "data": [ | |
| {"object": "embedding", "index": index, "embedding": encode(item.embedding)} | |
| for index, item in enumerate(results) | |
| ], | |
| "model": model, | |
| "usage": {"prompt_tokens": total_tokens, "total_tokens": total_tokens}, | |
| } | |
| def list_models(registry: ModelRegistry = Depends(resident_models)): | |
| available = [] | |
| if registry.e5 is not None: | |
| available.extend(("e5", E5_ID)) | |
| if registry.qwen is not None: | |
| available.extend(("qwen", QWEN_ID)) | |
| if registry.bge_m3 is not None: | |
| available.extend(("bge-m3", BGE_M3_ID)) | |
| return { | |
| "object": "list", | |
| "data": [ | |
| {"id": mid, "object": "model", "created": 0, "owned_by": "local"} | |
| for mid in available | |
| ], | |
| } | |