diff --git a/app/api/v1/api.py b/app/api/v1/api.py index fb97ecce..a63a703e 100644 --- a/app/api/v1/api.py +++ b/app/api/v1/api.py @@ -42,6 +42,7 @@ workspaces, workspace_iam, dashboard, + llm_gateway, ) api_router = APIRouter() @@ -87,3 +88,4 @@ api_router.include_router(workspaces.router) api_router.include_router(workspace_iam.router) api_router.include_router(dashboard.router) +api_router.include_router(llm_gateway.router) diff --git a/app/api/v1/routes/aiproviders.py b/app/api/v1/routes/aiproviders.py index 1cbe1d0f..0d44f90c 100644 --- a/app/api/v1/routes/aiproviders.py +++ b/app/api/v1/routes/aiproviders.py @@ -12,7 +12,7 @@ from fastapi import APIRouter, Depends, HTTPException, status from sqlalchemy import desc, func from sqlalchemy.orm import Session -from typing import List +from typing import List, Optional from uuid import UUID from app.dependencies import get_db, get_organization_id @@ -22,11 +22,16 @@ ) from app.core.encryption import encrypt_api_key from app.services.credentials.resolver import clear_other_defaults +from app.services.ai.llm_gateway import ( + GATEWAY_MANAGED_KEY_SENTINEL, + gateway_managed_credentials_enabled, + is_gateway_managed_stored_key, +) router = APIRouter(prefix="/aiproviders", tags=["aiproviders"]) -def _scrub_for_response(db: Session, instance: AIProvider) -> AIProvider: +def _scrub_for_response(db: Session, instance: AIProvider) -> AIProviderResponse: """Detach the row from the session before clearing ``api_key``. The same SQLAlchemy session is reused across requests in tests; if we @@ -34,9 +39,27 @@ def _scrub_for_response(db: Session, instance: AIProvider) -> AIProvider: treats it as a pending UPDATE and tries to flush ``api_key=NULL`` on the next request — which violates the column's NOT NULL constraint. """ + gateway_managed = is_gateway_managed_stored_key(instance.api_key) db.expunge(instance) instance.api_key = None - return instance + response = AIProviderResponse.model_validate(instance) + return response.model_copy(update={"gateway_managed": gateway_managed}) + + +def _encrypt_provider_api_key(api_key: Optional[str]) -> str: + trimmed = (api_key or "").strip() + if trimmed: + return encrypt_api_key(trimmed) + if gateway_managed_credentials_enabled(): + return encrypt_api_key(GATEWAY_MANAGED_KEY_SENTINEL) + raise HTTPException( + status_code=status.HTTP_400_BAD_REQUEST, + detail=( + "api_key is required. When LLM gateway-managed credentials " + "are enabled (passthrough_provider_keys: false), you can omit " + "the key and provider secrets will be resolved by the gateway." + ), + ) @router.post("", response_model=AIProviderResponse, status_code=status.HTTP_201_CREATED, operation_id="createAIProvider") @@ -63,7 +86,7 @@ async def create_aiprovider( requested_default = bool(aiprovider.is_default) will_be_default = requested_default or existing_default is None - encrypted_api_key = encrypt_api_key(aiprovider.api_key) + encrypted_api_key = _encrypt_provider_api_key(aiprovider.api_key) db_aiprovider = AIProvider( organization_id=organization_id, provider=provider_value, diff --git a/app/api/v1/routes/llm_gateway.py b/app/api/v1/routes/llm_gateway.py new file mode 100644 index 00000000..65924b8f --- /dev/null +++ b/app/api/v1/routes/llm_gateway.py @@ -0,0 +1,94 @@ +"""Organization-level LLM gateway settings.""" + +from typing import Literal, Optional +from uuid import UUID + +from fastapi import APIRouter, Depends +from pydantic import BaseModel, Field +from sqlalchemy.orm import Session + +from app.dependencies import get_db, get_organization_id +from app.services.ai.llm_gateway_settings import ( + GatewayMode, + GatewayTypeOverride, + get_org_settings, + set_org_settings, +) + +router = APIRouter(prefix="/organizations/llm-gateway", tags=["LLM Gateway"]) + + +class LLMGatewaySettingsResponse(BaseModel): + mode: GatewayMode + gateway_type: GatewayTypeOverride + base_url: Optional[str] = None + has_virtual_key: bool = False + has_master_key: bool = False + platform_enabled: bool = False + platform_gateway_type: Literal["bifrost", "litellm_proxy"] = "bifrost" + platform_base_url: Optional[str] = None + effective_routing: Literal["direct", "bifrost", "litellm_proxy"] + effective_gateway_type: Optional[Literal["bifrost", "litellm_proxy"]] = None + effective_base_url: Optional[str] = None + effective_has_virtual_key: bool = False + effective_has_master_key: bool = False + gateway_managed_credentials: bool = False + + +class LLMGatewaySettingsUpdate(BaseModel): + mode: GatewayMode = Field( + ..., + description="inherit: use platform default; enabled: force gateway; disabled: opt out", + ) + gateway_type: GatewayTypeOverride = Field( + default="inherit", + description="inherit: use platform gateway type; bifrost or litellm_proxy to override", + ) + base_url: Optional[str] = Field( + default=None, + description="Optional org-specific gateway URL override.", + ) + virtual_key: Optional[str] = Field( + default=None, + description="Optional Bifrost virtual key (x-bf-vk). Omit to keep existing.", + ) + master_key: Optional[str] = Field( + default=None, + description="Optional LiteLLM Proxy master key. Omit to keep existing.", + ) + clear_virtual_key: bool = Field( + default=False, + description="When true, remove the stored org virtual key.", + ) + clear_master_key: bool = Field( + default=False, + description="When true, remove the stored org master key.", + ) + + +@router.get("", response_model=LLMGatewaySettingsResponse) +def get_llm_gateway_settings( + organization_id: UUID = Depends(get_organization_id), + db: Session = Depends(get_db), +): + return LLMGatewaySettingsResponse(**get_org_settings(organization_id, db)) + + +@router.put("", response_model=LLMGatewaySettingsResponse) +def update_llm_gateway_settings( + body: LLMGatewaySettingsUpdate, + organization_id: UUID = Depends(get_organization_id), + db: Session = Depends(get_db), +): + result = set_org_settings( + organization_id, + db, + mode=body.mode, + gateway_type=body.gateway_type, + base_url=body.base_url, + virtual_key=body.virtual_key, + master_key=body.master_key, + clear_virtual_key=body.clear_virtual_key, + clear_master_key=body.clear_master_key, + ) + return LLMGatewaySettingsResponse(**result) diff --git a/app/config.py b/app/config.py index 6ec3ebd6..356176ae 100644 --- a/app/config.py +++ b/app/config.py @@ -163,6 +163,14 @@ class Settings(BaseSettings): JUDGE_ALIGNMENT_ENABLED: bool = True JUDGE_ALIGNMENT_CSV_MAX_ROWS: int = 5000 + # LLM gateway (optional platform-wide proxy for batch LLM calls). + LLM_GATEWAY_ENABLED: bool = False + LLM_GATEWAY_TYPE: str = "bifrost" # bifrost | litellm_proxy + LLM_GATEWAY_BASE_URL: Optional[str] = None + LLM_GATEWAY_VIRTUAL_KEY: Optional[str] = None + LLM_GATEWAY_MASTER_KEY: Optional[str] = None + LLM_GATEWAY_PASSTHROUGH_PROVIDER_KEYS: bool = True + model_config = SettingsConfigDict( env_file=".env", env_file_encoding="utf-8", @@ -600,6 +608,28 @@ def load_config_from_file(config_path: str) -> None: if "csv_max_rows" in ja_cfg: settings.JUDGE_ALIGNMENT_CSV_MAX_ROWS = int(ja_cfg["csv_max_rows"]) + def _apply_llm_gateway_settings(gateway_cfg: dict, *, gateway_type: str) -> None: + if "enabled" in gateway_cfg: + settings.LLM_GATEWAY_ENABLED = bool(gateway_cfg["enabled"]) + settings.LLM_GATEWAY_TYPE = gateway_type + if gateway_cfg.get("base_url"): + settings.LLM_GATEWAY_BASE_URL = gateway_cfg["base_url"] + if gateway_cfg.get("virtual_key"): + settings.LLM_GATEWAY_VIRTUAL_KEY = gateway_cfg["virtual_key"] + if gateway_cfg.get("master_key"): + settings.LLM_GATEWAY_MASTER_KEY = gateway_cfg["master_key"] + if "passthrough_provider_keys" in gateway_cfg: + settings.LLM_GATEWAY_PASSTHROUGH_PROVIDER_KEYS = bool( + gateway_cfg["passthrough_provider_keys"] + ) + + if "llm_gateway" in config_data: + llm_cfg = config_data["llm_gateway"] + gateway_type = (llm_cfg.get("type") or "bifrost").strip().lower() + if gateway_type not in ("bifrost", "litellm_proxy"): + gateway_type = "bifrost" + _apply_llm_gateway_settings(llm_cfg, gateway_type=gateway_type) + if "operational" in config_data: operational_config = config_data["operational"] if "public" in operational_config: diff --git a/app/migrations/050_llm_gateway_settings.py b/app/migrations/050_llm_gateway_settings.py new file mode 100644 index 00000000..dd3b72f2 --- /dev/null +++ b/app/migrations/050_llm_gateway_settings.py @@ -0,0 +1,48 @@ +""" +Migration: Add per-organization LLM gateway settings JSON column. +""" + +from sqlalchemy import text +from sqlalchemy.orm import Session + +description = "Add organizations.llm_gateway_settings JSON for per-org LLM gateway overrides" + + +def _column_exists(db: Session, column_name: str) -> bool: + row = db.execute( + text( + """ + SELECT 1 + FROM information_schema.columns + WHERE table_name = 'organizations' + AND column_name = :column_name + """ + ), + {"column_name": column_name}, + ).first() + return row is not None + + +def upgrade(db: Session): + if _column_exists(db, "llm_gateway_settings"): + print("Column llm_gateway_settings already exists on organizations, skipping...") + return + + db.execute( + text( + """ + ALTER TABLE organizations + ADD COLUMN llm_gateway_settings JSON + """ + ) + ) + + db.commit() + print("Added llm_gateway_settings column to organizations") + + +def downgrade(db: Session): + db.execute( + text("ALTER TABLE organizations DROP COLUMN IF EXISTS llm_gateway_settings") + ) + db.commit() diff --git a/app/migrations/051_rename_bifrost_gateway_settings_column.py b/app/migrations/051_rename_bifrost_gateway_settings_column.py new file mode 100644 index 00000000..bfc482d5 --- /dev/null +++ b/app/migrations/051_rename_bifrost_gateway_settings_column.py @@ -0,0 +1,72 @@ +""" +Migration: Rename legacy organizations.bifrost_gateway_settings column. + +For databases that applied an earlier 050 migration before the column was +renamed to llm_gateway_settings. +""" + +from sqlalchemy import text +from sqlalchemy.orm import Session + +description = ( + "Rename organizations.bifrost_gateway_settings to llm_gateway_settings" +) + + +def _column_exists(db: Session, column_name: str) -> bool: + row = db.execute( + text( + """ + SELECT 1 + FROM information_schema.columns + WHERE table_name = 'organizations' + AND column_name = :column_name + """ + ), + {"column_name": column_name}, + ).first() + return row is not None + + +def upgrade(db: Session): + if _column_exists(db, "llm_gateway_settings"): + print("Column llm_gateway_settings already exists, skipping rename...") + return + + if not _column_exists(db, "bifrost_gateway_settings"): + print("Legacy bifrost_gateway_settings column not found, skipping rename...") + return + + db.execute( + text( + """ + ALTER TABLE organizations + RENAME COLUMN bifrost_gateway_settings TO llm_gateway_settings + """ + ) + ) + db.commit() + print( + "Renamed organizations.bifrost_gateway_settings " + "to llm_gateway_settings" + ) + + +def downgrade(db: Session): + if _column_exists(db, "bifrost_gateway_settings"): + print("Column bifrost_gateway_settings already exists, skipping rename...") + return + + if not _column_exists(db, "llm_gateway_settings"): + print("Column llm_gateway_settings not found, skipping rename...") + return + + db.execute( + text( + """ + ALTER TABLE organizations + RENAME COLUMN llm_gateway_settings TO bifrost_gateway_settings + """ + ) + ) + db.commit() diff --git a/app/models/database.py b/app/models/database.py index 476fb7ed..baf9be8d 100644 --- a/app/models/database.py +++ b/app/models/database.py @@ -54,6 +54,8 @@ class Organization(Base): # Shape: {"min_labels_to_evaluate": int, "min_labels_to_optimize": int} # Falls back to system defaults (20 / 50) when null. judge_alignment_settings = Column(JSON, nullable=True) + # Per-org LLM gateway overrides (enabled, gateway_type, base_url, keys). + llm_gateway_settings = Column(JSON, nullable=True) created_at = Column(DateTime(timezone=True), server_default=func.now()) updated_at = Column(DateTime(timezone=True), server_default=func.now(), onupdate=func.now()) diff --git a/app/models/schemas.py b/app/models/schemas.py index eaf315c2..ace77f74 100644 --- a/app/models/schemas.py +++ b/app/models/schemas.py @@ -669,7 +669,13 @@ class S3UploadResponse(BaseModel): class AIProviderCreate(BaseModel): """Schema for creating an AI Provider.""" provider: ModelProvider - api_key: str = Field(..., min_length=1) + api_key: Optional[str] = Field( + None, + description=( + "Provider API key. Optional when the platform uses Bifrost " + "gateway-managed credentials (passthrough_provider_keys: false)." + ), + ) name: Optional[str] = None is_default: Optional[bool] = Field( None, @@ -679,6 +685,14 @@ class AIProviderCreate(BaseModel): ), ) + @field_validator("api_key") + @classmethod + def validate_api_key(cls, v: Optional[str]) -> Optional[str]: + if v is None: + return v + trimmed = v.strip() + return trimmed or None + class AIProviderUpdate(BaseModel): """Schema for updating an AI Provider.""" @@ -695,6 +709,7 @@ class AIProviderResponse(BaseModel): name: Optional[str] is_active: bool is_default: bool = False + gateway_managed: bool = False created_at: datetime updated_at: datetime last_tested_at: Optional[datetime] diff --git a/app/services/ai/llm_gateway.py b/app/services/ai/llm_gateway.py new file mode 100644 index 00000000..85d41b14 --- /dev/null +++ b/app/services/ai/llm_gateway.py @@ -0,0 +1,357 @@ +""" +LLM gateway resolver for batch/eval LiteLLM workloads. + +Supports Bifrost (/litellm proxy) and self-hosted LiteLLM Proxy gateways. +Platform defaults live in ``config.yml``; per-org overrides are stored in +``organizations.llm_gateway_settings`` JSON. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Any, Dict, Literal, Optional +from urllib.parse import urlparse +from uuid import UUID + +from loguru import logger +from sqlalchemy.orm import Session + +from app.config import settings + + +GatewayType = Literal["bifrost", "litellm_proxy"] +EffectiveRouting = Literal["direct", "bifrost", "litellm_proxy"] + +# Backward-compatible alias used by stored AI provider credentials. +GATEWAY_MANAGED_KEY_SENTINEL = "__bifrost_gateway_managed__" +# LiteLLM's OpenAI client path requires a non-empty api_key even when +# api_base points at a gateway. Gateways authenticate via headers or master keys. +LITELLM_GATEWAY_PLACEHOLDER_API_KEY = "gateway-managed" + + +def gateway_managed_credentials_enabled() -> bool: + """True when the platform expects provider keys to live in the gateway only.""" + return not bool(settings.LLM_GATEWAY_PASSTHROUGH_PROVIDER_KEYS) + + +def is_gateway_managed_stored_key(encrypted_api_key: str) -> bool: + """Return True when the stored credential is the gateway-managed placeholder.""" + try: + from app.core.encryption import decrypt_api_key + + return decrypt_api_key(encrypted_api_key) == GATEWAY_MANAGED_KEY_SENTINEL + except Exception: + return False + + +def resolve_litellm_api_key( + organization_id: UUID, + db: Session, + ai_provider: Any, +) -> Optional[str]: + """Decrypt the org credential, or return None when the gateway supplies the key.""" + from app.core.encryption import decrypt_api_key + + try: + raw_key = decrypt_api_key(ai_provider.api_key) + except Exception as exc: + raise RuntimeError( + f"Failed to decrypt API key for provider {ai_provider.provider}: {exc}" + ) from exc + + if raw_key != GATEWAY_MANAGED_KEY_SENTINEL: + return raw_key + + gateway = resolve_llm_gateway(organization_id, db) + if gateway and not gateway.passthrough_provider_keys: + return None + + raise RuntimeError( + "AI provider is configured for gateway-managed credentials, " + "but the LLM gateway is not active for this organization. " + "Enable the LLM Gateway or add a provider API key." + ) + + +@dataclass(frozen=True) +class LLMGatewayConfig: + """Resolved gateway settings for a single LiteLLM call.""" + + gateway_type: GatewayType + api_base: str + virtual_key: Optional[str] = None + master_key: Optional[str] = None + passthrough_provider_keys: bool = True + + +def normalize_bifrost_url(base_url: str) -> str: + """Ensure the Bifrost LiteLLM proxy URL is well-formed.""" + url = (base_url or "").strip().rstrip("/") + if not url: + return "" + parsed = urlparse(url) + if not parsed.scheme or not parsed.netloc: + raise ValueError( + "Bifrost base_url must be an absolute URL including scheme and host " + "(e.g. http://localhost:8080/litellm)." + ) + if not url.endswith("/litellm"): + logger.warning( + "Bifrost base_url '{}' does not end with '/litellm'; appending suffix.", + url, + ) + url = f"{url}/litellm" + return url + + +def normalize_litellm_proxy_url(base_url: str) -> str: + """Ensure a LiteLLM Proxy URL is well-formed.""" + url = (base_url or "").strip().rstrip("/") + if not url: + return "" + parsed = urlparse(url) + if not parsed.scheme or not parsed.netloc: + raise ValueError( + "LiteLLM proxy base_url must be an absolute URL including scheme and host " + "(e.g. http://localhost:4000)." + ) + return url + + +def _normalize_base_url(base_url: str, gateway_type: GatewayType) -> str: + if gateway_type == "bifrost": + return normalize_bifrost_url(base_url) + return normalize_litellm_proxy_url(base_url) + + +def _platform_config() -> Dict[str, Any]: + gateway_type = (settings.LLM_GATEWAY_TYPE or "bifrost").strip().lower() + if gateway_type not in ("bifrost", "litellm_proxy"): + gateway_type = "bifrost" + return { + "enabled": bool(settings.LLM_GATEWAY_ENABLED), + "gateway_type": gateway_type, + "base_url": (settings.LLM_GATEWAY_BASE_URL or "").strip() or None, + "virtual_key": (settings.LLM_GATEWAY_VIRTUAL_KEY or "").strip() or None, + "master_key": (settings.LLM_GATEWAY_MASTER_KEY or "").strip() or None, + "passthrough_provider_keys": bool(settings.LLM_GATEWAY_PASSTHROUGH_PROVIDER_KEYS), + } + + +def _get_org_raw_settings(organization_id: UUID, db: Session) -> Dict[str, Any]: + from app.models.database import Organization + + org = db.query(Organization).filter(Organization.id == organization_id).first() + if not org: + return {} + raw = org.llm_gateway_settings + return dict(raw) if isinstance(raw, dict) else {} + + +def _decrypt_org_secret(raw: Dict[str, Any], field: str) -> Optional[str]: + encrypted = raw.get(field) + if not encrypted: + return None + try: + from app.core.encryption import decrypt_api_key + + return decrypt_api_key(encrypted) + except Exception as exc: + logger.warning("Failed to decrypt org LLM gateway {}: {}", field, exc) + return None + + +def _decrypt_org_virtual_key(raw: Dict[str, Any]) -> Optional[str]: + return _decrypt_org_secret(raw, "virtual_key") + + +def _decrypt_org_master_key(raw: Dict[str, Any]) -> Optional[str]: + return _decrypt_org_secret(raw, "master_key") + + +def _resolve_gateway_type(org: Dict[str, Any], platform: Dict[str, Any]) -> GatewayType: + org_type = org.get("gateway_type") + if org_type in ("bifrost", "litellm_proxy"): + return org_type + platform_type = platform.get("gateway_type", "bifrost") + if platform_type in ("bifrost", "litellm_proxy"): + return platform_type + return "bifrost" + + +def resolve_llm_gateway( + organization_id: UUID, + db: Session, +) -> Optional[LLMGatewayConfig]: + """Return effective LLM gateway config, or ``None`` for direct routing.""" + platform = _platform_config() + org = _get_org_raw_settings(organization_id, db) + + org_enabled = org.get("enabled") + if org_enabled is False: + return None + + if org_enabled is True: + use_gateway = True + elif platform["enabled"]: + use_gateway = True + else: + return None + + if not use_gateway: + return None + + gateway_type = _resolve_gateway_type(org, platform) + + base_url = (org.get("base_url") or platform["base_url"] or "").strip() + if not base_url: + logger.warning( + "LLM gateway ({}) enabled for org {} but no base_url is configured; " + "falling back to direct provider routing.", + gateway_type, + organization_id, + ) + return None + + try: + api_base = _normalize_base_url(base_url, gateway_type) + except ValueError as exc: + logger.warning( + "Invalid LLM gateway base_url for org {} ({}): {}; falling back to direct routing.", + organization_id, + gateway_type, + exc, + ) + return None + + virtual_key = _decrypt_org_virtual_key(org) or platform["virtual_key"] + master_key = _decrypt_org_master_key(org) or platform["master_key"] + + return LLMGatewayConfig( + gateway_type=gateway_type, + api_base=api_base, + virtual_key=virtual_key, + master_key=master_key, + passthrough_provider_keys=platform["passthrough_provider_keys"], + ) + + +# Providers whose LiteLLM handlers build native API paths (e.g. Gemini +# ``:generateContent``) when ``api_base`` is set. Gateways like Bifrost and +# LiteLLM Proxy expect OpenAI-compatible ``/v1/chat/completions`` instead. +_NATIVE_GATEWAY_MODEL_PREFIXES = ( + "gemini/", + "google/", + "vertex/", + "vertex_ai/", +) + + +def _model_uses_native_provider_path(model: Any) -> bool: + model_str = str(model or "").lower() + return any(model_str.startswith(prefix) for prefix in _NATIVE_GATEWAY_MODEL_PREFIXES) + + +def _apply_proxy_compatible_routing( + call_kwargs: Dict[str, Any], + *, + routing_model: Optional[str] = None, +) -> Dict[str, Any]: + """Route native-path providers through the gateway's chat-completions API.""" + result = dict(call_kwargs) + model_for_routing = result.get("model") or routing_model + if _model_uses_native_provider_path(model_for_routing): + result["custom_llm_provider"] = "openai" + return result + + +def _strip_provider_keys(result: Dict[str, Any]) -> None: + stored_key = result.get("api_key") + if stored_key == GATEWAY_MANAGED_KEY_SENTINEL: + result.pop("api_key", None) + elif stored_key is not None: + result.pop("api_key", None) + + +def _apply_bifrost_gateway( + call_kwargs: Dict[str, Any], + config: LLMGatewayConfig, + *, + routing_model: Optional[str] = None, +) -> Dict[str, Any]: + result = dict(call_kwargs) + result["api_base"] = config.api_base + + if config.virtual_key: + extra_headers = dict(result.get("extra_headers") or {}) + extra_headers["x-bf-vk"] = config.virtual_key + result["extra_headers"] = extra_headers + + if not config.passthrough_provider_keys: + _strip_provider_keys(result) + + if not result.get("api_key"): + result["api_key"] = config.virtual_key or LITELLM_GATEWAY_PLACEHOLDER_API_KEY + + model = result.get("model") or routing_model or "" + if model and "/" not in str(model): + logger.warning( + "Routing model '{}' through Bifrost without a provider prefix; " + "ensure the model is supported by both LiteLLM and Bifrost.", + model, + ) + + return _apply_proxy_compatible_routing(result, routing_model=routing_model) + + +def _apply_litellm_proxy_gateway( + call_kwargs: Dict[str, Any], + config: LLMGatewayConfig, + *, + routing_model: Optional[str] = None, +) -> Dict[str, Any]: + result = dict(call_kwargs) + result["api_base"] = config.api_base + + if not config.passthrough_provider_keys: + _strip_provider_keys(result) + + if not result.get("api_key"): + result["api_key"] = config.master_key or LITELLM_GATEWAY_PLACEHOLDER_API_KEY + + return _apply_proxy_compatible_routing(result, routing_model=routing_model) + + +def apply_llm_gateway( + call_kwargs: Dict[str, Any], + *, + organization_id: UUID, + db: Session, + model: Optional[str] = None, +) -> Dict[str, Any]: + """Merge gateway proxy settings into LiteLLM ``completion`` kwargs. + + Pass ``model`` when the caller supplies the model separately (e.g. GEPA + ``DefaultAdapter``) so native-path routing still applies without putting + ``model`` in the returned kwargs. + """ + config = resolve_llm_gateway(organization_id, db) + if config is None: + return call_kwargs + + if config.gateway_type == "bifrost": + return _apply_bifrost_gateway(call_kwargs, config, routing_model=model) + return _apply_litellm_proxy_gateway(call_kwargs, config, routing_model=model) + + +def litellm_completion( + *, + organization_id: UUID, + db: Session, + **kwargs: Any, +): + """Call ``litellm.completion`` with LLM gateway settings applied.""" + import litellm + + kwargs = apply_llm_gateway(kwargs, organization_id=organization_id, db=db) + return litellm.completion(**kwargs) diff --git a/app/services/ai/llm_gateway_settings.py b/app/services/ai/llm_gateway_settings.py new file mode 100644 index 00000000..42662ddc --- /dev/null +++ b/app/services/ai/llm_gateway_settings.py @@ -0,0 +1,187 @@ +""" +Per-organization LLM gateway settings. + +Backed by ``organizations.llm_gateway_settings`` JSON. +Virtual keys and LiteLLM proxy master keys are encrypted at rest. +""" + +from __future__ import annotations + +from typing import Any, Dict, Literal, Optional +from uuid import UUID + +from fastapi import HTTPException +from sqlalchemy.orm import Session + +from app.models.database import Organization +from app.services.ai.llm_gateway import ( + EffectiveRouting, + GatewayType, + _platform_config, + gateway_managed_credentials_enabled, + normalize_bifrost_url, + normalize_litellm_proxy_url, + resolve_llm_gateway, +) + +GatewayMode = Literal["inherit", "enabled", "disabled"] +GatewayTypeOverride = Literal["inherit", "bifrost", "litellm_proxy"] + + +def _mode_from_enabled(enabled: Any) -> GatewayMode: + if enabled is True: + return "enabled" + if enabled is False: + return "disabled" + return "inherit" + + +def _enabled_from_mode(mode: GatewayMode) -> Optional[bool]: + if mode == "enabled": + return True + if mode == "disabled": + return False + return None + + +def _gateway_type_from_stored(raw: Dict[str, Any]) -> GatewayTypeOverride: + stored = raw.get("gateway_type") + if stored in ("bifrost", "litellm_proxy"): + return stored + return "inherit" + + +def _effective_routing(resolved: Any) -> EffectiveRouting: + if resolved is None: + return "direct" + return resolved.gateway_type + + +def _normalize_url_for_type(base_url: str, gateway_type: GatewayType) -> str: + if gateway_type == "bifrost": + return normalize_bifrost_url(base_url) + return normalize_litellm_proxy_url(base_url) + + +def _resolve_effective_gateway_type( + org_raw: Dict[str, Any], platform: Dict[str, Any] +) -> GatewayType: + org_type = org_raw.get("gateway_type") + if org_type in ("bifrost", "litellm_proxy"): + return org_type + platform_type = platform.get("gateway_type", "bifrost") + if platform_type in ("bifrost", "litellm_proxy"): + return platform_type + return "bifrost" + + +def get_org_settings(organization_id: UUID, db: Session) -> Dict[str, Any]: + """Return stored org settings merged with effective resolved routing.""" + org = db.query(Organization).filter(Organization.id == organization_id).first() + raw = (org.llm_gateway_settings or {}) if org else {} + + mode = _mode_from_enabled(raw.get("enabled")) + stored_base_url = raw.get("base_url") + has_virtual_key = bool(raw.get("virtual_key")) + has_master_key = bool(raw.get("master_key")) + gateway_type = _gateway_type_from_stored(raw) + + resolved = resolve_llm_gateway(organization_id, db) + platform = _platform_config() + effective_type = _resolve_effective_gateway_type(raw, platform) + + return { + "mode": mode, + "gateway_type": gateway_type, + "base_url": stored_base_url, + "has_virtual_key": has_virtual_key, + "has_master_key": has_master_key, + "platform_enabled": platform["enabled"], + "platform_gateway_type": platform["gateway_type"], + "platform_base_url": platform["base_url"], + "effective_routing": _effective_routing(resolved), + "effective_gateway_type": effective_type if resolved else None, + "effective_base_url": resolved.api_base if resolved else None, + "effective_has_virtual_key": bool(resolved and resolved.virtual_key), + "effective_has_master_key": bool(resolved and resolved.master_key), + "gateway_managed_credentials": gateway_managed_credentials_enabled(), + } + + +def set_org_settings( + organization_id: UUID, + db: Session, + *, + mode: GatewayMode, + gateway_type: GatewayTypeOverride = "inherit", + base_url: Optional[str] = None, + virtual_key: Optional[str] = None, + master_key: Optional[str] = None, + clear_virtual_key: bool = False, + clear_master_key: bool = False, +) -> Dict[str, Any]: + """Validate and persist org-level LLM gateway overrides.""" + org = db.query(Organization).filter(Organization.id == organization_id).first() + if not org: + raise HTTPException(status_code=404, detail="Organization not found") + + existing = dict(org.llm_gateway_settings or {}) + platform = _platform_config() + payload: Dict[str, Any] = { + "enabled": _enabled_from_mode(mode), + } + + if gateway_type == "inherit": + payload["gateway_type"] = None + else: + payload["gateway_type"] = gateway_type + + effective_type_for_validation: GatewayType + if gateway_type in ("bifrost", "litellm_proxy"): + effective_type_for_validation = gateway_type + else: + effective_type_for_validation = _resolve_effective_gateway_type({}, platform) + + if base_url is not None: + trimmed = base_url.strip() + if trimmed: + try: + _normalize_url_for_type(trimmed, effective_type_for_validation) + except ValueError as exc: + raise HTTPException(status_code=400, detail=str(exc)) from exc + payload["base_url"] = trimmed + else: + payload["base_url"] = None + elif "base_url" in existing: + payload["base_url"] = existing.get("base_url") + + if clear_virtual_key: + payload["virtual_key"] = None + elif virtual_key is not None: + trimmed_key = virtual_key.strip() + if trimmed_key: + from app.core.encryption import encrypt_api_key + + payload["virtual_key"] = encrypt_api_key(trimmed_key) + else: + payload["virtual_key"] = None + elif "virtual_key" in existing: + payload["virtual_key"] = existing.get("virtual_key") + + if clear_master_key: + payload["master_key"] = None + elif master_key is not None: + trimmed_key = master_key.strip() + if trimmed_key: + from app.core.encryption import encrypt_api_key + + payload["master_key"] = encrypt_api_key(trimmed_key) + else: + payload["master_key"] = None + elif "master_key" in existing: + payload["master_key"] = existing.get("master_key") + + org.llm_gateway_settings = payload + db.commit() + + return get_org_settings(organization_id, db) diff --git a/app/services/ai/llm_service.py b/app/services/ai/llm_service.py index 407ee51d..38b1228f 100644 --- a/app/services/ai/llm_service.py +++ b/app/services/ai/llm_service.py @@ -20,6 +20,7 @@ from app.models.database import ModelProvider, AIProvider from app.services.credentials import resolve_ai_provider from app.services.ai.llm_generation_config import build_litellm_kwargs +from app.services.ai.llm_gateway import apply_llm_gateway, resolve_litellm_api_key # LiteLLM will silently drop params the target provider doesn't support # rather than raising an error. @@ -216,14 +217,7 @@ def generate_response( f"AI provider {llm_provider} not configured for this organization." ) - from app.core.encryption import decrypt_api_key - - try: - api_key = decrypt_api_key(ai_provider.api_key) - except Exception as e: - raise RuntimeError( - f"Failed to decrypt API key for provider {llm_provider}: {e}" - ) + api_key = resolve_litellm_api_key(organization_id, db, ai_provider) # --- call LiteLLM -------------------------------------------------- model_str = self._litellm_model_name(llm_provider, llm_model) @@ -231,9 +225,10 @@ def generate_response( call_kwargs: Dict[str, Any] = { "model": model_str, "messages": messages, - "api_key": api_key, "temperature": temperature, } + if api_key is not None: + call_kwargs["api_key"] = api_key # Gemini "thinking" families (2.5 + 3.x) ship with reasoning # enabled by default. For structured-JSON workloads (the # diariser and evaluator) chain-of-thought is wasted output @@ -265,6 +260,12 @@ def generate_response( if config: call_kwargs.update(config) + call_kwargs = apply_llm_gateway( + call_kwargs, + organization_id=organization_id, + db=db, + ) + try: response = litellm.completion(**call_kwargs) except Exception as e: diff --git a/app/services/ai/stt_clients/google.py b/app/services/ai/stt_clients/google.py index 2c1b09d3..7e335818 100644 --- a/app/services/ai/stt_clients/google.py +++ b/app/services/ai/stt_clients/google.py @@ -24,6 +24,7 @@ import logging from pathlib import Path from typing import Any, Dict, Optional +from uuid import UUID logger = logging.getLogger(__name__) @@ -85,8 +86,11 @@ def _build_transcription_prompt(language: Optional[str]) -> str: def transcribe_google( audio_file_path: str, model: str, - api_key: str, + api_key: Optional[str] = None, language: Optional[str] = None, + *, + organization_id: Optional["UUID"] = None, + db: Optional[Any] = None, ) -> Dict[str, Any]: """Transcribe an audio file via Gemini through LiteLLM. @@ -122,12 +126,22 @@ def transcribe_google( ] try: - response = litellm.completion( - model=litellm_model, - messages=messages, - api_key=api_key, - temperature=0.0, - ) + call_kwargs: Dict[str, Any] = { + "model": litellm_model, + "messages": messages, + "temperature": 0.0, + } + if api_key is not None: + call_kwargs["api_key"] = api_key + if organization_id is not None and db is not None: + from app.services.ai.llm_gateway import apply_llm_gateway + + call_kwargs = apply_llm_gateway( + call_kwargs, + organization_id=organization_id, + db=db, + ) + response = litellm.completion(**call_kwargs) except Exception as e: logger.error( f"[transcribe_google] LiteLLM call failed for {litellm_model}: {e}" diff --git a/app/services/ai/transcription_service.py b/app/services/ai/transcription_service.py index b613a9c2..c5263dbc 100644 --- a/app/services/ai/transcription_service.py +++ b/app/services/ai/transcription_service.py @@ -47,6 +47,7 @@ from pathlib import Path from app.models.database import ModelProvider, AIProvider, Integration +from app.core.encryption import decrypt_api_key from app.services.credentials import resolve_ai_provider, resolve_integration from app.services.storage.s3_service import s3_service from app.core.exceptions import StorageError @@ -89,13 +90,13 @@ def _get_api_key_for_provider( preferred in either table; otherwise the row marked ``is_default`` wins, with a back-compat fallback to the most recent active row. """ - from app.core.encryption import decrypt_api_key + from app.services.ai.llm_gateway import resolve_litellm_api_key ai_provider = self._get_ai_provider( provider, db, organization_id, credential_id=credential_id ) if ai_provider: - return decrypt_api_key(ai_provider.api_key) + return resolve_litellm_api_key(organization_id, db, ai_provider) integration = resolve_integration( provider, db, organization_id, credential_id=credential_id @@ -507,7 +508,7 @@ def transcribe_text_only( api_key = self._get_api_key_for_provider( stt_provider, db, organization_id, credential_id=credential_id ) - if not api_key: + if not api_key and stt_provider != ModelProvider.GOOGLE: logger.warning( f"[TranscriptionService] No API key found for {stt_provider} " f"(checked AIProvider and Integration tables) for org {organization_id}" @@ -531,7 +532,10 @@ def transcribe_text_only( elif stt_provider == ModelProvider.ELEVENLABS: result = transcribe_elevenlabs(audio_file_path, stt_model, api_key, language) elif stt_provider == ModelProvider.GOOGLE: - result = transcribe_google(audio_file_path, stt_model, api_key, language) + result = transcribe_google( + audio_file_path, stt_model, api_key, language, + organization_id=organization_id, db=db, + ) elif stt_provider == ModelProvider.SARVAM: result = transcribe_sarvam(audio_file_path, stt_model, api_key, language) elif stt_provider == ModelProvider.SMALLEST: @@ -570,7 +574,7 @@ def transcribe( api_key = self._get_api_key_for_provider( stt_provider, db, organization_id, credential_id=credential_id ) - if not api_key: + if not api_key and stt_provider != ModelProvider.GOOGLE: raise RuntimeError( f"No API key found for {stt_provider} (checked AIProvider and Integration tables). " f"Please configure the provider in Settings." @@ -608,7 +612,10 @@ def transcribe( # are routed through LiteLLM as multimodal completions. # Legacy ``google-speech-v2`` would also land here, but # we don't yet have a Cloud Speech client wired up. - result = transcribe_google(temp_file_path, stt_model, api_key, language) + result = transcribe_google( + temp_file_path, stt_model, api_key, language, + organization_id=organization_id, db=db, + ) elif stt_provider == ModelProvider.AZURE: raise NotImplementedError("Azure Speech Services not yet implemented") elif stt_provider == ModelProvider.AWS: diff --git a/app/services/judge_alignment/gepa_bridge.py b/app/services/judge_alignment/gepa_bridge.py index 84d8deb1..c6fe9694 100644 --- a/app/services/judge_alignment/gepa_bridge.py +++ b/app/services/judge_alignment/gepa_bridge.py @@ -25,13 +25,13 @@ from typing import Any, Dict, List, Optional, Tuple from uuid import UUID -import litellm from loguru import logger from sqlalchemy.orm import Session -from app.core.encryption import decrypt_api_key +from app.services.ai.llm_gateway import apply_llm_gateway, litellm_completion, resolve_litellm_api_key +from app.services.credentials import resolve_ai_provider + from app.models.database import ( - AIProvider, Evaluator, JudgeDataset, JudgeRun, @@ -59,19 +59,11 @@ def _resolve_litellm_model(provider: str, model: str) -> str: def _resolve_api_key( provider: str, organization_id: UUID, db: Session -) -> str: - ai_provider = ( - db.query(AIProvider) - .filter( - AIProvider.organization_id == organization_id, - AIProvider.provider == provider, - AIProvider.is_active == True, # noqa: E712 - ) - .first() - ) +) -> Optional[str]: + ai_provider = resolve_ai_provider(provider, db, organization_id) if not ai_provider: raise RuntimeError(f"No active AIProvider for {provider!r}") - return decrypt_api_key(ai_provider.api_key) + return resolve_litellm_api_key(organization_id, db, ai_provider) def start_gepa_for_dataset( @@ -247,13 +239,17 @@ def _evaluate_prompt(candidate_prompt: str) -> float: }, ] try: - resp = litellm.completion( - model=lm_identifier, - messages=messages, - api_key=api_key, - temperature=0.0, - max_tokens=300, - ) + sample_kwargs: Dict[str, Any] = { + "organization_id": run.organization_id, + "db": db, + "model": lm_identifier, + "messages": messages, + "temperature": 0.0, + "max_tokens": 300, + } + if api_key is not None: + sample_kwargs["api_key"] = api_key + resp = litellm_completion(**sample_kwargs) text = resp.choices[0].message.content if resp.choices else "" parsed = _parse_judge_response(text) except Exception as exc: @@ -300,10 +296,20 @@ def evaluator_fn(data: Dict[str, Any], response: str) -> Any: } ] + batch_kwargs: Dict[str, Any] = {} + if api_key is not None: + batch_kwargs["api_key"] = api_key + batch_kwargs = apply_llm_gateway( + batch_kwargs, + organization_id=run.organization_id, + db=db, + model=lm_identifier, + ) + adapter = DefaultAdapter( model=lm_identifier, evaluator=evaluator_fn, - litellm_batch_completion_kwargs={"api_key": api_key}, + litellm_batch_completion_kwargs=batch_kwargs, ) # Custom evaluator wrapper to inject the candidate prompt into `data` @@ -321,11 +327,15 @@ def _capturing_evaluator(data: Dict[str, Any], response: str): adapter.evaluator = _capturing_evaluator # type: ignore[attr-defined] def reflection_lm(prompt: str) -> str: - resp = litellm.completion( - model=lm_identifier, - messages=[{"role": "user", "content": prompt}], - api_key=api_key, - ) + reflection_kwargs: Dict[str, Any] = { + "organization_id": run.organization_id, + "db": db, + "model": lm_identifier, + "messages": [{"role": "user", "content": prompt}], + } + if api_key is not None: + reflection_kwargs["api_key"] = api_key + resp = litellm_completion(**reflection_kwargs) return resp.choices[0].message.content run.status = PromptOptimizationStatus.RUNNING.value diff --git a/app/services/optimization/evaluator.py b/app/services/optimization/evaluator.py index 87ca1d53..bab71385 100644 --- a/app/services/optimization/evaluator.py +++ b/app/services/optimization/evaluator.py @@ -17,6 +17,58 @@ from app.models.database import AIProvider, Metric from app.workers.tasks.helpers.score_utils import get_metric_type_value +_DEFAULT_SCORING_MODELS = { + "openai": "gpt-4o-mini", + "anthropic": "claude-3-5-haiku-latest", + "google": "gemini-2.0-flash", +} + + +def _normalize_provider_prefix(prefix: str) -> str: + normalized = prefix.lower() + if normalized == "gemini": + return "google" + return normalized + + +def _resolve_scoring_lm( + lm_identifier: str | None, + ai_providers: List[AIProvider], +) -> tuple["ModelProvider", str] | tuple[None, None]: + """Pick provider + model for GEPA scoring calls.""" + from app.models.database import ModelProvider + + if lm_identifier and "/" in lm_identifier: + prefix, model = lm_identifier.split("/", 1) + provider_key = _normalize_provider_prefix(prefix) + provider = next( + ( + p + for p in ai_providers + if p.is_active and p.provider.lower() == provider_key + ), + None, + ) + if provider: + return ModelProvider(provider.provider.lower()), model + + provider = next( + (p for p in ai_providers if p.is_active and p.provider.lower() == "openai"), + None, + ) + if provider: + return ModelProvider.OPENAI, _DEFAULT_SCORING_MODELS["openai"] + + provider = next((p for p in ai_providers if p.is_active), None) + if not provider: + return None, None + + provider_key = provider.provider.lower() + return ( + ModelProvider(provider_key), + _DEFAULT_SCORING_MODELS.get(provider_key, _DEFAULT_SCORING_MODELS["openai"]), + ) + def _get_evaluation_result_class(): """Lazy-import EvaluationResult; gepa will already be installed by the time this is called.""" @@ -32,6 +84,7 @@ def build_evaluator( ai_providers: List[AIProvider], organization_id: UUID, db, + lm_identifier: str | None = None, ) -> Callable[[Dict[str, Any], str], Any]: """ Return a function with the signature GEPA's ``Evaluator`` protocol @@ -79,13 +132,9 @@ def evaluator_fn(data: Dict[str, Any], response: str) -> Any: ) from app.services.ai.llm_service import llm_service - from app.models.database import ModelProvider - provider = next( - (p for p in ai_providers if p.provider.lower() in ("openai", "anthropic")), - ai_providers[0] if ai_providers else None, - ) - if not provider: + llm_provider, llm_model = _resolve_scoring_lm(lm_identifier, ai_providers) + if not llm_provider or not llm_model: return EvaluationResult(score=0.5, feedback="No AI provider available") try: @@ -94,8 +143,8 @@ def evaluator_fn(data: Dict[str, Any], response: str) -> Any: {"role": "system", "content": "You are an expert voice AI evaluator. Respond with JSON only."}, {"role": "user", "content": eval_prompt}, ], - llm_provider=ModelProvider(provider.provider.lower()), - llm_model="gpt-4o", + llm_provider=llm_provider, + llm_model=llm_model, organization_id=organization_id, db=db, temperature=0.3, diff --git a/app/services/optimization/gepa_service.py b/app/services/optimization/gepa_service.py index 35a344b4..fe6bf2b1 100644 --- a/app/services/optimization/gepa_service.py +++ b/app/services/optimization/gepa_service.py @@ -9,7 +9,6 @@ from typing import Any, Dict, List, Optional from uuid import UUID -import litellm from loguru import logger from app.models.database import ( @@ -23,6 +22,7 @@ from app.services.optimization.data_preparation import build_trainset from app.services.optimization.evaluator import build_evaluator from app.services.optimization.lm_resolver import resolve_api_key, resolve_lm +from app.services.ai.llm_gateway import apply_llm_gateway, litellm_completion _gepa_install_attempted = False @@ -96,31 +96,51 @@ def run_optimization( seed_prompt = agent.provider_prompt or agent.description or "" lm_identifier = resolve_lm(voice_bundle, evaluator) - api_key = resolve_api_key(lm_identifier, ai_providers) + api_key = resolve_api_key(lm_identifier, ai_providers, organization_id, db) trainset = build_trainset(training_data, metrics) if not trainset: raise ValueError("No evaluator results with transcripts available for optimization") - evaluator_fn = build_evaluator(metrics, ai_providers, organization_id, db) + evaluator_fn = build_evaluator( + metrics, + ai_providers, + organization_id, + db, + lm_identifier=lm_identifier, + ) logger.info( f"[GEPA] Starting optimization for agent '{agent.name}' " f"with {len(trainset)} training examples, LM={lm_identifier}" ) + batch_kwargs: Dict[str, Any] = {} + if api_key is not None: + batch_kwargs["api_key"] = api_key + batch_kwargs = apply_llm_gateway( + batch_kwargs, + organization_id=organization_id, + db=db, + model=lm_identifier, + ) + adapter = DefaultAdapter( model=lm_identifier, evaluator=evaluator_fn, - litellm_batch_completion_kwargs={"api_key": api_key}, + litellm_batch_completion_kwargs=batch_kwargs, ) def reflection_lm(prompt: str) -> str: - resp = litellm.completion( - model=lm_identifier, - messages=[{"role": "user", "content": prompt}], - api_key=api_key, - ) + reflection_kwargs: Dict[str, Any] = { + "organization_id": organization_id, + "db": db, + "model": lm_identifier, + "messages": [{"role": "user", "content": prompt}], + } + if api_key is not None: + reflection_kwargs["api_key"] = api_key + resp = litellm_completion(**reflection_kwargs) return resp.choices[0].message.content result = gepa_optimize( diff --git a/app/services/optimization/lm_resolver.py b/app/services/optimization/lm_resolver.py index 55e79534..18acaa8e 100644 --- a/app/services/optimization/lm_resolver.py +++ b/app/services/optimization/lm_resolver.py @@ -7,9 +7,12 @@ """ from typing import List, Optional +from uuid import UUID + +from sqlalchemy.orm import Session -from app.core.encryption import decrypt_api_key from app.models.database import AIProvider, Evaluator, VoiceBundle +from app.services.ai.llm_gateway import resolve_litellm_api_key def resolve_lm( @@ -27,17 +30,22 @@ def resolve_lm( return "openai/gpt-4o" -def resolve_api_key(lm_identifier: str, ai_providers: List[AIProvider]) -> str: +def resolve_api_key( + lm_identifier: str, + ai_providers: List[AIProvider], + organization_id: UUID, + db: Session, +) -> Optional[str]: """ Given ``"openai/gpt-5.4"`` and the org's provider list, decrypt and - return the matching API key. + return the matching API key, or None when Bifrost manages provider keys. """ provider_prefix = lm_identifier.split("/")[0].lower() for p in ai_providers: if not p.is_active or not p.api_key: continue if p.provider.lower() == provider_prefix: - return decrypt_api_key(p.api_key) + return resolve_litellm_api_key(organization_id, db, p) raise RuntimeError( f"No active AI provider matching '{provider_prefix}' found. " "Add one in Settings > AI Providers." diff --git a/app/workers/tasks/helpers/llm_diarisation.py b/app/workers/tasks/helpers/llm_diarisation.py index ef9c3673..20040cde 100644 --- a/app/workers/tasks/helpers/llm_diarisation.py +++ b/app/workers/tasks/helpers/llm_diarisation.py @@ -750,6 +750,10 @@ def _upload_audio_to_gemini_files( / retry surface as the eventual ``litellm.completion`` call, and the returned ``file.id`` is already in the exact shape LiteLLM's Gemini transformation expects to see in a ``file`` content part. + * Bifrost gateway routing is **not** applied here — the Gemini Files + API upload goes directly to Google. Only the subsequent + ``litellm.completion`` diarisation call is proxied when Bifrost + is enabled (via ``LLMService``). * The ``(filename, bytes, content_type)`` tuple form is mandatory — passing raw bytes makes LiteLLM default the Content-Type header to ``application/octet-stream``, which Gemini then rejects as diff --git a/config.docker.yml b/config.docker.yml index a1a7e964..b48d5c77 100644 --- a/config.docker.yml +++ b/config.docker.yml @@ -115,6 +115,15 @@ judge_alignment: enabled: true csv_max_rows: 5000 +# LLM gateway (optional). Per-org overrides in Integrations UI. +# llm_gateway: +# enabled: false +# type: bifrost # bifrost | litellm_proxy +# base_url: "http://localhost:8080/litellm" +# virtual_key: null +# master_key: null +# passthrough_provider_keys: true + # Enterprise License (JWT signed with RS256). Unlocks gated features like # oidc_sso, mfa_enforce, audit_export, voice_playground, gepa_optimization. # license: diff --git a/config.yml.example b/config.yml.example index 32a1dd02..a61a79a6 100644 --- a/config.yml.example +++ b/config.yml.example @@ -166,6 +166,17 @@ judge_alignment: enabled: true # operator-level kill switch csv_max_rows: 5000 # safety limit for CSV uploads +# LLM gateway (optional). Routes batch/eval LiteLLM calls through a proxy +# instead of directly to providers. Supports Bifrost and self-hosted LiteLLM Proxy. +# Per-org overrides (type, URL, keys, opt-out) are configured in Integrations UI. +# llm_gateway: +# enabled: false +# type: bifrost # bifrost | litellm_proxy +# base_url: "http://localhost:8080/litellm" +# virtual_key: null # Bifrost only (x-bf-vk header) +# master_key: null # LiteLLM Proxy only (LITELLM_MASTER_KEY) +# passthrough_provider_keys: true # false = keys in gateway; AI integrations can omit api_key + # Enterprise License (JWT signed with RS256). Unlocks gated features like # oidc_sso, mfa_enforce, audit_export, voice_playground, gepa_optimization, # call_imports. diff --git a/docs-fumadocs/content/docs/getting-started/integrations.mdx b/docs-fumadocs/content/docs/getting-started/integrations.mdx index 9c52bbd3..a3df025f 100644 --- a/docs-fumadocs/content/docs/getting-started/integrations.mdx +++ b/docs-fumadocs/content/docs/getting-started/integrations.mdx @@ -88,6 +88,103 @@ Telephony integrations connect your phone-network provider to EfficientAI. These integrations handle telephony-specific setup like credentials, number sync, and routing-related configuration used by call workflows. +## LLM Gateway (batch/eval routing) + +Batch and evaluation LLM workloads (evaluators, GEPA, call-import LLM steps) can be routed through an optional **LLM Gateway** instead of calling providers directly. Real-time voice agents are **not** routed through the gateway. + +Configure a **platform default** in `config.yml` and optional **per-organization overrides** from **Configurations → Integrations → LLM Gateway**. + +Supported gateway types: + +- **Bifrost** — point at your Bifrost LiteLLM proxy (URL must include the `/litellm` path, e.g. `http://localhost:8080/litellm`). +- **LiteLLM Proxy** — point at a self-hosted LiteLLM Proxy instance (e.g. `http://localhost:4000`). + +### Platform default (`config.yml`) + +```yaml +llm_gateway: + enabled: true + type: bifrost # bifrost | litellm_proxy + base_url: "http://localhost:8080/litellm" + virtual_key: null # Bifrost only (sent as x-bf-vk) + master_key: null # LiteLLM Proxy only + passthrough_provider_keys: false # false = provider keys live in the gateway +``` + +When `passthrough_provider_keys` is `false`, AI provider integrations can omit API keys and use a **Gateway managed** placeholder instead. That setting is **platform-wide only** — it cannot be overridden per organization. + +### Organization settings (Integrations UI) + +Each organization can choose how it uses the gateway. The modal shows an **Effective routing** badge after save — that is the merged result actually used for LLM calls (`Bifrost`, `LiteLLM Proxy`, or `Direct`). + +#### Organization mode + +| Mode | Behavior | +|------|----------| +| **Inherit platform default** | Follows the platform `llm_gateway.enabled` flag. When the platform has the gateway enabled, this org uses the gateway. Choose **Disabled** to opt out, or override type, URL, or keys below. | +| **Enabled (use gateway)** | Forces gateway routing for this org even if the platform default is disabled. You must have a resolvable base URL (org override or platform default). | +| **Disabled (direct to providers)** | Opts this org out entirely. LLM calls go straight to provider APIs; gateway type, URL, and key overrides are ignored. | + +#### Partial overrides (while inheriting) + +When organization mode is **Inherit** or **Enabled**, you can override individual fields without replacing the whole configuration. Anything left blank or set to **Inherit platform** falls back to the platform default. + +| Field | Inherit / blank | Org override | +|-------|-----------------|--------------| +| **Gateway type** | Platform `type` | Org value (`Bifrost` or `LiteLLM Proxy`) | +| **Base URL** | Platform `base_url` | Org URL | +| **Virtual key** (Bifrost) | Platform `virtual_key` | Org-stored key (org wins if set) | +| **Master key** (LiteLLM Proxy) | Platform `master_key` | Org-stored key (org wins if set) | +| **Passthrough provider keys** | Platform setting only | Not configurable per org | + +**Important:** Inherit mode does **not** mean “ignore my overrides.” You can inherit platform enablement while still overriding gateway type, base URL, or keys. Those overrides are applied on top of the inherited on/off decision. + +If organization mode is **Inherit** but the platform has the gateway **disabled**, the org routes **direct** regardless of any stored type or URL overrides (overrides remain saved but inactive until the platform enables the gateway or you switch the org to **Enabled**). + +### How settings merge + +Resolution is field-by-field: + +1. **On/off** — org **Disabled** → direct. Org **Enabled** → gateway on. Org **Inherit** → follow platform `enabled`. +2. **Gateway type** — org override if set, else platform `type`. +3. **Base URL** — org override if set, else platform `base_url`. If no URL resolves, the org falls back to direct routing. +4. **Keys** — org virtual key / master key if stored, else platform keys. + +### Examples + +Platform default: + +```yaml +llm_gateway: + enabled: true + type: bifrost + base_url: "http://localhost:8080/litellm" + passthrough_provider_keys: false +``` + +| Org mode | Gateway type | Base URL override | Effective result | +|----------|--------------|-------------------|------------------| +| Inherit | Inherit | *(blank)* | Bifrost @ platform URL, platform keys | +| Inherit | Inherit | `http://customer:9090/litellm` | Bifrost @ org URL, platform keys | +| Inherit | LiteLLM Proxy | `http://org-proxy:4000` | LiteLLM Proxy @ org URL | +| Inherit | LiteLLM Proxy | *(blank)* | LiteLLM Proxy type + **platform URL** — override the URL too if the platform URL is Bifrost-specific | +| Enabled | Bifrost | `http://customer:9090/litellm` | Gateway forced on even if platform is disabled | +| Disabled | *(any)* | *(any)* | Direct to providers | + +### Practical guidance + +- **Different Bifrost instance for one customer?** Keep mode and gateway type on inherit; set **Base URL override** only. +- **Different gateway product for one org?** Override **both** gateway type and base URL so they match. +- **Different auth for one org?** Set an org **virtual key** (Bifrost) or **master key** (LiteLLM Proxy). Org keys take precedence over platform keys. +- **Confirm behavior** — after saving, check the **Effective routing** badge and resolved base URL shown in the LLM Gateway modal. + +### Scope and limitations + +- Routes batch/eval LiteLLM calls only (evaluators, GEPA, call-import LLM steps). +- Does not route real-time voice agent sessions. +- Gemini and other native-path providers are automatically sent through the gateway’s OpenAI-compatible chat-completions API (not provider-native URLs like `:generateContent`). +- File uploads used by some Gemini diarisation flows (`litellm.create_file`) are not proxied through the gateway. + ## What integrations enable in practice Once linked, an agent can: diff --git a/docs-fumadocs/content/docs/getting-started/llm-gateway.mdx b/docs-fumadocs/content/docs/getting-started/llm-gateway.mdx new file mode 100644 index 00000000..60a77c4b --- /dev/null +++ b/docs-fumadocs/content/docs/getting-started/llm-gateway.mdx @@ -0,0 +1,112 @@ +--- +id: llm-gateway +title: LLM Gateway +sidebar_position: 5 +--- + +# LLM Gateway + +EfficientAI routes batch and evaluation LLM workloads through [LiteLLM](https://docs.litellm.ai/). You can optionally point those calls at an enterprise or self-hosted gateway instead of calling provider APIs directly. + +Supported gateway types: + +| Type | Description | +|------|-------------| +| **Direct** | Default — LiteLLM calls providers using org AI Provider credentials | +| **Bifrost** | [Bifrost](https://docs.getbifrost.ai/) `/litellm` proxy with optional `x-bf-vk` virtual key | +| **LiteLLM Proxy** | Self-hosted [LiteLLM Proxy](https://docs.litellm.ai/docs/proxy/quick_start) with optional master key | + +## What is routed through the gateway + +When enabled, these workloads use the configured gateway: + +- Evaluators and metrics +- Call import transcription, diarisation, and evaluation LLM steps +- GEPA prompt optimization +- Gemini batch STT via LiteLLM + +**Not affected:** real-time voice agents (they use the EfficientAI SDK directly, not LiteLLM). + +## Platform configuration + +Operators enable a gateway globally in `config.yml`: + +```yaml +llm_gateway: + enabled: true + type: bifrost # bifrost | litellm_proxy + base_url: "http://localhost:8080/litellm" + virtual_key: null # Bifrost only (x-bf-vk header) + master_key: null # LiteLLM Proxy only (LITELLM_MASTER_KEY) + passthrough_provider_keys: false +``` + +| Setting | Description | +|---------|-------------| +| `enabled` | Platform-wide toggle for orgs that inherit the default | +| `type` | Gateway implementation: `bifrost` or `litellm_proxy` | +| `base_url` | Gateway endpoint (Bifrost: must end with `/litellm`; LiteLLM Proxy: e.g. `http://localhost:4000`) | +| `virtual_key` | Optional Bifrost `x-bf-vk` header | +| `master_key` | Optional LiteLLM Proxy master key | +| `passthrough_provider_keys` | When `true`, org AI provider API keys are still sent through the gateway | + +When routing through a gateway, Gemini models automatically use OpenAI-compatible proxy endpoints (`/v1/chat/completions`) instead of native Gemini `:generateContent` URLs. This is required for Bifrost and LiteLLM Proxy compatibility. + +### Gateway-managed credentials + +When `passthrough_provider_keys` is `false`, provider secrets live in the gateway. You can add AI provider integrations in EfficientAI **without an API key** — the integration row marks which provider to use, and the gateway supplies the secret at call time. Integrations created this way show a **Gateway managed** badge in the UI. + +Requirements: + +1. Platform `passthrough_provider_keys: false` +2. LLM gateway active for the organization (platform and/or org Integrations settings) +3. Real provider keys configured in the gateway + +EfficientAI sends a harmless placeholder `api_key` to satisfy LiteLLM's OpenAI client constructor; the gateway ignores it when provider keys are configured server-side. + +Voice agents and non-LiteLLM paths still require direct provider credentials. + +## Organization overrides (Integrations UI) + +On **Configurations → Integrations**, open **LLM Gateway** to configure per-organization settings: + +- **Inherit platform default** — follow `config.yml` +- **Enabled** — force routing through the gateway (optionally with org-specific type, URL, and keys) +- **Disabled** — opt out and call providers directly, even when the platform gateway is enabled + +Org virtual keys and LiteLLM Proxy master keys are encrypted at rest like AI provider credentials. + +## Provider compatibility + +Only providers supported by **both** LiteLLM and your gateway deployment work through the proxy. EfficientAI's LiteLLM integration covers OpenAI, Anthropic, Google/Gemini, Azure, AWS Bedrock, DeepSeek, Groq, xAI, and Fireworks. + +## Known limitations + +- **Gemini Files API uploads** (`litellm.create_file` used during LLM diarisation for large audio) go directly to Google. Only the subsequent `litellm.completion` diarisation call is proxied. +- **Voice agents** are not routed through the LLM gateway. + +## Example call shapes + +**Bifrost:** + +```python +litellm.completion( + model="openai/gpt-4o-mini", + messages=[{"role": "user", "content": "Hello!"}], + api_base="http://localhost:8080/litellm", + extra_headers={"x-bf-vk": "your-virtual-key"}, +) +``` + +**LiteLLM Proxy:** + +```python +litellm.completion( + model="openai/gpt-4o-mini", + messages=[{"role": "user", "content": "Hello!"}], + api_base="http://localhost:4000", + api_key="your-master-key", +) +``` + +EfficientAI injects `api_base`, headers, and keys automatically when the gateway is active. diff --git a/docs-fumadocs/content/docs/reference/configuration.mdx b/docs-fumadocs/content/docs/reference/configuration.mdx index df018df0..15e2778a 100644 --- a/docs-fumadocs/content/docs/reference/configuration.mdx +++ b/docs-fumadocs/content/docs/reference/configuration.mdx @@ -127,6 +127,16 @@ auth: # voice_playground, gepa_optimization, scim_provisioning, saml_sso. # license: # key: "eyJhbGciOi..." + +# LLM gateway (optional). Routes batch/eval LLM calls through a proxy. +# Per-org overrides in Integrations UI. +# llm_gateway: +# enabled: false +# type: bifrost # bifrost | litellm_proxy +# base_url: "http://localhost:8080/litellm" +# virtual_key: null +# master_key: null +# passthrough_provider_keys: true ``` :::tip diff --git a/frontend/public/geminiai.png b/frontend/public/geminiai.png index d8b9f5c3..cd78455b 100644 Binary files a/frontend/public/geminiai.png and b/frontend/public/geminiai.png differ diff --git a/frontend/src/lib/api.ts b/frontend/src/lib/api.ts index 83086ab6..d0cc42ab 100644 --- a/frontend/src/lib/api.ts +++ b/frontend/src/lib/api.ts @@ -169,6 +169,37 @@ export interface OrganizationSummary { name: string } +export type LLMGatewayMode = 'inherit' | 'enabled' | 'disabled' +export type LLMGatewayType = 'inherit' | 'bifrost' | 'litellm_proxy' +export type LLMGatewayRouting = 'direct' | 'bifrost' | 'litellm_proxy' + +export interface LLMGatewaySettings { + mode: LLMGatewayMode + gateway_type: LLMGatewayType + base_url?: string | null + has_virtual_key: boolean + has_master_key: boolean + platform_enabled: boolean + platform_gateway_type: 'bifrost' | 'litellm_proxy' + platform_base_url?: string | null + effective_routing: LLMGatewayRouting + effective_gateway_type?: 'bifrost' | 'litellm_proxy' | null + effective_base_url?: string | null + effective_has_virtual_key: boolean + effective_has_master_key: boolean + gateway_managed_credentials: boolean +} + +export interface LLMGatewaySettingsUpdate { + mode: LLMGatewayMode + gateway_type?: LLMGatewayType + base_url?: string | null + virtual_key?: string | null + master_key?: string | null + clear_virtual_key?: boolean + clear_master_key?: boolean +} + export interface TelephonyIntegrationResponse { id: string organization_id: string @@ -1003,6 +1034,16 @@ class ApiClient { return response.data } + async getLLMGatewaySettings(): Promise { + const response = await this.client.get('/api/v1/organizations/llm-gateway') + return response.data + } + + async updateLLMGatewaySettings(data: LLMGatewaySettingsUpdate): Promise { + const response = await this.client.put('/api/v1/organizations/llm-gateway', data) + return response.data + } + // IAM endpoints async getOrganization(): Promise { const response = await this.client.get('/api/v1/iam/organization') diff --git a/frontend/src/pages/configurations/Integrations.tsx b/frontend/src/pages/configurations/Integrations.tsx index bec11b6a..5a238fcb 100644 --- a/frontend/src/pages/configurations/Integrations.tsx +++ b/frontend/src/pages/configurations/Integrations.tsx @@ -3,8 +3,13 @@ import { apiClient } from '../../lib/api' import type { TelephonyIntegrationResponse } from '../../lib/api' import { useState, useEffect, useRef, type ReactNode } from 'react' import { createPortal } from 'react-dom' -import { Plus, Trash2, X, AlertCircle, Plug, Edit, Brain, ChevronDown, Phone, Star } from 'lucide-react' +import { Plus, Trash2, X, AlertCircle, Plug, Edit, Brain, ChevronDown, Phone, Star, Network } from 'lucide-react' import { IntegrationCreate, IntegrationPlatform, Integration, AIProvider, AIProviderCreate, ModelProvider, TelephonyProvider } from '../../types/api' +import type { + LLMGatewayMode, + LLMGatewaySettings, + LLMGatewayType, +} from '../../lib/api' import Button from '../../components/Button' import { useToast } from '../../hooks/useToast' import { @@ -69,6 +74,42 @@ export default function Integrations() { const [telephonySipDomain, setTelephonySipDomain] = useState('') const [telephonyProviderFilter, setTelephonyProviderFilter] = useState(TelephonyProvider.PLIVO) + const [llmGatewayMode, setLlmGatewayMode] = useState('inherit') + const [llmGatewayType, setLlmGatewayType] = useState('inherit') + const [llmGatewayBaseUrl, setLlmGatewayBaseUrl] = useState('') + const [llmGatewayVirtualKey, setLlmGatewayVirtualKey] = useState('') + const [llmGatewayMasterKey, setLlmGatewayMasterKey] = useState('') + const [clearLlmGatewayVirtualKey, setClearLlmGatewayVirtualKey] = useState(false) + const [clearLlmGatewayMasterKey, setClearLlmGatewayMasterKey] = useState(false) + const [showLlmGatewayModal, setShowLlmGatewayModal] = useState(false) + + const syncLlmGatewayFormFromSettings = () => { + if (!llmGatewaySettings) return + setLlmGatewayMode(llmGatewaySettings.mode) + setLlmGatewayType(llmGatewaySettings.gateway_type) + setLlmGatewayBaseUrl(llmGatewaySettings.base_url || '') + setLlmGatewayVirtualKey('') + setLlmGatewayMasterKey('') + setClearLlmGatewayVirtualKey(false) + setClearLlmGatewayMasterKey(false) + } + + const openLlmGatewayModal = () => { + syncLlmGatewayFormFromSettings() + setShowLlmGatewayModal(true) + } + + const closeLlmGatewayModal = () => { + setShowLlmGatewayModal(false) + syncLlmGatewayFormFromSettings() + } + + const llmGatewayRoutingLabel = (routing: LLMGatewaySettings['effective_routing']) => { + if (routing === 'bifrost') return 'Bifrost' + if (routing === 'litellm_proxy') return 'LiteLLM Proxy' + return 'Direct' + } + const renderModal = (content: ReactNode) => { if (typeof document === 'undefined') return null return createPortal(content, document.body) @@ -90,6 +131,51 @@ export default function Integrations() { retry: false, }) + const { data: llmGatewaySettings } = useQuery({ + queryKey: ['llm-gateway-settings'], + queryFn: () => apiClient.getLLMGatewaySettings(), + }) + + const effectiveLlmGatewayType = + llmGatewayType !== 'inherit' + ? llmGatewayType + : llmGatewaySettings?.platform_gateway_type || 'bifrost' + + const showLlmGatewayConfigOptions = llmGatewayMode !== 'disabled' + + useEffect(() => { + if (!llmGatewaySettings || showLlmGatewayModal) return + syncLlmGatewayFormFromSettings() + }, [llmGatewaySettings, showLlmGatewayModal]) + + const updateLlmGatewayMutation = useMutation({ + mutationFn: () => + apiClient.updateLLMGatewaySettings({ + mode: llmGatewayMode, + gateway_type: llmGatewayType, + base_url: llmGatewayBaseUrl.trim() || null, + virtual_key: llmGatewayVirtualKey.trim() || undefined, + master_key: llmGatewayMasterKey.trim() || undefined, + clear_virtual_key: clearLlmGatewayVirtualKey, + clear_master_key: clearLlmGatewayMasterKey, + }), + onSuccess: () => { + queryClient.invalidateQueries({ queryKey: ['llm-gateway-settings'] }) + showToast('LLM gateway settings saved', 'success') + setLlmGatewayVirtualKey('') + setLlmGatewayMasterKey('') + setClearLlmGatewayVirtualKey(false) + setClearLlmGatewayMasterKey(false) + setShowLlmGatewayModal(false) + }, + onError: (error: any) => { + showToast( + `Failed to save LLM gateway settings: ${error.response?.data?.detail || error.message}`, + 'error', + ) + }, + }) + // Pick the visible config: default-row preferred, then any row, scoped to // the dropdown filter at the top of the section. const telephonyConfigsForFilter = allTelephonyConfigs.filter( @@ -277,8 +363,19 @@ export default function Integrations() { if (Object.keys(updateData).length === 0) { resetForm(); return } updateAIProviderMutation.mutate({ id: selectedAIProvider.id, data: updateData }) } else { - if (!selectedProvider || !apiKey.trim()) { showToast('Please select a provider and enter an API key', 'error'); return } - createAIProviderMutation.mutate({ provider: selectedProvider, api_key: apiKey, name: name || null }) + if (!selectedProvider) { + showToast('Please select a provider', 'error') + return + } + if (!aiIntegrationGatewayManaged && !apiKey.trim()) { + showToast('Please enter an API key', 'error') + return + } + createAIProviderMutation.mutate({ + provider: selectedProvider, + api_key: apiKey.trim() || undefined, + name: name || null, + }) } } else if (integrationType === 'telephony_provider') { saveTelephonyConfigMutation.mutate() @@ -365,6 +462,10 @@ export default function Integrations() { aiIntegrationProviders.length > 0 || hasTelephony + const aiIntegrationGatewayManaged = Boolean( + llmGatewaySettings?.gateway_managed_credentials, + ) + const getPlatformInfo = (platformId: IntegrationPlatform) => { return platforms.find(p => p.id === platformId) } @@ -378,6 +479,24 @@ export default function Integrations() {

Connect with voice AI platforms, AI providers, and telephony providers

+
@@ -506,6 +625,11 @@ export default function Integrations() { Default )} + {provider.gateway_managed && ( + + Gateway managed + + )} {!provider.is_active && Inactive} @@ -623,6 +747,176 @@ export default function Integrations() { )} + {showLlmGatewayModal && renderModal( +
+
+
+
+ +

LLM Gateway

+
+ +
+
+

+ Route batch and evaluation LLM calls through Bifrost or a self-hosted LiteLLM Proxy. Real-time voice agents are unaffected. +

+ +
+ + Effective routing: {llmGatewayRoutingLabel(llmGatewaySettings?.effective_routing || 'direct')} + + {llmGatewaySettings?.effective_base_url && ( + + {llmGatewaySettings.effective_base_url} + + )} +
+ +
+ + +
+ + {showLlmGatewayConfigOptions ? ( + <> +
+ + +
+ +
+ + setLlmGatewayBaseUrl(e.target.value)} + placeholder={ + llmGatewaySettings?.platform_base_url || + (effectiveLlmGatewayType === 'litellm_proxy' + ? 'http://localhost:4000' + : 'http://localhost:8080/litellm') + } + className="w-full rounded-md border border-gray-300 px-3 py-2 text-sm" + /> +

+ {effectiveLlmGatewayType === 'litellm_proxy' + ? 'LiteLLM Proxy base URL (e.g. http://localhost:4000). Leave blank to inherit the platform URL.' + : 'Bifrost URL must include the /litellm path. Leave blank to inherit the platform URL.'} +

+
+ + {effectiveLlmGatewayType === 'bifrost' && ( +
+ + setLlmGatewayVirtualKey(e.target.value)} + placeholder={ + llmGatewaySettings?.has_virtual_key + ? '•••••••• (stored — enter new value to replace)' + : 'Optional Bifrost virtual key' + } + className="w-full rounded-md border border-gray-300 px-3 py-2 text-sm" + /> + {llmGatewaySettings?.has_virtual_key && ( + + )} +
+ )} + + {effectiveLlmGatewayType === 'litellm_proxy' && ( +
+ + setLlmGatewayMasterKey(e.target.value)} + placeholder={ + llmGatewaySettings?.has_master_key + ? '•••••••• (stored — enter new value to replace)' + : 'Optional LiteLLM Proxy master key' + } + className="w-full rounded-md border border-gray-300 px-3 py-2 text-sm" + /> + {llmGatewaySettings?.has_master_key && ( + + )} +
+ )} + + {llmGatewaySettings?.gateway_managed_credentials && ( +

+ Gateway-managed credentials are enabled. AI provider integrations can omit API keys when the gateway is active. +

+ )} + + ) : ( +

+ This organization will call LLM providers directly. Enable the gateway or inherit the platform default to configure gateway type, URL, and keys. +

+ )} + +
+ + +
+
+
+
+ )} + {showModal && renderModal(
@@ -761,11 +1055,34 @@ export default function Integrations() {
- setApiKey(e.target.value)} className="w-full px-3 py-2 border border-gray-300 rounded-lg focus:ring-2 focus:ring-primary-500" - placeholder={isEditMode ? "Enter new API key (optional)" : "Enter API key"} /> -

Your API key will be encrypted and stored securely

+ setApiKey(e.target.value)} + className="w-full px-3 py-2 border border-gray-300 rounded-lg focus:ring-2 focus:ring-primary-500" + placeholder={ + isEditMode + ? 'Enter new API key (optional)' + : aiIntegrationGatewayManaged + ? 'Leave blank to use gateway-managed credentials' + : 'Enter API key' + } + /> +

+ {aiIntegrationGatewayManaged + ? 'When blank, provider secrets are resolved by your LLM gateway. You can still override with a local key if needed.' + : 'Your API key will be encrypted and stored securely'} +

)} diff --git a/frontend/src/types/api.ts b/frontend/src/types/api.ts index 6aa38c77..2d6d8ecf 100644 --- a/frontend/src/types/api.ts +++ b/frontend/src/types/api.ts @@ -288,6 +288,8 @@ export interface AIProvider { is_active: boolean /** True if this row is the default credential for (org, provider). */ is_default?: boolean + /** True when provider secrets are resolved by the Bifrost gateway. */ + gateway_managed?: boolean created_at: string updated_at: string last_tested_at?: string | null @@ -295,7 +297,7 @@ export interface AIProvider { export interface AIProviderCreate { provider: ModelProvider - api_key: string + api_key?: string | null name?: string | null /** Mark the new credential as the default for (org, provider). */ is_default?: boolean diff --git a/tests/conftest.py b/tests/conftest.py index bf507bc0..895dde42 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -347,6 +347,7 @@ class _TaskResult: alerts, audio, auth, + llm_gateway, call_import_evaluations, call_import_schemas, call_import_tags, @@ -394,6 +395,7 @@ class _TaskResult: app.include_router(audio.router, prefix="/api/v1") app.include_router(integrations.router, prefix="/api/v1") app.include_router(aiproviders.router, prefix="/api/v1") + app.include_router(llm_gateway.router, prefix="/api/v1") app.include_router(metrics.router, prefix="/api/v1") app.include_router(evaluator_results.router, prefix="/api/v1") app.include_router(voicebundles.router, prefix="/api/v1") diff --git a/tests/test_api/test_gateway_managed_credentials.py b/tests/test_api/test_gateway_managed_credentials.py new file mode 100644 index 00000000..e7f17cb4 --- /dev/null +++ b/tests/test_api/test_gateway_managed_credentials.py @@ -0,0 +1,67 @@ +"""API tests for gateway-managed AI provider credentials.""" + +import pytest + +from app.config import settings +from app.models.database import AIProvider +from app.services.ai.llm_gateway import is_gateway_managed_stored_key + + +def _set_platform_gateway_passthrough(passthrough: bool): + settings.LLM_GATEWAY_PASSTHROUGH_PROVIDER_KEYS = passthrough + + +@pytest.fixture(autouse=True) +def _reset_gateway_settings(): + original = ( + settings.LLM_GATEWAY_ENABLED, + settings.LLM_GATEWAY_PASSTHROUGH_PROVIDER_KEYS, + ) + yield + ( + settings.LLM_GATEWAY_ENABLED, + settings.LLM_GATEWAY_PASSTHROUGH_PROVIDER_KEYS, + ) = original + + +def test_create_aiprovider_without_key_when_gateway_managed( + authenticated_client, db_session, org_id +): + _set_platform_gateway_passthrough(False) + + response = authenticated_client.post( + "/api/v1/aiproviders", + json={"provider": "openai", "name": "OpenAI via gateway"}, + ) + assert response.status_code == 201 + data = response.json() + assert data["provider"] == "openai" + assert data["gateway_managed"] is True + + row = ( + db_session.query(AIProvider) + .filter(AIProvider.organization_id == org_id) + .first() + ) + assert row is not None + assert is_gateway_managed_stored_key(row.api_key) is True + + +def test_create_aiprovider_without_key_rejected_when_passthrough_enabled( + authenticated_client, +): + _set_platform_gateway_passthrough(True) + + response = authenticated_client.post( + "/api/v1/aiproviders", + json={"provider": "openai"}, + ) + assert response.status_code == 400 + + +def test_llm_gateway_settings_expose_gateway_managed_flag(authenticated_client): + _set_platform_gateway_passthrough(False) + + response = authenticated_client.get("/api/v1/organizations/llm-gateway") + assert response.status_code == 200 + assert response.json()["gateway_managed_credentials"] is True diff --git a/tests/test_api/test_llm_gateway_settings.py b/tests/test_api/test_llm_gateway_settings.py new file mode 100644 index 00000000..947e66c6 --- /dev/null +++ b/tests/test_api/test_llm_gateway_settings.py @@ -0,0 +1,120 @@ +"""API tests for organization LLM gateway settings.""" + +import pytest + +from app.config import settings +from app.models.database import Organization + + +def _set_platform_gateway(**kwargs): + settings.LLM_GATEWAY_ENABLED = kwargs.get("enabled", False) + settings.LLM_GATEWAY_TYPE = kwargs.get("gateway_type", "bifrost") + settings.LLM_GATEWAY_BASE_URL = kwargs.get("base_url") + settings.LLM_GATEWAY_VIRTUAL_KEY = kwargs.get("virtual_key") + settings.LLM_GATEWAY_MASTER_KEY = kwargs.get("master_key") + settings.LLM_GATEWAY_PASSTHROUGH_PROVIDER_KEYS = kwargs.get("passthrough", True) + + +@pytest.fixture(autouse=True) +def _reset_gateway_settings(): + original = ( + settings.LLM_GATEWAY_ENABLED, + settings.LLM_GATEWAY_TYPE, + settings.LLM_GATEWAY_BASE_URL, + settings.LLM_GATEWAY_VIRTUAL_KEY, + settings.LLM_GATEWAY_MASTER_KEY, + settings.LLM_GATEWAY_PASSTHROUGH_PROVIDER_KEYS, + ) + yield + ( + settings.LLM_GATEWAY_ENABLED, + settings.LLM_GATEWAY_TYPE, + settings.LLM_GATEWAY_BASE_URL, + settings.LLM_GATEWAY_VIRTUAL_KEY, + settings.LLM_GATEWAY_MASTER_KEY, + settings.LLM_GATEWAY_PASSTHROUGH_PROVIDER_KEYS, + ) = original + + +def test_get_llm_gateway_settings_defaults(authenticated_client): + _set_platform_gateway( + enabled=True, + base_url="http://platform-bifrost/litellm", + ) + + response = authenticated_client.get("/api/v1/organizations/llm-gateway") + assert response.status_code == 200 + data = response.json() + assert data["mode"] == "inherit" + assert data["gateway_type"] == "inherit" + assert data["platform_enabled"] is True + assert data["platform_gateway_type"] == "bifrost" + assert data["platform_base_url"] == "http://platform-bifrost/litellm" + assert data["effective_routing"] == "bifrost" + + +def test_update_llm_gateway_settings_persists_bifrost_override( + authenticated_client, db_session, org_id +): + _set_platform_gateway(enabled=False) + + response = authenticated_client.put( + "/api/v1/organizations/llm-gateway", + json={ + "mode": "enabled", + "gateway_type": "bifrost", + "base_url": "http://org-bifrost:8080/litellm", + "virtual_key": "org-virtual-key", + }, + ) + assert response.status_code == 200 + data = response.json() + assert data["mode"] == "enabled" + assert data["gateway_type"] == "bifrost" + assert data["base_url"] == "http://org-bifrost:8080/litellm" + assert data["has_virtual_key"] is True + assert data["effective_routing"] == "bifrost" + assert data["effective_base_url"] == "http://org-bifrost:8080/litellm" + + org = db_session.query(Organization).filter(Organization.id == org_id).first() + assert org.llm_gateway_settings["enabled"] is True + assert org.llm_gateway_settings["gateway_type"] == "bifrost" + assert org.llm_gateway_settings["virtual_key"] + + +def test_update_llm_gateway_settings_persists_litellm_proxy_override( + authenticated_client, db_session, org_id +): + _set_platform_gateway(enabled=False) + + response = authenticated_client.put( + "/api/v1/organizations/llm-gateway", + json={ + "mode": "enabled", + "gateway_type": "litellm_proxy", + "base_url": "http://org-proxy:4000", + "master_key": "org-master-key", + }, + ) + assert response.status_code == 200 + data = response.json() + assert data["gateway_type"] == "litellm_proxy" + assert data["has_master_key"] is True + assert data["effective_routing"] == "litellm_proxy" + assert data["effective_base_url"] == "http://org-proxy:4000" + + org = db_session.query(Organization).filter(Organization.id == org_id).first() + assert org.llm_gateway_settings["gateway_type"] == "litellm_proxy" + assert org.llm_gateway_settings["master_key"] + + +def test_update_llm_gateway_settings_rejects_invalid_url(authenticated_client): + response = authenticated_client.put( + "/api/v1/organizations/llm-gateway", + json={ + "mode": "enabled", + "gateway_type": "bifrost", + "base_url": "not-a-valid-url", + }, + ) + assert response.status_code == 400 diff --git a/tests/test_services/test_ai/test_gateway_managed_credentials.py b/tests/test_services/test_ai/test_gateway_managed_credentials.py new file mode 100644 index 00000000..b5f68323 --- /dev/null +++ b/tests/test_services/test_ai/test_gateway_managed_credentials.py @@ -0,0 +1,83 @@ +"""Tests for gateway-managed AI provider credentials.""" + +from types import SimpleNamespace +from uuid import uuid4 + +import pytest +from fastapi import HTTPException + +from app.config import settings +from app.core.encryption import encrypt_api_key +from app.services.ai.llm_gateway import ( + GATEWAY_MANAGED_KEY_SENTINEL, + gateway_managed_credentials_enabled, + is_gateway_managed_stored_key, + resolve_litellm_api_key, +) +from app.api.v1.routes import aiproviders as aiproviders_routes + + +def _sync_gateway_settings(**kwargs): + settings.LLM_GATEWAY_PASSTHROUGH_PROVIDER_KEYS = kwargs.get("passthrough", True) + settings.LLM_GATEWAY_ENABLED = kwargs.get("enabled", False) + settings.LLM_GATEWAY_BASE_URL = kwargs.get("base_url") + + +@pytest.fixture(autouse=True) +def _reset_gateway_settings(): + original = ( + settings.LLM_GATEWAY_ENABLED, + settings.LLM_GATEWAY_BASE_URL, + settings.LLM_GATEWAY_PASSTHROUGH_PROVIDER_KEYS, + ) + yield + ( + settings.LLM_GATEWAY_ENABLED, + settings.LLM_GATEWAY_BASE_URL, + settings.LLM_GATEWAY_PASSTHROUGH_PROVIDER_KEYS, + ) = original + + +def test_gateway_managed_credentials_enabled_when_passthrough_disabled(): + _sync_gateway_settings(passthrough=False) + assert gateway_managed_credentials_enabled() is True + + +def test_is_gateway_managed_stored_key_detects_sentinel(): + encrypted = encrypt_api_key(GATEWAY_MANAGED_KEY_SENTINEL) + assert is_gateway_managed_stored_key(encrypted) is True + assert is_gateway_managed_stored_key(encrypt_api_key("sk-real")) is False + + +def test_resolve_litellm_api_key_returns_none_for_gateway_managed_credential(): + _sync_gateway_settings( + enabled=True, + base_url="http://localhost:8080/litellm", + passthrough=False, + ) + + org_id = uuid4() + provider = SimpleNamespace( + provider="openai", + api_key=encrypt_api_key(GATEWAY_MANAGED_KEY_SENTINEL), + ) + db = SimpleNamespace( + query=lambda *_args, **_kwargs: SimpleNamespace( + filter=lambda *_a, **_k: SimpleNamespace(first=lambda: SimpleNamespace(llm_gateway_settings={})) + ) + ) + + assert resolve_litellm_api_key(org_id, db, provider) is None + + +def test_encrypt_provider_api_key_uses_sentinel_when_missing_and_gateway_managed(): + _sync_gateway_settings(passthrough=False) + encrypted = aiproviders_routes._encrypt_provider_api_key(None) + assert is_gateway_managed_stored_key(encrypted) is True + + +def test_encrypt_provider_api_key_requires_value_when_passthrough_enabled(): + _sync_gateway_settings(passthrough=True) + with pytest.raises(HTTPException) as exc: + aiproviders_routes._encrypt_provider_api_key(None) + assert exc.value.status_code == 400 diff --git a/tests/test_services/test_ai/test_llm_gateway.py b/tests/test_services/test_ai/test_llm_gateway.py new file mode 100644 index 00000000..50a787a7 --- /dev/null +++ b/tests/test_services/test_ai/test_llm_gateway.py @@ -0,0 +1,332 @@ +"""Tests for multi-gateway LLM resolver (Bifrost + LiteLLM Proxy).""" + +from types import SimpleNamespace +from uuid import uuid4 + +import pytest + +from app.config import settings +from app.services.ai import llm_gateway as gateway_module +from app.services.ai.llm_gateway import ( + apply_llm_gateway, + LITELLM_GATEWAY_PLACEHOLDER_API_KEY, + resolve_llm_gateway, +) + + +def _set_platform_gateway( + *, + enabled=False, + gateway_type="bifrost", + base_url=None, + virtual_key=None, + master_key=None, + passthrough=True, +): + settings.LLM_GATEWAY_ENABLED = enabled + settings.LLM_GATEWAY_TYPE = gateway_type + settings.LLM_GATEWAY_BASE_URL = base_url + settings.LLM_GATEWAY_VIRTUAL_KEY = virtual_key + settings.LLM_GATEWAY_MASTER_KEY = master_key + settings.LLM_GATEWAY_PASSTHROUGH_PROVIDER_KEYS = passthrough + + +@pytest.fixture(autouse=True) +def _reset_gateway_settings(): + original = ( + settings.LLM_GATEWAY_ENABLED, + settings.LLM_GATEWAY_TYPE, + settings.LLM_GATEWAY_BASE_URL, + settings.LLM_GATEWAY_VIRTUAL_KEY, + settings.LLM_GATEWAY_MASTER_KEY, + settings.LLM_GATEWAY_PASSTHROUGH_PROVIDER_KEYS, + ) + yield + ( + settings.LLM_GATEWAY_ENABLED, + settings.LLM_GATEWAY_TYPE, + settings.LLM_GATEWAY_BASE_URL, + settings.LLM_GATEWAY_VIRTUAL_KEY, + settings.LLM_GATEWAY_MASTER_KEY, + settings.LLM_GATEWAY_PASSTHROUGH_PROVIDER_KEYS, + ) = original + + +def _org_db(org_settings=None): + org_id = uuid4() + org = SimpleNamespace( + id=org_id, + llm_gateway_settings=org_settings, + ) + + class _Query: + def filter(self, *_args, **_kwargs): + return self + + def first(self): + return org + + db = SimpleNamespace(query=lambda *_args, **_kwargs: _Query()) + return org_id, db + + +def test_resolve_returns_none_when_platform_and_org_disabled(): + _set_platform_gateway(enabled=False) + org_id, db = _org_db({"enabled": None}) + assert resolve_llm_gateway(org_id, db) is None + + +def test_resolve_bifrost_uses_platform_defaults(): + _set_platform_gateway( + enabled=True, + base_url="http://bifrost.example.com/litellm", + virtual_key="platform-vk", + ) + org_id, db = _org_db({"enabled": None}) + config = resolve_llm_gateway(org_id, db) + + assert config is not None + assert config.gateway_type == "bifrost" + assert config.api_base == "http://bifrost.example.com/litellm" + assert config.virtual_key == "platform-vk" + + +def test_org_opt_out_overrides_platform(): + _set_platform_gateway(enabled=True, base_url="http://bifrost.example.com/litellm") + org_id, db = _org_db({"enabled": False}) + assert resolve_llm_gateway(org_id, db) is None + + +def test_org_override_url_and_virtual_key(monkeypatch): + _set_platform_gateway(enabled=False) + org_id, db = _org_db( + { + "enabled": True, + "gateway_type": "bifrost", + "base_url": "http://customer-bifrost:9090/litellm", + "virtual_key": "encrypted-vk", + } + ) + monkeypatch.setattr( + gateway_module, + "_decrypt_org_virtual_key", + lambda _raw: "org-vk", + ) + + config = resolve_llm_gateway(org_id, db) + + assert config is not None + assert config.gateway_type == "bifrost" + assert config.api_base == "http://customer-bifrost:9090/litellm" + assert config.virtual_key == "org-vk" + + +def test_legacy_org_json_without_gateway_type_defaults_to_bifrost(): + _set_platform_gateway(enabled=True, base_url="http://bifrost.example.com/litellm") + org_id, db = _org_db({"enabled": True}) + config = resolve_llm_gateway(org_id, db) + assert config is not None + assert config.gateway_type == "bifrost" + + +def test_resolve_litellm_proxy_uses_platform_defaults(): + _set_platform_gateway( + enabled=True, + gateway_type="litellm_proxy", + base_url="http://proxy.example.com:4000", + master_key="platform-master", + ) + org_id, db = _org_db({"enabled": None}) + config = resolve_llm_gateway(org_id, db) + + assert config is not None + assert config.gateway_type == "litellm_proxy" + assert config.api_base == "http://proxy.example.com:4000" + assert config.master_key == "platform-master" + + +def test_org_gateway_type_overrides_platform_type(): + _set_platform_gateway( + enabled=True, + gateway_type="bifrost", + base_url="http://bifrost.example.com/litellm", + ) + org_id, db = _org_db( + { + "enabled": True, + "gateway_type": "litellm_proxy", + "base_url": "http://org-proxy:4000", + } + ) + config = resolve_llm_gateway(org_id, db) + assert config is not None + assert config.gateway_type == "litellm_proxy" + assert config.api_base == "http://org-proxy:4000" + + +def test_apply_bifrost_gateway_injects_placeholder_api_key_when_gateway_managed(): + _set_platform_gateway( + enabled=True, + base_url="http://localhost:8080/litellm", + passthrough=False, + ) + org_id, db = _org_db({"enabled": True}) + result = apply_llm_gateway( + {"model": "openai/gpt-4o-mini", "messages": []}, + organization_id=org_id, + db=db, + ) + + assert result["api_base"] == "http://localhost:8080/litellm" + assert result["api_key"] == LITELLM_GATEWAY_PLACEHOLDER_API_KEY + + +def test_apply_bifrost_gateway_injects_api_base_and_headers(): + _set_platform_gateway( + enabled=True, + base_url="http://localhost:8080/litellm", + virtual_key="vk-123", + ) + org_id, db = _org_db({"enabled": True}) + result = apply_llm_gateway( + {"model": "openai/gpt-4o-mini", "api_key": "sk-test", "messages": []}, + organization_id=org_id, + db=db, + ) + + assert result["api_base"] == "http://localhost:8080/litellm" + assert result["extra_headers"]["x-bf-vk"] == "vk-123" + assert result["api_key"] == "sk-test" + + +def test_apply_litellm_proxy_injects_master_key_when_not_passthrough(): + _set_platform_gateway( + enabled=True, + gateway_type="litellm_proxy", + base_url="http://localhost:4000", + master_key="proxy-master", + passthrough=False, + ) + org_id, db = _org_db({"enabled": True}) + result = apply_llm_gateway( + {"model": "openai/gpt-4o-mini", "api_key": "sk-test", "messages": []}, + organization_id=org_id, + db=db, + ) + + assert result["api_base"] == "http://localhost:4000" + assert result["api_key"] == "proxy-master" + + +def test_apply_litellm_proxy_keeps_org_key_when_passthrough(): + _set_platform_gateway( + enabled=True, + gateway_type="litellm_proxy", + base_url="http://localhost:4000", + master_key="proxy-master", + passthrough=True, + ) + org_id, db = _org_db({"enabled": True}) + result = apply_llm_gateway( + {"model": "openai/gpt-4o-mini", "api_key": "sk-test", "messages": []}, + organization_id=org_id, + db=db, + ) + + assert result["api_base"] == "http://localhost:4000" + assert result["api_key"] == "sk-test" + + +def test_apply_bifrost_gateway_forces_openai_compatible_routing_for_gemini(): + _set_platform_gateway( + enabled=True, + gateway_type="bifrost", + base_url="http://localhost:8080/litellm", + master_key="bifrost-key", + passthrough=False, + ) + org_id, db = _org_db({"enabled": True}) + result = apply_llm_gateway( + {"model": "gemini/gemini-2.5-flash", "api_key": "google-key", "messages": []}, + organization_id=org_id, + db=db, + ) + + assert result["api_base"] == "http://localhost:8080/litellm" + assert result["custom_llm_provider"] == "openai" + assert result["model"] == "gemini/gemini-2.5-flash" + + +def test_apply_litellm_proxy_forces_openai_compatible_routing_for_gemini(): + _set_platform_gateway( + enabled=True, + gateway_type="litellm_proxy", + base_url="http://localhost:4000", + master_key="proxy-master", + passthrough=False, + ) + org_id, db = _org_db({"enabled": True}) + result = apply_llm_gateway( + {"model": "gemini/gemini-2.5-flash", "api_key": "google-key", "messages": []}, + organization_id=org_id, + db=db, + ) + + assert result["custom_llm_provider"] == "openai" + assert result["model"] == "gemini/gemini-2.5-flash" + + +def test_apply_gateway_leaves_openai_models_unmodified(): + _set_platform_gateway( + enabled=True, + gateway_type="bifrost", + base_url="http://localhost:8080/litellm", + master_key="bifrost-key", + passthrough=False, + ) + org_id, db = _org_db({"enabled": True}) + result = apply_llm_gateway( + {"model": "openai/gpt-4o-mini", "api_key": "sk-test", "messages": []}, + organization_id=org_id, + db=db, + ) + + assert "custom_llm_provider" not in result + assert result["model"] == "openai/gpt-4o-mini" + + +def test_apply_gateway_without_model_skips_gemini_proxy_routing(): + """Without model in kwargs or routing param, native-path routing is skipped.""" + _set_platform_gateway( + enabled=True, + gateway_type="bifrost", + base_url="http://localhost:8080/litellm", + passthrough=False, + ) + org_id, db = _org_db({"enabled": True}) + result = apply_llm_gateway( + {"api_key": "google-key"}, + organization_id=org_id, + db=db, + ) + + assert "custom_llm_provider" not in result + + +def test_apply_gateway_routing_model_enables_gemini_proxy_routing_for_gepa_batch(): + _set_platform_gateway( + enabled=True, + gateway_type="bifrost", + base_url="http://localhost:8080/litellm", + passthrough=False, + ) + org_id, db = _org_db({"enabled": True}) + result = apply_llm_gateway( + {"api_key": "google-key"}, + organization_id=org_id, + db=db, + model="gemini/gemini-2.5-flash", + ) + + assert result["custom_llm_provider"] == "openai" + assert "model" not in result diff --git a/tests/test_services/test_ai/test_llm_service.py b/tests/test_services/test_ai/test_llm_service.py index 580bd2ad..7c6abd8a 100644 --- a/tests/test_services/test_ai/test_llm_service.py +++ b/tests/test_services/test_ai/test_llm_service.py @@ -12,6 +12,38 @@ llm_module = importlib.import_module("app.services.ai.llm_service") +def _mock_org_db(org_settings=None): + org = SimpleNamespace(llm_gateway_settings=org_settings) + return SimpleNamespace( + query=lambda *_args, **_kwargs: SimpleNamespace( + filter=lambda *_a, **_k: SimpleNamespace(first=lambda: org) + ) + ) + + +@pytest.fixture(autouse=True) +def _reset_llm_gateway_settings(): + from app.config import settings + + original = ( + settings.LLM_GATEWAY_ENABLED, + settings.LLM_GATEWAY_BASE_URL, + settings.LLM_GATEWAY_VIRTUAL_KEY, + settings.LLM_GATEWAY_PASSTHROUGH_PROVIDER_KEYS, + ) + settings.LLM_GATEWAY_ENABLED = False + settings.LLM_GATEWAY_BASE_URL = None + settings.LLM_GATEWAY_VIRTUAL_KEY = None + settings.LLM_GATEWAY_PASSTHROUGH_PROVIDER_KEYS = True + yield + ( + settings.LLM_GATEWAY_ENABLED, + settings.LLM_GATEWAY_BASE_URL, + settings.LLM_GATEWAY_VIRTUAL_KEY, + settings.LLM_GATEWAY_PASSTHROUGH_PROVIDER_KEYS, + ) = original + + def test_litellm_model_name_maps_known_provider_prefixes(): assert LLMService._litellm_model_name(ModelProvider.OPENAI, "gpt-4o") == "openai/gpt-4o" assert LLMService._litellm_model_name(ModelProvider.GOOGLE, "gemini-1.5-pro") == "gemini/gemini-1.5-pro" @@ -68,7 +100,7 @@ def test_generate_response_success_with_normalized_usage(monkeypatch): llm_provider=ModelProvider.OPENAI, llm_model="gpt-4o-mini", organization_id=uuid4(), - db=object(), + db=_mock_org_db(), temperature=0.2, ) @@ -78,6 +110,55 @@ def test_generate_response_success_with_normalized_usage(monkeypatch): assert result["processing_time"] >= 0 +def test_generate_response_applies_llm_gateway(monkeypatch): + from app.config import settings + + settings.LLM_GATEWAY_ENABLED = True + settings.LLM_GATEWAY_BASE_URL = "http://localhost:8080/litellm" + settings.LLM_GATEWAY_VIRTUAL_KEY = "test-vk" + + service = LLMService() + monkeypatch.setattr( + service, + "_get_ai_provider", + lambda *_args, **_kwargs: SimpleNamespace(api_key="encrypted-key"), + ) + + encryption_module = importlib.import_module("app.core.encryption") + monkeypatch.setattr(encryption_module, "decrypt_api_key", lambda value: f"decrypted::{value}") + + captured = {} + + def _fake_completion(**kwargs): + captured.update(kwargs) + return SimpleNamespace( + choices=[ + SimpleNamespace( + message=SimpleNamespace(content="via bifrost"), + finish_reason="stop", + ) + ], + usage=SimpleNamespace(prompt_tokens=1, completion_tokens=1, total_tokens=2), + ) + + monkeypatch.setattr(llm_module.litellm, "completion", _fake_completion) + + org_id = uuid4() + db = _mock_org_db() + + result = service.generate_response( + messages=[{"role": "user", "content": "hello"}], + llm_provider=ModelProvider.OPENAI, + llm_model="gpt-4o-mini", + organization_id=org_id, + db=db, + ) + + assert result["text"] == "via bifrost" + assert captured["api_base"] == "http://localhost:8080/litellm" + assert captured["extra_headers"]["x-bf-vk"] == "test-vk" + + def test_generate_response_wraps_litellm_errors(monkeypatch): service = LLMService() monkeypatch.setattr( @@ -100,5 +181,5 @@ def test_generate_response_wraps_litellm_errors(monkeypatch): llm_provider=ModelProvider.OPENAI, llm_model="gpt-4o-mini", organization_id=uuid4(), - db=object(), + db=_mock_org_db(), ) diff --git a/tests/test_services/test_optimization/test_evaluator.py b/tests/test_services/test_optimization/test_evaluator.py new file mode 100644 index 00000000..f65e95f2 --- /dev/null +++ b/tests/test_services/test_optimization/test_evaluator.py @@ -0,0 +1,34 @@ +"""Tests for GEPA optimization evaluator LM resolution.""" + +from types import SimpleNamespace + +from app.models.enums import ModelProvider +from app.services.optimization.evaluator import _resolve_scoring_lm + + +def _provider(name: str, *, active: bool = True): + return SimpleNamespace(provider=name, is_active=active, api_key="enc") + + +def test_resolve_scoring_lm_uses_optimization_lm_identifier(): + providers = [_provider("openai"), _provider("anthropic")] + llm_provider, llm_model = _resolve_scoring_lm("openai/gpt-4.1", providers) + + assert llm_provider == ModelProvider.OPENAI + assert llm_model == "gpt-4.1" + + +def test_resolve_scoring_lm_does_not_pair_anthropic_provider_with_openai_model(): + providers = [_provider("anthropic"), _provider("openai")] + llm_provider, llm_model = _resolve_scoring_lm("openai/gpt-4.1", providers) + + assert llm_provider == ModelProvider.OPENAI + assert llm_model == "gpt-4.1" + + +def test_resolve_scoring_lm_prefers_openai_when_lm_identifier_missing(): + providers = [_provider("anthropic"), _provider("openai")] + llm_provider, llm_model = _resolve_scoring_lm(None, providers) + + assert llm_provider == ModelProvider.OPENAI + assert llm_model == "gpt-4o-mini" diff --git a/tests/test_services/test_optimization/test_lm_resolver.py b/tests/test_services/test_optimization/test_lm_resolver.py index 74786dd4..75852bd7 100644 --- a/tests/test_services/test_optimization/test_lm_resolver.py +++ b/tests/test_services/test_optimization/test_lm_resolver.py @@ -2,6 +2,7 @@ import importlib from types import SimpleNamespace +from uuid import uuid4 import pytest @@ -23,12 +24,23 @@ def test_resolve_api_key_returns_decrypted_key(monkeypatch): SimpleNamespace(provider="openai", api_key="enc-1", is_active=True), SimpleNamespace(provider="anthropic", api_key="enc-2", is_active=True), ] - monkeypatch.setattr(resolver_module, "decrypt_api_key", lambda v: f"dec::{v}") + monkeypatch.setattr( + resolver_module, + "resolve_litellm_api_key", + lambda _org_id, _db, provider: f"dec::{provider.api_key}", + ) - assert resolver_module.resolve_api_key("openai/gpt-4o", providers) == "dec::enc-1" + org_id = uuid4() + db = SimpleNamespace() + assert ( + resolver_module.resolve_api_key("openai/gpt-4o", providers, org_id, db) + == "dec::enc-1" + ) def test_resolve_api_key_raises_for_missing_provider(): providers = [SimpleNamespace(provider="anthropic", api_key="enc-2", is_active=True)] + org_id = uuid4() + db = SimpleNamespace() with pytest.raises(RuntimeError, match="No active AI provider"): - resolver_module.resolve_api_key("openai/gpt-4o", providers) + resolver_module.resolve_api_key("openai/gpt-4o", providers, org_id, db)