diff --git a/.gitignore b/.gitignore index a8f8d4f2..cf2f92d4 100644 --- a/.gitignore +++ b/.gitignore @@ -111,4 +111,9 @@ frontend/dist/ # Enterprise license private keys enterprise/keys/*.pem +# Celery Beat persistent schedule (SQLite + WAL sidecars) +celerybeat-schedule +celerybeat-schedule-* +celerybeat-schedule.* + # \ No newline at end of file diff --git a/README.md b/README.md index cc48559e..e14527d3 100644 --- a/README.md +++ b/README.md @@ -54,8 +54,21 @@ There are two ways to run the application: This will automatically: - Pull pre-built images from GitHub Container Registry (no build required!) - - Start all services (database, Redis, API, worker) + - Start all services: `db`, `redis`, `api`, `media`, `worker`, `beat`, `worker-imports`, `worker-usage` - Run database migrations automatically on startup + + | Service | Purpose | + |---------|---------| + | `db` | PostgreSQL | + | `redis` | Redis (Celery broker + usage counters) | + | `api` | HTTP API + frontend | + | `media` | Live voice WebSocket media server | + | `worker` | Celery: `celery` (evaluator cron dispatch), `audio-metrics` queues | + | `beat` | Celery Beat scheduler + `platform` queue worker (alerts, FX, OSS prune) — **single replica** | + | `worker-imports` | Celery: `imports`, `diarization`, `eval-control`, `evaluations` | + | `worker-usage` | Celery: `usage` queue (flush Redis counters + cost recompute) | + + **Usage costs:** token/cost rollups stay stale without `beat`, `worker-usage`, and default `worker` (evaluator crons; or `eai start-all`). **Using a specific version:** ```bash @@ -137,20 +150,29 @@ docker compose up -d url: "redis://localhost:6379/0" ``` -4. **Start the application and worker** +4. **Start the application and workers** + + **Infra only (optional):** if Postgres/Redis run in Docker but the app runs locally: + ```bash + docker compose up -d db redis + ``` - **Option A: Start both together (Recommended)** + **Option A: Start everything together (Recommended)** ```bash eai start-all --config config.yml ``` - This single command will: - - Start the API server - - Start the Celery worker (for background task processing) - - Run database migrations automatically - - Build the frontend (if needed) + This single command spawns: + - API server (uvicorn) + - Telephony media server (`media` port, default 8001) + - Celery worker (`celery`, `audio-metrics`) + - Celery worker (`imports`, `diarization`, `eval-control`, `evaluations`) + - Celery worker (`usage` — flush + cost recompute) + - Celery Beat (platform schedules: usage flush, alerts, FX refresh, OSS prune) + + It also runs database migrations and builds the frontend when needed. - Press `Ctrl+C` to stop both services together. + Press `Ctrl+C` to stop all processes. **Option B: Start separately (for advanced use)** @@ -163,6 +185,11 @@ docker compose up -d ```bash eai worker --config config.yml ``` + + For platform periodic tasks (usage flush, alerts, etc.), start Celery Beat in a separate terminal (single replica): + ```bash + eai beat --config config.yml + ``` Or use the Celery command directly: ```bash @@ -247,7 +274,7 @@ make test-docker-db TEST_DB_HOST=localhost TEST_DB_PORT=5432 TEST_DB_NAME=effici ### Start Application and Worker Together (Recommended) ```bash -# Start both app and worker with default config.yml +# Start API + all workers with default config.yml eai start-all # Start with custom config @@ -264,9 +291,17 @@ eai start-all --no-reload --no-build-frontend # Customize worker log level eai start-all --worker-loglevel debug + +# Skip dedicated workers (not recommended for production) +eai start-all --no-imports-worker +eai start-all --no-usage-worker +eai start-all --no-telephony-worker + +# Tune usage worker concurrency (default: 4, thread pool) +eai start-all --usage-worker-concurrency 8 ``` -**Note:** This is the recommended way to run EfficientAI. It starts both the API server and Celery worker in a single command. Press `Ctrl+C` to stop both services. +**Note:** This is the recommended local-dev workflow. One command spawns the API, telephony media server, three Celery workers (`celery,audio-metrics` · `imports,…` · `usage`), and Celery Beat. Press `Ctrl+C` to stop all processes. For Docker deployments, use `docker compose up -d` instead (separate containers per role; see Quick Start). ### Start Application Only ```bash @@ -340,7 +375,50 @@ eai worker --loglevel debug celery -A app.workers.celery_app worker --loglevel=info ``` -**Note:** The worker is required for processing background tasks (transcription, evaluation, etc.). If you use `eai start-all`, the worker starts automatically. Only use this command if you need to run the worker separately. +**Note:** Workers are required for background tasks (transcription, evaluation, usage cost flush, etc.). If you use `eai start-all`, they start automatically. Only use `eai worker` if you need to run a worker separately (e.g. `eai worker --queues usage` for the usage queue only). + +### Usage Pricing Ops +Manage model pricing rates and backfill stored usage costs on `llm_usage_daily` rollups. Requires `beat`, `worker-usage`, and default `worker` (or `eai start-all`). + +```bash +# Upsert model_pricing_rates from app/config/models.json +eai usage seed-rates --config config.yml + +# Compare models.json pricing vs Postgres +eai usage diff-rates --config config.yml + +# Backfill costs in-process (all orgs; use after migrate or catalog change) +eai usage recompute --config config.yml --sync + +# Async recompute via usage queue (requires --organization-id) +eai usage recompute --config config.yml --organization-id + +# Optional scopes: --model, --usage-kind, --start-date, --end-date + +# Optional: fetch LiteLLM prices into pricing_catalog.json +eai usage sync-litellm --local +eai usage sync-litellm --local --write-models +``` + +**After migrations or catalog changes:** +```bash +eai migrate +eai usage seed-rates --config config.yml +eai usage recompute --config config.yml --sync +``` + +**Flush / Usage UI tuning** — set in `.env` (see `env.example`): + +| Variable | Default | Purpose | +|----------|---------|---------| +| `USAGE_FLUSH_BUCKET_BATCH_SIZE` | `500` | Buckets per DB transaction | +| `USAGE_FLUSH_MAX_BATCHES_PER_RUN` | `30` | Batches per flush tick (≤ **15,000** buckets/run) | +| `USAGE_FLUSH_BEAT_SECONDS` | `120` | Celery Beat flush interval (~2 min lag vs Redis) | +| `USAGE_FLUSH_LOCK_TTL_SECONDS` | `300` | Per-org flush lock TTL | +| `USAGE_READ_CACHE_TTL_SECONDS` | `90` | Redis cache TTL for usage summary/breakdown/filters | +| `CRON_DISPATCH_INTERVAL_SECONDS` | `30` | Evaluator cron dispatcher tick (default worker) | + +Usage UI reads Postgres only (summary/breakdown/filters); Redis counters flush on the Celery Beat schedule (~2 min eventual consistency). If Redis backlog grows, lower `USAGE_FLUSH_BEAT_SECONDS` or raise `USAGE_FLUSH_MAX_BATCHES_PER_RUN`. ### Generate Config File ```bash @@ -455,6 +533,10 @@ POSTGRES_PASSWORD=password POSTGRES_DB=efficientai SECRET_KEY=your-secret-key-here +# Usage cost flush (see README "Usage Pricing Ops"; full list in env.example) +# USAGE_FLUSH_BEAT_SECONDS=120 +# USAGE_FLUSH_MAX_BATCHES_PER_RUN=30 + # Optional: GCS blob storage (also set storage.blob_provider: gcs in config.yml) BLOB_STORAGE_PROVIDER=gcs GCS_BUCKET_NAME=your-gcs-bucket diff --git a/app/api/v1/api.py b/app/api/v1/api.py index bbf422e0..055bc147 100644 --- a/app/api/v1/api.py +++ b/app/api/v1/api.py @@ -47,6 +47,8 @@ dashboard, llm_gateway, platform_admin, + org_usage, + usage_pricing, ) api_router = APIRouter() @@ -97,3 +99,5 @@ api_router.include_router(dashboard.router) api_router.include_router(llm_gateway.router) api_router.include_router(platform_admin.router) +api_router.include_router(org_usage.router) +api_router.include_router(usage_pricing.router) diff --git a/app/api/v1/routes/agents.py b/app/api/v1/routes/agents.py index 0be8cd99..17094877 100644 --- a/app/api/v1/routes/agents.py +++ b/app/api/v1/routes/agents.py @@ -1,1025 +1,1080 @@ -""" -Agents API Routes -Complete CRUD operations for test agents -""" -from fastapi import APIRouter, Depends, HTTPException, status, Query -from fastapi.responses import JSONResponse, Response -from sqlalchemy.orm import Session -from typing import List, Optional -from uuid import UUID -import random -from pydantic import BaseModel -from loguru import logger - -from app.dependencies import get_db, get_organization_id, get_workspace_id, get_api_key -from app.models.database import ( - Agent, ConversationEvaluation, TestAgentConversation, VoiceBundle, - AIProvider, Integration, IntegrationPlatform, CallMediumEnum, - Evaluator, EvaluatorResult, CallRecording, Scenario, -) -from sqlalchemy import and_ -from app.models.schemas import ( - AgentCreate, - AgentUpdate, - AgentResponse, - AgentPhoneAssignmentCheckResponse, - AgentPhoneAssignmentConflict, - CallMediumEnum as CallMediumEnumSchema, - GenerateTestPromptRequest, - GenerateTestPromptResponse, - GenerateScenariosFromPromptRequest, - GenerateScenariosFromPromptResponse, - GenerateTestSetupRequest, - GenerateTestSetupResponse, - GeneratedScenarioDraftResponse, - TestPromptSectionResponse, -) - -router = APIRouter(prefix="/agents", tags=["agents"]) - - -def _validate_agent_phone_assignment( - db: Session, - *, - organization_id: UUID, - call_medium, - phone_number: Optional[str], - telephony_phone_number_id: Optional[UUID], - exclude_agent_id: Optional[UUID] = None, -) -> None: - """Raise HTTPException if phone assignment conflicts with another agent.""" - if call_medium != CallMediumEnum.PHONE_CALL: - return - if not phone_number and not telephony_phone_number_id: - return - - from app.services.telephony.phone_routing import find_agent_phone_assignment_conflict - - conflict = find_agent_phone_assignment_conflict( - db, - organization_id=organization_id, - phone_number=phone_number, - telephony_phone_number_id=telephony_phone_number_id, - exclude_agent_id=exclude_agent_id, - ) - if not conflict: - return - if conflict.get("error") == "telephony_not_found": - raise HTTPException(status_code=404, detail="Telephony phone number not found") - - agent_name = conflict["agent_name"] - raise HTTPException( - status_code=status.HTTP_409_CONFLICT, - detail={ - "message": f'This number is already assigned to agent "{agent_name}".', - "agent_id": str(conflict["agent_id"]), - "agent_name": agent_name, - "phone_number": conflict["phone_number"], - }, - ) - - -def resolve_agent_by_path_id( - db: Session, - *, - organization_id: UUID, - workspace_id: UUID, - agent_id: str, -) -> Agent: - """Resolve agent by UUID primary key or 6-digit agent_id (same as GET /agents/{id}).""" - try: - agent_uuid = UUID(agent_id) - agent = db.query(Agent).filter( - and_( - Agent.id == agent_uuid, - Agent.organization_id == organization_id, - Agent.workspace_id == workspace_id, - ) - ).first() - except ValueError: - agent = db.query(Agent).filter( - and_( - Agent.agent_id == agent_id, - Agent.organization_id == organization_id, - Agent.workspace_id == workspace_id, - ) - ).first() - if not agent: - raise HTTPException(status_code=404, detail=f"Agent {agent_id} not found") - return agent - - -# ====================================================================== -# AI Generation for agent descriptions -# ====================================================================== - -class GenerateAgentDescriptionRequest(BaseModel): - description: str - tone: Optional[str] = "professional" - format_style: Optional[str] = "structured" - provider: Optional[str] = None - model: Optional[str] = None - agent_id: Optional[UUID] = None - include_linked_scenarios: bool = True - append_scenarios_to_output: bool = False - - -GENERATE_AGENT_DESCRIPTION_SYSTEM = ( - "You are an expert at writing clear, well-structured descriptions for voice AI test agents. " - "The user will describe what they need the agent to do, and you will generate a comprehensive, " - "well-formatted agent description in markdown.\n\n" - "Guidelines:\n" - "- Use clear markdown structure: headings, bullet points, numbered lists\n" - "- Include sections for: Purpose, Behavior, Expected Interactions, Personality Traits, and Constraints\n" - "- Be specific about the agent's role, tone of voice, and how it should handle conversations\n" - "- Include example scenarios or edge cases where helpful\n" - "- Return ONLY the description in markdown, no preamble or explanation about what you did" -) - - -from app.services.ai.llm_resolver import get_llm_provider_and_model as _get_llm_provider_and_model - - -@router.post("/generate-description") -async def generate_agent_description( - data: GenerateAgentDescriptionRequest, - organization_id: UUID = Depends(get_organization_id), - workspace_id: UUID = Depends(get_workspace_id), - api_key: str = Depends(get_api_key), - db: Session = Depends(get_db), -): - """Generate an agent description using AI from a brief description.""" - from app.services.ai.llm_service import llm_service - from app.services.testing.test_agent_simulation_prompt import ( - format_scenarios_for_generation_context, - format_scenarios_reference_appendix, - load_linked_scenarios_for_agent, - merge_generated_description_with_scenario_appendix, - ) - - if not data.description.strip(): - raise HTTPException(400, "Description is required") - - provider_enum, model_str = _get_llm_provider_and_model( - organization_id, db, data.provider, data.model - ) - - linked_scenarios = [] - if data.agent_id and data.include_linked_scenarios: - agent = db.query(Agent).filter( - Agent.id == data.agent_id, - Agent.organization_id == organization_id, - Agent.workspace_id == workspace_id, - ).first() - if not agent: - raise HTTPException(status_code=404, detail=f"Agent {data.agent_id} not found") - linked_scenarios = load_linked_scenarios_for_agent( - db, - organization_id=organization_id, - workspace_id=workspace_id, - agent_id=agent.id, - ) - - user_prompt_parts = [ - "Create a detailed agent description for the following:", - "", - f"Description: {data.description}", - f"Tone: {data.tone or 'professional'}", - f"Format: {data.format_style or 'structured'}", - ] - scenario_context = format_scenarios_for_generation_context(linked_scenarios) - if scenario_context: - user_prompt_parts.extend(["", scenario_context]) - user_prompt_parts.extend([ - "", - "Generate a comprehensive, well-formatted agent description in markdown.", - ]) - user_prompt = "\n".join(user_prompt_parts) - - messages = [ - {"role": "system", "content": GENERATE_AGENT_DESCRIPTION_SYSTEM}, - {"role": "user", "content": user_prompt}, - ] - - try: - result = llm_service.generate_response( - messages=messages, - llm_provider=provider_enum, - llm_model=model_str, - organization_id=organization_id, - db=db, - temperature=0.7, - max_tokens=4000, - ) - content = result["text"] - if data.append_scenarios_to_output and linked_scenarios: - appendix = format_scenarios_reference_appendix(linked_scenarios) - content = merge_generated_description_with_scenario_appendix(content, appendix) - return {"content": content, "provider": provider_enum.value, "model": model_str} - except Exception as e: - logger.error(f"[Agents] AI description generation failed: {repr(e)}") - raise HTTPException(500, f"AI generation failed: {str(e)}") - - -def _test_prompt_section_responses(sections) -> list[TestPromptSectionResponse]: - return [ - TestPromptSectionResponse(key=s.key, title=s.title, content=s.content) - for s in sections - ] - - -def _scenario_draft_responses(scenarios) -> list[GeneratedScenarioDraftResponse]: - return [ - GeneratedScenarioDraftResponse(name=s.name, description=s.description, goal=s.goal) - for s in scenarios - ] - - -@router.post("/generate-test-prompt", response_model=GenerateTestPromptResponse) -async def generate_test_prompt( - data: GenerateTestPromptRequest, - organization_id: UUID = Depends(get_organization_id), - api_key: str = Depends(get_api_key), - db: Session = Depends(get_db), -): - """Stage 1: generate foundational test agent prompt from production prompt.""" - from app.services.testing.agent_test_setup_generation import ( - generate_test_prompt_from_production, - ) - - if not data.production_prompt.strip(): - raise HTTPException(400, "Production prompt is required") - - provider_enum, model_str = _get_llm_provider_and_model( - organization_id, db, data.provider, data.model, data.credential_id - ) - - try: - result = generate_test_prompt_from_production( - data.production_prompt, - agent_name=data.agent_name, - language=data.language, - call_type=data.call_type, - additional_context=data.additional_context, - llm_provider=provider_enum, - llm_model=model_str, - organization_id=organization_id, - db=db, - llm_config=data.llm_config, - credential_id=data.credential_id, - ) - return GenerateTestPromptResponse( - sections=_test_prompt_section_responses(result.sections), - test_agent_prompt=result.test_agent_prompt, - provider=result.provider, - model=result.model, - ) - except ValueError as e: - raise HTTPException(400, str(e)) from e - except Exception as e: - logger.error(f"[Agents] Test prompt generation failed: {repr(e)}") - raise HTTPException(500, f"AI generation failed: {str(e)}") from e - - -@router.post("/generate-scenarios-from-prompt", response_model=GenerateScenariosFromPromptResponse) -async def generate_scenarios_from_prompt( - data: GenerateScenariosFromPromptRequest, - organization_id: UUID = Depends(get_organization_id), - api_key: str = Depends(get_api_key), - db: Session = Depends(get_db), -): - """Stage 2: generate scenario drafts from a test agent prompt.""" - from app.services.testing.agent_test_setup_generation import ( - generate_scenarios_from_test_prompt, - ) - - if not data.test_agent_prompt.strip(): - raise HTTPException(400, "Test agent prompt is required") - - provider_enum, model_str = _get_llm_provider_and_model( - organization_id, db, data.provider, data.model, data.credential_id - ) - - try: - result = generate_scenarios_from_test_prompt( - data.test_agent_prompt, - agent_name=data.agent_name, - scenario_count=data.scenario_count, - language=data.language, - call_type=data.call_type, - additional_context=data.additional_context, - llm_provider=provider_enum, - llm_model=model_str, - organization_id=organization_id, - db=db, - llm_config=data.llm_config, - credential_id=data.credential_id, - ) - return GenerateScenariosFromPromptResponse( - scenarios=_scenario_draft_responses(result.scenarios), - provider=result.provider, - model=result.model, - ) - except ValueError as e: - raise HTTPException(400, str(e)) from e - except Exception as e: - logger.error(f"[Agents] Scenario generation failed: {repr(e)}") - raise HTTPException(500, f"AI generation failed: {str(e)}") from e - - -@router.post("/generate-test-setup", response_model=GenerateTestSetupResponse) -async def generate_test_setup( - data: GenerateTestSetupRequest, - organization_id: UUID = Depends(get_organization_id), - api_key: str = Depends(get_api_key), - db: Session = Depends(get_db), -): - """Run stage 1 then stage 2: foundational test prompt + scenario drafts.""" - from app.services.testing.agent_test_setup_generation import ( - generate_scenarios_from_test_prompt, - generate_test_prompt_from_production, - ) - - if not data.production_prompt.strip(): - raise HTTPException(400, "Production prompt is required") - - provider_enum, model_str = _get_llm_provider_and_model( - organization_id, db, data.provider, data.model, data.credential_id - ) - - try: - prompt_result = generate_test_prompt_from_production( - data.production_prompt, - agent_name=data.agent_name, - language=data.language, - call_type=data.call_type, - additional_context=data.additional_context, - llm_provider=provider_enum, - llm_model=model_str, - organization_id=organization_id, - db=db, - llm_config=data.llm_config, - credential_id=data.credential_id, - ) - scenario_result = generate_scenarios_from_test_prompt( - prompt_result.test_agent_prompt, - agent_name=data.agent_name, - scenario_count=data.scenario_count, - language=data.language, - call_type=data.call_type, - additional_context=data.additional_context, - llm_provider=provider_enum, - llm_model=model_str, - organization_id=organization_id, - db=db, - llm_config=data.llm_config, - credential_id=data.credential_id, - ) - return GenerateTestSetupResponse( - sections=_test_prompt_section_responses(prompt_result.sections), - test_agent_prompt=prompt_result.test_agent_prompt, - scenarios=_scenario_draft_responses(scenario_result.scenarios), - provider=prompt_result.provider, - model=prompt_result.model, - ) - except ValueError as e: - raise HTTPException(400, str(e)) from e - except Exception as e: - logger.error(f"[Agents] Test setup generation failed: {repr(e)}") - raise HTTPException(500, f"AI generation failed: {str(e)}") from e - - -def generate_unique_agent_id(db: Session) -> str: - """Generate a unique 6-digit agent ID.""" - max_attempts = 100 - for _ in range(max_attempts): - agent_id = f"{random.randint(100000, 999999)}" - existing = db.query(Agent).filter(Agent.agent_id == agent_id).first() - if not existing: - return agent_id - raise HTTPException( - status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, - detail="Failed to generate unique agent ID" - ) - - -def get_agent_dependencies(db: Session, organization_id: UUID, agent_uuid: UUID) -> dict: - """Return dependency counts that block non-force delete.""" - evaluators_count = db.query(Evaluator).filter( - Evaluator.agent_id == agent_uuid, - Evaluator.organization_id == organization_id, - ).count() - - evaluator_results_count = db.query(EvaluatorResult).filter( - EvaluatorResult.agent_id == agent_uuid, - EvaluatorResult.organization_id == organization_id, - ).count() - - call_recordings_count = db.query(CallRecording).filter( - CallRecording.agent_id == agent_uuid, - CallRecording.organization_id == organization_id, - ).count() - - conversation_evaluations_count = db.query(ConversationEvaluation).filter( - ConversationEvaluation.agent_id == agent_uuid, - ConversationEvaluation.organization_id == organization_id, - ).count() - - test_conversations_count = db.query(TestAgentConversation).filter( - TestAgentConversation.agent_id == agent_uuid, - TestAgentConversation.organization_id == organization_id, - ).count() - - dependencies = {} - if evaluators_count > 0: - dependencies["evaluators"] = evaluators_count - if evaluator_results_count > 0: - dependencies["evaluator_results"] = evaluator_results_count - if call_recordings_count > 0: - dependencies["call_recordings"] = call_recordings_count - if conversation_evaluations_count > 0: - dependencies["conversation_evaluations"] = conversation_evaluations_count - if test_conversations_count > 0: - dependencies["test_conversations"] = test_conversations_count - - return dependencies - - -@router.post("", response_model=AgentResponse, status_code=status.HTTP_201_CREATED) -async def create_agent( - agent: AgentCreate, - organization_id: UUID = Depends(get_organization_id), - workspace_id: UUID = Depends(get_workspace_id), - db: Session = Depends(get_db) -): - """Create a new test agent. - - The agent is stamped with the active workspace from the - ``X-Workspace-Id`` header (falling back to the org's Default). - """ - # Validate phone_number is provided when call_medium is phone_call - if agent.call_medium == CallMediumEnumSchema.PHONE_CALL and not agent.phone_number: - raise HTTPException( - status_code=status.HTTP_400_BAD_REQUEST, - detail="phone_number is required when call_medium is phone_call" - ) - - # Validate voice_bundle_id exists, is active, and belongs to organization - voice_bundle = db.query(VoiceBundle).filter( - and_( - VoiceBundle.id == agent.voice_bundle_id, - VoiceBundle.organization_id == organization_id, - VoiceBundle.is_active == True, - ) - ).first() - if not voice_bundle: - raise HTTPException( - status_code=status.HTTP_400_BAD_REQUEST, - detail="Active voice bundle not found", - ) - - # Validate voice_ai_integration_id exists and belongs to organization - if agent.voice_ai_integration_id: - integration = db.query(Integration).filter( - and_( - Integration.id == agent.voice_ai_integration_id, - Integration.organization_id == organization_id, - Integration.is_active == True - ) - ).first() - if not integration: - raise HTTPException(status_code=404, detail="Integration not found or inactive") - - if integration.platform not in [ - IntegrationPlatform.RETELL, - IntegrationPlatform.VAPI, - IntegrationPlatform.ELEVENLABS, - IntegrationPlatform.SMALLEST, - ]: - raise HTTPException( - status_code=status.HTTP_400_BAD_REQUEST, - detail=( - f"Integration platform {integration.platform.value} is not supported for Voice AI agents. " - "Only Retell, Vapi, ElevenLabs, and Smallest are supported." - ) - ) - - if not agent.voice_ai_agent_id: - raise HTTPException( - status_code=status.HTTP_400_BAD_REQUEST, - detail="voice_ai_agent_id is required when voice_ai_integration_id is provided" - ) - - _validate_agent_phone_assignment( - db, - organization_id=organization_id, - call_medium=agent.call_medium, - phone_number=agent.phone_number, - telephony_phone_number_id=agent.telephony_phone_number_id, - ) - - # Generate unique 6-digit agent_id - agent_id = generate_unique_agent_id(db) - - db_agent = Agent( - agent_id=agent_id, - organization_id=organization_id, - workspace_id=workspace_id, - name=agent.name, - phone_number=agent.phone_number, - language=agent.language, - description=agent.description, - call_type=agent.call_type, - call_medium=agent.call_medium, - telephony_phone_number_id=agent.telephony_phone_number_id, - voice_bundle_id=agent.voice_bundle_id, - ai_provider_id=agent.ai_provider_id, - voice_ai_integration_id=agent.voice_ai_integration_id, - voice_ai_agent_id=agent.voice_ai_agent_id, - provider_prompt=agent.provider_prompt, - silence_hangup_secs=agent.silence_hangup_secs, - ) - db.add(db_agent) - db.commit() - db.refresh(db_agent) - - from app.services.telephony.phone_routing import sync_agent_telephony_number_link - - sync_agent_telephony_number_link(db, db_agent) - db.refresh(db_agent) - - has_provider_prompt = isinstance(agent.provider_prompt, str) and bool(agent.provider_prompt.strip()) - if agent.voice_ai_integration_id and agent.voice_ai_agent_id and not has_provider_prompt: - try: - from app.services.voice_providers.prompt_sync import sync_provider_prompt - integration = db.query(Integration).filter(Integration.id == agent.voice_ai_integration_id).first() - if integration: - sync_provider_prompt(db_agent, integration, db) - db.refresh(db_agent) - except Exception as e: - logger.warning(f"[Agents] Best-effort provider prompt sync failed on create: {e}") - - return db_agent - - -@router.get("", response_model=List[AgentResponse]) -async def list_agents( - skip: int = 0, - limit: int = 100, - organization_id: UUID = Depends(get_organization_id), - workspace_id: UUID = Depends(get_workspace_id), - db: Session = Depends(get_db) -): - """Get list of all agents for the active workspace. - - Scoped to (organization_id, workspace_id) so users only see agents - in the workspace they're currently in. - """ - agents = db.query(Agent).filter( - Agent.organization_id == organization_id, - Agent.workspace_id == workspace_id, - ).offset(skip).limit(limit).all() - return agents - - -@router.get("/check-phone-assignment", response_model=AgentPhoneAssignmentCheckResponse) -async def check_phone_assignment( - phone_number: Optional[str] = Query(None), - telephony_phone_number_id: Optional[UUID] = Query(None), - exclude_agent_id: Optional[UUID] = Query(None), - organization_id: UUID = Depends(get_organization_id), - db: Session = Depends(get_db), -): - """Check whether a phone number is available for agent assignment in this org.""" - if not phone_number and not telephony_phone_number_id: - raise HTTPException( - status_code=status.HTTP_400_BAD_REQUEST, - detail="phone_number or telephony_phone_number_id is required", - ) - - from app.services.telephony.phone_routing import find_agent_phone_assignment_conflict - - conflict = find_agent_phone_assignment_conflict( - db, - organization_id=organization_id, - phone_number=phone_number, - telephony_phone_number_id=telephony_phone_number_id, - exclude_agent_id=exclude_agent_id, - ) - if conflict and conflict.get("error") == "telephony_not_found": - raise HTTPException(status_code=404, detail="Telephony phone number not found") - if conflict: - return AgentPhoneAssignmentCheckResponse( - available=False, - phone_number=conflict["phone_number"], - conflict=AgentPhoneAssignmentConflict(**conflict), - ) - - resolved_phone = phone_number - if telephony_phone_number_id: - from app.models.database import TelephonyPhoneNumber - - row = ( - db.query(TelephonyPhoneNumber) - .filter( - TelephonyPhoneNumber.id == telephony_phone_number_id, - TelephonyPhoneNumber.organization_id == organization_id, - ) - .first() - ) - if row: - resolved_phone = row.phone_number - - return AgentPhoneAssignmentCheckResponse( - available=True, - phone_number=resolved_phone, - ) - - -@router.get("/{agent_id}", response_model=AgentResponse) -async def get_agent( - agent_id: str, - organization_id: UUID = Depends(get_organization_id), - workspace_id: UUID = Depends(get_workspace_id), - db: Session = Depends(get_db) -): - """Get a specific agent by ID (UUID) or agent_id (6-digit) within the active workspace.""" - try: - # Try as UUID first - agent_uuid = UUID(agent_id) - agent = db.query(Agent).filter( - and_( - Agent.id == agent_uuid, - Agent.organization_id == organization_id, - Agent.workspace_id == workspace_id, - ) - ).first() - except ValueError: - # Try as 6-digit ID - agent = db.query(Agent).filter( - and_( - Agent.agent_id == agent_id, - Agent.organization_id == organization_id, - Agent.workspace_id == workspace_id, - ) - ).first() - - if not agent: - raise HTTPException(status_code=404, detail=f"Agent {agent_id} not found") - return agent - - -@router.put("/{agent_id}", response_model=AgentResponse) -async def update_agent( - agent_id: str, - agent_update: AgentUpdate, - organization_id: UUID = Depends(get_organization_id), - workspace_id: UUID = Depends(get_workspace_id), - db: Session = Depends(get_db) -): - """Update an existing agent by ID (UUID) or agent_id (6-digit) within the active workspace.""" - try: - # Try as UUID first - agent_uuid = UUID(agent_id) - db_agent = db.query(Agent).filter( - and_( - Agent.id == agent_uuid, - Agent.organization_id == organization_id, - Agent.workspace_id == workspace_id, - ) - ).first() - except ValueError: - # Try as 6-digit ID - db_agent = db.query(Agent).filter( - and_( - Agent.agent_id == agent_id, - Agent.organization_id == organization_id, - Agent.workspace_id == workspace_id, - ) - ).first() - - if not db_agent: - raise HTTPException(status_code=404, detail=f"Agent {agent_id} not found") - - # Determine the call_medium to validate - call_medium = agent_update.call_medium if agent_update.call_medium is not None else db_agent.call_medium - - # Validate phone_number is provided when call_medium is phone_call - if call_medium == CallMediumEnum.PHONE_CALL: - phone_number = agent_update.phone_number if agent_update.phone_number is not None else db_agent.phone_number - if not phone_number: - raise HTTPException( - status_code=status.HTTP_400_BAD_REQUEST, - detail="phone_number is required when call_medium is phone_call" - ) - - # Validate voice_bundle_id if provided - if agent_update.voice_bundle_id: - voice_bundle = db.query(VoiceBundle).filter( - and_( - VoiceBundle.id == agent_update.voice_bundle_id, - VoiceBundle.organization_id == organization_id - ) - ).first() - if not voice_bundle: - raise HTTPException(status_code=404, detail="Voice bundle not found") - - # Validate voice_ai_integration_id if provided - if agent_update.voice_ai_integration_id: - integration = db.query(Integration).filter( - and_( - Integration.id == agent_update.voice_ai_integration_id, - Integration.organization_id == organization_id, - Integration.is_active == True - ) - ).first() - if not integration: - raise HTTPException(status_code=404, detail="Integration not found or inactive") - - if integration.platform not in [ - IntegrationPlatform.RETELL, - IntegrationPlatform.VAPI, - IntegrationPlatform.ELEVENLABS, - IntegrationPlatform.SMALLEST, - ]: - raise HTTPException( - status_code=status.HTTP_400_BAD_REQUEST, - detail=( - f"Integration platform {integration.platform.value} is not supported for Voice AI agents. " - "Only Retell, Vapi, ElevenLabs, and Smallest are supported." - ) - ) - - if not agent_update.voice_ai_agent_id: - raise HTTPException( - status_code=status.HTTP_400_BAD_REQUEST, - detail="voice_ai_agent_id is required when voice_ai_integration_id is provided" - ) - - update_data = agent_update.model_dump(exclude_unset=True, exclude_none=False) - - effective_call_medium = ( - agent_update.call_medium if agent_update.call_medium is not None else db_agent.call_medium - ) - effective_phone_number = ( - agent_update.phone_number - if "phone_number" in update_data - else db_agent.phone_number - ) - effective_telephony_id = ( - agent_update.telephony_phone_number_id - if "telephony_phone_number_id" in update_data - else db_agent.telephony_phone_number_id - ) - - _validate_agent_phone_assignment( - db, - organization_id=organization_id, - call_medium=effective_call_medium, - phone_number=effective_phone_number, - telephony_phone_number_id=effective_telephony_id, - exclude_agent_id=db_agent.id, - ) - - # Convert the update model to dict, handling None values properly - # Use model_dump with exclude_unset to only get fields that were explicitly provided - - # Apply updates - for field, value in update_data.items(): - setattr(db_agent, field, value) - - db.commit() - db.refresh(db_agent) - - from app.services.telephony.phone_routing import sync_agent_telephony_number_link - - if "telephony_phone_number_id" in update_data or "phone_number" in update_data: - sync_agent_telephony_number_link(db, db_agent) - db.refresh(db_agent) - - if "voice_ai_agent_id" in update_data or "voice_ai_integration_id" in update_data: - integration_id = db_agent.voice_ai_integration_id - provider_prompt_updated = "provider_prompt" in update_data - if integration_id and db_agent.voice_ai_agent_id and not provider_prompt_updated: - try: - from app.services.voice_providers.prompt_sync import sync_provider_prompt - integration = db.query(Integration).filter(Integration.id == integration_id).first() - if integration: - sync_provider_prompt(db_agent, integration, db) - db.refresh(db_agent) - except Exception as e: - logger.warning(f"[Agents] Best-effort provider prompt sync failed on update: {e}") - - return db_agent - - -@router.post("/{agent_id}/sync-provider-prompt") -async def sync_agent_provider_prompt( - agent_id: str, - organization_id: UUID = Depends(get_organization_id), - workspace_id: UUID = Depends(get_workspace_id), - api_key: str = Depends(get_api_key), - db: Session = Depends(get_db), -): - """Fetch and store the current system prompt from the voice provider.""" - try: - agent_uuid = UUID(agent_id) - db_agent = db.query(Agent).filter( - and_( - Agent.id == agent_uuid, - Agent.organization_id == organization_id, - Agent.workspace_id == workspace_id, - ) - ).first() - except ValueError: - db_agent = db.query(Agent).filter( - and_( - Agent.agent_id == agent_id, - Agent.organization_id == organization_id, - Agent.workspace_id == workspace_id, - ) - ).first() - - if not db_agent: - raise HTTPException(status_code=404, detail=f"Agent {agent_id} not found") - - if not db_agent.voice_ai_integration_id or not db_agent.voice_ai_agent_id: - raise HTTPException( - status_code=400, - detail="Agent is not linked to an external voice provider", - ) - - integration = db.query(Integration).filter( - and_( - Integration.id == db_agent.voice_ai_integration_id, - Integration.organization_id == organization_id, - Integration.is_active == True, - ) - ).first() - if not integration: - raise HTTPException(status_code=404, detail="Integration not found or inactive") - - try: - from app.services.voice_providers.prompt_sync import sync_provider_prompt - prompt = sync_provider_prompt(db_agent, integration, db) - except Exception as e: - raise HTTPException(status_code=502, detail=f"Failed to fetch prompt from provider: {str(e)}") - - db.refresh(db_agent) - synced = isinstance(prompt, str) and bool(prompt.strip()) - if not synced: - raise HTTPException( - status_code=422, - detail="Provider returned no prompt. Verify the external agent has a system prompt configured.", - ) - return { - "synced": synced, - "provider_prompt": db_agent.provider_prompt, - "provider_prompt_synced_at": db_agent.provider_prompt_synced_at.isoformat() if db_agent.provider_prompt_synced_at else None, - } - - -@router.delete("/{agent_id}") -async def delete_agent( - agent_id: str, - force: bool = Query(False, description="Force delete with all dependent records"), - organization_id: UUID = Depends(get_organization_id), - workspace_id: UUID = Depends(get_workspace_id), - db: Session = Depends(get_db) -): - """Delete an agent (scoped to the active workspace). Returns 409 if dependent records exist unless force=true.""" - try: - agent_uuid = UUID(agent_id) - db_agent = db.query(Agent).filter( - and_( - Agent.id == agent_uuid, - Agent.organization_id == organization_id, - Agent.workspace_id == workspace_id, - ) - ).first() - except ValueError: - db_agent = db.query(Agent).filter( - and_( - Agent.agent_id == agent_id, - Agent.organization_id == organization_id, - Agent.workspace_id == workspace_id, - ) - ).first() - - if not db_agent: - raise HTTPException(status_code=404, detail=f"Agent {agent_id} not found") - - agent_uuid = db_agent.id - dependencies = get_agent_dependencies(db, organization_id, agent_uuid) - - if dependencies and not force: - parts = [] - if dependencies.get("evaluators"): - parts.append(f"{dependencies['evaluators']} evaluator(s)") - if dependencies.get("evaluator_results"): - parts.append(f"{dependencies['evaluator_results']} evaluator result(s)") - if dependencies.get("call_recordings"): - parts.append(f"{dependencies['call_recordings']} call recording(s)") - if dependencies.get("conversation_evaluations"): - parts.append(f"{dependencies['conversation_evaluations']} conversation evaluation(s)") - if dependencies.get("test_conversations"): - parts.append(f"{dependencies['test_conversations']} test conversation(s)") - - raise HTTPException( - status_code=status.HTTP_409_CONFLICT, - detail={ - "message": f"Cannot delete agent. It is referenced by: {', '.join(parts)}.", - "dependencies": dependencies, - "hint": "Use force=true to delete this agent and all its dependent records.", - }, - ) - - if dependencies: - # Delete in FK-safe order: - # 1. EvaluatorResults (references evaluators and agents) - db.query(EvaluatorResult).filter( - EvaluatorResult.agent_id == agent_uuid, - ).delete(synchronize_session=False) - - # 2. Evaluators (references agents) - db.query(Evaluator).filter( - Evaluator.agent_id == agent_uuid, - ).delete(synchronize_session=False) - - # 3. Nullify call recordings (keep recordings, unlink agent) - db.query(CallRecording).filter( - CallRecording.agent_id == agent_uuid, - ).update({CallRecording.agent_id: None}, synchronize_session=False) - - # 4. ConversationEvaluations - db.query(ConversationEvaluation).filter( - ConversationEvaluation.agent_id == agent_uuid, - ).delete(synchronize_session=False) - - # 5. TestAgentConversations - db.query(TestAgentConversation).filter( - TestAgentConversation.agent_id == agent_uuid, - ).delete(synchronize_session=False) - - db.delete(db_agent) - db.commit() - - if dependencies: - return JSONResponse( - status_code=200, - content={ - "message": "Agent and all dependent records deleted successfully.", - "deleted": dependencies, - }, - ) - - return Response(status_code=204) - - -@router.get("/{agent_id}/delete-impact") -async def get_agent_delete_impact( - agent_id: str, - organization_id: UUID = Depends(get_organization_id), - workspace_id: UUID = Depends(get_workspace_id), - db: Session = Depends(get_db) -): - """Preview dependent records that would be affected by force delete (scoped to the active workspace).""" - try: - agent_uuid = UUID(agent_id) - db_agent = db.query(Agent).filter( - and_( - Agent.id == agent_uuid, - Agent.organization_id == organization_id, - Agent.workspace_id == workspace_id, - ) - ).first() - except ValueError: - db_agent = db.query(Agent).filter( - and_( - Agent.agent_id == agent_id, - Agent.organization_id == organization_id, - Agent.workspace_id == workspace_id, - ) - ).first() - - if not db_agent: - raise HTTPException(status_code=404, detail=f"Agent {agent_id} not found") - - dependencies = get_agent_dependencies(db, organization_id, db_agent.id) - return { - "agent_id": str(db_agent.id), - "agent_name": db_agent.name, - "dependencies": dependencies, - "can_delete_without_force": len(dependencies) == 0, - } - - -from app.core.auth.capabilities import SIM_MANAGE, SIM_VIEW -from app.core.auth.workspace_route_capabilities import apply_workspace_route_capabilities - -apply_workspace_route_capabilities( - router, - view_capability=SIM_VIEW, - manage_capability=SIM_MANAGE, -) - +""" +Agents API Routes +Complete CRUD operations for test agents +""" +from fastapi import APIRouter, Depends, HTTPException, status, Query +from fastapi.responses import JSONResponse, Response +from sqlalchemy.orm import Session +from typing import List, Optional +from uuid import UUID +import random +from pydantic import BaseModel +from loguru import logger + +from app.dependencies import get_db, get_organization_id, get_workspace_id, get_api_key +from app.models.database import ( + Agent, ConversationEvaluation, TestAgentConversation, VoiceBundle, + AIProvider, Integration, IntegrationPlatform, CallMediumEnum, + Evaluator, EvaluatorResult, CallRecording, Scenario, +) +from sqlalchemy import and_ +from app.models.schemas import ( + AgentCreate, + AgentUpdate, + AgentResponse, + AgentPhoneAssignmentCheckResponse, + AgentPhoneAssignmentConflict, + CallMediumEnum as CallMediumEnumSchema, + GenerateTestPromptRequest, + GenerateTestPromptResponse, + GenerateScenariosFromPromptRequest, + GenerateScenariosFromPromptResponse, + GenerateTestSetupRequest, + GenerateTestSetupResponse, + GeneratedScenarioDraftResponse, + TestPromptSectionResponse, +) + +router = APIRouter(prefix="/agents", tags=["agents"]) + + +def _validate_agent_phone_assignment( + db: Session, + *, + organization_id: UUID, + call_medium, + phone_number: Optional[str], + telephony_phone_number_id: Optional[UUID], + exclude_agent_id: Optional[UUID] = None, +) -> None: + """Raise HTTPException if phone assignment conflicts with another agent.""" + if call_medium != CallMediumEnum.PHONE_CALL: + return + if not phone_number and not telephony_phone_number_id: + return + + from app.services.telephony.phone_routing import find_agent_phone_assignment_conflict + + conflict = find_agent_phone_assignment_conflict( + db, + organization_id=organization_id, + phone_number=phone_number, + telephony_phone_number_id=telephony_phone_number_id, + exclude_agent_id=exclude_agent_id, + ) + if not conflict: + return + if conflict.get("error") == "telephony_not_found": + raise HTTPException(status_code=404, detail="Telephony phone number not found") + + agent_name = conflict["agent_name"] + raise HTTPException( + status_code=status.HTTP_409_CONFLICT, + detail={ + "message": f'This number is already assigned to agent "{agent_name}".', + "agent_id": str(conflict["agent_id"]), + "agent_name": agent_name, + "phone_number": conflict["phone_number"], + }, + ) + + +def resolve_agent_by_path_id( + db: Session, + *, + organization_id: UUID, + workspace_id: UUID, + agent_id: str, +) -> Agent: + """Resolve agent by UUID primary key or 6-digit agent_id (same as GET /agents/{id}).""" + try: + agent_uuid = UUID(agent_id) + agent = db.query(Agent).filter( + and_( + Agent.id == agent_uuid, + Agent.organization_id == organization_id, + Agent.workspace_id == workspace_id, + ) + ).first() + except ValueError: + agent = db.query(Agent).filter( + and_( + Agent.agent_id == agent_id, + Agent.organization_id == organization_id, + Agent.workspace_id == workspace_id, + ) + ).first() + if not agent: + raise HTTPException(status_code=404, detail=f"Agent {agent_id} not found") + return agent + + +# ====================================================================== +# AI Generation for agent descriptions +# ====================================================================== + +class GenerateAgentDescriptionRequest(BaseModel): + description: str + tone: Optional[str] = "professional" + format_style: Optional[str] = "structured" + provider: Optional[str] = None + model: Optional[str] = None + agent_id: Optional[UUID] = None + include_linked_scenarios: bool = True + append_scenarios_to_output: bool = False + + +GENERATE_AGENT_DESCRIPTION_SYSTEM = ( + "You are an expert at writing clear, well-structured descriptions for voice AI test agents. " + "The user will describe what they need the agent to do, and you will generate a comprehensive, " + "well-formatted agent description in markdown.\n\n" + "Guidelines:\n" + "- Use clear markdown structure: headings, bullet points, numbered lists\n" + "- Include sections for: Purpose, Behavior, Expected Interactions, Personality Traits, and Constraints\n" + "- Be specific about the agent's role, tone of voice, and how it should handle conversations\n" + "- Include example scenarios or edge cases where helpful\n" + "- Return ONLY the description in markdown, no preamble or explanation about what you did" +) + + +from app.services.ai.llm_resolver import get_llm_provider_and_model as _get_llm_provider_and_model + + +@router.post("/generate-description") +async def generate_agent_description( + data: GenerateAgentDescriptionRequest, + organization_id: UUID = Depends(get_organization_id), + workspace_id: UUID = Depends(get_workspace_id), + api_key: str = Depends(get_api_key), + db: Session = Depends(get_db), +): + """Generate an agent description using AI from a brief description.""" + from contextlib import nullcontext + + from app.services.ai.llm_service import llm_service + from app.services.usage.context import llm_usage_context, usage_context_for_agent + from app.services.testing.test_agent_simulation_prompt import ( + format_scenarios_for_generation_context, + format_scenarios_reference_appendix, + load_linked_scenarios_for_agent, + merge_generated_description_with_scenario_appendix, + ) + + if not data.description.strip(): + raise HTTPException(400, "Description is required") + + provider_enum, model_str = _get_llm_provider_and_model( + organization_id, db, data.provider, data.model + ) + + linked_scenarios = [] + if data.agent_id and data.include_linked_scenarios: + agent = db.query(Agent).filter( + Agent.id == data.agent_id, + Agent.organization_id == organization_id, + Agent.workspace_id == workspace_id, + ).first() + if not agent: + raise HTTPException(status_code=404, detail=f"Agent {data.agent_id} not found") + linked_scenarios = load_linked_scenarios_for_agent( + db, + organization_id=organization_id, + workspace_id=workspace_id, + agent_id=agent.id, + ) + + user_prompt_parts = [ + "Create a detailed agent description for the following:", + "", + f"Description: {data.description}", + f"Tone: {data.tone or 'professional'}", + f"Format: {data.format_style or 'structured'}", + ] + scenario_context = format_scenarios_for_generation_context(linked_scenarios) + if scenario_context: + user_prompt_parts.extend(["", scenario_context]) + user_prompt_parts.extend([ + "", + "Generate a comprehensive, well-formatted agent description in markdown.", + ]) + user_prompt = "\n".join(user_prompt_parts) + + messages = [ + {"role": "system", "content": GENERATE_AGENT_DESCRIPTION_SYSTEM}, + {"role": "user", "content": user_prompt}, + ] + + try: + usage_ctx = nullcontext() + if data.agent_id: + agent_for_usage = db.query(Agent).filter( + Agent.id == data.agent_id, + Agent.organization_id == organization_id, + Agent.workspace_id == workspace_id, + ).first() + if agent_for_usage: + usage_ctx = llm_usage_context( + usage_context_for_agent(agent_for_usage, workspace_id=workspace_id) + ) + + with usage_ctx: + result = llm_service.generate_response( + messages=messages, + llm_provider=provider_enum, + llm_model=model_str, + organization_id=organization_id, + db=db, + temperature=0.7, + max_tokens=4000, + ) + content = result["text"] + if data.append_scenarios_to_output and linked_scenarios: + appendix = format_scenarios_reference_appendix(linked_scenarios) + content = merge_generated_description_with_scenario_appendix(content, appendix) + return {"content": content, "provider": provider_enum.value, "model": model_str} + except Exception as e: + logger.error(f"[Agents] AI description generation failed: {repr(e)}") + raise HTTPException(500, f"AI generation failed: {str(e)}") + + +def _test_prompt_section_responses(sections) -> list[TestPromptSectionResponse]: + return [ + TestPromptSectionResponse(key=s.key, title=s.title, content=s.content) + for s in sections + ] + + +def _scenario_draft_responses(scenarios) -> list[GeneratedScenarioDraftResponse]: + return [ + GeneratedScenarioDraftResponse(name=s.name, description=s.description, goal=s.goal) + for s in scenarios + ] + + +@router.post("/generate-test-prompt", response_model=GenerateTestPromptResponse) +async def generate_test_prompt( + data: GenerateTestPromptRequest, + organization_id: UUID = Depends(get_organization_id), + workspace_id: UUID = Depends(get_workspace_id), + api_key: str = Depends(get_api_key), + db: Session = Depends(get_db), +): + """Stage 1: generate foundational test agent prompt from production prompt.""" + from app.services.testing.agent_test_setup_generation import ( + generate_test_prompt_from_production, + ) + from app.services.usage.context import ( + LLMUsageContext, + LLMUsageProductSection, + llm_usage_context, + ) + + if not data.production_prompt.strip(): + raise HTTPException(400, "Production prompt is required") + + provider_enum, model_str = _get_llm_provider_and_model( + organization_id, db, data.provider, data.model, data.credential_id + ) + + try: + with llm_usage_context( + LLMUsageContext( + organization_id=organization_id, + workspace_id=workspace_id, + product_section=LLMUsageProductSection.AGENTS, + ) + ): + result = generate_test_prompt_from_production( + data.production_prompt, + agent_name=data.agent_name, + language=data.language, + call_type=data.call_type, + additional_context=data.additional_context, + llm_provider=provider_enum, + llm_model=model_str, + organization_id=organization_id, + db=db, + llm_config=data.llm_config, + credential_id=data.credential_id, + ) + return GenerateTestPromptResponse( + sections=_test_prompt_section_responses(result.sections), + test_agent_prompt=result.test_agent_prompt, + provider=result.provider, + model=result.model, + ) + except ValueError as e: + raise HTTPException(400, str(e)) from e + except Exception as e: + logger.error(f"[Agents] Test prompt generation failed: {repr(e)}") + raise HTTPException(500, f"AI generation failed: {str(e)}") from e + + +@router.post("/generate-scenarios-from-prompt", response_model=GenerateScenariosFromPromptResponse) +async def generate_scenarios_from_prompt( + data: GenerateScenariosFromPromptRequest, + organization_id: UUID = Depends(get_organization_id), + workspace_id: UUID = Depends(get_workspace_id), + api_key: str = Depends(get_api_key), + db: Session = Depends(get_db), +): + """Stage 2: generate scenario drafts from a test agent prompt.""" + from app.services.testing.agent_test_setup_generation import ( + generate_scenarios_from_test_prompt, + ) + from app.services.usage.context import ( + LLMUsageContext, + LLMUsageProductSection, + llm_usage_context, + ) + + if not data.test_agent_prompt.strip(): + raise HTTPException(400, "Test agent prompt is required") + + provider_enum, model_str = _get_llm_provider_and_model( + organization_id, db, data.provider, data.model, data.credential_id + ) + + try: + with llm_usage_context( + LLMUsageContext( + organization_id=organization_id, + workspace_id=workspace_id, + product_section=LLMUsageProductSection.AGENTS, + ) + ): + result = generate_scenarios_from_test_prompt( + data.test_agent_prompt, + agent_name=data.agent_name, + scenario_count=data.scenario_count, + language=data.language, + call_type=data.call_type, + additional_context=data.additional_context, + llm_provider=provider_enum, + llm_model=model_str, + organization_id=organization_id, + db=db, + llm_config=data.llm_config, + credential_id=data.credential_id, + ) + return GenerateScenariosFromPromptResponse( + scenarios=_scenario_draft_responses(result.scenarios), + provider=result.provider, + model=result.model, + ) + except ValueError as e: + raise HTTPException(400, str(e)) from e + except Exception as e: + logger.error(f"[Agents] Scenario generation failed: {repr(e)}") + raise HTTPException(500, f"AI generation failed: {str(e)}") from e + + +@router.post("/generate-test-setup", response_model=GenerateTestSetupResponse) +async def generate_test_setup( + data: GenerateTestSetupRequest, + organization_id: UUID = Depends(get_organization_id), + workspace_id: UUID = Depends(get_workspace_id), + api_key: str = Depends(get_api_key), + db: Session = Depends(get_db), +): + """Run stage 1 then stage 2: foundational test prompt + scenario drafts.""" + from app.services.testing.agent_test_setup_generation import ( + generate_scenarios_from_test_prompt, + generate_test_prompt_from_production, + ) + from app.services.usage.context import ( + LLMUsageContext, + LLMUsageProductSection, + llm_usage_context, + ) + + if not data.production_prompt.strip(): + raise HTTPException(400, "Production prompt is required") + + provider_enum, model_str = _get_llm_provider_and_model( + organization_id, db, data.provider, data.model, data.credential_id + ) + + try: + with llm_usage_context( + LLMUsageContext( + organization_id=organization_id, + workspace_id=workspace_id, + product_section=LLMUsageProductSection.AGENTS, + ) + ): + prompt_result = generate_test_prompt_from_production( + data.production_prompt, + agent_name=data.agent_name, + language=data.language, + call_type=data.call_type, + additional_context=data.additional_context, + llm_provider=provider_enum, + llm_model=model_str, + organization_id=organization_id, + db=db, + llm_config=data.llm_config, + credential_id=data.credential_id, + ) + scenario_result = generate_scenarios_from_test_prompt( + prompt_result.test_agent_prompt, + agent_name=data.agent_name, + scenario_count=data.scenario_count, + language=data.language, + call_type=data.call_type, + additional_context=data.additional_context, + llm_provider=provider_enum, + llm_model=model_str, + organization_id=organization_id, + db=db, + llm_config=data.llm_config, + credential_id=data.credential_id, + ) + return GenerateTestSetupResponse( + sections=_test_prompt_section_responses(prompt_result.sections), + test_agent_prompt=prompt_result.test_agent_prompt, + scenarios=_scenario_draft_responses(scenario_result.scenarios), + provider=prompt_result.provider, + model=prompt_result.model, + ) + except ValueError as e: + raise HTTPException(400, str(e)) from e + except Exception as e: + logger.error(f"[Agents] Test setup generation failed: {repr(e)}") + raise HTTPException(500, f"AI generation failed: {str(e)}") from e + + +def generate_unique_agent_id(db: Session) -> str: + """Generate a unique 6-digit agent ID.""" + max_attempts = 100 + for _ in range(max_attempts): + agent_id = f"{random.randint(100000, 999999)}" + existing = db.query(Agent).filter(Agent.agent_id == agent_id).first() + if not existing: + return agent_id + raise HTTPException( + status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, + detail="Failed to generate unique agent ID" + ) + + +def get_agent_dependencies(db: Session, organization_id: UUID, agent_uuid: UUID) -> dict: + """Return dependency counts that block non-force delete.""" + evaluators_count = db.query(Evaluator).filter( + Evaluator.agent_id == agent_uuid, + Evaluator.organization_id == organization_id, + ).count() + + evaluator_results_count = db.query(EvaluatorResult).filter( + EvaluatorResult.agent_id == agent_uuid, + EvaluatorResult.organization_id == organization_id, + ).count() + + call_recordings_count = db.query(CallRecording).filter( + CallRecording.agent_id == agent_uuid, + CallRecording.organization_id == organization_id, + ).count() + + conversation_evaluations_count = db.query(ConversationEvaluation).filter( + ConversationEvaluation.agent_id == agent_uuid, + ConversationEvaluation.organization_id == organization_id, + ).count() + + test_conversations_count = db.query(TestAgentConversation).filter( + TestAgentConversation.agent_id == agent_uuid, + TestAgentConversation.organization_id == organization_id, + ).count() + + dependencies = {} + if evaluators_count > 0: + dependencies["evaluators"] = evaluators_count + if evaluator_results_count > 0: + dependencies["evaluator_results"] = evaluator_results_count + if call_recordings_count > 0: + dependencies["call_recordings"] = call_recordings_count + if conversation_evaluations_count > 0: + dependencies["conversation_evaluations"] = conversation_evaluations_count + if test_conversations_count > 0: + dependencies["test_conversations"] = test_conversations_count + + return dependencies + + +@router.post("", response_model=AgentResponse, status_code=status.HTTP_201_CREATED) +async def create_agent( + agent: AgentCreate, + organization_id: UUID = Depends(get_organization_id), + workspace_id: UUID = Depends(get_workspace_id), + db: Session = Depends(get_db) +): + """Create a new test agent. + + The agent is stamped with the active workspace from the + ``X-Workspace-Id`` header (falling back to the org's Default). + """ + # Validate phone_number is provided when call_medium is phone_call + if agent.call_medium == CallMediumEnumSchema.PHONE_CALL and not agent.phone_number: + raise HTTPException( + status_code=status.HTTP_400_BAD_REQUEST, + detail="phone_number is required when call_medium is phone_call" + ) + + # Validate voice_bundle_id exists, is active, and belongs to organization + voice_bundle = db.query(VoiceBundle).filter( + and_( + VoiceBundle.id == agent.voice_bundle_id, + VoiceBundle.organization_id == organization_id, + VoiceBundle.is_active == True, + ) + ).first() + if not voice_bundle: + raise HTTPException( + status_code=status.HTTP_400_BAD_REQUEST, + detail="Active voice bundle not found", + ) + + # Validate voice_ai_integration_id exists and belongs to organization + if agent.voice_ai_integration_id: + integration = db.query(Integration).filter( + and_( + Integration.id == agent.voice_ai_integration_id, + Integration.organization_id == organization_id, + Integration.is_active == True + ) + ).first() + if not integration: + raise HTTPException(status_code=404, detail="Integration not found or inactive") + + if integration.platform not in [ + IntegrationPlatform.RETELL, + IntegrationPlatform.VAPI, + IntegrationPlatform.ELEVENLABS, + IntegrationPlatform.SMALLEST, + ]: + raise HTTPException( + status_code=status.HTTP_400_BAD_REQUEST, + detail=( + f"Integration platform {integration.platform.value} is not supported for Voice AI agents. " + "Only Retell, Vapi, ElevenLabs, and Smallest are supported." + ) + ) + + if not agent.voice_ai_agent_id: + raise HTTPException( + status_code=status.HTTP_400_BAD_REQUEST, + detail="voice_ai_agent_id is required when voice_ai_integration_id is provided" + ) + + _validate_agent_phone_assignment( + db, + organization_id=organization_id, + call_medium=agent.call_medium, + phone_number=agent.phone_number, + telephony_phone_number_id=agent.telephony_phone_number_id, + ) + + # Generate unique 6-digit agent_id + agent_id = generate_unique_agent_id(db) + + db_agent = Agent( + agent_id=agent_id, + organization_id=organization_id, + workspace_id=workspace_id, + name=agent.name, + phone_number=agent.phone_number, + language=agent.language, + description=agent.description, + call_type=agent.call_type, + call_medium=agent.call_medium, + telephony_phone_number_id=agent.telephony_phone_number_id, + voice_bundle_id=agent.voice_bundle_id, + ai_provider_id=agent.ai_provider_id, + voice_ai_integration_id=agent.voice_ai_integration_id, + voice_ai_agent_id=agent.voice_ai_agent_id, + provider_prompt=agent.provider_prompt, + silence_hangup_secs=agent.silence_hangup_secs, + ) + db.add(db_agent) + db.commit() + db.refresh(db_agent) + + from app.services.telephony.phone_routing import sync_agent_telephony_number_link + + sync_agent_telephony_number_link(db, db_agent) + db.refresh(db_agent) + + has_provider_prompt = isinstance(agent.provider_prompt, str) and bool(agent.provider_prompt.strip()) + if agent.voice_ai_integration_id and agent.voice_ai_agent_id and not has_provider_prompt: + try: + from app.services.voice_providers.prompt_sync import sync_provider_prompt + integration = db.query(Integration).filter(Integration.id == agent.voice_ai_integration_id).first() + if integration: + sync_provider_prompt(db_agent, integration, db) + db.refresh(db_agent) + except Exception as e: + logger.warning(f"[Agents] Best-effort provider prompt sync failed on create: {e}") + + return db_agent + + +@router.get("", response_model=List[AgentResponse]) +async def list_agents( + skip: int = 0, + limit: int = 100, + organization_id: UUID = Depends(get_organization_id), + workspace_id: UUID = Depends(get_workspace_id), + db: Session = Depends(get_db) +): + """Get list of all agents for the active workspace. + + Scoped to (organization_id, workspace_id) so users only see agents + in the workspace they're currently in. + """ + agents = db.query(Agent).filter( + Agent.organization_id == organization_id, + Agent.workspace_id == workspace_id, + ).offset(skip).limit(limit).all() + return agents + + +@router.get("/check-phone-assignment", response_model=AgentPhoneAssignmentCheckResponse) +async def check_phone_assignment( + phone_number: Optional[str] = Query(None), + telephony_phone_number_id: Optional[UUID] = Query(None), + exclude_agent_id: Optional[UUID] = Query(None), + organization_id: UUID = Depends(get_organization_id), + db: Session = Depends(get_db), +): + """Check whether a phone number is available for agent assignment in this org.""" + if not phone_number and not telephony_phone_number_id: + raise HTTPException( + status_code=status.HTTP_400_BAD_REQUEST, + detail="phone_number or telephony_phone_number_id is required", + ) + + from app.services.telephony.phone_routing import find_agent_phone_assignment_conflict + + conflict = find_agent_phone_assignment_conflict( + db, + organization_id=organization_id, + phone_number=phone_number, + telephony_phone_number_id=telephony_phone_number_id, + exclude_agent_id=exclude_agent_id, + ) + if conflict and conflict.get("error") == "telephony_not_found": + raise HTTPException(status_code=404, detail="Telephony phone number not found") + if conflict: + return AgentPhoneAssignmentCheckResponse( + available=False, + phone_number=conflict["phone_number"], + conflict=AgentPhoneAssignmentConflict(**conflict), + ) + + resolved_phone = phone_number + if telephony_phone_number_id: + from app.models.database import TelephonyPhoneNumber + + row = ( + db.query(TelephonyPhoneNumber) + .filter( + TelephonyPhoneNumber.id == telephony_phone_number_id, + TelephonyPhoneNumber.organization_id == organization_id, + ) + .first() + ) + if row: + resolved_phone = row.phone_number + + return AgentPhoneAssignmentCheckResponse( + available=True, + phone_number=resolved_phone, + ) + + +@router.get("/{agent_id}", response_model=AgentResponse) +async def get_agent( + agent_id: str, + organization_id: UUID = Depends(get_organization_id), + workspace_id: UUID = Depends(get_workspace_id), + db: Session = Depends(get_db) +): + """Get a specific agent by ID (UUID) or agent_id (6-digit) within the active workspace.""" + try: + # Try as UUID first + agent_uuid = UUID(agent_id) + agent = db.query(Agent).filter( + and_( + Agent.id == agent_uuid, + Agent.organization_id == organization_id, + Agent.workspace_id == workspace_id, + ) + ).first() + except ValueError: + # Try as 6-digit ID + agent = db.query(Agent).filter( + and_( + Agent.agent_id == agent_id, + Agent.organization_id == organization_id, + Agent.workspace_id == workspace_id, + ) + ).first() + + if not agent: + raise HTTPException(status_code=404, detail=f"Agent {agent_id} not found") + return agent + + +@router.put("/{agent_id}", response_model=AgentResponse) +async def update_agent( + agent_id: str, + agent_update: AgentUpdate, + organization_id: UUID = Depends(get_organization_id), + workspace_id: UUID = Depends(get_workspace_id), + db: Session = Depends(get_db) +): + """Update an existing agent by ID (UUID) or agent_id (6-digit) within the active workspace.""" + try: + # Try as UUID first + agent_uuid = UUID(agent_id) + db_agent = db.query(Agent).filter( + and_( + Agent.id == agent_uuid, + Agent.organization_id == organization_id, + Agent.workspace_id == workspace_id, + ) + ).first() + except ValueError: + # Try as 6-digit ID + db_agent = db.query(Agent).filter( + and_( + Agent.agent_id == agent_id, + Agent.organization_id == organization_id, + Agent.workspace_id == workspace_id, + ) + ).first() + + if not db_agent: + raise HTTPException(status_code=404, detail=f"Agent {agent_id} not found") + + # Determine the call_medium to validate + call_medium = agent_update.call_medium if agent_update.call_medium is not None else db_agent.call_medium + + # Validate phone_number is provided when call_medium is phone_call + if call_medium == CallMediumEnum.PHONE_CALL: + phone_number = agent_update.phone_number if agent_update.phone_number is not None else db_agent.phone_number + if not phone_number: + raise HTTPException( + status_code=status.HTTP_400_BAD_REQUEST, + detail="phone_number is required when call_medium is phone_call" + ) + + # Validate voice_bundle_id if provided + if agent_update.voice_bundle_id: + voice_bundle = db.query(VoiceBundle).filter( + and_( + VoiceBundle.id == agent_update.voice_bundle_id, + VoiceBundle.organization_id == organization_id + ) + ).first() + if not voice_bundle: + raise HTTPException(status_code=404, detail="Voice bundle not found") + + # Validate voice_ai_integration_id if provided + if agent_update.voice_ai_integration_id: + integration = db.query(Integration).filter( + and_( + Integration.id == agent_update.voice_ai_integration_id, + Integration.organization_id == organization_id, + Integration.is_active == True + ) + ).first() + if not integration: + raise HTTPException(status_code=404, detail="Integration not found or inactive") + + if integration.platform not in [ + IntegrationPlatform.RETELL, + IntegrationPlatform.VAPI, + IntegrationPlatform.ELEVENLABS, + IntegrationPlatform.SMALLEST, + ]: + raise HTTPException( + status_code=status.HTTP_400_BAD_REQUEST, + detail=( + f"Integration platform {integration.platform.value} is not supported for Voice AI agents. " + "Only Retell, Vapi, ElevenLabs, and Smallest are supported." + ) + ) + + if not agent_update.voice_ai_agent_id: + raise HTTPException( + status_code=status.HTTP_400_BAD_REQUEST, + detail="voice_ai_agent_id is required when voice_ai_integration_id is provided" + ) + + update_data = agent_update.model_dump(exclude_unset=True, exclude_none=False) + + effective_call_medium = ( + agent_update.call_medium if agent_update.call_medium is not None else db_agent.call_medium + ) + effective_phone_number = ( + agent_update.phone_number + if "phone_number" in update_data + else db_agent.phone_number + ) + effective_telephony_id = ( + agent_update.telephony_phone_number_id + if "telephony_phone_number_id" in update_data + else db_agent.telephony_phone_number_id + ) + + _validate_agent_phone_assignment( + db, + organization_id=organization_id, + call_medium=effective_call_medium, + phone_number=effective_phone_number, + telephony_phone_number_id=effective_telephony_id, + exclude_agent_id=db_agent.id, + ) + + # Convert the update model to dict, handling None values properly + # Use model_dump with exclude_unset to only get fields that were explicitly provided + + # Apply updates + for field, value in update_data.items(): + setattr(db_agent, field, value) + + db.commit() + db.refresh(db_agent) + + from app.services.telephony.phone_routing import sync_agent_telephony_number_link + + if "telephony_phone_number_id" in update_data or "phone_number" in update_data: + sync_agent_telephony_number_link(db, db_agent) + db.refresh(db_agent) + + if "voice_ai_agent_id" in update_data or "voice_ai_integration_id" in update_data: + integration_id = db_agent.voice_ai_integration_id + provider_prompt_updated = "provider_prompt" in update_data + if integration_id and db_agent.voice_ai_agent_id and not provider_prompt_updated: + try: + from app.services.voice_providers.prompt_sync import sync_provider_prompt + integration = db.query(Integration).filter(Integration.id == integration_id).first() + if integration: + sync_provider_prompt(db_agent, integration, db) + db.refresh(db_agent) + except Exception as e: + logger.warning(f"[Agents] Best-effort provider prompt sync failed on update: {e}") + + return db_agent + + +@router.post("/{agent_id}/sync-provider-prompt") +async def sync_agent_provider_prompt( + agent_id: str, + organization_id: UUID = Depends(get_organization_id), + workspace_id: UUID = Depends(get_workspace_id), + api_key: str = Depends(get_api_key), + db: Session = Depends(get_db), +): + """Fetch and store the current system prompt from the voice provider.""" + try: + agent_uuid = UUID(agent_id) + db_agent = db.query(Agent).filter( + and_( + Agent.id == agent_uuid, + Agent.organization_id == organization_id, + Agent.workspace_id == workspace_id, + ) + ).first() + except ValueError: + db_agent = db.query(Agent).filter( + and_( + Agent.agent_id == agent_id, + Agent.organization_id == organization_id, + Agent.workspace_id == workspace_id, + ) + ).first() + + if not db_agent: + raise HTTPException(status_code=404, detail=f"Agent {agent_id} not found") + + if not db_agent.voice_ai_integration_id or not db_agent.voice_ai_agent_id: + raise HTTPException( + status_code=400, + detail="Agent is not linked to an external voice provider", + ) + + integration = db.query(Integration).filter( + and_( + Integration.id == db_agent.voice_ai_integration_id, + Integration.organization_id == organization_id, + Integration.is_active == True, + ) + ).first() + if not integration: + raise HTTPException(status_code=404, detail="Integration not found or inactive") + + try: + from app.services.voice_providers.prompt_sync import sync_provider_prompt + prompt = sync_provider_prompt(db_agent, integration, db) + except Exception as e: + raise HTTPException(status_code=502, detail=f"Failed to fetch prompt from provider: {str(e)}") + + db.refresh(db_agent) + synced = isinstance(prompt, str) and bool(prompt.strip()) + if not synced: + raise HTTPException( + status_code=422, + detail="Provider returned no prompt. Verify the external agent has a system prompt configured.", + ) + return { + "synced": synced, + "provider_prompt": db_agent.provider_prompt, + "provider_prompt_synced_at": db_agent.provider_prompt_synced_at.isoformat() if db_agent.provider_prompt_synced_at else None, + } + + +@router.delete("/{agent_id}") +async def delete_agent( + agent_id: str, + force: bool = Query(False, description="Force delete with all dependent records"), + organization_id: UUID = Depends(get_organization_id), + workspace_id: UUID = Depends(get_workspace_id), + db: Session = Depends(get_db) +): + """Delete an agent (scoped to the active workspace). Returns 409 if dependent records exist unless force=true.""" + try: + agent_uuid = UUID(agent_id) + db_agent = db.query(Agent).filter( + and_( + Agent.id == agent_uuid, + Agent.organization_id == organization_id, + Agent.workspace_id == workspace_id, + ) + ).first() + except ValueError: + db_agent = db.query(Agent).filter( + and_( + Agent.agent_id == agent_id, + Agent.organization_id == organization_id, + Agent.workspace_id == workspace_id, + ) + ).first() + + if not db_agent: + raise HTTPException(status_code=404, detail=f"Agent {agent_id} not found") + + agent_uuid = db_agent.id + dependencies = get_agent_dependencies(db, organization_id, agent_uuid) + + if dependencies and not force: + parts = [] + if dependencies.get("evaluators"): + parts.append(f"{dependencies['evaluators']} evaluator(s)") + if dependencies.get("evaluator_results"): + parts.append(f"{dependencies['evaluator_results']} evaluator result(s)") + if dependencies.get("call_recordings"): + parts.append(f"{dependencies['call_recordings']} call recording(s)") + if dependencies.get("conversation_evaluations"): + parts.append(f"{dependencies['conversation_evaluations']} conversation evaluation(s)") + if dependencies.get("test_conversations"): + parts.append(f"{dependencies['test_conversations']} test conversation(s)") + + raise HTTPException( + status_code=status.HTTP_409_CONFLICT, + detail={ + "message": f"Cannot delete agent. It is referenced by: {', '.join(parts)}.", + "dependencies": dependencies, + "hint": "Use force=true to delete this agent and all its dependent records.", + }, + ) + + if dependencies: + # Delete in FK-safe order: + # 1. EvaluatorResults (references evaluators and agents) + db.query(EvaluatorResult).filter( + EvaluatorResult.agent_id == agent_uuid, + ).delete(synchronize_session=False) + + # 2. Evaluators (references agents) + db.query(Evaluator).filter( + Evaluator.agent_id == agent_uuid, + ).delete(synchronize_session=False) + + # 3. Nullify call recordings (keep recordings, unlink agent) + db.query(CallRecording).filter( + CallRecording.agent_id == agent_uuid, + ).update({CallRecording.agent_id: None}, synchronize_session=False) + + # 4. ConversationEvaluations + db.query(ConversationEvaluation).filter( + ConversationEvaluation.agent_id == agent_uuid, + ).delete(synchronize_session=False) + + # 5. TestAgentConversations + db.query(TestAgentConversation).filter( + TestAgentConversation.agent_id == agent_uuid, + ).delete(synchronize_session=False) + + db.delete(db_agent) + db.commit() + + if dependencies: + return JSONResponse( + status_code=200, + content={ + "message": "Agent and all dependent records deleted successfully.", + "deleted": dependencies, + }, + ) + + return Response(status_code=204) + + +@router.get("/{agent_id}/delete-impact") +async def get_agent_delete_impact( + agent_id: str, + organization_id: UUID = Depends(get_organization_id), + workspace_id: UUID = Depends(get_workspace_id), + db: Session = Depends(get_db) +): + """Preview dependent records that would be affected by force delete (scoped to the active workspace).""" + try: + agent_uuid = UUID(agent_id) + db_agent = db.query(Agent).filter( + and_( + Agent.id == agent_uuid, + Agent.organization_id == organization_id, + Agent.workspace_id == workspace_id, + ) + ).first() + except ValueError: + db_agent = db.query(Agent).filter( + and_( + Agent.agent_id == agent_id, + Agent.organization_id == organization_id, + Agent.workspace_id == workspace_id, + ) + ).first() + + if not db_agent: + raise HTTPException(status_code=404, detail=f"Agent {agent_id} not found") + + dependencies = get_agent_dependencies(db, organization_id, db_agent.id) + return { + "agent_id": str(db_agent.id), + "agent_name": db_agent.name, + "dependencies": dependencies, + "can_delete_without_force": len(dependencies) == 0, + } + + +from app.core.auth.capabilities import SIM_MANAGE, SIM_VIEW +from app.core.auth.workspace_route_capabilities import apply_workspace_route_capabilities + +apply_workspace_route_capabilities( + router, + view_capability=SIM_VIEW, + manage_capability=SIM_MANAGE, +) + diff --git a/app/api/v1/routes/aiproviders.py b/app/api/v1/routes/aiproviders.py index b7f8b9e6..7f8f8e01 100644 --- a/app/api/v1/routes/aiproviders.py +++ b/app/api/v1/routes/aiproviders.py @@ -198,6 +198,7 @@ async def create_aiprovider( else None ), gateway_extra_headers=aiprovider.gateway_extra_headers, + enabled_models=aiprovider.enabled_models, ) db.add(db_aiprovider) db.flush() diff --git a/app/api/v1/routes/call_import_evaluations.py b/app/api/v1/routes/call_import_evaluations.py index a57e2e58..39675b13 100644 --- a/app/api/v1/routes/call_import_evaluations.py +++ b/app/api/v1/routes/call_import_evaluations.py @@ -3074,10 +3074,19 @@ def _explain_period_deltas( from app.services.ai.llm_resolver import get_llm_provider_and_model from app.services.call_import_user_insights import _call_llm, _parse_json_object + from app.services.usage.call_import_context import ( + call_import_evaluation_usage_context, + ) provider_enum, model_str = get_llm_provider_and_model( organization_id, db, provider_hint, model_hint ) + usage_ctx = call_import_evaluation_usage_context( + organization_id=organization_id, + workspace_id=evaluation.workspace_id, + evaluation_id=evaluation.id, + call_import_id=evaluation.call_import_id, + ) try: text = _call_llm( db, @@ -3097,6 +3106,7 @@ def _explain_period_deltas( ], temperature=0.3, max_tokens=900, + usage_ctx=usage_ctx, ) except Exception as exc: logger.warning("[PeriodDeltaExplain] LLM call failed: {}", exc) @@ -3748,15 +3758,29 @@ async def generate_call_import_evaluation_pdf_report( ) return _pdf_report_response_from_row(cached_pdf_report, cache_hit=True) - narrative = _generate_report_narrative( - db, - organization_id, - metric_aggregates=metric_aggregates, - insight_aggregates=insight_aggregates if is_internal else [], - period_delta_by_metric=period_delta_by_metric, - evidence_samples=evidence_samples if is_internal else {}, - report_config=report_config, + from app.services.usage.call_import_context import ( + call_import_evaluation_usage_context, ) + from app.services.usage.context import llm_usage_context + + with llm_usage_context( + call_import_evaluation_usage_context( + organization_id=organization_id, + workspace_id=getattr(call_import, "workspace_id", None) + or evaluation.workspace_id, + evaluation_id=evaluation.id, + call_import_id=evaluation.call_import_id, + ) + ): + narrative = _generate_report_narrative( + db, + organization_id, + metric_aggregates=metric_aggregates, + insight_aggregates=insight_aggregates if is_internal else [], + period_delta_by_metric=period_delta_by_metric, + evidence_samples=evidence_samples if is_internal else {}, + report_config=report_config, + ) generated_at = datetime.now(timezone.utc) try: @@ -5467,29 +5491,41 @@ def _generate_and_persist_tldr_summary( from app.services.ai.llm_resolver import get_llm_provider_and_model from app.services.ai.llm_service import llm_service + from app.services.usage.call_import_context import ( + call_import_evaluation_usage_context, + ) + from app.services.usage.context import llm_usage_context provider_enum, model_str = get_llm_provider_and_model( organization_id, db, provider, model ) try: - llm_result = llm_service.generate_response( - messages=messages, - llm_provider=provider_enum, - llm_model=model_str, - organization_id=organization_id, - db=db, - temperature=0.4, - max_tokens=1400, - ) + with llm_usage_context( + call_import_evaluation_usage_context( + organization_id=organization_id, + workspace_id=evaluation.workspace_id, + evaluation_id=evaluation.id, + call_import_id=evaluation.call_import_id, + ) + ): + llm_result = llm_service.generate_response( + messages=messages, + llm_provider=provider_enum, + llm_model=model_str, + organization_id=organization_id, + db=db, + temperature=0.4, + max_tokens=1400, + ) except Exception as e: logger.error(f"[CallImportInsights] LLM call failed: {e}") raise HTTPException( status_code=502, detail=f"LLM call failed: {e}" ) from e - summary = _parse_insights_response(llm_result.get("text", "")) total = int(evaluation.total_rows or 0) + summary = _parse_insights_response(llm_result.get("text", "")) ui_completed = min(int(evaluation.completed_rows or 0), total) if total else int( evaluation.completed_rows or 0 ) diff --git a/app/api/v1/routes/call_imports.py b/app/api/v1/routes/call_imports.py index 7cb60d2a..094845ae 100644 --- a/app/api/v1/routes/call_imports.py +++ b/app/api/v1/routes/call_imports.py @@ -1,4 +1,4 @@ -"""CSV-driven call import routes. +"""CSV-driven call import routes. Users upload a CSV plus a per-batch column mapping (CSV header -> system field). The backend persists a CallImport batch + one CallImportRow per diff --git a/app/api/v1/routes/chat.py b/app/api/v1/routes/chat.py index a8317844..6fc026f8 100644 --- a/app/api/v1/routes/chat.py +++ b/app/api/v1/routes/chat.py @@ -47,20 +47,32 @@ async def chat_completion( ): """Generate a chat completion using the specified AI provider and model.""" try: - # Convert ChatMessage to dict format expected by LLM service - messages = [{"role": msg.role, "content": msg.content} for msg in request.messages] - - result = llm_service.generate_response( - messages=messages, - llm_provider=request.provider, - llm_model=request.model, - organization_id=organization_id, - db=db, - llm_config=request.llm_config, - temperature=request.temperature, - max_tokens=request.max_tokens, - task_defaults={"temperature": 0.7}, + from app.services.usage.context import ( + LLMUsageContext, + LLMUsageProductSection, + llm_usage_context, ) + + messages = [{"role": msg.role, "content": msg.content} for msg in request.messages] + + with llm_usage_context( + LLMUsageContext( + organization_id=organization_id, + workspace_id=workspace_id, + product_section=LLMUsageProductSection.CHAT, + ) + ): + result = llm_service.generate_response( + messages=messages, + llm_provider=request.provider, + llm_model=request.model, + organization_id=organization_id, + db=db, + llm_config=request.llm_config, + temperature=request.temperature, + max_tokens=request.max_tokens, + task_defaults={"temperature": 0.7}, + ) background_tasks.add_task( record_chat_completion, diff --git a/app/api/v1/routes/cron_jobs.py b/app/api/v1/routes/cron_jobs.py index 6bc6b197..c09ad7f4 100644 --- a/app/api/v1/routes/cron_jobs.py +++ b/app/api/v1/routes/cron_jobs.py @@ -140,6 +140,8 @@ def create_cron_job( cron_job = CronJob( organization_id=organization_id, name=cron_job_data.name, + job_type="evaluator_run", + is_system=False, cron_expression=cron_job_data.cron_expression, timezone=cron_job_data.timezone, max_runs=cron_job_data.max_runs, @@ -162,7 +164,10 @@ def list_cron_jobs( db: Session = Depends(get_db), ): """List all cron jobs for the organization.""" - query = db.query(CronJob).filter(CronJob.organization_id == organization_id) + query = db.query(CronJob).filter( + CronJob.organization_id == organization_id, + CronJob.is_system.is_(False), + ) if status_filter: query = query.filter(CronJob.status == status_filter.value) @@ -209,6 +214,12 @@ def update_cron_job( if not cron_job: raise HTTPException(status_code=404, detail="Cron job not found") + if cron_job.is_system: + raise HTTPException( + status_code=status.HTTP_400_BAD_REQUEST, + detail="System cron jobs cannot be modified", + ) + # Update fields if provided if cron_job_data.name is not None: # Check for name conflicts @@ -309,6 +320,12 @@ def delete_cron_job( if not cron_job: raise HTTPException(status_code=404, detail="Cron job not found") + if cron_job.is_system: + raise HTTPException( + status_code=status.HTTP_400_BAD_REQUEST, + detail="System cron jobs cannot be deleted", + ) + db.delete(cron_job) db.commit() @@ -332,6 +349,12 @@ def toggle_cron_job_status( if not cron_job: raise HTTPException(status_code=404, detail="Cron job not found") + if cron_job.is_system: + raise HTTPException( + status_code=status.HTTP_400_BAD_REQUEST, + detail="System cron jobs cannot be toggled", + ) + # Don't allow toggling completed jobs if cron_job.status == CronJobStatus.COMPLETED.value: raise HTTPException( diff --git a/app/api/v1/routes/org_usage.py b/app/api/v1/routes/org_usage.py new file mode 100644 index 00000000..53dbbb27 --- /dev/null +++ b/app/api/v1/routes/org_usage.py @@ -0,0 +1,1717 @@ +"""Org-scoped LLM Usage API (tokens + call counts).""" + +from __future__ import annotations + +from datetime import date, datetime, timezone +from app.services.usage.access import UsageAccessPolicy +from app.services.usage.dates import usage_date_filter_bounds, usage_local_today +from typing import List, Literal, Optional +from uuid import UUID + +from fastapi import APIRouter, Depends, HTTPException, Query +from pydantic import BaseModel, Field +from sqlalchemy import and_, case, cast, func, or_, select, String +from sqlalchemy.dialects.postgresql import UUID as PG_UUID +from sqlalchemy.orm import Session + +from app.database import get_db +from app.dependencies import get_organization_id +from app.models.database import ( + CallImport, + CallImportEvaluation, + CallImportEvaluationRow, + CallImportRow, + CallImportTag, + CallImportTagAssignment, + Agent, + LLMUsageDaily, + Workspace, +) +from app.services.usage.read_cache import ( + cache_key_for, + get_cached_response, + set_cached_response, +) +from app.services.usage.fx_rates import get_usd_inr_rate +from app.services.usage.usage_costs import costs_from_micro +from app.services.usage.usage_labels import ( + labels_for_call_import_ids, + labels_for_resource_buckets, + usage_kind_label, + UsageNameResolver, + parse_uuid, +) + +router = APIRouter(prefix="/organizations/usage", tags=["Usage"]) + +GroupBy = Literal[ + "workspace", + "product_section", + "model", + "resource", + "usage_kind", + "call_import", +] + +SECTION_LABELS = { + "call_import_evaluations": "Call Import Evaluations", + "call_imports": "Call Imports", + "playground": "Playground", + "voice_playground": "Voice Playground", + "evaluators": "Evaluators", + "metrics": "Metrics", + "chat": "Chat", + "judge_alignment": "Judge Alignment", + "prompt_optimization": "Prompt Optimization", + "personas": "Personas", + "agents": "Agents", + "prompt_partials": "Prompt Partials", + "conversation_evaluations": "Conversation Evaluations", + "telephony": "Telephony", + "test_agent": "Test Agent", + "other": "Other", +} + +_LABEL_ROW_LIMIT = 5000 + + +def _usage_row_weight(): + return ( + LLMUsageDaily.prompt_tokens + + LLMUsageDaily.completion_tokens + + LLMUsageDaily.cache_read_tokens + + LLMUsageDaily.cache_creation_tokens + + LLMUsageDaily.reasoning_tokens + + LLMUsageDaily.audio_seconds + + LLMUsageDaily.tts_characters + ) + + +def _label_row_order(): + return ( + _usage_row_weight().desc(), + LLMUsageDaily.call_count.desc(), + ) + + +class UsageCosts(BaseModel): + input_cost_usd: float = 0 + output_cost_usd: float = 0 + cache_read_cost_usd: float = 0 + cache_write_cost_usd: float = 0 + reasoning_cost_usd: float = 0 + audio_cost_usd: float = 0 + tts_cost_usd: float = 0 + total_cost_usd: float = 0 + currency: str = "USD" + has_unpriced_usage: bool = False + + +class UsageTotals(BaseModel): + prompt_tokens: int = 0 + completion_tokens: int = 0 + total_tokens: int = 0 + cache_read_tokens: int = 0 + cache_creation_tokens: int = 0 + reasoning_tokens: int = 0 + audio_seconds: int = 0 + tts_characters: int = 0 + call_count: int = 0 + input_cost_micro_usd: int = 0 + output_cost_micro_usd: int = 0 + cache_read_cost_micro_usd: int = 0 + cache_creation_cost_micro_usd: int = 0 + reasoning_cost_micro_usd: int = 0 + audio_cost_micro_usd: int = 0 + tts_cost_micro_usd: int = 0 + total_cost_micro_usd: int = 0 + costs: UsageCosts = Field(default_factory=UsageCosts) + + +class UsageFxRateResponse(BaseModel): + base: str = "USD" + quote: str = "INR" + rate: float + as_of: datetime + source: str + + +class UsagePolicyMeta(BaseModel): + extended_history: bool + max_history_days: Optional[int] = None + range_clamped: bool = False + + +class UsageSummaryResponse(BaseModel): + start: date + end: date + totals: UsageTotals + usage_policy: UsagePolicyMeta + last_updated_at: Optional[datetime] = None + + +class UsageBreakdownRow(BaseModel): + workspace_id: Optional[UUID] = None + workspace_name: Optional[str] = None + product_section: Optional[str] = None + product_section_label: Optional[str] = None + model: Optional[str] = None + resource_id: Optional[UUID] = None + resource_type: Optional[str] = None + resource_label: Optional[str] = None + call_import_id: Optional[UUID] = None + call_import_label: Optional[str] = None + usage_kind: Optional[str] = None + prompt_tokens: int = 0 + completion_tokens: int = 0 + total_tokens: int = 0 + cache_read_tokens: int = 0 + cache_creation_tokens: int = 0 + reasoning_tokens: int = 0 + audio_seconds: int = 0 + tts_characters: int = 0 + call_count: int = 0 + input_cost_micro_usd: int = 0 + output_cost_micro_usd: int = 0 + cache_read_cost_micro_usd: int = 0 + cache_creation_cost_micro_usd: int = 0 + reasoning_cost_micro_usd: int = 0 + audio_cost_micro_usd: int = 0 + tts_cost_micro_usd: int = 0 + total_cost_micro_usd: int = 0 + costs: UsageCosts = Field(default_factory=UsageCosts) + + +class UsageBreakdownResponse(BaseModel): + start: date + end: date + group_by: GroupBy + rows: List[UsageBreakdownRow] + total_count: int + truncated_at_limit: bool = False + usage_policy: UsagePolicyMeta + last_updated_at: Optional[datetime] = None + + +class UsageFiltersResponse(BaseModel): + workspaces: List[dict] = Field(default_factory=list) + product_sections: List[dict] = Field(default_factory=list) + call_imports: List[dict] = Field(default_factory=list) + evaluations: List[dict] = Field(default_factory=list) + models: List[str] = Field(default_factory=list) + resources: List[dict] = Field(default_factory=list) + usage_kinds: List[dict] = Field(default_factory=list) + datasets: List[str] = Field(default_factory=list) + tags: List[dict] = Field(default_factory=list) + + +def _parse_usage_range( + start: Optional[date], + end: Optional[date], + tz: Optional[str], +) -> tuple[date, date, date, date]: + """Return display_start, display_end, filter_start, filter_end.""" + today = usage_local_today(tz) + display_start = start or today + display_end = end or today + filter_start, filter_end = usage_date_filter_bounds( + display_start, display_end, tz + ) + return display_start, display_end, filter_start, filter_end + + +def _usage_policy_meta(access) -> UsagePolicyMeta: + return UsagePolicyMeta( + extended_history=access.policy.extended_history, + max_history_days=access.policy.max_history_days, + range_clamped=access.range_clamped, + ) + + +def _evaluation_id_expr(): + return func.coalesce( + LLMUsageDaily.context["evaluation_id"].astext, + case( + ( + LLMUsageDaily.context["resource_type"].astext + == "call_import_evaluation", + LLMUsageDaily.context["resource_id"].astext, + ), + else_=None, + ), + ) + + +def _resource_id_expr(): + """Resource rollup key: explicit resource_id or agent_id fallback.""" + return func.coalesce( + LLMUsageDaily.context["resource_id"].astext, + LLMUsageDaily.context["agent_id"].astext, + ) + + +def _resource_type_expr(): + rid_expr = _resource_id_expr() + return func.coalesce( + LLMUsageDaily.context["resource_type"].astext, + case( + (LLMUsageDaily.context["agent_id"].astext.isnot(None), "agent"), + ( + and_( + LLMUsageDaily.product_section == "agents", + rid_expr.isnot(None), + ), + "agent", + ), + else_=None, + ), + ) + + +def _call_import_group_expr(): + """Resolve call import id from context, resource row, evaluation, or import row.""" + eval_call_import = ( + select(CallImportEvaluation.call_import_id) + .where( + CallImportEvaluation.id == cast(_evaluation_id_expr(), PG_UUID) + ) + .correlate(LLMUsageDaily) + .scalar_subquery() + ) + row_call_import = ( + select(CallImportRow.call_import_id) + .where( + CallImportRow.id + == cast(LLMUsageDaily.context["call_import_row_id"].astext, PG_UUID) + ) + .correlate(LLMUsageDaily) + .scalar_subquery() + ) + eval_row_call_import = ( + select(CallImportEvaluation.call_import_id) + .select_from(CallImportEvaluationRow) + .join( + CallImportEvaluation, + CallImportEvaluation.id == CallImportEvaluationRow.evaluation_id, + ) + .where( + CallImportEvaluationRow.id + == cast(LLMUsageDaily.context["evaluation_row_id"].astext, PG_UUID) + ) + .correlate(LLMUsageDaily) + .scalar_subquery() + ) + return cast( + func.coalesce( + LLMUsageDaily.context["call_import_id"].astext, + case( + ( + LLMUsageDaily.context["resource_type"].astext == "call_import", + LLMUsageDaily.context["resource_id"].astext, + ), + else_=None, + ), + cast(eval_call_import, String), + cast(row_call_import, String), + cast(eval_row_call_import, String), + ), + String, + ) + + +def _resource_scope_filter(resource_id: UUID): + """Match usage attributed to a product resource (agent, simulation, etc.).""" + rid = str(resource_id) + return or_( + LLMUsageDaily.context["resource_id"].astext == rid, + LLMUsageDaily.context["agent_id"].astext == rid, + ) + + +def _evaluation_scope_filter( + evaluation_id: UUID, + organization_id: UUID, + db: Session, +): + eid = str(evaluation_id) + row_ids = [ + str(row[0]) + for row in db.query(CallImportEvaluationRow.id) + .filter( + CallImportEvaluationRow.evaluation_id == evaluation_id, + ) + .all() + ] + clauses = [ + LLMUsageDaily.context["evaluation_id"].astext == eid, + and_( + LLMUsageDaily.context["resource_id"].astext == eid, + LLMUsageDaily.context["resource_type"].astext == "call_import_evaluation", + ), + ] + if row_ids: + clauses.append(LLMUsageDaily.context["evaluation_row_id"].astext.in_(row_ids)) + return or_(*clauses) + + +def _call_import_scope_filter( + call_import_id: UUID, + organization_id: UUID, + db: Session, +): + """Match usage tied to a call import (direct context, resource row, or eval runs).""" + cid = str(call_import_id) + eval_ids = [ + str(row[0]) + for row in db.query(CallImportEvaluation.id) + .filter( + CallImportEvaluation.organization_id == organization_id, + CallImportEvaluation.call_import_id == call_import_id, + ) + .all() + ] + clauses = [ + LLMUsageDaily.context["call_import_id"].astext == cid, + and_( + LLMUsageDaily.context["resource_id"].astext == cid, + LLMUsageDaily.context["resource_type"].astext == "call_import", + ), + ] + if eval_ids: + clauses.append(LLMUsageDaily.context["evaluation_id"].astext.in_(eval_ids)) + clauses.append( + and_( + LLMUsageDaily.context["resource_id"].astext.in_(eval_ids), + LLMUsageDaily.context["resource_type"].astext == "call_import_evaluation", + ) + ) + return or_(*clauses) + + +def _call_import_ids_for_filters( + db: Session, + *, + organization_id: UUID, + workspace_id: Optional[UUID] = None, + dataset: Optional[str] = None, + tag_id: Optional[UUID] = None, +) -> List[UUID]: + query = db.query(CallImport.id).filter( + CallImport.organization_id == organization_id, + ) + if workspace_id is not None: + query = query.filter(CallImport.workspace_id == workspace_id) + if dataset: + query = query.filter(CallImport.dataset == dataset) + if tag_id is not None: + query = query.filter( + CallImport.id.in_( + db.query(CallImportTagAssignment.call_import_id).filter( + CallImportTagAssignment.tag_id == tag_id, + ) + ) + ) + return [row[0] for row in query.all()] + + +def _call_import_ids_scope_filter( + allowed_import_ids: List[UUID], +): + if not allowed_import_ids: + return LLMUsageDaily.id.is_(None) + allowed = [str(uid) for uid in allowed_import_ids] + return cast(_call_import_group_expr(), String).in_(allowed) + + +def _workspace_scope_filter( + workspace_id: UUID, + organization_id: UUID, +): + """Match workspace-scoped rows plus legacy call-import usage with null workspace_id.""" + ws_call_import_ids = ( + select(cast(CallImport.id, String)) + .where( + CallImport.organization_id == organization_id, + CallImport.workspace_id == workspace_id, + ) + ) + legacy_call_import = and_( + LLMUsageDaily.workspace_id.is_(None), + _call_import_group_expr().in_(ws_call_import_ids), + ) + return or_(LLMUsageDaily.workspace_id == workspace_id, legacy_call_import) + + +def _apply_filters( + query, + *, + organization_id: UUID, + start: date, + end: date, + enforced_floor: Optional[date] = None, + workspace_id: Optional[UUID], + product_section: Optional[str], + model: Optional[str], + resource_id: Optional[UUID], + usage_kind: Optional[str] = None, + call_import_id: Optional[UUID] = None, + evaluation_id: Optional[UUID] = None, + evaluation_row_id: Optional[UUID] = None, + dataset: Optional[str] = None, + tag_id: Optional[UUID] = None, + db: Optional[Session] = None, +): + effective_start = start + if enforced_floor is not None and effective_start < enforced_floor: + effective_start = enforced_floor + query = query.filter( + LLMUsageDaily.organization_id == organization_id, + LLMUsageDaily.usage_date >= effective_start, + LLMUsageDaily.usage_date <= end, + ) + if workspace_id is not None: + query = query.filter( + _workspace_scope_filter(workspace_id, organization_id) + ) + if product_section: + query = query.filter(LLMUsageDaily.product_section == product_section) + if model: + query = query.filter(LLMUsageDaily.model == model) + if resource_id is not None: + rid_filter = _resource_scope_filter(resource_id) + if db is not None: + query = query.filter( + or_( + _evaluation_scope_filter(resource_id, organization_id, db), + rid_filter, + ) + ) + else: + query = query.filter( + or_( + LLMUsageDaily.context["resource_id"].astext == str(resource_id), + LLMUsageDaily.context["agent_id"].astext == str(resource_id), + LLMUsageDaily.context["evaluation_id"].astext == str(resource_id), + LLMUsageDaily.context["evaluation_row_id"].astext == str(resource_id), + ) + ) + if call_import_id is not None: + if db is not None: + query = query.filter( + _call_import_scope_filter(call_import_id, organization_id, db) + ) + else: + query = query.filter( + LLMUsageDaily.context["call_import_id"].astext == str(call_import_id) + ) + if evaluation_id is not None and evaluation_id != resource_id: + if db is not None: + query = query.filter( + _evaluation_scope_filter(evaluation_id, organization_id, db) + ) + else: + query = query.filter( + LLMUsageDaily.context["evaluation_id"].astext == str(evaluation_id) + ) + if evaluation_row_id is not None: + query = query.filter( + LLMUsageDaily.context["evaluation_row_id"].astext + == str(evaluation_row_id) + ) + if usage_kind: + query = query.filter(LLMUsageDaily.usage_kind == usage_kind) + if dataset or tag_id is not None: + if db is None: + raise ValueError("db required for dataset/tag filters") + allowed = _call_import_ids_for_filters( + db, + organization_id=organization_id, + workspace_id=workspace_id, + dataset=dataset, + tag_id=tag_id, + ) + query = query.filter(_call_import_ids_scope_filter(allowed)) + return query + + +def _filtered_query( + db: Session, + *, + organization_id: UUID, + start: date, + end: date, + enforced_floor: Optional[date] = None, + workspace_id: Optional[UUID] = None, + product_section: Optional[str] = None, + model: Optional[str] = None, + resource_id: Optional[UUID] = None, + usage_kind: Optional[str] = None, + call_import_id: Optional[UUID] = None, + evaluation_id: Optional[UUID] = None, + evaluation_row_id: Optional[UUID] = None, + dataset: Optional[str] = None, + tag_id: Optional[UUID] = None, +): + return _apply_filters( + db.query(LLMUsageDaily), + organization_id=organization_id, + start=start, + end=end, + enforced_floor=enforced_floor, + workspace_id=workspace_id, + product_section=product_section, + model=model, + resource_id=resource_id, + usage_kind=usage_kind, + call_import_id=call_import_id, + evaluation_id=evaluation_id, + evaluation_row_id=evaluation_row_id, + dataset=dataset, + tag_id=tag_id, + db=db, + ) + + +def _infer_agent_types_for_label_buckets( + db: Session, + organization_id: UUID, + grouped: dict[str, tuple[Optional[str], list]], +) -> None: + """When context JSON omits resource_type, infer agent rows from Agent.id.""" + candidate_ids: list[UUID] = [] + for rid, (rtype, _) in grouped.items(): + if rtype: + continue + uid = parse_uuid(rid) + if uid: + candidate_ids.append(uid) + if not candidate_ids: + return + known_agent_ids = { + row.id + for row in db.query(Agent.id).filter( + Agent.organization_id == organization_id, + Agent.id.in_(candidate_ids), + ).all() + } + for rid in grouped: + uid = parse_uuid(rid) + if uid and uid in known_agent_ids: + rtype, contexts = grouped[rid] + if not rtype: + grouped[rid] = ("agent", contexts) + + +def _resource_label_map( + db: Session, + organization_id: UUID, + query, +) -> dict[str, str]: + """Build resource_id -> hierarchical label from usage rows in query.""" + rows = ( + query.with_entities( + _resource_id_expr(), + _resource_type_expr(), + LLMUsageDaily.context, + ) + .filter( + or_( + LLMUsageDaily.context["resource_id"].astext.isnot(None), + LLMUsageDaily.context["agent_id"].astext.isnot(None), + ) + ) + .order_by(*_label_row_order()) + .limit(_LABEL_ROW_LIMIT) + .all() + ) + grouped: dict[str, tuple[Optional[str], list]] = {} + for raw_id, rtype, ctx in rows: + if not raw_id: + continue + key = str(raw_id) + if key not in grouped: + grouped[key] = (rtype, []) + grouped[key][1].append(ctx) + + _infer_agent_types_for_label_buckets(db, organization_id, grouped) + + buckets = [(rid, rtype, contexts) for rid, (rtype, contexts) in grouped.items()] + resolver = UsageNameResolver(db, organization_id) + contexts_for_preload = [] + for rid, rtype, contexts in buckets: + for ctx in contexts: + merged = dict(ctx or {}) + if rid: + merged.setdefault("resource_id", rid) + if rtype: + merged.setdefault("resource_type", rtype) + contexts_for_preload.append(merged) + resolver.preload(contexts_for_preload) + return labels_for_resource_buckets(buckets, resolver) + + +def _resource_filter_meta_map( + db: Session, + organization_id: UUID, + query, +) -> dict[str, dict[str, Optional[str]]]: + """Map resource id -> {type, product_section} for filter dropdowns.""" + rows = ( + query.with_entities( + _resource_id_expr(), + _resource_type_expr(), + LLMUsageDaily.product_section, + ) + .filter( + or_( + LLMUsageDaily.context["resource_id"].astext.isnot(None), + LLMUsageDaily.context["agent_id"].astext.isnot(None), + ) + ) + .distinct() + .limit(_LABEL_ROW_LIMIT) + .all() + ) + meta: dict[str, dict[str, Optional[str]]] = {} + for raw_id, rtype, section in rows: + if not raw_id: + continue + key = str(raw_id) + entry = meta.setdefault(key, {"type": None, "product_section": None}) + if rtype: + entry["type"] = str(rtype) + if section: + entry["product_section"] = str(section) + if not entry["type"] and section == "agents": + entry["type"] = "agent" + return meta + + +def _breakdown_resource_label( + raw_res_id: Optional[str], + res_type: Optional[str], + section: Optional[str], + resource_labels: dict[str, str], + db: Session, + organization_id: UUID, +) -> str: + if not raw_res_id: + return "Unscoped" + key = str(raw_res_id) + label = resource_labels.get(key) + if label and label != "Unscoped": + return label + effective_type = res_type or ("agent" if section == "agents" else None) + if effective_type == "agent": + resolver = UsageNameResolver(db, organization_id) + resolver.preload([{"resource_id": key, "resource_type": "agent"}]) + return resolver.agent_name(key) + return label or "Unscoped" + + +def _collect_call_import_ids_from_usage_rows( + db: Session, + organization_id: UUID, + rows: list, +) -> set[UUID]: + """Extract call import ids referenced in usage row context tuples.""" + import_ids: set[UUID] = set() + eval_ids: set[UUID] = set() + row_ids: set[UUID] = set() + eval_row_ids: set[UUID] = set() + + for row in rows: + cid_raw = row[0] if len(row) > 0 else None + res_raw = row[1] if len(row) > 1 else None + rtype = row[2] if len(row) > 2 else None + eval_raw = row[3] if len(row) > 3 else None + row_id_raw = row[4] if len(row) > 4 else None + eval_row_raw = row[5] if len(row) > 5 else None + + if cid_raw: + uid = parse_uuid(cid_raw) + if uid: + import_ids.add(uid) + if res_raw and rtype == "call_import": + uid = parse_uuid(res_raw) + if uid: + import_ids.add(uid) + if eval_raw: + uid = parse_uuid(eval_raw) + if uid: + eval_ids.add(uid) + if res_raw and rtype == "call_import_evaluation": + uid = parse_uuid(res_raw) + if uid: + eval_ids.add(uid) + if row_id_raw: + uid = parse_uuid(row_id_raw) + if uid: + row_ids.add(uid) + if eval_row_raw: + uid = parse_uuid(eval_row_raw) + if uid: + eval_row_ids.add(uid) + + if eval_ids: + for ev in ( + db.query(CallImportEvaluation) + .filter( + CallImportEvaluation.organization_id == organization_id, + CallImportEvaluation.id.in_(eval_ids), + ) + .all() + ): + import_ids.add(ev.call_import_id) + + if row_ids: + for cid in ( + db.query(CallImportRow.call_import_id) + .filter(CallImportRow.id.in_(row_ids)) + .distinct() + .all() + ): + import_ids.add(cid[0]) + + if eval_row_ids: + for cid in ( + db.query(CallImportEvaluation.call_import_id) + .join( + CallImportEvaluationRow, + CallImportEvaluationRow.evaluation_id == CallImportEvaluation.id, + ) + .filter(CallImportEvaluationRow.id.in_(eval_row_ids)) + .distinct() + .all() + ): + import_ids.add(cid[0]) + + return import_ids + + +def _call_import_label_map( + db: Session, + organization_id: UUID, + query, +) -> dict[str, str]: + rows = ( + query.with_entities( + LLMUsageDaily.context["call_import_id"].astext, + LLMUsageDaily.context["resource_id"].astext, + LLMUsageDaily.context["resource_type"].astext, + LLMUsageDaily.context["evaluation_id"].astext, + LLMUsageDaily.context["call_import_row_id"].astext, + LLMUsageDaily.context["evaluation_row_id"].astext, + ) + .order_by(*_label_row_order()) + .limit(_LABEL_ROW_LIMIT) + .all() + ) + import_ids = _collect_call_import_ids_from_usage_rows( + db, organization_id, rows + ) + + if not import_ids: + return {} + resolver = UsageNameResolver(db, organization_id) + resolver.preload( + [{"call_import_id": str(uid)} for uid in import_ids] + ) + return labels_for_call_import_ids(list(import_ids), resolver) + + +def _call_import_filter_labels( + db: Session, + organization_id: UUID, + query, + workspace_id: Optional[UUID] = None, + dataset: Optional[str] = None, + tag_id: Optional[UUID] = None, +) -> dict[str, str]: + """Call imports for filters: usage in range plus workspace imports when scoped.""" + import_ids = _collect_call_import_ids_from_usage_rows( + db, + organization_id, + query.with_entities( + LLMUsageDaily.context["call_import_id"].astext, + LLMUsageDaily.context["resource_id"].astext, + LLMUsageDaily.context["resource_type"].astext, + LLMUsageDaily.context["evaluation_id"].astext, + LLMUsageDaily.context["call_import_row_id"].astext, + LLMUsageDaily.context["evaluation_row_id"].astext, + ).order_by(*_label_row_order()).limit(_LABEL_ROW_LIMIT).all(), + ) + + if workspace_id is not None: + scoped_ids = _call_import_ids_for_filters( + db, + organization_id=organization_id, + workspace_id=workspace_id, + dataset=dataset, + tag_id=tag_id, + ) + import_ids.update(scoped_ids) + + if not import_ids: + return {} + resolver = UsageNameResolver(db, organization_id) + resolver.preload([{"call_import_id": str(uid)} for uid in import_ids]) + return labels_for_call_import_ids(list(import_ids), resolver) + + +def _evaluation_label_map( + db: Session, + organization_id: UUID, + query, +) -> dict[str, str]: + """Evaluation id -> short label (name + id suffix) for filter dropdowns.""" + rows = ( + query.with_entities( + LLMUsageDaily.context["evaluation_id"].astext, + LLMUsageDaily.context["resource_id"].astext, + LLMUsageDaily.context["resource_type"].astext, + LLMUsageDaily.context, + ) + .filter( + or_( + LLMUsageDaily.context["evaluation_id"].astext.isnot(None), + LLMUsageDaily.context["resource_type"].astext == "call_import_evaluation", + ) + ) + .order_by(*_label_row_order()) + .limit(_LABEL_ROW_LIMIT) + .all() + ) + grouped: dict[str, tuple[Optional[str], list]] = {} + for eval_raw, res_raw, rtype, ctx in rows: + key_raw = eval_raw + if not key_raw and rtype == "call_import_evaluation" and res_raw: + key_raw = res_raw + if not key_raw: + continue + key = str(key_raw) + if key not in grouped: + grouped[key] = (rtype, []) + grouped[key][1].append(ctx) + + buckets = [(rid, rtype, contexts) for rid, (rtype, contexts) in grouped.items()] + resolver = UsageNameResolver(db, organization_id) + contexts_for_preload = [] + for _, rtype, contexts in buckets: + for ctx in contexts: + merged = dict(ctx or {}) + if rtype and "resource_type" not in merged: + merged["resource_type"] = rtype + contexts_for_preload.append(merged) + resolver.preload(contexts_for_preload) + + labels: dict[str, str] = {} + for raw_id, rtype, contexts in buckets: + ctx = max([dict(c or {}) for c in contexts], key=len, default={}) + if rtype and "resource_type" not in ctx: + ctx["resource_type"] = rtype + eval_key = ctx.get("evaluation_id") or raw_id + labels[str(raw_id)] = resolver.evaluation_name(str(eval_key)) + return labels + + +def _unpriced_usage_condition(): + return and_( + LLMUsageDaily.pricing_rate_id.is_(None), + _usage_row_weight() > 0, + ) + + +def _has_unpriced_usage_column(): + return func.coalesce(func.bool_or(_unpriced_usage_condition()), False).label( + "has_unpriced_usage" + ) + + +def _usage_cost_sum_columns(): + return ( + func.coalesce(func.sum(LLMUsageDaily.input_cost_micro_usd), 0).label( + "input_cost_micro_usd" + ), + func.coalesce(func.sum(LLMUsageDaily.output_cost_micro_usd), 0).label( + "output_cost_micro_usd" + ), + func.coalesce(func.sum(LLMUsageDaily.cache_read_cost_micro_usd), 0).label( + "cache_read_cost_micro_usd" + ), + func.coalesce(func.sum(LLMUsageDaily.cache_creation_cost_micro_usd), 0).label( + "cache_creation_cost_micro_usd" + ), + func.coalesce(func.sum(LLMUsageDaily.reasoning_cost_micro_usd), 0).label( + "reasoning_cost_micro_usd" + ), + func.coalesce(func.sum(LLMUsageDaily.audio_cost_micro_usd), 0).label( + "audio_cost_micro_usd" + ), + func.coalesce(func.sum(LLMUsageDaily.tts_cost_micro_usd), 0).label( + "tts_cost_micro_usd" + ), + func.coalesce(func.sum(LLMUsageDaily.total_cost_micro_usd), 0).label( + "total_cost_micro_usd" + ), + ) + + +def _attach_costs(metrics: dict, *, has_unpriced_usage: bool = False) -> dict: + costs = costs_from_micro( + input_cost_micro_usd=metrics.get("input_cost_micro_usd", 0), + output_cost_micro_usd=metrics.get("output_cost_micro_usd", 0), + cache_read_cost_micro_usd=metrics.get("cache_read_cost_micro_usd", 0), + cache_creation_cost_micro_usd=metrics.get("cache_creation_cost_micro_usd", 0), + reasoning_cost_micro_usd=metrics.get("reasoning_cost_micro_usd", 0), + audio_cost_micro_usd=metrics.get("audio_cost_micro_usd", 0), + tts_cost_micro_usd=metrics.get("tts_cost_micro_usd", 0), + total_cost_micro_usd=metrics.get("total_cost_micro_usd", 0), + has_unpriced_usage=has_unpriced_usage, + ) + return {**metrics, "costs": costs} + + +def _usage_totals_from_row(row) -> dict: + prompt = int(row.prompt_tokens) + completion = int(row.completion_tokens) + metrics = { + "prompt_tokens": prompt, + "completion_tokens": completion, + "total_tokens": prompt + completion, + "cache_read_tokens": int(row.cache_read_tokens), + "cache_creation_tokens": int(row.cache_creation_tokens), + "reasoning_tokens": int(row.reasoning_tokens), + "audio_seconds": int(row.audio_seconds), + "tts_characters": int(row.tts_characters), + "call_count": int(row.call_count), + "input_cost_micro_usd": int(getattr(row, "input_cost_micro_usd", 0) or 0), + "output_cost_micro_usd": int(getattr(row, "output_cost_micro_usd", 0) or 0), + "cache_read_cost_micro_usd": int( + getattr(row, "cache_read_cost_micro_usd", 0) or 0 + ), + "cache_creation_cost_micro_usd": int( + getattr(row, "cache_creation_cost_micro_usd", 0) or 0 + ), + "reasoning_cost_micro_usd": int( + getattr(row, "reasoning_cost_micro_usd", 0) or 0 + ), + "audio_cost_micro_usd": int(getattr(row, "audio_cost_micro_usd", 0) or 0), + "tts_cost_micro_usd": int(getattr(row, "tts_cost_micro_usd", 0) or 0), + "total_cost_micro_usd": int(getattr(row, "total_cost_micro_usd", 0) or 0), + } + return _attach_costs( + metrics, + has_unpriced_usage=bool(getattr(row, "has_unpriced_usage", False)), + ) + + +def _breakdown_metrics_from_tuple(metrics: tuple) -> dict: + prompt = int(metrics[0]) + completion = int(metrics[1]) + base = { + "prompt_tokens": prompt, + "completion_tokens": completion, + "total_tokens": prompt + completion, + "cache_read_tokens": int(metrics[2]), + "cache_creation_tokens": int(metrics[3]), + "reasoning_tokens": int(metrics[4]), + "audio_seconds": int(metrics[5]), + "tts_characters": int(metrics[6]), + "call_count": int(metrics[7]), + "input_cost_micro_usd": int(metrics[8]), + "output_cost_micro_usd": int(metrics[9]), + "cache_read_cost_micro_usd": int(metrics[10]), + "cache_creation_cost_micro_usd": int(metrics[11]), + "reasoning_cost_micro_usd": int(metrics[12]), + "audio_cost_micro_usd": int(metrics[13]), + "tts_cost_micro_usd": int(metrics[14]), + "total_cost_micro_usd": int(metrics[15]), + } + has_unpriced = bool(metrics[16]) if len(metrics) > 16 else False + return _attach_costs(base, has_unpriced_usage=has_unpriced) + + +def _last_updated(db: Session, organization_id: UUID) -> Optional[datetime]: + return ( + db.query(func.max(LLMUsageDaily.updated_at)) + .filter(LLMUsageDaily.organization_id == organization_id) + .scalar() + ) + + +def _summary_aggregate_query( + db: Session, + *, + organization_id: UUID, + start: date, + end: date, + enforced_floor: Optional[date] = None, + workspace_id: Optional[UUID], + product_section: Optional[str], + model: Optional[str], + resource_id: Optional[UUID], + usage_kind: Optional[str], + call_import_id: Optional[UUID], + evaluation_id: Optional[UUID], + evaluation_row_id: Optional[UUID], + dataset: Optional[str] = None, + tag_id: Optional[UUID] = None, +): + return _apply_filters( + db.query( + func.coalesce(func.sum(LLMUsageDaily.prompt_tokens), 0).label("prompt_tokens"), + func.coalesce(func.sum(LLMUsageDaily.completion_tokens), 0).label( + "completion_tokens" + ), + func.coalesce(func.sum(LLMUsageDaily.cache_read_tokens), 0).label( + "cache_read_tokens" + ), + func.coalesce(func.sum(LLMUsageDaily.cache_creation_tokens), 0).label( + "cache_creation_tokens" + ), + func.coalesce(func.sum(LLMUsageDaily.reasoning_tokens), 0).label( + "reasoning_tokens" + ), + func.coalesce(func.sum(LLMUsageDaily.audio_seconds), 0).label("audio_seconds"), + func.coalesce(func.sum(LLMUsageDaily.tts_characters), 0).label("tts_characters"), + func.coalesce(func.sum(LLMUsageDaily.call_count), 0).label("call_count"), + *_usage_cost_sum_columns(), + _has_unpriced_usage_column(), + ), + organization_id=organization_id, + start=start, + end=end, + enforced_floor=enforced_floor, + workspace_id=workspace_id, + product_section=product_section, + model=model, + resource_id=resource_id, + usage_kind=usage_kind, + call_import_id=call_import_id, + evaluation_id=evaluation_id, + evaluation_row_id=evaluation_row_id, + dataset=dataset, + tag_id=tag_id, + db=db, + ) + + +@router.get("/fx-rate", response_model=UsageFxRateResponse) +def get_usage_fx_rate(): + payload = get_usd_inr_rate() + return UsageFxRateResponse( + base=payload["base"], + quote=payload["quote"], + rate=float(payload["rate"]), + as_of=datetime.fromisoformat(payload["as_of"]), + source=payload["source"], + ) + + +@router.get("/summary", response_model=UsageSummaryResponse) +def get_usage_summary( + start: Optional[date] = Query(None), + end: Optional[date] = Query(None), + tz: Optional[str] = Query( + None, + description="IANA timezone for interpreting start/end calendar dates", + ), + workspace_id: Optional[UUID] = Query(None), + product_section: Optional[str] = Query(None), + model: Optional[str] = Query(None), + resource_id: Optional[UUID] = Query(None), + usage_kind: Optional[str] = Query(None), + call_import_id: Optional[UUID] = Query(None), + evaluation_id: Optional[UUID] = Query(None), + evaluation_row_id: Optional[UUID] = Query(None), + dataset: Optional[str] = Query(None), + tag_id: Optional[UUID] = Query(None), + organization_id: UUID = Depends(get_organization_id), + db: Session = Depends(get_db), +): + access = UsageAccessPolicy.resolve(organization_id, start, end, tz) + display_start = access.display_start + display_end = access.display_end + if display_end < display_start: + raise HTTPException(status_code=400, detail="end must be >= start") + + cache_key = cache_key_for( + access, + workspace_id=workspace_id, + product_section=product_section, + model=model, + resource_id=resource_id, + usage_kind=usage_kind, + call_import_id=call_import_id, + evaluation_id=evaluation_id, + evaluation_row_id=evaluation_row_id, + dataset=dataset, + tag_id=tag_id, + ) + cached = get_cached_response(organization_id, "summary", cache_key) + if cached is not None: + try: + return UsageSummaryResponse.model_validate(cached) + except Exception: + pass + + row = _summary_aggregate_query( + db, + organization_id=organization_id, + start=access.filter_start, + end=access.filter_end, + enforced_floor=access.enforced_filter_floor, + workspace_id=workspace_id, + product_section=product_section, + model=model, + resource_id=resource_id, + usage_kind=usage_kind, + call_import_id=call_import_id, + evaluation_id=evaluation_id, + evaluation_row_id=evaluation_row_id, + dataset=dataset, + tag_id=tag_id, + ).one() + totals = _usage_totals_from_row(row) + response = UsageSummaryResponse( + start=display_start, + end=display_end, + totals=UsageTotals(**totals), + usage_policy=_usage_policy_meta(access), + last_updated_at=_last_updated(db, organization_id), + ) + set_cached_response( + organization_id, + "summary", + cache_key, + response.model_dump(mode="json"), + ) + return response + + +@router.get("/breakdown", response_model=UsageBreakdownResponse) +def get_usage_breakdown( + start: Optional[date] = Query(None), + end: Optional[date] = Query(None), + tz: Optional[str] = Query( + None, + description="IANA timezone for interpreting start/end calendar dates", + ), + group_by: GroupBy = Query("workspace"), + workspace_id: Optional[UUID] = Query(None), + product_section: Optional[str] = Query(None), + model: Optional[str] = Query(None), + resource_id: Optional[UUID] = Query(None), + usage_kind: Optional[str] = Query(None), + call_import_id: Optional[UUID] = Query(None), + evaluation_id: Optional[UUID] = Query(None), + evaluation_row_id: Optional[UUID] = Query(None), + dataset: Optional[str] = Query(None), + tag_id: Optional[UUID] = Query(None), + limit: int = Query(100, ge=1, le=500), + offset: int = Query(0, ge=0), + organization_id: UUID = Depends(get_organization_id), + db: Session = Depends(get_db), +): + access = UsageAccessPolicy.resolve(organization_id, start, end, tz) + display_start = access.display_start + display_end = access.display_end + if display_end < display_start: + raise HTTPException(status_code=400, detail="end must be >= start") + + cache_key = cache_key_for( + access, + group_by=group_by, + workspace_id=workspace_id, + product_section=product_section, + model=model, + resource_id=resource_id, + usage_kind=usage_kind, + call_import_id=call_import_id, + evaluation_id=evaluation_id, + evaluation_row_id=evaluation_row_id, + dataset=dataset, + tag_id=tag_id, + limit=limit, + offset=offset, + ) + cached = get_cached_response(organization_id, "breakdown", cache_key) + if cached is not None: + try: + return UsageBreakdownResponse.model_validate(cached) + except Exception: + pass + + dim = { + "workspace": LLMUsageDaily.workspace_id, + "product_section": LLMUsageDaily.product_section, + "model": LLMUsageDaily.model, + "resource": cast(_resource_id_expr(), String), + "usage_kind": LLMUsageDaily.usage_kind, + "call_import": _call_import_group_expr(), + }[group_by] + + aggregates = [ + func.coalesce(func.sum(LLMUsageDaily.prompt_tokens), 0).label("prompt_tokens"), + func.coalesce(func.sum(LLMUsageDaily.completion_tokens), 0).label( + "completion_tokens" + ), + func.coalesce(func.sum(LLMUsageDaily.cache_read_tokens), 0).label( + "cache_read_tokens" + ), + func.coalesce(func.sum(LLMUsageDaily.cache_creation_tokens), 0).label( + "cache_creation_tokens" + ), + func.coalesce(func.sum(LLMUsageDaily.reasoning_tokens), 0).label( + "reasoning_tokens" + ), + func.coalesce(func.sum(LLMUsageDaily.audio_seconds), 0).label("audio_seconds"), + func.coalesce(func.sum(LLMUsageDaily.tts_characters), 0).label("tts_characters"), + func.coalesce(func.sum(LLMUsageDaily.call_count), 0).label("call_count"), + *_usage_cost_sum_columns(), + _has_unpriced_usage_column(), + ] + + select_cols = [dim] + group_cols = [dim] + if group_by == "resource": + select_cols.append(cast(_resource_type_expr(), String)) + group_cols.append(cast(_resource_type_expr(), String)) + select_cols.append(LLMUsageDaily.product_section) + group_cols.append(LLMUsageDaily.product_section) + + query = _apply_filters( + db.query(*select_cols, *aggregates), + organization_id=organization_id, + start=access.filter_start, + end=access.filter_end, + enforced_floor=access.enforced_filter_floor, + workspace_id=workspace_id, + product_section=product_section, + model=model, + resource_id=resource_id, + usage_kind=usage_kind, + call_import_id=call_import_id, + evaluation_id=evaluation_id, + evaluation_row_id=evaluation_row_id, + dataset=dataset, + tag_id=tag_id, + db=db, + ).group_by(*group_cols) + + results = ( + query.order_by(func.sum(LLMUsageDaily.call_count).desc()) + .offset(offset) + .limit(limit) + .all() + ) + + workspace_names: dict = {} + if group_by == "workspace": + ws_ids = {r[0] for r in results if r[0] is not None} + if ws_ids: + workspace_names = { + w.id: w.name + for w in db.query(Workspace) + .filter( + Workspace.organization_id == organization_id, + Workspace.id.in_(ws_ids), + ) + .all() + } + resource_labels: dict[str, str] = {} + call_import_labels: dict[str, str] = {} + if group_by == "resource": + label_query = _filtered_query( + db, + organization_id=organization_id, + start=access.filter_start, + end=access.filter_end, + enforced_floor=access.enforced_filter_floor, + workspace_id=workspace_id, + product_section=product_section, + model=model, + resource_id=resource_id, + usage_kind=usage_kind, + call_import_id=call_import_id, + evaluation_id=evaluation_id, + evaluation_row_id=evaluation_row_id, + ) + resource_labels = _resource_label_map(db, organization_id, label_query) + elif group_by == "call_import": + label_query = _filtered_query( + db, + organization_id=organization_id, + start=access.filter_start, + end=access.filter_end, + enforced_floor=access.enforced_filter_floor, + workspace_id=workspace_id, + product_section=product_section, + model=model, + resource_id=resource_id, + usage_kind=usage_kind, + call_import_id=call_import_id, + evaluation_id=evaluation_id, + evaluation_row_id=evaluation_row_id, + ) + call_import_labels = _call_import_label_map(db, organization_id, label_query) + + rows: List[UsageBreakdownRow] = [] + for result in results: + if group_by == "workspace": + ws_id = result[0] + metrics = result[1:] + rows.append( + UsageBreakdownRow( + workspace_id=ws_id, + workspace_name=workspace_names.get(ws_id) if ws_id else "Unknown", + **_breakdown_metrics_from_tuple(metrics), + ) + ) + elif group_by == "product_section": + section = result[0] + metrics = result[1:] + rows.append( + UsageBreakdownRow( + product_section=section, + product_section_label=SECTION_LABELS.get(section or "", section), + **_breakdown_metrics_from_tuple(metrics), + ) + ) + elif group_by == "model": + model_name = result[0] + metrics = result[1:] + rows.append( + UsageBreakdownRow( + model=model_name, + **_breakdown_metrics_from_tuple(metrics), + ) + ) + elif group_by == "usage_kind": + kind = result[0] + metrics = result[1:] + rows.append( + UsageBreakdownRow( + usage_kind=kind, + **_breakdown_metrics_from_tuple(metrics), + ) + ) + elif group_by == "call_import": + raw_cid = result[0] + metrics = result[1:] + cid = None + if raw_cid: + try: + cid = UUID(str(raw_cid)) + except (ValueError, TypeError): + cid = None + label = ( + call_import_labels.get(str(raw_cid), "Unscoped") + if raw_cid + else "Unscoped" + ) + rows.append( + UsageBreakdownRow( + call_import_id=cid, + call_import_label=label, + **_breakdown_metrics_from_tuple(metrics), + ) + ) + else: + raw_res_id, res_type, section = result[0], result[1], result[2] + metrics = result[3:] + res_id = None + if raw_res_id: + try: + res_id = UUID(str(raw_res_id)) + except (ValueError, TypeError): + res_id = None + rows.append( + UsageBreakdownRow( + resource_id=res_id, + resource_type=res_type or ( + "agent" if section == "agents" and raw_res_id else None + ), + resource_label=_breakdown_resource_label( + raw_res_id, + res_type, + section, + resource_labels, + db, + organization_id, + ), + product_section=section, + product_section_label=SECTION_LABELS.get(section or "", section), + **_breakdown_metrics_from_tuple(metrics), + ) + ) + + response = UsageBreakdownResponse( + start=display_start, + end=display_end, + group_by=group_by, + rows=rows, + total_count=len(rows), + truncated_at_limit=len(rows) >= limit, + usage_policy=_usage_policy_meta(access), + last_updated_at=_last_updated(db, organization_id), + ) + set_cached_response( + organization_id, + "breakdown", + cache_key, + response.model_dump(mode="json"), + ) + return response + + +@router.get("/filters", response_model=UsageFiltersResponse) +def get_usage_filters( + start: Optional[date] = Query(None), + end: Optional[date] = Query(None), + tz: Optional[str] = Query( + None, + description="IANA timezone for interpreting start/end calendar dates", + ), + workspace_id: Optional[UUID] = Query(None), + product_section: Optional[str] = Query(None), + model: Optional[str] = Query(None), + resource_id: Optional[UUID] = Query(None), + usage_kind: Optional[str] = Query(None), + call_import_id: Optional[UUID] = Query(None), + evaluation_id: Optional[UUID] = Query(None), + dataset: Optional[str] = Query(None), + tag_id: Optional[UUID] = Query(None), + q: Optional[str] = Query(None, description="Optional resource label search"), + organization_id: UUID = Depends(get_organization_id), + db: Session = Depends(get_db), +): + access = UsageAccessPolicy.resolve(organization_id, start, end, tz) + filter_start = access.filter_start + filter_end = access.filter_end + enforced_floor = access.enforced_filter_floor + + cache_key = cache_key_for( + access, + workspace_id=workspace_id, + product_section=product_section, + model=model, + resource_id=resource_id, + usage_kind=usage_kind, + call_import_id=call_import_id, + evaluation_id=evaluation_id, + dataset=dataset, + tag_id=tag_id, + q=q, + ) + cached = get_cached_response(organization_id, "filters", cache_key) + if cached is not None: + try: + return UsageFiltersResponse.model_validate(cached) + except Exception: + pass + + scoped_resource_id = resource_id or evaluation_id + + workspace_base = _filtered_query( + db, + organization_id=organization_id, + start=filter_start, + end=filter_end, + enforced_floor=enforced_floor, + dataset=dataset, + tag_id=tag_id, + ) + section_base = _filtered_query( + db, + organization_id=organization_id, + start=filter_start, + end=filter_end, + enforced_floor=enforced_floor, + workspace_id=workspace_id, + dataset=dataset, + tag_id=tag_id, + ) + kind_base = _filtered_query( + db, + organization_id=organization_id, + start=filter_start, + end=filter_end, + enforced_floor=enforced_floor, + workspace_id=workspace_id, + product_section=product_section, + resource_id=scoped_resource_id, + call_import_id=call_import_id, + evaluation_id=evaluation_id, + dataset=dataset, + tag_id=tag_id, + ) + model_base = _filtered_query( + db, + organization_id=organization_id, + start=filter_start, + end=filter_end, + enforced_floor=enforced_floor, + workspace_id=workspace_id, + product_section=product_section, + resource_id=scoped_resource_id, + usage_kind=usage_kind, + call_import_id=call_import_id, + evaluation_id=evaluation_id, + dataset=dataset, + tag_id=tag_id, + ) + call_import_base = _filtered_query( + db, + organization_id=organization_id, + start=filter_start, + end=filter_end, + enforced_floor=enforced_floor, + workspace_id=workspace_id, + product_section=product_section, + dataset=dataset, + tag_id=tag_id, + ) + resource_base = _filtered_query( + db, + organization_id=organization_id, + start=filter_start, + end=filter_end, + enforced_floor=enforced_floor, + workspace_id=workspace_id, + product_section=product_section, + model=model, + usage_kind=usage_kind, + call_import_id=call_import_id, + dataset=dataset, + tag_id=tag_id, + ) + + workspace_ids = { + row[0] + for row in workspace_base.with_entities(LLMUsageDaily.workspace_id).distinct().all() + if row[0] is not None + } + workspaces = [ + {"id": str(w.id), "name": w.name} + for w in db.query(Workspace) + .filter(Workspace.id.in_(workspace_ids) if workspace_ids else False) + .order_by(Workspace.name) + .all() + ] + + sections = sorted( + { + row[0] + for row in section_base.with_entities(LLMUsageDaily.product_section).distinct().all() + if row[0] + } + ) + models = sorted( + { + row[0] + for row in model_base.with_entities(LLMUsageDaily.model).distinct().all() + if row[0] + } + ) + kinds = sorted( + { + row[0] + for row in kind_base.with_entities(LLMUsageDaily.usage_kind).distinct().all() + if row[0] + } + ) + + call_import_labels = _call_import_filter_labels( + db, + organization_id, + call_import_base, + workspace_id=workspace_id, + dataset=dataset, + tag_id=tag_id, + ) + call_imports = [ + {"id": cid, "label": label} + for cid, label in sorted(call_import_labels.items(), key=lambda x: x[1].lower()) + ] + + evaluation_labels = _evaluation_label_map(db, organization_id, resource_base) + needle = (q or "").strip().lower() + evaluations = [] + for rid, label in sorted(evaluation_labels.items(), key=lambda x: x[1].lower()): + if needle and needle not in label.lower(): + continue + evaluations.append({"id": rid, "label": label}) + + resource_labels = _resource_label_map(db, organization_id, resource_base) + resource_meta = _resource_filter_meta_map(db, organization_id, resource_base) + resources = [] + for rid, label in sorted(resource_labels.items(), key=lambda x: x[1].lower()): + if needle and needle not in label.lower(): + continue + info = resource_meta.get(rid, {}) + rtype = info.get("type") + if rtype == "call_import_evaluation": + continue + resources.append( + { + "id": rid, + "label": label, + "type": rtype, + "product_section": info.get("product_section"), + } + ) + + dataset_query = db.query(CallImport.dataset).filter( + CallImport.organization_id == organization_id, + CallImport.dataset.isnot(None), + CallImport.dataset != "", + ) + if workspace_id is not None: + dataset_query = dataset_query.filter(CallImport.workspace_id == workspace_id) + if dataset or tag_id is not None: + scoped_import_ids = _call_import_ids_for_filters( + db, + organization_id=organization_id, + workspace_id=workspace_id, + dataset=dataset, + tag_id=tag_id, + ) + if scoped_import_ids: + dataset_query = dataset_query.filter(CallImport.id.in_(scoped_import_ids)) + else: + dataset_query = dataset_query.filter(CallImport.id.is_(None)) + datasets = sorted({row[0] for row in dataset_query.distinct().all() if row[0]}) + + tags = [ + {"id": str(tag.id), "label": tag.name} + for tag in ( + db.query(CallImportTag) + .filter(CallImportTag.organization_id == organization_id) + .order_by(CallImportTag.name) + .all() + ) + ] + + response = UsageFiltersResponse( + workspaces=workspaces, + product_sections=[ + {"id": s, "label": SECTION_LABELS.get(s, s)} for s in sections + ], + call_imports=call_imports, + evaluations=evaluations, + models=models, + resources=resources, + usage_kinds=[{"id": k, "label": usage_kind_label(k)} for k in kinds], + datasets=datasets, + tags=tags, + ) + set_cached_response( + organization_id, + "filters", + cache_key, + response.model_dump(mode="json"), + ) + return response + diff --git a/app/api/v1/routes/playground.py b/app/api/v1/routes/playground.py index 8671d7ba..31b1a3c2 100644 --- a/app/api/v1/routes/playground.py +++ b/app/api/v1/routes/playground.py @@ -38,215 +38,9 @@ from app.services.evaluators.evaluator_result_call_data import slim_call_data_for_evaluator_result from app.utils.call_recordings import generate_unique_call_short_id -router = APIRouter(prefix="/playground", tags=["playground"]) - - -def extract_transcript_from_call_data(call_data: Dict[str, Any], provider_platform: str) -> tuple: - """ - Extract transcript and speaker segments from provider call_data. - - Args: - call_data: Full call data from voice provider - provider_platform: The provider platform ("vapi", "retell", "elevenlabs", "smallest") - - Returns: - Tuple of (transcript_text, speaker_segments) - - transcript_text: Plain text transcript - - speaker_segments: List of segments with speaker labels - """ - transcript_text = "" - speaker_segments = [] - - if not call_data: - return transcript_text, speaker_segments - - provider_platform_lower = provider_platform.lower() if provider_platform else "" - - if provider_platform_lower == "vapi": - # Vapi: keep provider payload raw and derive transcript from transcript/messages. - transcript_text = call_data.get("transcript", "") - - # Get structured messages for speaker segments - transcript_object = call_data.get("transcript_object", []) - if not transcript_object: - # Try messages array - artifact = call_data.get("artifact", {}) if isinstance(call_data, dict) else {} - messages = call_data.get("messages", []) or artifact.get("messages", []) - for msg in messages: - role = msg.get("role", "unknown") - content = msg.get("message", "") or msg.get("content", "") - - if not content or role == "system": - continue - - # Map roles - if role in ["bot", "assistant"]: - normalized_role = "agent" - elif role == "user": - normalized_role = "user" - else: - continue - - speaker_segments.append({ - "speaker": "Agent" if normalized_role == "agent" else "User", - "text": content, - "start": msg.get("secondsFromStart", 0), - "end": msg.get("secondsFromStart", 0) + (msg.get("duration", 0) / 1000), - }) - else: - for entry in transcript_object: - role = entry.get("role", "unknown") - content = entry.get("content", "") - - if not content: - continue - - speaker_segments.append({ - "speaker": "Agent" if role == "agent" else "User", - "text": content, - "start": entry.get("seconds_from_start", 0), - "end": entry.get("seconds_from_start", 0) + (entry.get("duration_ms", 0) / 1000), - }) - - # Build transcript text from segments if not available - if not transcript_text and speaker_segments: - transcript_text = "\n".join([ - f"{seg['speaker']}: {seg['text']}" for seg in speaker_segments - ]) - - elif provider_platform_lower == "elevenlabs": - raw_transcript = call_data.get("transcript") - transcript_obj = call_data.get("transcript_object", []) - - # retrieve_call_metrics already processes the transcript into a - # formatted string + speaker_segments list, so handle both the - # pre-processed shape and the raw ElevenLabs API shape. - if isinstance(raw_transcript, str) and raw_transcript: - transcript_text = raw_transcript - if isinstance(transcript_obj, list): - for seg in transcript_obj: - speaker_segments.append({ - "speaker": seg.get("speaker", "Unknown"), - "text": seg.get("text", ""), - "start": seg.get("start", 0), - "end": seg.get("end", 0), - }) - elif isinstance(raw_transcript, list): - for entry in raw_transcript: - role = entry.get("role", "unknown") - content = entry.get("message", "") or entry.get("text", "") - if not content: - continue - speaker = "Agent" if role in ("agent", "assistant", "ai") else "User" - speaker_segments.append({ - "speaker": speaker, - "text": content, - "start": entry.get("time_in_call_secs", 0) or entry.get("start", 0), - "end": entry.get("time_in_call_secs", 0) or entry.get("end", 0), - }) - transcript_text = "\n".join( - f"{seg['speaker']}: {seg['text']}" for seg in speaker_segments - ) - - elif provider_platform_lower == "smallest": - transcript_raw = call_data.get("transcript") - transcript_object = call_data.get("transcript_object", []) - if isinstance(transcript_object, list) and transcript_object: - for entry in transcript_object: - if not isinstance(entry, dict): - continue - text = entry.get("text", "") - if not text: - continue - speaker = entry.get("speaker", "Unknown") - speaker_segments.append( - { - "speaker": speaker, - "text": text, - "start": entry.get("start", 0), - "end": entry.get("end", entry.get("start", 0)), - } - ) - if not transcript_text: - transcript_text = "\n".join( - f"{seg['speaker']}: {seg['text']}" for seg in speaker_segments - ) - elif isinstance(transcript_raw, list): - for entry in transcript_raw: - if not isinstance(entry, dict): - continue - role = str(entry.get("speaker") or entry.get("role") or "").lower() - speaker = "Agent" if role in ("agent", "assistant", "ai", "bot") else "User" - text = entry.get("text", "") or entry.get("message", "") or entry.get("content", "") - if not text: - continue - ts = entry.get("timeInCallSecs", 0) or entry.get("start", 0) or entry.get("timestamp", 0) - speaker_segments.append( - { - "speaker": speaker, - "text": text, - "start": ts, - "end": entry.get("end", ts), - } - ) - transcript_text = "\n".join( - f"{seg['speaker']}: {seg['text']}" for seg in speaker_segments - ) - elif isinstance(transcript_raw, str): - transcript_text = transcript_raw - - elif provider_platform_lower == "retell": - # Retell: transcript can be a string or list of objects - transcript_raw = call_data.get("transcript", "") - - if isinstance(transcript_raw, str): - transcript_text = transcript_raw - # Parse transcript text into speaker segments if it has pattern like "Agent: text\nUser: text" - lines = transcript_raw.split("\n") if transcript_raw else [] - for line in lines: - line = line.strip() - if not line: - continue - if line.startswith("Agent:") or line.startswith("agent:"): - speaker_segments.append({ - "speaker": "Agent", - "text": line.split(":", 1)[1].strip() if ":" in line else line, - "start": 0, - "end": 0, - }) - elif line.startswith("User:") or line.startswith("user:"): - speaker_segments.append({ - "speaker": "User", - "text": line.split(":", 1)[1].strip() if ":" in line else line, - "start": 0, - "end": 0, - }) - elif isinstance(transcript_raw, list): - # Retell sometimes returns transcript as array of objects - for item in transcript_raw: - if isinstance(item, dict): - role = item.get("role", "") - content = item.get("content", "") or item.get("text", "") - - if not content: - continue - - speaker = "Agent" if role in ["agent", "assistant", "bot"] else "User" - speaker_segments.append({ - "speaker": speaker, - "text": content, - "start": item.get("start_time", 0) or item.get("timestamp", 0), - "end": item.get("end_time", 0), - }) - - # Build transcript text from segments - transcript_text = "\n".join([ - f"{seg['speaker']}: {seg['text']}" for seg in speaker_segments - ]) - - return transcript_text, speaker_segments - +from app.services.evaluators.call_data_transcript import extract_transcript_from_call_data +router = APIRouter(prefix="/playground", tags=["playground"]) def generate_unique_result_id(db: Session) -> str: """Generate a unique 6-digit result ID for EvaluatorResult.""" max_attempts = 100 @@ -411,7 +205,8 @@ def poll_call_metrics( or recording_urls.get("stereo_url") ) if audio_url: - resp = _http.get(audio_url, timeout=120) + vapi_headers = {"Authorization": f"Bearer {integration_api_key}"} + resp = _http.get(audio_url, headers=vapi_headers, timeout=120) if resp.status_code == 200: audio_bytes = resp.content @@ -656,13 +451,32 @@ async def create_web_call( # For now, we'll skip it for Retell. Other providers can handle it in their implementation. if integration.platform != "retell" and web_call_data.custom_sip_headers: call_params["custom_sip_headers"] = web_call_data.custom_sip_headers - - web_call_response = provider.create_web_call(**call_params) - - # Store call recording in database + + platform_value = ( + integration.platform.value + if hasattr(integration.platform, "value") + else integration.platform + ) + plat_lower = str(platform_value).lower() + + # Vapi Web SDK creates the call in the browser; server-side /call/web + # would spawn a second call that never receives the user's microphone. + if plat_lower == "vapi": + from app.services.voice_providers.vapi import VAPI_SAMPLE_RATE + + web_call_response = { + "call_type": "web_call", + "agent_id": agent.voice_ai_agent_id, + "metadata": web_call_data.metadata or {}, + "sample_rate": VAPI_SAMPLE_RATE, + "client_sdk_creates_call": True, + } + provider_call_id = None + else: + web_call_response = provider.create_web_call(**call_params) + provider_call_id = web_call_response.get("call_id") + call_short_id = generate_unique_call_short_id(db) - provider_call_id = web_call_response.get("call_id") - call_recording = CallRecording( organization_id=organization_id, workspace_id=workspace_id, @@ -702,17 +516,13 @@ async def create_web_call( # Add call_short_id to response for frontend response = web_call_response.copy() response["call_short_id"] = call_short_id - - platform_value = integration.platform.value if hasattr(integration.platform, 'value') else integration.platform - - # For Vapi, include the public key in the response (needed for frontend SDK) - if platform_value.lower() == "vapi" and integration.public_key: + + if plat_lower == "vapi" and integration.public_key: response["public_key"] = integration.public_key - - # For ElevenLabs, pass through the signed_url (frontend SDK connects directly) - if platform_value.lower() == "elevenlabs": + + if plat_lower == "elevenlabs": response["signed_url"] = web_call_response.get("signed_url") - + return response except Exception as e: raise HTTPException( @@ -1867,7 +1677,10 @@ async def summarize_transcript( 1. The voice bundle of ``agent_id`` (or the agent on ``call_short_id``). 2. Any configured AIProvider matching the fallback preference list. """ + from contextlib import nullcontext + from app.services.ai.llm_service import llm_service + from app.services.usage.context import llm_usage_context, usage_context_for_agent transcript_text = (payload.transcript or "").strip() if not transcript_text and payload.entries: @@ -1964,15 +1777,27 @@ async def summarize_transcript( ] try: - result = llm_service.generate_response( - messages=messages, - llm_provider=llm_provider, - llm_model=llm_model, - organization_id=organization_id, - db=db, - temperature=0.3, - max_tokens=400, - ) + usage_ctx = nullcontext() + if agent_uuid: + agent_row = db.query(Agent).filter( + Agent.id == agent_uuid, + Agent.organization_id == organization_id, + ).first() + if agent_row: + usage_ctx = llm_usage_context( + usage_context_for_agent(agent_row, workspace_id=workspace_id) + ) + + with usage_ctx: + result = llm_service.generate_response( + messages=messages, + llm_provider=llm_provider, + llm_model=llm_model, + organization_id=organization_id, + db=db, + temperature=0.3, + max_tokens=400, + ) except Exception as e: logger.error(f"[summarize-transcript] LLM call failed: {e}") raise HTTPException( diff --git a/app/api/v1/routes/settings.py b/app/api/v1/routes/settings.py index 3c49c0cc..83225380 100644 --- a/app/api/v1/routes/settings.py +++ b/app/api/v1/routes/settings.py @@ -32,6 +32,7 @@ is_feature_enabled, ENTERPRISE_FEATURES, ) +from app.core.usage_entitlement import get_usage_policy REPORT_LOGO_CONTENT_TYPES = { @@ -85,12 +86,14 @@ def license_info(organization_id: UUID = Depends(get_organization_id)): data = get_license_info() all_licensed = data.get("features", []) if isinstance(data.get("features"), list) else [] enabled_for_org = [f for f in all_licensed if is_feature_enabled(f, organization_id)] + usage_policy = get_usage_policy(organization_id) return { "is_enterprise": bool(enabled_for_org), "enabled_features": enabled_for_org, "all_enterprise_features": ENTERPRISE_FEATURES, "feature_catalog": get_feature_catalog(), "organization": data.get("org_id"), + "usage_policy": usage_policy.as_dict(), } diff --git a/app/api/v1/routes/usage_pricing.py b/app/api/v1/routes/usage_pricing.py new file mode 100644 index 00000000..dc868b19 --- /dev/null +++ b/app/api/v1/routes/usage_pricing.py @@ -0,0 +1,271 @@ +"""Org-scoped usage pricing overrides and recompute API.""" + +from __future__ import annotations + +from datetime import date, datetime +from typing import List, Optional +from uuid import UUID + +from fastapi import APIRouter, Depends, HTTPException, Query, status +from pydantic import BaseModel +from sqlalchemy.orm import Session + +from app.core.auth.rbac import require_admin +from app.database import get_db +from app.dependencies import get_organization_id, require_enterprise_entitlement +from app.services.usage.pricing_jobs import ( + create_recompute_job, + enqueue_recompute_job, + get_recompute_job, + job_to_dict, +) +from app.services.usage.pricing_overrides import ( + delete_override, + get_effective_rate, + list_effective_pricing, + list_overrides, + upsert_override, +) + +from app.services.usage.access import UsageAccessPolicy + +router = APIRouter( + prefix="/organizations/usage/pricing", + tags=["Usage"], + dependencies=[Depends(require_admin), Depends(require_enterprise_entitlement())], +) + + +class PricingRatesUsd(BaseModel): + input_per_1m: Optional[float] = None + output_per_1m: Optional[float] = None + cache_read_per_1m: Optional[float] = None + cache_write_per_1m: Optional[float] = None + reasoning_per_1m: Optional[float] = None + audio_per_minute: Optional[float] = None + tts_per_1m_characters: Optional[float] = None + + +class PricingOverrideResponse(BaseModel): + id: str + organization_id: str + model: str + usage_kind: str + effective_from: date + effective_to: Optional[date] = None + rates: PricingRatesUsd + created_at: Optional[datetime] = None + updated_at: Optional[datetime] = None + recompute_enqueued: Optional[bool] = None + recompute_job_id: Optional[str] = None + + +class PricingOverrideUpsertRequest(BaseModel): + usage_kind: str = "llm" + effective_from: date + effective_to: Optional[date] = None + rates: PricingRatesUsd + recompute: bool = False + + +class EffectivePricingResponse(BaseModel): + model: str + usage_kind: str + as_of: date + catalog_rates: Optional[PricingRatesUsd] = None + catalog_rate_id: Optional[str] = None + override: Optional[PricingOverrideResponse] = None + effective_rates: Optional[PricingRatesUsd] = None + effective_source: Optional[str] = None + effective_rate_id: Optional[str] = None + has_override: bool = False + + +class PricingOverrideDeleteResponse(BaseModel): + deleted: bool + model: str + usage_kind: str + recompute_enqueued: bool = False + recompute_job_id: Optional[str] = None + + +class UsageRecomputeRequest(BaseModel): + start_date: Optional[date] = None + end_date: Optional[date] = None + model: Optional[str] = None + usage_kind: Optional[str] = None + + +class UsageRecomputeJobResponse(BaseModel): + id: UUID + organization_id: UUID + status: str + model: Optional[str] = None + usage_kind: Optional[str] = None + start_date: Optional[date] = None + end_date: Optional[date] = None + updated_rows: int = 0 + error_message: Optional[str] = None + celery_task_id: Optional[str] = None + created_at: Optional[datetime] = None + updated_at: Optional[datetime] = None + completed_at: Optional[datetime] = None + + +class AvailableModelsResponse(BaseModel): + models: List[str] + + +@router.get("/available-models", response_model=AvailableModelsResponse) +def list_pricing_available_models( + organization_id: UUID = Depends(get_organization_id), + db: Session = Depends(get_db), +): + from app.services.usage.enabled_models import org_pricing_eligible_models + + return AvailableModelsResponse( + models=org_pricing_eligible_models(db, organization_id), + ) + + +@router.get("", response_model=List[EffectivePricingResponse]) +def list_effective_usage_pricing( + usage_kind: Optional[str] = Query(None), + model: Optional[str] = Query(None), + as_of: Optional[date] = Query(None), + limit: int = Query(200, ge=1, le=1000), + organization_id: UUID = Depends(get_organization_id), + db: Session = Depends(get_db), +): + day = as_of or date.today() + rows = list_effective_pricing( + db, + organization_id=organization_id, + usage_kind=usage_kind, + model=model, + as_of=day, + limit=limit, + ) + return rows + + +@router.get("/overrides", response_model=List[PricingOverrideResponse]) +def list_usage_pricing_overrides( + model: Optional[str] = Query(None), + usage_kind: Optional[str] = Query(None), + organization_id: UUID = Depends(get_organization_id), + db: Session = Depends(get_db), +): + return list_overrides( + db, + organization_id=organization_id, + model=model, + usage_kind=usage_kind, + ) + + +@router.get("/overrides/{model}", response_model=EffectivePricingResponse) +def get_usage_pricing_override_effective( + model: str, + usage_kind: str = Query("llm"), + as_of: Optional[date] = Query(None), + organization_id: UUID = Depends(get_organization_id), + db: Session = Depends(get_db), +): + return get_effective_rate( + db, + organization_id=organization_id, + model=model, + usage_kind=usage_kind, + as_of=as_of or date.today(), + ) + + +@router.put("/overrides/{model}", response_model=PricingOverrideResponse) +def upsert_usage_pricing_override( + model: str, + body: PricingOverrideUpsertRequest, + organization_id: UUID = Depends(get_organization_id), + db: Session = Depends(get_db), +): + return upsert_override( + db, + organization_id=organization_id, + model=model, + usage_kind=body.usage_kind, + effective_from=body.effective_from, + effective_to=body.effective_to, + rates=body.rates.model_dump(exclude_unset=True), + recompute=body.recompute, + ) + + +@router.delete("/overrides/{model}", response_model=PricingOverrideDeleteResponse) +def delete_usage_pricing_override( + model: str, + usage_kind: str = Query("llm"), + effective_from: Optional[date] = Query(None), + recompute: bool = Query(False), + organization_id: UUID = Depends(get_organization_id), + db: Session = Depends(get_db), +): + return delete_override( + db, + organization_id=organization_id, + model=model, + usage_kind=usage_kind, + effective_from=effective_from, + recompute=recompute, + ) + + +@router.post( + "/recompute", + response_model=UsageRecomputeJobResponse, + status_code=status.HTTP_202_ACCEPTED, +) +def trigger_usage_cost_recompute( + body: UsageRecomputeRequest, + organization_id: UUID = Depends(get_organization_id), + db: Session = Depends(get_db), +): + if ( + body.start_date is not None + and body.end_date is not None + and body.end_date < body.start_date + ): + raise HTTPException(status_code=400, detail="end_date must be >= start_date") + + clamped_start = body.start_date + clamped_end = body.end_date + if body.start_date is not None or body.end_date is not None: + access = UsageAccessPolicy.resolve( + organization_id, + body.start_date, + body.end_date, + None, + ) + clamped_start = access.display_start if body.start_date is not None else None + clamped_end = access.display_end if body.end_date is not None else None + + job = create_recompute_job( + db, + organization_id=organization_id, + model=body.model, + usage_kind=body.usage_kind, + start_date=clamped_start, + end_date=clamped_end, + ) + enqueue_recompute_job(db, job) + db.refresh(job) + return UsageRecomputeJobResponse(**job_to_dict(job)) + + +@router.get("/recompute/{job_id}", response_model=UsageRecomputeJobResponse) +def get_usage_cost_recompute_job( + job_id: UUID, + organization_id: UUID = Depends(get_organization_id), + db: Session = Depends(get_db), +): + job = get_recompute_job(db, organization_id=organization_id, job_id=job_id) + return UsageRecomputeJobResponse(**job_to_dict(job)) diff --git a/app/api/v1/routes/vobiz_telephony.py b/app/api/v1/routes/vobiz_telephony.py index 084d03b7..052d03b9 100644 --- a/app/api/v1/routes/vobiz_telephony.py +++ b/app/api/v1/routes/vobiz_telephony.py @@ -698,6 +698,7 @@ async def vobiz_media_websocket(websocket: WebSocket): websocket, context.system_instruction, str(context.organization_id), + str(context.workspace_id) if context.workspace_id else None, agent_id, persona_id, scenario_id, diff --git a/app/api/v1/routes/voice_agent.py b/app/api/v1/routes/voice_agent.py index 47f7350e..e6d553b5 100644 --- a/app/api/v1/routes/voice_agent.py +++ b/app/api/v1/routes/voice_agent.py @@ -517,6 +517,7 @@ def resolve_azure_endpoint_for_provider(provider: ModelProvider) -> str | None: websocket, system_instruction, str(organization_id), + str(workspace_id) if workspace_id else None, agent_id, persona_id, scenario_id, diff --git a/app/api/v1/routes/voice_playground.py b/app/api/v1/routes/voice_playground.py index 8cc4519d..40477c25 100644 --- a/app/api/v1/routes/voice_playground.py +++ b/app/api/v1/routes/voice_playground.py @@ -398,6 +398,7 @@ def _serialize_custom_voice(voice: CustomTTSVoice) -> Dict[str, Any]: async def generate_sample_texts( data: GenerateSamplesRequest, organization_id: UUID = Depends(get_organization_id), + workspace_id: UUID = Depends(get_workspace_id), api_key: str = Depends(get_api_key), db: Session = Depends(get_db), ): @@ -460,16 +461,29 @@ async def generate_sample_texts( {"role": "user", "content": user_prompt}, ] + from app.services.usage.context import ( + LLMUsageContext, + LLMUsageProductSection, + llm_usage_context, + ) + try: - result = llm_service.generate_response( - messages=messages, - llm_provider=provider_enum, - llm_model=llm_model_str, - organization_id=organization_id, - db=db, - llm_config=request_llm_config, - task_defaults={"temperature": 0.8, "max_tokens": max_tokens}, - ) + with llm_usage_context( + LLMUsageContext( + organization_id=organization_id, + workspace_id=workspace_id, + product_section=LLMUsageProductSection.VOICE_PLAYGROUND, + ) + ): + result = llm_service.generate_response( + messages=messages, + llm_provider=provider_enum, + llm_model=llm_model_str, + organization_id=organization_id, + db=db, + llm_config=request_llm_config, + task_defaults={"temperature": 0.8, "max_tokens": max_tokens}, + ) except Exception as e: logger.error(f"[VoicePlayground] LLM generation failed: {e}") raise HTTPException(500, f"LLM generation failed: {str(e)}") diff --git a/app/app_factory.py b/app/app_factory.py index 2b6b8d61..c2af3e09 100644 --- a/app/app_factory.py +++ b/app/app_factory.py @@ -19,6 +19,7 @@ from app.core.operational_access_middleware import OperationalAccessMiddleware from app.core.rbac_middleware import ReaderReadOnlyMiddleware from app.core.security_headers_middleware import SecurityHeadersMiddleware +from app.core.usage_context_middleware import LLMUsageContextMiddleware from app.database import init_db logger = logging.getLogger(__name__) @@ -98,6 +99,7 @@ def _add_common_middleware(app: FastAPI) -> None: if _includes_http_routes(): app.add_middleware(MigrationCheckMiddleware) app.add_middleware(ReaderReadOnlyMiddleware) + app.add_middleware(LLMUsageContextMiddleware) if settings.OBSERVABILITY_ENABLED and settings.LOKI_ENABLED and settings.LOKI_MULTI_TENANT: from app.core.observability_middleware import OrgLoggingMiddleware diff --git a/app/cli.py b/app/cli.py index 3078344f..7fbefbff 100644 --- a/app/cli.py +++ b/app/cli.py @@ -423,6 +423,91 @@ def worker(config: str, loglevel: str, queues: Optional[str], concurrency: Optio sys.exit(1) +@main.command("beat") +@click.option( + "--config", + "-c", + default="config.yml", + help="Path to configuration file", +) +@click.option( + "--loglevel", + "-l", + default="info", + type=click.Choice(["debug", "info", "warning", "error", "critical"], case_sensitive=False), + help="Log level for Celery beat", +) +@click.option( + "--platform-worker-concurrency", + default=2, + type=int, + help="Concurrency for the co-located platform task worker (default: 2; thread pool).", +) +def beat(config: str, loglevel: str, platform_worker_concurrency: int): + """Start Celery Beat + platform task worker (alerts, FX, prune) — single replica only.""" + from app.config import load_config_from_file + + config_path = Path(config) + if not config_path.exists(): + click.echo(f"❌ Config file not found: {config}", err=True) + sys.exit(1) + + try: + load_config_from_file(str(config_path)) + click.echo(f"✅ Loaded configuration from {config_path}") + except Exception as e: + click.echo(f"❌ Error loading config: {e}", err=True) + sys.exit(1) + + from app.workers.config import PLATFORM_WORKER_QUEUE + + click.echo( + "🚀 Starting Celery Beat + platform worker " + f"(queue={PLATFORM_WORKER_QUEUE}, concurrency={platform_worker_concurrency})" + ) + platform_proc = None + try: + import subprocess + + platform_proc = subprocess.Popen( + [ + "celery", + "-A", + "app.workers.celery_app", + "worker", + f"--queues={PLATFORM_WORKER_QUEUE}", + "--pool=threads", + f"--concurrency={platform_worker_concurrency}", + f"--loglevel={loglevel}", + ], + ) + subprocess.run( + [ + "celery", + "-A", + "app.workers.celery_app", + "beat", + f"--loglevel={loglevel}", + ], + check=True, + ) + except KeyboardInterrupt: + click.echo("\n👋 Celery Beat stopped") + except subprocess.CalledProcessError as e: + click.echo(f"❌ Celery Beat failed: {e}", err=True) + sys.exit(1) + except FileNotFoundError: + click.echo("❌ Celery not found. Please install it: pip install celery", err=True) + sys.exit(1) + finally: + if platform_proc is not None and platform_proc.poll() is None: + platform_proc.terminate() + try: + platform_proc.wait(timeout=5) + except subprocess.TimeoutExpired: + platform_proc.kill() + + @main.command("telephony-worker") @click.option( "--config", @@ -623,6 +708,34 @@ def _handle_signal(sig, frame): "is enabled; 32 threads can exhaust per-shard SQLAlchemy pools." ), ) +@click.option( + "--usage-worker/--no-usage-worker", + default=True, + help=( + "Also start a dedicated worker for the `usage` queue (flush + cost recompute; " + "default: True)." + ), +) +@click.option( + "--usage-worker-concurrency", + default=4, + type=int, + help="Concurrency for the usage worker (default: 4; thread pool).", +) +@click.option( + "--beat/--no-beat", + default=True, + help=( + "Start Celery Beat for platform periodic tasks (usage flush, alerts, etc.; " + "default: True — single replica only in production)." + ), +) +@click.option( + "--beat-loglevel", + default=None, + type=click.Choice(["debug", "info", "warning", "error", "critical"], case_sensitive=False), + help="Log level for Celery Beat (defaults to --worker-loglevel).", +) @click.option( "--telephony-worker/--no-telephony-worker", default=True, @@ -646,6 +759,10 @@ def start_all( worker_loglevel: str, imports_worker: bool, imports_worker_concurrency: int, + usage_worker: bool, + usage_worker_concurrency: int, + beat: bool, + beat_loglevel: Optional[str], telephony_worker: bool, media_port: Optional[int], ): @@ -692,6 +809,9 @@ def start_all( # Store worker processes for cleanup. worker_process = None worker_imports_process = None + worker_usage_process = None + beat_process = None + platform_worker_process = None telephony_process = None def _terminate(proc, label: str): @@ -711,10 +831,13 @@ def _terminate(proc, label: str): def cleanup_processes(): """Clean up spawned processes.""" - nonlocal worker_process, worker_imports_process, telephony_process + nonlocal worker_process, worker_imports_process, worker_usage_process, beat_process, platform_worker_process, telephony_process _terminate(telephony_process, "Telephony media server") + _terminate(beat_process, "Celery Beat") + _terminate(platform_worker_process, "Celery platform worker") _terminate(worker_process, "Celery worker (default)") _terminate(worker_imports_process, "Celery worker (imports)") + _terminate(worker_usage_process, "Celery worker (usage)") # Register cleanup on exit atexit.register(cleanup_processes) @@ -867,6 +990,64 @@ def _stream_telephony(): ), prefix="[WORKER-IMPORTS]", ) + + if usage_worker: + from app.workers.config import USAGE_WORKER_QUEUE + + worker_usage_process = _spawn_worker( + [ + "celery", + "-A", + "app.workers.celery_app", + "worker", + f"--loglevel={worker_loglevel}", + "-Q", + USAGE_WORKER_QUEUE, + "-P", + "threads", + "-c", + str(usage_worker_concurrency), + ], + label=( + f"Celery worker ({USAGE_WORKER_QUEUE} queue, " + f"pool=threads, concurrency={usage_worker_concurrency})" + ), + prefix="[WORKER-USAGE]", + ) + + if beat: + from app.workers.config import PLATFORM_WORKER_QUEUE + + beat_level = beat_loglevel or worker_loglevel + platform_worker_process = _spawn_worker( + [ + "celery", + "-A", + "app.workers.celery_app", + "worker", + f"--loglevel={beat_level}", + "-Q", + PLATFORM_WORKER_QUEUE, + "-P", + "threads", + "-c", + "2", + ], + label=f"Celery platform worker ({PLATFORM_WORKER_QUEUE} queue)", + prefix="[BEAT-WORKER]", + ) + beat_process = _spawn_worker( + [ + "celery", + "-A", + "app.workers.celery_app", + "beat", + f"--loglevel={beat_level}", + ], + label=f"Celery Beat (scheduler, loglevel={beat_level})", + prefix="[BEAT]", + ) + except FileNotFoundError: click.echo("❌ Celery not found. Please install it: pip install celery", err=True) sys.exit(1) @@ -906,6 +1087,19 @@ def _stream_telephony(): ) else: click.echo(" Workers: default queue only (--no-imports-worker)") + if usage_worker: + from app.workers.config import USAGE_WORKER_QUEUE + + click.echo( + f" Usage worker: {USAGE_WORKER_QUEUE} queue " + f"(concurrency={usage_worker_concurrency})" + ) + else: + click.echo(" Usage worker: disabled (--no-usage-worker)") + if beat: + click.echo(" Celery Beat: scheduler + platform worker (alerts, FX, prune); flush on worker-usage") + else: + click.echo(" Celery Beat: disabled (--no-beat)") if telephony_worker: telephony_public = (settings.VOBIZ_WEBHOOK_BASE_URL or "").strip() click.echo(f" Telephony edge: http://localhost:{bind_media_port} (local)") @@ -1184,6 +1378,244 @@ def sharding_rebalance_slices( catalog.close() +@click.group() +def usage(): + """Usage pricing ops (seed rates, diff catalog, recompute costs).""" + pass + + +main.add_command(usage) + + +def _load_cli_config(config: str) -> Path: + config_path = Path(config) + if not config_path.exists(): + click.echo(f"❌ Config file not found: {config}", err=True) + sys.exit(1) + from app.config import load_config_from_file + + load_config_from_file(str(config_path)) + return config_path + + +@usage.command("seed-rates") +@click.option( + "--config", + "-c", + type=click.Path(exists=True, readable=True), + default="config.yml", + help="Path to configuration YAML file", +) +@click.option( + "--effective-from", + type=click.DateTime(formats=["%Y-%m-%d"]), + default=None, + help="Effective date for seeded rates (default: 2020-01-01)", +) +def usage_seed_rates(config: str, effective_from): + """Upsert model_pricing_rates from models.json pricing blocks.""" + _load_cli_config(config) + from app.database import SessionLocal + from app.services.usage.pricing_ops import seed_rates_from_models_json + + day = effective_from.date() if effective_from else None + db = SessionLocal() + try: + count = seed_rates_from_models_json(db, effective_from=day) + db.commit() + click.echo(f"✅ Seeded/updated {count} pricing rate row(s)") + finally: + db.close() + + +@usage.command("diff-rates") +@click.option( + "--config", + "-c", + type=click.Path(exists=True, readable=True), + default="config.yml", + help="Path to configuration YAML file", +) +@click.option( + "--effective-from", + type=click.DateTime(formats=["%Y-%m-%d"]), + default=None, + help="Compare rates at this effective_from date (default: 2020-01-01)", +) +@click.option("--json", "as_json", is_flag=True, help="Print machine-readable JSON") +def usage_diff_rates(config: str, effective_from, as_json: bool): + """Diff models.json pricing blocks vs model_pricing_rates in Postgres.""" + import json as json_module + + _load_cli_config(config) + from app.database import SessionLocal + from app.services.usage.pricing_ops import diff_models_json_vs_db + + day = effective_from.date() if effective_from else None + db = SessionLocal() + try: + report = diff_models_json_vs_db(db, effective_from=day) + finally: + db.close() + + if as_json: + click.echo(json_module.dumps(report, indent=2, default=str)) + return + + click.echo(f"effective_from: {report['effective_from']}") + click.echo( + f"models.json priced: {report['models_json_count']} | " + f"database rows: {report['database_count']} | " + f"in_sync: {report['in_sync']}" + ) + if report["only_in_models_json"]: + click.echo(f"\nOnly in models.json ({len(report['only_in_models_json'])}):") + for item in report["only_in_models_json"][:20]: + click.echo(f" - {item['model']} ({item['usage_kind']})") + if report["only_in_database"]: + click.echo(f"\nOnly in database ({len(report['only_in_database'])}):") + for item in report["only_in_database"][:20]: + click.echo(f" - {item['model']} ({item['usage_kind']})") + if report["mismatches"]: + click.echo(f"\nMismatched rates ({len(report['mismatches'])}):") + for item in report["mismatches"][:20]: + click.echo(f" - {item['model']} ({item['usage_kind']})") + for field, values in item["fields"].items(): + click.echo( + f" {field}: json={values['models_json']} db={values['database']}" + ) + missing = report["missing_pricing_blocks"] + if missing: + click.echo(f"\nmodels.json entries missing pricing blocks ({len(missing)}):") + for model in missing[:20]: + click.echo(f" - {model}") + unresolved = report["litellm_unresolved"] + if unresolved: + click.echo(f"\nLiteLLM unresolved ({len(unresolved)}):") + for item in unresolved[:20]: + click.echo(f" - {item.get('model')} ({item.get('reason', 'unresolved')})") + + +@usage.command("recompute") +@click.option( + "--config", + "-c", + type=click.Path(exists=True, readable=True), + default="config.yml", + help="Path to configuration YAML file", +) +@click.option("--organization-id", default=None, help="Scope recompute to one org UUID") +@click.option("--model", default=None, help="Scope recompute to one model") +@click.option("--usage-kind", default=None, help="Scope recompute to llm/stt/tts") +@click.option("--start-date", type=click.DateTime(formats=["%Y-%m-%d"]), default=None) +@click.option("--end-date", type=click.DateTime(formats=["%Y-%m-%d"]), default=None) +@click.option( + "--async/--sync", + "run_async", + default=True, + help="Enqueue Celery task (default) or run synchronously in this process", +) +def usage_recompute( + config: str, + organization_id: Optional[str], + model: Optional[str], + usage_kind: Optional[str], + start_date, + end_date, + run_async: bool, +): + """Backfill or recompute stored usage costs on llm_usage_daily rollups.""" + from uuid import UUID + + _load_cli_config(config) + start = start_date.date() if start_date else None + end = end_date.date() if end_date else None + org_uuid = UUID(organization_id) if organization_id else None + + if run_async: + if org_uuid is None: + click.echo( + "❌ --organization-id is required for async recompute (creates a tracked job).", + err=True, + ) + click.echo( + "💡 Use --sync to recompute in this process without an org scope, or pass --organization-id.", + err=True, + ) + sys.exit(1) + + from app.database import SessionLocal + from app.services.usage.pricing_jobs import ( + create_recompute_job, + enqueue_recompute_job, + job_to_dict, + ) + + db = SessionLocal() + try: + job = create_recompute_job( + db, + organization_id=org_uuid, + model=model, + usage_kind=usage_kind, + start_date=start, + end_date=end, + ) + enqueue_recompute_job(db, job) + db.refresh(job) + payload = job_to_dict(job) + click.echo(f"✅ Enqueued recompute job {payload['id']} (status={payload['status']})") + if payload.get("celery_task_id"): + click.echo(f" Celery task: {payload['celery_task_id']}") + except Exception as exc: + click.echo(f"❌ Failed to enqueue recompute job: {exc}", err=True) + sys.exit(1) + finally: + db.close() + return + + from app.database import SessionLocal + from app.services.usage.pricing import recompute_usage_costs + + db = SessionLocal() + try: + updated = recompute_usage_costs( + db, + organization_id=org_uuid, + model=model, + usage_kind=usage_kind, + start_date=start, + end_date=end, + ) + click.echo(f"✅ Recomputed costs for {updated} rollup row(s)") + finally: + db.close() + + +@usage.command("sync-litellm") +@click.option("--local", is_flag=True, help="Use bundled LiteLLM model_cost JSON") +@click.option( + "--write-models", + is_flag=True, + help="Merge generated pricing into app/config/models.json", +) +@click.option("--stdout", is_flag=True, help="Print pricing_catalog.json to stdout") +def usage_sync_litellm(local: bool, write_models: bool, stdout: bool): + """Fetch LiteLLM prices and regenerate pricing_catalog.json.""" + import subprocess + import sys as sys_module + + script = Path(__file__).resolve().parent.parent / "scripts" / "sync_pricing_catalog_from_litellm.py" + cmd = [sys_module.executable, str(script)] + if local: + cmd.append("--local") + if write_models: + cmd.append("--write-models") + if stdout: + cmd.append("--stdout") + subprocess.run(cmd, check=True) + + if __name__ == "__main__": main() diff --git a/app/config/models.json b/app/config/models.json index d727a14d..38f24cca 100644 --- a/app/config/models.json +++ b/app/config/models.json @@ -1,1029 +1,2025 @@ -{ - "whisper-1": { - "provider": "openai", - "model_type": "stt", - "description": "General-purpose speech recognition (verbose_json + word/segment timestamps)" - }, - "gpt-4o-transcribe": { - "provider": "openai", - "model_type": "stt", - "description": "GPT-4o speech-to-text; higher accuracy, no granular word timestamps" - }, - "gpt-4o-mini-transcribe": { - "provider": "openai", - "model_type": "stt", - "description": "Cheaper / faster GPT-4o transcription; no granular word timestamps" - }, - "gpt-4o-transcribe-diarize": { - "provider": "openai", - "model_type": "stt", - "description": "Transcription model that identifies who's speaking when" - }, - "gpt-realtime-whisper": { - "provider": "openai", - "model_type": "stt", - "description": "Streaming speech-to-text for realtime transcription" - }, - "gpt-5.6": { - "provider": "openai", - "model_type": "llm", - "description": "Preview frontier model for select partners; broad availability coming soon" - }, - "gpt-5.5": { - "provider": "openai", - "model_type": "llm", - "description": "OpenAI flagship — complex reasoning, agentic coding, 1M context" - }, - "gpt-5.5-pro": { - "provider": "openai", - "model_type": "llm", - "description": "Higher-accuracy GPT-5.5 variant with parallel test-time compute" - }, - "gpt-5.4": { - "provider": "openai", - "model_type": "llm", - "description": "Affordable frontier model for coding and professional work" - }, - "gpt-5.4-pro": { - "provider": "openai", - "model_type": "llm", - "description": "Higher-accuracy GPT-5.4 variant" - }, - "gpt-5.4-mini": { - "provider": "openai", - "model_type": "llm", - "description": "Strong mini model for coding, computer use, and subagents" - }, - "gpt-5.4-nano": { - "provider": "openai", - "model_type": "llm", - "description": "Cheapest GPT-5.4-class model for simple high-volume tasks" - }, - "gpt-5.3-codex": { - "provider": "openai", - "model_type": "llm", - "description": "Most capable agentic coding model" - }, - "gpt-5.2": { - "provider": "openai", - "model_type": "llm", - "description": "Previous frontier model for professional work with configurable reasoning" - }, - "gpt-5.2-pro": { - "provider": "openai", - "model_type": "llm", - "description": "Previous pro model for professional work" - }, - "gpt-5.2-codex": { - "provider": "openai", - "model_type": "llm", - "description": "Intelligent coding model optimized for long-horizon agentic tasks (deprecated)" - }, - "gpt-5.1": { - "provider": "openai", - "model_type": "llm", - "description": "Best model for coding and agentic tasks with configurable reasoning effort" - }, - "gpt-5.1-codex": { - "provider": "openai", - "model_type": "llm", - "description": "GPT-5.1 optimized for agentic coding in Codex (deprecated)" - }, - "gpt-5.1-codex-max": { - "provider": "openai", - "model_type": "llm", - "description": "GPT-5.1-codex optimized for long-running tasks (deprecated)" - }, - "gpt-5.1-codex-mini": { - "provider": "openai", - "model_type": "llm", - "description": "Smaller, cost-effective GPT-5.1-Codex variant (deprecated)" - }, - "gpt-5": { - "provider": "openai", - "model_type": "llm", - "description": "Intelligent reasoning model for coding and agentic tasks" - }, - "gpt-5-pro": { - "provider": "openai", - "model_type": "llm", - "description": "Higher-accuracy GPT-5 variant" - }, - "gpt-5-mini": { - "provider": "openai", - "model_type": "llm", - "description": "Near-frontier intelligence for cost-sensitive, low-latency workloads" - }, - "gpt-5-nano": { - "provider": "openai", - "model_type": "llm", - "description": "Fastest, most cost-efficient GPT-5 variant" - }, - "gpt-5-codex": { - "provider": "openai", - "model_type": "llm", - "description": "GPT-5 optimized for agentic coding in Codex (deprecated)" - }, - "codex-mini-latest": { - "provider": "openai", - "model_type": "llm", - "description": "Fast reasoning model optimized for the Codex CLI (deprecated)" - }, - "o3-pro": { - "provider": "openai", - "model_type": "llm", - "description": "o3 with more compute for better responses" - }, - "o3": { - "provider": "openai", - "model_type": "llm", - "description": "Reasoning model for complex tasks (succeeded by GPT-5)" - }, - "o3-mini": { - "provider": "openai", - "model_type": "llm", - "description": "Small o3 alternative (deprecated)" - }, - "o3-deep-research": { - "provider": "openai", - "model_type": "llm", - "description": "OpenAI flagship \u2014 complex reasoning, agentic coding, 1M context" - }, - "o4-mini": { - "provider": "openai", - "model_type": "llm", - "description": "Fast, cost-efficient reasoning model (succeeded by GPT-5 mini)" - }, - "o4-mini-deep-research": { - "provider": "openai", - "model_type": "llm", - "description": "Faster, more affordable deep research model (deprecated)" - }, - "o1-pro": { - "provider": "openai", - "model_type": "llm", - "description": "o1 with more compute for better responses (deprecated)" - }, - "o1": { - "provider": "openai", - "model_type": "llm", - "description": "Previous full o-series reasoning model (deprecated)" - }, - "o1-mini": { - "provider": "openai", - "model_type": "llm", - "description": "Small o1 alternative (deprecated)" - }, - "o1-preview": { - "provider": "openai", - "model_type": "llm", - "description": "Preview of the first o-series reasoning model (deprecated)" - }, - "gpt-4.1": { - "provider": "openai", - "model_type": "llm", - "description": "Smartest non-reasoning GPT-4.x model" - }, - "gpt-4.1-mini": { - "provider": "openai", - "model_type": "llm", - "description": "Smaller, faster GPT-4.1 variant" - }, - "gpt-4.1-nano": { - "provider": "openai", - "model_type": "llm", - "description": "Fastest, most cost-efficient GPT-4.1 variant (deprecated)" - }, - "gpt-4o": { - "provider": "openai", - "model_type": "llm", - "description": "Fast, intelligent, flexible GPT model (deprecated)" - }, - "gpt-4o-mini": { - "provider": "openai", - "model_type": "llm", - "description": "Fast, affordable small model for focused tasks (deprecated)" - }, - "gpt-4o-search-preview": { - "provider": "openai", - "model_type": "llm", - "description": "GPT-4o with web search in Chat Completions (deprecated)" - }, - "gpt-4o-mini-search-preview": { - "provider": "openai", - "model_type": "llm", - "description": "GPT-4o mini with web search in Chat Completions (deprecated)" - }, - "gpt-4-turbo": { - "provider": "openai", - "model_type": "llm", - "description": "Older high-intelligence GPT model (deprecated)" - }, - "gpt-4-turbo-preview": { - "provider": "openai", - "model_type": "llm", - "description": "Older fast GPT model preview (deprecated)" - }, - "gpt-4": { - "provider": "openai", - "model_type": "llm", - "description": "Older high-intelligence GPT model (deprecated)" - }, - "gpt-4.5-preview": { - "provider": "openai", - "model_type": "llm", - "description": "Deprecated large GPT preview model" - }, - "gpt-3.5-turbo": { - "provider": "openai", - "model_type": "llm", - "description": "Legacy GPT model for cheaper chat tasks (deprecated)" - }, - "computer-use-preview": { - "provider": "openai", - "model_type": "llm", - "description": "Specialized model for computer use tool (deprecated)" - }, - "gpt-4o-audio-preview": { - "provider": "openai", - "model_type": "llm", - "description": "GPT-4o with audio input/output on Chat Completions (deprecated)" - }, - "gpt-4o-mini-audio-preview": { - "provider": "openai", - "model_type": "llm", - "description": "GPT-4o mini with audio input/output on Chat Completions (deprecated)" - }, - "gpt-audio-1.5": { - "provider": "openai", - "model_type": "llm", - "description": "Best voice model for audio in/out with Chat Completions" - }, - "gpt-audio": { - "provider": "openai", - "model_type": "llm", - "description": "Audio inputs and outputs with Chat Completions API" - }, - "gpt-audio-mini": { - "provider": "openai", - "model_type": "llm", - "description": "Cost-efficient audio in/out with Chat Completions (deprecated)" - }, - "chat-latest": { - "provider": "openai", - "model_type": "llm", - "description": "Latest Instant model used in ChatGPT; not recommended for most API use" - }, - "gpt-5.3-chat": { - "provider": "openai", - "model_type": "llm", - "description": "GPT-5.3 Instant model used in ChatGPT (deprecated)" - }, - "gpt-5.2-chat": { - "provider": "openai", - "model_type": "llm", - "description": "GPT-5.2 model used in ChatGPT (deprecated)" - }, - "gpt-5.1-chat": { - "provider": "openai", - "model_type": "llm", - "description": "GPT-5.1 model used in ChatGPT (deprecated)" - }, - "gpt-5-chat": { - "provider": "openai", - "model_type": "llm", - "description": "GPT-5 model used in ChatGPT (deprecated)" - }, - "chatgpt-4o-latest": { - "provider": "openai", - "model_type": "llm", - "description": "GPT-4o model used in ChatGPT (deprecated)" - }, - "gpt-oss-120b": { - "provider": "openai", - "model_type": "llm", - "description": "Most powerful OpenAI open-weight model (Apache 2.0)" - }, - "gpt-oss-20b": { - "provider": "openai", - "model_type": "llm", - "description": "Medium-sized OpenAI open-weight model for low latency (Apache 2.0)" - }, - "gpt-5.6-sol": { - "provider": "openai", - "model_type": "llm", - "description": "GPT-5.6 flagship — frontier reasoning, agentic coding, and computer use (1M context)", - "featured": true, - "featured_rank": 1, - "highlights": ["Flagship", "1M context", "Agentic coding"] - }, - "gpt-5.6-terra": { - "provider": "openai", - "model_type": "llm", - "description": "GPT-5.6 balanced tier — strong everyday performance at lower cost than Sol" - }, - "gpt-5.6-luna": { - "provider": "openai", - "model_type": "llm", - "description": "GPT-5.6 fast tier — lowest-cost model for high-volume classification and extraction" - }, - "gpt-5.6": { - "provider": "openai", - "model_type": "llm", - "description": "OpenAI API alias for gpt-5.6-sol" - }, - "gpt-4o-mini-tts": { - "provider": "openai", - "model_type": "tts", - "description": "Text-to-speech powered by GPT-4o mini (deprecated)", - "featured": true, - "featured_rank": 3, - "highlights": [ - "Natural prosody", - "Multiple voices", - "Studio quality" - ] - }, - "claude-sonnet-4.5": { - "provider": "anthropic", - "model_type": "llm" - }, - "claude-opus-4.5": { - "provider": "anthropic", - "model_type": "llm" - }, - "claude-haiku-4.5": { - "provider": "anthropic", - "model_type": "llm" - }, - "claude-sonnet-4.6": { - "provider": "anthropic", - "model_type": "llm" - }, - "claude-opus-4.6": { - "provider": "anthropic", - "model_type": "llm" - }, - "claude-opus-4-7": { - "provider": "anthropic", - "model_type": "llm", - "description": "Claude Opus 4.7 \u2014 stronger coding, vision, and complex multi-step tasks" - }, - "claude-opus-4-8": { - "provider": "anthropic", - "model_type": "llm", - "description": "Claude Opus 4.8 \u2014 most capable Opus-tier model; 1M context, adaptive thinking" - }, - "claude-fable-5": { - "provider": "anthropic", - "model_type": "llm", - "description": "Claude Fable 5 \u2014 Anthropic's most capable widely released model for demanding agentic work" - }, - "claude-mythos-5": { - "provider": "anthropic", - "model_type": "llm", - "description": "Claude Mythos 5 \u2014 Fable 5 capabilities; limited availability via Project Glasswing" - }, - "claude-sonnet-5": { - "provider": "anthropic", - "model_type": "llm", - "description": "Claude Sonnet 5 \u2014 latest Sonnet tier; near-Opus quality with lower latency and cost", - "featured": true, - "featured_rank": 4, - "highlights": ["Latest Sonnet", "Agentic workflows", "Cost-efficient"] - }, - "grok-4.3": { - "provider": "xai", - "model_type": "llm", - "description": "xAI flagship \u2014 low hallucination, agentic tool calling, 1M context" - }, - "grok-build-0.1": { - "provider": "xai", - "model_type": "llm", - "description": "xAI fast coding model trained for agentic coding (early access)" - }, - "grok-4.20-0309-reasoning": { - "provider": "xai", - "model_type": "llm", - "description": "Grok 4.20 reasoning snapshot (Mar 2026)" - }, - "grok-4.20-0309-non-reasoning": { - "provider": "xai", - "model_type": "llm", - "description": "Grok 4.20 non-reasoning snapshot (Mar 2026)" - }, - "grok-4.20-multi-agent-0309": { - "provider": "xai", - "model_type": "llm", - "description": "Grok 4.20 multi-agent snapshot (Mar 2026)" - }, - "deepseek-v4-pro": { - "provider": "fireworks", - "model_type": "llm", - "description": "DeepSeek V4 Pro \u2014 frontier MoE reasoning and coding (1M context)" - }, - "deepseek-v4-flash": { - "provider": "fireworks", - "model_type": "llm", - "description": "DeepSeek V4 Flash \u2014 fast extraction, classification, and search" - }, - "kimi-k2p6": { - "provider": "fireworks", - "model_type": "llm", - "description": "Kimi K2.6 \u2014 native multimodal agentic model for long-horizon coding" - }, - "kimi-k2p5": { - "provider": "fireworks", - "model_type": "llm", - "description": "Kimi K2.5 \u2014 unified vision/text agentic model with controllable reasoning" - }, - "glm-5p1": { - "provider": "fireworks", - "model_type": "llm", - "description": "GLM 5.1 \u2014 frontier open model for reasoning and agentic workflows" - }, - "minimax-m2p7": { - "provider": "fireworks", - "model_type": "llm", - "description": "MiniMax M2.7 \u2014 MoE model for complex agent harnesses and productivity tasks" - }, - "minimax-m2p5": { - "provider": "fireworks", - "model_type": "llm", - "description": "MiniMax M2.5 \u2014 fast coding and agentic tool use at low cost" - }, - "qwen3p6-plus": { - "provider": "fireworks", - "model_type": "llm", - "description": "Qwen 3.6 Plus \u2014 flagship multimodal model (Fireworks exclusive outside Alibaba)" - }, - "gpt-oss-120b": { - "provider": "fireworks", - "model_type": "llm", - "description": "OpenAI gpt-oss-120b \u2014 high-quality open-weight model for general reasoning" - }, - "gpt-oss-20b": { - "provider": "fireworks", - "model_type": "llm", - "description": "OpenAI gpt-oss-20b \u2014 fast open-weight model for chat and classification" - }, - "firefunction-v2": { - "provider": "fireworks", - "model_type": "llm", - "description": "FireFunction V2 \u2014 Fireworks function-calling optimized model" - }, - "google-speech-v2": { - "provider": "google", - "model_type": "stt" - }, - "gemini-2.5-pro-stt": { - "provider": "google", - "model_type": "stt", - "description": "Gemini 2.5 Pro used for audio transcription via LiteLLM proxy (audio input -> text)" - }, - "gemini-2.5-flash-stt": { - "provider": "google", - "model_type": "stt", - "description": "Gemini 2.5 Flash used for audio transcription via LiteLLM proxy (fast, cost-efficient)", - "featured": true, - "featured_rank": 6, - "highlights": [ - "Fast transcription", - "Cost-efficient", - "Gemini 2.5" - ] - }, - "gemini-2.5-flash-lite-stt": { - "provider": "google", - "model_type": "stt", - "description": "Gemini 2.5 Flash Lite used for audio transcription via LiteLLM proxy (lowest cost)" - }, - "gemini-2.5-pro": { - "provider": "google", - "model_type": "llm" - }, - "gemini-2.5-flash": { - "provider": "google", - "model_type": "llm" - }, - "gemini-2.5-flash-lite": { - "provider": "google", - "model_type": "llm", - "description": "Gemini 2.5 Flash Lite \u2014 lowest-latency 2.5 model; ideal for diarisation, classification, and other low-reasoning multimodal tasks" - }, - "gemini-3-pro-preview": { - "provider": "google", - "model_type": "llm", - "description": "Gemini 3 Pro \u2014 advanced reasoning and multimodal understanding" - }, - "gemini-3-flash-preview": { - "provider": "google", - "model_type": "llm", - "description": "Gemini 3 Flash \u2014 frontier-class performance at lower cost" - }, - "gemini-3.1-pro-preview": { - "provider": "google", - "model_type": "llm", - "description": "Gemini 3.1 Pro \u2014 latest reasoning model with 1M context" - }, - "gemini-3.5-flash": { - "provider": "google", - "model_type": "llm", - "description": "Gemini 3.5 Flash \u2014 GA flagship Flash model for agentic coding and long-horizon tasks", - "featured": true, - "featured_rank": 7, - "highlights": ["GA stable", "1M context", "Agentic coding"] - }, - "gemini-3.5-flash-lite": { - "provider": "google", - "model_type": "llm", - "description": "Gemini 3.5 Flash Lite \u2014 fastest, lowest-cost 3.5 model for high-throughput extraction, classification, and subagent workflows" - }, - "gemini-3.1-flash-lite": { - "provider": "google", - "model_type": "llm", - "description": "Gemini 3.1 Flash Lite \u2014 cost-efficient stable model for high-volume lightweight tasks" - }, - "azure-speech-v1": { - "provider": "azure", - "model_type": "stt", - "description": "Azure Speech-to-text (batch and realtime)" - }, - "azure-openai-gpt4": { - "provider": "azure", - "model_type": "llm", - "description": "Legacy alias — maps to gpt-4 deployment; prefer azure-gpt-5-mini if that is your deployment name" - }, - "azure-gpt-4o": { - "provider": "azure", - "model_type": "llm", - "description": "Azure OpenAI GPT-4o — catalog key maps to deployment name gpt-4o" - }, - "azure-gpt-4o-mini": { - "provider": "azure", - "model_type": "llm", - "description": "Azure OpenAI GPT-4o mini" - }, - "azure-gpt-4.1": { - "provider": "azure", - "model_type": "llm", - "description": "Azure OpenAI GPT-4.1" - }, - "azure-gpt-4.1-mini": { - "provider": "azure", - "model_type": "llm", - "description": "Azure OpenAI GPT-4.1 mini" - }, - "azure-gpt-4.1-nano": { - "provider": "azure", - "model_type": "llm", - "description": "Azure OpenAI GPT-4.1 nano" - }, - "azure-gpt-5": { - "provider": "azure", - "model_type": "llm", - "description": "Azure OpenAI GPT-5" - }, - "azure-gpt-5-mini": { - "provider": "azure", - "model_type": "llm", - "description": "Azure OpenAI GPT-5 mini — pick this if your deployment is named gpt-5-mini" - }, - "azure-gpt-5-nano": { - "provider": "azure", - "model_type": "llm", - "description": "Azure OpenAI GPT-5 nano" - }, - "azure-gpt-5.1": { - "provider": "azure", - "model_type": "llm", - "description": "Azure OpenAI GPT-5.1" - }, - "azure-gpt-5.2": { - "provider": "azure", - "model_type": "llm", - "description": "Azure OpenAI GPT-5.2" - }, - "azure-o3": { - "provider": "azure", - "model_type": "llm", - "description": "Azure OpenAI o3 reasoning model" - }, - "azure-o3-mini": { - "provider": "azure", - "model_type": "llm", - "description": "Azure OpenAI o3 mini" - }, - "azure-o4-mini": { - "provider": "azure", - "model_type": "llm", - "description": "Azure OpenAI o4 mini" - }, - "azure-tts-v1": { - "provider": "azure", - "model_type": "tts", - "description": "Azure neural text-to-speech" - }, - "aws-transcribe": { - "provider": "aws", - "model_type": "stt" - }, - "aws-bedrock-claude": { - "provider": "aws", - "model_type": "llm" - }, - "aws-polly": { - "provider": "aws", - "model_type": "tts" - }, - "deepgram-flux": { - "provider": "deepgram", - "model_type": "stt", - "description": "Conversational STT with integrated turn detection for voice agents", - "featured": true, - "featured_rank": 4, - "highlights": [ - "Turn detection", - "Voice agents", - "Low latency" - ] - }, - "deepgram-nova-3": { - "provider": "deepgram", - "model_type": "stt", - "description": "Highest-performing general-purpose ASR model" - }, - "deepgram-nova-3-general": { - "provider": "deepgram", - "model_type": "stt" - }, - "deepgram-nova-3-general-preview-12-2025": { - "provider": "deepgram", - "model_type": "stt" - }, - "deepgram-nova-2": { - "provider": "deepgram", - "model_type": "stt", - "description": "Available for languages not yet supported by Nova-3" - }, - "pulse-v4": { - "provider": "smallest", - "model_type": "stt", - "description": "Smallest Pulse v4 speech-to-text (batch and realtime)" - }, - "cartesia-sonic-3": { - "provider": "cartesia", - "model_type": "tts", - "description": "Flagship streaming TTS with laughter, volume/speed/emotion controls, 42 languages", - "featured": true, - "featured_rank": 1, - "highlights": [ - "Ultra-low latency", - "42 languages", - "Emotion controls" - ] - }, - "cartesia-sonic-3-mini": { - "provider": "cartesia", - "model_type": "tts" - }, - "cartesia-sonic-3-nano": { - "provider": "cartesia", - "model_type": "tts" - }, - "eleven_v3": { - "provider": "elevenlabs", - "model_type": "tts", - "description": "Human-like and expressive speech generation", - "languages": "70+ languages" - }, - "eleven_ttv_v3": { - "provider": "elevenlabs", - "model_type": "tts", - "description": "Human-like and expressive voice design model (Text to Voice)", - "languages": "70+ languages" - }, - "eleven_multilingual_v2": { - "provider": "elevenlabs", - "model_type": "tts", - "description": "Most lifelike model with rich emotional expression", - "languages": "en, ja, zh, de, hi, fr, ko, pt, it, es, id, nl, tr, fil, pl, sv, bg, ro, ar, cs, el, fi, hr, ms, sk, da, ta, uk, ru" - }, - "eleven_flash_v2_5": { - "provider": "elevenlabs", - "model_type": "tts", - "description": "Ultra-fast model optimized for real-time use (~75ms)", - "languages": "en, ja, zh, de, hi, fr, ko, pt, it, es, id, nl, tr, fil, pl, sv, bg, ro, ar, cs, el, fi, hr, ms, sk, da, ta, uk, ru, hu, no, vi", - "featured": true, - "featured_rank": 2, - "highlights": [ - "~75ms latency", - "Real-time", - "Multilingual" - ] - }, - "eleven_turbo_v2_5": { - "provider": "elevenlabs", - "model_type": "tts", - "description": "High quality, low-latency model (~250ms-300ms)", - "languages": "en, ja, zh, de, hi, fr, ko, pt, it, es, id, nl, tr, fil, pl, sv, bg, ro, ar, cs, el, fi, hr, ms, sk, da, ta, uk, ru, hu, no, vi" - }, - "eleven_multilingual_sts_v2": { - "provider": "elevenlabs", - "model_type": "sts", - "description": "State-of-the-art multilingual voice changer model (Speech to Speech)", - "languages": "en, ja, zh, de, hi, fr, ko, pt, it, es, id, nl, tr, fil, pl, sv, bg, ro, ar, cs, el, fi, hr, ms, sk, da, ta, uk, ru" - }, - "eleven_multilingual_ttv_v2": { - "provider": "elevenlabs", - "model_type": "tts", - "description": "State-of-the-art multilingual voice designer model (Text to Voice)", - "languages": "en, ja, zh, de, hi, fr, ko, pt, it, es, id, nl, tr, fil, pl, sv, bg, ro, ar, cs, el, fi, hr, ms, sk, da, ta, uk, ru" - }, - "eleven_english_sts_v2": { - "provider": "elevenlabs", - "model_type": "sts", - "description": "English-only voice changer model (Speech to Speech)", - "languages": "en" - }, - "scribe_v2_realtime": { - "provider": "elevenlabs", - "model_type": "stt", - "description": "Real-time speech recognition model", - "languages": "90+ languages", - "featured": true, - "featured_rank": 5, - "highlights": [ - "Real-time STT", - "90+ languages", - "Streaming" - ] - }, - "scribe_v2": { - "provider": "elevenlabs", - "model_type": "stt", - "description": "Most accurate transcription model with keyterm prompting and entity detection", - "languages": "90+ languages" - }, - "eleven_text_to_sound_v2": { - "provider": "elevenlabs", - "model_type": "sound_effects", - "description": "Sound effects generation from text prompts" - }, - "music_v1": { - "provider": "elevenlabs", - "model_type": "music", - "description": "Studio-grade music generation from text prompts", - "languages": "en, es, de, ja, and more" - }, - "murf-falcon": { - "provider": "murf", - "model_type": "tts", - "description": "Murf FALCON \u2013 ultra-low latency streaming TTS", - "languages": "en, es, de, fr, hi, ja, ko, pt, it, zh, and 10+ more", - "voices_source_file": "murf_falcon_voices.json" - }, - "murf-gen2": { - "provider": "murf", - "model_type": "tts", - "description": "Murf GEN2 \u2013 high-quality, natural-sounding TTS", - "languages": "en, es, de, fr, hi, ja, ko, pt, it, zh, and 10+ more", - "voices_source_file": "murf_gen2_voices.json" - }, - "voicemaker-ai3": { - "provider": "voicemaker", - "model_type": "tts", - "description": "VoiceMaker AI3 neural TTS voices", - "languages": "en-US, en-GB, en-AU, es-ES, fr-FR, de-DE, ja-JP, zh-CN", - "voices_source_file": "voicemaker_voices.json" - }, - "voicemaker-proplus": { - "provider": "voicemaker", - "model_type": "tts", - "description": "VoiceMaker ProPlus expressive multilingual voices", - "languages": "multi-lang + regional accents", - "voices_source_file": "voicemaker_voices.json" - }, - "lightning-v3.1": { - "provider": "smallest", - "model_type": "tts", - "description": "Smallest Lightning v3.1 low-latency neural TTS" - }, - "sarvam-30b": { - "provider": "sarvam", - "model_type": "llm", - "description": "Sarvam 30B \u2014 balanced Indian-language + English chat model (64K context)" - }, - "sarvam-105b": { - "provider": "sarvam", - "model_type": "llm", - "description": "Sarvam 105B \u2014 flagship MoE for complex reasoning, coding, and agentic workflows (128K context)" - }, - "saaras:v3": { - "provider": "sarvam", - "model_type": "stt", - "description": "State-of-the-art Sarvam STT model with transcribe/translate modes" - }, - "bulbul:v3": { - "provider": "sarvam", - "model_type": "tts", - "description": "High quality Indian language TTS", - "voices": [ - { - "id": "aditya", - "name": "Aditya", - "gender": "Male" - }, - { - "id": "ritu", - "name": "Ritu", - "gender": "Female" - }, - { - "id": "ashutosh", - "name": "Ashutosh", - "gender": "Male" - }, - { - "id": "priya", - "name": "Priya", - "gender": "Female" - }, - { - "id": "neha", - "name": "Neha", - "gender": "Female" - }, - { - "id": "rahul", - "name": "Rahul", - "gender": "Male" - }, - { - "id": "pooja", - "name": "Pooja", - "gender": "Female" - }, - { - "id": "rohan", - "name": "Rohan", - "gender": "Male" - }, - { - "id": "simran", - "name": "Simran", - "gender": "Female" - }, - { - "id": "kavya", - "name": "Kavya", - "gender": "Female" - }, - { - "id": "amit", - "name": "Amit", - "gender": "Male" - }, - { - "id": "dev", - "name": "Dev", - "gender": "Male" - }, - { - "id": "ishita", - "name": "Ishita", - "gender": "Female" - }, - { - "id": "shreya", - "name": "Shreya", - "gender": "Female" - }, - { - "id": "ratan", - "name": "Ratan", - "gender": "Male" - }, - { - "id": "varun", - "name": "Varun", - "gender": "Male" - }, - { - "id": "manan", - "name": "Manan", - "gender": "Male" - }, - { - "id": "sumit", - "name": "Sumit", - "gender": "Male" - }, - { - "id": "roopa", - "name": "Roopa", - "gender": "Female" - }, - { - "id": "kabir", - "name": "Kabir", - "gender": "Male" - }, - { - "id": "aayan", - "name": "Aayan", - "gender": "Male" - }, - { - "id": "shubh", - "name": "Shubh", - "gender": "Male" - }, - { - "id": "advait", - "name": "Advait", - "gender": "Male" - }, - { - "id": "amelia", - "name": "Amelia", - "gender": "Female" - }, - { - "id": "sophia", - "name": "Sophia", - "gender": "Female" - }, - { - "id": "anand", - "name": "Anand", - "gender": "Male" - }, - { - "id": "tanya", - "name": "Tanya", - "gender": "Female" - }, - { - "id": "tarun", - "name": "Tarun", - "gender": "Male" - }, - { - "id": "sunny", - "name": "Sunny", - "gender": "Male" - }, - { - "id": "mani", - "name": "Mani", - "gender": "Male" - }, - { - "id": "gokul", - "name": "Gokul", - "gender": "Male" - }, - { - "id": "vijay", - "name": "Vijay", - "gender": "Male" - }, - { - "id": "shruti", - "name": "Shruti", - "gender": "Female" - }, - { - "id": "suhani", - "name": "Suhani", - "gender": "Female" - }, - { - "id": "mohit", - "name": "Mohit", - "gender": "Male" - }, - { - "id": "kavitha", - "name": "Kavitha", - "gender": "Female" - }, - { - "id": "rehan", - "name": "Rehan", - "gender": "Male" - }, - { - "id": "soham", - "name": "Soham", - "gender": "Male" - }, - { - "id": "rupali", - "name": "Rupali", - "gender": "Female" - } - ] - } -} +{ + "whisper-1": { + "provider": "openai", + "model_type": "stt", + "description": "General-purpose speech recognition (verbose_json + word/segment timestamps)", + "pricing": { + "source": "litellm_import", + "usage_kind": "stt", + "audio_per_minute": 0.006 + } + }, + "gpt-4o-transcribe": { + "provider": "openai", + "model_type": "stt", + "description": "GPT-4o speech-to-text; higher accuracy, no granular word timestamps", + "pricing": { + "source": "litellm_import", + "usage_kind": "stt", + "audio_per_minute": 0.00372 + } + }, + "gpt-4o-mini-transcribe": { + "provider": "openai", + "model_type": "stt", + "description": "Cheaper / faster GPT-4o transcription; no granular word timestamps", + "pricing": { + "source": "litellm_import", + "usage_kind": "stt", + "audio_per_minute": 0.00186 + } + }, + "gpt-4o-transcribe-diarize": { + "provider": "openai", + "model_type": "stt", + "description": "Transcription model that identifies who's speaking when", + "pricing": { + "source": "litellm_import", + "usage_kind": "stt", + "audio_per_minute": 0.00372 + } + }, + "gpt-realtime-whisper": { + "provider": "openai", + "model_type": "stt", + "description": "Streaming speech-to-text for realtime transcription", + "pricing": { + "source": "litellm_import", + "usage_kind": "stt", + "audio_per_minute": 0.01698 + } + }, + "gpt-5.6": { + "provider": "openai", + "model_type": "llm", + "description": "OpenAI API alias for gpt-5.6-sol", + "pricing": { + "source": "litellm_import", + "usage_kind": "llm", + "input_per_1m": 5.0, + "output_per_1m": 30.0, + "cache_read_per_1m": 0.5, + "cache_write_per_1m": 6.25 + } + }, + "gpt-5.5": { + "provider": "openai", + "model_type": "llm", + "description": "OpenAI flagship — complex reasoning, agentic coding, 1M context", + "pricing": { + "source": "litellm_import", + "usage_kind": "llm", + "input_per_1m": 5.0, + "output_per_1m": 30.0, + "cache_read_per_1m": 0.5 + } + }, + "gpt-5.5-pro": { + "provider": "openai", + "model_type": "llm", + "description": "Higher-accuracy GPT-5.5 variant with parallel test-time compute", + "pricing": { + "source": "litellm_import", + "usage_kind": "llm", + "input_per_1m": 30.0, + "output_per_1m": 180.0, + "cache_read_per_1m": 3.0 + } + }, + "gpt-5.4": { + "provider": "openai", + "model_type": "llm", + "description": "Affordable frontier model for coding and professional work", + "pricing": { + "source": "litellm_import", + "usage_kind": "llm", + "input_per_1m": 2.5, + "output_per_1m": 15.0, + "cache_read_per_1m": 0.25 + } + }, + "gpt-5.4-pro": { + "provider": "openai", + "model_type": "llm", + "description": "Higher-accuracy GPT-5.4 variant", + "pricing": { + "source": "litellm_import", + "usage_kind": "llm", + "input_per_1m": 30.0, + "output_per_1m": 180.0, + "cache_read_per_1m": 3.0 + } + }, + "gpt-5.4-mini": { + "provider": "openai", + "model_type": "llm", + "description": "Strong mini model for coding, computer use, and subagents", + "pricing": { + "source": "litellm_import", + "usage_kind": "llm", + "input_per_1m": 0.75, + "output_per_1m": 4.5, + "cache_read_per_1m": 0.075 + } + }, + "gpt-5.4-nano": { + "provider": "openai", + "model_type": "llm", + "description": "Cheapest GPT-5.4-class model for simple high-volume tasks", + "pricing": { + "source": "litellm_import", + "usage_kind": "llm", + "input_per_1m": 0.2, + "output_per_1m": 1.25, + "cache_read_per_1m": 0.02 + } + }, + "gpt-5.3-codex": { + "provider": "openai", + "model_type": "llm", + "description": "Most capable agentic coding model", + "pricing": { + "source": "litellm_import", + "usage_kind": "llm", + "input_per_1m": 1.75, + "output_per_1m": 14.0, + "cache_read_per_1m": 0.175 + } + }, + "gpt-5.2": { + "provider": "openai", + "model_type": "llm", + "description": "Previous frontier model for professional work with configurable reasoning", + "pricing": { + "source": "litellm_import", + "usage_kind": "llm", + "input_per_1m": 1.75, + "output_per_1m": 14.0, + "cache_read_per_1m": 0.175 + } + }, + "gpt-5.2-pro": { + "provider": "openai", + "model_type": "llm", + "description": "Previous pro model for professional work", + "pricing": { + "source": "litellm_import", + "usage_kind": "llm", + "input_per_1m": 21.0, + "output_per_1m": 168.0 + } + }, + "gpt-5.2-codex": { + "provider": "openai", + "model_type": "llm", + "description": "Intelligent coding model optimized for long-horizon agentic tasks (deprecated)", + "pricing": { + "source": "litellm_import", + "usage_kind": "llm", + "input_per_1m": 1.75, + "output_per_1m": 14.0, + "cache_read_per_1m": 0.175 + } + }, + "gpt-5.1": { + "provider": "openai", + "model_type": "llm", + "description": "Best model for coding and agentic tasks with configurable reasoning effort", + "pricing": { + "source": "litellm_import", + "usage_kind": "llm", + "input_per_1m": 1.25, + "output_per_1m": 10.0, + "cache_read_per_1m": 0.125 + } + }, + "gpt-5.1-codex": { + "provider": "openai", + "model_type": "llm", + "description": "GPT-5.1 optimized for agentic coding in Codex (deprecated)", + "pricing": { + "source": "litellm_import", + "usage_kind": "llm", + "input_per_1m": 1.25, + "output_per_1m": 10.0, + "cache_read_per_1m": 0.125 + } + }, + "gpt-5.1-codex-max": { + "provider": "openai", + "model_type": "llm", + "description": "GPT-5.1-codex optimized for long-running tasks (deprecated)", + "pricing": { + "source": "litellm_import", + "usage_kind": "llm", + "input_per_1m": 1.25, + "output_per_1m": 10.0, + "cache_read_per_1m": 0.125 + } + }, + "gpt-5.1-codex-mini": { + "provider": "openai", + "model_type": "llm", + "description": "Smaller, cost-effective GPT-5.1-Codex variant (deprecated)", + "pricing": { + "source": "litellm_import", + "usage_kind": "llm", + "input_per_1m": 0.25, + "output_per_1m": 2.0, + "cache_read_per_1m": 0.025 + } + }, + "gpt-5": { + "provider": "openai", + "model_type": "llm", + "description": "Intelligent reasoning model for coding and agentic tasks", + "pricing": { + "source": "litellm_import", + "usage_kind": "llm", + "input_per_1m": 1.25, + "output_per_1m": 10.0, + "cache_read_per_1m": 0.125 + } + }, + "gpt-5-pro": { + "provider": "openai", + "model_type": "llm", + "description": "Higher-accuracy GPT-5 variant", + "pricing": { + "source": "litellm_import", + "usage_kind": "llm", + "input_per_1m": 15.0, + "output_per_1m": 120.0 + } + }, + "gpt-5-mini": { + "provider": "openai", + "model_type": "llm", + "description": "Near-frontier intelligence for cost-sensitive, low-latency workloads", + "pricing": { + "source": "litellm_import", + "usage_kind": "llm", + "input_per_1m": 0.25, + "output_per_1m": 2.0, + "cache_read_per_1m": 0.025 + } + }, + "gpt-5-nano": { + "provider": "openai", + "model_type": "llm", + "description": "Fastest, most cost-efficient GPT-5 variant", + "pricing": { + "source": "litellm_import", + "usage_kind": "llm", + "input_per_1m": 0.05, + "output_per_1m": 0.4, + "cache_read_per_1m": 0.005 + } + }, + "gpt-5-codex": { + "provider": "openai", + "model_type": "llm", + "description": "GPT-5 optimized for agentic coding in Codex (deprecated)", + "pricing": { + "source": "litellm_import", + "usage_kind": "llm", + "input_per_1m": 1.25, + "output_per_1m": 10.0, + "cache_read_per_1m": 0.125 + } + }, + "codex-mini-latest": { + "provider": "openai", + "model_type": "llm", + "description": "Fast reasoning model optimized for the Codex CLI (deprecated)", + "pricing": { + "source": "litellm_import", + "usage_kind": "llm", + "input_per_1m": 1.5, + "output_per_1m": 6.0, + "cache_read_per_1m": 0.375 + } + }, + "o3-pro": { + "provider": "openai", + "model_type": "llm", + "description": "o3 with more compute for better responses", + "pricing": { + "source": "litellm_import", + "usage_kind": "llm", + "input_per_1m": 20.0, + "output_per_1m": 80.0 + } + }, + "o3": { + "provider": "openai", + "model_type": "llm", + "description": "Reasoning model for complex tasks (succeeded by GPT-5)", + "pricing": { + "source": "litellm_import", + "usage_kind": "llm", + "input_per_1m": 2.0, + "output_per_1m": 8.0, + "cache_read_per_1m": 0.5 + } + }, + "o3-mini": { + "provider": "openai", + "model_type": "llm", + "description": "Small o3 alternative (deprecated)", + "pricing": { + "source": "litellm_import", + "usage_kind": "llm", + "input_per_1m": 1.1, + "output_per_1m": 4.4, + "cache_read_per_1m": 0.55 + } + }, + "o3-deep-research": { + "provider": "openai", + "model_type": "llm", + "description": "OpenAI flagship — complex reasoning, agentic coding, 1M context", + "pricing": { + "source": "litellm_import", + "usage_kind": "llm", + "input_per_1m": 10.0, + "output_per_1m": 40.0, + "cache_read_per_1m": 2.5 + } + }, + "o4-mini": { + "provider": "openai", + "model_type": "llm", + "description": "Fast, cost-efficient reasoning model (succeeded by GPT-5 mini)", + "pricing": { + "source": "litellm_import", + "usage_kind": "llm", + "input_per_1m": 1.1, + "output_per_1m": 4.4, + "cache_read_per_1m": 0.275 + } + }, + "o4-mini-deep-research": { + "provider": "openai", + "model_type": "llm", + "description": "Faster, more affordable deep research model (deprecated)", + "pricing": { + "source": "litellm_import", + "usage_kind": "llm", + "input_per_1m": 2.0, + "output_per_1m": 8.0, + "cache_read_per_1m": 0.5 + } + }, + "o1-pro": { + "provider": "openai", + "model_type": "llm", + "description": "o1 with more compute for better responses (deprecated)", + "pricing": { + "source": "litellm_import", + "usage_kind": "llm", + "input_per_1m": 150.0, + "output_per_1m": 600.0 + } + }, + "o1": { + "provider": "openai", + "model_type": "llm", + "description": "Previous full o-series reasoning model (deprecated)", + "pricing": { + "source": "litellm_import", + "usage_kind": "llm", + "input_per_1m": 15.0, + "output_per_1m": 60.0, + "cache_read_per_1m": 7.5 + } + }, + "o1-mini": { + "provider": "openai", + "model_type": "llm", + "description": "Small o1 alternative (deprecated)", + "pricing": { + "source": "litellm_import", + "usage_kind": "llm", + "input_per_1m": 1.21, + "output_per_1m": 4.84, + "cache_read_per_1m": 0.605 + } + }, + "o1-preview": { + "provider": "openai", + "model_type": "llm", + "description": "Preview of the first o-series reasoning model (deprecated)", + "pricing": { + "source": "litellm_import", + "usage_kind": "llm", + "input_per_1m": 15.0, + "output_per_1m": 60.0, + "cache_read_per_1m": 7.5 + } + }, + "gpt-4.1": { + "provider": "openai", + "model_type": "llm", + "description": "Smartest non-reasoning GPT-4.x model", + "pricing": { + "source": "litellm_import", + "usage_kind": "llm", + "input_per_1m": 2.0, + "output_per_1m": 8.0, + "cache_read_per_1m": 0.5 + } + }, + "gpt-4.1-mini": { + "provider": "openai", + "model_type": "llm", + "description": "Smaller, faster GPT-4.1 variant", + "pricing": { + "source": "litellm_import", + "usage_kind": "llm", + "input_per_1m": 0.4, + "output_per_1m": 1.6, + "cache_read_per_1m": 0.1 + } + }, + "gpt-4.1-nano": { + "provider": "openai", + "model_type": "llm", + "description": "Fastest, most cost-efficient GPT-4.1 variant (deprecated)", + "pricing": { + "source": "litellm_import", + "usage_kind": "llm", + "input_per_1m": 0.1, + "output_per_1m": 0.4, + "cache_read_per_1m": 0.025 + } + }, + "gpt-4o": { + "provider": "openai", + "model_type": "llm", + "description": "Fast, intelligent, flexible GPT model (deprecated)", + "pricing": { + "source": "litellm_import", + "usage_kind": "llm", + "input_per_1m": 2.5, + "output_per_1m": 10.0, + "cache_read_per_1m": 1.25 + } + }, + "gpt-4o-mini": { + "provider": "openai", + "model_type": "llm", + "description": "Fast, affordable small model for focused tasks (deprecated)", + "pricing": { + "source": "litellm_import", + "usage_kind": "llm", + "input_per_1m": 0.15, + "output_per_1m": 0.6, + "cache_read_per_1m": 0.075 + } + }, + "gpt-4o-search-preview": { + "provider": "openai", + "model_type": "llm", + "description": "GPT-4o with web search in Chat Completions (deprecated)", + "pricing": { + "source": "litellm_import", + "usage_kind": "llm", + "input_per_1m": 2.5, + "output_per_1m": 10.0, + "cache_read_per_1m": 1.25 + } + }, + "gpt-4o-mini-search-preview": { + "provider": "openai", + "model_type": "llm", + "description": "GPT-4o mini with web search in Chat Completions (deprecated)", + "pricing": { + "source": "litellm_import", + "usage_kind": "llm", + "input_per_1m": 0.15, + "output_per_1m": 0.6, + "cache_read_per_1m": 0.075 + } + }, + "gpt-4-turbo": { + "provider": "openai", + "model_type": "llm", + "description": "Older high-intelligence GPT model (deprecated)", + "pricing": { + "source": "litellm_import", + "usage_kind": "llm", + "input_per_1m": 10.0, + "output_per_1m": 30.0 + } + }, + "gpt-4-turbo-preview": { + "provider": "openai", + "model_type": "llm", + "description": "Older fast GPT model preview (deprecated)", + "pricing": { + "source": "litellm_import", + "usage_kind": "llm", + "input_per_1m": 10.0, + "output_per_1m": 30.0 + } + }, + "gpt-4": { + "provider": "openai", + "model_type": "llm", + "description": "Older high-intelligence GPT model (deprecated)", + "pricing": { + "source": "litellm_import", + "usage_kind": "llm", + "input_per_1m": 30.0, + "output_per_1m": 60.0 + } + }, + "gpt-4.5-preview": { + "provider": "openai", + "model_type": "llm", + "description": "Deprecated large GPT preview model", + "pricing": { + "source": "litellm_import", + "usage_kind": "llm", + "input_per_1m": 75.0, + "output_per_1m": 150.0, + "cache_read_per_1m": 37.5 + } + }, + "gpt-3.5-turbo": { + "provider": "openai", + "model_type": "llm", + "description": "Legacy GPT model for cheaper chat tasks (deprecated)", + "pricing": { + "source": "litellm_import", + "usage_kind": "llm", + "input_per_1m": 0.5, + "output_per_1m": 1.5 + } + }, + "computer-use-preview": { + "provider": "openai", + "model_type": "llm", + "description": "Specialized model for computer use tool (deprecated)", + "pricing": { + "source": "litellm_import", + "usage_kind": "llm", + "input_per_1m": 3.0, + "output_per_1m": 12.0 + } + }, + "gpt-4o-audio-preview": { + "provider": "openai", + "model_type": "llm", + "description": "GPT-4o with audio input/output on Chat Completions (deprecated)", + "pricing": { + "source": "litellm_import", + "usage_kind": "llm", + "input_per_1m": 2.5, + "output_per_1m": 10.0 + } + }, + "gpt-4o-mini-audio-preview": { + "provider": "openai", + "model_type": "llm", + "description": "GPT-4o mini with audio input/output on Chat Completions (deprecated)", + "pricing": { + "source": "litellm_import", + "usage_kind": "llm", + "input_per_1m": 0.15, + "output_per_1m": 0.6 + } + }, + "gpt-audio-1.5": { + "provider": "openai", + "model_type": "llm", + "description": "Best voice model for audio in/out with Chat Completions", + "pricing": { + "source": "litellm_import", + "usage_kind": "llm", + "input_per_1m": 2.5, + "output_per_1m": 10.0 + } + }, + "gpt-audio": { + "provider": "openai", + "model_type": "llm", + "description": "Audio inputs and outputs with Chat Completions API", + "pricing": { + "source": "litellm_import", + "usage_kind": "llm", + "input_per_1m": 2.5, + "output_per_1m": 10.0 + } + }, + "gpt-audio-mini": { + "provider": "openai", + "model_type": "llm", + "description": "Cost-efficient audio in/out with Chat Completions (deprecated)", + "pricing": { + "source": "litellm_import", + "usage_kind": "llm", + "input_per_1m": 0.6, + "output_per_1m": 2.4 + } + }, + "chat-latest": { + "provider": "openai", + "model_type": "llm", + "description": "Latest Instant model used in ChatGPT; not recommended for most API use", + "pricing": { + "source": "litellm_import", + "usage_kind": "llm", + "input_per_1m": 1.25, + "output_per_1m": 10.0, + "cache_read_per_1m": 0.125 + } + }, + "gpt-5.3-chat": { + "provider": "openai", + "model_type": "llm", + "description": "GPT-5.3 Instant model used in ChatGPT (deprecated)", + "pricing": { + "source": "litellm_import", + "usage_kind": "llm", + "input_per_1m": 1.75, + "output_per_1m": 14.0, + "cache_read_per_1m": 0.175 + } + }, + "gpt-5.2-chat": { + "provider": "openai", + "model_type": "llm", + "description": "GPT-5.2 model used in ChatGPT (deprecated)", + "pricing": { + "source": "litellm_import", + "usage_kind": "llm", + "input_per_1m": 1.75, + "output_per_1m": 14.0, + "cache_read_per_1m": 0.175 + } + }, + "gpt-5.1-chat": { + "provider": "openai", + "model_type": "llm", + "description": "GPT-5.1 model used in ChatGPT (deprecated)", + "pricing": { + "source": "litellm_import", + "usage_kind": "llm", + "input_per_1m": 1.38, + "output_per_1m": 11.0, + "cache_read_per_1m": 0.14 + } + }, + "gpt-5-chat": { + "provider": "openai", + "model_type": "llm", + "description": "GPT-5 model used in ChatGPT (deprecated)", + "pricing": { + "source": "litellm_import", + "usage_kind": "llm", + "input_per_1m": 1.25, + "output_per_1m": 10.0, + "cache_read_per_1m": 0.125 + } + }, + "chatgpt-4o-latest": { + "provider": "openai", + "model_type": "llm", + "description": "GPT-4o model used in ChatGPT (deprecated)", + "pricing": { + "source": "litellm_import", + "usage_kind": "llm", + "input_per_1m": 5.0, + "output_per_1m": 15.0 + } + }, + "gpt-oss-120b": { + "provider": "fireworks", + "model_type": "llm", + "description": "OpenAI gpt-oss-120b — high-quality open-weight model for general reasoning", + "pricing": { + "source": "litellm_import", + "usage_kind": "llm", + "input_per_1m": 0.15, + "output_per_1m": 0.6, + "cache_read_per_1m": 0.015 + } + }, + "gpt-oss-20b": { + "provider": "fireworks", + "model_type": "llm", + "description": "OpenAI gpt-oss-20b — fast open-weight model for chat and classification", + "pricing": { + "source": "litellm_import", + "usage_kind": "llm", + "input_per_1m": 0.07, + "output_per_1m": 0.3, + "cache_read_per_1m": 0.035 + } + }, + "gpt-5.6-sol": { + "provider": "openai", + "model_type": "llm", + "description": "GPT-5.6 flagship — frontier reasoning, agentic coding, and computer use (1M context)", + "featured": true, + "featured_rank": 1, + "highlights": [ + "Flagship", + "1M context", + "Agentic coding" + ], + "pricing": { + "source": "litellm_import", + "usage_kind": "llm", + "input_per_1m": 5.0, + "output_per_1m": 30.0, + "cache_read_per_1m": 0.5, + "cache_write_per_1m": 6.25 + } + }, + "gpt-5.6-terra": { + "provider": "openai", + "model_type": "llm", + "description": "GPT-5.6 balanced tier — strong everyday performance at lower cost than Sol", + "pricing": { + "source": "litellm_import", + "usage_kind": "llm", + "input_per_1m": 2.0, + "output_per_1m": 12.0, + "cache_read_per_1m": 0.2, + "cache_write_per_1m": 2.5 + } + }, + "gpt-5.6-luna": { + "provider": "openai", + "model_type": "llm", + "description": "GPT-5.6 fast tier — lowest-cost model for high-volume classification and extraction", + "pricing": { + "source": "litellm_import", + "usage_kind": "llm", + "input_per_1m": 0.2, + "output_per_1m": 1.2, + "cache_read_per_1m": 0.02, + "cache_write_per_1m": 0.25 + } + }, + "gpt-4o-mini-tts": { + "provider": "openai", + "model_type": "tts", + "description": "Text-to-speech powered by GPT-4o mini (deprecated)", + "featured": true, + "featured_rank": 3, + "highlights": [ + "Natural prosody", + "Multiple voices", + "Studio quality" + ], + "pricing": { + "source": "litellm_import", + "usage_kind": "tts", + "tts_per_1m_characters": 2.5 + } + }, + "claude-sonnet-4.5": { + "provider": "anthropic", + "model_type": "llm", + "pricing": { + "source": "litellm_import", + "usage_kind": "llm", + "input_per_1m": 3.0, + "output_per_1m": 15.0 + } + }, + "claude-opus-4.5": { + "provider": "anthropic", + "model_type": "llm", + "pricing": { + "source": "litellm_import", + "usage_kind": "llm", + "input_per_1m": 5.0, + "output_per_1m": 25.0 + } + }, + "claude-haiku-4.5": { + "provider": "anthropic", + "model_type": "llm", + "pricing": { + "source": "litellm_import", + "usage_kind": "llm", + "input_per_1m": 1.0, + "output_per_1m": 5.0, + "cache_read_per_1m": 0.1, + "cache_write_per_1m": 1.25 + } + }, + "claude-sonnet-4.6": { + "provider": "anthropic", + "model_type": "llm", + "pricing": { + "source": "litellm_import", + "usage_kind": "llm", + "input_per_1m": 3.0, + "output_per_1m": 15.0, + "cache_read_per_1m": 0.3, + "cache_write_per_1m": 3.75 + } + }, + "claude-opus-4.6": { + "provider": "anthropic", + "model_type": "llm", + "pricing": { + "source": "litellm_import", + "usage_kind": "llm", + "input_per_1m": 5.0, + "output_per_1m": 25.0, + "cache_read_per_1m": 0.5, + "cache_write_per_1m": 6.25 + } + }, + "claude-opus-4-7": { + "provider": "anthropic", + "model_type": "llm", + "description": "Claude Opus 4.7 — stronger coding, vision, and complex multi-step tasks", + "pricing": { + "source": "litellm_import", + "usage_kind": "llm", + "input_per_1m": 5.0, + "output_per_1m": 25.0, + "cache_read_per_1m": 0.5, + "cache_write_per_1m": 6.25 + } + }, + "claude-opus-4-8": { + "provider": "anthropic", + "model_type": "llm", + "description": "Claude Opus 4.8 — most capable Opus-tier model; 1M context, adaptive thinking", + "pricing": { + "source": "litellm_import", + "usage_kind": "llm", + "input_per_1m": 5.0, + "output_per_1m": 25.0, + "cache_read_per_1m": 0.5, + "cache_write_per_1m": 6.25 + } + }, + "claude-fable-5": { + "provider": "anthropic", + "model_type": "llm", + "description": "Claude Fable 5 — Anthropic's most capable widely released model for demanding agentic work", + "pricing": { + "source": "litellm_import", + "usage_kind": "llm", + "input_per_1m": 10.0, + "output_per_1m": 50.0, + "cache_read_per_1m": 1.0, + "cache_write_per_1m": 12.5 + } + }, + "claude-mythos-5": { + "provider": "anthropic", + "model_type": "llm", + "description": "Claude Mythos 5 — Fable 5 capabilities; limited availability via Project Glasswing", + "pricing": { + "source": "litellm_import", + "usage_kind": "llm", + "input_per_1m": 10.0, + "output_per_1m": 50.0, + "cache_read_per_1m": 1.0, + "cache_write_per_1m": 12.5 + } + }, + "claude-sonnet-5": { + "provider": "anthropic", + "model_type": "llm", + "description": "Claude Sonnet 5 — latest Sonnet tier; near-Opus quality with lower latency and cost", + "featured": true, + "featured_rank": 4, + "highlights": [ + "Latest Sonnet", + "Agentic workflows", + "Cost-efficient" + ], + "pricing": { + "source": "litellm_import", + "usage_kind": "llm", + "input_per_1m": 2.0, + "output_per_1m": 10.0, + "cache_read_per_1m": 0.2, + "cache_write_per_1m": 2.5 + } + }, + "grok-4.3": { + "provider": "xai", + "model_type": "llm", + "description": "xAI flagship — low hallucination, agentic tool calling, 1M context", + "pricing": { + "source": "litellm_import", + "usage_kind": "llm", + "input_per_1m": 1.25, + "output_per_1m": 2.5, + "cache_read_per_1m": 0.2 + } + }, + "grok-build-0.1": { + "provider": "xai", + "model_type": "llm", + "description": "xAI fast coding model trained for agentic coding (early access)", + "pricing": { + "source": "litellm_import", + "usage_kind": "llm", + "input_per_1m": 3.0, + "output_per_1m": 15.0, + "cache_read_per_1m": 0.75 + } + }, + "grok-4.20-0309-reasoning": { + "provider": "xai", + "model_type": "llm", + "description": "Grok 4.20 reasoning snapshot (Mar 2026)", + "pricing": { + "source": "litellm_import", + "usage_kind": "llm", + "input_per_1m": 1.25, + "output_per_1m": 2.5, + "cache_read_per_1m": 0.2 + } + }, + "grok-4.20-0309-non-reasoning": { + "provider": "xai", + "model_type": "llm", + "description": "Grok 4.20 non-reasoning snapshot (Mar 2026)", + "pricing": { + "source": "litellm_import", + "usage_kind": "llm", + "input_per_1m": 0.2, + "output_per_1m": 0.5, + "cache_read_per_1m": 0.05 + } + }, + "grok-4.20-multi-agent-0309": { + "provider": "xai", + "model_type": "llm", + "description": "Grok 4.20 multi-agent snapshot (Mar 2026)", + "pricing": { + "source": "litellm_import", + "usage_kind": "llm", + "input_per_1m": 3.0, + "output_per_1m": 15.0 + } + }, + "deepseek-v4-pro": { + "provider": "fireworks", + "model_type": "llm", + "description": "DeepSeek V4 Pro — frontier MoE reasoning and coding (1M context)", + "pricing": { + "source": "litellm_import", + "usage_kind": "llm", + "input_per_1m": 1.74, + "output_per_1m": 3.48, + "cache_read_per_1m": 0.145 + } + }, + "deepseek-v4-flash": { + "provider": "fireworks", + "model_type": "llm", + "description": "DeepSeek V4 Flash — fast extraction, classification, and search", + "pricing": { + "source": "litellm_import", + "usage_kind": "llm", + "input_per_1m": 0.14, + "output_per_1m": 0.28, + "cache_read_per_1m": 0.028 + } + }, + "kimi-k2p6": { + "provider": "fireworks", + "model_type": "llm", + "description": "Kimi K2.6 — native multimodal agentic model for long-horizon coding", + "pricing": { + "source": "litellm_import", + "usage_kind": "llm", + "input_per_1m": 0.95, + "output_per_1m": 4.0, + "cache_read_per_1m": 0.16 + } + }, + "kimi-k2p5": { + "provider": "fireworks", + "model_type": "llm", + "description": "Kimi K2.5 — unified vision/text agentic model with controllable reasoning", + "pricing": { + "source": "litellm_import", + "usage_kind": "llm", + "input_per_1m": 0.6, + "output_per_1m": 3.0, + "cache_read_per_1m": 0.1 + } + }, + "glm-5p1": { + "provider": "fireworks", + "model_type": "llm", + "description": "GLM 5.1 — frontier open model for reasoning and agentic workflows", + "pricing": { + "source": "litellm_import", + "usage_kind": "llm", + "input_per_1m": 1.4, + "output_per_1m": 4.4, + "cache_read_per_1m": 0.26 + } + }, + "minimax-m2p7": { + "provider": "fireworks", + "model_type": "llm", + "description": "MiniMax M2.7 — MoE model for complex agent harnesses and productivity tasks", + "pricing": { + "source": "litellm_import", + "usage_kind": "llm", + "input_per_1m": 0.3, + "output_per_1m": 1.2, + "cache_read_per_1m": 0.06 + } + }, + "minimax-m2p5": { + "provider": "fireworks", + "model_type": "llm", + "description": "MiniMax M2.5 — fast coding and agentic tool use at low cost", + "pricing": { + "source": "litellm_import", + "usage_kind": "llm", + "input_per_1m": 0.3, + "output_per_1m": 1.2, + "cache_read_per_1m": 0.06 + } + }, + "qwen3p6-plus": { + "provider": "fireworks", + "model_type": "llm", + "description": "Qwen 3.6 Plus — flagship multimodal model (Fireworks exclusive outside Alibaba)", + "pricing": { + "source": "litellm_import", + "usage_kind": "llm", + "input_per_1m": 0.325, + "output_per_1m": 1.95 + } + }, + "firefunction-v2": { + "provider": "fireworks", + "model_type": "llm", + "description": "FireFunction V2 — Fireworks function-calling optimized model", + "pricing": { + "source": "litellm_import", + "usage_kind": "llm", + "input_per_1m": 0.9, + "output_per_1m": 0.9 + } + }, + "google-speech-v2": { + "provider": "google", + "model_type": "stt", + "pricing": { + "source": "litellm_import", + "usage_kind": "stt", + "audio_per_minute": 0.0015 + } + }, + "gemini-2.5-pro-stt": { + "provider": "google", + "model_type": "stt", + "description": "Gemini 2.5 Pro used for audio transcription via LiteLLM proxy (audio input -> text)", + "pricing": { + "source": "litellm_import", + "usage_kind": "stt", + "audio_per_minute": 0.0015 + } + }, + "gemini-2.5-flash-stt": { + "provider": "google", + "model_type": "stt", + "description": "Gemini 2.5 Flash used for audio transcription via LiteLLM proxy (fast, cost-efficient)", + "featured": true, + "featured_rank": 6, + "highlights": [ + "Fast transcription", + "Cost-efficient", + "Gemini 2.5" + ], + "pricing": { + "source": "litellm_import", + "usage_kind": "stt", + "audio_per_minute": 0.0015 + } + }, + "gemini-2.5-flash-lite-stt": { + "provider": "google", + "model_type": "stt", + "description": "Gemini 2.5 Flash Lite used for audio transcription via LiteLLM proxy (lowest cost)", + "pricing": { + "source": "litellm_import", + "usage_kind": "stt", + "audio_per_minute": 0.00042 + } + }, + "gemini-2.5-pro": { + "provider": "google", + "model_type": "llm", + "pricing": { + "source": "litellm_import", + "usage_kind": "llm", + "input_per_1m": 1.25, + "output_per_1m": 10.0, + "cache_read_per_1m": 0.125 + } + }, + "gemini-2.5-flash": { + "provider": "google", + "model_type": "llm", + "pricing": { + "source": "litellm_import", + "usage_kind": "llm", + "input_per_1m": 0.3, + "output_per_1m": 2.5, + "cache_read_per_1m": 0.03, + "reasoning_per_1m": 2.5 + } + }, + "gemini-2.5-flash-lite": { + "provider": "google", + "model_type": "llm", + "description": "Gemini 2.5 Flash Lite — lowest-latency 2.5 model; ideal for diarisation, classification, and other low-reasoning multimodal tasks", + "pricing": { + "source": "litellm_import", + "usage_kind": "llm", + "input_per_1m": 0.1, + "output_per_1m": 0.4, + "cache_read_per_1m": 0.01, + "reasoning_per_1m": 0.4 + } + }, + "gemini-3-pro-preview": { + "provider": "google", + "model_type": "llm", + "description": "Gemini 3 Pro — advanced reasoning and multimodal understanding", + "pricing": { + "source": "litellm_import", + "usage_kind": "llm", + "input_per_1m": 2.0, + "output_per_1m": 12.0, + "cache_read_per_1m": 0.2 + } + }, + "gemini-3-flash-preview": { + "provider": "google", + "model_type": "llm", + "description": "Gemini 3 Flash — frontier-class performance at lower cost", + "pricing": { + "source": "litellm_import", + "usage_kind": "llm", + "input_per_1m": 0.5, + "output_per_1m": 3.0, + "cache_read_per_1m": 0.05, + "reasoning_per_1m": 3.0 + } + }, + "gemini-3.1-pro-preview": { + "provider": "google", + "model_type": "llm", + "description": "Gemini 3.1 Pro — latest reasoning model with 1M context", + "pricing": { + "source": "litellm_import", + "usage_kind": "llm", + "input_per_1m": 2.0, + "output_per_1m": 12.0, + "cache_read_per_1m": 0.2 + } + }, + "gemini-3.5-flash": { + "provider": "google", + "model_type": "llm", + "description": "Gemini 3.5 Flash — GA flagship Flash model for agentic coding and long-horizon tasks", + "featured": true, + "featured_rank": 7, + "highlights": [ + "GA stable", + "1M context", + "Agentic coding" + ], + "pricing": { + "source": "litellm_import", + "usage_kind": "llm", + "input_per_1m": 1.5, + "output_per_1m": 9.0, + "cache_read_per_1m": 0.15, + "reasoning_per_1m": 9.0 + } + }, + "gemini-3.5-flash-lite": { + "provider": "google", + "model_type": "llm", + "description": "Gemini 3.5 Flash Lite — fastest, lowest-cost 3.5 model for high-throughput extraction, classification, and subagent workflows", + "pricing": { + "source": "litellm_import", + "usage_kind": "llm", + "input_per_1m": 0.3, + "output_per_1m": 2.5, + "cache_read_per_1m": 0.03, + "reasoning_per_1m": 2.5 + } + }, + "gemini-3.1-flash-lite": { + "provider": "google", + "model_type": "llm", + "description": "Gemini 3.1 Flash Lite — cost-efficient stable model for high-volume lightweight tasks", + "pricing": { + "source": "litellm_import", + "usage_kind": "llm", + "input_per_1m": 0.25, + "output_per_1m": 1.5, + "cache_read_per_1m": 0.025, + "reasoning_per_1m": 1.5 + } + }, + "azure-speech-v1": { + "provider": "azure", + "model_type": "stt", + "description": "Azure Speech-to-text (batch and realtime)", + "pricing": { + "source": "litellm_import", + "usage_kind": "stt", + "audio_per_minute": 0.01668 + } + }, + "azure-openai-gpt4": { + "provider": "azure", + "model_type": "llm", + "description": "Legacy alias — maps to gpt-4 deployment; prefer azure-gpt-5-mini if that is your deployment name", + "pricing": { + "source": "litellm_import", + "usage_kind": "llm", + "input_per_1m": 30.0, + "output_per_1m": 60.0 + } + }, + "azure-gpt-4o": { + "provider": "azure", + "model_type": "llm", + "description": "Azure OpenAI GPT-4o — catalog key maps to deployment name gpt-4o", + "pricing": { + "source": "litellm_import", + "usage_kind": "llm", + "input_per_1m": 2.5, + "output_per_1m": 10.0, + "cache_read_per_1m": 1.25 + } + }, + "azure-gpt-4o-mini": { + "provider": "azure", + "model_type": "llm", + "description": "Azure OpenAI GPT-4o mini", + "pricing": { + "source": "litellm_import", + "usage_kind": "llm", + "input_per_1m": 0.165, + "output_per_1m": 0.66, + "cache_read_per_1m": 0.075 + } + }, + "azure-gpt-4.1": { + "provider": "azure", + "model_type": "llm", + "description": "Azure OpenAI GPT-4.1", + "pricing": { + "source": "litellm_import", + "usage_kind": "llm", + "input_per_1m": 2.0, + "output_per_1m": 8.0, + "cache_read_per_1m": 0.5 + } + }, + "azure-gpt-4.1-mini": { + "provider": "azure", + "model_type": "llm", + "description": "Azure OpenAI GPT-4.1 mini", + "pricing": { + "source": "litellm_import", + "usage_kind": "llm", + "input_per_1m": 0.4, + "output_per_1m": 1.6, + "cache_read_per_1m": 0.1 + } + }, + "azure-gpt-4.1-nano": { + "provider": "azure", + "model_type": "llm", + "description": "Azure OpenAI GPT-4.1 nano", + "pricing": { + "source": "litellm_import", + "usage_kind": "llm", + "input_per_1m": 0.1, + "output_per_1m": 0.4, + "cache_read_per_1m": 0.025 + } + }, + "azure-gpt-5": { + "provider": "azure", + "model_type": "llm", + "description": "Azure OpenAI GPT-5", + "pricing": { + "source": "litellm_import", + "usage_kind": "llm", + "input_per_1m": 1.25, + "output_per_1m": 10.0, + "cache_read_per_1m": 0.125 + } + }, + "azure-gpt-5-mini": { + "provider": "azure", + "model_type": "llm", + "description": "Azure OpenAI GPT-5 mini — pick this if your deployment is named gpt-5-mini", + "pricing": { + "source": "litellm_import", + "usage_kind": "llm", + "input_per_1m": 0.25, + "output_per_1m": 2.0, + "cache_read_per_1m": 0.025 + } + }, + "azure-gpt-5-nano": { + "provider": "azure", + "model_type": "llm", + "description": "Azure OpenAI GPT-5 nano", + "pricing": { + "source": "litellm_import", + "usage_kind": "llm", + "input_per_1m": 0.05, + "output_per_1m": 0.4, + "cache_read_per_1m": 0.005 + } + }, + "azure-gpt-5.1": { + "provider": "azure", + "model_type": "llm", + "description": "Azure OpenAI GPT-5.1", + "pricing": { + "source": "litellm_import", + "usage_kind": "llm", + "input_per_1m": 1.25, + "output_per_1m": 10.0, + "cache_read_per_1m": 0.125 + } + }, + "azure-gpt-5.2": { + "provider": "azure", + "model_type": "llm", + "description": "Azure OpenAI GPT-5.2", + "pricing": { + "source": "litellm_import", + "usage_kind": "llm", + "input_per_1m": 1.75, + "output_per_1m": 14.0, + "cache_read_per_1m": 0.175 + } + }, + "azure-o3": { + "provider": "azure", + "model_type": "llm", + "description": "Azure OpenAI o3 reasoning model", + "pricing": { + "source": "litellm_import", + "usage_kind": "llm", + "input_per_1m": 2.0, + "output_per_1m": 8.0, + "cache_read_per_1m": 0.5 + } + }, + "azure-o3-mini": { + "provider": "azure", + "model_type": "llm", + "description": "Azure OpenAI o3 mini", + "pricing": { + "source": "litellm_import", + "usage_kind": "llm", + "input_per_1m": 1.1, + "output_per_1m": 4.4, + "cache_read_per_1m": 0.55 + } + }, + "azure-o4-mini": { + "provider": "azure", + "model_type": "llm", + "description": "Azure OpenAI o4 mini", + "pricing": { + "source": "litellm_import", + "usage_kind": "llm", + "input_per_1m": 1.1, + "output_per_1m": 4.4, + "cache_read_per_1m": 0.275 + } + }, + "azure-tts-v1": { + "provider": "azure", + "model_type": "tts", + "description": "Azure neural text-to-speech", + "pricing": { + "source": "litellm_import", + "usage_kind": "tts", + "tts_per_1m_characters": 15.0 + } + }, + "aws-transcribe": { + "provider": "aws", + "model_type": "stt", + "pricing": { + "source": "litellm_import", + "usage_kind": "stt", + "audio_per_minute": 0.006 + } + }, + "aws-bedrock-claude": { + "provider": "aws", + "model_type": "llm", + "pricing": { + "source": "litellm_import", + "usage_kind": "llm", + "input_per_1m": 3.0, + "output_per_1m": 15.0, + "cache_read_per_1m": 0.3, + "cache_write_per_1m": 3.75 + } + }, + "aws-polly": { + "provider": "aws", + "model_type": "tts", + "pricing": { + "source": "litellm_import", + "usage_kind": "tts", + "tts_per_1m_characters": 16.0 + } + }, + "deepgram-flux": { + "provider": "deepgram", + "model_type": "stt", + "description": "Conversational STT with integrated turn detection for voice agents", + "featured": true, + "featured_rank": 4, + "highlights": [ + "Turn detection", + "Voice agents", + "Low latency" + ], + "pricing": { + "source": "litellm_import", + "usage_kind": "stt", + "audio_per_minute": 0.00432 + } + }, + "deepgram-nova-3": { + "provider": "deepgram", + "model_type": "stt", + "description": "Highest-performing general-purpose ASR model", + "pricing": { + "source": "litellm_import", + "usage_kind": "stt", + "audio_per_minute": 0.00432 + } + }, + "deepgram-nova-3-general": { + "provider": "deepgram", + "model_type": "stt", + "pricing": { + "source": "litellm_import", + "usage_kind": "stt", + "audio_per_minute": 0.00432 + } + }, + "deepgram-nova-3-general-preview-12-2025": { + "provider": "deepgram", + "model_type": "stt", + "pricing": { + "source": "litellm_import", + "usage_kind": "stt", + "audio_per_minute": 0.00432 + } + }, + "deepgram-nova-2": { + "provider": "deepgram", + "model_type": "stt", + "description": "Available for languages not yet supported by Nova-3", + "pricing": { + "source": "litellm_import", + "usage_kind": "stt", + "audio_per_minute": 0.00432 + } + }, + "pulse-v4": { + "provider": "smallest", + "model_type": "stt", + "description": "Smallest Pulse v4 speech-to-text (batch and realtime)", + "pricing": { + "source": "litellm_import", + "usage_kind": "stt", + "audio_per_minute": 0.006 + } + }, + "cartesia-sonic-3": { + "provider": "cartesia", + "model_type": "tts", + "description": "Flagship streaming TTS with laughter, volume/speed/emotion controls, 42 languages", + "featured": true, + "featured_rank": 1, + "highlights": [ + "Ultra-low latency", + "42 languages", + "Emotion controls" + ], + "pricing": { + "source": "cartesia.ai credits ~$50/1M chars (Pro tier, 1 credit/char)", + "usage_kind": "tts", + "tts_per_1m_characters": 50.0 + } + }, + "cartesia-sonic-3-mini": { + "provider": "cartesia", + "model_type": "tts", + "pricing": { + "source": "cartesia.ai ~$40/1M chars (mid-tier estimate)", + "usage_kind": "tts", + "tts_per_1m_characters": 40.0 + } + }, + "cartesia-sonic-3-nano": { + "provider": "cartesia", + "model_type": "tts", + "pricing": { + "source": "cartesia.ai Scale ~$37/1M chars", + "usage_kind": "tts", + "tts_per_1m_characters": 37.0 + } + }, + "eleven_v3": { + "provider": "elevenlabs", + "model_type": "tts", + "description": "Human-like and expressive speech generation", + "languages": "70+ languages", + "pricing": { + "source": "litellm_import", + "usage_kind": "tts", + "tts_per_1m_characters": 180.0 + } + }, + "eleven_ttv_v3": { + "provider": "elevenlabs", + "model_type": "tts", + "description": "Human-like and expressive voice design model (Text to Voice)", + "languages": "70+ languages", + "pricing": { + "source": "litellm_import", + "usage_kind": "tts", + "tts_per_1m_characters": 180.0 + } + }, + "eleven_multilingual_v2": { + "provider": "elevenlabs", + "model_type": "tts", + "description": "Most lifelike model with rich emotional expression", + "languages": "en, ja, zh, de, hi, fr, ko, pt, it, es, id, nl, tr, fil, pl, sv, bg, ro, ar, cs, el, fi, hr, ms, sk, da, ta, uk, ru", + "pricing": { + "source": "litellm_import", + "usage_kind": "tts", + "tts_per_1m_characters": 180.0 + } + }, + "eleven_flash_v2_5": { + "provider": "elevenlabs", + "model_type": "tts", + "description": "Ultra-fast model optimized for real-time use (~75ms)", + "languages": "en, ja, zh, de, hi, fr, ko, pt, it, es, id, nl, tr, fil, pl, sv, bg, ro, ar, cs, el, fi, hr, ms, sk, da, ta, uk, ru, hu, no, vi", + "featured": true, + "featured_rank": 2, + "highlights": [ + "~75ms latency", + "Real-time", + "Multilingual" + ], + "pricing": { + "source": "litellm_import", + "usage_kind": "tts", + "tts_per_1m_characters": 180.0 + } + }, + "eleven_turbo_v2_5": { + "provider": "elevenlabs", + "model_type": "tts", + "description": "High quality, low-latency model (~250ms-300ms)", + "languages": "en, ja, zh, de, hi, fr, ko, pt, it, es, id, nl, tr, fil, pl, sv, bg, ro, ar, cs, el, fi, hr, ms, sk, da, ta, uk, ru, hu, no, vi", + "pricing": { + "source": "litellm_import", + "usage_kind": "tts", + "tts_per_1m_characters": 180.0 + } + }, + "eleven_multilingual_sts_v2": { + "provider": "elevenlabs", + "model_type": "sts", + "description": "State-of-the-art multilingual voice changer model (Speech to Speech)", + "languages": "en, ja, zh, de, hi, fr, ko, pt, it, es, id, nl, tr, fil, pl, sv, bg, ro, ar, cs, el, fi, hr, ms, sk, da, ta, uk, ru", + "pricing": { + "source": "litellm_import", + "usage_kind": "tts", + "tts_per_1m_characters": 180.0 + } + }, + "eleven_multilingual_ttv_v2": { + "provider": "elevenlabs", + "model_type": "tts", + "description": "State-of-the-art multilingual voice designer model (Text to Voice)", + "languages": "en, ja, zh, de, hi, fr, ko, pt, it, es, id, nl, tr, fil, pl, sv, bg, ro, ar, cs, el, fi, hr, ms, sk, da, ta, uk, ru", + "pricing": { + "source": "litellm_import", + "usage_kind": "tts", + "tts_per_1m_characters": 180.0 + } + }, + "eleven_english_sts_v2": { + "provider": "elevenlabs", + "model_type": "sts", + "description": "English-only voice changer model (Speech to Speech)", + "languages": "en", + "pricing": { + "source": "litellm_import", + "usage_kind": "tts", + "tts_per_1m_characters": 180.0 + } + }, + "scribe_v2_realtime": { + "provider": "elevenlabs", + "model_type": "stt", + "description": "Real-time speech recognition model", + "languages": "90+ languages", + "featured": true, + "featured_rank": 5, + "highlights": [ + "Real-time STT", + "90+ languages", + "Streaming" + ], + "pricing": { + "source": "litellm_import", + "usage_kind": "stt", + "audio_per_minute": 0.00366 + } + }, + "scribe_v2": { + "provider": "elevenlabs", + "model_type": "stt", + "description": "Most accurate transcription model with keyterm prompting and entity detection", + "languages": "90+ languages", + "pricing": { + "source": "litellm_import", + "usage_kind": "stt", + "audio_per_minute": 0.00366 + } + }, + "eleven_text_to_sound_v2": { + "provider": "elevenlabs", + "model_type": "sound_effects", + "description": "Sound effects generation from text prompts", + "pricing": { + "source": "litellm_import", + "usage_kind": "tts", + "tts_per_1m_characters": 180.0 + } + }, + "music_v1": { + "provider": "elevenlabs", + "model_type": "music", + "description": "Studio-grade music generation from text prompts", + "languages": "en, es, de, ja, and more", + "pricing": { + "source": "litellm_import", + "usage_kind": "tts", + "tts_per_1m_characters": 180.0 + } + }, + "murf-falcon": { + "provider": "murf", + "model_type": "tts", + "description": "Murf FALCON – ultra-low latency streaming TTS", + "languages": "en, es, de, fr, hi, ja, ko, pt, it, zh, and 10+ more", + "voices_source_file": "murf_falcon_voices.json", + "pricing": { + "source": "murf.ai API $0.01/1K chars", + "usage_kind": "tts", + "tts_per_1m_characters": 10.0 + } + }, + "murf-gen2": { + "provider": "murf", + "model_type": "tts", + "description": "Murf GEN2 – high-quality, natural-sounding TTS", + "languages": "en, es, de, fr, hi, ja, ko, pt, it, zh, and 10+ more", + "voices_source_file": "murf_gen2_voices.json", + "pricing": { + "source": "murf.ai API $0.03/1K chars", + "usage_kind": "tts", + "tts_per_1m_characters": 30.0 + } + }, + "voicemaker-ai3": { + "provider": "voicemaker", + "model_type": "tts", + "description": "VoiceMaker AI3 neural TTS voices", + "languages": "en-US, en-GB, en-AU, es-ES, fr-FR, de-DE, ja-JP, zh-CN", + "voices_source_file": "voicemaker_voices.json", + "pricing": { + "source": "developer.voicemaker.in $25/1M chars × 1× (AI3)", + "usage_kind": "tts", + "tts_per_1m_characters": 25.0 + } + }, + "voicemaker-proplus": { + "provider": "voicemaker", + "model_type": "tts", + "description": "VoiceMaker ProPlus expressive multilingual voices", + "languages": "multi-lang + regional accents", + "voices_source_file": "voicemaker_voices.json", + "pricing": { + "source": "developer.voicemaker.in $25/1M chars × 2× (ProPlus Turbo)", + "usage_kind": "tts", + "tts_per_1m_characters": 50.0 + } + }, + "lightning-v3.1": { + "provider": "smallest", + "model_type": "tts", + "description": "Smallest Lightning v3.1 low-latency neural TTS", + "pricing": { + "source": "smallest.ai $0.175/10K chars", + "usage_kind": "tts", + "tts_per_1m_characters": 17.5 + } + }, + "sarvam-30b": { + "provider": "sarvam", + "model_type": "llm", + "description": "Sarvam 30B — balanced Indian-language + English chat model (64K context)", + "pricing": { + "source": "sarvam.ai ₹2.5 / ₹1.5 / ₹10 per 1M tokens @ 83 INR/USD", + "usage_kind": "llm", + "input_per_1m": 0.03012, + "output_per_1m": 0.120482, + "cache_read_per_1m": 0.018072 + } + }, + "sarvam-105b": { + "provider": "sarvam", + "model_type": "llm", + "description": "Sarvam 105B — flagship MoE for complex reasoning, coding, and agentic workflows (128K context)", + "pricing": { + "source": "sarvam.ai ₹4 / ₹2.5 / ₹16 per 1M tokens @ 83 INR/USD", + "usage_kind": "llm", + "input_per_1m": 0.048193, + "output_per_1m": 0.192771, + "cache_read_per_1m": 0.03012 + } + }, + "saaras:v3": { + "provider": "sarvam", + "model_type": "stt", + "description": "State-of-the-art Sarvam STT model with transcribe/translate modes", + "pricing": { + "source": "sarvam.ai ₹30/hour audio @ 83 INR/USD", + "usage_kind": "stt", + "audio_per_minute": 0.006 + } + }, + "bulbul:v3": { + "provider": "sarvam", + "model_type": "tts", + "description": "High quality Indian language TTS", + "voices": [ + { + "id": "aditya", + "name": "Aditya", + "gender": "Male" + }, + { + "id": "ritu", + "name": "Ritu", + "gender": "Female" + }, + { + "id": "ashutosh", + "name": "Ashutosh", + "gender": "Male" + }, + { + "id": "priya", + "name": "Priya", + "gender": "Female" + }, + { + "id": "neha", + "name": "Neha", + "gender": "Female" + }, + { + "id": "rahul", + "name": "Rahul", + "gender": "Male" + }, + { + "id": "pooja", + "name": "Pooja", + "gender": "Female" + }, + { + "id": "rohan", + "name": "Rohan", + "gender": "Male" + }, + { + "id": "simran", + "name": "Simran", + "gender": "Female" + }, + { + "id": "kavya", + "name": "Kavya", + "gender": "Female" + }, + { + "id": "amit", + "name": "Amit", + "gender": "Male" + }, + { + "id": "dev", + "name": "Dev", + "gender": "Male" + }, + { + "id": "ishita", + "name": "Ishita", + "gender": "Female" + }, + { + "id": "shreya", + "name": "Shreya", + "gender": "Female" + }, + { + "id": "ratan", + "name": "Ratan", + "gender": "Male" + }, + { + "id": "varun", + "name": "Varun", + "gender": "Male" + }, + { + "id": "manan", + "name": "Manan", + "gender": "Male" + }, + { + "id": "sumit", + "name": "Sumit", + "gender": "Male" + }, + { + "id": "roopa", + "name": "Roopa", + "gender": "Female" + }, + { + "id": "kabir", + "name": "Kabir", + "gender": "Male" + }, + { + "id": "aayan", + "name": "Aayan", + "gender": "Male" + }, + { + "id": "shubh", + "name": "Shubh", + "gender": "Male" + }, + { + "id": "advait", + "name": "Advait", + "gender": "Male" + }, + { + "id": "amelia", + "name": "Amelia", + "gender": "Female" + }, + { + "id": "sophia", + "name": "Sophia", + "gender": "Female" + }, + { + "id": "anand", + "name": "Anand", + "gender": "Male" + }, + { + "id": "tanya", + "name": "Tanya", + "gender": "Female" + }, + { + "id": "tarun", + "name": "Tarun", + "gender": "Male" + }, + { + "id": "sunny", + "name": "Sunny", + "gender": "Male" + }, + { + "id": "mani", + "name": "Mani", + "gender": "Male" + }, + { + "id": "gokul", + "name": "Gokul", + "gender": "Male" + }, + { + "id": "vijay", + "name": "Vijay", + "gender": "Male" + }, + { + "id": "shruti", + "name": "Shruti", + "gender": "Female" + }, + { + "id": "suhani", + "name": "Suhani", + "gender": "Female" + }, + { + "id": "mohit", + "name": "Mohit", + "gender": "Male" + }, + { + "id": "kavitha", + "name": "Kavitha", + "gender": "Female" + }, + { + "id": "rehan", + "name": "Rehan", + "gender": "Male" + }, + { + "id": "soham", + "name": "Soham", + "gender": "Male" + }, + { + "id": "rupali", + "name": "Rupali", + "gender": "Female" + } + ], + "pricing": { + "source": "sarvam.ai ₹30/10K chars @ 83 INR/USD", + "usage_kind": "tts", + "tts_per_1m_characters": 36.144578 + } + }, + "voice-agent-call": { + "provider": "telephony", + "model_type": "llm", + "description": "External voice-agent call duration (Vapi, Retell, ElevenLabs, Smallest); billed per minute", + "pricing": { + "source": "manual", + "usage_kind": "llm", + "audio_per_minute": 0.05 + } + }, + "unknown": { + "provider": "internal", + "model_type": "llm", + "description": "Fallback when model name is missing from usage records", + "pricing": { + "source": "manual", + "usage_kind": "llm", + "input_per_1m": 0, + "output_per_1m": 0 + } + } +} diff --git a/app/core/usage_context_middleware.py b/app/core/usage_context_middleware.py new file mode 100644 index 00000000..c9531c98 --- /dev/null +++ b/app/core/usage_context_middleware.py @@ -0,0 +1,43 @@ +"""Clear LLM usage ContextVar at request boundaries to avoid thread reuse leaks.""" + +from __future__ import annotations + +from uuid import UUID + +from starlette.middleware.base import BaseHTTPMiddleware +from starlette.requests import Request + +from app.services.usage.context import ( + infer_product_section_from_path, + reset_usage_context, + reset_usage_hints, + set_usage_context, + set_usage_hints, +) + + +class LLMUsageContextMiddleware(BaseHTTPMiddleware): + async def dispatch(self, request: Request, call_next): + section = infer_product_section_from_path(request.url.path) + request.state.usage_product_section = section + + workspace_hint = None + raw_ws = request.headers.get("x-workspace-id") or request.query_params.get( + "workspace_id" + ) + if raw_ws: + try: + workspace_hint = UUID(raw_ws) + except (TypeError, ValueError): + workspace_hint = None + + ctx_token = set_usage_context(None) + hint_tokens = set_usage_hints( + workspace_id=workspace_hint, + product_section=section, + ) + try: + return await call_next(request) + finally: + reset_usage_hints(hint_tokens) + reset_usage_context(ctx_token) diff --git a/app/core/usage_entitlement.py b/app/core/usage_entitlement.py new file mode 100644 index 00000000..cbedbf7d --- /dev/null +++ b/app/core/usage_entitlement.py @@ -0,0 +1,55 @@ +"""Enterprise entitlement for usage history (any catalog feature, not per-feature gates).""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Optional +from uuid import UUID + +from app.core.license import get_enabled_features, get_license_info + +OSS_USAGE_HISTORY_DAYS = 7 + + +@dataclass(frozen=True) +class UsagePolicySnapshot: + extended_history: bool + max_history_days: Optional[int] = None + + def as_dict(self) -> dict: + return { + "extended_history": self.extended_history, + "max_history_days": self.max_history_days, + } + + +def has_enterprise_entitlement(organization_id: Optional[UUID] = None) -> bool: + """Valid license with at least one catalog feature; org-scoped licenses must match.""" + if not get_enabled_features(): + return False + + info = get_license_info() + licensed_org = info.get("org_id") + if licensed_org is None: + return True + + if organization_id is None: + return False + + return str(organization_id) == str(licensed_org) + + +def deployment_has_entitlement() -> bool: + """Deployment-wide entitlement (license without org_id scoping).""" + if not get_enabled_features(): + return False + return get_license_info().get("org_id") is None + + +def get_usage_policy(organization_id: UUID) -> UsagePolicySnapshot: + if has_enterprise_entitlement(organization_id): + return UsagePolicySnapshot(extended_history=True, max_history_days=None) + return UsagePolicySnapshot( + extended_history=False, + max_history_days=OSS_USAGE_HISTORY_DAYS, + ) diff --git a/app/dependencies.py b/app/dependencies.py index 0dea6a91..b1824a48 100644 --- a/app/dependencies.py +++ b/app/dependencies.py @@ -21,6 +21,7 @@ from app.core.auth import Principal, get_principal # noqa: F401 - re-exported from app.core.auth.rbac import get_org_role from app.core.license import is_feature_enabled +from app.core.usage_entitlement import has_enterprise_entitlement from app.database import get_db from app.models.database import RoleEnum, Workspace, WorkspaceMember from app.core.auth.capabilities import capability_denied_message @@ -201,6 +202,22 @@ def get_workspace_context( is_org_admin=org_role == RoleEnum.ADMIN, ) request.state.workspace_context = ctx + + from app.services.usage.context import ( + LLMUsageProductSection, + ensure_usage_context, + ) + + section = getattr( + request.state, "usage_product_section", LLMUsageProductSection.OTHER + ) + ensure_usage_context( + organization_id, + workspace_id=workspace.id, + product_section=section + if isinstance(section, LLMUsageProductSection) + else LLMUsageProductSection.OTHER, + ) return ctx @@ -283,3 +300,26 @@ def _check( ) return _check + + +def require_enterprise_entitlement(): + """ + FastAPI dependency: valid enterprise license with any catalog feature. + Distinct from require_enterprise_feature (per-feature product gates). + """ + + def _check(organization_id: UUID = Depends(get_organization_id)): + if not has_enterprise_entitlement(organization_id): + raise HTTPException( + status_code=403, + detail={ + "error": "enterprise_license_required", + "message": ( + "This capability requires a valid EfficientAI Enterprise license. " + "Set EFFICIENTAI_LICENSE in your environment with any enterprise " + "feature enabled. Contact sales@efficientai.com for a license key." + ), + }, + ) + + return _check diff --git a/app/migrations/062_llm_usage_daily.py b/app/migrations/062_llm_usage_daily.py new file mode 100644 index 00000000..e160c955 --- /dev/null +++ b/app/migrations/062_llm_usage_daily.py @@ -0,0 +1,92 @@ +"""Migration: LLM usage daily rollups for org-scoped Usage reporting.""" + +from sqlalchemy import text +from sqlalchemy.orm import Session + +description = "Add llm_usage_daily table for LLM token/call rollups" + + +def upgrade(db: Session): + exists = db.execute( + text( + """ + SELECT 1 FROM information_schema.tables + WHERE table_name = 'llm_usage_daily' + """ + ) + ).first() + if exists: + print("llm_usage_daily already exists, skipping 062") + db.commit() + return + + db.execute( + text( + """ + CREATE TABLE llm_usage_daily ( + id UUID PRIMARY KEY DEFAULT gen_random_uuid(), + organization_id UUID NOT NULL REFERENCES organizations(id) ON DELETE CASCADE, + workspace_id UUID REFERENCES workspaces(id) ON DELETE SET NULL, + product_section VARCHAR(64) NOT NULL, + model VARCHAR(255) NOT NULL, + resource_id UUID, + resource_type VARCHAR(64), + usage_date DATE NOT NULL, + prompt_tokens BIGINT NOT NULL DEFAULT 0, + completion_tokens BIGINT NOT NULL DEFAULT 0, + cache_read_tokens BIGINT NOT NULL DEFAULT 0, + cache_creation_tokens BIGINT NOT NULL DEFAULT 0, + reasoning_tokens BIGINT NOT NULL DEFAULT 0, + call_count BIGINT NOT NULL DEFAULT 0, + created_at TIMESTAMPTZ NOT NULL DEFAULT now(), + updated_at TIMESTAMPTZ NOT NULL DEFAULT now() + ) + """ + ) + ) + db.execute( + text( + """ + CREATE INDEX ix_llm_usage_daily_org_date + ON llm_usage_daily (organization_id, usage_date) + """ + ) + ) + db.execute( + text( + """ + CREATE INDEX ix_llm_usage_daily_org_workspace_date + ON llm_usage_daily (organization_id, workspace_id, usage_date) + """ + ) + ) + db.execute( + text( + """ + CREATE INDEX ix_llm_usage_daily_org_resource_date + ON llm_usage_daily (organization_id, resource_id, usage_date) + """ + ) + ) + db.execute( + text( + """ + CREATE UNIQUE INDEX uq_llm_usage_daily_bucket + ON llm_usage_daily ( + organization_id, + COALESCE(workspace_id, '00000000-0000-0000-0000-000000000000'::uuid), + product_section, + model, + COALESCE(resource_id, '00000000-0000-0000-0000-000000000000'::uuid), + usage_date + ) + """ + ) + ) + db.commit() + print("Created llm_usage_daily table") + + +def downgrade(db: Session): + db.execute(text("DROP TABLE IF EXISTS llm_usage_daily")) + db.commit() diff --git a/app/migrations/063_usage_kind_stt_and_buffer.py b/app/migrations/063_usage_kind_stt_and_buffer.py new file mode 100644 index 00000000..fb1d9e4c --- /dev/null +++ b/app/migrations/063_usage_kind_stt_and_buffer.py @@ -0,0 +1,277 @@ +"""Migration: STT usage_kind/audio_seconds + Redis-fallback pending buffer.""" + +from sqlalchemy import text +from sqlalchemy.orm import Session + +description = ( + "Add usage_kind/audio_seconds to llm_usage_daily and usage_pending_buffer " + "for durable STT + LLM usage when Redis is unavailable" +) + +_ZERO_UUID = "00000000-0000-0000-0000-000000000000" + + +def _column_exists(db: Session, table: str, column: str) -> bool: + return ( + db.execute( + text( + """ + SELECT 1 FROM information_schema.columns + WHERE table_name = :table_name AND column_name = :column_name + """ + ), + {"table_name": table, "column_name": column}, + ).first() + is not None + ) + + +def _table_exists(db: Session, table: str) -> bool: + return ( + db.execute( + text( + """ + SELECT 1 FROM information_schema.tables + WHERE table_name = :table_name + """ + ), + {"table_name": table}, + ).first() + is not None + ) + + +def _dedupe_llm_usage_daily(db: Session) -> int: + """Merge duplicate bucket rows so the unique index can be created. + + Keeps the oldest row (MIN id), sums metrics into it, deletes extras. + """ + db.execute( + text( + f""" + WITH dupes AS ( + SELECT + organization_id, + COALESCE( + workspace_id, + '{_ZERO_UUID}'::uuid + ) AS ws_key, + product_section, + model, + COALESCE( + resource_id, + '{_ZERO_UUID}'::uuid + ) AS rid_key, + usage_date, + COALESCE(usage_kind, 'llm') AS kind_key, + MIN(id::text)::uuid AS keep_id, + SUM(prompt_tokens)::bigint AS prompt_tokens, + SUM(completion_tokens)::bigint AS completion_tokens, + SUM(cache_read_tokens)::bigint AS cache_read_tokens, + SUM(cache_creation_tokens)::bigint AS cache_creation_tokens, + SUM(reasoning_tokens)::bigint AS reasoning_tokens, + SUM(COALESCE(audio_seconds, 0))::bigint AS audio_seconds, + SUM(call_count)::bigint AS call_count + FROM llm_usage_daily + GROUP BY 1, 2, 3, 4, 5, 6, 7 + HAVING COUNT(*) > 1 + ) + UPDATE llm_usage_daily AS u SET + prompt_tokens = d.prompt_tokens, + completion_tokens = d.completion_tokens, + cache_read_tokens = d.cache_read_tokens, + cache_creation_tokens = d.cache_creation_tokens, + reasoning_tokens = d.reasoning_tokens, + audio_seconds = d.audio_seconds, + call_count = d.call_count, + updated_at = now() + FROM dupes AS d + WHERE u.id = d.keep_id + """ + ) + ) + + result = db.execute( + text( + f""" + WITH keepers AS ( + SELECT MIN(id::text)::uuid AS keep_id + FROM llm_usage_daily + GROUP BY + organization_id, + COALESCE(workspace_id, '{_ZERO_UUID}'::uuid), + product_section, + model, + COALESCE(resource_id, '{_ZERO_UUID}'::uuid), + usage_date, + COALESCE(usage_kind, 'llm') + HAVING COUNT(*) > 1 + ) + DELETE FROM llm_usage_daily AS u + USING keepers AS k, + llm_usage_daily AS peer + WHERE peer.id = k.keep_id + AND u.id <> k.keep_id + AND u.organization_id = peer.organization_id + AND u.product_section = peer.product_section + AND u.model = peer.model + AND u.usage_date = peer.usage_date + AND COALESCE(u.usage_kind, 'llm') = COALESCE(peer.usage_kind, 'llm') + AND u.workspace_id IS NOT DISTINCT FROM peer.workspace_id + AND u.resource_id IS NOT DISTINCT FROM peer.resource_id + """ + ) + ) + return int(result.rowcount or 0) + + +def _ensure_unique_bucket_index(db: Session) -> None: + """Drop legacy/broken unique index and recreate with usage_kind.""" + db.execute(text("DROP INDEX IF EXISTS uq_llm_usage_daily_bucket")) + db.execute( + text( + f""" + CREATE UNIQUE INDEX uq_llm_usage_daily_bucket + ON llm_usage_daily ( + organization_id, + COALESCE(workspace_id, '{_ZERO_UUID}'::uuid), + product_section, + model, + COALESCE(resource_id, '{_ZERO_UUID}'::uuid), + usage_date, + usage_kind + ) + """ + ) + ) + + +def _ensure_llm_usage_daily_base_columns(db: Session) -> None: + """Align pre-062 tables with the 062 schema before dedupe/index steps.""" + if not _column_exists(db, "llm_usage_daily", "resource_id"): + db.execute(text("ALTER TABLE llm_usage_daily ADD COLUMN resource_id UUID")) + if not _column_exists(db, "llm_usage_daily", "resource_type"): + db.execute( + text("ALTER TABLE llm_usage_daily ADD COLUMN resource_type VARCHAR(64)") + ) + if not _column_exists(db, "llm_usage_daily", "updated_at"): + db.execute( + text( + """ + ALTER TABLE llm_usage_daily + ADD COLUMN updated_at TIMESTAMPTZ NOT NULL DEFAULT now() + """ + ) + ) + + +def upgrade(db: Session): + if not _table_exists(db, "llm_usage_daily"): + print("llm_usage_daily missing; run 062 first — skipping 063") + db.commit() + return + + _ensure_llm_usage_daily_base_columns(db) + + if not _column_exists(db, "llm_usage_daily", "usage_kind"): + db.execute( + text( + """ + ALTER TABLE llm_usage_daily + ADD COLUMN usage_kind VARCHAR(16) NOT NULL DEFAULT 'llm' + """ + ) + ) + + if not _column_exists(db, "llm_usage_daily", "audio_seconds"): + db.execute( + text( + """ + ALTER TABLE llm_usage_daily + ADD COLUMN audio_seconds BIGINT NOT NULL DEFAULT 0 + """ + ) + ) + + removed = _dedupe_llm_usage_daily(db) + if removed: + print(f"Merged/removed {removed} duplicate llm_usage_daily row(s)") + + _ensure_unique_bucket_index(db) + + db.execute( + text( + """ + CREATE INDEX IF NOT EXISTS ix_llm_usage_daily_org_kind_date + ON llm_usage_daily (organization_id, usage_kind, usage_date) + """ + ) + ) + print("Ensured usage_kind + audio_seconds + unique bucket index") + + if not _table_exists(db, "usage_pending_buffer"): + db.execute( + text( + """ + CREATE TABLE usage_pending_buffer ( + id UUID PRIMARY KEY DEFAULT gen_random_uuid(), + organization_id UUID NOT NULL, + workspace_id UUID, + product_section VARCHAR(64) NOT NULL, + model VARCHAR(255) NOT NULL, + context JSONB NOT NULL DEFAULT '{}'::jsonb, + usage_date DATE NOT NULL, + usage_kind VARCHAR(16) NOT NULL DEFAULT 'llm', + prompt_tokens BIGINT NOT NULL DEFAULT 0, + completion_tokens BIGINT NOT NULL DEFAULT 0, + cache_read_tokens BIGINT NOT NULL DEFAULT 0, + cache_creation_tokens BIGINT NOT NULL DEFAULT 0, + reasoning_tokens BIGINT NOT NULL DEFAULT 0, + audio_seconds BIGINT NOT NULL DEFAULT 0, + call_count BIGINT NOT NULL DEFAULT 0, + created_at TIMESTAMPTZ NOT NULL DEFAULT now() + ) + """ + ) + ) + db.execute( + text( + """ + CREATE INDEX IF NOT EXISTS ix_usage_pending_buffer_org_created + ON usage_pending_buffer (organization_id, created_at) + """ + ) + ) + print("Created usage_pending_buffer") + + if not _table_exists(db, "usage_committed_claims"): + db.execute( + text( + """ + CREATE TABLE usage_committed_claims ( + claim_key TEXT PRIMARY KEY, + organization_id UUID NOT NULL, + committed_at TIMESTAMPTZ NOT NULL DEFAULT now() + ) + """ + ) + ) + db.execute( + text( + """ + CREATE INDEX IF NOT EXISTS ix_usage_committed_claims_committed_at + ON usage_committed_claims (committed_at) + """ + ) + ) + print("Created usage_committed_claims") + + db.commit() + + +def downgrade(db: Session): + db.execute(text("DROP TABLE IF EXISTS usage_committed_claims")) + db.execute(text("DROP TABLE IF EXISTS usage_pending_buffer")) + db.execute(text("DROP INDEX IF EXISTS ix_llm_usage_daily_org_kind_date")) + # Keep usage_kind/audio_seconds columns on downgrade to avoid data loss. + db.commit() diff --git a/app/migrations/064_usage_kind_tts_characters.py b/app/migrations/064_usage_kind_tts_characters.py new file mode 100644 index 00000000..57d42da0 --- /dev/null +++ b/app/migrations/064_usage_kind_tts_characters.py @@ -0,0 +1,74 @@ +"""Migration: TTS usage_kind + tts_characters metric.""" + +from sqlalchemy import text +from sqlalchemy.orm import Session + +description = "Add tts_characters to llm_usage_daily and usage_pending_buffer for TTS usage" + + +def _column_exists(db: Session, table: str, column: str) -> bool: + return ( + db.execute( + text( + """ + SELECT 1 FROM information_schema.columns + WHERE table_name = :table_name AND column_name = :column_name + """ + ), + {"table_name": table, "column_name": column}, + ).first() + is not None + ) + + +def _table_exists(db: Session, table: str) -> bool: + return ( + db.execute( + text( + """ + SELECT 1 FROM information_schema.tables + WHERE table_name = :table_name + """ + ), + {"table_name": table}, + ).first() + is not None + ) + + +def upgrade(db: Session): + if _table_exists(db, "llm_usage_daily") and not _column_exists( + db, "llm_usage_daily", "tts_characters" + ): + db.execute( + text( + """ + ALTER TABLE llm_usage_daily + ADD COLUMN tts_characters BIGINT NOT NULL DEFAULT 0 + """ + ) + ) + print("Added tts_characters to llm_usage_daily") + + if _table_exists(db, "usage_pending_buffer") and not _column_exists( + db, "usage_pending_buffer", "tts_characters" + ): + db.execute( + text( + """ + ALTER TABLE usage_pending_buffer + ADD COLUMN tts_characters BIGINT NOT NULL DEFAULT 0 + """ + ) + ) + print("Added tts_characters to usage_pending_buffer") + + db.commit() + + +def downgrade(db: Session): + if _column_exists(db, "llm_usage_daily", "tts_characters"): + db.execute(text("ALTER TABLE llm_usage_daily DROP COLUMN tts_characters")) + if _column_exists(db, "usage_pending_buffer", "tts_characters"): + db.execute(text("ALTER TABLE usage_pending_buffer DROP COLUMN tts_characters")) + db.commit() diff --git a/app/migrations/065_usage_context_jsonb.py b/app/migrations/065_usage_context_jsonb.py new file mode 100644 index 00000000..4dd62378 --- /dev/null +++ b/app/migrations/065_usage_context_jsonb.py @@ -0,0 +1,216 @@ +"""Migration: JSONB context column for usage attribution metadata.""" + +from sqlalchemy import text +from sqlalchemy.orm import Session + +description = ( + "Move resource_id/resource_type into llm_usage_daily.context JSONB and " + "recreate bucket uniqueness on context keys" +) + +_ZERO_UUID = "00000000-0000-0000-0000-000000000000" + + +def _column_exists(db: Session, table: str, column: str) -> bool: + return ( + db.execute( + text( + """ + SELECT 1 FROM information_schema.columns + WHERE table_name = :table_name AND column_name = :column_name + """ + ), + {"table_name": table, "column_name": column}, + ).first() + is not None + ) + + +def _table_exists(db: Session, table: str) -> bool: + return ( + db.execute( + text( + """ + SELECT 1 FROM information_schema.tables + WHERE table_name = :table_name + """ + ), + {"table_name": table}, + ).first() + is not None + ) + + +def _ensure_unique_bucket_index(db: Session) -> None: + db.execute(text("DROP INDEX IF EXISTS uq_llm_usage_daily_bucket")) + db.execute( + text( + f""" + CREATE UNIQUE INDEX uq_llm_usage_daily_bucket + ON llm_usage_daily ( + organization_id, + COALESCE(workspace_id, '{_ZERO_UUID}'::uuid), + product_section, + model, + usage_date, + usage_kind, + context + ) + """ + ) + ) + + +def _migrate_table_context(db: Session, table: str) -> None: + if not _table_exists(db, table): + return + if not _column_exists(db, table, "context"): + db.execute( + text( + f""" + ALTER TABLE {table} + ADD COLUMN context JSONB NOT NULL DEFAULT '{{}}'::jsonb + """ + ) + ) + + if _column_exists(db, table, "resource_id") or _column_exists(db, table, "resource_type"): + db.execute( + text( + f""" + UPDATE {table} + SET context = COALESCE(context, '{{}}'::jsonb) + || CASE + WHEN resource_id IS NOT NULL THEN + jsonb_build_object('resource_id', resource_id::text) + ELSE '{{}}'::jsonb + END + || CASE + WHEN resource_type IS NOT NULL AND resource_type <> '' THEN + jsonb_build_object('resource_type', resource_type) + ELSE '{{}}'::jsonb + END + WHERE resource_id IS NOT NULL + OR (resource_type IS NOT NULL AND resource_type <> '') + """ + ) + ) + if _column_exists(db, table, "resource_id"): + db.execute(text(f"DROP INDEX IF EXISTS ix_llm_usage_daily_org_resource_date")) + db.execute(text(f"ALTER TABLE {table} DROP COLUMN resource_id")) + if _column_exists(db, table, "resource_type"): + db.execute(text(f"ALTER TABLE {table} DROP COLUMN resource_type")) + + +def upgrade(db: Session): + if not _table_exists(db, "llm_usage_daily"): + print("llm_usage_daily missing; run 062 first — skipping 065") + db.commit() + return + + _migrate_table_context(db, "llm_usage_daily") + _migrate_table_context(db, "usage_pending_buffer") + + _ensure_unique_bucket_index(db) + + db.execute( + text( + """ + CREATE INDEX IF NOT EXISTS ix_llm_usage_daily_context_gin + ON llm_usage_daily USING gin (context) + """ + ) + ) + db.execute( + text( + """ + CREATE INDEX IF NOT EXISTS ix_llm_usage_daily_context_resource_id + ON llm_usage_daily ((context->>'resource_id')) + WHERE context ? 'resource_id' + """ + ) + ) + db.execute( + text( + """ + CREATE INDEX IF NOT EXISTS ix_llm_usage_daily_context_call_import_id + ON llm_usage_daily ((context->>'call_import_id')) + WHERE context ? 'call_import_id' + """ + ) + ) + db.execute( + text( + """ + CREATE INDEX IF NOT EXISTS ix_llm_usage_daily_context_evaluation_id + ON llm_usage_daily ((context->>'evaluation_id')) + WHERE context ? 'evaluation_id' + """ + ) + ) + db.execute( + text( + """ + CREATE INDEX IF NOT EXISTS ix_llm_usage_daily_context_evaluation_row_id + ON llm_usage_daily ((context->>'evaluation_row_id')) + WHERE context ? 'evaluation_row_id' + """ + ) + ) + print("Added context JSONB + migrated resource attribution") + db.commit() + + +def downgrade(db: Session): + if not _table_exists(db, "llm_usage_daily"): + db.commit() + return + + if not _column_exists(db, "llm_usage_daily", "resource_id"): + db.execute( + text( + """ + ALTER TABLE llm_usage_daily + ADD COLUMN resource_id UUID, + ADD COLUMN resource_type VARCHAR(64) + """ + ) + ) + db.execute( + text( + """ + UPDATE llm_usage_daily + SET resource_id = NULLIF(context->>'resource_id', '')::uuid, + resource_type = NULLIF(context->>'resource_type', '') + WHERE context IS NOT NULL + """ + ) + ) + + db.execute(text("DROP INDEX IF EXISTS ix_llm_usage_daily_context_evaluation_row_id")) + db.execute(text("DROP INDEX IF EXISTS ix_llm_usage_daily_context_evaluation_id")) + db.execute(text("DROP INDEX IF EXISTS ix_llm_usage_daily_context_call_import_id")) + db.execute(text("DROP INDEX IF EXISTS ix_llm_usage_daily_context_resource_id")) + db.execute(text("DROP INDEX IF EXISTS ix_llm_usage_daily_context_gin")) + db.execute(text("DROP INDEX IF EXISTS uq_llm_usage_daily_bucket")) + + if _column_exists(db, "llm_usage_daily", "context"): + db.execute(text("ALTER TABLE llm_usage_daily DROP COLUMN context")) + + db.execute( + text( + f""" + CREATE UNIQUE INDEX uq_llm_usage_daily_bucket + ON llm_usage_daily ( + organization_id, + COALESCE(workspace_id, '{_ZERO_UUID}'::uuid), + product_section, + model, + COALESCE(resource_id, '{_ZERO_UUID}'::uuid), + usage_date, + usage_kind + ) + """ + ) + ) + db.commit() diff --git a/app/migrations/066_backfill_llm_usage_workspace_from_call_import.py b/app/migrations/066_backfill_llm_usage_workspace_from_call_import.py new file mode 100644 index 00000000..478afd4f --- /dev/null +++ b/app/migrations/066_backfill_llm_usage_workspace_from_call_import.py @@ -0,0 +1,279 @@ +"""Migration: Backfill llm_usage_daily.workspace_id from call-import attribution.""" + +from sqlalchemy import text +from sqlalchemy.orm import Session + +description = ( + "Backfill workspace_id on llm_usage_daily and usage_pending_buffer rows " + "that have call-import context but were recorded with NULL workspace_id" +) + +_ZERO_UUID = "00000000-0000-0000-0000-000000000000" +_UUID_RE = "^[0-9a-fA-F]{8}-[0-9a-fA-F]{4}-[0-9a-fA-F]{4}-[0-9a-fA-F]{4}-[0-9a-fA-F]{12}$" + + +def _column_exists(db: Session, table: str, column: str) -> bool: + return ( + db.execute( + text( + """ + SELECT 1 FROM information_schema.columns + WHERE table_name = :table_name AND column_name = :column_name + """ + ), + {"table_name": table, "column_name": column}, + ).first() + is not None + ) + + +def _table_exists(db: Session, table: str) -> bool: + return ( + db.execute( + text( + """ + SELECT 1 FROM information_schema.tables + WHERE table_name = :table_name + """ + ), + {"table_name": table}, + ).first() + is not None + ) + + +def _dedupe_llm_usage_daily(db: Session) -> int: + """Merge duplicate buckets after workspace backfill.""" + tts_col = ( + "SUM(COALESCE(tts_characters, 0))::bigint AS tts_characters," + if _column_exists(db, "llm_usage_daily", "tts_characters") + else "" + ) + tts_set = ( + "tts_characters = d.tts_characters," + if _column_exists(db, "llm_usage_daily", "tts_characters") + else "" + ) + + db.execute( + text( + f""" + WITH dupes AS ( + SELECT + organization_id, + COALESCE(workspace_id, '{_ZERO_UUID}'::uuid) AS ws_key, + product_section, + model, + usage_date, + COALESCE(usage_kind, 'llm') AS kind_key, + context, + MIN(id::text)::uuid AS keep_id, + SUM(prompt_tokens)::bigint AS prompt_tokens, + SUM(completion_tokens)::bigint AS completion_tokens, + SUM(cache_read_tokens)::bigint AS cache_read_tokens, + SUM(cache_creation_tokens)::bigint AS cache_creation_tokens, + SUM(reasoning_tokens)::bigint AS reasoning_tokens, + SUM(COALESCE(audio_seconds, 0))::bigint AS audio_seconds, + {tts_col} + SUM(call_count)::bigint AS call_count + FROM llm_usage_daily + GROUP BY 1, 2, 3, 4, 5, 6, 7 + HAVING COUNT(*) > 1 + ) + UPDATE llm_usage_daily AS u SET + prompt_tokens = d.prompt_tokens, + completion_tokens = d.completion_tokens, + cache_read_tokens = d.cache_read_tokens, + cache_creation_tokens = d.cache_creation_tokens, + reasoning_tokens = d.reasoning_tokens, + audio_seconds = d.audio_seconds, + {tts_set} + call_count = d.call_count, + updated_at = now() + FROM dupes AS d + WHERE u.id = d.keep_id + """ + ) + ) + + result = db.execute( + text( + f""" + WITH keepers AS ( + SELECT MIN(id::text)::uuid AS keep_id + FROM llm_usage_daily + GROUP BY + organization_id, + COALESCE(workspace_id, '{_ZERO_UUID}'::uuid), + product_section, + model, + usage_date, + COALESCE(usage_kind, 'llm'), + context + HAVING COUNT(*) > 1 + ) + DELETE FROM llm_usage_daily AS u + USING keepers AS k, + llm_usage_daily AS peer + WHERE peer.id = k.keep_id + AND u.id <> k.keep_id + AND u.organization_id = peer.organization_id + AND u.product_section = peer.product_section + AND u.model = peer.model + AND u.usage_date = peer.usage_date + AND COALESCE(u.usage_kind, 'llm') = COALESCE(peer.usage_kind, 'llm') + AND u.workspace_id IS NOT DISTINCT FROM peer.workspace_id + AND u.context IS NOT DISTINCT FROM peer.context + """ + ) + ) + return int(result.rowcount or 0) + + +def _backfill_table_workspace(db: Session, table: str) -> int: + if not _table_exists(db, table): + return 0 + if not _column_exists(db, table, "workspace_id"): + return 0 + if not _column_exists(db, table, "context"): + return 0 + + total = 0 + + # context.call_import_id + result = db.execute( + text( + f""" + UPDATE {table} AS u + SET workspace_id = ci.workspace_id + FROM call_imports AS ci + WHERE u.workspace_id IS NULL + AND u.context ? 'call_import_id' + AND u.context->>'call_import_id' ~ :uuid_re + AND ci.id = (u.context->>'call_import_id')::uuid + AND ci.organization_id = u.organization_id + """ + ), + {"uuid_re": _UUID_RE}, + ) + total += int(result.rowcount or 0) + + # resource_type = call_import + result = db.execute( + text( + f""" + UPDATE {table} AS u + SET workspace_id = ci.workspace_id + FROM call_imports AS ci + WHERE u.workspace_id IS NULL + AND u.context->>'resource_type' = 'call_import' + AND u.context->>'resource_id' ~ :uuid_re + AND ci.id = (u.context->>'resource_id')::uuid + AND ci.organization_id = u.organization_id + """ + ), + {"uuid_re": _UUID_RE}, + ) + total += int(result.rowcount or 0) + + # context.evaluation_id + result = db.execute( + text( + f""" + UPDATE {table} AS u + SET workspace_id = e.workspace_id + FROM call_import_evaluations AS e + WHERE u.workspace_id IS NULL + AND u.context ? 'evaluation_id' + AND u.context->>'evaluation_id' ~ :uuid_re + AND e.id = (u.context->>'evaluation_id')::uuid + AND e.organization_id = u.organization_id + """ + ), + {"uuid_re": _UUID_RE}, + ) + total += int(result.rowcount or 0) + + # resource_type = call_import_evaluation + result = db.execute( + text( + f""" + UPDATE {table} AS u + SET workspace_id = e.workspace_id + FROM call_import_evaluations AS e + WHERE u.workspace_id IS NULL + AND u.context->>'resource_type' = 'call_import_evaluation' + AND u.context->>'resource_id' ~ :uuid_re + AND e.id = (u.context->>'resource_id')::uuid + AND e.organization_id = u.organization_id + """ + ), + {"uuid_re": _UUID_RE}, + ) + total += int(result.rowcount or 0) + + # context.call_import_row_id + result = db.execute( + text( + f""" + UPDATE {table} AS u + SET workspace_id = cir.workspace_id + FROM call_import_rows AS cir + WHERE u.workspace_id IS NULL + AND u.context ? 'call_import_row_id' + AND u.context->>'call_import_row_id' ~ :uuid_re + AND cir.id = (u.context->>'call_import_row_id')::uuid + AND cir.organization_id = u.organization_id + """ + ), + {"uuid_re": _UUID_RE}, + ) + total += int(result.rowcount or 0) + + # context.evaluation_row_id + result = db.execute( + text( + f""" + UPDATE {table} AS u + SET workspace_id = e.workspace_id + FROM call_import_evaluation_rows AS er + JOIN call_import_evaluations AS e ON e.id = er.evaluation_id + WHERE u.workspace_id IS NULL + AND u.context ? 'evaluation_row_id' + AND u.context->>'evaluation_row_id' ~ :uuid_re + AND er.id = (u.context->>'evaluation_row_id')::uuid + AND e.organization_id = u.organization_id + """ + ), + {"uuid_re": _UUID_RE}, + ) + total += int(result.rowcount or 0) + + return total + + +def upgrade(db: Session): + if not _table_exists(db, "llm_usage_daily"): + print("llm_usage_daily missing; skipping 066") + db.commit() + return + + daily_updated = _backfill_table_workspace(db, "llm_usage_daily") + print(f"Backfilled workspace_id on {daily_updated} llm_usage_daily row(s)") + + removed = _dedupe_llm_usage_daily(db) + if removed: + print(f"Removed {removed} duplicate llm_usage_daily row(s) after merge") + + buffer_updated = _backfill_table_workspace(db, "usage_pending_buffer") + if buffer_updated: + print( + f"Backfilled workspace_id on {buffer_updated} usage_pending_buffer row(s)" + ) + + db.commit() + + +def downgrade(db: Session): + # No-op: cannot distinguish backfilled workspace_id from originally recorded values. + db.commit() diff --git a/app/migrations/067_fix_llm_usage_bucket_unique_index.py b/app/migrations/067_fix_llm_usage_bucket_unique_index.py new file mode 100644 index 00000000..98a78e5d --- /dev/null +++ b/app/migrations/067_fix_llm_usage_bucket_unique_index.py @@ -0,0 +1,264 @@ +"""Migration: Reconcile llm_usage_daily bucket unique index with full JSONB context.""" + +from sqlalchemy import text +from sqlalchemy.orm import Session + +description = ( + "Merge legacy resource-scoped usage buckets and recreate uq_llm_usage_daily_bucket " + "on full context JSONB (per-row attribution)" +) + +_ZERO_UUID = "00000000-0000-0000-0000-000000000000" + + +def _table_exists(db: Session, table: str) -> bool: + return ( + db.execute( + text( + """ + SELECT 1 FROM information_schema.tables + WHERE table_name = :table_name + """ + ), + {"table_name": table}, + ).first() + is not None + ) + + +def _column_exists(db: Session, table: str, column: str) -> bool: + return ( + db.execute( + text( + """ + SELECT 1 FROM information_schema.columns + WHERE table_name = :table_name AND column_name = :column_name + """ + ), + {"table_name": table, "column_name": column}, + ).first() + is not None + ) + + +def _dedupe_by_context_resource_keys(db: Session) -> int: + """Merge rows that share context resource_id/resource_type (legacy unique index).""" + tts_col = ( + "SUM(COALESCE(tts_characters, 0))::bigint AS tts_characters," + if _column_exists(db, "llm_usage_daily", "tts_characters") + else "" + ) + tts_set = ( + "tts_characters = d.tts_characters," + if _column_exists(db, "llm_usage_daily", "tts_characters") + else "" + ) + db.execute( + text( + f""" + WITH dupes AS ( + SELECT + organization_id, + COALESCE(workspace_id, '{_ZERO_UUID}'::uuid) AS ws_key, + product_section, + model, + usage_date, + COALESCE(usage_kind, 'llm') AS kind_key, + COALESCE(context->>'resource_id', '') AS res_id_key, + COALESCE(context->>'resource_type', '') AS res_type_key, + MIN(id::text)::uuid AS keep_id, + SUM(prompt_tokens)::bigint AS prompt_tokens, + SUM(completion_tokens)::bigint AS completion_tokens, + SUM(cache_read_tokens)::bigint AS cache_read_tokens, + SUM(cache_creation_tokens)::bigint AS cache_creation_tokens, + SUM(reasoning_tokens)::bigint AS reasoning_tokens, + SUM(COALESCE(audio_seconds, 0))::bigint AS audio_seconds, + {tts_col} + SUM(call_count)::bigint AS call_count + FROM llm_usage_daily + GROUP BY 1, 2, 3, 4, 5, 6, 7, 8 + HAVING COUNT(*) > 1 + ) + UPDATE llm_usage_daily AS u SET + prompt_tokens = d.prompt_tokens, + completion_tokens = d.completion_tokens, + cache_read_tokens = d.cache_read_tokens, + cache_creation_tokens = d.cache_creation_tokens, + reasoning_tokens = d.reasoning_tokens, + audio_seconds = d.audio_seconds, + {tts_set} + call_count = d.call_count, + updated_at = now() + FROM dupes AS d + WHERE u.id = d.keep_id + """ + ) + ) + result = db.execute( + text( + f""" + WITH keepers AS ( + SELECT MIN(id::text)::uuid AS keep_id + FROM llm_usage_daily + GROUP BY + organization_id, + COALESCE(workspace_id, '{_ZERO_UUID}'::uuid), + product_section, + model, + usage_date, + COALESCE(usage_kind, 'llm'), + COALESCE(context->>'resource_id', ''), + COALESCE(context->>'resource_type', '') + HAVING COUNT(*) > 1 + ) + DELETE FROM llm_usage_daily AS u + USING keepers AS k, + llm_usage_daily AS peer + WHERE peer.id = k.keep_id + AND u.id <> k.keep_id + AND u.organization_id = peer.organization_id + AND u.product_section = peer.product_section + AND u.model = peer.model + AND u.usage_date = peer.usage_date + AND COALESCE(u.usage_kind, 'llm') = COALESCE(peer.usage_kind, 'llm') + AND u.workspace_id IS NOT DISTINCT FROM peer.workspace_id + AND COALESCE(u.context->>'resource_id', '') = + COALESCE(peer.context->>'resource_id', '') + AND COALESCE(u.context->>'resource_type', '') = + COALESCE(peer.context->>'resource_type', '') + """ + ) + ) + return int(result.rowcount or 0) + + +def _dedupe_by_full_context(db: Session) -> int: + tts_col = ( + "SUM(COALESCE(tts_characters, 0))::bigint AS tts_characters," + if _column_exists(db, "llm_usage_daily", "tts_characters") + else "" + ) + tts_set = ( + "tts_characters = d.tts_characters," + if _column_exists(db, "llm_usage_daily", "tts_characters") + else "" + ) + db.execute( + text( + f""" + WITH dupes AS ( + SELECT + organization_id, + COALESCE(workspace_id, '{_ZERO_UUID}'::uuid) AS ws_key, + product_section, + model, + usage_date, + COALESCE(usage_kind, 'llm') AS kind_key, + context, + MIN(id::text)::uuid AS keep_id, + SUM(prompt_tokens)::bigint AS prompt_tokens, + SUM(completion_tokens)::bigint AS completion_tokens, + SUM(cache_read_tokens)::bigint AS cache_read_tokens, + SUM(cache_creation_tokens)::bigint AS cache_creation_tokens, + SUM(reasoning_tokens)::bigint AS reasoning_tokens, + SUM(COALESCE(audio_seconds, 0))::bigint AS audio_seconds, + {tts_col} + SUM(call_count)::bigint AS call_count + FROM llm_usage_daily + GROUP BY 1, 2, 3, 4, 5, 6, 7 + HAVING COUNT(*) > 1 + ) + UPDATE llm_usage_daily AS u SET + prompt_tokens = d.prompt_tokens, + completion_tokens = d.completion_tokens, + cache_read_tokens = d.cache_read_tokens, + cache_creation_tokens = d.cache_creation_tokens, + reasoning_tokens = d.reasoning_tokens, + audio_seconds = d.audio_seconds, + {tts_set} + call_count = d.call_count, + updated_at = now() + FROM dupes AS d + WHERE u.id = d.keep_id + """ + ) + ) + result = db.execute( + text( + f""" + WITH keepers AS ( + SELECT MIN(id::text)::uuid AS keep_id + FROM llm_usage_daily + GROUP BY + organization_id, + COALESCE(workspace_id, '{_ZERO_UUID}'::uuid), + product_section, + model, + usage_date, + COALESCE(usage_kind, 'llm'), + context + HAVING COUNT(*) > 1 + ) + DELETE FROM llm_usage_daily AS u + USING keepers AS k, + llm_usage_daily AS peer + WHERE peer.id = k.keep_id + AND u.id <> k.keep_id + AND u.organization_id = peer.organization_id + AND u.product_section = peer.product_section + AND u.model = peer.model + AND u.usage_date = peer.usage_date + AND COALESCE(u.usage_kind, 'llm') = COALESCE(peer.usage_kind, 'llm') + AND u.workspace_id IS NOT DISTINCT FROM peer.workspace_id + AND u.context IS NOT DISTINCT FROM peer.context + """ + ) + ) + return int(result.rowcount or 0) + + +def _ensure_context_bucket_index(db: Session) -> None: + db.execute(text("DROP INDEX IF EXISTS uq_llm_usage_daily_bucket")) + db.execute( + text( + f""" + CREATE UNIQUE INDEX uq_llm_usage_daily_bucket + ON llm_usage_daily ( + organization_id, + COALESCE(workspace_id, '{_ZERO_UUID}'::uuid), + product_section, + model, + usage_date, + usage_kind, + context + ) + """ + ) + ) + + +def upgrade(db: Session): + if not _table_exists(db, "llm_usage_daily"): + print("llm_usage_daily missing; skipping 067") + db.commit() + return + if not _column_exists(db, "llm_usage_daily", "context"): + print("llm_usage_daily.context missing; run 065 first — skipping 067") + db.commit() + return + + legacy_removed = _dedupe_by_context_resource_keys(db) + if legacy_removed: + print(f"Merged {legacy_removed} legacy resource-scoped duplicate row(s)") + + context_removed = _dedupe_by_full_context(db) + if context_removed: + print(f"Merged {context_removed} duplicate full-context row(s)") + + _ensure_context_bucket_index(db) + print("Recreated uq_llm_usage_daily_bucket on full context JSONB") + db.commit() + + +def downgrade(db: Session): + db.commit() diff --git a/app/migrations/068_usage_committed_claims.py b/app/migrations/068_usage_committed_claims.py new file mode 100644 index 00000000..25c59fb6 --- /dev/null +++ b/app/migrations/068_usage_committed_claims.py @@ -0,0 +1,57 @@ +"""Migration: usage_committed_claims for durable flush idempotency.""" + +from sqlalchemy import text +from sqlalchemy.orm import Session + +description = ( + "Create usage_committed_claims when missing (063 may have run before this table was added)" +) + + +def _table_exists(db: Session, table: str) -> bool: + return ( + db.execute( + text( + """ + SELECT 1 FROM information_schema.tables + WHERE table_name = :table_name + """ + ), + {"table_name": table}, + ).first() + is not None + ) + + +def upgrade(db: Session): + if _table_exists(db, "usage_committed_claims"): + print("usage_committed_claims already exists, skipping 068") + db.commit() + return + + db.execute( + text( + """ + CREATE TABLE usage_committed_claims ( + claim_key TEXT PRIMARY KEY, + organization_id UUID NOT NULL, + committed_at TIMESTAMPTZ NOT NULL DEFAULT now() + ) + """ + ) + ) + db.execute( + text( + """ + CREATE INDEX ix_usage_committed_claims_committed_at + ON usage_committed_claims (committed_at) + """ + ) + ) + print("Created usage_committed_claims") + db.commit() + + +def downgrade(db: Session): + db.execute(text("DROP TABLE IF EXISTS usage_committed_claims")) + db.commit() diff --git a/app/migrations/069_usage_pricing.py b/app/migrations/069_usage_pricing.py new file mode 100644 index 00000000..0d8f8a76 --- /dev/null +++ b/app/migrations/069_usage_pricing.py @@ -0,0 +1,190 @@ +"""Migration: usage pricing catalog, org overrides, and cost columns on rollups.""" + +from __future__ import annotations + +from sqlalchemy import text +from sqlalchemy.orm import Session + +description = ( + "Add model_pricing_catalog, org_model_pricing_overrides, usage_pricing_mode, " + "and cost columns on llm_usage_daily" +) + + +def _table_exists(db: Session, table: str) -> bool: + return ( + db.execute( + text( + """ + SELECT 1 FROM information_schema.tables + WHERE table_name = :table_name + """ + ), + {"table_name": table}, + ).first() + is not None + ) + + +def _column_exists(db: Session, table: str, column: str) -> bool: + return ( + db.execute( + text( + """ + SELECT 1 FROM information_schema.columns + WHERE table_name = :table_name AND column_name = :column_name + """ + ), + {"table_name": table, "column_name": column}, + ).first() + is not None + ) + + +def _seed_pricing_catalog(db: Session) -> int: + from app.services.usage.pricing import DEFAULT_CATALOG_EFFECTIVE_FROM, seed_pricing_catalog + + return seed_pricing_catalog(db, effective_from=DEFAULT_CATALOG_EFFECTIVE_FROM) + + +def upgrade(db: Session): + if not _column_exists(db, "organizations", "usage_pricing_mode"): + db.execute( + text( + """ + ALTER TABLE organizations + ADD COLUMN usage_pricing_mode VARCHAR(32) NOT NULL DEFAULT 'platform_managed' + """ + ) + ) + print("Added organizations.usage_pricing_mode") + + if not _table_exists(db, "model_pricing_catalog"): + db.execute( + text( + """ + CREATE TABLE model_pricing_catalog ( + id UUID PRIMARY KEY DEFAULT gen_random_uuid(), + model VARCHAR(255) NOT NULL, + usage_kind VARCHAR(16) NOT NULL DEFAULT 'llm', + effective_from DATE NOT NULL DEFAULT CURRENT_DATE, + effective_to DATE, + input_micro_usd_per_million BIGINT NOT NULL DEFAULT 0, + output_micro_usd_per_million BIGINT NOT NULL DEFAULT 0, + cache_read_micro_usd_per_million BIGINT NOT NULL DEFAULT 0, + cache_creation_micro_usd_per_million BIGINT NOT NULL DEFAULT 0, + reasoning_micro_usd_per_million BIGINT NOT NULL DEFAULT 0, + audio_micro_usd_per_second BIGINT NOT NULL DEFAULT 0, + tts_micro_usd_per_million_chars BIGINT NOT NULL DEFAULT 0, + created_at TIMESTAMPTZ NOT NULL DEFAULT now(), + updated_at TIMESTAMPTZ NOT NULL DEFAULT now(), + CONSTRAINT uq_model_pricing_catalog_model_kind_from + UNIQUE (model, usage_kind, effective_from) + ) + """ + ) + ) + db.execute( + text( + """ + CREATE INDEX ix_model_pricing_catalog_lookup + ON model_pricing_catalog (model, usage_kind, effective_from DESC) + """ + ) + ) + print("Created model_pricing_catalog") + + if not _table_exists(db, "org_model_pricing_overrides"): + db.execute( + text( + """ + CREATE TABLE org_model_pricing_overrides ( + id UUID PRIMARY KEY DEFAULT gen_random_uuid(), + organization_id UUID NOT NULL + REFERENCES organizations(id) ON DELETE CASCADE, + model VARCHAR(255) NOT NULL, + usage_kind VARCHAR(16) NOT NULL DEFAULT 'llm', + effective_from DATE NOT NULL DEFAULT CURRENT_DATE, + effective_to DATE, + input_micro_usd_per_million BIGINT, + output_micro_usd_per_million BIGINT, + cache_read_micro_usd_per_million BIGINT, + cache_creation_micro_usd_per_million BIGINT, + reasoning_micro_usd_per_million BIGINT, + audio_micro_usd_per_second BIGINT, + tts_micro_usd_per_million_chars BIGINT, + created_at TIMESTAMPTZ NOT NULL DEFAULT now(), + updated_at TIMESTAMPTZ NOT NULL DEFAULT now(), + CONSTRAINT uq_org_model_pricing_override + UNIQUE (organization_id, model, usage_kind, effective_from) + ) + """ + ) + ) + db.execute( + text( + """ + CREATE INDEX ix_org_model_pricing_overrides_lookup + ON org_model_pricing_overrides ( + organization_id, model, usage_kind, effective_from DESC + ) + """ + ) + ) + print("Created org_model_pricing_overrides") + + cost_columns = [ + ("input_cost_micro_usd", "BIGINT NOT NULL DEFAULT 0"), + ("output_cost_micro_usd", "BIGINT NOT NULL DEFAULT 0"), + ("cache_read_cost_micro_usd", "BIGINT NOT NULL DEFAULT 0"), + ("cache_creation_cost_micro_usd", "BIGINT NOT NULL DEFAULT 0"), + ("reasoning_cost_micro_usd", "BIGINT NOT NULL DEFAULT 0"), + ("audio_cost_micro_usd", "BIGINT NOT NULL DEFAULT 0"), + ("tts_cost_micro_usd", "BIGINT NOT NULL DEFAULT 0"), + ("total_cost_micro_usd", "BIGINT NOT NULL DEFAULT 0"), + ("pricing_rate_source", "VARCHAR(16)"), + ("pricing_rate_id", "UUID"), + ] + if _table_exists(db, "llm_usage_daily"): + for column, col_type in cost_columns: + if not _column_exists(db, "llm_usage_daily", column): + db.execute( + text( + f"ALTER TABLE llm_usage_daily ADD COLUMN {column} {col_type}" + ) + ) + print("Added cost columns to llm_usage_daily") + + db.commit() + + if _table_exists(db, "model_pricing_catalog"): + seeded = _seed_pricing_catalog(db) + if seeded: + print(f"Seeded {seeded} model pricing catalog row(s)") + db.commit() + + +def downgrade(db: Session): + if _table_exists(db, "llm_usage_daily"): + for column in ( + "input_cost_micro_usd", + "output_cost_micro_usd", + "cache_read_cost_micro_usd", + "cache_creation_cost_micro_usd", + "reasoning_cost_micro_usd", + "audio_cost_micro_usd", + "tts_cost_micro_usd", + "total_cost_micro_usd", + "pricing_rate_source", + "pricing_rate_id", + ): + if _column_exists(db, "llm_usage_daily", column): + db.execute(text(f"ALTER TABLE llm_usage_daily DROP COLUMN {column}")) + + db.execute(text("DROP TABLE IF EXISTS org_model_pricing_overrides")) + db.execute(text("DROP TABLE IF EXISTS model_pricing_catalog")) + + if _column_exists(db, "organizations", "usage_pricing_mode"): + db.execute(text("ALTER TABLE organizations DROP COLUMN usage_pricing_mode")) + + db.commit() diff --git a/app/migrations/070_usage_margin_multiplier.py b/app/migrations/070_usage_margin_multiplier.py new file mode 100644 index 00000000..9a6c3e03 --- /dev/null +++ b/app/migrations/070_usage_margin_multiplier.py @@ -0,0 +1,45 @@ +"""Migration: org-level usage margin multiplier for priced rollups.""" + +from __future__ import annotations + +from sqlalchemy import text +from sqlalchemy.orm import Session + +description = "Add organizations.usage_margin_multiplier for usage cost markup" + + +def _column_exists(db: Session, table: str, column: str) -> bool: + return ( + db.execute( + text( + """ + SELECT 1 FROM information_schema.columns + WHERE table_name = :table_name AND column_name = :column_name + """ + ), + {"table_name": table, "column_name": column}, + ).first() + is not None + ) + + +def upgrade(db: Session): + if not _column_exists(db, "organizations", "usage_margin_multiplier"): + db.execute( + text( + """ + ALTER TABLE organizations + ADD COLUMN usage_margin_multiplier DOUBLE PRECISION NOT NULL DEFAULT 1.0 + """ + ) + ) + print("Added organizations.usage_margin_multiplier") + db.commit() + + +def downgrade(db: Session): + if _column_exists(db, "organizations", "usage_margin_multiplier"): + db.execute( + text("ALTER TABLE organizations DROP COLUMN usage_margin_multiplier") + ) + db.commit() diff --git a/app/migrations/071_reseed_pricing_catalog.py b/app/migrations/071_reseed_pricing_catalog.py new file mode 100644 index 00000000..ad17a360 --- /dev/null +++ b/app/migrations/071_reseed_pricing_catalog.py @@ -0,0 +1,37 @@ +"""Migration: re-seed model_pricing_rates from models.json after effective_from fix.""" + +from __future__ import annotations + +from sqlalchemy import text +from sqlalchemy.orm import Session + +from app.services.usage.pricing import DEFAULT_RATES_EFFECTIVE_FROM, seed_pricing_rates + +description = "Re-seed model_pricing_rates from models.json pricing blocks" + + +def upgrade(db: Session): + table = "model_pricing_rates" + exists = db.execute( + text("SELECT to_regclass('public.model_pricing_rates')") + ).scalar() + if not exists: + table = "model_pricing_catalog" + db.execute( + text( + f""" + UPDATE {table} + SET effective_from = CAST(:effective_from AS date) + WHERE effective_from > CAST(:effective_from AS date) + """ + ), + {"effective_from": DEFAULT_RATES_EFFECTIVE_FROM.isoformat()}, + ) + seeded = seed_pricing_rates(db, effective_from=DEFAULT_RATES_EFFECTIVE_FROM) + db.commit() + if seeded: + print(f"Re-seeded {seeded} model pricing rate row(s)") + + +def downgrade(db: Session): + pass diff --git a/app/migrations/072_usage_pricing_phase1.py b/app/migrations/072_usage_pricing_phase1.py new file mode 100644 index 00000000..a4a64ad9 --- /dev/null +++ b/app/migrations/072_usage_pricing_phase1.py @@ -0,0 +1,224 @@ +"""Migration: Phase 1 plan alignment — model_pricing_rates, buffer costs, strip extras.""" + +from __future__ import annotations + +from sqlalchemy import text +from sqlalchemy.orm import Session + +description = ( + "Rename model_pricing_catalog to model_pricing_rates, add currency/source, " + "cost columns on usage_pending_buffer, drop margin/BYOK org columns" +) + + +def _table_exists(db: Session, table: str) -> bool: + return ( + db.execute( + text( + """ + SELECT 1 FROM information_schema.tables + WHERE table_name = :table_name + """ + ), + {"table_name": table}, + ).first() + is not None + ) + + +def _column_exists(db: Session, table: str, column: str) -> bool: + return ( + db.execute( + text( + """ + SELECT 1 FROM information_schema.columns + WHERE table_name = :table_name AND column_name = :column_name + """ + ), + {"table_name": table, "column_name": column}, + ).first() + is not None + ) + + +def _ensure_model_pricing_rates(db: Session) -> None: + if _table_exists(db, "model_pricing_catalog") and not _table_exists( + db, "model_pricing_rates" + ): + db.execute( + text("ALTER TABLE model_pricing_catalog RENAME TO model_pricing_rates") + ) + db.execute( + text( + """ + ALTER INDEX IF EXISTS uq_model_pricing_catalog_model_kind_from + RENAME TO uq_model_pricing_rates_model_kind_from + """ + ) + ) + db.execute( + text( + """ + ALTER INDEX IF EXISTS ix_model_pricing_catalog_lookup + RENAME TO ix_model_pricing_rates_lookup + """ + ) + ) + print("Renamed model_pricing_catalog -> model_pricing_rates") + + if not _table_exists(db, "model_pricing_rates"): + db.execute( + text( + """ + CREATE TABLE model_pricing_rates ( + id UUID PRIMARY KEY DEFAULT gen_random_uuid(), + model VARCHAR(255) NOT NULL, + usage_kind VARCHAR(16) NOT NULL DEFAULT 'llm', + effective_from DATE NOT NULL DEFAULT CURRENT_DATE, + effective_to DATE, + currency VARCHAR(8) NOT NULL DEFAULT 'USD', + source VARCHAR(255) NOT NULL DEFAULT 'catalog', + input_micro_usd_per_million BIGINT NOT NULL DEFAULT 0, + output_micro_usd_per_million BIGINT NOT NULL DEFAULT 0, + cache_read_micro_usd_per_million BIGINT NOT NULL DEFAULT 0, + cache_creation_micro_usd_per_million BIGINT NOT NULL DEFAULT 0, + reasoning_micro_usd_per_million BIGINT NOT NULL DEFAULT 0, + audio_micro_usd_per_second BIGINT NOT NULL DEFAULT 0, + tts_micro_usd_per_million_chars BIGINT NOT NULL DEFAULT 0, + created_at TIMESTAMPTZ NOT NULL DEFAULT now(), + updated_at TIMESTAMPTZ NOT NULL DEFAULT now(), + CONSTRAINT uq_model_pricing_rates_model_kind_from + UNIQUE (model, usage_kind, effective_from) + ) + """ + ) + ) + db.execute( + text( + """ + CREATE INDEX ix_model_pricing_rates_lookup + ON model_pricing_rates (model, usage_kind, effective_from DESC) + """ + ) + ) + print("Created model_pricing_rates") + + if not _column_exists(db, "model_pricing_rates", "currency"): + db.execute( + text( + """ + ALTER TABLE model_pricing_rates + ADD COLUMN currency VARCHAR(8) NOT NULL DEFAULT 'USD' + """ + ) + ) + if not _column_exists(db, "model_pricing_rates", "source"): + db.execute( + text( + """ + ALTER TABLE model_pricing_rates + ADD COLUMN source VARCHAR(255) NOT NULL DEFAULT 'catalog' + """ + ) + ) + + +def _widen_source_column(db: Session) -> None: + if not _table_exists(db, "model_pricing_rates"): + return + if not _column_exists(db, "model_pricing_rates", "source"): + return + db.execute( + text( + """ + ALTER TABLE model_pricing_rates + ALTER COLUMN source TYPE VARCHAR(255) + """ + ) + ) + print("Widened model_pricing_rates.source to VARCHAR(255)") + + +def _add_buffer_cost_columns(db: Session) -> None: + if not _table_exists(db, "usage_pending_buffer"): + return + cost_columns = [ + ("input_cost_micro_usd", "BIGINT NOT NULL DEFAULT 0"), + ("output_cost_micro_usd", "BIGINT NOT NULL DEFAULT 0"), + ("cache_read_cost_micro_usd", "BIGINT NOT NULL DEFAULT 0"), + ("cache_creation_cost_micro_usd", "BIGINT NOT NULL DEFAULT 0"), + ("reasoning_cost_micro_usd", "BIGINT NOT NULL DEFAULT 0"), + ("audio_cost_micro_usd", "BIGINT NOT NULL DEFAULT 0"), + ("tts_cost_micro_usd", "BIGINT NOT NULL DEFAULT 0"), + ("total_cost_micro_usd", "BIGINT NOT NULL DEFAULT 0"), + ("pricing_rate_source", "VARCHAR(16)"), + ("pricing_rate_id", "UUID"), + ] + for column, col_type in cost_columns: + if not _column_exists(db, "usage_pending_buffer", column): + db.execute( + text( + f"ALTER TABLE usage_pending_buffer ADD COLUMN {column} {col_type}" + ) + ) + print("Ensured cost columns on usage_pending_buffer") + + +def _strip_org_extras(db: Session) -> None: + if _column_exists(db, "organizations", "usage_margin_multiplier"): + db.execute( + text("ALTER TABLE organizations DROP COLUMN usage_margin_multiplier") + ) + print("Dropped organizations.usage_margin_multiplier") + if _column_exists(db, "organizations", "usage_pricing_mode"): + db.execute(text("ALTER TABLE organizations DROP COLUMN usage_pricing_mode")) + print("Dropped organizations.usage_pricing_mode") + + +def upgrade(db: Session): + _ensure_model_pricing_rates(db) + _widen_source_column(db) + _add_buffer_cost_columns(db) + _strip_org_extras(db) + db.commit() + + from app.services.usage.pricing import ( + DEFAULT_RATES_EFFECTIVE_FROM, + clear_rates_table_cache, + seed_pricing_rates, + ) + + clear_rates_table_cache() + seeded = seed_pricing_rates(db, effective_from=DEFAULT_RATES_EFFECTIVE_FROM) + db.commit() + if seeded: + print(f"Seeded {seeded} model_pricing_rates row(s)") + + +def downgrade(db: Session): + if _table_exists(db, "usage_pending_buffer"): + for column in ( + "input_cost_micro_usd", + "output_cost_micro_usd", + "cache_read_cost_micro_usd", + "cache_creation_cost_micro_usd", + "reasoning_cost_micro_usd", + "audio_cost_micro_usd", + "tts_cost_micro_usd", + "total_cost_micro_usd", + "pricing_rate_source", + "pricing_rate_id", + ): + if _column_exists(db, "usage_pending_buffer", column): + db.execute( + text(f"ALTER TABLE usage_pending_buffer DROP COLUMN {column}") + ) + + if _table_exists(db, "model_pricing_rates") and not _table_exists( + db, "model_pricing_catalog" + ): + db.execute( + text("ALTER TABLE model_pricing_rates RENAME TO model_pricing_catalog") + ) + + db.commit() diff --git a/app/migrations/073_usage_cost_recompute_jobs.py b/app/migrations/073_usage_cost_recompute_jobs.py new file mode 100644 index 00000000..55f18192 --- /dev/null +++ b/app/migrations/073_usage_cost_recompute_jobs.py @@ -0,0 +1,72 @@ +"""Migration: async usage cost recompute job tracking.""" + +from __future__ import annotations + +from sqlalchemy import text +from sqlalchemy.orm import Session + +description = "Add usage_cost_recompute_jobs table for async cost backfill/recompute" + + +def _table_exists(db: Session, table: str) -> bool: + return ( + db.execute( + text( + """ + SELECT 1 FROM information_schema.tables + WHERE table_name = :table_name + """ + ), + {"table_name": table}, + ).first() + is not None + ) + + +def upgrade(db: Session) -> None: + if _table_exists(db, "usage_cost_recompute_jobs"): + print("usage_cost_recompute_jobs already exists, skipping...") + return + + db.execute( + text( + """ + CREATE TABLE usage_cost_recompute_jobs ( + id UUID PRIMARY KEY DEFAULT gen_random_uuid(), + organization_id UUID NOT NULL REFERENCES organizations(id) ON DELETE CASCADE, + status VARCHAR(32) NOT NULL DEFAULT 'pending', + model VARCHAR(255), + usage_kind VARCHAR(16), + start_date DATE, + end_date DATE, + updated_rows BIGINT NOT NULL DEFAULT 0, + error_message TEXT, + celery_task_id VARCHAR(255), + created_at TIMESTAMPTZ NOT NULL DEFAULT now(), + updated_at TIMESTAMPTZ NOT NULL DEFAULT now(), + completed_at TIMESTAMPTZ + ) + """ + ) + ) + db.execute( + text( + """ + CREATE INDEX IF NOT EXISTS ix_usage_cost_recompute_jobs_organization_id + ON usage_cost_recompute_jobs(organization_id) + """ + ) + ) + db.execute( + text( + """ + CREATE INDEX IF NOT EXISTS ix_usage_cost_recompute_jobs_org_status + ON usage_cost_recompute_jobs(organization_id, status) + """ + ) + ) + print("Created usage_cost_recompute_jobs table") + + +def downgrade(db: Session) -> None: + db.execute(text("DROP TABLE IF EXISTS usage_cost_recompute_jobs")) diff --git a/app/migrations/074_pricing_rates_source_and_effective_from.py b/app/migrations/074_pricing_rates_source_and_effective_from.py new file mode 100644 index 00000000..6f78f8b4 --- /dev/null +++ b/app/migrations/074_pricing_rates_source_and_effective_from.py @@ -0,0 +1,81 @@ +"""Migration: widen pricing source column and normalize effective_from baseline.""" + +from __future__ import annotations + +from sqlalchemy import text +from sqlalchemy.orm import Session + +from app.services.usage.pricing import DEFAULT_RATES_EFFECTIVE_FROM, seed_pricing_rates + +description = ( + "Widen model_pricing_rates.source to VARCHAR(255) and normalize effective_from " + "to 2020-01-01 baseline" +) + +_BASELINE = DEFAULT_RATES_EFFECTIVE_FROM.isoformat() + + +def _table_exists(db: Session, table: str) -> bool: + return ( + db.execute( + text("SELECT to_regclass(:table_name)"), + {"table_name": f"public.{table}"}, + ).scalar() + is not None + ) + + +def upgrade(db: Session) -> None: + table = "model_pricing_rates" + if not _table_exists(db, table): + table = "model_pricing_catalog" + if not _table_exists(db, table): + print("No pricing rates table found, skipping") + return + + db.execute( + text( + f""" + ALTER TABLE {table} + ALTER COLUMN source TYPE VARCHAR(255) + """ + ) + ) + print(f"Widened {table}.source to VARCHAR(255)") + + db.execute( + text( + f""" + DELETE FROM {table} newer + USING {table} baseline + WHERE newer.model = baseline.model + AND newer.usage_kind = baseline.usage_kind + AND newer.effective_from > CAST(:baseline AS date) + AND baseline.effective_from = CAST(:baseline AS date) + """ + ), + {"baseline": _BASELINE}, + ) + db.execute( + text( + f""" + UPDATE {table} + SET effective_from = CAST(:baseline AS date) + WHERE effective_from > CAST(:baseline AS date) + """ + ), + {"baseline": _BASELINE}, + ) + print(f"Normalized {table} effective_from to {_BASELINE}") + + seeded = seed_pricing_rates(db, effective_from=DEFAULT_RATES_EFFECTIVE_FROM) + from app.services.usage.pricing_cache import invalidate_all_pricing_cache + + invalidate_all_pricing_cache() + db.commit() + if seeded: + print(f"Re-seeded {seeded} pricing rate row(s) at {_BASELINE}") + + +def downgrade(db: Session) -> None: + pass diff --git a/app/migrations/075_cron_job_types_and_system_jobs.py b/app/migrations/075_cron_job_types_and_system_jobs.py new file mode 100644 index 00000000..abe8c572 --- /dev/null +++ b/app/migrations/075_cron_job_types_and_system_jobs.py @@ -0,0 +1,159 @@ +"""Migration: cron job types, system scheduled jobs, nullable org for platform tasks.""" + +from __future__ import annotations + +import json +import os +import uuid +from datetime import datetime, timezone + +from croniter import croniter +import pytz +from sqlalchemy import text +from sqlalchemy.orm import Session + +description = "Add cron job_type/is_system/config and seed platform system jobs" + +_SYSTEM_JOB_IDS = { + "usage_flush": uuid.UUID("00000000-0000-0000-0000-000000000001"), + "alert_evaluate": uuid.UUID("00000000-0000-0000-0000-000000000002"), + "oss_usage_prune": uuid.UUID("00000000-0000-0000-0000-000000000003"), + "fx_rate_refresh": uuid.UUID("00000000-0000-0000-0000-000000000004"), +} + + +def _column_exists(db: Session, table: str, column: str) -> bool: + return ( + db.execute( + text( + """ + SELECT 1 FROM information_schema.columns + WHERE table_name = :table_name AND column_name = :column_name + """ + ), + {"table_name": table, "column_name": column}, + ).first() + is not None + ) + + +def _next_run(cron_expression: str, tz_name: str) -> datetime: + tz = pytz.timezone(tz_name) + now = datetime.now(tz) + cron = croniter(cron_expression, now) + return cron.get_next(datetime).astimezone(timezone.utc) + + +def upgrade(db: Session) -> None: + if not _column_exists(db, "cron_jobs", "job_type"): + db.execute( + text( + """ + ALTER TABLE cron_jobs + ADD COLUMN job_type VARCHAR(64) NOT NULL DEFAULT 'evaluator_run' + """ + ) + ) + if not _column_exists(db, "cron_jobs", "is_system"): + db.execute( + text( + """ + ALTER TABLE cron_jobs + ADD COLUMN is_system BOOLEAN NOT NULL DEFAULT false + """ + ) + ) + if not _column_exists(db, "cron_jobs", "config"): + db.execute( + text( + """ + ALTER TABLE cron_jobs + ADD COLUMN config JSONB NOT NULL DEFAULT '{}'::jsonb + """ + ) + ) + + db.execute(text("ALTER TABLE cron_jobs ALTER COLUMN organization_id DROP NOT NULL")) + + flush_cron = os.environ.get("USAGE_FLUSH_BEAT_SECONDS", "120") + if flush_cron.isdigit(): + flush_expr = f"*/{max(1, int(flush_cron) // 60)} * * * *" + else: + flush_expr = "*/2 * * * *" + + system_jobs = [ + ( + _SYSTEM_JOB_IDS["usage_flush"], + "__system_usage_flush", + flush_expr, + "usage_flush", + ), + ( + _SYSTEM_JOB_IDS["alert_evaluate"], + "__system_alert_evaluate", + "*/5 * * * *", + "alert_evaluate", + ), + ( + _SYSTEM_JOB_IDS["oss_usage_prune"], + "__system_oss_usage_prune", + "0 3 * * *", + "oss_usage_prune", + ), + ( + _SYSTEM_JOB_IDS["fx_rate_refresh"], + "__system_fx_rate_refresh", + "0 6 * * *", + "fx_rate_refresh", + ), + ] + + for job_id, name, cron_expression, job_type in system_jobs: + exists = db.execute( + text("SELECT 1 FROM cron_jobs WHERE id = CAST(:id AS uuid)"), + {"id": str(job_id)}, + ).first() + if exists: + continue + next_run = _next_run(cron_expression, "UTC") + db.execute( + text( + """ + INSERT INTO cron_jobs ( + id, organization_id, name, cron_expression, timezone, + max_runs, current_runs, evaluator_ids, status, + next_run_at, job_type, is_system, config + ) VALUES ( + CAST(:id AS uuid), NULL, :name, :cron_expression, 'UTC', + 2147483647, 0, CAST(:evaluator_ids AS jsonb), 'active', + :next_run_at, :job_type, true, CAST('{}' AS jsonb) + ) + """ + ), + { + "id": str(job_id), + "name": name, + "cron_expression": cron_expression, + "evaluator_ids": json.dumps([]), + "next_run_at": next_run, + "job_type": job_type, + }, + ) + print(f"Seeded system cron job {job_type}") + + db.commit() + + +def downgrade(db: Session) -> None: + for job_id in _SYSTEM_JOB_IDS.values(): + db.execute( + text("DELETE FROM cron_jobs WHERE id = CAST(:id AS uuid)"), + {"id": str(job_id)}, + ) + if _column_exists(db, "cron_jobs", "config"): + db.execute(text("ALTER TABLE cron_jobs DROP COLUMN config")) + if _column_exists(db, "cron_jobs", "is_system"): + db.execute(text("ALTER TABLE cron_jobs DROP COLUMN is_system")) + if _column_exists(db, "cron_jobs", "job_type"): + db.execute(text("ALTER TABLE cron_jobs DROP COLUMN job_type")) + db.commit() diff --git a/app/migrations/076_ai_provider_enabled_models.py b/app/migrations/076_ai_provider_enabled_models.py new file mode 100644 index 00000000..c1714f28 --- /dev/null +++ b/app/migrations/076_ai_provider_enabled_models.py @@ -0,0 +1,38 @@ +"""Migration: per-credential enabled model allowlist for integrations.""" + +from sqlalchemy import text +from sqlalchemy.orm import Session + +description = "Add enabled_models JSONB to aiproviders" + + +def _column_exists(db: Session, table: str, column: str) -> bool: + return ( + db.execute( + text( + """ + SELECT 1 FROM information_schema.columns + WHERE table_name = :table_name AND column_name = :column_name + """ + ), + {"table_name": table, "column_name": column}, + ).first() + is not None + ) + + +def upgrade(db: Session) -> None: + if not _column_exists(db, "aiproviders", "enabled_models"): + db.execute( + text( + """ + ALTER TABLE aiproviders + ADD COLUMN enabled_models JSONB NULL + """ + ) + ) + + +def downgrade(db: Session) -> None: + if _column_exists(db, "aiproviders", "enabled_models"): + db.execute(text("ALTER TABLE aiproviders DROP COLUMN enabled_models")) diff --git a/app/migrations/077_remove_system_cron_jobs.py b/app/migrations/077_remove_system_cron_jobs.py new file mode 100644 index 00000000..5aadeccb --- /dev/null +++ b/app/migrations/077_remove_system_cron_jobs.py @@ -0,0 +1,19 @@ +"""Remove platform system cron rows; platform jobs use Celery Beat.""" + +from sqlalchemy import text +from sqlalchemy.orm import Session + +description = "Delete is_system cron_jobs; platform scheduling moved to Celery Beat" + + +def upgrade(db: Session) -> None: + result = db.execute( + text("DELETE FROM cron_jobs WHERE is_system = true") + ) + db.commit() + print(f"Removed {result.rowcount} system cron job row(s)") + + +def downgrade(db: Session) -> None: + # System jobs are re-seeded by re-running migration 075 logic if needed. + print("Downgrade: re-run 075 upgrade to restore system cron jobs if required") diff --git a/app/models/database.py b/app/models/database.py index 2db26e68..92350139 100644 --- a/app/models/database.py +++ b/app/models/database.py @@ -19,7 +19,7 @@ select, text, ) -from sqlalchemy.dialects.postgresql import UUID +from sqlalchemy.dialects.postgresql import JSONB, UUID from sqlalchemy.orm import relationship from sqlalchemy.sql import func import uuid @@ -702,6 +702,8 @@ class AIProvider(Base): gateway_auth_secret = Column(String, nullable=True) # Arbitrary HTTP headers sent with gateway-routed LiteLLM calls gateway_extra_headers = Column(JSON, nullable=True) + # Non-empty list restricts model pickers; null/empty = all catalog models for provider. + enabled_models = 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()) last_tested_at = Column(DateTime(timezone=True), nullable=True) # When API key was last validated @@ -1262,10 +1264,13 @@ class CronJob(Base): __tablename__ = "cron_jobs" id = Column(UUID(as_uuid=True), primary_key=True, default=uuid.uuid4) - organization_id = Column(UUID(as_uuid=True), ForeignKey("organizations.id"), nullable=False, index=True) + organization_id = Column(UUID(as_uuid=True), ForeignKey("organizations.id"), nullable=True, index=True) # Basic information name = Column(String(255), nullable=False) + job_type = Column(String(64), nullable=False, default="evaluator_run") + is_system = Column(Boolean, nullable=False, default=False) + config = Column(JSON, nullable=False, default=dict) cron_expression = Column(String(100), nullable=False) # e.g., "0 9 * * 1-5" timezone = Column(String(100), nullable=False, default="UTC") @@ -2935,3 +2940,83 @@ class JudgeRun(Base): created_by = Column(String, nullable=True) dataset = relationship("JudgeDataset", back_populates="runs") + + +class UsageCostRecomputeJob(Base): + """Async job tracking for retroactive usage cost recompute.""" + + __tablename__ = "usage_cost_recompute_jobs" + + id = Column(UUID(as_uuid=True), primary_key=True, default=uuid.uuid4) + organization_id = Column( + UUID(as_uuid=True), + ForeignKey("organizations.id", ondelete="CASCADE"), + nullable=False, + index=True, + ) + status = Column(String(32), nullable=False, default="pending", server_default="pending") + model = Column(String(255), nullable=True) + usage_kind = Column(String(16), nullable=True) + start_date = Column(Date, nullable=True) + end_date = Column(Date, nullable=True) + updated_rows = Column(BigInteger, nullable=False, default=0, server_default="0") + error_message = Column(String, nullable=True) + celery_task_id = Column(String(255), nullable=True, index=True) + created_at = Column(DateTime(timezone=True), server_default=func.now()) + updated_at = Column(DateTime(timezone=True), server_default=func.now(), onupdate=func.now()) + completed_at = Column(DateTime(timezone=True), nullable=True) + + +class LLMUsageDaily(Base): + """Daily LLM/STT usage rollups for org-scoped Usage reporting.""" + + __tablename__ = "llm_usage_daily" + + id = Column(UUID(as_uuid=True), primary_key=True, default=uuid.uuid4) + organization_id = Column( + UUID(as_uuid=True), + ForeignKey("organizations.id", ondelete="CASCADE"), + nullable=False, + index=True, + ) + workspace_id = Column( + UUID(as_uuid=True), + ForeignKey("workspaces.id", ondelete="SET NULL"), + nullable=True, + index=True, + ) + product_section = Column(String(64), nullable=False, index=True) + model = Column(String(255), nullable=False, index=True) + context = Column(JSONB, nullable=False, server_default="{}", default=dict) + usage_date = Column(Date, nullable=False, index=True) + usage_kind = Column(String(16), nullable=False, default="llm", server_default="llm") + prompt_tokens = Column(BigInteger, nullable=False, default=0, server_default="0") + completion_tokens = Column(BigInteger, nullable=False, default=0, server_default="0") + cache_read_tokens = Column(BigInteger, nullable=False, default=0, server_default="0") + cache_creation_tokens = Column( + BigInteger, nullable=False, default=0, server_default="0" + ) + reasoning_tokens = Column(BigInteger, nullable=False, default=0, server_default="0") + audio_seconds = Column(BigInteger, nullable=False, default=0, server_default="0") + tts_characters = Column(BigInteger, nullable=False, default=0, server_default="0") + call_count = Column(BigInteger, nullable=False, default=0, server_default="0") + input_cost_micro_usd = Column(BigInteger, nullable=False, default=0, server_default="0") + output_cost_micro_usd = Column(BigInteger, nullable=False, default=0, server_default="0") + cache_read_cost_micro_usd = Column( + BigInteger, nullable=False, default=0, server_default="0" + ) + cache_creation_cost_micro_usd = Column( + BigInteger, nullable=False, default=0, server_default="0" + ) + reasoning_cost_micro_usd = Column( + BigInteger, nullable=False, default=0, server_default="0" + ) + audio_cost_micro_usd = Column(BigInteger, nullable=False, default=0, server_default="0") + tts_cost_micro_usd = Column(BigInteger, nullable=False, default=0, server_default="0") + total_cost_micro_usd = Column(BigInteger, nullable=False, default=0, server_default="0") + pricing_rate_source = Column(String(16), nullable=True) + pricing_rate_id = Column(UUID(as_uuid=True), 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 25208f30..25241e77 100644 --- a/app/models/schemas.py +++ b/app/models/schemas.py @@ -970,6 +970,13 @@ class AIProviderCreate(BaseModel): None, description="Arbitrary HTTP headers sent with gateway-routed LiteLLM calls.", ) + enabled_models: Optional[List[str]] = Field( + None, + description=( + "Allowlisted model names for this credential. " + "Null or empty means all catalog models for the provider." + ), + ) is_default: Optional[bool] = Field( None, description=( @@ -1053,6 +1060,21 @@ def validate_gateway_auth_secret(cls, v: Optional[str]) -> Optional[str]: def validate_gateway_extra_headers(cls, v: Optional[Dict[str, Any]]) -> Optional[Dict[str, str]]: return _validate_gateway_extra_headers(v) + @field_validator("enabled_models") + @classmethod + def validate_enabled_models_create(cls, v: Optional[List[str]]) -> Optional[List[str]]: + if v is None: + return None + seen: set[str] = set() + out: list[str] = [] + for item in v: + name = str(item).strip() + if not name or name in seen: + continue + seen.add(name) + out.append(name) + return out or None + class AIProviderUpdate(BaseModel): """Schema for updating an AI Provider.""" @@ -1069,6 +1091,7 @@ class AIProviderUpdate(BaseModel): gateway_auth_secret: Optional[str] = None clear_gateway_auth_secret: bool = False gateway_extra_headers: Optional[Dict[str, str]] = None + enabled_models: Optional[List[str]] = None @field_validator("gateway_model") @classmethod @@ -1127,6 +1150,21 @@ def validate_gateway_extra_headers_update( ) -> Optional[Dict[str, str]]: return _validate_gateway_extra_headers(v) + @field_validator("enabled_models") + @classmethod + def validate_enabled_models_update(cls, v: Optional[List[str]]) -> Optional[List[str]]: + if v is None: + return None + seen: set[str] = set() + out: list[str] = [] + for item in v: + name = str(item).strip() + if not name or name in seen: + continue + seen.add(name) + out.append(name) + return out or None + class AIProviderResponse(BaseModel): """Schema for AI Provider response.""" @@ -1145,6 +1183,7 @@ class AIProviderResponse(BaseModel): gateway_auth_secret_env: Optional[str] = None has_gateway_auth_secret: bool = False gateway_extra_headers: Optional[Dict[str, str]] = None + enabled_models: Optional[List[str]] = None gateway_managed: bool = False effective_routing: Literal["inherit", "direct", "gateway", "bifrost", "litellm_proxy"] = "inherit" effective_gateway_interface: Literal["litellm_shim", "native_openai"] = "litellm_shim" @@ -2591,6 +2630,8 @@ class CronJobResponse(BaseModel): id: UUID organization_id: UUID name: str + job_type: str = "evaluator_run" + is_system: bool = False cron_expression: str timezone: str max_runs: int diff --git a/app/services/ai/llm_gateway.py b/app/services/ai/llm_gateway.py index fc2b3a71..2283a3a5 100644 --- a/app/services/ai/llm_gateway.py +++ b/app/services/ai/llm_gateway.py @@ -10,6 +10,7 @@ from __future__ import annotations import os +from contextlib import contextmanager from dataclasses import dataclass from typing import Any, Dict, Literal, Optional, Tuple from urllib.parse import urlparse @@ -577,10 +578,15 @@ def get_credential_effective_routing_label( db: Session, credential: Any, ) -> EffectiveRouting: - """Resolved routing label for API responses.""" + """Resolved routing label for API responses (never raises on misconfigured gateway).""" ctx = _credential_routing_context(credential) - _, effective = resolve_effective_routing(organization_id, db, ctx) - return effective + try: + _, effective = resolve_effective_routing(organization_id, db, ctx) + return effective + except RuntimeError: + if ctx.routing_mode == "gateway": + return "gateway" + raise # Gateways speak OpenAI-compatible ``/v1/chat/completions``. LiteLLM @@ -716,4 +722,92 @@ def litellm_completion( db=db, credential=credential, ) - return litellm.completion(**kwargs) + response = litellm.completion(**kwargs) + try: + from app.services.usage.context import ( + LLMUsageProductSection, + ensure_usage_context, + reset_usage_context, + ) + from app.services.usage.llm_usage import record_llm_usage + from app.services.usage.normalize import ( + normalize_llm_usage, + usage_snapshot_is_billable, + ) + + usage_token = ensure_usage_context( + organization_id, + product_section=LLMUsageProductSection.OTHER, + ) + try: + model_name = str(kwargs.get("model") or "unknown") + if "/" in model_name: + model_name = model_name.rsplit("/", 1)[-1] + snapshot = normalize_llm_usage(raw_response=response) + if usage_snapshot_is_billable(snapshot): + record_llm_usage( + model_name, + snapshot, + organization_id=organization_id, + ) + finally: + if usage_token is not None: + reset_usage_context(usage_token) + except Exception as exc: + logger.debug("litellm_completion usage record skipped: {}", exc) + return response + + +@contextmanager +def litellm_batch_completion_recording( + *, + organization_id: UUID, + db: Session, + model: Optional[str] = None, + credential: Optional[CredentialRoutingContext] = None, +): + """Temporarily wrap litellm.batch_completion to record usage for each response.""" + import litellm + + from app.services.usage.llm_usage import record_llm_usage + from app.services.usage.normalize import ( + normalize_llm_usage, + usage_snapshot_is_billable, + ) + + original = litellm.batch_completion + + def _recording_batch(**kwargs: Any): + kwargs = apply_llm_gateway( + kwargs, + organization_id=organization_id, + db=db, + model=model or kwargs.get("model"), + credential=credential, + ) + responses = original(**kwargs) + model_name = str(kwargs.get("model") or model or "unknown") + if "/" in model_name: + model_name = model_name.rsplit("/", 1)[-1] + items = responses if isinstance(responses, list) else [responses] + for resp in items: + if resp is None: + continue + try: + snapshot = normalize_llm_usage(raw_response=resp) + if not usage_snapshot_is_billable(snapshot): + continue + record_llm_usage( + model_name, + snapshot, + organization_id=organization_id, + ) + except Exception as exc: + logger.debug("litellm batch usage record skipped: {}", exc) + return responses + + litellm.batch_completion = _recording_batch + try: + yield + finally: + litellm.batch_completion = original diff --git a/app/services/ai/llm_service.py b/app/services/ai/llm_service.py index fb22285d..95e64488 100644 --- a/app/services/ai/llm_service.py +++ b/app/services/ai/llm_service.py @@ -490,22 +490,6 @@ def generate_response( credential=credential_ctx, ) - # #region agent log - try: - import json as _json, time as _time - _msg_types = [] - for _m in messages: - _c = _m.get("content") - if isinstance(_c, list): - _msg_types.extend( - p.get("type") for p in _c if isinstance(p, dict) - ) - with open("debug-bfc313.log", "a", encoding="utf-8") as _f: - _f.write(_json.dumps({"sessionId": "bfc313", "runId": "post-fix", "hypothesisId": "C", "location": "llm_service.py:generate_response", "message": "pre-completion routing", "data": {"effective_routing": effective_routing, "credential_mode": getattr(credential_ctx, "routing_mode", None), "credential_id": str(credential_id) if credential_id else None, "model": model_str, "has_api_base": "api_base" in call_kwargs, "custom_llm_provider": call_kwargs.get("custom_llm_provider"), "content_part_types": _msg_types}, "timestamp": int(_time.time() * 1000)}) + "\n") - except Exception: - pass - # #endregion - try: response = litellm.completion(**call_kwargs) except Exception as e: @@ -549,6 +533,36 @@ def generate_response( "raw_response": response, "processing_time": time.time() - start_time, } + try: + from app.services.usage.context import ( + LLMUsageProductSection, + ensure_usage_context, + reset_usage_context, + ) + from app.services.usage.normalize import ( + normalize_llm_usage, + usage_snapshot_is_billable, + ) + from app.services.usage.llm_usage import record_llm_usage + + usage_token = ensure_usage_context( + organization_id, + product_section=LLMUsageProductSection.OTHER, + ) + try: + snapshot = normalize_llm_usage(raw_response=response) + result["usage"]["cache_read_tokens"] = snapshot.cache_read_tokens + result["usage"]["cache_creation_tokens"] = snapshot.cache_creation_tokens + result["usage"]["reasoning_tokens"] = snapshot.reasoning_tokens + if usage_snapshot_is_billable(snapshot): + record_llm_usage( + llm_model, snapshot, organization_id=organization_id + ) + finally: + if usage_token is not None: + reset_usage_context(usage_token) + except Exception as exc: + logger.debug("llm usage record skipped: {}", exc) return result diff --git a/app/services/ai/stt_clients/google.py b/app/services/ai/stt_clients/google.py index 6d86ec30..84caf40d 100644 --- a/app/services/ai/stt_clients/google.py +++ b/app/services/ai/stt_clients/google.py @@ -83,6 +83,28 @@ def _build_transcription_prompt(language: Optional[str]) -> str: ) +def _gemini_stt_usage_ctx(config_model: str): + """Merge Gemini multimodal STT tags into the active usage context.""" + from app.services.usage.context import LLMUsageContext, get_usage_context + + tags = { + "stt_backend": "gemini_multimodal", + "config_model": (config_model or "").strip(), + "usage_split": "llm_tokens_and_stt_seconds", + } + base = get_usage_context() + if base is None: + return None + return LLMUsageContext( + organization_id=base.organization_id, + workspace_id=base.workspace_id, + product_section=base.product_section, + resource_id=base.resource_id, + resource_type=base.resource_type, + extra={**(base.extra or {}), **tags}, + ) + + def transcribe_google( audio_file_path: str, model: str, @@ -152,6 +174,33 @@ def transcribe_google( f"Gemini transcription failed for {litellm_model}: {e}" ) + if organization_id is not None: + try: + from app.services.usage.llm_usage import record_llm_usage + from app.services.usage.normalize import normalize_llm_usage + + gemini_model = _strip_stt_suffix(model) + usage_ctx = _gemini_stt_usage_ctx(model) + record_llm_usage( + gemini_model, + normalize_llm_usage(raw_response=response), + organization_id=organization_id, + ctx=usage_ctx, + ) + from app.services.usage.llm_usage import probe_audio_seconds, record_stt_usage + + audio_seconds = probe_audio_seconds(audio_file_path) + if audio_seconds > 0: + record_stt_usage( + gemini_model, + audio_seconds=audio_seconds, + organization_id=organization_id, + ctx=usage_ctx, + count_call=False, + ) + except Exception as exc: + logger.debug("[transcribe_google] llm usage record skipped: %s", exc) + text = "" try: text = (response.choices[0].message.content or "").strip() diff --git a/app/services/ai/transcription_service.py b/app/services/ai/transcription_service.py index 955e1a7b..9620e1ca 100644 --- a/app/services/ai/transcription_service.py +++ b/app/services/ai/transcription_service.py @@ -1,735 +1,780 @@ -""" -Transcription service for converting audio to text using various STT providers. - -Response Format: -All providers return a standardized format: -{ - "text": str, # Full transcript text - "language": str, # Language code (e.g., "en", "es") - "segments": [ # List of segments with timestamps - { - "start": float, # Start time in seconds - "end": float, # End time in seconds - "text": str # Text for this segment - } - ], - "speaker_segments": [ # Segments with speaker labels (if diarization enabled) - { - "speaker": str, # Speaker label (e.g., "Speaker 1") - "text": str, - "start": float, - "end": float - } - ], - "processing_time": float, - "raw_output": dict # Original provider response -} - -Provider-Specific Formats: -- OpenAI Whisper API: Uses verbose_json format which includes segments with timestamps -- Local Whisper: Returns segments by default with word-level timestamps available -- Google/Azure/AWS: Formats documented but not yet implemented - -Speaker Diarization: -- Whisper does NOT provide speaker diarization natively (only transcription) -- We use pyannote.audio (speaker-diarization-3.1) for accurate ML-based diarization -- Whisper provides word-level timestamps, pyannote identifies speakers, then we align -- Requires: pyannote.audio installed + diarization.huggingface_token in config.yml -- Falls back to unreliable gap-based heuristics if pyannote is unavailable -""" - -import time -import tempfile -import os -import logging -from typing import Optional, Dict, Any, List -from uuid import UUID -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 -from sqlalchemy.orm import Session - -logger = logging.getLogger(__name__) - - -class TranscriptionService: - """Service for transcribing audio files using various STT providers.""" - - def __init__(self): - """Initialize transcription service.""" - self._pyannote_pipeline = None - - def _get_ai_provider( - self, - provider: ModelProvider, - db: Session, - organization_id: UUID, - credential_id: Optional[UUID] = None, - ) -> Optional[AIProvider]: - """Resolve the AIProvider row to use, honoring ``credential_id``.""" - return resolve_ai_provider( - provider, db, organization_id, credential_id=credential_id - ) - - def _get_api_key_for_provider( - self, - provider: ModelProvider, - db: Session, - organization_id: UUID, - credential_id: Optional[UUID] = None, - ) -> Optional[str]: - """Resolve and decrypt API key from AIProvider or Integration tables. - - Checks AIProvider first (LLM-style providers like OpenAI), then falls - back to the Integration table (voice platforms like Deepgram, - ElevenLabs). When ``credential_id`` is given the explicit row is - preferred in either table; otherwise the row marked ``is_default`` - wins, with a back-compat fallback to the most recent active row. - """ - from app.services.ai.llm_gateway import ( - resolve_litellm_api_key, - routing_context_from_ai_provider, - ) - - ai_provider = self._get_ai_provider( - provider, db, organization_id, credential_id=credential_id - ) - if ai_provider: - credential_ctx = routing_context_from_ai_provider(ai_provider) - return resolve_litellm_api_key( - organization_id, - db, - ai_provider, - credential=credential_ctx, - ) - - integration = resolve_integration( - provider, db, organization_id, credential_id=credential_id - ) - if integration: - return decrypt_api_key(integration.api_key) - - return None - - def _get_credential_context_for_provider( - self, - provider: ModelProvider, - db: Session, - organization_id: UUID, - credential_id: Optional[UUID] = None, - ): - from app.services.ai.llm_gateway import routing_context_from_ai_provider - - ai_provider = self._get_ai_provider( - provider, db, organization_id, credential_id=credential_id - ) - if ai_provider: - return routing_context_from_ai_provider(ai_provider) - return None - - def _download_audio_to_temp(self, audio_file_key: str, db: Optional[Session] = None) -> str: - """ - Download audio from S3 to temporary file, or use local file if S3 is not available. - """ - import os - - # First, try S3 if enabled - if s3_service.is_enabled(): - try: - # Download from S3 - audio_bytes = s3_service.download_file_by_key(audio_file_key) - - # Determine file extension from key - file_ext = Path(audio_file_key).suffix.lstrip(".") or "wav" - - # Create temporary file - with tempfile.NamedTemporaryFile(delete=False, suffix=f".{file_ext}") as temp_file: - temp_file.write(audio_bytes) - return temp_file.name - except Exception as e: - # If S3 download fails, fall through to local file check - logger.warning(f"S3 download failed for {audio_file_key}: {e}, trying local file...") - - # Fallback: Try to find local file - # Check if it's already a local file path - if os.path.exists(audio_file_key): - # It's a local file path, return it directly - return audio_file_key - - # Try to look up in database if db session is provided - if db: - try: - from app.models.database import AudioFile - # Try to find by S3 key or file path - audio_file = db.query(AudioFile).filter( - (AudioFile.file_path == audio_file_key) | (AudioFile.file_path.like(f"%{audio_file_key}%")) - ).first() - - if audio_file and os.path.exists(audio_file.file_path): - return audio_file.file_path - except Exception as e: - logger.warning(f"Database lookup failed for {audio_file_key}: {e}") - - # If all else fails, raise error - raise StorageError(f"Failed to download audio file: Cloud blob storage is not enabled and local file not found for key: {audio_file_key}") - - # Provider-specific transcription is delegated to app.services.ai.stt_clients - - @staticmethod - def _transcribe_with_whisper_local(audio_file_path: str, model_name: str = "base") -> Dict[str, Any]: - """Transcribe audio using local Whisper model (not a remote API).""" - try: - import whisper - - model = whisper.load_model(model_name) - result = model.transcribe(audio_file_path) - - return { - "text": result.get("text", ""), - "language": result.get("language", "en"), - "segments": [ - {"start": seg.get("start", 0), "end": seg.get("end", 0), "text": seg.get("text", "")} - for seg in result.get("segments", []) - ], - } - except ImportError: - raise RuntimeError("Whisper library not installed. Install with: pip install openai-whisper") - except Exception as e: - raise RuntimeError(f"Whisper transcription failed: {str(e)}") - - def _get_pyannote_pipeline(self): - """Load and cache the pyannote diarization pipeline.""" - if self._pyannote_pipeline is not None: - return self._pyannote_pipeline - - # Compatibility shim: list_audio_backends was removed in torchaudio 2.4+ - import torchaudio - if not hasattr(torchaudio, "list_audio_backends"): - torchaudio.list_audio_backends = lambda: ["ffmpeg"] - - from pyannote.audio import Pipeline - from app.config import settings - - hf_token = settings.HUGGINGFACE_TOKEN - if not hf_token: - raise RuntimeError( - "HUGGINGFACE_TOKEN not configured. Set it under 'diarization.huggingface_token' " - "in config.yml. Required for pyannote speaker diarization." - ) - - logger.info("Loading pyannote speaker-diarization-3.1 pipeline (first call, will be cached)...") - self._pyannote_pipeline = Pipeline.from_pretrained("pyannote/speaker-diarization-3.1", token=hf_token) - logger.info("Pyannote pipeline loaded successfully") - return self._pyannote_pipeline - - def _detect_speakers_with_pyannote( - self, - audio_file_path: str, - segments: List[Dict[str, Any]], - words: Optional[List[Dict[str, Any]]] = None, - num_speakers: Optional[int] = 2, - min_speakers: Optional[int] = None, - max_speakers: Optional[int] = None, - ) -> List[Dict[str, Any]]: - """ - Use pyannote.audio for ML-based speaker diarization, aligned with - Whisper word timestamps for accurate speaker-text mapping. - """ - pipeline = self._get_pyannote_pipeline() - - # Build pipeline kwargs for speaker count hints - pipeline_kwargs = {} - if num_speakers is not None: - pipeline_kwargs["num_speakers"] = num_speakers - if min_speakers is not None: - pipeline_kwargs["min_speakers"] = min_speakers - if max_speakers is not None: - pipeline_kwargs["max_speakers"] = max_speakers - - logger.info( - f"Running pyannote diarization on {audio_file_path} " f"(speaker hints: {pipeline_kwargs or 'auto'})" - ) - raw_output = pipeline(audio_file_path, **pipeline_kwargs) - - # Handle different return types across pyannote versions: - # - Older versions: Annotation directly (has itertracks) - # - Newer versions: DiarizeOutput with .speaker_diarization attribute - if hasattr(raw_output, "itertracks"): - annotation = raw_output - elif hasattr(raw_output, "speaker_diarization"): - annotation = raw_output.speaker_diarization - elif hasattr(raw_output, "annotation"): - annotation = raw_output.annotation - elif isinstance(raw_output, tuple): - annotation = raw_output[0] - else: - attrs = [a for a in dir(raw_output) if not a.startswith("_")] - raise TypeError( - f"Unexpected pyannote output type: {type(raw_output).__name__}. " f"Available attributes: {attrs}" - ) - - # Collect diarization turns as a sorted list for fast lookup - diar_turns = [] - raw_labels = set() - for turn, _, speaker_label in annotation.itertracks(yield_label=True): - diar_turns.append((turn.start, turn.end, speaker_label)) - raw_labels.add(speaker_label) - - if not diar_turns: - logger.warning("Pyannote returned no speaker turns") - return [] - - sorted_labels = sorted(raw_labels) - label_map = {lbl: f"Speaker {i + 1}" for i, lbl in enumerate(sorted_labels)} - logger.info(f"Pyannote detected {len(sorted_labels)} speakers, {len(diar_turns)} turns") - - def find_speaker(midpoint: float) -> str: - """Find which speaker is active at a given timestamp.""" - for t_start, t_end, lbl in diar_turns: - if t_start <= midpoint <= t_end: - return label_map[lbl] - # No exact match -- find the closest turn - min_dist = float("inf") - closest_label = label_map[diar_turns[0][2]] - for t_start, t_end, lbl in diar_turns: - dist = min(abs(midpoint - t_start), abs(midpoint - t_end)) - if dist < min_dist: - min_dist = dist - closest_label = label_map[lbl] - return closest_label - - if words and len(words) > 0: - speaker_segments = [] - current_speaker = None - current_words: List[str] = [] - current_start = 0.0 - current_end = 0.0 - - for w in words: - word_text = w.get("word", "").strip() - w_start = w.get("start", 0) - w_end = w.get("end", 0) - if not word_text: - continue - - midpoint = (w_start + w_end) / 2.0 - speaker = find_speaker(midpoint) - - if speaker != current_speaker: - if current_words and current_speaker: - speaker_segments.append( - { - "speaker": current_speaker, - "text": " ".join(current_words).strip(), - "start": round(current_start, 3), - "end": round(current_end, 3), - } - ) - current_speaker = speaker - current_words = [word_text] - current_start = w_start - current_end = w_end - else: - current_words.append(word_text) - current_end = w_end - - if current_words and current_speaker: - speaker_segments.append( - { - "speaker": current_speaker, - "text": " ".join(current_words).strip(), - "start": round(current_start, 3), - "end": round(current_end, 3), - } - ) - - return speaker_segments - - # Fallback: align at segment level when word timestamps are not available - speaker_segments = [] - for seg in segments: - seg_start = seg.get("start", 0) - seg_end = seg.get("end", 0) - seg_text = seg.get("text", "").strip() - if not seg_text: - continue - - midpoint = (seg_start + seg_end) / 2.0 - speaker = find_speaker(midpoint) - - speaker_segments.append( - { - "speaker": speaker, - "text": seg_text, - "start": round(seg_start, 3), - "end": round(seg_end, 3), - } - ) - - return speaker_segments - - def _detect_speakers_heuristic(self, segments: List[Dict[str, Any]]) -> List[Dict[str, Any]]: - """ - Heuristic-based speaker diarization fallback. Uses gap-based detection - which is unreliable -- pyannote.audio should be used for accurate results. - """ - logger.warning( - "Using heuristic speaker diarization (unreliable). For accurate results, " - "install pyannote.audio and configure diarization.huggingface_token in config.yml" - ) - if not segments: - return [] - - speaker_segments = [] - current_speaker = "Speaker 1" - - # Thresholds for speaker change detection - GAP_THRESHOLD_LARGE = 0.8 # seconds - large gap definitely suggests change - GAP_THRESHOLD_MEDIUM = 0.4 # seconds - medium gap suggests change - GAP_THRESHOLD_SMALL = 0.2 # seconds - small gap, but combined with other factors - MIN_SEGMENT_DURATION = 0.2 # seconds - minimum segment to consider - - # Calculate average gap size to adapt thresholds - gaps = [] - for i in range(1, len(segments)): - gap = segments[i].get("start", 0) - segments[i - 1].get("end", 0) - if gap > 0: - gaps.append(gap) - - avg_gap = sum(gaps) / len(gaps) if gaps else 0.5 - # Adaptive threshold based on conversation pace - adaptive_threshold = min(max(avg_gap * 1.5, 0.3), 1.0) - - # Track recent speaker assignments for pattern detection - recent_assignments = [] - assignment_window = 5 # Look at last N segments for patterns - - for i, seg in enumerate(segments): - seg_start = seg.get("start", 0) - seg_end = seg.get("end", 0) - seg_duration = seg_end - seg_start - seg_text = seg.get("text", "").strip() - - # Skip very short segments (likely noise or artifacts) - if seg_duration < MIN_SEGMENT_DURATION or not seg_text: - continue - - # Check for speaker change indicators - should_switch = False - gap = 0 - - if i > 0: - prev_seg = segments[i - 1] - gap = seg_start - prev_seg.get("end", 0) - - # Large gap definitely suggests speaker change - if gap > GAP_THRESHOLD_LARGE: - should_switch = True - # Medium gap suggests change - elif gap > GAP_THRESHOLD_MEDIUM: - should_switch = True - # Small gap but check for alternating pattern - elif gap > GAP_THRESHOLD_SMALL: - # If we've been alternating, continue the pattern - if len(recent_assignments) >= 2: - # Check if last two were the same speaker (suggests we should switch) - if recent_assignments[-1] == recent_assignments[-2] == current_speaker: - should_switch = True - # Or if we see a pattern of quick back-and-forth - elif len(recent_assignments) >= 3: - # If pattern is A-A-A, switch to B - if all(a == current_speaker for a in recent_assignments[-3:]): - should_switch = True - # Very small or no gap - use alternating pattern if established - elif gap >= 0: - # If we have an established alternating pattern, continue it - if len(recent_assignments) >= 2: - # Check if we should alternate based on recent pattern - if recent_assignments[-1] == current_speaker: - # If last segment was same speaker, consider switching - # But only if we have a pattern suggesting alternation - if len(recent_assignments) >= 4: - # Check for A-B-A-B pattern - pattern = recent_assignments[-4:] - if pattern[0] != pattern[1] and pattern[1] != pattern[2] and pattern[2] != pattern[3]: - # We have alternating pattern, continue it - should_switch = True - - # Additional heuristics for first few segments - if i < 3 and i > 0: - # Early in conversation, be more aggressive about switching - if gap > 0.1: # Any noticeable gap - should_switch = True - - # If we have multiple consecutive segments from same speaker, force alternation - if len(recent_assignments) >= 2: - # If last 2 segments were both the same speaker (and it's the current speaker), switch - # This prevents one speaker from getting too many consecutive segments - if recent_assignments[-1] == current_speaker and recent_assignments[-2] == current_speaker: - should_switch = True - - # Switch speaker if needed - if should_switch: - current_speaker = "Speaker 2" if current_speaker == "Speaker 1" else "Speaker 1" - - speaker_segments.append( - { - "speaker": current_speaker, - "text": seg_text, - "start": seg_start, - "end": seg_end, - } - ) - - # Track recent assignments for pattern detection - recent_assignments.append(current_speaker) - if len(recent_assignments) > assignment_window: - recent_assignments.pop(0) - - # Post-process: Balance speakers if one dominates too much - speaker_1_count = sum(1 for s in speaker_segments if s["speaker"] == "Speaker 1") - speaker_2_count = sum(1 for s in speaker_segments if s["speaker"] == "Speaker 2") - total_segments = len(speaker_segments) - - # If one speaker has more than 70% of segments, redistribute by alternating - if total_segments > 0: - speaker_1_ratio = speaker_1_count / total_segments - if speaker_1_ratio > 0.7: - # Redistribute: alternate segments starting from index 1 - # This assumes the first speaker is correct, then alternates - for i in range(1, len(speaker_segments)): - expected_speaker = "Speaker 2" if i % 2 == 1 else "Speaker 1" - if speaker_segments[i]["speaker"] != expected_speaker: - speaker_segments[i]["speaker"] = expected_speaker - elif speaker_1_ratio < 0.3: - # Speaker 2 dominates, redistribute - for i in range(1, len(speaker_segments)): - expected_speaker = "Speaker 1" if i % 2 == 1 else "Speaker 2" - if speaker_segments[i]["speaker"] != expected_speaker: - speaker_segments[i]["speaker"] = expected_speaker - - return speaker_segments - - def transcribe_text_only( - self, - audio_file_path: str, - stt_provider: ModelProvider, - stt_model: str, - organization_id: UUID, - db: Session, - language: Optional[str] = None, - credential_id: Optional[UUID] = None, - ) -> Optional[str]: - """Transcribe a local audio file and return just the text. - - Lightweight alternative to `transcribe()` -- skips S3 download, - diarization, and segment extraction. Designed for WER/CER - evaluation where only the transcript string is needed. - """ - api_key = self._get_api_key_for_provider( - stt_provider, db, organization_id, credential_id=credential_id - ) - 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}" - ) - return None - - credential_ctx = ( - self._get_credential_context_for_provider( - stt_provider, db, organization_id, credential_id=credential_id - ) - if stt_provider == ModelProvider.GOOGLE - else None - ) - - from app.services.ai.stt_clients import ( - transcribe_openai, - transcribe_deepgram, - transcribe_elevenlabs, - transcribe_google, - transcribe_sarvam, - transcribe_smallest, - ) - - try: - if stt_provider == ModelProvider.OPENAI: - result = transcribe_openai(audio_file_path, stt_model, api_key, language) - elif stt_provider == ModelProvider.DEEPGRAM: - result = transcribe_deepgram(audio_file_path, stt_model, api_key, language) - 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, - organization_id=organization_id, db=db, - credential=credential_ctx, - ) - elif stt_provider == ModelProvider.SARVAM: - result = transcribe_sarvam(audio_file_path, stt_model, api_key, language) - elif stt_provider == ModelProvider.SMALLEST: - result = transcribe_smallest(audio_file_path, stt_model, api_key, language) - else: - logger.warning(f"[TranscriptionService] Unsupported STT provider for text-only: {stt_provider}") - return None - - text = (result.get("text") or "").strip() - return text or None - except Exception as e: - logger.error(f"[TranscriptionService] text-only transcription failed ({stt_provider}/{stt_model}): {e}") - return None - - def transcribe( - self, - audio_file_key: str, - stt_provider: ModelProvider, - stt_model: str, - organization_id: UUID, - db: Session, - language: Optional[str] = None, - enable_speaker_diarization: bool = True, - credential_id: Optional[UUID] = None, - ) -> Dict[str, Any]: - """ - Transcribe audio file from S3. - """ - start_time = time.time() - temp_file_path = None - - try: - # Download audio to temporary file - temp_file_path = self._download_audio_to_temp(audio_file_key, db=db) - - api_key = self._get_api_key_for_provider( - stt_provider, db, organization_id, credential_id=credential_id - ) - credential_ctx = self._get_credential_context_for_provider( - stt_provider, db, organization_id, credential_id=credential_id - ) - 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." - ) - - from app.services.ai.stt_clients import ( - transcribe_openai, - transcribe_deepgram, - transcribe_elevenlabs, - transcribe_google, - transcribe_sarvam, - transcribe_smallest, - ) - - if stt_provider == ModelProvider.OPENAI: - # All OpenAI STT models in app/config/models.json - # (``whisper-1``, ``gpt-4o-transcribe``, - # ``gpt-4o-mini-transcribe``) are hosted-API models, so - # always go through the OpenAI client. ``transcribe_openai`` - # handles the per-model differences in ``response_format`` - # / ``timestamp_granularities`` internally. Local Whisper - # is reserved for the unknown-provider fallback below. - result = transcribe_openai(temp_file_path, stt_model, api_key, language) - elif stt_provider == ModelProvider.DEEPGRAM: - result = transcribe_deepgram(temp_file_path, stt_model, api_key, language) - elif stt_provider == ModelProvider.ELEVENLABS: - result = transcribe_elevenlabs(temp_file_path, stt_model, api_key, language) - elif stt_provider == ModelProvider.SARVAM: - result = transcribe_sarvam(temp_file_path, stt_model, api_key, language) - elif stt_provider == ModelProvider.SMALLEST: - result = transcribe_smallest(temp_file_path, stt_model, api_key, language) - elif stt_provider == ModelProvider.GOOGLE: - # Gemini STT models (``gemini-2.5-pro-stt``, - # ``gemini-2.5-flash-stt``, ``gemini-2.5-flash-lite-stt``) - # 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, - organization_id=organization_id, db=db, - credential=credential_ctx, - ) - elif stt_provider == ModelProvider.AZURE: - raise NotImplementedError("Azure Speech Services not yet implemented") - elif stt_provider == ModelProvider.AWS: - raise NotImplementedError("AWS Transcribe not yet implemented") - else: - result = self._transcribe_with_whisper_local(temp_file_path, "base") - - # Apply speaker diarization if enabled - speaker_segments = None - if enable_speaker_diarization: - segments = result.get("segments", []) - - # If no segments from transcription, create a single segment from the full text. - # The local audio file is bound to ``temp_file_path`` (downloaded - # from S3 a few lines above); the previous name ``audio_file_path`` - # was a typo that NameError'd as soon as a provider that doesn't - # surface segments (e.g. Gemini STT) was used with diarisation - # enabled. - if not segments and result.get("text"): - estimated_duration = 0.0 - if hasattr(result, "duration") and result.get("duration"): - estimated_duration = result.get("duration") - elif temp_file_path and os.path.exists(temp_file_path): - try: - import librosa - duration = librosa.get_duration(path=temp_file_path) - estimated_duration = duration - except Exception: - pass - - segments = [ - { - "start": 0.0, - "end": estimated_duration if estimated_duration > 0 else 10.0, # Default to 10s if unknown - "text": result.get("text", ""), - } - ] - - if segments: - words = result.get("words", []) - # Check if words actually have valid timestamps - valid_word_count = sum(1 for w in words if w.get("start", 0) > 0 or w.get("end", 0) > 0) - if words and valid_word_count == 0: - words = [] - - try: - from app.config import settings as app_settings - num_spk = getattr(app_settings, "DIARIZATION_NUM_SPEAKERS", 2) - speaker_segments = self._detect_speakers_with_pyannote( - temp_file_path, segments, words, num_speakers=num_spk - ) - if not speaker_segments: - speaker_segments = self._detect_speakers_heuristic(segments) - except Exception as e: - logger.warning(f"Pyannote diarization failed: {e}, falling back to heuristic") - speaker_segments = self._detect_speakers_heuristic(segments) - - processing_time = time.time() - start_time - - return { - "transcript": result["text"], - "language": result.get("language", language), - "speaker_segments": speaker_segments, - "segments": result.get("segments", []), - "processing_time": processing_time, - "raw_output": result, - } - - finally: - # Clean up temporary file - if temp_file_path and os.path.exists(temp_file_path): - try: - os.unlink(temp_file_path) - except Exception: - pass - - -# Singleton instance -transcription_service = TranscriptionService() +""" +Transcription service for converting audio to text using various STT providers. + +Response Format: +All providers return a standardized format: +{ + "text": str, # Full transcript text + "language": str, # Language code (e.g., "en", "es") + "segments": [ # List of segments with timestamps + { + "start": float, # Start time in seconds + "end": float, # End time in seconds + "text": str # Text for this segment + } + ], + "speaker_segments": [ # Segments with speaker labels (if diarization enabled) + { + "speaker": str, # Speaker label (e.g., "Speaker 1") + "text": str, + "start": float, + "end": float + } + ], + "processing_time": float, + "raw_output": dict # Original provider response +} + +Provider-Specific Formats: +- OpenAI Whisper API: Uses verbose_json format which includes segments with timestamps +- Local Whisper: Returns segments by default with word-level timestamps available +- Google/Azure/AWS: Formats documented but not yet implemented + +Speaker Diarization: +- Whisper does NOT provide speaker diarization natively (only transcription) +- We use pyannote.audio (speaker-diarization-3.1) for accurate ML-based diarization +- Whisper provides word-level timestamps, pyannote identifies speakers, then we align +- Requires: pyannote.audio installed + diarization.huggingface_token in config.yml +- Falls back to unreliable gap-based heuristics if pyannote is unavailable +""" + +import time +import tempfile +import os +import logging +from typing import Optional, Dict, Any, List +from uuid import UUID +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 +from sqlalchemy.orm import Session + +logger = logging.getLogger(__name__) + + +class TranscriptionService: + """Service for transcribing audio files using various STT providers.""" + + def __init__(self): + """Initialize transcription service.""" + self._pyannote_pipeline = None + + def _get_ai_provider( + self, + provider: ModelProvider, + db: Session, + organization_id: UUID, + credential_id: Optional[UUID] = None, + ) -> Optional[AIProvider]: + """Resolve the AIProvider row to use, honoring ``credential_id``.""" + return resolve_ai_provider( + provider, db, organization_id, credential_id=credential_id + ) + + def _get_api_key_for_provider( + self, + provider: ModelProvider, + db: Session, + organization_id: UUID, + credential_id: Optional[UUID] = None, + ) -> Optional[str]: + """Resolve and decrypt API key from AIProvider or Integration tables. + + Checks AIProvider first (LLM-style providers like OpenAI), then falls + back to the Integration table (voice platforms like Deepgram, + ElevenLabs). When ``credential_id`` is given the explicit row is + preferred in either table; otherwise the row marked ``is_default`` + wins, with a back-compat fallback to the most recent active row. + """ + from app.services.ai.llm_gateway import ( + resolve_litellm_api_key, + routing_context_from_ai_provider, + ) + + ai_provider = self._get_ai_provider( + provider, db, organization_id, credential_id=credential_id + ) + if ai_provider: + credential_ctx = routing_context_from_ai_provider(ai_provider) + return resolve_litellm_api_key( + organization_id, + db, + ai_provider, + credential=credential_ctx, + ) + + integration = resolve_integration( + provider, db, organization_id, credential_id=credential_id + ) + if integration: + return decrypt_api_key(integration.api_key) + + return None + + def _get_credential_context_for_provider( + self, + provider: ModelProvider, + db: Session, + organization_id: UUID, + credential_id: Optional[UUID] = None, + ): + from app.services.ai.llm_gateway import routing_context_from_ai_provider + + ai_provider = self._get_ai_provider( + provider, db, organization_id, credential_id=credential_id + ) + if ai_provider: + return routing_context_from_ai_provider(ai_provider) + return None + + def _download_audio_to_temp(self, audio_file_key: str, db: Optional[Session] = None) -> str: + """ + Download audio from S3 to temporary file, or use local file if S3 is not available. + """ + import os + + # First, try S3 if enabled + if s3_service.is_enabled(): + try: + # Download from S3 + audio_bytes = s3_service.download_file_by_key(audio_file_key) + + # Determine file extension from key + file_ext = Path(audio_file_key).suffix.lstrip(".") or "wav" + + # Create temporary file + with tempfile.NamedTemporaryFile(delete=False, suffix=f".{file_ext}") as temp_file: + temp_file.write(audio_bytes) + return temp_file.name + except Exception as e: + # If S3 download fails, fall through to local file check + logger.warning(f"S3 download failed for {audio_file_key}: {e}, trying local file...") + + # Fallback: Try to find local file + # Check if it's already a local file path + if os.path.exists(audio_file_key): + # It's a local file path, return it directly + return audio_file_key + + # Try to look up in database if db session is provided + if db: + try: + from app.models.database import AudioFile + # Try to find by S3 key or file path + audio_file = db.query(AudioFile).filter( + (AudioFile.file_path == audio_file_key) | (AudioFile.file_path.like(f"%{audio_file_key}%")) + ).first() + + if audio_file and os.path.exists(audio_file.file_path): + return audio_file.file_path + except Exception as e: + logger.warning(f"Database lookup failed for {audio_file_key}: {e}") + + # If all else fails, raise error + raise StorageError(f"Failed to download audio file: Cloud blob storage is not enabled and local file not found for key: {audio_file_key}") + + # Provider-specific transcription is delegated to app.services.ai.stt_clients + + @staticmethod + def _transcribe_with_whisper_local(audio_file_path: str, model_name: str = "base") -> Dict[str, Any]: + """Transcribe audio using local Whisper model (not a remote API).""" + try: + import whisper + + model = whisper.load_model(model_name) + result = model.transcribe(audio_file_path) + + return { + "text": result.get("text", ""), + "language": result.get("language", "en"), + "segments": [ + {"start": seg.get("start", 0), "end": seg.get("end", 0), "text": seg.get("text", "")} + for seg in result.get("segments", []) + ], + } + except ImportError: + raise RuntimeError("Whisper library not installed. Install with: pip install openai-whisper") + except Exception as e: + raise RuntimeError(f"Whisper transcription failed: {str(e)}") + + def _get_pyannote_pipeline(self): + """Load and cache the pyannote diarization pipeline.""" + if self._pyannote_pipeline is not None: + return self._pyannote_pipeline + + # Compatibility shim: list_audio_backends was removed in torchaudio 2.4+ + import torchaudio + if not hasattr(torchaudio, "list_audio_backends"): + torchaudio.list_audio_backends = lambda: ["ffmpeg"] + + from pyannote.audio import Pipeline + from app.config import settings + + hf_token = settings.HUGGINGFACE_TOKEN + if not hf_token: + raise RuntimeError( + "HUGGINGFACE_TOKEN not configured. Set it under 'diarization.huggingface_token' " + "in config.yml. Required for pyannote speaker diarization." + ) + + logger.info("Loading pyannote speaker-diarization-3.1 pipeline (first call, will be cached)...") + self._pyannote_pipeline = Pipeline.from_pretrained("pyannote/speaker-diarization-3.1", token=hf_token) + logger.info("Pyannote pipeline loaded successfully") + return self._pyannote_pipeline + + def _detect_speakers_with_pyannote( + self, + audio_file_path: str, + segments: List[Dict[str, Any]], + words: Optional[List[Dict[str, Any]]] = None, + num_speakers: Optional[int] = 2, + min_speakers: Optional[int] = None, + max_speakers: Optional[int] = None, + ) -> List[Dict[str, Any]]: + """ + Use pyannote.audio for ML-based speaker diarization, aligned with + Whisper word timestamps for accurate speaker-text mapping. + """ + pipeline = self._get_pyannote_pipeline() + + # Build pipeline kwargs for speaker count hints + pipeline_kwargs = {} + if num_speakers is not None: + pipeline_kwargs["num_speakers"] = num_speakers + if min_speakers is not None: + pipeline_kwargs["min_speakers"] = min_speakers + if max_speakers is not None: + pipeline_kwargs["max_speakers"] = max_speakers + + logger.info( + f"Running pyannote diarization on {audio_file_path} " f"(speaker hints: {pipeline_kwargs or 'auto'})" + ) + raw_output = pipeline(audio_file_path, **pipeline_kwargs) + + # Handle different return types across pyannote versions: + # - Older versions: Annotation directly (has itertracks) + # - Newer versions: DiarizeOutput with .speaker_diarization attribute + if hasattr(raw_output, "itertracks"): + annotation = raw_output + elif hasattr(raw_output, "speaker_diarization"): + annotation = raw_output.speaker_diarization + elif hasattr(raw_output, "annotation"): + annotation = raw_output.annotation + elif isinstance(raw_output, tuple): + annotation = raw_output[0] + else: + attrs = [a for a in dir(raw_output) if not a.startswith("_")] + raise TypeError( + f"Unexpected pyannote output type: {type(raw_output).__name__}. " f"Available attributes: {attrs}" + ) + + # Collect diarization turns as a sorted list for fast lookup + diar_turns = [] + raw_labels = set() + for turn, _, speaker_label in annotation.itertracks(yield_label=True): + diar_turns.append((turn.start, turn.end, speaker_label)) + raw_labels.add(speaker_label) + + if not diar_turns: + logger.warning("Pyannote returned no speaker turns") + return [] + + sorted_labels = sorted(raw_labels) + label_map = {lbl: f"Speaker {i + 1}" for i, lbl in enumerate(sorted_labels)} + logger.info(f"Pyannote detected {len(sorted_labels)} speakers, {len(diar_turns)} turns") + + def find_speaker(midpoint: float) -> str: + """Find which speaker is active at a given timestamp.""" + for t_start, t_end, lbl in diar_turns: + if t_start <= midpoint <= t_end: + return label_map[lbl] + # No exact match -- find the closest turn + min_dist = float("inf") + closest_label = label_map[diar_turns[0][2]] + for t_start, t_end, lbl in diar_turns: + dist = min(abs(midpoint - t_start), abs(midpoint - t_end)) + if dist < min_dist: + min_dist = dist + closest_label = label_map[lbl] + return closest_label + + if words and len(words) > 0: + speaker_segments = [] + current_speaker = None + current_words: List[str] = [] + current_start = 0.0 + current_end = 0.0 + + for w in words: + word_text = w.get("word", "").strip() + w_start = w.get("start", 0) + w_end = w.get("end", 0) + if not word_text: + continue + + midpoint = (w_start + w_end) / 2.0 + speaker = find_speaker(midpoint) + + if speaker != current_speaker: + if current_words and current_speaker: + speaker_segments.append( + { + "speaker": current_speaker, + "text": " ".join(current_words).strip(), + "start": round(current_start, 3), + "end": round(current_end, 3), + } + ) + current_speaker = speaker + current_words = [word_text] + current_start = w_start + current_end = w_end + else: + current_words.append(word_text) + current_end = w_end + + if current_words and current_speaker: + speaker_segments.append( + { + "speaker": current_speaker, + "text": " ".join(current_words).strip(), + "start": round(current_start, 3), + "end": round(current_end, 3), + } + ) + + return speaker_segments + + # Fallback: align at segment level when word timestamps are not available + speaker_segments = [] + for seg in segments: + seg_start = seg.get("start", 0) + seg_end = seg.get("end", 0) + seg_text = seg.get("text", "").strip() + if not seg_text: + continue + + midpoint = (seg_start + seg_end) / 2.0 + speaker = find_speaker(midpoint) + + speaker_segments.append( + { + "speaker": speaker, + "text": seg_text, + "start": round(seg_start, 3), + "end": round(seg_end, 3), + } + ) + + return speaker_segments + + def _detect_speakers_heuristic(self, segments: List[Dict[str, Any]]) -> List[Dict[str, Any]]: + """ + Heuristic-based speaker diarization fallback. Uses gap-based detection + which is unreliable -- pyannote.audio should be used for accurate results. + """ + logger.warning( + "Using heuristic speaker diarization (unreliable). For accurate results, " + "install pyannote.audio and configure diarization.huggingface_token in config.yml" + ) + if not segments: + return [] + + speaker_segments = [] + current_speaker = "Speaker 1" + + # Thresholds for speaker change detection + GAP_THRESHOLD_LARGE = 0.8 # seconds - large gap definitely suggests change + GAP_THRESHOLD_MEDIUM = 0.4 # seconds - medium gap suggests change + GAP_THRESHOLD_SMALL = 0.2 # seconds - small gap, but combined with other factors + MIN_SEGMENT_DURATION = 0.2 # seconds - minimum segment to consider + + # Calculate average gap size to adapt thresholds + gaps = [] + for i in range(1, len(segments)): + gap = segments[i].get("start", 0) - segments[i - 1].get("end", 0) + if gap > 0: + gaps.append(gap) + + avg_gap = sum(gaps) / len(gaps) if gaps else 0.5 + # Adaptive threshold based on conversation pace + adaptive_threshold = min(max(avg_gap * 1.5, 0.3), 1.0) + + # Track recent speaker assignments for pattern detection + recent_assignments = [] + assignment_window = 5 # Look at last N segments for patterns + + for i, seg in enumerate(segments): + seg_start = seg.get("start", 0) + seg_end = seg.get("end", 0) + seg_duration = seg_end - seg_start + seg_text = seg.get("text", "").strip() + + # Skip very short segments (likely noise or artifacts) + if seg_duration < MIN_SEGMENT_DURATION or not seg_text: + continue + + # Check for speaker change indicators + should_switch = False + gap = 0 + + if i > 0: + prev_seg = segments[i - 1] + gap = seg_start - prev_seg.get("end", 0) + + # Large gap definitely suggests speaker change + if gap > GAP_THRESHOLD_LARGE: + should_switch = True + # Medium gap suggests change + elif gap > GAP_THRESHOLD_MEDIUM: + should_switch = True + # Small gap but check for alternating pattern + elif gap > GAP_THRESHOLD_SMALL: + # If we've been alternating, continue the pattern + if len(recent_assignments) >= 2: + # Check if last two were the same speaker (suggests we should switch) + if recent_assignments[-1] == recent_assignments[-2] == current_speaker: + should_switch = True + # Or if we see a pattern of quick back-and-forth + elif len(recent_assignments) >= 3: + # If pattern is A-A-A, switch to B + if all(a == current_speaker for a in recent_assignments[-3:]): + should_switch = True + # Very small or no gap - use alternating pattern if established + elif gap >= 0: + # If we have an established alternating pattern, continue it + if len(recent_assignments) >= 2: + # Check if we should alternate based on recent pattern + if recent_assignments[-1] == current_speaker: + # If last segment was same speaker, consider switching + # But only if we have a pattern suggesting alternation + if len(recent_assignments) >= 4: + # Check for A-B-A-B pattern + pattern = recent_assignments[-4:] + if pattern[0] != pattern[1] and pattern[1] != pattern[2] and pattern[2] != pattern[3]: + # We have alternating pattern, continue it + should_switch = True + + # Additional heuristics for first few segments + if i < 3 and i > 0: + # Early in conversation, be more aggressive about switching + if gap > 0.1: # Any noticeable gap + should_switch = True + + # If we have multiple consecutive segments from same speaker, force alternation + if len(recent_assignments) >= 2: + # If last 2 segments were both the same speaker (and it's the current speaker), switch + # This prevents one speaker from getting too many consecutive segments + if recent_assignments[-1] == current_speaker and recent_assignments[-2] == current_speaker: + should_switch = True + + # Switch speaker if needed + if should_switch: + current_speaker = "Speaker 2" if current_speaker == "Speaker 1" else "Speaker 1" + + speaker_segments.append( + { + "speaker": current_speaker, + "text": seg_text, + "start": seg_start, + "end": seg_end, + } + ) + + # Track recent assignments for pattern detection + recent_assignments.append(current_speaker) + if len(recent_assignments) > assignment_window: + recent_assignments.pop(0) + + # Post-process: Balance speakers if one dominates too much + speaker_1_count = sum(1 for s in speaker_segments if s["speaker"] == "Speaker 1") + speaker_2_count = sum(1 for s in speaker_segments if s["speaker"] == "Speaker 2") + total_segments = len(speaker_segments) + + # If one speaker has more than 70% of segments, redistribute by alternating + if total_segments > 0: + speaker_1_ratio = speaker_1_count / total_segments + if speaker_1_ratio > 0.7: + # Redistribute: alternate segments starting from index 1 + # This assumes the first speaker is correct, then alternates + for i in range(1, len(speaker_segments)): + expected_speaker = "Speaker 2" if i % 2 == 1 else "Speaker 1" + if speaker_segments[i]["speaker"] != expected_speaker: + speaker_segments[i]["speaker"] = expected_speaker + elif speaker_1_ratio < 0.3: + # Speaker 2 dominates, redistribute + for i in range(1, len(speaker_segments)): + expected_speaker = "Speaker 1" if i % 2 == 1 else "Speaker 2" + if speaker_segments[i]["speaker"] != expected_speaker: + speaker_segments[i]["speaker"] = expected_speaker + + return speaker_segments + + def transcribe_text_only( + self, + audio_file_path: str, + stt_provider: ModelProvider, + stt_model: str, + organization_id: UUID, + db: Session, + language: Optional[str] = None, + credential_id: Optional[UUID] = None, + ) -> Optional[str]: + """Transcribe a local audio file and return just the text. + + Lightweight alternative to `transcribe()` -- skips S3 download, + diarization, and segment extraction. Designed for WER/CER + evaluation where only the transcript string is needed. + """ + api_key = self._get_api_key_for_provider( + stt_provider, db, organization_id, credential_id=credential_id + ) + 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}" + ) + return None + + credential_ctx = ( + self._get_credential_context_for_provider( + stt_provider, db, organization_id, credential_id=credential_id + ) + if stt_provider == ModelProvider.GOOGLE + else None + ) + + from app.services.ai.stt_clients import ( + transcribe_openai, + transcribe_deepgram, + transcribe_elevenlabs, + transcribe_google, + transcribe_sarvam, + transcribe_smallest, + ) + + try: + if stt_provider == ModelProvider.OPENAI: + result = transcribe_openai(audio_file_path, stt_model, api_key, language) + elif stt_provider == ModelProvider.DEEPGRAM: + result = transcribe_deepgram(audio_file_path, stt_model, api_key, language) + 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, + organization_id=organization_id, db=db, + credential=credential_ctx, + ) + elif stt_provider == ModelProvider.SARVAM: + result = transcribe_sarvam(audio_file_path, stt_model, api_key, language) + elif stt_provider == ModelProvider.SMALLEST: + result = transcribe_smallest(audio_file_path, stt_model, api_key, language) + else: + logger.warning(f"[TranscriptionService] Unsupported STT provider for text-only: {stt_provider}") + return None + + text = (result.get("text") or "").strip() + try: + from app.services.usage.llm_usage import ( + probe_audio_seconds, + record_stt_usage, + ) + + audio_seconds = probe_audio_seconds(audio_file_path) + if stt_provider != ModelProvider.GOOGLE: + record_stt_usage( + stt_model or "unknown", + audio_seconds=audio_seconds, + organization_id=organization_id, + ) + except Exception as exc: + logger.debug( + "[TranscriptionService] text-only stt usage record skipped: %s", + exc, + ) + return text or None + except Exception as e: + logger.error(f"[TranscriptionService] text-only transcription failed ({stt_provider}/{stt_model}): {e}") + return None + + def transcribe( + self, + audio_file_key: str, + stt_provider: ModelProvider, + stt_model: str, + organization_id: UUID, + db: Session, + language: Optional[str] = None, + enable_speaker_diarization: bool = True, + credential_id: Optional[UUID] = None, + ) -> Dict[str, Any]: + """ + Transcribe audio file from S3. + """ + start_time = time.time() + temp_file_path = None + + try: + # Download audio to temporary file + temp_file_path = self._download_audio_to_temp(audio_file_key, db=db) + + api_key = self._get_api_key_for_provider( + stt_provider, db, organization_id, credential_id=credential_id + ) + credential_ctx = self._get_credential_context_for_provider( + stt_provider, db, organization_id, credential_id=credential_id + ) + 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." + ) + + from app.services.ai.stt_clients import ( + transcribe_openai, + transcribe_deepgram, + transcribe_elevenlabs, + transcribe_google, + transcribe_sarvam, + transcribe_smallest, + ) + + if stt_provider == ModelProvider.OPENAI: + # All OpenAI STT models in app/config/models.json + # (``whisper-1``, ``gpt-4o-transcribe``, + # ``gpt-4o-mini-transcribe``) are hosted-API models, so + # always go through the OpenAI client. ``transcribe_openai`` + # handles the per-model differences in ``response_format`` + # / ``timestamp_granularities`` internally. Local Whisper + # is reserved for the unknown-provider fallback below. + result = transcribe_openai(temp_file_path, stt_model, api_key, language) + elif stt_provider == ModelProvider.DEEPGRAM: + result = transcribe_deepgram(temp_file_path, stt_model, api_key, language) + elif stt_provider == ModelProvider.ELEVENLABS: + result = transcribe_elevenlabs(temp_file_path, stt_model, api_key, language) + elif stt_provider == ModelProvider.SARVAM: + result = transcribe_sarvam(temp_file_path, stt_model, api_key, language) + elif stt_provider == ModelProvider.SMALLEST: + result = transcribe_smallest(temp_file_path, stt_model, api_key, language) + elif stt_provider == ModelProvider.GOOGLE: + # Gemini STT models (``gemini-2.5-pro-stt``, + # ``gemini-2.5-flash-stt``, ``gemini-2.5-flash-lite-stt``) + # 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, + organization_id=organization_id, db=db, + credential=credential_ctx, + ) + elif stt_provider == ModelProvider.AZURE: + raise NotImplementedError("Azure Speech Services not yet implemented") + elif stt_provider == ModelProvider.AWS: + raise NotImplementedError("AWS Transcribe not yet implemented") + else: + result = self._transcribe_with_whisper_local(temp_file_path, "base") + + try: + from app.services.usage.llm_usage import ( + probe_audio_seconds, + record_stt_usage, + ) + + audio_seconds = 0 + if isinstance(result, dict): + raw_duration = result.get("duration") + if raw_duration is not None: + try: + audio_seconds = int(float(raw_duration)) + except (TypeError, ValueError): + audio_seconds = 0 + if audio_seconds <= 0 and temp_file_path: + audio_seconds = probe_audio_seconds(temp_file_path) + if stt_provider != ModelProvider.GOOGLE: + record_stt_usage( + stt_model or "unknown", + audio_seconds=audio_seconds, + organization_id=organization_id, + ) + except Exception as exc: + logger.debug( + "[TranscriptionService] stt usage record skipped: %s", exc + ) + + # Apply speaker diarization if enabled + speaker_segments = None + if enable_speaker_diarization: + segments = result.get("segments", []) + + # If no segments from transcription, create a single segment from the full text. + # The local audio file is bound to ``temp_file_path`` (downloaded + # from S3 a few lines above); the previous name ``audio_file_path`` + # was a typo that NameError'd as soon as a provider that doesn't + # surface segments (e.g. Gemini STT) was used with diarisation + # enabled. + if not segments and result.get("text"): + estimated_duration = 0.0 + if hasattr(result, "duration") and result.get("duration"): + estimated_duration = result.get("duration") + elif temp_file_path and os.path.exists(temp_file_path): + try: + import librosa + duration = librosa.get_duration(path=temp_file_path) + estimated_duration = duration + except Exception: + pass + + segments = [ + { + "start": 0.0, + "end": estimated_duration if estimated_duration > 0 else 10.0, # Default to 10s if unknown + "text": result.get("text", ""), + } + ] + + if segments: + words = result.get("words", []) + # Check if words actually have valid timestamps + valid_word_count = sum(1 for w in words if w.get("start", 0) > 0 or w.get("end", 0) > 0) + if words and valid_word_count == 0: + words = [] + + try: + from app.config import settings as app_settings + num_spk = getattr(app_settings, "DIARIZATION_NUM_SPEAKERS", 2) + speaker_segments = self._detect_speakers_with_pyannote( + temp_file_path, segments, words, num_speakers=num_spk + ) + if not speaker_segments: + speaker_segments = self._detect_speakers_heuristic(segments) + except Exception as e: + logger.warning(f"Pyannote diarization failed: {e}, falling back to heuristic") + speaker_segments = self._detect_speakers_heuristic(segments) + + processing_time = time.time() - start_time + + return { + "transcript": result["text"], + "language": result.get("language", language), + "speaker_segments": speaker_segments, + "segments": result.get("segments", []), + "processing_time": processing_time, + "raw_output": result, + } + + finally: + # Clean up temporary file + if temp_file_path and os.path.exists(temp_file_path): + try: + os.unlink(temp_file_path) + except Exception: + pass + + +# Singleton instance +transcription_service = TranscriptionService() diff --git a/app/services/ai/tts_service.py b/app/services/ai/tts_service.py index 60d41255..4d1162f4 100644 --- a/app/services/ai/tts_service.py +++ b/app/services/ai/tts_service.py @@ -1,312 +1,349 @@ -""" -TTS service for converting text to speech using various TTS providers. -""" - -import time -from typing import Optional, Dict, Any, Tuple -from uuid import UUID - -from app.models.database import ModelProvider, AIProvider, Integration -from app.services.storage.s3_service import s3_service -from efficientai.services.cartesia.http_tts import synthesize_cartesia_bytes -from efficientai.services.deepgram.http_tts import synthesize_deepgram_bytes -from efficientai.services.elevenlabs.http_tts import synthesize_elevenlabs_bytes -from efficientai.services.google.http_tts import synthesize_google_bytes -from efficientai.services.murf.tts import synthesize_murf_stream_bytes -from efficientai.services.openai.http_tts import synthesize_openai_bytes -from efficientai.services.sarvam.http_tts import synthesize_sarvam_bytes -from efficientai.services.smallest.http_tts import synthesize_smallest_bytes -from efficientai.services.voicemaker.http_tts import synthesize_voicemaker_bytes -from sqlalchemy.orm import Session - - -ELEVENLABS_HZ_TO_OUTPUT_FORMAT = { - 8000: "pcm_8000", - 16000: "pcm_16000", - 22050: "mp3_22050_32", - 24000: "pcm_24000", - 44100: "mp3_44100_128", -} - -PROVIDER_SUPPORTED_SAMPLE_RATES: Dict[str, list] = { - "elevenlabs": [8000, 16000, 22050, 24000, 44100], - "cartesia": [8000, 16000, 22050, 24000, 44100], - "deepgram": [8000, 16000, 24000, 48000], - "sarvam": [8000, 16000, 22050], - "murf": [8000, 16000, 24000, 44100, 48000], - "smallest": [8000, 16000, 24000], - "voicemaker": [8000, 16000, 22050, 24000, 44100, 48000], -} - - -def get_audio_file_extension(provider: str, sample_rate_hz: Optional[int] = None) -> str: - """Determine audio file extension based on provider and requested sample rate.""" - if provider == "sarvam": - # Sarvam HTTP TTS returns base64 WAV audio. - return "wav" - if provider == "smallest": - return "wav" - if provider == "elevenlabs" and sample_rate_hz: - fmt = ELEVENLABS_HZ_TO_OUTPUT_FORMAT.get(sample_rate_hz, "") - if fmt.startswith(("pcm_", "ulaw_")): - return "wav" - return "mp3" - - -class TTSService: - """Service for converting text to speech using various TTS providers.""" - - def __init__(self): - self._provider_handlers: Dict[str, Any] = {} - - def _get_ai_provider(self, provider: ModelProvider, db: Session, organization_id: UUID) -> Optional[AIProvider]: - """Get AI provider configuration from database.""" - from sqlalchemy import func - - provider_value = provider.value if hasattr(provider, "value") else provider - - ai_provider = db.query(AIProvider).filter( - AIProvider.provider == provider_value, - AIProvider.organization_id == organization_id, - AIProvider.is_active == True, - ).first() - - if not ai_provider: - ai_provider = db.query(AIProvider).filter( - func.lower(AIProvider.provider) == provider_value.lower(), - AIProvider.organization_id == organization_id, - AIProvider.is_active == True, - ).first() - - return ai_provider - - def _get_api_key_for_provider( - self, provider: ModelProvider, db: Session, organization_id: UUID - ) -> str: - """Resolve and decrypt API key from AIProvider or Integration tables.""" - from app.core.encryption import decrypt_api_key - from sqlalchemy import func - - ai_provider = self._get_ai_provider(provider, db, organization_id) - if ai_provider: - return decrypt_api_key(ai_provider.api_key) - - # Fallback: check Integration table for cartesia/elevenlabs/deepgram - provider_value = provider.value if hasattr(provider, "value") else provider - integration = db.query(Integration).filter( - func.lower(Integration.platform) == provider_value.lower(), - Integration.organization_id == organization_id, - Integration.is_active == True, - ).first() - if integration: - return decrypt_api_key(integration.api_key) - - raise RuntimeError(f"No API key configured for provider {provider_value}") - - # ------------------------------------------------------------------ - # OpenAI - # ------------------------------------------------------------------ - - def _synthesize_with_openai( - self, text: str, model: str, api_key: str, - voice: Optional[str] = "alloy", config: Optional[Dict[str, Any]] = None - ) -> Tuple[bytes, float]: - return synthesize_openai_bytes(text=text, model=model, api_key=api_key, voice=voice, config=config) - - # ------------------------------------------------------------------ - # Google (unary RPC - no streaming; TTFB ~= total API time) - # ------------------------------------------------------------------ - - def _synthesize_with_google( - self, text: str, model: str, api_key: str, - voice: Optional[str] = None, config: Optional[Dict[str, Any]] = None - ) -> Tuple[bytes, float]: - return synthesize_google_bytes(text=text, model=model, api_key=api_key, voice=voice, config=config) - - # ------------------------------------------------------------------ - # ElevenLabs - # ------------------------------------------------------------------ - - def _synthesize_with_elevenlabs( - self, text: str, model: str, api_key: str, - voice: Optional[str] = None, config: Optional[Dict[str, Any]] = None - ) -> Tuple[bytes, float]: - return synthesize_elevenlabs_bytes(text=text, model=model, api_key=api_key, voice=voice, config=config) - - # ------------------------------------------------------------------ - # Cartesia - # ------------------------------------------------------------------ - - def _synthesize_with_cartesia( - self, text: str, model: str, api_key: str, - voice: Optional[str] = None, config: Optional[Dict[str, Any]] = None - ) -> Tuple[bytes, float]: - return synthesize_cartesia_bytes(text=text, model=model, api_key=api_key, voice=voice, config=config) - - # ------------------------------------------------------------------ - # Deepgram - # ------------------------------------------------------------------ - - def _synthesize_with_deepgram( - self, text: str, model: str, api_key: str, - voice: Optional[str] = None, config: Optional[Dict[str, Any]] = None - ) -> Tuple[bytes, float]: - return synthesize_deepgram_bytes(text=text, model=model, api_key=api_key, voice=voice, config=config) - - # ------------------------------------------------------------------ - # Sarvam - # ------------------------------------------------------------------ - - def _synthesize_with_sarvam( - self, text: str, model: str, api_key: str, - voice: Optional[str] = None, config: Optional[Dict[str, Any]] = None - ) -> Tuple[bytes, float]: - return synthesize_sarvam_bytes(text=text, model=model, api_key=api_key, voice=voice, config=config) - - # ------------------------------------------------------------------ - # VoiceMaker - # ------------------------------------------------------------------ - - def _synthesize_with_voicemaker( - self, text: str, model: str, api_key: str, - voice: Optional[str] = None, config: Optional[Dict[str, Any]] = None - ) -> Tuple[bytes, float]: - return synthesize_voicemaker_bytes(text=text, model=model, api_key=api_key, voice=voice, config=config) - - # ------------------------------------------------------------------ - # Smallest.ai - # ------------------------------------------------------------------ - - def _synthesize_with_smallest( - self, text: str, model: str, api_key: str, - voice: Optional[str] = None, config: Optional[Dict[str, Any]] = None - ) -> Tuple[bytes, float]: - return synthesize_smallest_bytes(text=text, model=model, api_key=api_key, voice=voice, config=config) - - # ------------------------------------------------------------------ - # Main synthesis entry point - # ------------------------------------------------------------------ - - def register_tts_provider(self, provider: str, handler: Any) -> None: - """Register/override a TTS provider handler dynamically.""" - self._provider_handlers[provider.strip().lower()] = handler - - def _get_tts_handler(self, tts_provider: ModelProvider): - provider_key = (tts_provider.value if hasattr(tts_provider, "value") else str(tts_provider)).lower() - - # Prefer explicitly registered handlers. - handler = self._provider_handlers.get(provider_key) - if handler: - return handler - - # Fallback convention: _synthesize_with_. - handler = getattr(self, f"_synthesize_with_{provider_key}", None) - if callable(handler): - # Cache resolved handler to avoid repeated getattr lookups. - self._provider_handlers[provider_key] = handler - return handler - - if not handler: - raise NotImplementedError(f"TTS provider {tts_provider} not supported") - return handler - - def _synthesize_with_murf( - self, - text: str, - model: str, - api_key: str, - voice: Optional[str] = None, - config: Optional[Dict[str, Any]] = None - ) -> Tuple[bytes, float]: - """Delegate Murf synthesis to shared efficientai Murf service helper.""" - try: - return synthesize_murf_stream_bytes( - text=text, - model=model, - api_key=api_key, - voice=voice, - config=config, - ) - except ImportError: - raise RuntimeError("requests library not installed. Install with: pip install requests") - except Exception as e: - import traceback - error_details = traceback.format_exc() - raise RuntimeError(f"Murf TTS synthesis failed: {str(e)}\nDetails: {error_details}") - - def synthesize( - self, - text: str, - tts_provider: ModelProvider, - tts_model: str, - organization_id: UUID, - db: Session, - voice: Optional[str] = None, - config: Optional[Dict[str, Any]] = None, - ) -> bytes: - """ - Synthesize speech from text. - - Args: - text: Text to convert to speech - tts_provider: TTS provider to use - tts_model: TTS model name - organization_id: Organization ID - db: Database session - voice: Voice selection (if applicable) - config: Additional provider-specific configuration - - Returns: - Audio bytes (MP3 format) - """ - api_key = self._get_api_key_for_provider(tts_provider, db, organization_id) - handler = self._get_tts_handler(tts_provider) - audio_bytes, _ttfb_ms = handler(text, tts_model, api_key, voice, config) - return audio_bytes - - def synthesize_timed( - self, - text: str, - tts_provider: ModelProvider, - tts_model: str, - organization_id: UUID, - db: Session, - voice: Optional[str] = None, - config: Optional[Dict[str, Any]] = None, - ) -> Tuple[bytes, float, float]: - """Synthesize and return (audio_bytes, total_latency_ms, ttfb_ms).""" - api_key = self._get_api_key_for_provider(tts_provider, db, organization_id) - handler = self._get_tts_handler(tts_provider) - start = time.time() - audio_bytes, ttfb_ms = handler(text, tts_model, api_key, voice, config) - total_latency_ms = (time.time() - start) * 1000 - return audio_bytes, total_latency_ms, ttfb_ms - - def synthesize_and_upload( - self, - text: str, - tts_provider: ModelProvider, - tts_model: str, - organization_id: UUID, - db: Session, - voice: Optional[str] = None, - config: Optional[Dict[str, Any]] = None, - file_prefix: str = "tts_output", - ) -> str: - """Synthesize speech and upload to S3. Returns S3 key.""" - audio_bytes = self.synthesize(text, tts_provider, tts_model, organization_id, db, voice, config) - - import uuid as _uuid - - file_id = _uuid.uuid4() - s3_key = s3_service.upload_file( - file_id=file_id, - file_content=audio_bytes, - file_format="mp3", - organization_id=organization_id, - ) - return s3_key - - -# Singleton instance -tts_service = TTSService() +""" +TTS service for converting text to speech using various TTS providers. +""" + +import time +from typing import Optional, Dict, Any, Tuple +from uuid import UUID + +from app.models.database import ModelProvider, AIProvider, Integration +from app.services.storage.s3_service import s3_service +from efficientai.services.cartesia.http_tts import synthesize_cartesia_bytes +from efficientai.services.deepgram.http_tts import synthesize_deepgram_bytes +from efficientai.services.elevenlabs.http_tts import synthesize_elevenlabs_bytes +from efficientai.services.google.http_tts import synthesize_google_bytes +from efficientai.services.murf.tts import synthesize_murf_stream_bytes +from efficientai.services.openai.http_tts import synthesize_openai_bytes +from efficientai.services.sarvam.http_tts import synthesize_sarvam_bytes +from efficientai.services.smallest.http_tts import synthesize_smallest_bytes +from efficientai.services.voicemaker.http_tts import synthesize_voicemaker_bytes +from sqlalchemy.orm import Session + + +ELEVENLABS_HZ_TO_OUTPUT_FORMAT = { + 8000: "pcm_8000", + 16000: "pcm_16000", + 22050: "mp3_22050_32", + 24000: "pcm_24000", + 44100: "mp3_44100_128", +} + +PROVIDER_SUPPORTED_SAMPLE_RATES: Dict[str, list] = { + "elevenlabs": [8000, 16000, 22050, 24000, 44100], + "cartesia": [8000, 16000, 22050, 24000, 44100], + "deepgram": [8000, 16000, 24000, 48000], + "sarvam": [8000, 16000, 22050], + "murf": [8000, 16000, 24000, 44100, 48000], + "smallest": [8000, 16000, 24000], + "voicemaker": [8000, 16000, 22050, 24000, 44100, 48000], +} + + +def get_audio_file_extension(provider: str, sample_rate_hz: Optional[int] = None) -> str: + """Determine audio file extension based on provider and requested sample rate.""" + if provider == "sarvam": + # Sarvam HTTP TTS returns base64 WAV audio. + return "wav" + if provider == "smallest": + return "wav" + if provider == "elevenlabs" and sample_rate_hz: + fmt = ELEVENLABS_HZ_TO_OUTPUT_FORMAT.get(sample_rate_hz, "") + if fmt.startswith(("pcm_", "ulaw_")): + return "wav" + return "mp3" + + +class TTSService: + """Service for converting text to speech using various TTS providers.""" + + def __init__(self): + self._provider_handlers: Dict[str, Any] = {} + + def _get_ai_provider(self, provider: ModelProvider, db: Session, organization_id: UUID) -> Optional[AIProvider]: + """Get AI provider configuration from database.""" + from sqlalchemy import func + + provider_value = provider.value if hasattr(provider, "value") else provider + + ai_provider = db.query(AIProvider).filter( + AIProvider.provider == provider_value, + AIProvider.organization_id == organization_id, + AIProvider.is_active == True, + ).first() + + if not ai_provider: + ai_provider = db.query(AIProvider).filter( + func.lower(AIProvider.provider) == provider_value.lower(), + AIProvider.organization_id == organization_id, + AIProvider.is_active == True, + ).first() + + return ai_provider + + def _get_api_key_for_provider( + self, provider: ModelProvider, db: Session, organization_id: UUID + ) -> str: + """Resolve and decrypt API key from AIProvider or Integration tables.""" + from app.core.encryption import decrypt_api_key + from sqlalchemy import func + + ai_provider = self._get_ai_provider(provider, db, organization_id) + if ai_provider: + return decrypt_api_key(ai_provider.api_key) + + # Fallback: check Integration table for cartesia/elevenlabs/deepgram + provider_value = provider.value if hasattr(provider, "value") else provider + integration = db.query(Integration).filter( + func.lower(Integration.platform) == provider_value.lower(), + Integration.organization_id == organization_id, + Integration.is_active == True, + ).first() + if integration: + return decrypt_api_key(integration.api_key) + + raise RuntimeError(f"No API key configured for provider {provider_value}") + + # ------------------------------------------------------------------ + # OpenAI + # ------------------------------------------------------------------ + + def _synthesize_with_openai( + self, text: str, model: str, api_key: str, + voice: Optional[str] = "alloy", config: Optional[Dict[str, Any]] = None + ) -> Tuple[bytes, float]: + return synthesize_openai_bytes(text=text, model=model, api_key=api_key, voice=voice, config=config) + + # ------------------------------------------------------------------ + # Google (unary RPC - no streaming; TTFB ~= total API time) + # ------------------------------------------------------------------ + + def _synthesize_with_google( + self, text: str, model: str, api_key: str, + voice: Optional[str] = None, config: Optional[Dict[str, Any]] = None + ) -> Tuple[bytes, float]: + return synthesize_google_bytes(text=text, model=model, api_key=api_key, voice=voice, config=config) + + # ------------------------------------------------------------------ + # ElevenLabs + # ------------------------------------------------------------------ + + def _synthesize_with_elevenlabs( + self, text: str, model: str, api_key: str, + voice: Optional[str] = None, config: Optional[Dict[str, Any]] = None + ) -> Tuple[bytes, float]: + return synthesize_elevenlabs_bytes(text=text, model=model, api_key=api_key, voice=voice, config=config) + + # ------------------------------------------------------------------ + # Cartesia + # ------------------------------------------------------------------ + + def _synthesize_with_cartesia( + self, text: str, model: str, api_key: str, + voice: Optional[str] = None, config: Optional[Dict[str, Any]] = None + ) -> Tuple[bytes, float]: + return synthesize_cartesia_bytes(text=text, model=model, api_key=api_key, voice=voice, config=config) + + # ------------------------------------------------------------------ + # Deepgram + # ------------------------------------------------------------------ + + def _synthesize_with_deepgram( + self, text: str, model: str, api_key: str, + voice: Optional[str] = None, config: Optional[Dict[str, Any]] = None + ) -> Tuple[bytes, float]: + return synthesize_deepgram_bytes(text=text, model=model, api_key=api_key, voice=voice, config=config) + + # ------------------------------------------------------------------ + # Sarvam + # ------------------------------------------------------------------ + + def _synthesize_with_sarvam( + self, text: str, model: str, api_key: str, + voice: Optional[str] = None, config: Optional[Dict[str, Any]] = None + ) -> Tuple[bytes, float]: + return synthesize_sarvam_bytes(text=text, model=model, api_key=api_key, voice=voice, config=config) + + # ------------------------------------------------------------------ + # VoiceMaker + # ------------------------------------------------------------------ + + def _synthesize_with_voicemaker( + self, text: str, model: str, api_key: str, + voice: Optional[str] = None, config: Optional[Dict[str, Any]] = None + ) -> Tuple[bytes, float]: + return synthesize_voicemaker_bytes(text=text, model=model, api_key=api_key, voice=voice, config=config) + + # ------------------------------------------------------------------ + # Smallest.ai + # ------------------------------------------------------------------ + + def _synthesize_with_smallest( + self, text: str, model: str, api_key: str, + voice: Optional[str] = None, config: Optional[Dict[str, Any]] = None + ) -> Tuple[bytes, float]: + return synthesize_smallest_bytes(text=text, model=model, api_key=api_key, voice=voice, config=config) + + # ------------------------------------------------------------------ + # Main synthesis entry point + # ------------------------------------------------------------------ + + def register_tts_provider(self, provider: str, handler: Any) -> None: + """Register/override a TTS provider handler dynamically.""" + self._provider_handlers[provider.strip().lower()] = handler + + def _get_tts_handler(self, tts_provider: ModelProvider): + provider_key = (tts_provider.value if hasattr(tts_provider, "value") else str(tts_provider)).lower() + + # Prefer explicitly registered handlers. + handler = self._provider_handlers.get(provider_key) + if handler: + return handler + + # Fallback convention: _synthesize_with_. + handler = getattr(self, f"_synthesize_with_{provider_key}", None) + if callable(handler): + # Cache resolved handler to avoid repeated getattr lookups. + self._provider_handlers[provider_key] = handler + return handler + + if not handler: + raise NotImplementedError(f"TTS provider {tts_provider} not supported") + return handler + + def _synthesize_with_murf( + self, + text: str, + model: str, + api_key: str, + voice: Optional[str] = None, + config: Optional[Dict[str, Any]] = None + ) -> Tuple[bytes, float]: + """Delegate Murf synthesis to shared efficientai Murf service helper.""" + try: + return synthesize_murf_stream_bytes( + text=text, + model=model, + api_key=api_key, + voice=voice, + config=config, + ) + except ImportError: + raise RuntimeError("requests library not installed. Install with: pip install requests") + except Exception as e: + import traceback + error_details = traceback.format_exc() + raise RuntimeError(f"Murf TTS synthesis failed: {str(e)}\nDetails: {error_details}") + + def synthesize( + self, + text: str, + tts_provider: ModelProvider, + tts_model: str, + organization_id: UUID, + db: Session, + voice: Optional[str] = None, + config: Optional[Dict[str, Any]] = None, + ) -> bytes: + """ + Synthesize speech from text. + + Args: + text: Text to convert to speech + tts_provider: TTS provider to use + tts_model: TTS model name + organization_id: Organization ID + db: Database session + voice: Voice selection (if applicable) + config: Additional provider-specific configuration + + Returns: + Audio bytes (MP3 format) + """ + api_key = self._get_api_key_for_provider(tts_provider, db, organization_id) + handler = self._get_tts_handler(tts_provider) + audio_bytes, _ttfb_ms = handler(text, tts_model, api_key, voice, config) + self._record_tts_usage( + text=text, + tts_model=tts_model, + organization_id=organization_id, + ) + return audio_bytes + + def synthesize_timed( + self, + text: str, + tts_provider: ModelProvider, + tts_model: str, + organization_id: UUID, + db: Session, + voice: Optional[str] = None, + config: Optional[Dict[str, Any]] = None, + ) -> Tuple[bytes, float, float]: + """Synthesize and return (audio_bytes, total_latency_ms, ttfb_ms).""" + api_key = self._get_api_key_for_provider(tts_provider, db, organization_id) + handler = self._get_tts_handler(tts_provider) + start = time.time() + audio_bytes, ttfb_ms = handler(text, tts_model, api_key, voice, config) + total_latency_ms = (time.time() - start) * 1000 + self._record_tts_usage( + text=text, + tts_model=tts_model, + organization_id=organization_id, + ) + return audio_bytes, total_latency_ms, ttfb_ms + + def _record_tts_usage( + self, + *, + text: str, + tts_model: str, + organization_id: UUID, + ) -> None: + try: + from app.services.usage.context import ( + ensure_usage_context, + reset_usage_context, + ) + from app.services.usage.llm_usage import record_tts_usage + + usage_token = ensure_usage_context(organization_id) + try: + record_tts_usage( + tts_model, + characters=len(text or ""), + organization_id=organization_id, + ) + finally: + if usage_token is not None: + reset_usage_context(usage_token) + except Exception: + pass + + def synthesize_and_upload( + self, + text: str, + tts_provider: ModelProvider, + tts_model: str, + organization_id: UUID, + db: Session, + voice: Optional[str] = None, + config: Optional[Dict[str, Any]] = None, + file_prefix: str = "tts_output", + ) -> str: + """Synthesize speech and upload to S3. Returns S3 key.""" + audio_bytes = self.synthesize(text, tts_provider, tts_model, organization_id, db, voice, config) + + import uuid as _uuid + + file_id = _uuid.uuid4() + s3_key = s3_service.upload_file( + file_id=file_id, + file_content=audio_bytes, + file_format="mp3", + organization_id=organization_id, + ) + return s3_key + + +# Singleton instance +tts_service = TTSService() diff --git a/app/services/call_import_user_insights.py b/app/services/call_import_user_insights.py index bf723e6a..7484a34b 100644 --- a/app/services/call_import_user_insights.py +++ b/app/services/call_import_user_insights.py @@ -21,6 +21,8 @@ Metric, ModelProvider, ) +from app.services.usage.context import LLMUsageContext + from app.models.schemas import ( EvaluationUserInsightItem, EvaluationUserInsightsState, @@ -222,17 +224,28 @@ def _call_llm( *, temperature: float, max_tokens: int, + usage_ctx: Optional[LLMUsageContext] = None, ) -> str: - result = llm_service.generate_response( - messages=messages, - llm_provider=provider, - llm_model=model, - organization_id=organization_id, - db=db, - temperature=temperature, - max_tokens=max_tokens, - ) - return str(result.get("text") or "") + from app.services.usage.context import get_usage_context, llm_usage_context + + effective_ctx = usage_ctx or get_usage_context() + + def _run() -> str: + result = llm_service.generate_response( + messages=messages, + llm_provider=provider, + llm_model=model, + organization_id=organization_id, + db=db, + temperature=temperature, + max_tokens=max_tokens, + ) + return str(result.get("text") or "") + + if effective_ctx is not None: + with llm_usage_context(effective_ctx): + return _run() + return _run() def run_extraction_batch( diff --git a/app/services/cron/dispatcher_lock.py b/app/services/cron/dispatcher_lock.py new file mode 100644 index 00000000..26ac638f --- /dev/null +++ b/app/services/cron/dispatcher_lock.py @@ -0,0 +1,81 @@ +"""Redis locks for singleton cron dispatcher.""" + +from __future__ import annotations + +import os + +import redis +from loguru import logger + +from app.config import settings + +_DISPATCHER_LOCK_KEY = "cron:dispatcher:lock" +_DISPATCHER_LEADER_KEY = "cron:dispatcher:leader" +_redis: redis.Redis | None = None + + +def _client() -> redis.Redis: + global _redis + if _redis is None: + _redis = redis.from_url(settings.REDIS_URL, decode_responses=True) + return _redis + + +def _lock_ttl_seconds() -> int: + raw = os.environ.get("CRON_DISPATCH_LOCK_TTL_SECONDS", "55") + try: + return max(10, int(raw)) + except (TypeError, ValueError): + return 55 + + +def _leader_ttl_seconds() -> int: + raw = os.environ.get("CRON_DISPATCH_LEADER_TTL_SECONDS", "300") + try: + return max(60, int(raw)) + except (TypeError, ValueError): + return 300 + + +def try_acquire_dispatcher_leader() -> bool: + try: + return bool( + _client().set( + _DISPATCHER_LEADER_KEY, + "1", + nx=True, + ex=_leader_ttl_seconds(), + ) + ) + except redis.RedisError as exc: + logger.warning("cron dispatcher leader lock skipped: {}", exc) + return False + + +def refresh_dispatcher_leader() -> None: + try: + _client().expire(_DISPATCHER_LEADER_KEY, _leader_ttl_seconds()) + except redis.RedisError: + pass + + +def acquire_dispatcher_run_lock() -> bool: + try: + return bool( + _client().set( + _DISPATCHER_LOCK_KEY, + "1", + nx=True, + ex=_lock_ttl_seconds(), + ) + ) + except redis.RedisError as exc: + logger.warning("cron dispatcher run lock skipped: {}", exc) + return True + + +def release_dispatcher_run_lock() -> None: + try: + _client().delete(_DISPATCHER_LOCK_KEY) + except redis.RedisError: + pass diff --git a/app/services/cron/job_dispatch.py b/app/services/cron/job_dispatch.py new file mode 100644 index 00000000..9b905079 --- /dev/null +++ b/app/services/cron/job_dispatch.py @@ -0,0 +1,112 @@ +"""Dispatch due cron jobs to Celery workers.""" + +from __future__ import annotations + +from collections import defaultdict +from datetime import datetime, timezone +from typing import Any, Dict, List +from uuid import UUID + +from loguru import logger +from sqlalchemy.orm import Session + +from app.models.database import CronJob, Evaluator +from app.models.enums import CronJobStatus +from app.services.cron.scheduling import calculate_next_run + + +def list_due_cron_jobs(db: Session, *, now: datetime | None = None) -> List[CronJob]: + moment = now or datetime.now(timezone.utc) + return ( + db.query(CronJob) + .filter( + CronJob.status == CronJobStatus.ACTIVE.value, + CronJob.is_system.is_(False), + CronJob.next_run_at.isnot(None), + CronJob.next_run_at <= moment, + ) + .order_by(CronJob.next_run_at.asc()) + .all() + ) + + +def advance_cron_job(db: Session, job: CronJob, *, now: datetime | None = None) -> None: + moment = now or datetime.now(timezone.utc) + job.last_run_at = moment + job.current_runs = int(job.current_runs or 0) + 1 + if job.current_runs >= int(job.max_runs or 0): + job.status = CronJobStatus.COMPLETED.value + job.next_run_at = None + else: + job.next_run_at = calculate_next_run(job.cron_expression, job.timezone) + db.add(job) + + +def enqueue_cron_job(job: CronJob) -> Dict[str, Any]: + job_type = (job.job_type or "evaluator_run").strip() + + if job_type == "evaluator_run": + from app.workers.tasks.run_cron_evaluator_job import run_cron_evaluator_job_task + + result = run_cron_evaluator_job_task.delay(str(job.id)) + return {"task": "run_cron_evaluator_job", "celery_task_id": result.id} + + logger.warning("Unknown cron job_type {} for job {}", job_type, job.id) + return {"task": "unknown", "job_type": job_type} + + +def run_evaluator_cron_job(db: Session, job: CronJob) -> Dict[str, Any]: + from app.services.evaluators.evaluator_run_service import queue_evaluator_runs + + if job.organization_id is None: + return {"error": "evaluator_run requires organization_id"} + + raw_ids = job.evaluator_ids or [] + evaluator_ids: List[UUID] = [] + for value in raw_ids: + try: + evaluator_ids.append(UUID(str(value))) + except (TypeError, ValueError): + continue + + if not evaluator_ids: + return {"error": "no evaluator_ids configured"} + + by_workspace: dict[UUID, List[UUID]] = defaultdict(list) + for evaluator_id in evaluator_ids: + evaluator = ( + db.query(Evaluator) + .filter( + Evaluator.id == evaluator_id, + Evaluator.organization_id == job.organization_id, + ) + .first() + ) + if evaluator is None: + logger.warning( + "cron evaluator_run skipped missing evaluator {} for job {}", + evaluator_id, + job.id, + ) + continue + by_workspace[evaluator.workspace_id].append(evaluator_id) + + task_ids: List[str] = [] + for workspace_id, ids in by_workspace.items(): + try: + ws_task_ids, _results = queue_evaluator_runs( + db, + job.organization_id, + workspace_id, + ids, + ) + task_ids.extend(ws_task_ids) + except Exception as exc: + logger.warning( + "cron evaluator_run failed for job {} workspace {}: {}", + job.id, + workspace_id, + exc, + ) + + return {"evaluator_tasks": len(task_ids), "celery_task_ids": task_ids[:20]} diff --git a/app/services/cron/scheduling.py b/app/services/cron/scheduling.py new file mode 100644 index 00000000..833fa392 --- /dev/null +++ b/app/services/cron/scheduling.py @@ -0,0 +1,19 @@ +"""Shared cron scheduling helpers.""" + +from __future__ import annotations + +from datetime import datetime, timezone +from typing import Optional + +import pytz +from croniter import croniter + + +def calculate_next_run(cron_expression: str, tz_name: str) -> Optional[datetime]: + try: + tz = pytz.timezone(tz_name) + now = datetime.now(tz) + cron = croniter(cron_expression, now) + return cron.get_next(datetime).astimezone(timezone.utc) + except Exception: + return None diff --git a/app/services/evaluators/call_data_transcript.py b/app/services/evaluators/call_data_transcript.py new file mode 100644 index 00000000..4c7afe2e --- /dev/null +++ b/app/services/evaluators/call_data_transcript.py @@ -0,0 +1,195 @@ +"""Extract transcript text from voice-provider call payloads.""" + +from __future__ import annotations + +from typing import Any, Dict, List, Tuple + + +def extract_transcript_from_call_data( + call_data: Dict[str, Any], + provider_platform: str, +) -> Tuple[str, List[dict]]: + """Return plain-text transcript and speaker segments from provider call_data.""" + transcript_text = "" + speaker_segments: List[dict] = [] + + if not call_data: + return transcript_text, speaker_segments + + provider_platform_lower = provider_platform.lower() if provider_platform else "" + + if provider_platform_lower == "vapi": + transcript_text = call_data.get("transcript", "") or "" + transcript_object = call_data.get("transcript_object", []) + if not transcript_object: + artifact = call_data.get("artifact", {}) if isinstance(call_data, dict) else {} + messages = call_data.get("messages", []) or artifact.get("messages", []) + for msg in messages: + role = msg.get("role", "unknown") + content = msg.get("message", "") or msg.get("content", "") + if not content or role == "system": + continue + if role in ("bot", "assistant"): + normalized_role = "agent" + elif role == "user": + normalized_role = "user" + else: + continue + speaker_segments.append( + { + "speaker": "Agent" if normalized_role == "agent" else "User", + "text": content, + "start": msg.get("secondsFromStart", 0), + "end": msg.get("secondsFromStart", 0) + + (msg.get("duration", 0) / 1000), + } + ) + else: + for entry in transcript_object: + role = entry.get("role", "unknown") + content = entry.get("content", "") + if not content: + continue + speaker_segments.append( + { + "speaker": "Agent" if role == "agent" else "User", + "text": content, + "start": entry.get("seconds_from_start", 0), + "end": entry.get("seconds_from_start", 0) + + (entry.get("duration_ms", 0) / 1000), + } + ) + if not transcript_text and speaker_segments: + transcript_text = "\n".join( + f"{seg['speaker']}: {seg['text']}" for seg in speaker_segments + ) + + elif provider_platform_lower == "elevenlabs": + raw_transcript = call_data.get("transcript") + transcript_obj = call_data.get("transcript_object", []) + if isinstance(raw_transcript, str) and raw_transcript: + transcript_text = raw_transcript + if isinstance(transcript_obj, list): + for seg in transcript_obj: + speaker_segments.append( + { + "speaker": seg.get("speaker", "Unknown"), + "text": seg.get("text", ""), + "start": seg.get("start", 0), + "end": seg.get("end", 0), + } + ) + elif isinstance(raw_transcript, list): + for entry in raw_transcript: + role = entry.get("role", "unknown") + content = entry.get("message", "") or entry.get("text", "") + if not content: + continue + speaker = "Agent" if role in ("agent", "assistant", "ai") else "User" + speaker_segments.append( + { + "speaker": speaker, + "text": content, + "start": entry.get("time_in_call_secs", 0) or entry.get("start", 0), + "end": entry.get("time_in_call_secs", 0) or entry.get("end", 0), + } + ) + transcript_text = "\n".join( + f"{seg['speaker']}: {seg['text']}" for seg in speaker_segments + ) + + elif provider_platform_lower == "smallest": + transcript_raw = call_data.get("transcript") + transcript_object = call_data.get("transcript_object", []) + if isinstance(transcript_object, list) and transcript_object: + for entry in transcript_object: + if not isinstance(entry, dict): + continue + text = entry.get("text", "") + if not text: + continue + speaker_segments.append( + { + "speaker": entry.get("speaker", "Unknown"), + "text": text, + "start": entry.get("start", 0), + "end": entry.get("end", entry.get("start", 0)), + } + ) + if not transcript_text: + transcript_text = "\n".join( + f"{seg['speaker']}: {seg['text']}" for seg in speaker_segments + ) + elif isinstance(transcript_raw, list): + for entry in transcript_raw: + if not isinstance(entry, dict): + continue + role = str(entry.get("speaker") or entry.get("role") or "").lower() + speaker = "Agent" if role in ("agent", "assistant", "ai", "bot") else "User" + text = entry.get("text", "") or entry.get("message", "") or entry.get("content", "") + if not text: + continue + ts = entry.get("timeInCallSecs", 0) or entry.get("start", 0) or entry.get("timestamp", 0) + speaker_segments.append( + { + "speaker": speaker, + "text": text, + "start": ts, + "end": entry.get("end", ts), + } + ) + transcript_text = "\n".join( + f"{seg['speaker']}: {seg['text']}" for seg in speaker_segments + ) + elif isinstance(transcript_raw, str): + transcript_text = transcript_raw + + elif provider_platform_lower == "retell": + transcript_raw = call_data.get("transcript", "") + if isinstance(transcript_raw, str): + transcript_text = transcript_raw + lines = transcript_raw.split("\n") if transcript_raw else [] + for line in lines: + line = line.strip() + if not line: + continue + if line.startswith("Agent:") or line.startswith("agent:"): + speaker_segments.append( + { + "speaker": "Agent", + "text": line.split(":", 1)[1].strip() if ":" in line else line, + "start": 0, + "end": 0, + } + ) + elif line.startswith("User:") or line.startswith("user:"): + speaker_segments.append( + { + "speaker": "User", + "text": line.split(":", 1)[1].strip() if ":" in line else line, + "start": 0, + "end": 0, + } + ) + elif isinstance(transcript_raw, list): + for item in transcript_raw: + if not isinstance(item, dict): + continue + role = item.get("role", "") + content = item.get("content", "") or item.get("text", "") + if not content: + continue + speaker = "Agent" if role in ["agent", "assistant", "bot"] else "User" + speaker_segments.append( + { + "speaker": speaker, + "text": content, + "start": item.get("start_time", 0) or item.get("timestamp", 0), + "end": item.get("end_time", 0), + } + ) + transcript_text = "\n".join( + f"{seg['speaker']}: {seg['text']}" for seg in speaker_segments + ) + + return transcript_text, speaker_segments diff --git a/app/services/judge_alignment/gepa_bridge.py b/app/services/judge_alignment/gepa_bridge.py index 37926fa1..6a12506b 100644 --- a/app/services/judge_alignment/gepa_bridge.py +++ b/app/services/judge_alignment/gepa_bridge.py @@ -219,6 +219,31 @@ def execute_judge_gepa(run_id: str, db: Session) -> Dict[str, Any]: if not evaluator: raise RuntimeError("Evaluator vanished between dispatch and execution") + from app.services.usage.context import ( + reset_usage_context, + set_usage_context, + usage_context_for_prompt_optimization_run, + ) + + usage_token = set_usage_context(usage_context_for_prompt_optimization_run(run)) + try: + return _execute_judge_gepa_with_context( + run=run, + cfg=cfg, + evaluator=evaluator, + db=db, + ) + finally: + reset_usage_context(usage_token) + + +def _execute_judge_gepa_with_context( + *, + run: PromptOptimizationRun, + cfg: Dict[str, Any], + evaluator: Evaluator, + db: Session, +) -> Dict[str, Any]: dev_ids: List[str] = cfg.get("dev_sample_ids", []) if not dev_ids: raise RuntimeError("Optimisation run has no dev_sample_ids in config") @@ -370,19 +395,27 @@ def reflection_lm(prompt: str) -> str: resp = litellm_completion(**reflection_kwargs, credential=credential_ctx) return resp.choices[0].message.content + from app.services.ai.llm_gateway import litellm_batch_completion_recording + run.status = PromptOptimizationStatus.RUNNING.value db.commit() try: - result = gepa_optimize( - seed_candidate={"system_prompt": evaluator.custom_prompt}, - trainset=trainset, - adapter=adapter, - reflection_lm=reflection_lm, - max_metric_calls=int(cfg.get("max_metric_calls", 20)), - reflection_minibatch_size=1, - candidate_selection_strategy="pareto", - ) + with litellm_batch_completion_recording( + organization_id=run.organization_id, + db=db, + model=lm_identifier, + credential=credential_ctx, + ): + result = gepa_optimize( + seed_candidate={"system_prompt": evaluator.custom_prompt}, + trainset=trainset, + adapter=adapter, + reflection_lm=reflection_lm, + max_metric_calls=int(cfg.get("max_metric_calls", 20)), + reflection_minibatch_size=1, + candidate_selection_strategy="pareto", + ) except Exception as exc: run.status = PromptOptimizationStatus.FAILED.value run.error_message = str(exc) diff --git a/app/services/optimization/gepa_service.py b/app/services/optimization/gepa_service.py index 5328dba9..d92895ac 100644 --- a/app/services/optimization/gepa_service.py +++ b/app/services/optimization/gepa_service.py @@ -149,15 +149,23 @@ def reflection_lm(prompt: str) -> str: resp = litellm_completion(**reflection_kwargs, credential=credential_ctx) return resp.choices[0].message.content - result = gepa_optimize( - seed_candidate={"system_prompt": seed_prompt}, - trainset=trainset, - adapter=adapter, - reflection_lm=reflection_lm, - max_metric_calls=max_metric_calls, - reflection_minibatch_size=min(minibatch_size, len(trainset)), - candidate_selection_strategy="pareto", - ) + from app.services.ai.llm_gateway import litellm_batch_completion_recording + + with litellm_batch_completion_recording( + organization_id=organization_id, + db=db, + model=lm_identifier, + credential=credential_ctx, + ): + result = gepa_optimize( + seed_candidate={"system_prompt": seed_prompt}, + trainset=trainset, + adapter=adapter, + reflection_lm=reflection_lm, + max_metric_calls=max_metric_calls, + reflection_minibatch_size=min(minibatch_size, len(trainset)), + candidate_selection_strategy="pareto", + ) return _format_result(result, seed_prompt) diff --git a/app/services/testing/test_agent_bridge_service.py b/app/services/testing/test_agent_bridge_service.py index f8e8d41a..4a05cdb0 100644 --- a/app/services/testing/test_agent_bridge_service.py +++ b/app/services/testing/test_agent_bridge_service.py @@ -613,6 +613,8 @@ def resolve_api_key_for_provider( ) test_agent_config = TestAgentConfig( + organization_id=organization_id, + workspace_id=getattr(agent, "workspace_id", None), agent_name=agent.name or "Voice AI Agent", agent_description=agent.description or "A voice AI assistant", test_agent_simulation_prompt=simulation_prompt, diff --git a/app/services/usage/__init__.py b/app/services/usage/__init__.py new file mode 100644 index 00000000..4ebc75ca --- /dev/null +++ b/app/services/usage/__init__.py @@ -0,0 +1,50 @@ +"""LLM/STT/TTS usage tracking with Redis buffer, PG fallback, and catalog rollups.""" + +from app.services.usage.context import ( + LLMUsageContext, + LLMUsageProductSection, + ensure_usage_context, + infer_product_section_from_path, + llm_usage_context, + reset_usage_context, + reset_usage_hints, + set_usage_context, + set_usage_hints, + usage_context_for_judge_run, + usage_context_for_prompt_optimization_run, + usage_context_for_prompt_partial, +) +from app.services.usage.llm_usage import ( + flush_all_usage_to_catalog, + flush_usage_to_catalog, + probe_audio_seconds, + record_call_usage, + record_llm_usage, + record_stt_usage, + record_tts_usage, +) +from app.services.usage.normalize import UsageSnapshot, normalize_llm_usage + +__all__ = [ + "LLMUsageContext", + "LLMUsageProductSection", + "usage_context_for_judge_run", + "usage_context_for_prompt_optimization_run", + "usage_context_for_prompt_partial", + "UsageSnapshot", + "ensure_usage_context", + "flush_all_usage_to_catalog", + "flush_usage_to_catalog", + "infer_product_section_from_path", + "llm_usage_context", + "normalize_llm_usage", + "probe_audio_seconds", + "record_call_usage", + "record_llm_usage", + "record_stt_usage", + "record_tts_usage", + "reset_usage_context", + "reset_usage_hints", + "set_usage_context", + "set_usage_hints", +] diff --git a/app/services/usage/access.py b/app/services/usage/access.py new file mode 100644 index 00000000..9591bf2d --- /dev/null +++ b/app/services/usage/access.py @@ -0,0 +1,83 @@ +"""Usage read access policy — date window clamping and SQL floor for OSS tier.""" + +from __future__ import annotations + +from dataclasses import dataclass +from datetime import date, timedelta +from typing import Optional +from uuid import UUID + +from loguru import logger + +from app.core.usage_entitlement import ( + OSS_USAGE_HISTORY_DAYS, + UsagePolicySnapshot, + get_usage_policy, + has_enterprise_entitlement, +) +from app.services.usage.dates import usage_date_filter_bounds, usage_local_today + + +@dataclass(frozen=True) +class UsageAccessResult: + display_start: date + display_end: date + filter_start: date + filter_end: date + enforced_filter_floor: Optional[date] + policy: UsagePolicySnapshot + range_clamped: bool + + +def oss_usage_min_local_date(tz: Optional[str]) -> date: + """Earliest inclusive local calendar day allowed for OSS tier.""" + today = usage_local_today(tz) + return today - timedelta(days=OSS_USAGE_HISTORY_DAYS - 1) + + +class UsageAccessPolicy: + @staticmethod + def resolve( + organization_id: UUID, + start: Optional[date], + end: Optional[date], + tz: Optional[str], + ) -> UsageAccessResult: + today = usage_local_today(tz) + display_start = start or today + display_end = end or today + policy = get_usage_policy(organization_id) + range_clamped = False + enforced_floor: Optional[date] = None + + if not has_enterprise_entitlement(organization_id): + oss_min = oss_usage_min_local_date(tz) + if display_start < oss_min: + display_start = oss_min + range_clamped = True + enforced_floor = usage_date_filter_bounds(oss_min, oss_min, tz)[0] + + filter_start, filter_end = usage_date_filter_bounds( + display_start, display_end, tz + ) + + if enforced_floor is not None and filter_start < enforced_floor: + filter_start = enforced_floor + + if range_clamped: + logger.info( + "usage_history_clamped org_id={} effective_start={} effective_end={}", + organization_id, + display_start, + display_end, + ) + + return UsageAccessResult( + display_start=display_start, + display_end=display_end, + filter_start=filter_start, + filter_end=filter_end, + enforced_filter_floor=enforced_floor, + policy=policy, + range_clamped=range_clamped, + ) diff --git a/app/services/usage/bucket_context.py b/app/services/usage/bucket_context.py new file mode 100644 index 00000000..64187a39 --- /dev/null +++ b/app/services/usage/bucket_context.py @@ -0,0 +1,115 @@ +"""Canonical usage bucket context (JSONB) helpers.""" + +from __future__ import annotations + +import json +from typing import Any, Dict, Optional +from uuid import UUID + +_NONE = "__none__" + +# Stable keys stored in llm_usage_daily.context (extend without new columns). +# evaluation_row_id / call_import_row_id may appear on legacy rows only. +KNOWN_CONTEXT_KEYS = frozenset( + { + "resource_id", + "resource_type", + "call_import_id", + "evaluation_id", + "evaluation_row_id", + "call_import_row_id", + "credential_id", + "agent_id", + "job_id", + "user_id", + "trace_id", + } +) + +# Never store roll-up counters in JSONB — they stay as BIGINT columns for SUM(). +_FORBIDDEN_CONTEXT_KEYS = frozenset( + { + "prompt_tokens", + "completion_tokens", + "cache_read_tokens", + "cache_creation_tokens", + "reasoning_tokens", + "audio_seconds", + "tts_characters", + "call_count", + "total_tokens", + } +) + + +def build_bucket_context( + *, + resource_id: Optional[UUID] = None, + resource_type: Optional[str] = None, + extra: Optional[Dict[str, Any]] = None, +) -> Dict[str, str]: + """Build a normalized string-keyed context dict for rollup buckets.""" + ctx: Dict[str, str] = {} + if resource_id is not None: + ctx["resource_id"] = str(resource_id) + if resource_type: + ctx["resource_type"] = resource_type + if extra: + for key, value in extra.items(): + if value is None or key in ctx: + continue + if key in _FORBIDDEN_CONTEXT_KEYS: + continue + ctx[key] = str(value) + return ctx + + +def context_bucket_token(context: Optional[Dict[str, Any]]) -> str: + """Stable Redis bucket token for a context dict.""" + if not context: + return _NONE + normalized = {k: str(v) for k, v in sorted(context.items()) if v is not None} + if not normalized: + return _NONE + return json.dumps(normalized, sort_keys=True, separators=(",", ":")) + + +def parse_context_bucket_token(token: str) -> Dict[str, str]: + if not token or token == _NONE: + return {} + try: + parsed = json.loads(token) + except json.JSONDecodeError: + return {} + if not isinstance(parsed, dict): + return {} + return {str(k): str(v) for k, v in parsed.items() if v is not None} + + +def legacy_resource_context( + resource_id: Optional[UUID], + resource_type: Optional[str], +) -> Dict[str, str]: + return build_bucket_context( + resource_id=resource_id, + resource_type=resource_type, + ) + + +def resource_id_from_context(context: Optional[Dict[str, Any]]) -> Optional[UUID]: + if not context: + return None + raw = context.get("resource_id") + if not raw: + return None + try: + return UUID(str(raw)) + except (ValueError, TypeError): + return None + + +def resource_type_from_context(context: Optional[Dict[str, Any]]) -> Optional[str]: + if not context: + return None + raw = context.get("resource_type") + return str(raw) if raw else None diff --git a/app/services/usage/call_import_context.py b/app/services/usage/call_import_context.py new file mode 100644 index 00000000..1ceeb723 --- /dev/null +++ b/app/services/usage/call_import_context.py @@ -0,0 +1,187 @@ +"""Usage context builders for call-import evaluation pipelines.""" + +from __future__ import annotations + +from typing import Any, Optional +from uuid import UUID + +from app.services.usage.context import LLMUsageContext, LLMUsageProductSection + + +def _parse_uuid(raw: Any) -> Optional[UUID]: + try: + return UUID(str(raw)) + except (TypeError, ValueError): + return None + + +def call_import_ids_from_usage_context(ctx: LLMUsageContext) -> dict[str, Optional[UUID]]: + """Extract call-import linkage ids from a usage context (no DB).""" + extra = ctx.extra or {} + ids: dict[str, Optional[UUID]] = { + "call_import_id": _parse_uuid(extra.get("call_import_id")), + "evaluation_id": _parse_uuid(extra.get("evaluation_id")), + "call_import_row_id": _parse_uuid(extra.get("call_import_row_id")), + "evaluation_row_id": _parse_uuid(extra.get("evaluation_row_id")), + } + if ctx.resource_type == "call_import" and ctx.resource_id: + ids["call_import_id"] = ids["call_import_id"] or ctx.resource_id + if ctx.resource_type == "call_import_evaluation" and ctx.resource_id: + ids["evaluation_id"] = ids["evaluation_id"] or ctx.resource_id + return ids + + +def resolve_workspace_id_for_usage_context(ctx: LLMUsageContext) -> Optional[UUID]: + """Resolve workspace from call-import entities when context omitted workspace_id.""" + if ctx.workspace_id is not None: + return ctx.workspace_id + + ids = call_import_ids_from_usage_context(ctx) + if not any(ids.values()): + return None + + from app.database import SessionLocal + from app.models.database import ( + CallImport, + CallImportEvaluation, + CallImportEvaluationRow, + CallImportRow, + ) + + org_id = ctx.organization_id + db = SessionLocal() + try: + if ids["call_import_id"]: + ws = ( + db.query(CallImport.workspace_id) + .filter( + CallImport.id == ids["call_import_id"], + CallImport.organization_id == org_id, + ) + .scalar() + ) + if ws: + return ws + + if ids["evaluation_id"]: + ws = ( + db.query(CallImportEvaluation.workspace_id) + .filter( + CallImportEvaluation.id == ids["evaluation_id"], + CallImportEvaluation.organization_id == org_id, + ) + .scalar() + ) + if ws: + return ws + + if ids["call_import_row_id"]: + ws = ( + db.query(CallImportRow.workspace_id) + .filter( + CallImportRow.id == ids["call_import_row_id"], + CallImportRow.organization_id == org_id, + ) + .scalar() + ) + if ws: + return ws + + if ids["evaluation_row_id"]: + ws = ( + db.query(CallImportEvaluation.workspace_id) + .join( + CallImportEvaluationRow, + CallImportEvaluationRow.evaluation_id == CallImportEvaluation.id, + ) + .filter( + CallImportEvaluationRow.id == ids["evaluation_row_id"], + CallImportEvaluation.organization_id == org_id, + ) + .scalar() + ) + if ws: + return ws + + return None + finally: + db.close() + + +def enrich_usage_context_workspace(ctx: LLMUsageContext) -> LLMUsageContext: + """Fill workspace_id from call-import linkage when missing at record time.""" + if ctx.workspace_id is not None: + return ctx + try: + resolved = resolve_workspace_id_for_usage_context(ctx) + except Exception: + return ctx + if resolved is None: + return ctx + return LLMUsageContext( + organization_id=ctx.organization_id, + workspace_id=resolved, + product_section=ctx.product_section, + resource_id=ctx.resource_id, + resource_type=ctx.resource_type, + extra=ctx.extra, + ) + + +def call_import_evaluation_usage_context( + *, + organization_id: UUID, + workspace_id: Optional[UUID], + evaluation_id: UUID, + call_import_id: UUID, +) -> LLMUsageContext: + """Usage rollup for an eval run (model-level drilldown; not per recording).""" + extra: dict[str, str] = { + "call_import_id": str(call_import_id), + "evaluation_id": str(evaluation_id), + } + return LLMUsageContext( + organization_id=organization_id, + workspace_id=workspace_id, + product_section=LLMUsageProductSection.CALL_IMPORT_EVALUATIONS, + resource_id=evaluation_id, + resource_type="call_import_evaluation", + extra=extra, + ) + + +def call_import_row_usage_context( + *, + organization_id: UUID, + workspace_id: Optional[UUID], + call_import_id: UUID, + evaluation_id: Optional[UUID] = None, +) -> LLMUsageContext: + """Usage rollup for diarisation / STT (call-import level, not per recording).""" + if evaluation_id is not None: + return call_import_evaluation_usage_context( + organization_id=organization_id, + workspace_id=workspace_id, + evaluation_id=evaluation_id, + call_import_id=call_import_id, + ) + return LLMUsageContext( + organization_id=organization_id, + workspace_id=workspace_id, + product_section=LLMUsageProductSection.CALL_IMPORTS, + resource_id=call_import_id, + resource_type="call_import", + extra={"call_import_id": str(call_import_id)}, + ) + + +def usage_context_for_evaluation( + evaluation: Any, +) -> LLMUsageContext: + """Usage context from a loaded CallImportEvaluation row.""" + return call_import_evaluation_usage_context( + organization_id=evaluation.organization_id, + workspace_id=evaluation.workspace_id, + evaluation_id=evaluation.id, + call_import_id=evaluation.call_import_id, + ) diff --git a/app/services/usage/context.py b/app/services/usage/context.py new file mode 100644 index 00000000..8bd895e1 --- /dev/null +++ b/app/services/usage/context.py @@ -0,0 +1,286 @@ +"""Usage attribution context (org, workspace, product section, resource).""" + +from __future__ import annotations + +from contextlib import contextmanager +from contextvars import ContextVar, Token +from dataclasses import dataclass +from enum import Enum +from typing import Iterator, Optional, Any +from uuid import UUID + + +class LLMUsageProductSection(str, Enum): + CALL_IMPORT_EVALUATIONS = "call_import_evaluations" + CALL_IMPORTS = "call_imports" + PLAYGROUND = "playground" + VOICE_PLAYGROUND = "voice_playground" + EVALUATORS = "evaluators" + METRICS = "metrics" + CHAT = "chat" + JUDGE_ALIGNMENT = "judge_alignment" + PROMPT_OPTIMIZATION = "prompt_optimization" + PERSONAS = "personas" + AGENTS = "agents" + PROMPT_PARTIALS = "prompt_partials" + CONVERSATION_EVALUATIONS = "conversation_evaluations" + TELEPHONY = "telephony" + TEST_AGENT = "test_agent" + OTHER = "other" + + +@dataclass(frozen=True) +class LLMUsageContext: + organization_id: UUID + workspace_id: Optional[UUID] = None + product_section: LLMUsageProductSection = LLMUsageProductSection.OTHER + resource_id: Optional[UUID] = None + resource_type: Optional[str] = None + extra: Optional[dict[str, str]] = None + + +_usage_context_var: ContextVar[Optional[LLMUsageContext]] = ContextVar( + "llm_usage_context", default=None +) +_usage_workspace_hint_var: ContextVar[Optional[UUID]] = ContextVar( + "llm_usage_workspace_hint", default=None +) +_usage_section_hint_var: ContextVar[LLMUsageProductSection] = ContextVar( + "llm_usage_section_hint", default=LLMUsageProductSection.OTHER +) + +# Path fragments under /api/v1 → product section (longest match wins). +_PATH_SECTION_RULES: tuple[tuple[str, LLMUsageProductSection], ...] = ( + ("call-import-evaluations", LLMUsageProductSection.CALL_IMPORT_EVALUATIONS), + ("call-imports", LLMUsageProductSection.CALL_IMPORTS), + ("voice-playground", LLMUsageProductSection.VOICE_PLAYGROUND), + ("playground", LLMUsageProductSection.PLAYGROUND), + ("evaluators", LLMUsageProductSection.EVALUATORS), + ("metrics", LLMUsageProductSection.METRICS), + ("chat", LLMUsageProductSection.CHAT), + ("judge-alignment", LLMUsageProductSection.JUDGE_ALIGNMENT), + ("prompt-optimization", LLMUsageProductSection.PROMPT_OPTIMIZATION), + ("personas", LLMUsageProductSection.PERSONAS), + ("agents", LLMUsageProductSection.AGENTS), + ("prompt-partials", LLMUsageProductSection.PROMPT_PARTIALS), + ("conversation-evaluations", LLMUsageProductSection.CONVERSATION_EVALUATIONS), + ("telephony", LLMUsageProductSection.TELEPHONY), + ("test-agent", LLMUsageProductSection.TEST_AGENT), +) + + +def get_usage_context() -> Optional[LLMUsageContext]: + return _usage_context_var.get() + + +def set_usage_context(ctx: Optional[LLMUsageContext]) -> Token: + return _usage_context_var.set(ctx) + + +def reset_usage_context(token: Token) -> None: + _usage_context_var.reset(token) + + +def set_usage_hints( + *, + workspace_id: Optional[UUID] = None, + product_section: Optional[LLMUsageProductSection] = None, +) -> tuple[Token, Token]: + ws_token = _usage_workspace_hint_var.set(workspace_id) + section = product_section or LLMUsageProductSection.OTHER + section_token = _usage_section_hint_var.set(section) + return ws_token, section_token + + +def reset_usage_hints(tokens: tuple[Token, Token]) -> None: + _usage_workspace_hint_var.reset(tokens[0]) + _usage_section_hint_var.reset(tokens[1]) + + +def get_usage_workspace_hint() -> Optional[UUID]: + return _usage_workspace_hint_var.get() + + +def get_usage_section_hint() -> LLMUsageProductSection: + return _usage_section_hint_var.get() + + +@contextmanager +def llm_usage_context(ctx: LLMUsageContext) -> Iterator[None]: + token = set_usage_context(ctx) + try: + yield + finally: + reset_usage_context(token) + + +def infer_product_section_from_path(path: str) -> LLMUsageProductSection: + normalized = (path or "").lower() + for fragment, section in _PATH_SECTION_RULES: + if f"/{fragment}" in normalized or normalized.endswith(fragment): + return section + return LLMUsageProductSection.OTHER + + +def ensure_usage_context( + organization_id: UUID, + *, + workspace_id: Optional[UUID] = None, + product_section: LLMUsageProductSection = LLMUsageProductSection.OTHER, + resource_id: Optional[UUID] = None, + resource_type: Optional[str] = None, + extra: Optional[dict[str, str]] = None, +) -> Token | None: + """Set or enrich usage context. Returns token to reset, or None if unchanged.""" + resolved_workspace = workspace_id or get_usage_workspace_hint() + resolved_section = product_section + if resolved_section == LLMUsageProductSection.OTHER: + hint = get_usage_section_hint() + if hint != LLMUsageProductSection.OTHER: + resolved_section = hint + + current = get_usage_context() + if current is None: + return set_usage_context( + LLMUsageContext( + organization_id=organization_id, + workspace_id=resolved_workspace, + product_section=resolved_section, + resource_id=resource_id, + resource_type=resource_type, + extra=extra, + ) + ) + + upgraded_workspace = current.workspace_id or resolved_workspace + upgraded_section = ( + resolved_section + if current.product_section == LLMUsageProductSection.OTHER + and resolved_section != LLMUsageProductSection.OTHER + else current.product_section + ) + upgraded_resource_id = current.resource_id or resource_id + upgraded_resource_type = current.resource_type or resource_type + upgraded_extra = {**(current.extra or {}), **(extra or {})} or None + if ( + upgraded_workspace == current.workspace_id + and upgraded_section == current.product_section + and upgraded_resource_id == current.resource_id + and upgraded_resource_type == current.resource_type + and upgraded_extra == current.extra + ): + return None + + return set_usage_context( + LLMUsageContext( + organization_id=current.organization_id, + workspace_id=upgraded_workspace, + product_section=upgraded_section, + resource_id=upgraded_resource_id, + resource_type=upgraded_resource_type, + extra=upgraded_extra, + ) + ) + + +def usage_context_for_agent( + agent: Any, + *, + workspace_id: Optional[UUID] = None, + extra: Optional[dict[str, str]] = None, +) -> LLMUsageContext: + """Usage context for agent-scoped LLM work (simulations, setup, summaries).""" + merged: dict[str, str] = dict(extra or {}) + merged.setdefault("agent_id", str(agent.id)) + short = getattr(agent, "agent_id", None) + if short: + merged.setdefault("agent_short_id", str(short)) + return LLMUsageContext( + organization_id=agent.organization_id, + workspace_id=workspace_id or agent.workspace_id, + product_section=LLMUsageProductSection.AGENTS, + resource_id=agent.id, + resource_type="agent", + extra=merged, + ) + + +def usage_context_for_evaluator_result(result: Any) -> LLMUsageContext: + """Usage context for processing an evaluator result (Vapi / playground runs).""" + if result.agent_id: + return LLMUsageContext( + organization_id=result.organization_id, + workspace_id=result.workspace_id, + product_section=LLMUsageProductSection.AGENTS, + resource_id=result.agent_id, + resource_type="agent", + extra={"agent_id": str(result.agent_id)}, + ) + extra: dict[str, str] = {"evaluator_result_id": str(result.id)} + if getattr(result, "result_id", None): + extra["result_short_id"] = str(result.result_id) + if result.evaluator_id: + extra["evaluator_id"] = str(result.evaluator_id) + return LLMUsageContext( + organization_id=result.organization_id, + workspace_id=result.workspace_id, + product_section=LLMUsageProductSection.EVALUATORS, + resource_id=result.id, + resource_type="evaluator_result", + extra=extra, + ) + + +def usage_context_for_prompt_optimization_run(run: Any) -> LLMUsageContext: + """Usage context for a GEPA prompt optimization run.""" + cfg = run.config if isinstance(run.config, dict) else {} + is_judge = cfg.get("source") == "judge_alignment" + extra: dict[str, str] = { + "optimization_run_id": str(run.id), + "agent_id": str(run.agent_id), + } + if run.evaluator_id: + extra["evaluator_id"] = str(run.evaluator_id) + if is_judge: + extra["source"] = "judge_alignment" + if cfg.get("judge_dataset_id"): + extra["judge_dataset_id"] = str(cfg["judge_dataset_id"]) + return LLMUsageContext( + organization_id=run.organization_id, + workspace_id=run.workspace_id, + product_section=( + LLMUsageProductSection.JUDGE_ALIGNMENT + if is_judge + else LLMUsageProductSection.PROMPT_OPTIMIZATION + ), + resource_id=run.agent_id, + resource_type="agent", + extra=extra, + ) + + +def usage_context_for_judge_run(run: Any) -> LLMUsageContext: + """Usage context for a judge alignment scoring run.""" + return LLMUsageContext( + organization_id=run.organization_id, + workspace_id=run.workspace_id, + product_section=LLMUsageProductSection.JUDGE_ALIGNMENT, + resource_id=run.evaluator_id, + resource_type="evaluator", + extra={ + "judge_run_id": str(run.id), + "judge_dataset_id": str(run.dataset_id), + }, + ) + + +def usage_context_for_prompt_partial(partial: Any) -> LLMUsageContext: + """Usage context for prompt partial / agent flowchart LLM work.""" + return LLMUsageContext( + organization_id=partial.organization_id, + workspace_id=partial.workspace_id, + product_section=LLMUsageProductSection.PROMPT_PARTIALS, + resource_id=partial.id, + resource_type="prompt_partial", + extra={"prompt_partial_id": str(partial.id)}, + ) diff --git a/app/services/usage/dates.py b/app/services/usage/dates.py new file mode 100644 index 00000000..340ba9eb --- /dev/null +++ b/app/services/usage/dates.py @@ -0,0 +1,38 @@ +"""Local calendar dates vs UTC usage_date bucket bounds.""" + +from __future__ import annotations + +from datetime import date, datetime, time, timedelta, timezone +from zoneinfo import ZoneInfo, ZoneInfoNotFoundError + + +def resolve_usage_timezone(tz: str | None) -> ZoneInfo: + if not tz: + return ZoneInfo("UTC") + try: + return ZoneInfo(tz) + except ZoneInfoNotFoundError: + return ZoneInfo("UTC") + + +def usage_local_today(tz: str | None) -> date: + return datetime.now(resolve_usage_timezone(tz)).date() + + +def usage_date_filter_bounds( + start: date, + end: date, + tz: str | None, +) -> tuple[date, date]: + """Map inclusive local calendar days to usage_date (UTC-day) filter bounds.""" + if not tz: + return start, end + + zone = resolve_usage_timezone(tz) + start_local = datetime.combine(start, time.min, tzinfo=zone) + end_exclusive = datetime.combine(end + timedelta(days=1), time.min, tzinfo=zone) + filter_start = start_local.astimezone(timezone.utc).date() + filter_end = ( + end_exclusive - timedelta(seconds=1) + ).astimezone(timezone.utc).date() + return filter_start, filter_end diff --git a/app/services/usage/enabled_models.py b/app/services/usage/enabled_models.py new file mode 100644 index 00000000..728356c5 --- /dev/null +++ b/app/services/usage/enabled_models.py @@ -0,0 +1,132 @@ +"""Org- and credential-level enabled model allowlists.""" + +from __future__ import annotations + +from typing import Iterable, List, Optional, Set +from uuid import UUID + +from sqlalchemy import text +from sqlalchemy.orm import Session + +from app.models.database import AIProvider, ModelProvider +from app.services.ai.model_config_service import ModelConfigService + + +def normalize_enabled_models(raw: Optional[Iterable[str]]) -> Optional[List[str]]: + if raw is None: + return None + seen: set[str] = set() + normalized: list[str] = [] + for item in raw: + if item is None: + continue + name = str(item).strip() + if not name or name in seen: + continue + seen.add(name) + normalized.append(name) + return normalized or None + + +def catalog_models_for_provider(provider: str) -> List[str]: + service = ModelConfigService() + try: + provider_enum = ModelProvider(provider.lower()) + except ValueError: + return [] + options = service.get_model_options_by_provider(provider_enum) + models: list[str] = [] + for key in ("llm", "stt", "tts", "s2s"): + models.extend(options.get(key) or []) + return sorted({m for m in models if m}) + + +def effective_enabled_models_for_credential(credential: AIProvider) -> Optional[List[str]]: + """Return explicit allowlist, or None meaning unrestricted (full provider catalog).""" + return normalize_enabled_models(credential.enabled_models) + + +def filter_models_by_credential( + credential: Optional[AIProvider], + catalog_models: List[str], +) -> List[str]: + if credential is None: + return catalog_models + allowlist = effective_enabled_models_for_credential(credential) + if not allowlist: + return catalog_models + allowed = set(allowlist) + filtered = [m for m in catalog_models if m in allowed] + if not filtered: + filtered = list(allowlist) + gateway = (credential.gateway_model or "").strip() + if gateway and gateway not in filtered: + filtered = [gateway, *filtered] + return filtered + + +def _usage_models_for_org(db: Session, organization_id: UUID) -> Set[str]: + rows = db.execute( + text( + """ + SELECT DISTINCT model + FROM llm_usage_daily + WHERE organization_id = CAST(:organization_id AS uuid) + AND model IS NOT NULL + AND model <> '' + """ + ), + {"organization_id": str(organization_id)}, + ).scalars().all() + return {str(row).strip() for row in rows if row} + + +def _override_models_for_org(db: Session, organization_id: UUID) -> Set[str]: + rows = db.execute( + text( + """ + SELECT DISTINCT model + FROM org_model_pricing_overrides + WHERE organization_id = CAST(:organization_id AS uuid) + """ + ), + {"organization_id": str(organization_id)}, + ).scalars().all() + return {str(row).strip() for row in rows if row} + + +def org_pricing_eligible_models(db: Session, organization_id: UUID) -> List[str]: + """Models org admins may set pricing overrides for.""" + models: set[str] = set() + providers = ( + db.query(AIProvider) + .filter( + AIProvider.organization_id == organization_id, + AIProvider.is_active.is_(True), + ) + .all() + ) + any_explicit_allowlist = False + for provider in providers: + gateway = (provider.gateway_model or "").strip() + if gateway: + models.add(gateway) + allowlist = effective_enabled_models_for_credential(provider) + if allowlist: + any_explicit_allowlist = True + models.update(allowlist) + else: + models.update(catalog_models_for_provider(provider.provider)) + + if not any_explicit_allowlist and not models: + for provider in providers: + models.update(catalog_models_for_provider(provider.provider)) + + models.update(_usage_models_for_org(db, organization_id)) + models.update(_override_models_for_org(db, organization_id)) + return sorted(models) + + +def org_union_enabled_models(db: Session, organization_id: UUID) -> List[str]: + """All models enabled on any active integration credential.""" + return org_pricing_eligible_models(db, organization_id) diff --git a/app/services/usage/fx_rates.py b/app/services/usage/fx_rates.py new file mode 100644 index 00000000..927026b1 --- /dev/null +++ b/app/services/usage/fx_rates.py @@ -0,0 +1,125 @@ +"""USD/INR display FX rate (Frankfurter v2, cached in Redis).""" + +from __future__ import annotations + +import json +from datetime import datetime, timezone +from typing import Any, Optional + +import httpx +import redis +from loguru import logger + +from app.config import settings + +_REDIS_KEY = "usage:fx:USD_INR" +_FRANKFURTER_URL = "https://api.frankfurter.dev/v2/rate/USD/INR" +_DEFAULT_RATE = 95.0 +_CACHE_TTL_SECONDS = 25 * 3600 +_redis: redis.Redis | None = None + + +def _client() -> redis.Redis: + global _redis + if _redis is None: + _redis = redis.from_url(settings.REDIS_URL, decode_responses=True) + return _redis + + +def _cache_payload(rate: float, as_of: datetime, source: str) -> dict[str, Any]: + return { + "base": "USD", + "quote": "INR", + "rate": rate, + "as_of": as_of.astimezone(timezone.utc).isoformat(), + "source": source, + } + + +def _read_cached() -> Optional[dict[str, Any]]: + try: + raw = _client().get(_REDIS_KEY) + if not raw: + return None + payload = json.loads(raw) + if payload.get("source") != "frankfurter": + return None + if not isinstance(payload.get("rate"), (int, float)): + return None + return payload + except (redis.RedisError, json.JSONDecodeError) as exc: + logger.warning("USD/INR FX cache read failed: {}", exc) + return None + + +def _write_cached(payload: dict[str, Any]) -> None: + try: + _client().set(_REDIS_KEY, json.dumps(payload), ex=_CACHE_TTL_SECONDS) + except redis.RedisError as exc: + logger.warning("USD/INR FX cache write failed: {}", exc) + + +def _parse_frankfurter_payload(data: Any) -> tuple[float, datetime]: + if not isinstance(data, dict): + raise ValueError(f"Frankfurter response is not an object: {type(data).__name__}") + + base = data.get("base") + quote = data.get("quote") + if base != "USD" or quote != "INR": + raise ValueError(f"Unexpected Frankfurter pair: {base}/{quote}") + + rate_raw = data.get("rate") + if not isinstance(rate_raw, (int, float)): + raise ValueError(f"Frankfurter rate missing or invalid: {rate_raw!r}") + + rate = float(rate_raw) + if rate <= 0: + raise ValueError(f"Frankfurter rate must be positive: {rate}") + + date_raw = data.get("date") + if not isinstance(date_raw, str) or not date_raw.strip(): + raise ValueError(f"Frankfurter date missing or invalid: {date_raw!r}") + + as_of = datetime.fromisoformat(date_raw).replace(tzinfo=timezone.utc) + return rate, as_of + + +def _fallback_payload(reason: str, *, response_body: Any = None) -> dict[str, Any]: + logger.error( + "USD/INR FX using hardcoded fallback rate {:.2f} — INR costs may be wrong. reason={} response={}", + _DEFAULT_RATE, + reason, + response_body, + ) + return _cache_payload(_DEFAULT_RATE, datetime.now(timezone.utc), "default") + + +def get_usd_inr_rate() -> dict[str, Any]: + cached = _read_cached() + if cached is not None: + return cached + return refresh_usd_inr_rate() + + +def refresh_usd_inr_rate() -> dict[str, Any]: + try: + with httpx.Client(timeout=15.0) as client: + response = client.get(_FRANKFURTER_URL) + response.raise_for_status() + data = response.json() + except Exception as exc: + return _fallback_payload(f"Frankfurter request failed: {exc}") + + try: + rate, as_of = _parse_frankfurter_payload(data) + except Exception as exc: + return _fallback_payload(f"Frankfurter response parse failed: {exc}", response_body=data) + + payload = _cache_payload(rate, as_of, "frankfurter") + _write_cached(payload) + logger.info( + "USD/INR FX refreshed from Frankfurter: rate={} as_of={}", + rate, + as_of.date().isoformat(), + ) + return payload diff --git a/app/services/usage/llm_usage.py b/app/services/usage/llm_usage.py new file mode 100644 index 00000000..bbeac232 --- /dev/null +++ b/app/services/usage/llm_usage.py @@ -0,0 +1,1540 @@ +"""Redis-buffered LLM/STT usage counters with catalog rollup flush.""" + +from __future__ import annotations + +import json +import math +import os +import time +import uuid +from datetime import date, datetime, timezone +from typing import Any, Dict, Iterable, List, Optional, Tuple +from uuid import UUID + +import redis +from loguru import logger +from sqlalchemy import text +from sqlalchemy.exc import IntegrityError +from sqlalchemy.orm import Session + +from app.config import settings +from app.services.usage.bucket_context import ( + build_bucket_context, + context_bucket_token, + legacy_resource_context, + parse_context_bucket_token, +) +from app.services.usage.context import ( + LLMUsageContext, + LLMUsageProductSection, + get_usage_context, +) +from app.services.usage.normalize import UsageSnapshot + +_redis: redis.Redis | None = None + +_NONE = "__none__" +_PENDING_TTL_SECONDS = 14 * 24 * 60 * 60 +_DEFAULT_FLUSH_BUCKET_BATCH_SIZE = 500 +_DEFAULT_FLUSH_MAX_BATCHES_PER_RUN = 30 +_DEFAULT_FLUSH_LOCK_TTL_SECONDS = 300 +_FLUSH_LOCK_WAIT_SECONDS = 3.0 +USAGE_KIND_LLM = "llm" +USAGE_KIND_STT = "stt" +USAGE_KIND_TTS = "tts" +_METRIC_FIELDS = ( + "prompt_tokens", + "completion_tokens", + "cache_read_tokens", + "cache_creation_tokens", + "reasoning_tokens", + "audio_seconds", + "tts_characters", + "call_count", +) + +_CLAIM_LUA = """ +if redis.call('EXISTS', KEYS[1]) == 0 then + return 0 +end +redis.call('RENAME', KEYS[1], KEYS[2]) +return 1 +""" + + +def _env_int(name: str, default: int, *, minimum: int = 1) -> int: + raw = os.environ.get(name) + if raw is None: + return default + try: + return max(minimum, int(raw)) + except (TypeError, ValueError): + logger.warning("invalid {}={!r}; using default {}", name, raw, default) + return default + + +def flush_bucket_batch_size() -> int: + """Max rollup buckets persisted per DB transaction during Redis flush.""" + return _env_int( + "USAGE_FLUSH_BUCKET_BATCH_SIZE", + _DEFAULT_FLUSH_BUCKET_BATCH_SIZE, + ) + + +def flush_max_batches_per_run() -> int: + """Max Redis flush batches per org per flush_usage_to_catalog invocation.""" + return _env_int( + "USAGE_FLUSH_MAX_BATCHES_PER_RUN", + _DEFAULT_FLUSH_MAX_BATCHES_PER_RUN, + ) + + +def _flush_lock_ttl_seconds() -> int: + return _env_int( + "USAGE_FLUSH_LOCK_TTL_SECONDS", + _DEFAULT_FLUSH_LOCK_TTL_SECONDS, + ) + + +def _split_buckets_for_flush( + buckets: Dict[str, Dict[str, int]], + batch_size: int, +) -> Tuple[Dict[str, Dict[str, int]], Dict[str, Dict[str, int]]]: + if batch_size <= 0 or len(buckets) <= batch_size: + return buckets, {} + prefixes = sorted(buckets.keys())[:batch_size] + batch = {prefix: buckets[prefix] for prefix in prefixes} + remainder = { + prefix: metrics for prefix, metrics in buckets.items() if prefix not in batch + } + return batch, remainder + + +def _client() -> redis.Redis: + global _redis + if _redis is None: + _redis = redis.from_url(settings.REDIS_URL, decode_responses=True) + return _redis + + +def _token(value: Optional[UUID]) -> str: + return str(value) if value else _NONE + + +def _pending_hash_key(organization_id: UUID) -> str: + return f"usage:pending:{organization_id}" + + +def _flush_lock_key(organization_id: UUID) -> str: + return f"usage:flush:lock:{organization_id}" + + +def _claim_hash_key(organization_id: UUID, claim_id: str) -> str: + return f"usage:flushing:{organization_id}:{claim_id}" + + +def _bucket_prefix( + *, + workspace_id: Optional[UUID], + product_section: str, + model: str, + context: Optional[Dict[str, Any]], + usage_date: date, + usage_kind: str, +) -> str: + return "|".join( + [ + _token(workspace_id), + product_section, + model, + context_bucket_token(context), + usage_date.isoformat(), + usage_kind or USAGE_KIND_LLM, + ] + ) + + +def _parse_bucket_prefix(prefix: str) -> Optional[Dict[str, Any]]: + parts = prefix.split("|") + # New: 6 parts with JSON context token + usage_kind. + if len(parts) == 6: + ws_token, section, model, context_token, day_str, usage_kind = parts + context = parse_context_bucket_token(context_token) + elif len(parts) == 7: + # Legacy Redis keys: resource_id + resource_type before date/kind. + ( + ws_token, + section, + model, + resource_token, + resource_type, + day_str, + usage_kind, + ) = parts + resource_id = None if resource_token == _NONE else UUID(resource_token) + resolved_resource_type = None if resource_type == _NONE else resource_type + context = legacy_resource_context(resource_id, resolved_resource_type) + elif len(parts) == 5: + ws_token, section, model, context_token, day_str = parts + usage_kind = USAGE_KIND_LLM + context = parse_context_bucket_token(context_token) + else: + return None + try: + usage_date = date.fromisoformat(day_str) + except ValueError: + return None + workspace_id = None if ws_token == _NONE else UUID(ws_token) + return { + "workspace_id": workspace_id, + "product_section": section, + "model": model, + "context": context, + "usage_date": usage_date, + "usage_kind": usage_kind or USAGE_KIND_LLM, + } + + +def _bucket_from_context( + context: LLMUsageContext, + *, + model: str, + usage_date: date, + usage_kind: str, +) -> Dict[str, Any]: + return { + "workspace_id": context.workspace_id, + "product_section": context.product_section.value, + "model": model, + "context": build_bucket_context( + resource_id=context.resource_id, + resource_type=context.resource_type, + extra=context.extra, + ), + "usage_date": usage_date, + "usage_kind": usage_kind, + } + + +def _resolve_context( + *, + organization_id: Optional[UUID] = None, + ctx: Optional[LLMUsageContext] = None, +) -> Optional[LLMUsageContext]: + if ctx is not None: + return ctx + current = get_usage_context() + if current is not None: + return current + if organization_id is not None: + from app.services.usage.context import ( + get_usage_section_hint, + get_usage_workspace_hint, + ) + + return LLMUsageContext( + organization_id=organization_id, + workspace_id=get_usage_workspace_hint(), + product_section=get_usage_section_hint(), + ) + return None + + +_USAGE_METRIC_KEYS = ( + "prompt_tokens", + "completion_tokens", + "cache_read_tokens", + "cache_creation_tokens", + "reasoning_tokens", + "audio_seconds", + "tts_characters", +) + + +def _has_billable_usage_deltas(deltas: Dict[str, int]) -> bool: + return any(int(deltas.get(key, 0) or 0) > 0 for key in _USAGE_METRIC_KEYS) + + +def _deltas_from_usage(usage: UsageSnapshot) -> Dict[str, int]: + billable = ( + usage.prompt_tokens > 0 + or usage.completion_tokens > 0 + or usage.cache_read_tokens > 0 + or usage.cache_creation_tokens > 0 + or usage.reasoning_tokens > 0 + ) + if not billable: + return { + "prompt_tokens": 0, + "completion_tokens": 0, + "cache_read_tokens": 0, + "cache_creation_tokens": 0, + "reasoning_tokens": 0, + "audio_seconds": 0, + "tts_characters": 0, + "call_count": 0, + } + return { + "prompt_tokens": usage.prompt_tokens, + "completion_tokens": usage.completion_tokens, + "cache_read_tokens": usage.cache_read_tokens, + "cache_creation_tokens": usage.cache_creation_tokens, + "reasoning_tokens": usage.reasoning_tokens, + "audio_seconds": 0, + "tts_characters": 0, + "call_count": 1, + } + + +def _deltas_from_stt(audio_seconds: int, *, count_call: bool = True) -> Dict[str, int]: + return { + "prompt_tokens": 0, + "completion_tokens": 0, + "cache_read_tokens": 0, + "cache_creation_tokens": 0, + "reasoning_tokens": 0, + "audio_seconds": max(0, int(audio_seconds)), + "tts_characters": 0, + "call_count": 1 if count_call else 0, + } + + +def _deltas_from_tts(characters: int) -> Dict[str, int]: + return { + "prompt_tokens": 0, + "completion_tokens": 0, + "cache_read_tokens": 0, + "cache_creation_tokens": 0, + "reasoning_tokens": 0, + "audio_seconds": 0, + "tts_characters": max(0, int(characters)), + "call_count": 1, + } + + +def _cost_fields_for_pending_deltas( + organization_id: UUID, + bucket: Dict[str, Any], + deltas: Dict[str, int], + *, + db: Optional[Session] = None, +) -> Dict[str, Any]: + if not _has_billable_usage_deltas(deltas): + return {} + try: + from app.services.usage.pricing import cost_fields_from_deltas + except Exception as exc: + logger.debug("usage pending cost stamp skipped: {}", exc) + return {} + + usage_date = bucket["usage_date"] + if isinstance(usage_date, str): + usage_date = date.fromisoformat(usage_date) + usage_kind = bucket.get("usage_kind") or USAGE_KIND_LLM + + if db is not None: + try: + return cost_fields_from_deltas( + deltas, + organization_id=organization_id, + model=bucket["model"], + usage_kind=usage_kind, + usage_date=usage_date, + db=db, + ) + except Exception as exc: + logger.debug("usage pending cost stamp failed: {}", exc) + return {} + + try: + from app.database import SessionLocal + except Exception as exc: + logger.debug("usage pending cost stamp skipped: {}", exc) + return {} + + session = SessionLocal() + try: + return cost_fields_from_deltas( + deltas, + organization_id=organization_id, + model=bucket["model"], + usage_kind=usage_kind, + usage_date=usage_date, + db=session, + ) + except Exception as exc: + logger.debug("usage pending cost stamp failed: {}", exc) + return {} + finally: + session.close() + + +def _buffer_to_postgres( + organization_id: UUID, + bucket: Dict[str, Any], + deltas: Dict[str, int], +) -> None: + """Durable fallback when Redis is unavailable.""" + try: + from app.database import SessionLocal + except Exception as exc: + logger.warning("usage postgres fallback unavailable: {}", exc) + return + + usage_date = bucket["usage_date"] + if isinstance(usage_date, str): + usage_date = date.fromisoformat(usage_date) + usage_kind = bucket.get("usage_kind") or USAGE_KIND_LLM + params = { + "organization_id": str(organization_id), + "workspace_id": str(bucket["workspace_id"]) if bucket.get("workspace_id") else None, + "product_section": bucket["product_section"], + "model": bucket["model"], + "context": json.dumps(bucket.get("context") or {}), + "usage_date": usage_date.isoformat(), + "usage_kind": usage_kind, + "prompt_tokens": int(deltas.get("prompt_tokens", 0)), + "completion_tokens": int(deltas.get("completion_tokens", 0)), + "cache_read_tokens": int(deltas.get("cache_read_tokens", 0)), + "cache_creation_tokens": int(deltas.get("cache_creation_tokens", 0)), + "reasoning_tokens": int(deltas.get("reasoning_tokens", 0)), + "audio_seconds": int(deltas.get("audio_seconds", 0)), + "tts_characters": int(deltas.get("tts_characters", 0)), + "call_count": int(deltas.get("call_count", 0)), + } + db = SessionLocal() + try: + params.update( + _cost_fields_for_pending_deltas( + organization_id, bucket, deltas, db=db + ) + ) + db.execute( + text( + """ + INSERT INTO usage_pending_buffer ( + id, organization_id, workspace_id, product_section, model, + context, usage_date, usage_kind, + prompt_tokens, completion_tokens, cache_read_tokens, + cache_creation_tokens, reasoning_tokens, audio_seconds, + tts_characters, call_count, + input_cost_micro_usd, output_cost_micro_usd, + cache_read_cost_micro_usd, cache_creation_cost_micro_usd, + reasoning_cost_micro_usd, audio_cost_micro_usd, tts_cost_micro_usd, + total_cost_micro_usd, pricing_rate_source, pricing_rate_id, + created_at + ) VALUES ( + gen_random_uuid(), CAST(:organization_id AS uuid), + CAST(:workspace_id AS uuid), :product_section, :model, + CAST(:context AS jsonb), + CAST(:usage_date AS date), :usage_kind, + :prompt_tokens, :completion_tokens, :cache_read_tokens, + :cache_creation_tokens, :reasoning_tokens, :audio_seconds, + :tts_characters, :call_count, + :input_cost_micro_usd, :output_cost_micro_usd, + :cache_read_cost_micro_usd, :cache_creation_cost_micro_usd, + :reasoning_cost_micro_usd, :audio_cost_micro_usd, :tts_cost_micro_usd, + :total_cost_micro_usd, :pricing_rate_source, + CAST(:pricing_rate_id AS uuid), + now() + ) + """ + ), + params, + ) + db.commit() + except Exception as exc: + db.rollback() + logger.warning("usage postgres fallback insert failed: {}", exc) + finally: + db.close() + + +def _incr_pending( + organization_id: UUID, + prefix: str, + deltas: Dict[str, int], + bucket: Dict[str, Any], +) -> None: + hash_key = _pending_hash_key(organization_id) + try: + client = _client() + pipe = client.pipeline() + for metric, delta in deltas.items(): + if delta: + pipe.hincrby(hash_key, f"{prefix}|{metric}", int(delta)) + pipe.sadd("usage:pending:orgs", str(organization_id)) + pipe.expire(hash_key, _PENDING_TTL_SECONDS) + pipe.execute() + except redis.RedisError as exc: + logger.warning("usage redis counter failed, buffering to postgres: {}", exc) + _buffer_to_postgres(organization_id, bucket, deltas) + + +def _context_for_record( + *, + organization_id: Optional[UUID] = None, + ctx: Optional[LLMUsageContext] = None, +) -> Optional[LLMUsageContext]: + context = _resolve_context(organization_id=organization_id, ctx=ctx) + if context is None: + return None + from app.services.usage.call_import_context import enrich_usage_context_workspace + + return enrich_usage_context_workspace(context) + + +def record_llm_usage( + model: str, + usage: UsageSnapshot, + *, + organization_id: Optional[UUID] = None, + ctx: Optional[LLMUsageContext] = None, + usage_date: Optional[date] = None, +) -> None: + """Increment counters for one LLM call (best-effort, never raises).""" + if not model: + model = "unknown" + context = _context_for_record(organization_id=organization_id, ctx=ctx) + if context is None: + logger.warning("llm usage record skipped: missing organization_id") + return + + deltas = _deltas_from_usage(usage) + if not _has_billable_usage_deltas(deltas): + return + + day = usage_date or datetime.now(timezone.utc).date() + bucket = _bucket_from_context( + context, + model=model, + usage_date=day, + usage_kind=USAGE_KIND_LLM, + ) + prefix = _bucket_prefix(**bucket) + _incr_pending(context.organization_id, prefix, deltas, bucket) + + +def record_stt_usage( + model: str, + *, + audio_seconds: float | int = 0, + organization_id: Optional[UUID] = None, + ctx: Optional[LLMUsageContext] = None, + usage_date: Optional[date] = None, + count_call: bool = True, +) -> None: + """Increment counters for one STT call (audio seconds + optional call_count).""" + if not model: + model = "unknown" + context = _context_for_record(organization_id=organization_id, ctx=ctx) + if context is None: + logger.warning("stt usage record skipped: missing organization_id") + return + + seconds = int(max(0, math.ceil(float(audio_seconds or 0)))) + if seconds <= 0: + return + deltas = _deltas_from_stt(seconds, count_call=count_call) + day = usage_date or datetime.now(timezone.utc).date() + bucket = _bucket_from_context( + context, + model=model, + usage_date=day, + usage_kind=USAGE_KIND_STT, + ) + prefix = _bucket_prefix(**bucket) + _incr_pending(context.organization_id, prefix, deltas, bucket) + + +def record_call_usage( + model: str = "voice-call", + *, + organization_id: Optional[UUID] = None, + ctx: Optional[LLMUsageContext] = None, + usage_date: Optional[date] = None, + audio_seconds: int = 0, +) -> None: + """Record one completed call session (call_count + optional duration). + + Best-effort, never raises. Uses the same Redis-buffered path as other + usage counters (one pipelined HINCRBY batch per call). + """ + if not model: + model = "voice-call" + context = _context_for_record(organization_id=organization_id, ctx=ctx) + if context is None: + logger.warning("call usage record skipped: missing organization_id") + return + + seconds = max(0, int(audio_seconds or 0)) + deltas = { + "prompt_tokens": 0, + "completion_tokens": 0, + "cache_read_tokens": 0, + "cache_creation_tokens": 0, + "reasoning_tokens": 0, + "audio_seconds": seconds, + "tts_characters": 0, + "call_count": 1, + } + day = usage_date or datetime.now(timezone.utc).date() + bucket = _bucket_from_context( + context, + model=model, + usage_date=day, + usage_kind=USAGE_KIND_LLM, + ) + prefix = _bucket_prefix(**bucket) + _incr_pending(context.organization_id, prefix, deltas, bucket) + + +def record_tts_usage( + model: str, + *, + characters: int = 0, + organization_id: Optional[UUID] = None, + ctx: Optional[LLMUsageContext] = None, + usage_date: Optional[date] = None, +) -> None: + """Increment counters for one TTS call (characters + call_count).""" + if not model: + model = "unknown" + context = _context_for_record(organization_id=organization_id, ctx=ctx) + if context is None: + logger.warning("tts usage record skipped: missing organization_id") + return + + chars = max(0, int(characters or 0)) + if chars <= 0: + return + deltas = _deltas_from_tts(chars) + day = usage_date or datetime.now(timezone.utc).date() + bucket = _bucket_from_context( + context, + model=model, + usage_date=day, + usage_kind=USAGE_KIND_TTS, + ) + prefix = _bucket_prefix(**bucket) + _incr_pending(context.organization_id, prefix, deltas, bucket) + + +def probe_audio_seconds(audio_file_path: str) -> int: + """Best-effort audio duration in whole seconds.""" + if not audio_file_path: + return 0 + try: + from pydub import AudioSegment + + return max(0, int(math.ceil(AudioSegment.from_file(audio_file_path).duration_seconds))) + except Exception: + pass + try: + import librosa + + return max(0, int(math.ceil(librosa.get_duration(path=audio_file_path)))) + except Exception: + return 0 + + +def _parse_hash_to_buckets(raw: Dict[str, str]) -> Dict[str, Dict[str, int]]: + buckets: Dict[str, Dict[str, int]] = {} + for field, value in raw.items(): + if "|" not in field: + continue + prefix, metric = field.rsplit("|", 1) + if metric not in _METRIC_FIELDS: + continue + amount = int(value or 0) + if not amount: + continue + bucket = buckets.setdefault(prefix, {}) + bucket[metric] = bucket.get(metric, 0) + amount + return buckets + + +def _read_hash_buckets(hash_key: str) -> Dict[str, Dict[str, int]]: + try: + client = _client() + raw = client.hgetall(hash_key) + except redis.RedisError as exc: + logger.warning("llm usage read hash failed: {}", exc) + return {} + return _parse_hash_to_buckets(raw) + + +def _restore_buckets_to_pending( + organization_id: UUID, buckets: Dict[str, Dict[str, int]] +) -> None: + try: + client = _client() + pipe = client.pipeline() + hash_key = _pending_hash_key(organization_id) + for prefix, metrics in buckets.items(): + for metric, amount in metrics.items(): + if amount: + pipe.hincrby(hash_key, f"{prefix}|{metric}", amount) + pipe.sadd("usage:pending:orgs", str(organization_id)) + pipe.expire(hash_key, _PENDING_TTL_SECONDS) + pipe.execute() + except redis.RedisError as exc: + logger.warning("llm usage redis restore failed, buffering: {}", exc) + for prefix, metrics in buckets.items(): + parsed = _parse_bucket_prefix(prefix) + if parsed: + _buffer_to_postgres(organization_id, parsed, metrics) + + +_CLAIM_COMMITTED_TTL_SECONDS = 24 * 60 * 60 +_CLAIM_COMMITTED_REDIS_RETRIES = 3 +_ORPHAN_RECOVERY_INTERVAL_SEC = 30 +_COMMITTED_CLAIMS_PRUNE_INTERVAL_SEC = 3600 +_COMMITTED_CLAIMS_RETENTION_DAYS = 2 + + +def _claim_committed_key(claim_key: str) -> str: + return f"usage:claim_done:{claim_key}" + + +def _redis_interval_gate(key: str, interval_sec: int) -> bool: + """Return True when the gated work should run (interval lock acquired).""" + try: + return bool(_client().set(key, "1", nx=True, ex=interval_sec)) + except redis.RedisError: + return False + + +def _record_claim_committed_pg( + db: Session, claim_key: str, organization_id: UUID +) -> None: + db.execute( + text( + """ + INSERT INTO usage_committed_claims (claim_key, organization_id) + VALUES (:claim_key, CAST(:organization_id AS uuid)) + ON CONFLICT (claim_key) DO NOTHING + """ + ), + { + "claim_key": claim_key, + "organization_id": str(organization_id), + }, + ) + + +def _committed_claim_keys_in_pg(claim_keys: List[str]) -> set[str]: + if not claim_keys: + return set() + try: + from app.database import SessionLocal + + db = SessionLocal() + try: + rows = db.execute( + text( + """ + SELECT claim_key + FROM usage_committed_claims + WHERE claim_key = ANY(CAST(:claim_keys AS text[])) + """ + ), + {"claim_keys": claim_keys}, + ).all() + return {row[0] for row in rows} + finally: + db.close() + except Exception: + return set() + + +def _prune_old_committed_claims() -> None: + try: + from app.database import SessionLocal + + db = SessionLocal() + try: + db.execute( + text( + f""" + DELETE FROM usage_committed_claims + WHERE committed_at < now() - interval '{_COMMITTED_CLAIMS_RETENTION_DAYS} days' + """ + ) + ) + db.commit() + except Exception as exc: + db.rollback() + logger.debug("usage committed claim prune skipped: {}", exc) + finally: + db.close() + except Exception: + pass + + +def _discard_committed_claim(claim_key: str, organization_id: UUID) -> None: + """Best-effort Redis cleanup when usage was not (or must not be) persisted.""" + _mark_claim_committed(claim_key, fast=True) + _ack_claim(claim_key, organization_id) + + +def _is_claim_committed_pg(claim_key: str) -> bool: + return claim_key in _committed_claim_keys_in_pg([claim_key]) + + +def _mark_claim_committed(claim_key: str, *, fast: bool = False) -> bool: + max_attempts = 2 if fast else _CLAIM_COMMITTED_REDIS_RETRIES + for attempt in range(max_attempts): + try: + _client().set( + _claim_committed_key(claim_key), + "1", + ex=_CLAIM_COMMITTED_TTL_SECONDS, + ) + return True + except redis.RedisError: + if attempt + 1 < max_attempts and not fast: + time.sleep(0.05 * (attempt + 1)) + return False + + +def _is_claim_committed(claim_key: str) -> bool: + try: + if _client().exists(_claim_committed_key(claim_key)): + return True + except redis.RedisError: + pass + return _is_claim_committed_pg(claim_key) + + +def _has_pending_usage(organization_id: UUID) -> bool: + try: + client = _client() + if client.exists(_pending_hash_key(organization_id)): + return True + return bool(client.sismember("usage:pending:orgs", str(organization_id))) + except redis.RedisError: + return False + + +def _acquire_flush_lock(organization_id: UUID) -> bool: + try: + client = _client() + lock_key = _flush_lock_key(organization_id) + lock_ttl = _flush_lock_ttl_seconds() + if client.set(lock_key, "1", nx=True, ex=lock_ttl): + return True + deadline = time.monotonic() + _FLUSH_LOCK_WAIT_SECONDS + while time.monotonic() < deadline: + time.sleep(0.05) + if client.set(lock_key, "1", nx=True, ex=lock_ttl): + return True + if client.get(lock_key) is None: + continue + return False + except redis.RedisError as exc: + logger.warning("llm usage flush lock failed: {}", exc) + return False + + +def _release_flush_lock(organization_id: UUID) -> None: + try: + _client().delete(_flush_lock_key(organization_id)) + except redis.RedisError: + pass + + +def _claim_pending( + organization_id: UUID, +) -> Tuple[Optional[str], Dict[str, Dict[str, int]]]: + claim_id = str(uuid.uuid4()) + pending_key = _pending_hash_key(organization_id) + claim_key = _claim_hash_key(organization_id, claim_id) + try: + client = _client() + claimed = int(client.eval(_CLAIM_LUA, 2, pending_key, claim_key) or 0) + if not claimed: + return None, {} + buckets = _read_hash_buckets(claim_key) + if not buckets: + client.delete(claim_key) + return None, {} + return claim_key, buckets + except redis.RedisError as exc: + logger.warning("llm usage claim failed: {}", exc) + return None, {} + + +def _ack_claim(claim_key: str, organization_id: UUID) -> None: + try: + client = _client() + client.delete(claim_key) + pending_key = _pending_hash_key(organization_id) + if not client.exists(pending_key): + client.srem("usage:pending:orgs", str(organization_id)) + except redis.RedisError: + pass + + +def _stamp_cost_params( + db: Session, + organization_id: UUID, + bucket: Dict[str, Any], + deltas: Dict[str, int], + *, + pricing_resolver: Any = None, + stamped_costs: Optional[Dict[str, Any]] = None, +) -> Dict[str, Any]: + if stamped_costs is not None: + return { + "input_cost_micro_usd": int(stamped_costs.get("input_cost_micro_usd") or 0), + "output_cost_micro_usd": int(stamped_costs.get("output_cost_micro_usd") or 0), + "cache_read_cost_micro_usd": int( + stamped_costs.get("cache_read_cost_micro_usd") or 0 + ), + "cache_creation_cost_micro_usd": int( + stamped_costs.get("cache_creation_cost_micro_usd") or 0 + ), + "reasoning_cost_micro_usd": int( + stamped_costs.get("reasoning_cost_micro_usd") or 0 + ), + "audio_cost_micro_usd": int(stamped_costs.get("audio_cost_micro_usd") or 0), + "tts_cost_micro_usd": int(stamped_costs.get("tts_cost_micro_usd") or 0), + "total_cost_micro_usd": int(stamped_costs.get("total_cost_micro_usd") or 0), + "pricing_rate_source": stamped_costs.get("pricing_rate_source"), + "pricing_rate_id": stamped_costs.get("pricing_rate_id"), + } + + from app.services.usage.pricing import cost_fields_from_deltas, PricingResolver + + usage_date = bucket["usage_date"] + if isinstance(usage_date, str): + usage_date = date.fromisoformat(usage_date) + resolver = pricing_resolver or PricingResolver(db) + return cost_fields_from_deltas( + deltas, + organization_id=organization_id, + model=bucket["model"], + usage_kind=bucket.get("usage_kind") or USAGE_KIND_LLM, + usage_date=usage_date, + db=db, + resolver=resolver, + ) + + +def _upsert_bucket( + db: Session, + organization_id: UUID, + bucket: Dict[str, Any], + deltas: Dict[str, int], + *, + pricing_resolver: Any = None, + stamped_costs: Optional[Dict[str, Any]] = None, +) -> None: + context = bucket.get("context") or {} + params = { + "organization_id": str(organization_id), + "workspace_id": str(bucket["workspace_id"]) if bucket["workspace_id"] else None, + "product_section": bucket["product_section"], + "model": bucket["model"], + "context": json.dumps(context), + "context_resource_id": str(context.get("resource_id") or ""), + "context_resource_type": str(context.get("resource_type") or ""), + "usage_date": bucket["usage_date"].isoformat(), + "usage_kind": bucket.get("usage_kind") or USAGE_KIND_LLM, + "prompt_tokens": int(deltas.get("prompt_tokens", 0)), + "completion_tokens": int(deltas.get("completion_tokens", 0)), + "cache_read_tokens": int(deltas.get("cache_read_tokens", 0)), + "cache_creation_tokens": int(deltas.get("cache_creation_tokens", 0)), + "reasoning_tokens": int(deltas.get("reasoning_tokens", 0)), + "audio_seconds": int(deltas.get("audio_seconds", 0)), + "tts_characters": int(deltas.get("tts_characters", 0)), + "call_count": int(deltas.get("call_count", 0)), + } + apply_cost_increment = stamped_costs is not None or _has_billable_usage_deltas(deltas) + if apply_cost_increment: + cost_fields = _stamp_cost_params( + db, + organization_id, + bucket, + deltas, + pricing_resolver=pricing_resolver, + stamped_costs=stamped_costs, + ) + params.update( + { + "input_cost_micro_usd": int(cost_fields["input_cost_micro_usd"]), + "output_cost_micro_usd": int(cost_fields["output_cost_micro_usd"]), + "cache_read_cost_micro_usd": int( + cost_fields["cache_read_cost_micro_usd"] + ), + "cache_creation_cost_micro_usd": int( + cost_fields["cache_creation_cost_micro_usd"] + ), + "reasoning_cost_micro_usd": int( + cost_fields["reasoning_cost_micro_usd"] + ), + "audio_cost_micro_usd": int(cost_fields["audio_cost_micro_usd"]), + "tts_cost_micro_usd": int(cost_fields["tts_cost_micro_usd"]), + "total_cost_micro_usd": int(cost_fields["total_cost_micro_usd"]), + "pricing_rate_source": cost_fields.get("pricing_rate_source"), + "pricing_rate_id": cost_fields.get("pricing_rate_id"), + } + ) + cost_update_set = "" + if apply_cost_increment: + cost_update_set = """ + input_cost_micro_usd = input_cost_micro_usd + :input_cost_micro_usd, + output_cost_micro_usd = output_cost_micro_usd + :output_cost_micro_usd, + cache_read_cost_micro_usd = cache_read_cost_micro_usd + :cache_read_cost_micro_usd, + cache_creation_cost_micro_usd = cache_creation_cost_micro_usd + :cache_creation_cost_micro_usd, + reasoning_cost_micro_usd = reasoning_cost_micro_usd + :reasoning_cost_micro_usd, + audio_cost_micro_usd = audio_cost_micro_usd + :audio_cost_micro_usd, + tts_cost_micro_usd = tts_cost_micro_usd + :tts_cost_micro_usd, + total_cost_micro_usd = total_cost_micro_usd + :total_cost_micro_usd, + pricing_rate_source = COALESCE(:pricing_rate_source, pricing_rate_source), + pricing_rate_id = CASE + WHEN CAST(:pricing_rate_id AS text) IS NOT NULL + THEN CAST(:pricing_rate_id AS uuid) + ELSE pricing_rate_id + END, + """ + update_set = ( + """ + prompt_tokens = prompt_tokens + :prompt_tokens, + completion_tokens = completion_tokens + :completion_tokens, + cache_read_tokens = cache_read_tokens + :cache_read_tokens, + cache_creation_tokens = cache_creation_tokens + :cache_creation_tokens, + reasoning_tokens = reasoning_tokens + :reasoning_tokens, + audio_seconds = audio_seconds + :audio_seconds, + tts_characters = tts_characters + :tts_characters, + call_count = call_count + :call_count, + """ + + cost_update_set + + """ + context = CASE + WHEN context = '{}'::jsonb THEN CAST(:context AS jsonb) + ELSE context || CAST(:context AS jsonb) + END, + updated_at = now() + """ + ) + where_base = """ + WHERE organization_id = CAST(:organization_id AS uuid) + AND product_section = :product_section + AND model = :model + AND usage_date = CAST(:usage_date AS date) + AND usage_kind = :usage_kind + AND workspace_id IS NOT DISTINCT FROM CAST(:workspace_id AS uuid) + """ + exact_context_where = where_base + " AND context = CAST(:context AS jsonb)" + legacy_context_where = ( + where_base + + """ + AND COALESCE(context->>'resource_id', '') = :context_resource_id + AND COALESCE(context->>'resource_type', '') = :context_resource_type + """ + ) + + result = db.execute( + text(f"UPDATE llm_usage_daily SET {update_set} {exact_context_where}"), + params, + ) + if not result.rowcount: + result = db.execute( + text(f"UPDATE llm_usage_daily SET {update_set} {legacy_context_where}"), + params, + ) + if not result.rowcount: + db.execute(text("SAVEPOINT llm_usage_bucket_insert")) + try: + if apply_cost_increment: + db.execute( + text( + """ + INSERT INTO llm_usage_daily ( + id, organization_id, workspace_id, product_section, model, + context, usage_date, usage_kind, + prompt_tokens, completion_tokens, cache_read_tokens, + cache_creation_tokens, reasoning_tokens, audio_seconds, + tts_characters, call_count, + input_cost_micro_usd, output_cost_micro_usd, + cache_read_cost_micro_usd, cache_creation_cost_micro_usd, + reasoning_cost_micro_usd, audio_cost_micro_usd, tts_cost_micro_usd, + total_cost_micro_usd, pricing_rate_source, pricing_rate_id, + created_at, updated_at + ) VALUES ( + gen_random_uuid(), CAST(:organization_id AS uuid), + CAST(:workspace_id AS uuid), :product_section, :model, + CAST(:context AS jsonb), CAST(:usage_date AS date), + :usage_kind, + :prompt_tokens, :completion_tokens, :cache_read_tokens, + :cache_creation_tokens, :reasoning_tokens, :audio_seconds, + :tts_characters, :call_count, + :input_cost_micro_usd, :output_cost_micro_usd, + :cache_read_cost_micro_usd, :cache_creation_cost_micro_usd, + :reasoning_cost_micro_usd, :audio_cost_micro_usd, :tts_cost_micro_usd, + :total_cost_micro_usd, :pricing_rate_source, + CAST(:pricing_rate_id AS uuid), + now(), now() + ) + """ + ), + params, + ) + else: + db.execute( + text( + """ + INSERT INTO llm_usage_daily ( + id, organization_id, workspace_id, product_section, model, + context, usage_date, usage_kind, + prompt_tokens, completion_tokens, cache_read_tokens, + cache_creation_tokens, reasoning_tokens, audio_seconds, + tts_characters, call_count, + created_at, updated_at + ) VALUES ( + gen_random_uuid(), CAST(:organization_id AS uuid), + CAST(:workspace_id AS uuid), :product_section, :model, + CAST(:context AS jsonb), CAST(:usage_date AS date), + :usage_kind, + :prompt_tokens, :completion_tokens, :cache_read_tokens, + :cache_creation_tokens, :reasoning_tokens, :audio_seconds, + :tts_characters, :call_count, + now(), now() + ) + """ + ), + params, + ) + db.execute(text("RELEASE SAVEPOINT llm_usage_bucket_insert")) + except IntegrityError as exc: + if not _is_unique_violation(exc): + db.execute(text("ROLLBACK TO SAVEPOINT llm_usage_bucket_insert")) + db.execute(text("RELEASE SAVEPOINT llm_usage_bucket_insert")) + raise + db.execute(text("ROLLBACK TO SAVEPOINT llm_usage_bucket_insert")) + result = db.execute( + text(f"UPDATE llm_usage_daily SET {update_set} {legacy_context_where}"), + params, + ) + if not result.rowcount: + result = db.execute( + text(f"UPDATE llm_usage_daily SET {update_set} {exact_context_where}"), + params, + ) + db.execute(text("RELEASE SAVEPOINT llm_usage_bucket_insert")) + if not result.rowcount: + logger.warning( + "llm usage upsert unique conflict but no matching bucket for org {}", + organization_id, + ) + + +def _upsert_claimed_buckets( + db: Session, + organization_id: UUID, + buckets: Dict[str, Dict[str, int]], + *, + pricing_resolver: Any, +) -> Tuple[int, Dict[str, Dict[str, int]]]: + """Persist claimed Redis buckets; return flushed count and unparseable buckets.""" + flushed = 0 + skipped: Dict[str, Dict[str, int]] = {} + for prefix, deltas in buckets.items(): + parsed = _parse_bucket_prefix(prefix) + if not parsed: + skipped[prefix] = deltas + logger.warning( + "llm usage skipped unparseable bucket prefix for org {}", + organization_id, + ) + continue + _upsert_bucket( + db, + organization_id, + parsed, + deltas, + pricing_resolver=pricing_resolver, + ) + flushed += 1 + return flushed, skipped + + +def _flush_redis_pending_to_catalog( + db: Session, + organization_id: UUID, + *, + pricing_resolver: Any, +) -> int: + """Drain Redis pending hash into llm_usage_daily in bounded batches.""" + batch_size = flush_bucket_batch_size() + max_batches = flush_max_batches_per_run() + flushed = 0 + + for _ in range(max_batches): + if not _has_pending_usage(organization_id): + break + + claim_key, claimed_buckets = _claim_pending(organization_id) + if not claim_key or not claimed_buckets: + break + + batch, remainder = _split_buckets_for_flush(claimed_buckets, batch_size) + skipped: Dict[str, Dict[str, int]] = {} + try: + batch_flushed, skipped = _upsert_claimed_buckets( + db, + organization_id, + batch, + pricing_resolver=pricing_resolver, + ) + if claim_key: + _record_claim_committed_pg(db, claim_key, organization_id) + db.commit() + flushed += batch_flushed + if remainder: + _restore_buckets_to_pending(organization_id, remainder) + if skipped: + _restore_buckets_to_pending(organization_id, skipped) + if claim_key: + _mark_claim_committed(claim_key, fast=True) + _ack_claim(claim_key, organization_id) + except Exception as exc: + db.rollback() + if _is_missing_organization_fk(exc): + logger.warning( + "llm usage flush dropped for unknown organization {}: {}", + organization_id, + exc, + ) + if claim_key: + try: + _record_claim_committed_pg(db, claim_key, organization_id) + db.commit() + except Exception: + db.rollback() + _discard_committed_claim(claim_key, organization_id) + return flushed + logger.warning("llm usage catalog flush failed, restoring redis: {}", exc) + restore_buckets = dict(claimed_buckets) + _restore_buckets_to_pending(organization_id, restore_buckets) + if claim_key: + try: + _client().delete(claim_key) + except redis.RedisError: + pass + return flushed + + return flushed + + +def _is_unique_violation(exc: BaseException) -> bool: + text_blob = " ".join( + str(part) + for part in ( + exc, + getattr(exc, "orig", None), + getattr(getattr(exc, "orig", None), "pgcode", None), + ) + if part is not None + ).lower() + return "uniqueviolation" in text_blob or "duplicate key" in text_blob + + +def _is_missing_organization_fk(exc: BaseException) -> bool: + text_blob = " ".join( + str(part) + for part in (exc, getattr(exc, "orig", None), getattr(exc, "args", None)) + if part is not None + ).lower() + return "llm_usage_daily_organization_id_fkey" in text_blob + + +def _flush_pending_buffer( + db: Session, organization_id: UUID, *, pricing_resolver: Any = None +) -> int: + """Drain Postgres write-ahead rows into llm_usage_daily.""" + from app.services.usage.pricing import PricingResolver + + resolver = pricing_resolver or PricingResolver(db) + try: + rows = db.execute( + text( + """ + SELECT id, workspace_id, product_section, model, context, + usage_date, usage_kind, + prompt_tokens, completion_tokens, cache_read_tokens, + cache_creation_tokens, reasoning_tokens, audio_seconds, + tts_characters, call_count, + input_cost_micro_usd, output_cost_micro_usd, + cache_read_cost_micro_usd, cache_creation_cost_micro_usd, + reasoning_cost_micro_usd, audio_cost_micro_usd, tts_cost_micro_usd, + total_cost_micro_usd, pricing_rate_source, pricing_rate_id + FROM usage_pending_buffer + WHERE organization_id = CAST(:organization_id AS uuid) + ORDER BY created_at ASC + LIMIT :batch_limit + """ + ), + { + "organization_id": str(organization_id), + "batch_limit": flush_bucket_batch_size(), + }, + ).mappings().all() + except Exception as exc: + db.rollback() + logger.debug("usage buffer read skipped: {}", exc) + return 0 + if not rows: + return 0 + + flushed = 0 + ids: List[str] = [] + try: + for row in rows: + row_context = row["context"] or {} + if isinstance(row_context, str): + row_context = json.loads(row_context) + bucket = { + "workspace_id": row["workspace_id"], + "product_section": row["product_section"], + "model": row["model"], + "context": row_context, + "usage_date": row["usage_date"], + "usage_kind": row["usage_kind"] or USAGE_KIND_LLM, + } + deltas = { + "prompt_tokens": int(row["prompt_tokens"] or 0), + "completion_tokens": int(row["completion_tokens"] or 0), + "cache_read_tokens": int(row["cache_read_tokens"] or 0), + "cache_creation_tokens": int(row["cache_creation_tokens"] or 0), + "reasoning_tokens": int(row["reasoning_tokens"] or 0), + "audio_seconds": int(row["audio_seconds"] or 0), + "tts_characters": int(row.get("tts_characters") or 0), + "call_count": int(row["call_count"] or 0), + } + stamped_costs = { + "input_cost_micro_usd": row.get("input_cost_micro_usd"), + "output_cost_micro_usd": row.get("output_cost_micro_usd"), + "cache_read_cost_micro_usd": row.get("cache_read_cost_micro_usd"), + "cache_creation_cost_micro_usd": row.get("cache_creation_cost_micro_usd"), + "reasoning_cost_micro_usd": row.get("reasoning_cost_micro_usd"), + "audio_cost_micro_usd": row.get("audio_cost_micro_usd"), + "tts_cost_micro_usd": row.get("tts_cost_micro_usd"), + "total_cost_micro_usd": row.get("total_cost_micro_usd"), + "pricing_rate_source": row.get("pricing_rate_source"), + "pricing_rate_id": ( + str(row["pricing_rate_id"]) if row.get("pricing_rate_id") else None + ), + } + _upsert_bucket( + db, + organization_id, + bucket, + deltas, + pricing_resolver=resolver, + stamped_costs=stamped_costs, + ) + ids.append(str(row["id"])) + flushed += 1 + if ids: + from sqlalchemy import bindparam + + db.execute( + text( + "DELETE FROM usage_pending_buffer WHERE id IN :ids" + ).bindparams(bindparam("ids", expanding=True)), + {"ids": ids}, + ) + db.commit() + return flushed + except Exception as exc: + db.rollback() + if _is_missing_organization_fk(exc): + logger.warning( + "usage buffer dropped for unknown organization {}: {}", + organization_id, + exc, + ) + if ids: + try: + db.execute( + text( + """ + DELETE FROM usage_pending_buffer + WHERE organization_id = CAST(:organization_id AS uuid) + """ + ), + {"organization_id": str(organization_id)}, + ) + db.commit() + except Exception: + db.rollback() + return 0 + logger.warning("usage buffer flush failed: {}", exc) + return 0 + + +def flush_usage_to_catalog(db: Session, organization_id: UUID) -> int: + """Claim Redis deltas + drain PG buffer into llm_usage_daily.""" + from app.services.usage.pricing import PricingResolver + from app.services.usage.read_cache import invalidate_org_usage_read_cache + + _recover_orphaned_claims() + flushed = 0 + pricing_resolver = PricingResolver(db) + redis_locked = _acquire_flush_lock(organization_id) + try: + if redis_locked: + flushed += _flush_redis_pending_to_catalog( + db, + organization_id, + pricing_resolver=pricing_resolver, + ) + finally: + if redis_locked: + _release_flush_lock(organization_id) + + flushed += _flush_pending_buffer(db, organization_id, pricing_resolver=pricing_resolver) + invalidate_org_usage_read_cache(organization_id) + return flushed + + +def _recover_orphaned_claims() -> None: + if not _redis_interval_gate( + "usage:orphan_recovery:due", _ORPHAN_RECOVERY_INTERVAL_SEC + ): + return + try: + client = _client() + candidates: List[Tuple[str, UUID]] = [] + for claim_key in client.scan_iter(match="usage:flushing:*", count=100): + parts = claim_key.split(":") + if len(parts) < 4: + continue + try: + org_id = UUID(parts[2]) + except ValueError: + continue + if client.exists(_flush_lock_key(org_id)): + continue + candidates.append((claim_key, org_id)) + + if not candidates: + return + + pipe = client.pipeline() + for claim_key, _org_id in candidates: + pipe.exists(_claim_committed_key(claim_key)) + redis_committed_flags = pipe.execute() + + needs_pg: List[Tuple[str, UUID]] = [] + for index, (claim_key, org_id) in enumerate(candidates): + if redis_committed_flags[index]: + client.delete(claim_key) + continue + needs_pg.append((claim_key, org_id)) + + pg_committed = _committed_claim_keys_in_pg([key for key, _ in needs_pg]) + for claim_key, org_id in needs_pg: + if claim_key in pg_committed: + client.delete(claim_key) + continue + buckets = _read_hash_buckets(claim_key) + if buckets: + _restore_buckets_to_pending(org_id, buckets) + client.delete(claim_key) + + if _redis_interval_gate( + "usage:committed_claims:prune", _COMMITTED_CLAIMS_PRUNE_INTERVAL_SEC + ): + _prune_old_committed_claims() + except redis.RedisError as exc: + logger.warning("llm usage orphan claim recovery failed: {}", exc) + + +def list_pending_organization_ids() -> List[UUID]: + result: List[UUID] = [] + seen: set[UUID] = set() + try: + client = _client() + raw_ids = client.smembers("usage:pending:orgs") + stale: List[str] = [] + for value in raw_ids: + try: + org_id = UUID(value) + except ValueError: + stale.append(value) + continue + if client.exists(_pending_hash_key(org_id)): + result.append(org_id) + seen.add(org_id) + else: + stale.append(value) + if stale: + client.srem("usage:pending:orgs", *stale) + except (redis.RedisError, ValueError): + pass + + try: + from app.database import SessionLocal + + db = SessionLocal() + try: + rows = db.execute( + text( + """ + SELECT DISTINCT organization_id + FROM usage_pending_buffer + LIMIT 5000 + """ + ) + ).all() + for row in rows: + org_id = row[0] if not isinstance(row, dict) else row["organization_id"] + if isinstance(org_id, str): + org_id = UUID(org_id) + if org_id not in seen: + result.append(org_id) + seen.add(org_id) + finally: + db.close() + except Exception as exc: + logger.debug("usage buffer org scan skipped: {}", exc) + + return result + + +def flush_all_usage_to_catalog(db_factory) -> int: + _recover_orphaned_claims() + total = 0 + for org_id in list_pending_organization_ids(): + db = db_factory() + try: + total += flush_usage_to_catalog(db, org_id) + except Exception as exc: + db.rollback() + logger.warning("flush_all usage failed for {}: {}", org_id, exc) + finally: + db.close() + return total + + +def merge_usage_totals( + rows: Iterable[Any], +) -> Dict[str, int]: + totals = {field: 0 for field in _METRIC_FIELDS} + for row in rows: + totals["prompt_tokens"] += int(getattr(row, "prompt_tokens", 0) or 0) + totals["completion_tokens"] += int(getattr(row, "completion_tokens", 0) or 0) + totals["cache_read_tokens"] += int(getattr(row, "cache_read_tokens", 0) or 0) + totals["cache_creation_tokens"] += int( + getattr(row, "cache_creation_tokens", 0) or 0 + ) + totals["reasoning_tokens"] += int(getattr(row, "reasoning_tokens", 0) or 0) + totals["audio_seconds"] += int(getattr(row, "audio_seconds", 0) or 0) + totals["tts_characters"] += int(getattr(row, "tts_characters", 0) or 0) + totals["call_count"] += int(getattr(row, "call_count", 0) or 0) + totals["total_tokens"] = totals["prompt_tokens"] + totals["completion_tokens"] + return totals diff --git a/app/services/usage/normalize.py b/app/services/usage/normalize.py new file mode 100644 index 00000000..7b2b8e88 --- /dev/null +++ b/app/services/usage/normalize.py @@ -0,0 +1,109 @@ +"""Normalize LiteLLM / provider usage objects into UsageSnapshot.""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Any, Mapping, Optional + + +@dataclass(frozen=True) +class UsageSnapshot: + prompt_tokens: int + completion_tokens: int + cache_read_tokens: int = 0 + cache_creation_tokens: int = 0 + reasoning_tokens: int = 0 + + @property + def total_tokens(self) -> int: + return self.prompt_tokens + self.completion_tokens + + +def _as_int(value: Any) -> int: + try: + return int(value or 0) + except (TypeError, ValueError): + return 0 + + +def _usage_mapping(raw: Any) -> Mapping[str, Any]: + if raw is None: + return {} + if isinstance(raw, Mapping): + return raw + if hasattr(raw, "model_dump"): + try: + return raw.model_dump() + except Exception: + pass + if hasattr(raw, "__dict__"): + return { + key: value + for key, value in vars(raw).items() + if not key.startswith("_") + } + return {} + + +def _details_dict(usage: Mapping[str, Any], key: str) -> Mapping[str, Any]: + details = usage.get(key) + if details is None: + return {} + if isinstance(details, Mapping): + return details + if hasattr(details, "model_dump"): + try: + return details.model_dump() + except Exception: + pass + if hasattr(details, "__dict__"): + return { + k: v for k, v in vars(details).items() if not k.startswith("_") + } + return {} + + +def normalize_llm_usage(raw_response: Any = None, *, usage: Any = None) -> UsageSnapshot: + """Extract token buckets from a LiteLLM response or raw usage object.""" + usage_obj = usage + if usage_obj is None and raw_response is not None: + usage_obj = getattr(raw_response, "usage", None) + + data = _usage_mapping(usage_obj) + prompt_tokens = _as_int(data.get("prompt_tokens") or data.get("input_tokens")) + completion_tokens = _as_int( + data.get("completion_tokens") or data.get("output_tokens") + ) + + cache_read = _as_int(data.get("cache_read_input_tokens")) + cache_creation = _as_int(data.get("cache_creation_input_tokens")) + + prompt_details = _details_dict(data, "prompt_tokens_details") + if not cache_read: + cache_read = _as_int(prompt_details.get("cached_tokens")) + if not cache_creation: + cache_creation = _as_int( + prompt_details.get("cache_write_tokens") + or prompt_details.get("cache_creation_tokens") + ) + + completion_details = _details_dict(data, "completion_tokens_details") + reasoning_tokens = _as_int(completion_details.get("reasoning_tokens")) + + return UsageSnapshot( + prompt_tokens=prompt_tokens, + completion_tokens=completion_tokens, + cache_read_tokens=cache_read, + cache_creation_tokens=cache_creation, + reasoning_tokens=reasoning_tokens, + ) + + +def usage_snapshot_is_billable(snapshot: UsageSnapshot) -> bool: + return ( + snapshot.prompt_tokens > 0 + or snapshot.completion_tokens > 0 + or snapshot.cache_read_tokens > 0 + or snapshot.cache_creation_tokens > 0 + or snapshot.reasoning_tokens > 0 + ) diff --git a/app/services/usage/pricing.py b/app/services/usage/pricing.py new file mode 100644 index 00000000..cb9dde02 --- /dev/null +++ b/app/services/usage/pricing.py @@ -0,0 +1,923 @@ +"""Usage pricing: rate resolution, cost computation, seed, and rollup backfill.""" + +from __future__ import annotations + +import json +from dataclasses import dataclass +from datetime import date +from pathlib import Path +from typing import Any, Callable, Dict, List, Optional, Tuple +from uuid import UUID + +from sqlalchemy import text +from sqlalchemy.orm import Session + +from app.services.usage.pricing_cache import ( + get_cached_rate_payload, + pricing_cache_key, + set_cached_rate_payload, +) + +DEFAULT_RATES_EFFECTIVE_FROM = date(2020, 1, 1) +# Back-compat alias for older migrations. +DEFAULT_CATALOG_EFFECTIVE_FROM = DEFAULT_RATES_EFFECTIVE_FROM + +USAGE_KIND_LLM = "llm" +USAGE_KIND_STT = "stt" +USAGE_KIND_TTS = "tts" + +MICRO_USD_PER_UNIT = 1_000_000 +RATE_SOURCE_CATALOG = "catalog" +RATE_SOURCE_OVERRIDE = "override" + +_MODELS_JSON_PATH = ( + Path(__file__).resolve().parent.parent.parent / "config" / "models.json" +) + +_RATES_TABLE = "model_pricing_rates" +_RATES_TABLE_CACHE: Optional[str] = None + + +def clear_rates_table_cache() -> None: + global _RATES_TABLE_CACHE + _RATES_TABLE_CACHE = None + + +def _rates_table(db: Session) -> str: + global _RATES_TABLE_CACHE + if _RATES_TABLE_CACHE: + exists = db.execute( + text("SELECT to_regclass(:table_name)"), + {"table_name": f"public.{_RATES_TABLE_CACHE}"}, + ).scalar() + if exists: + return _RATES_TABLE_CACHE + clear_rates_table_cache() + row = db.execute(text("SELECT to_regclass('public.model_pricing_rates')")).scalar() + if row: + _RATES_TABLE_CACHE = "model_pricing_rates" + return _RATES_TABLE_CACHE + row = db.execute(text("SELECT to_regclass('public.model_pricing_catalog')")).scalar() + if row: + _RATES_TABLE_CACHE = "model_pricing_catalog" + return _RATES_TABLE_CACHE + _RATES_TABLE_CACHE = _RATES_TABLE + return _RATES_TABLE_CACHE + + +@dataclass(frozen=True) +class RateCard: + source: str + rate_id: UUID + input_micro_usd_per_million: int = 0 + output_micro_usd_per_million: int = 0 + cache_read_micro_usd_per_million: int = 0 + cache_creation_micro_usd_per_million: int = 0 + reasoning_micro_usd_per_million: int = 0 + audio_micro_usd_per_second: int = 0 + tts_micro_usd_per_million_chars: int = 0 + + +@dataclass(frozen=True) +class CostBreakdown: + input_cost_micro_usd: int = 0 + output_cost_micro_usd: int = 0 + cache_read_cost_micro_usd: int = 0 + cache_creation_cost_micro_usd: int = 0 + reasoning_cost_micro_usd: int = 0 + audio_cost_micro_usd: int = 0 + tts_cost_micro_usd: int = 0 + total_cost_micro_usd: int = 0 + pricing_rate_source: Optional[str] = None + pricing_rate_id: Optional[UUID] = None + + +@dataclass(frozen=True) +class UsageMetrics: + prompt_tokens: int = 0 + completion_tokens: int = 0 + cache_read_tokens: int = 0 + cache_creation_tokens: int = 0 + reasoning_tokens: int = 0 + audio_seconds: int = 0 + tts_characters: int = 0 + + +def _int(value: Any, default: int = 0) -> int: + try: + return int(value or 0) + except (TypeError, ValueError): + return default + + +def _normalize_rate_source(value: Any) -> str: + source = str(value or "catalog").strip() or "catalog" + if len(source) > 255: + return source[:255] + return source + + +def _usage_kind_for_model_type(model_type: Optional[str]) -> str: + if model_type == "stt": + return USAGE_KIND_STT + if model_type in {"tts", "sts", "sound_effects", "music"}: + return USAGE_KIND_TTS + return USAGE_KIND_LLM + + +def _usd_per_million_to_micro(value: Any) -> int: + if value is None: + return 0 + try: + return int(round(float(value) * MICRO_USD_PER_UNIT)) + except (TypeError, ValueError): + return 0 + + +def _usd_per_minute_to_micro_per_second(value: Any) -> int: + if value is None: + return 0 + try: + return int(round(float(value) * MICRO_USD_PER_UNIT / 60.0)) + except (TypeError, ValueError): + return 0 + + +def _normalize_pricing_block( + pricing: Dict[str, Any], *, model_type: Optional[str] +) -> Dict[str, Any]: + """Convert models.json pricing block (plan USD fields or legacy micro fields) to DB shape.""" + usage_kind = pricing.get("usage_kind") or _usage_kind_for_model_type(model_type) + source = str(pricing.get("source") or pricing.get("_price_source") or "catalog") + + def micro(field_micro: str, field_usd: str) -> int: + if pricing.get(field_micro) is not None: + return _int(pricing.get(field_micro)) + return _usd_per_million_to_micro(pricing.get(field_usd)) + + audio_micro = _int(pricing.get("audio_micro_usd_per_second")) + if not audio_micro: + audio_micro = _usd_per_minute_to_micro_per_second(pricing.get("audio_per_minute")) + + tts_micro = _int(pricing.get("tts_micro_usd_per_million_chars")) + if not tts_micro: + tts_micro = _usd_per_million_to_micro(pricing.get("tts_per_1m_characters")) + + return { + "usage_kind": usage_kind, + "source": source, + "currency": str(pricing.get("currency") or "USD"), + "input_micro_usd_per_million": micro( + "input_micro_usd_per_million", "input_per_1m" + ), + "output_micro_usd_per_million": micro( + "output_micro_usd_per_million", "output_per_1m" + ), + "cache_read_micro_usd_per_million": micro( + "cache_read_micro_usd_per_million", "cache_read_per_1m" + ), + "cache_creation_micro_usd_per_million": micro( + "cache_creation_micro_usd_per_million", "cache_write_per_1m" + ), + "reasoning_micro_usd_per_million": micro( + "reasoning_micro_usd_per_million", "reasoning_per_1m" + ), + "audio_micro_usd_per_second": audio_micro, + "tts_micro_usd_per_million_chars": tts_micro, + } + + +def _catalog_lookup_models(model: str, usage_kind: str) -> Tuple[str, ...]: + candidates: List[str] = [] + for item in ( + model, + f"azure-{model}" if not model.startswith("azure-") else model[len("azure-") :], + ): + if item and item not in candidates: + candidates.append(item) + return tuple(candidates) + + +def _scaled_cost(units: int, rate_per_million: int) -> int: + if not units or not rate_per_million: + return 0 + return (int(units) * int(rate_per_million)) // MICRO_USD_PER_UNIT + + +def metrics_from_deltas(deltas: Dict[str, Any]) -> UsageMetrics: + return UsageMetrics( + prompt_tokens=_int(deltas.get("prompt_tokens")), + completion_tokens=_int(deltas.get("completion_tokens")), + cache_read_tokens=_int(deltas.get("cache_read_tokens")), + cache_creation_tokens=_int(deltas.get("cache_creation_tokens")), + reasoning_tokens=_int(deltas.get("reasoning_tokens")), + audio_seconds=_int(deltas.get("audio_seconds")), + tts_characters=_int(deltas.get("tts_characters")), + ) + + +def cost_fields_from_deltas( + deltas: Dict[str, Any], + *, + organization_id: UUID, + model: str, + usage_kind: str, + usage_date: date, + db: Session, + resolver: Optional["PricingResolver"] = None, +) -> Dict[str, Any]: + """Compute persisted cost columns for one pending-buffer delta row.""" + pricing = resolver or PricingResolver(db) + rate = pricing.resolve_rate( + organization_id=organization_id, + model=model, + usage_kind=usage_kind, + usage_date=usage_date, + ) + costs = compute_cost(metrics_from_deltas(deltas), rate) + return { + "input_cost_micro_usd": costs.input_cost_micro_usd, + "output_cost_micro_usd": costs.output_cost_micro_usd, + "cache_read_cost_micro_usd": costs.cache_read_cost_micro_usd, + "cache_creation_cost_micro_usd": costs.cache_creation_cost_micro_usd, + "reasoning_cost_micro_usd": costs.reasoning_cost_micro_usd, + "audio_cost_micro_usd": costs.audio_cost_micro_usd, + "tts_cost_micro_usd": costs.tts_cost_micro_usd, + "total_cost_micro_usd": costs.total_cost_micro_usd, + "pricing_rate_source": costs.pricing_rate_source, + "pricing_rate_id": str(costs.pricing_rate_id) + if costs.pricing_rate_id + else None, + } + + +def compute_cost(metrics: UsageMetrics, rate: Optional[RateCard]) -> CostBreakdown: + if rate is None: + return CostBreakdown() + + input_cost = _scaled_cost(metrics.prompt_tokens, rate.input_micro_usd_per_million) + output_cost = _scaled_cost( + metrics.completion_tokens, rate.output_micro_usd_per_million + ) + cache_read_cost = _scaled_cost( + metrics.cache_read_tokens, rate.cache_read_micro_usd_per_million + ) + cache_creation_cost = _scaled_cost( + metrics.cache_creation_tokens, rate.cache_creation_micro_usd_per_million + ) + reasoning_cost = _scaled_cost( + metrics.reasoning_tokens, rate.reasoning_micro_usd_per_million + ) + audio_cost = 0 + if metrics.audio_seconds and rate.audio_micro_usd_per_second: + audio_cost = int(metrics.audio_seconds) * int(rate.audio_micro_usd_per_second) + tts_cost = _scaled_cost( + metrics.tts_characters, rate.tts_micro_usd_per_million_chars + ) + total = ( + input_cost + + output_cost + + cache_read_cost + + cache_creation_cost + + reasoning_cost + + audio_cost + + tts_cost + ) + return CostBreakdown( + input_cost_micro_usd=input_cost, + output_cost_micro_usd=output_cost, + cache_read_cost_micro_usd=cache_read_cost, + cache_creation_cost_micro_usd=cache_creation_cost, + reasoning_cost_micro_usd=reasoning_cost, + audio_cost_micro_usd=audio_cost, + tts_cost_micro_usd=tts_cost, + total_cost_micro_usd=total, + pricing_rate_source=rate.source, + pricing_rate_id=rate.rate_id, + ) + + +def _rate_card_from_row(row: Any, *, source: str) -> RateCard: + return RateCard( + source=source, + rate_id=row["id"], + input_micro_usd_per_million=_int(row["input_micro_usd_per_million"]), + output_micro_usd_per_million=_int(row["output_micro_usd_per_million"]), + cache_read_micro_usd_per_million=_int(row["cache_read_micro_usd_per_million"]), + cache_creation_micro_usd_per_million=_int( + row["cache_creation_micro_usd_per_million"] + ), + reasoning_micro_usd_per_million=_int(row["reasoning_micro_usd_per_million"]), + audio_micro_usd_per_second=_int(row["audio_micro_usd_per_second"]), + tts_micro_usd_per_million_chars=_int(row["tts_micro_usd_per_million_chars"]), + ) + + +def _merge_override_with_catalog( + override_row: Any, catalog: Optional[RateCard] +) -> RateCard: + def pick(column: str, attr: str) -> int: + value = override_row.get(column) + if value is not None: + return _int(value) + if catalog is not None: + return getattr(catalog, attr) + return 0 + + return RateCard( + source=RATE_SOURCE_OVERRIDE, + rate_id=override_row["id"], + input_micro_usd_per_million=pick( + "input_micro_usd_per_million", "input_micro_usd_per_million" + ), + output_micro_usd_per_million=pick( + "output_micro_usd_per_million", "output_micro_usd_per_million" + ), + cache_read_micro_usd_per_million=pick( + "cache_read_micro_usd_per_million", "cache_read_micro_usd_per_million" + ), + cache_creation_micro_usd_per_million=pick( + "cache_creation_micro_usd_per_million", + "cache_creation_micro_usd_per_million", + ), + reasoning_micro_usd_per_million=pick( + "reasoning_micro_usd_per_million", "reasoning_micro_usd_per_million" + ), + audio_micro_usd_per_second=pick( + "audio_micro_usd_per_second", "audio_micro_usd_per_second" + ), + tts_micro_usd_per_million_chars=pick( + "tts_micro_usd_per_million_chars", "tts_micro_usd_per_million_chars" + ), + ) + + +def _rate_card_to_cache_payload(card: RateCard) -> Dict[str, Any]: + return { + "source": card.source, + "rate_id": str(card.rate_id), + "input_micro_usd_per_million": card.input_micro_usd_per_million, + "output_micro_usd_per_million": card.output_micro_usd_per_million, + "cache_read_micro_usd_per_million": card.cache_read_micro_usd_per_million, + "cache_creation_micro_usd_per_million": card.cache_creation_micro_usd_per_million, + "reasoning_micro_usd_per_million": card.reasoning_micro_usd_per_million, + "audio_micro_usd_per_second": card.audio_micro_usd_per_second, + "tts_micro_usd_per_million_chars": card.tts_micro_usd_per_million_chars, + } + + +def _rate_card_from_cache_payload(payload: Dict[str, Any]) -> RateCard: + return RateCard( + source=str(payload["source"]), + rate_id=UUID(str(payload["rate_id"])), + input_micro_usd_per_million=_int(payload.get("input_micro_usd_per_million")), + output_micro_usd_per_million=_int(payload.get("output_micro_usd_per_million")), + cache_read_micro_usd_per_million=_int( + payload.get("cache_read_micro_usd_per_million") + ), + cache_creation_micro_usd_per_million=_int( + payload.get("cache_creation_micro_usd_per_million") + ), + reasoning_micro_usd_per_million=_int( + payload.get("reasoning_micro_usd_per_million") + ), + audio_micro_usd_per_second=_int(payload.get("audio_micro_usd_per_second")), + tts_micro_usd_per_million_chars=_int( + payload.get("tts_micro_usd_per_million_chars") + ), + ) + + +class PricingResolver: + """Cached pricing lookups for flush/recompute batches.""" + + def __init__(self, db: Session): + self._db = db + self._memory_cache: Dict[Tuple[str, ...], Optional[RateCard]] = {} + + def resolve_rate( + self, + *, + organization_id: UUID, + model: str, + usage_kind: str, + usage_date: date, + ) -> Optional[RateCard]: + kind = usage_kind or USAGE_KIND_LLM + memory_key = (str(organization_id), model, kind, usage_date.isoformat()) + if memory_key in self._memory_cache: + return self._memory_cache[memory_key] + + redis_key = pricing_cache_key( + organization_id=organization_id, + model=model, + usage_kind=kind, + usage_date=usage_date, + ) + cached = get_cached_rate_payload(redis_key) + if cached: + card = _rate_card_from_cache_payload(cached) + self._memory_cache[memory_key] = card + return card + + override = self._load_override_rate(organization_id, model, kind, usage_date) + if override is not None: + set_cached_rate_payload(redis_key, _rate_card_to_cache_payload(override)) + self._memory_cache[memory_key] = override + return override + + catalog = self._resolve_catalog(model, kind, usage_date) + set_cached_rate_payload( + redis_key, + _rate_card_to_cache_payload(catalog) if catalog else None, + ) + self._memory_cache[memory_key] = catalog + return catalog + + def _resolve_catalog( + self, model: str, usage_kind: str, usage_date: date + ) -> Optional[RateCard]: + for candidate in _catalog_lookup_models(model, usage_kind): + card = self._load_catalog_rate(candidate, usage_kind, usage_date) + if card is not None: + return card + return None + + def _load_catalog_rate( + self, model: str, usage_kind: str, usage_date: date + ) -> Optional[RateCard]: + table = _rates_table(self._db) + row = self._db.execute( + text( + f""" + SELECT + id, + input_micro_usd_per_million, + output_micro_usd_per_million, + cache_read_micro_usd_per_million, + cache_creation_micro_usd_per_million, + reasoning_micro_usd_per_million, + audio_micro_usd_per_second, + tts_micro_usd_per_million_chars + FROM {table} + WHERE model = :model + AND usage_kind = :usage_kind + AND effective_from <= :usage_date + AND (effective_to IS NULL OR effective_to >= :usage_date) + ORDER BY effective_from DESC + LIMIT 1 + """ + ), + { + "model": model, + "usage_kind": usage_kind, + "usage_date": usage_date.isoformat(), + }, + ).mappings().first() + if not row: + return None + return _rate_card_from_row(row, source=RATE_SOURCE_CATALOG) + + def _load_override_rate( + self, + organization_id: UUID, + model: str, + usage_kind: str, + usage_date: date, + ) -> Optional[RateCard]: + row = self._db.execute( + text( + """ + SELECT + id, + input_micro_usd_per_million, + output_micro_usd_per_million, + cache_read_micro_usd_per_million, + cache_creation_micro_usd_per_million, + reasoning_micro_usd_per_million, + audio_micro_usd_per_second, + tts_micro_usd_per_million_chars + FROM org_model_pricing_overrides + WHERE organization_id = CAST(:organization_id AS uuid) + AND model = :model + AND usage_kind = :usage_kind + AND effective_from <= :usage_date + AND (effective_to IS NULL OR effective_to >= :usage_date) + ORDER BY effective_from DESC + LIMIT 1 + """ + ), + { + "organization_id": str(organization_id), + "model": model, + "usage_kind": usage_kind, + "usage_date": usage_date.isoformat(), + }, + ).mappings().first() + if not row: + return None + catalog = self._resolve_catalog(model, usage_kind, usage_date) + return _merge_override_with_catalog(row, catalog) + + +def _pricing_entries_from_models_json() -> Dict[str, Dict[str, Any]]: + if not _MODELS_JSON_PATH.exists(): + return {} + try: + with open(_MODELS_JSON_PATH, "r", encoding="utf-8") as handle: + payload = json.load(handle) + except (OSError, json.JSONDecodeError): + return {} + entries: Dict[str, Dict[str, Any]] = {} + for model_name, config in payload.items(): + if model_name.startswith("_") or not isinstance(config, dict): + continue + pricing = config.get("pricing") + if not isinstance(pricing, dict): + continue + entries[model_name] = _normalize_pricing_block( + pricing, model_type=config.get("model_type") + ) + return entries + + +def seed_pricing_rates(db: Session, *, effective_from: Optional[date] = None) -> int: + """Upsert global rates from models.json pricing blocks (seed only; DB is runtime truth).""" + day = effective_from or DEFAULT_RATES_EFFECTIVE_FROM + table = _rates_table(db) + has_currency = ( + db.execute( + text( + """ + SELECT 1 FROM information_schema.columns + WHERE table_name = :table_name AND column_name = 'currency' + """ + ), + {"table_name": table}, + ).first() + is not None + ) + inserted = 0 + for model_name, pricing in _pricing_entries_from_models_json().items(): + usage_kind = pricing.get("usage_kind") or USAGE_KIND_LLM + base_params = { + "model": model_name, + "usage_kind": usage_kind, + "effective_from": day.isoformat(), + "input_micro_usd_per_million": _int( + pricing.get("input_micro_usd_per_million") + ), + "output_micro_usd_per_million": _int( + pricing.get("output_micro_usd_per_million") + ), + "cache_read_micro_usd_per_million": _int( + pricing.get("cache_read_micro_usd_per_million") + ), + "cache_creation_micro_usd_per_million": _int( + pricing.get("cache_creation_micro_usd_per_million") + ), + "reasoning_micro_usd_per_million": _int( + pricing.get("reasoning_micro_usd_per_million") + ), + "audio_micro_usd_per_second": _int( + pricing.get("audio_micro_usd_per_second") + ), + "tts_micro_usd_per_million_chars": _int( + pricing.get("tts_micro_usd_per_million_chars") + ), + } + if has_currency: + sql = f""" + INSERT INTO {table} ( + id, model, usage_kind, effective_from, currency, source, + input_micro_usd_per_million, + output_micro_usd_per_million, + cache_read_micro_usd_per_million, + cache_creation_micro_usd_per_million, + reasoning_micro_usd_per_million, + audio_micro_usd_per_second, + tts_micro_usd_per_million_chars, + created_at, updated_at + ) VALUES ( + gen_random_uuid(), :model, :usage_kind, CAST(:effective_from AS date), + :currency, :source, + :input_micro_usd_per_million, + :output_micro_usd_per_million, + :cache_read_micro_usd_per_million, + :cache_creation_micro_usd_per_million, + :reasoning_micro_usd_per_million, + :audio_micro_usd_per_second, + :tts_micro_usd_per_million_chars, + now(), now() + ) + ON CONFLICT (model, usage_kind, effective_from) DO UPDATE SET + currency = EXCLUDED.currency, + source = EXCLUDED.source, + input_micro_usd_per_million = EXCLUDED.input_micro_usd_per_million, + output_micro_usd_per_million = EXCLUDED.output_micro_usd_per_million, + cache_read_micro_usd_per_million = EXCLUDED.cache_read_micro_usd_per_million, + cache_creation_micro_usd_per_million = EXCLUDED.cache_creation_micro_usd_per_million, + reasoning_micro_usd_per_million = EXCLUDED.reasoning_micro_usd_per_million, + audio_micro_usd_per_second = EXCLUDED.audio_micro_usd_per_second, + tts_micro_usd_per_million_chars = EXCLUDED.tts_micro_usd_per_million_chars, + updated_at = now() + """ + params = { + **base_params, + "currency": pricing.get("currency") or "USD", + "source": _normalize_rate_source(pricing.get("source")), + } + else: + sql = f""" + INSERT INTO {table} ( + id, model, usage_kind, effective_from, + input_micro_usd_per_million, + output_micro_usd_per_million, + cache_read_micro_usd_per_million, + cache_creation_micro_usd_per_million, + reasoning_micro_usd_per_million, + audio_micro_usd_per_second, + tts_micro_usd_per_million_chars, + created_at, updated_at + ) VALUES ( + gen_random_uuid(), :model, :usage_kind, CAST(:effective_from AS date), + :input_micro_usd_per_million, + :output_micro_usd_per_million, + :cache_read_micro_usd_per_million, + :cache_creation_micro_usd_per_million, + :reasoning_micro_usd_per_million, + :audio_micro_usd_per_second, + :tts_micro_usd_per_million_chars, + now(), now() + ) + ON CONFLICT (model, usage_kind, effective_from) DO UPDATE SET + input_micro_usd_per_million = EXCLUDED.input_micro_usd_per_million, + output_micro_usd_per_million = EXCLUDED.output_micro_usd_per_million, + cache_read_micro_usd_per_million = EXCLUDED.cache_read_micro_usd_per_million, + cache_creation_micro_usd_per_million = EXCLUDED.cache_creation_micro_usd_per_million, + reasoning_micro_usd_per_million = EXCLUDED.reasoning_micro_usd_per_million, + audio_micro_usd_per_second = EXCLUDED.audio_micro_usd_per_second, + tts_micro_usd_per_million_chars = EXCLUDED.tts_micro_usd_per_million_chars, + updated_at = now() + """ + params = base_params + result = db.execute(text(sql), params) + if result.rowcount: + inserted += 1 + if inserted: + from app.services.usage.pricing_cache import invalidate_all_pricing_cache + + invalidate_all_pricing_cache() + return inserted + + +def seed_pricing_catalog(db: Session, *, effective_from: Optional[date] = None) -> int: + """Back-compat alias.""" + return seed_pricing_rates(db, effective_from=effective_from) + + +def apply_cost_to_bucket( + db: Session, + *, + organization_id: UUID, + bucket: Dict[str, Any], + resolver: Optional[PricingResolver] = None, +) -> bool: + """Recompute and persist cost columns for one rollup bucket from row totals.""" + context = bucket.get("context") or {} + workspace_id = bucket.get("workspace_id") + usage_kind = bucket.get("usage_kind") or USAGE_KIND_LLM + usage_date = bucket["usage_date"] + if isinstance(usage_date, str): + usage_date = date.fromisoformat(usage_date) + + base_params = { + "organization_id": str(organization_id), + "workspace_id": str(workspace_id) if workspace_id else None, + "product_section": bucket["product_section"], + "model": bucket["model"], + "usage_date": usage_date.isoformat(), + "usage_kind": usage_kind, + } + exact_params = { + **base_params, + "context": json.dumps(context), + "context_resource_id": str(context.get("resource_id") or ""), + "context_resource_type": str(context.get("resource_type") or ""), + } + row = db.execute( + text( + """ + SELECT + prompt_tokens, completion_tokens, cache_read_tokens, + cache_creation_tokens, reasoning_tokens, audio_seconds, + tts_characters + FROM llm_usage_daily + WHERE organization_id = CAST(:organization_id AS uuid) + AND product_section = :product_section + AND model = :model + AND usage_date = CAST(:usage_date AS date) + AND usage_kind = :usage_kind + AND workspace_id IS NOT DISTINCT FROM CAST(:workspace_id AS uuid) + AND context = CAST(:context AS jsonb) + """ + ), + exact_params, + ).mappings().first() + if not row: + row = db.execute( + text( + """ + SELECT + prompt_tokens, completion_tokens, cache_read_tokens, + cache_creation_tokens, reasoning_tokens, audio_seconds, + tts_characters + FROM llm_usage_daily + WHERE organization_id = CAST(:organization_id AS uuid) + AND product_section = :product_section + AND model = :model + AND usage_date = CAST(:usage_date AS date) + AND usage_kind = :usage_kind + AND workspace_id IS NOT DISTINCT FROM CAST(:workspace_id AS uuid) + AND COALESCE(context->>'resource_id', '') = :context_resource_id + AND COALESCE(context->>'resource_type', '') = :context_resource_type + """ + ), + exact_params, + ).mappings().first() + if not row: + return False + + pricing = resolver or PricingResolver(db) + rate = pricing.resolve_rate( + organization_id=organization_id, + model=bucket["model"], + usage_kind=usage_kind, + usage_date=usage_date, + ) + costs = compute_cost( + UsageMetrics( + prompt_tokens=_int(row["prompt_tokens"]), + completion_tokens=_int(row["completion_tokens"]), + cache_read_tokens=_int(row["cache_read_tokens"]), + cache_creation_tokens=_int(row["cache_creation_tokens"]), + reasoning_tokens=_int(row["reasoning_tokens"]), + audio_seconds=_int(row["audio_seconds"]), + tts_characters=_int(row["tts_characters"]), + ), + rate, + ) + update_params = { + **exact_params, + "input_cost_micro_usd": costs.input_cost_micro_usd, + "output_cost_micro_usd": costs.output_cost_micro_usd, + "cache_read_cost_micro_usd": costs.cache_read_cost_micro_usd, + "cache_creation_cost_micro_usd": costs.cache_creation_cost_micro_usd, + "reasoning_cost_micro_usd": costs.reasoning_cost_micro_usd, + "audio_cost_micro_usd": costs.audio_cost_micro_usd, + "tts_cost_micro_usd": costs.tts_cost_micro_usd, + "total_cost_micro_usd": costs.total_cost_micro_usd, + "pricing_rate_source": costs.pricing_rate_source, + "pricing_rate_id": str(costs.pricing_rate_id) + if costs.pricing_rate_id + else None, + } + result = db.execute( + text( + """ + UPDATE llm_usage_daily SET + input_cost_micro_usd = :input_cost_micro_usd, + output_cost_micro_usd = :output_cost_micro_usd, + cache_read_cost_micro_usd = :cache_read_cost_micro_usd, + cache_creation_cost_micro_usd = :cache_creation_cost_micro_usd, + reasoning_cost_micro_usd = :reasoning_cost_micro_usd, + audio_cost_micro_usd = :audio_cost_micro_usd, + tts_cost_micro_usd = :tts_cost_micro_usd, + total_cost_micro_usd = :total_cost_micro_usd, + pricing_rate_source = :pricing_rate_source, + pricing_rate_id = CAST(:pricing_rate_id AS uuid), + updated_at = now() + WHERE organization_id = CAST(:organization_id AS uuid) + AND product_section = :product_section + AND model = :model + AND usage_date = CAST(:usage_date AS date) + AND usage_kind = :usage_kind + AND workspace_id IS NOT DISTINCT FROM CAST(:workspace_id AS uuid) + AND context = CAST(:context AS jsonb) + """ + ), + update_params, + ) + if result.rowcount: + return True + result = db.execute( + text( + """ + UPDATE llm_usage_daily SET + input_cost_micro_usd = :input_cost_micro_usd, + output_cost_micro_usd = :output_cost_micro_usd, + cache_read_cost_micro_usd = :cache_read_cost_micro_usd, + cache_creation_cost_micro_usd = :cache_creation_cost_micro_usd, + reasoning_cost_micro_usd = :reasoning_cost_micro_usd, + audio_cost_micro_usd = :audio_cost_micro_usd, + tts_cost_micro_usd = :tts_cost_micro_usd, + total_cost_micro_usd = :total_cost_micro_usd, + pricing_rate_source = :pricing_rate_source, + pricing_rate_id = CAST(:pricing_rate_id AS uuid), + updated_at = now() + WHERE organization_id = CAST(:organization_id AS uuid) + AND product_section = :product_section + AND model = :model + AND usage_date = CAST(:usage_date AS date) + AND usage_kind = :usage_kind + AND workspace_id IS NOT DISTINCT FROM CAST(:workspace_id AS uuid) + AND COALESCE(context->>'resource_id', '') = :context_resource_id + AND COALESCE(context->>'resource_type', '') = :context_resource_type + """ + ), + update_params, + ) + return bool(result.rowcount) + + +def recompute_usage_costs( + db: Session, + *, + organization_id: Optional[UUID] = None, + model: Optional[str] = None, + usage_kind: Optional[str] = None, + start_date: Optional[date] = None, + end_date: Optional[date] = None, + batch_size: int = 500, + on_progress: Optional[Callable[[int], None]] = None, +) -> int: + """Recompute stored costs for existing rollup rows (backfill / override changes).""" + resolver = PricingResolver(db) + updated = 0 + last_id: Optional[str] = None + + while True: + params: Dict[str, Any] = {"batch_size": batch_size} + filters = ["1=1"] + if organization_id is not None: + filters.append("organization_id = CAST(:organization_id AS uuid)") + params["organization_id"] = str(organization_id) + if model is not None: + filters.append("model = :model") + params["model"] = model + if usage_kind is not None: + filters.append("usage_kind = :usage_kind") + params["usage_kind"] = usage_kind + if start_date is not None: + filters.append("usage_date >= CAST(:start_date AS date)") + params["start_date"] = start_date.isoformat() + if end_date is not None: + filters.append("usage_date <= CAST(:end_date AS date)") + params["end_date"] = end_date.isoformat() + if last_id is not None: + filters.append("id > CAST(:last_id AS uuid)") + params["last_id"] = last_id + + rows = db.execute( + text( + f""" + SELECT + id, organization_id, workspace_id, product_section, model, + context, usage_date, usage_kind + FROM llm_usage_daily + WHERE {' AND '.join(filters)} + ORDER BY id + LIMIT :batch_size + """ + ), + params, + ).mappings().all() + if not rows: + break + + for row in rows: + bucket = { + "workspace_id": row["workspace_id"], + "product_section": row["product_section"], + "model": row["model"], + "context": row["context"] or {}, + "usage_date": row["usage_date"], + "usage_kind": row["usage_kind"] or USAGE_KIND_LLM, + } + if apply_cost_to_bucket( + db, + organization_id=row["organization_id"], + bucket=bucket, + resolver=resolver, + ): + updated += 1 + last_id = str(row["id"]) + + db.commit() + if on_progress is not None: + on_progress(updated) + + if on_progress is not None: + on_progress(updated) + + return updated diff --git a/app/services/usage/pricing_cache.py b/app/services/usage/pricing_cache.py new file mode 100644 index 00000000..dd2ebb0f --- /dev/null +++ b/app/services/usage/pricing_cache.py @@ -0,0 +1,87 @@ +"""Redis cache for resolved usage pricing rates.""" + +from __future__ import annotations + +import json +from datetime import date +from typing import Optional +from uuid import UUID + +import redis +from loguru import logger + +from app.config import settings + +PRICING_CACHE_TTL_SEC = 3600 +PRICING_NULL_CACHE_TTL_SEC = 300 +PRICING_CACHE_PREFIX = "usage:pricing" + +_redis: redis.Redis | None = None + + +def _client() -> redis.Redis: + global _redis + if _redis is None: + _redis = redis.from_url(settings.REDIS_URL, decode_responses=True) + return _redis + + +def pricing_cache_key( + *, + organization_id: UUID, + model: str, + usage_kind: str, + usage_date: date, +) -> str: + return ( + f"{PRICING_CACHE_PREFIX}:{organization_id}:{model}:" + f"{usage_kind}:{usage_date.isoformat()}" + ) + + +def get_cached_rate_payload(key: str) -> Optional[dict]: + try: + raw = _client().get(key) + except redis.RedisError as exc: + logger.debug("pricing cache read skipped: {}", exc) + return None + if raw is None: + return None + if raw == "__null__": + return {} + try: + payload = json.loads(raw) + except json.JSONDecodeError: + return None + return payload if isinstance(payload, dict) else None + + +def set_cached_rate_payload(key: str, payload: Optional[dict]) -> None: + try: + value = "__null__" if not payload else json.dumps(payload) + ttl = PRICING_NULL_CACHE_TTL_SEC if not payload else PRICING_CACHE_TTL_SEC + _client().setex(key, ttl, value) + except redis.RedisError as exc: + logger.debug("pricing cache write skipped: {}", exc) + + +def invalidate_org_pricing_cache(organization_id: UUID) -> None: + _invalidate_pricing_cache_pattern(f"{PRICING_CACHE_PREFIX}:{organization_id}:*") + + +def invalidate_all_pricing_cache() -> None: + _invalidate_pricing_cache_pattern(f"{PRICING_CACHE_PREFIX}:*") + + +def _invalidate_pricing_cache_pattern(pattern: str) -> None: + try: + client = _client() + cursor = 0 + while True: + cursor, keys = client.scan(cursor=cursor, match=pattern, count=200) + if keys: + client.delete(*keys) + if cursor == 0: + break + except redis.RedisError as exc: + logger.debug("pricing cache invalidate skipped: {}", exc) diff --git a/app/services/usage/pricing_jobs.py b/app/services/usage/pricing_jobs.py new file mode 100644 index 00000000..5e01947a --- /dev/null +++ b/app/services/usage/pricing_jobs.py @@ -0,0 +1,145 @@ +"""Usage cost recompute job lifecycle.""" + +from __future__ import annotations + +from datetime import date, datetime, timezone +from typing import Any, Dict, Optional +from uuid import UUID + +from fastapi import HTTPException +from sqlalchemy.orm import Session + +from app.models.database import UsageCostRecomputeJob + +ACTIVE_JOB_STATUSES = ("pending", "running") +TERMINAL_JOB_STATUSES = ("completed", "failed") + + +def _utcnow() -> datetime: + return datetime.now(timezone.utc) + + +def get_recompute_job( + db: Session, + *, + organization_id: UUID, + job_id: UUID, +) -> UsageCostRecomputeJob: + job = ( + db.query(UsageCostRecomputeJob) + .filter( + UsageCostRecomputeJob.id == job_id, + UsageCostRecomputeJob.organization_id == organization_id, + ) + .first() + ) + if job is None: + raise HTTPException(status_code=404, detail="Recompute job not found") + return job + + +def create_recompute_job( + db: Session, + *, + organization_id: UUID, + model: Optional[str] = None, + usage_kind: Optional[str] = None, + start_date: Optional[date] = None, + end_date: Optional[date] = None, +) -> UsageCostRecomputeJob: + active = ( + db.query(UsageCostRecomputeJob) + .filter( + UsageCostRecomputeJob.organization_id == organization_id, + UsageCostRecomputeJob.status.in_(ACTIVE_JOB_STATUSES), + ) + .first() + ) + if active is not None: + raise HTTPException( + status_code=409, + detail="A usage cost recompute job is already in progress", + ) + + job = UsageCostRecomputeJob( + organization_id=organization_id, + status="pending", + model=model, + usage_kind=usage_kind, + start_date=start_date, + end_date=end_date, + ) + db.add(job) + db.commit() + db.refresh(job) + return job + + +def enqueue_recompute_job(db: Session, job: UsageCostRecomputeJob) -> str: + from app.workers.tasks import recompute_usage_costs_task + + result = recompute_usage_costs_task.delay(job_id=str(job.id)) + job.celery_task_id = result.id + job.updated_at = _utcnow() + db.commit() + return result.id + + +def mark_job_running(db: Session, job_id: UUID) -> None: + job = db.query(UsageCostRecomputeJob).filter(UsageCostRecomputeJob.id == job_id).first() + if job is None: + return + job.status = "running" + job.updated_at = _utcnow() + db.commit() + + +def update_job_progress(db: Session, job_id: UUID, updated_rows: int) -> None: + job = db.query(UsageCostRecomputeJob).filter(UsageCostRecomputeJob.id == job_id).first() + if job is None: + return + job.updated_rows = updated_rows + job.updated_at = _utcnow() + db.commit() + + +def mark_job_completed(db: Session, job_id: UUID, updated_rows: int) -> None: + job = db.query(UsageCostRecomputeJob).filter(UsageCostRecomputeJob.id == job_id).first() + if job is None: + return + now = _utcnow() + job.status = "completed" + job.updated_rows = updated_rows + job.updated_at = now + job.completed_at = now + db.commit() + + +def mark_job_failed(db: Session, job_id: UUID, error_message: str) -> None: + job = db.query(UsageCostRecomputeJob).filter(UsageCostRecomputeJob.id == job_id).first() + if job is None: + return + now = _utcnow() + job.status = "failed" + job.error_message = error_message[:4000] + job.updated_at = now + job.completed_at = now + db.commit() + + +def job_to_dict(job: UsageCostRecomputeJob) -> Dict[str, Any]: + return { + "id": job.id, + "organization_id": job.organization_id, + "status": job.status, + "model": job.model, + "usage_kind": job.usage_kind, + "start_date": job.start_date, + "end_date": job.end_date, + "updated_rows": int(job.updated_rows or 0), + "error_message": job.error_message, + "celery_task_id": job.celery_task_id, + "created_at": job.created_at, + "updated_at": job.updated_at, + "completed_at": job.completed_at, + } diff --git a/app/services/usage/pricing_ops.py b/app/services/usage/pricing_ops.py new file mode 100644 index 00000000..abd55d4b --- /dev/null +++ b/app/services/usage/pricing_ops.py @@ -0,0 +1,137 @@ +"""Ops helpers for usage pricing catalog maintenance.""" + +from __future__ import annotations + +import json +from datetime import date +from pathlib import Path +from typing import Any, Dict, List, Optional, Tuple + +from sqlalchemy import text +from sqlalchemy.orm import Session + +from app.services.usage.pricing import ( + DEFAULT_RATES_EFFECTIVE_FROM, + _pricing_entries_from_models_json, + _rates_table, + seed_pricing_rates, +) + +_MODELS_JSON_PATH = ( + Path(__file__).resolve().parent.parent.parent / "config" / "models.json" +) +_CATALOG_JSON_PATH = ( + Path(__file__).resolve().parent.parent.parent / "config" / "pricing_catalog.json" +) + +_RATE_COMPARE_COLUMNS = ( + "input_micro_usd_per_million", + "output_micro_usd_per_million", + "cache_read_micro_usd_per_million", + "cache_creation_micro_usd_per_million", + "reasoning_micro_usd_per_million", + "audio_micro_usd_per_second", + "tts_micro_usd_per_million_chars", +) + + +def models_missing_pricing_blocks() -> List[str]: + if not _MODELS_JSON_PATH.exists(): + return [] + payload = json.loads(_MODELS_JSON_PATH.read_text(encoding="utf-8")) + missing: List[str] = [] + for model_name, config in payload.items(): + if model_name.startswith("_") or not isinstance(config, dict): + continue + pricing = config.get("pricing") + if not isinstance(pricing, dict): + missing.append(model_name) + return sorted(missing) + + +def litellm_unresolved_models() -> List[Dict[str, Any]]: + if not _CATALOG_JSON_PATH.exists(): + return [] + payload = json.loads(_CATALOG_JSON_PATH.read_text(encoding="utf-8")) + meta = payload.get("_metadata") + if not isinstance(meta, dict): + return [] + unresolved = meta.get("unresolved") + return unresolved if isinstance(unresolved, list) else [] + + +def _load_db_rates( + db: Session, *, effective_from: date +) -> Dict[Tuple[str, str], Dict[str, Any]]: + table = _rates_table(db) + rows = db.execute( + text( + f""" + SELECT model, usage_kind, {_rate_columns_sql()} + FROM {table} + WHERE effective_from = CAST(:effective_from AS date) + """ + ), + {"effective_from": effective_from.isoformat()}, + ).mappings().all() + return {(row["model"], row["usage_kind"]): dict(row) for row in rows} + + +def _rate_columns_sql() -> str: + return ", ".join(_RATE_COMPARE_COLUMNS) + + +def diff_models_json_vs_db( + db: Session, *, effective_from: Optional[date] = None +) -> Dict[str, Any]: + day = effective_from or DEFAULT_RATES_EFFECTIVE_FROM + json_rates = _pricing_entries_from_models_json() + db_rates = _load_db_rates(db, effective_from=day) + + json_keys = {(model, entry.get("usage_kind") or "llm") for model, entry in json_rates.items()} + db_keys = set(db_rates.keys()) + + only_in_json = sorted(json_keys - db_keys) + only_in_db = sorted(db_keys - json_keys) + mismatches: List[Dict[str, Any]] = [] + + for key in sorted(json_keys & db_keys): + model, usage_kind = key + expected = json_rates[model] + actual = db_rates[key] + field_diffs: Dict[str, Dict[str, int]] = {} + for column in _RATE_COMPARE_COLUMNS: + left = int(expected.get(column) or 0) + right = int(actual.get(column) or 0) + if left != right: + field_diffs[column] = {"models_json": left, "database": right} + if field_diffs: + mismatches.append( + { + "model": model, + "usage_kind": usage_kind, + "fields": field_diffs, + } + ) + + return { + "effective_from": day.isoformat(), + "models_json_count": len(json_rates), + "database_count": len(db_rates), + "only_in_models_json": [ + {"model": model, "usage_kind": kind} for model, kind in only_in_json + ], + "only_in_database": [ + {"model": model, "usage_kind": kind} for model, kind in only_in_db + ], + "mismatches": mismatches, + "missing_pricing_blocks": models_missing_pricing_blocks(), + "litellm_unresolved": litellm_unresolved_models(), + "in_sync": not only_in_json and not only_in_db and not mismatches, + } + + +def seed_rates_from_models_json( + db: Session, *, effective_from: Optional[date] = None +) -> int: + return seed_pricing_rates(db, effective_from=effective_from) diff --git a/app/services/usage/pricing_overrides.py b/app/services/usage/pricing_overrides.py new file mode 100644 index 00000000..77b84c2b --- /dev/null +++ b/app/services/usage/pricing_overrides.py @@ -0,0 +1,542 @@ +"""Org-level usage pricing overrides.""" + +from __future__ import annotations + +from datetime import date, timedelta +from typing import Any, Dict, List, Optional, Set +from uuid import UUID + +from fastapi import HTTPException +from sqlalchemy import text +from sqlalchemy.orm import Session + +from app.services.usage.pricing import ( + PricingResolver, + USAGE_KIND_LLM, + USAGE_KIND_STT, + USAGE_KIND_TTS, + _int, + _pricing_entries_from_models_json, + _rates_table, + _usd_per_million_to_micro, + _usd_per_minute_to_micro_per_second, +) +from app.services.usage.pricing_cache import invalidate_org_pricing_cache +from app.services.usage.pricing_jobs import create_recompute_job, enqueue_recompute_job +from app.services.usage.usage_costs import micro_to_usd + +RATE_COLUMNS = ( + "input_micro_usd_per_million", + "output_micro_usd_per_million", + "cache_read_micro_usd_per_million", + "cache_creation_micro_usd_per_million", + "reasoning_micro_usd_per_million", + "audio_micro_usd_per_second", + "tts_micro_usd_per_million_chars", +) + +USD_RATE_FIELDS = { + "input_per_1m": "input_micro_usd_per_million", + "output_per_1m": "output_micro_usd_per_million", + "cache_read_per_1m": "cache_read_micro_usd_per_million", + "cache_write_per_1m": "cache_creation_micro_usd_per_million", + "reasoning_per_1m": "reasoning_micro_usd_per_million", + "audio_per_minute": "audio_micro_usd_per_second", + "tts_per_1m_characters": "tts_micro_usd_per_million_chars", +} + + +def _known_models(db: Session) -> Set[str]: + models = set(_pricing_entries_from_models_json().keys()) + table = _rates_table(db) + rows = db.execute(text(f"SELECT DISTINCT model FROM {table}")).scalars().all() + models.update(rows) + return models + + +def validate_model_name( + db: Session, + model: str, + *, + organization_id: Optional[UUID] = None, +) -> None: + from app.services.usage.enabled_models import org_pricing_eligible_models + + if organization_id is not None: + eligible = set(org_pricing_eligible_models(db, organization_id)) + if model in eligible: + return + if model in _known_models(db): + return + raise HTTPException(status_code=400, detail=f"Unknown model: {model}") + + +def _validate_usage_kind(usage_kind: str) -> str: + kind = usage_kind or USAGE_KIND_LLM + if kind not in {USAGE_KIND_LLM, USAGE_KIND_STT, USAGE_KIND_TTS}: + raise HTTPException(status_code=400, detail=f"Invalid usage_kind: {kind}") + return kind + + +def _micro_to_optional_usd_per_1m(micro: Optional[int]) -> Optional[float]: + if micro is None: + return None + return micro_to_usd(_int(micro)) + + +def _micro_to_optional_usd_per_minute(micro_per_second: Optional[int]) -> Optional[float]: + if micro_per_second is None: + return None + return micro_to_usd(_int(micro_per_second) * 60) + + +def _rates_usd_from_row(row: Dict[str, Any]) -> Dict[str, Optional[float]]: + return { + "input_per_1m": _micro_to_optional_usd_per_1m(row.get("input_micro_usd_per_million")), + "output_per_1m": _micro_to_optional_usd_per_1m(row.get("output_micro_usd_per_million")), + "cache_read_per_1m": _micro_to_optional_usd_per_1m( + row.get("cache_read_micro_usd_per_million") + ), + "cache_write_per_1m": _micro_to_optional_usd_per_1m( + row.get("cache_creation_micro_usd_per_million") + ), + "reasoning_per_1m": _micro_to_optional_usd_per_1m( + row.get("reasoning_micro_usd_per_million") + ), + "audio_per_minute": _micro_to_optional_usd_per_minute( + row.get("audio_micro_usd_per_second") + ), + "tts_per_1m_characters": _micro_to_optional_usd_per_1m( + row.get("tts_micro_usd_per_million_chars") + ), + } + + +def _rate_card_to_usd_dict(card) -> Dict[str, float]: + return { + "input_per_1m": micro_to_usd(card.input_micro_usd_per_million), + "output_per_1m": micro_to_usd(card.output_micro_usd_per_million), + "cache_read_per_1m": micro_to_usd(card.cache_read_micro_usd_per_million), + "cache_write_per_1m": micro_to_usd(card.cache_creation_micro_usd_per_million), + "reasoning_per_1m": micro_to_usd(card.reasoning_micro_usd_per_million), + "audio_per_minute": micro_to_usd(card.audio_micro_usd_per_second * 60), + "tts_per_1m_characters": micro_to_usd(card.tts_micro_usd_per_million_chars), + } + + +def _override_row_to_dict(row: Any) -> Dict[str, Any]: + payload = dict(row) + for key in ("id", "organization_id"): + if payload.get(key) is not None: + payload[key] = str(payload[key]) + rates = _rates_usd_from_row(payload) + return { + "id": payload["id"], + "organization_id": payload["organization_id"], + "model": payload["model"], + "usage_kind": payload["usage_kind"], + "effective_from": payload["effective_from"], + "effective_to": payload.get("effective_to"), + "rates": rates, + "created_at": payload.get("created_at"), + "updated_at": payload.get("updated_at"), + } + + +def _usd_payload_to_micro_columns(payload: Dict[str, Any]) -> Dict[str, Optional[int]]: + columns: Dict[str, Optional[int]] = {} + for usd_field, micro_column in USD_RATE_FIELDS.items(): + if usd_field not in payload: + continue + value = payload.get(usd_field) + if value is None: + columns[micro_column] = None + continue + if usd_field == "audio_per_minute": + columns[micro_column] = _usd_per_minute_to_micro_per_second(value) + else: + columns[micro_column] = _usd_per_million_to_micro(value) + return columns + + +def list_overrides( + db: Session, + *, + organization_id: UUID, + model: Optional[str] = None, + usage_kind: Optional[str] = None, +) -> List[Dict[str, Any]]: + filters = ["organization_id = CAST(:organization_id AS uuid)"] + params: Dict[str, Any] = {"organization_id": str(organization_id)} + if model is not None: + filters.append("model = :model") + params["model"] = model + if usage_kind is not None: + filters.append("usage_kind = :usage_kind") + params["usage_kind"] = _validate_usage_kind(usage_kind) + + rows = db.execute( + text( + f""" + SELECT * + FROM org_model_pricing_overrides + WHERE {' AND '.join(filters)} + ORDER BY model ASC, usage_kind ASC, effective_from DESC + """ + ), + params, + ).mappings().all() + return [_override_row_to_dict(row) for row in rows] + + +def get_effective_rate( + db: Session, + *, + organization_id: UUID, + model: str, + usage_kind: str, + as_of: date, +) -> Dict[str, Any]: + validate_model_name(db, model, organization_id=organization_id) + kind = _validate_usage_kind(usage_kind) + resolver = PricingResolver(db) + effective = resolver.resolve_rate( + organization_id=organization_id, + model=model, + usage_kind=kind, + usage_date=as_of, + ) + catalog = resolver._resolve_catalog(model, kind, as_of) + override_row = db.execute( + text( + """ + SELECT * + FROM org_model_pricing_overrides + WHERE organization_id = CAST(:organization_id AS uuid) + AND model = :model + AND usage_kind = :usage_kind + AND effective_from <= :as_of + AND (effective_to IS NULL OR effective_to >= :as_of) + ORDER BY effective_from DESC + LIMIT 1 + """ + ), + { + "organization_id": str(organization_id), + "model": model, + "usage_kind": kind, + "as_of": as_of.isoformat(), + }, + ).mappings().first() + + return { + "model": model, + "usage_kind": kind, + "as_of": as_of, + "catalog_rates": _rate_card_to_usd_dict(catalog) if catalog else None, + "catalog_rate_id": str(catalog.rate_id) if catalog else None, + "override": _override_row_to_dict(override_row) if override_row else None, + "effective_rates": _rate_card_to_usd_dict(effective) if effective else None, + "effective_source": effective.source if effective else None, + "effective_rate_id": str(effective.rate_id) if effective else None, + "has_override": override_row is not None, + } + + +def list_effective_pricing( + db: Session, + *, + organization_id: UUID, + usage_kind: Optional[str] = None, + model: Optional[str] = None, + as_of: date, + limit: int = 200, +) -> List[Dict[str, Any]]: + if model: + return [ + get_effective_rate( + db, + organization_id=organization_id, + model=model, + usage_kind=usage_kind or USAGE_KIND_LLM, + as_of=as_of, + ) + ] + + keys: Set[tuple[str, str]] = set() + for entry in list_overrides(db, organization_id=organization_id, usage_kind=usage_kind): + keys.add((entry["model"], entry["usage_kind"])) + + table = _rates_table(db) + rate_filters = ["effective_from <= :as_of", "(effective_to IS NULL OR effective_to >= :as_of)"] + params: Dict[str, Any] = {"as_of": as_of.isoformat(), "limit": limit} + if usage_kind is not None: + rate_filters.append("usage_kind = :usage_kind") + params["usage_kind"] = _validate_usage_kind(usage_kind) + rate_rows = db.execute( + text( + f""" + SELECT DISTINCT model, usage_kind + FROM {table} + WHERE {' AND '.join(rate_filters)} + ORDER BY model ASC + LIMIT :limit + """ + ), + params, + ).mappings().all() + for row in rate_rows: + keys.add((row["model"], row["usage_kind"])) + + results: List[Dict[str, Any]] = [] + for model_name, kind in sorted(keys)[:limit]: + results.append( + get_effective_rate( + db, + organization_id=organization_id, + model=model_name, + usage_kind=kind, + as_of=as_of, + ) + ) + return results + + +def upsert_override( + db: Session, + *, + organization_id: UUID, + model: str, + usage_kind: str, + effective_from: date, + effective_to: Optional[date] = None, + rates: Dict[str, Any], + recompute: bool = False, +) -> Dict[str, Any]: + validate_model_name(db, model, organization_id=organization_id) + kind = _validate_usage_kind(usage_kind) + if effective_to is not None and effective_to < effective_from: + raise HTTPException(status_code=400, detail="effective_to must be >= effective_from") + + micro_columns = _usd_payload_to_micro_columns(rates) + if not micro_columns: + raise HTTPException(status_code=400, detail="At least one rate field is required") + + params: Dict[str, Any] = { + "organization_id": str(organization_id), + "model": model, + "usage_kind": kind, + "effective_from": effective_from.isoformat(), + "effective_to": effective_to.isoformat() if effective_to else None, + } + for column in RATE_COLUMNS: + params[column] = micro_columns.get(column) + + row = db.execute( + text( + """ + INSERT INTO org_model_pricing_overrides ( + organization_id, model, usage_kind, effective_from, effective_to, + input_micro_usd_per_million, output_micro_usd_per_million, + cache_read_micro_usd_per_million, cache_creation_micro_usd_per_million, + reasoning_micro_usd_per_million, audio_micro_usd_per_second, + tts_micro_usd_per_million_chars + ) VALUES ( + CAST(:organization_id AS uuid), :model, :usage_kind, + CAST(:effective_from AS date), CAST(:effective_to AS date), + :input_micro_usd_per_million, :output_micro_usd_per_million, + :cache_read_micro_usd_per_million, :cache_creation_micro_usd_per_million, + :reasoning_micro_usd_per_million, :audio_micro_usd_per_second, + :tts_micro_usd_per_million_chars + ) + ON CONFLICT (organization_id, model, usage_kind, effective_from) + DO UPDATE SET + effective_to = EXCLUDED.effective_to, + input_micro_usd_per_million = COALESCE( + EXCLUDED.input_micro_usd_per_million, + org_model_pricing_overrides.input_micro_usd_per_million + ), + output_micro_usd_per_million = COALESCE( + EXCLUDED.output_micro_usd_per_million, + org_model_pricing_overrides.output_micro_usd_per_million + ), + cache_read_micro_usd_per_million = COALESCE( + EXCLUDED.cache_read_micro_usd_per_million, + org_model_pricing_overrides.cache_read_micro_usd_per_million + ), + cache_creation_micro_usd_per_million = COALESCE( + EXCLUDED.cache_creation_micro_usd_per_million, + org_model_pricing_overrides.cache_creation_micro_usd_per_million + ), + reasoning_micro_usd_per_million = COALESCE( + EXCLUDED.reasoning_micro_usd_per_million, + org_model_pricing_overrides.reasoning_micro_usd_per_million + ), + audio_micro_usd_per_second = COALESCE( + EXCLUDED.audio_micro_usd_per_second, + org_model_pricing_overrides.audio_micro_usd_per_second + ), + tts_micro_usd_per_million_chars = COALESCE( + EXCLUDED.tts_micro_usd_per_million_chars, + org_model_pricing_overrides.tts_micro_usd_per_million_chars + ), + updated_at = now() + RETURNING * + """ + ), + params, + ).mappings().first() + db.commit() + + invalidate_org_pricing_cache(organization_id) + + recompute_job_id = None + recompute_enqueued = False + if recompute: + recompute_job_id = _enqueue_override_recompute( + db, + organization_id=organization_id, + model=model, + usage_kind=kind, + start_date=effective_from, + end_date=effective_to, + ) + recompute_enqueued = recompute_job_id is not None + + result = _override_row_to_dict(row) + result["recompute_enqueued"] = recompute_enqueued + result["recompute_job_id"] = recompute_job_id + return result + + +def delete_override( + db: Session, + *, + organization_id: UUID, + model: str, + usage_kind: str, + effective_from: Optional[date] = None, + recompute: bool = False, +) -> Dict[str, Any]: + validate_model_name(db, model, organization_id=organization_id) + kind = _validate_usage_kind(usage_kind) + filters = [ + "organization_id = CAST(:organization_id AS uuid)", + "model = :model", + "usage_kind = :usage_kind", + ] + params: Dict[str, Any] = { + "organization_id": str(organization_id), + "model": model, + "usage_kind": kind, + } + + if effective_from is not None: + filters.append("effective_from = CAST(:effective_from AS date)") + params["effective_from"] = effective_from.isoformat() + result = db.execute( + text( + f""" + DELETE FROM org_model_pricing_overrides + WHERE {' AND '.join(filters)} + RETURNING effective_from, effective_to + """ + ), + params, + ).mappings().first() + if not result: + raise HTTPException(status_code=404, detail="Override not found") + recompute_from = effective_from + recompute_to = result.get("effective_to") + db.commit() + else: + active = db.execute( + text( + f""" + SELECT id, effective_from, effective_to + FROM org_model_pricing_overrides + WHERE {' AND '.join(filters)} + AND effective_from <= CURRENT_DATE + AND (effective_to IS NULL OR effective_to >= CURRENT_DATE) + ORDER BY effective_from DESC + LIMIT 1 + """ + ), + params, + ).mappings().first() + if not active: + raise HTTPException(status_code=404, detail="No active override found") + recompute_from = active["effective_from"] + recompute_to = active.get("effective_to") + yesterday = date.today() - timedelta(days=1) + if active["effective_from"] > yesterday: + db.execute( + text( + """ + DELETE FROM org_model_pricing_overrides + WHERE id = CAST(:id AS uuid) + """ + ), + {"id": str(active["id"])}, + ) + else: + db.execute( + text( + """ + UPDATE org_model_pricing_overrides + SET effective_to = CAST(:effective_to AS date), updated_at = now() + WHERE id = CAST(:id AS uuid) + """ + ), + {"id": str(active["id"]), "effective_to": yesterday.isoformat()}, + ) + db.commit() + + invalidate_org_pricing_cache(organization_id) + + recompute_job_id = None + recompute_enqueued = False + if recompute: + recompute_job_id = _enqueue_override_recompute( + db, + organization_id=organization_id, + model=model, + usage_kind=kind, + start_date=recompute_from, + end_date=recompute_to, + ) + recompute_enqueued = recompute_job_id is not None + + return { + "deleted": True, + "model": model, + "usage_kind": kind, + "recompute_enqueued": recompute_enqueued, + "recompute_job_id": recompute_job_id, + } + + +def _enqueue_override_recompute( + db: Session, + *, + organization_id: UUID, + model: str, + usage_kind: str, + start_date: date, + end_date: Optional[date], +) -> Optional[str]: + try: + job = create_recompute_job( + db, + organization_id=organization_id, + model=model, + usage_kind=usage_kind, + start_date=start_date, + end_date=end_date, + ) + enqueue_recompute_job(db, job) + return str(job.id) + except HTTPException as exc: + if exc.status_code == 409: + return None + raise diff --git a/app/services/usage/read_cache.py b/app/services/usage/read_cache.py new file mode 100644 index 00000000..5db010f0 --- /dev/null +++ b/app/services/usage/read_cache.py @@ -0,0 +1,151 @@ +"""Redis cache for org usage read endpoints (summary, breakdown, filters).""" + +from __future__ import annotations + +import hashlib +import json +import os +from typing import Any, Optional +from uuid import UUID + +import redis +from loguru import logger + +from app.config import settings +from app.services.usage.access import UsageAccessResult + +_redis: redis.Redis | None = None + + +def _client() -> redis.Redis: + global _redis + if _redis is None: + _redis = redis.from_url(settings.REDIS_URL, decode_responses=True) + return _redis + + +def cache_ttl_seconds() -> int: + raw = os.environ.get("USAGE_READ_CACHE_TTL_SECONDS", "90") + try: + return max(0, int(raw)) + except (TypeError, ValueError): + return 90 + + +def cache_key_for(access: UsageAccessResult, **params: Any) -> str: + payload: dict[str, Any] = { + "display_start": access.display_start.isoformat(), + "display_end": access.display_end.isoformat(), + "filter_start": access.filter_start.isoformat(), + "filter_end": access.filter_end.isoformat(), + "enforced_floor": ( + access.enforced_filter_floor.isoformat() + if access.enforced_filter_floor + else None + ), + "range_clamped": access.range_clamped, + "policy": access.policy.as_dict(), + } + for key, value in sorted(params.items()): + if value is None: + payload[key] = None + elif isinstance(value, UUID): + payload[key] = str(value) + else: + payload[key] = value + digest = hashlib.sha256( + json.dumps(payload, sort_keys=True, default=str).encode() + ).hexdigest() + return digest[:32] + + +def _redis_key(organization_id: UUID, endpoint: str, cache_key: str) -> str: + return f"usage:read:{organization_id}:{endpoint}:{cache_key}" + + +def _is_empty_usage_payload(endpoint: str, payload: dict[str, Any]) -> bool: + if endpoint == "summary": + totals = payload.get("totals") or {} + if int(totals.get("call_count") or 0) > 0: + return False + return not any( + int(totals.get(field) or 0) > 0 + for field in ( + "prompt_tokens", + "completion_tokens", + "audio_seconds", + "tts_characters", + ) + ) + if endpoint == "breakdown": + rows = payload.get("rows") or [] + if not rows: + return True + return not any( + int(row.get("call_count") or 0) > 0 + or int(row.get("total_tokens") or 0) > 0 + or int(row.get("audio_seconds") or 0) > 0 + or int(row.get("tts_characters") or 0) > 0 + for row in rows + ) + return False + + +def get_cached_response( + organization_id: UUID, + endpoint: str, + cache_key: str, +) -> Optional[dict[str, Any]]: + ttl = cache_ttl_seconds() + if ttl <= 0: + return None + try: + raw = _client().get(_redis_key(organization_id, endpoint, cache_key)) + if not raw: + return None + payload = json.loads(raw) + if _is_empty_usage_payload(endpoint, payload): + return None + return payload + except (redis.RedisError, json.JSONDecodeError) as exc: + logger.debug("usage read cache miss/error org={} endpoint={}: {}", organization_id, endpoint, exc) + return None + + +def set_cached_response( + organization_id: UUID, + endpoint: str, + cache_key: str, + payload: dict[str, Any], +) -> None: + ttl = cache_ttl_seconds() + if ttl <= 0: + return + if _is_empty_usage_payload(endpoint, payload): + return + try: + _client().set( + _redis_key(organization_id, endpoint, cache_key), + json.dumps(payload, default=str), + ex=ttl, + ) + except redis.RedisError as exc: + logger.debug("usage read cache set failed org={}: {}", organization_id, exc) + + +def invalidate_org_usage_read_cache(organization_id: UUID) -> int: + pattern = f"usage:read:{organization_id}:*" + deleted = 0 + try: + client = _client() + batch: list[str] = [] + for key in client.scan_iter(match=pattern, count=200): + batch.append(key) + if len(batch) >= 200: + deleted += client.delete(*batch) + batch.clear() + if batch: + deleted += client.delete(*batch) + except redis.RedisError as exc: + logger.warning("usage read cache invalidate failed org={}: {}", organization_id, exc) + return deleted diff --git a/app/services/usage/retention.py b/app/services/usage/retention.py new file mode 100644 index 00000000..763cbbcc --- /dev/null +++ b/app/services/usage/retention.py @@ -0,0 +1,41 @@ +"""OSS usage rollup retention — per-org flush beyond history window.""" + +from __future__ import annotations + +from datetime import date, timedelta + +from sqlalchemy.orm import Session + +from loguru import logger + +from app.core.license import get_enabled_features, get_license_info +from app.core.usage_entitlement import OSS_USAGE_HISTORY_DAYS +from app.models.database import LLMUsageDaily + + +def oss_usage_cutoff_date() -> date: + return date.today() - timedelta(days=OSS_USAGE_HISTORY_DAYS - 1) + + +def prune_oss_usage_history(db: Session) -> dict: + """ + Delete rollup rows older than OSS window for non-entitled orgs. + Deployment-wide license → no deletes. Org-scoped license → skip licensed org. + """ + if get_enabled_features() and get_license_info().get("org_id") is None: + return {"deleted": 0} + + cutoff = oss_usage_cutoff_date() + licensed_org = get_license_info().get("org_id") + + query = db.query(LLMUsageDaily).filter(LLMUsageDaily.usage_date < cutoff) + if get_enabled_features() and licensed_org is not None: + query = query.filter(LLMUsageDaily.organization_id != licensed_org) + + deleted = query.delete(synchronize_session=False) + db.commit() + + if deleted: + logger.info("pruned_oss_usage_history deleted_rows={} cutoff={}", deleted, cutoff) + + return {"deleted": deleted, "cutoff": cutoff.isoformat()} diff --git a/app/services/usage/usage_costs.py b/app/services/usage/usage_costs.py new file mode 100644 index 00000000..914de0b1 --- /dev/null +++ b/app/services/usage/usage_costs.py @@ -0,0 +1,51 @@ +"""Usage cost presentation helpers for API responses.""" + +from __future__ import annotations + +from typing import Any, Dict, Optional + +MICRO_USD_PER_DOLLAR = 1_000_000 + + +def micro_to_usd(micro: int) -> float: + return micro / MICRO_USD_PER_DOLLAR + + +def costs_from_micro( + *, + input_cost_micro_usd: int = 0, + output_cost_micro_usd: int = 0, + cache_read_cost_micro_usd: int = 0, + cache_creation_cost_micro_usd: int = 0, + reasoning_cost_micro_usd: int = 0, + audio_cost_micro_usd: int = 0, + tts_cost_micro_usd: int = 0, + total_cost_micro_usd: Optional[int] = None, + has_unpriced_usage: bool = False, + currency: str = "USD", +) -> Dict[str, Any]: + total_micro = ( + total_cost_micro_usd + if total_cost_micro_usd is not None + else ( + input_cost_micro_usd + + output_cost_micro_usd + + cache_read_cost_micro_usd + + cache_creation_cost_micro_usd + + reasoning_cost_micro_usd + + audio_cost_micro_usd + + tts_cost_micro_usd + ) + ) + return { + "input_cost_usd": micro_to_usd(input_cost_micro_usd), + "output_cost_usd": micro_to_usd(output_cost_micro_usd), + "cache_read_cost_usd": micro_to_usd(cache_read_cost_micro_usd), + "cache_write_cost_usd": micro_to_usd(cache_creation_cost_micro_usd), + "reasoning_cost_usd": micro_to_usd(reasoning_cost_micro_usd), + "audio_cost_usd": micro_to_usd(audio_cost_micro_usd), + "tts_cost_usd": micro_to_usd(tts_cost_micro_usd), + "total_cost_usd": micro_to_usd(total_micro), + "currency": currency, + "has_unpriced_usage": has_unpriced_usage, + } diff --git a/app/services/usage/usage_labels.py b/app/services/usage/usage_labels.py new file mode 100644 index 00000000..4cee5167 --- /dev/null +++ b/app/services/usage/usage_labels.py @@ -0,0 +1,424 @@ +"""Human-readable labels for usage attribution context (no raw UUIDs in UI).""" + +from __future__ import annotations + +from typing import Any, Dict, Optional, Set +from uuid import UUID + +from sqlalchemy.orm import Session + +from app.models.database import ( + Agent, + CallImport, + CallImportEvaluation, + CallImportRow, + CallImportTag, + CallImportTagAssignment, + EvaluatorResult, + TTSComparison, +) + +RESOURCE_TYPE_LABELS = { + "call_import_evaluation": "Evaluation", + "call_import": "Import", + "tts_comparison": "Simulation", + "evaluator_result": "Evaluator result", + "agent": "Agent", + "metric": "Metric", +} + +_USAGE_KIND_LABELS = { + "llm": "LLM", + "stt": "STT", + "tts": "TTS", +} + + +def usage_kind_label(kind: Optional[str]) -> str: + if not kind: + return "—" + return _USAGE_KIND_LABELS.get(kind, kind) + + +def short_entity_id(uid: UUID) -> str: + return str(uid)[:8] + + +def format_entity_label( + custom_name: Optional[str], + uid: UUID, + default_prefix: str, +) -> str: + """Human label: name-shortId (e.g. unauthenticated sheet.xlsx-3111d376).""" + short = short_entity_id(uid) + text = (custom_name or "").strip() + if text: + return f"{text}-{short}" + return f"{default_prefix}-{short}" + + +class UsageNameResolver: + """Batch-resolve entity names for usage context JSONB keys.""" + + def __init__(self, db: Session, organization_id: UUID) -> None: + self._db = db + self._organization_id = organization_id + self._evaluations: Dict[UUID, str] = {} + self._call_imports: Dict[UUID, str] = {} + self._call_import_rows: Dict[UUID, str] = {} + self._tts_comparisons: Dict[UUID, str] = {} + self._agents: Dict[UUID, str] = {} + + def preload(self, contexts: list[Dict[str, Any]]) -> None: + eval_ids: Set[UUID] = set() + import_ids: Set[UUID] = set() + row_ids: Set[UUID] = set() + comparison_ids: Set[UUID] = set() + agent_ids: Set[UUID] = set() + evaluator_result_ids: Set[UUID] = set() + + for ctx in contexts: + if not ctx: + continue + norm = _normalize_context(ctx) + rtype = norm.get("resource_type") + resource_id = norm.get("resource_id") + if resource_id: + uid = parse_uuid(resource_id) + if uid: + if rtype == "call_import": + import_ids.add(uid) + elif rtype == "call_import_evaluation": + eval_ids.add(uid) + elif rtype == "tts_comparison": + comparison_ids.add(uid) + elif rtype == "agent": + agent_ids.add(uid) + elif rtype == "evaluator_result": + eval_ids.add(uid) + else: + eval_ids.add(uid) + import_ids.add(uid) + for key, bucket in ( + ("evaluation_id", eval_ids), + ("call_import_id", import_ids), + ("call_import_row_id", row_ids), + ("agent_id", agent_ids), + ("evaluator_result_id", evaluator_result_ids), + ): + raw = norm.get(key) + if not raw: + continue + uid = parse_uuid(raw) + if uid: + bucket.add(uid) + + if eval_ids: + for row in ( + self._db.query(CallImportEvaluation) + .filter( + CallImportEvaluation.organization_id == self._organization_id, + CallImportEvaluation.id.in_(eval_ids), + ) + .all() + ): + self._evaluations[row.id] = format_entity_label( + row.name, row.id, "Evaluation" + ) + + tag_names_by_import: Dict[UUID, list[str]] = {} + if import_ids: + for cid, tag_name in ( + self._db.query( + CallImportTagAssignment.call_import_id, + CallImportTag.name, + ) + .join( + CallImportTag, + CallImportTag.id == CallImportTagAssignment.tag_id, + ) + .filter(CallImportTagAssignment.call_import_id.in_(import_ids)) + .all() + ): + tag_names_by_import.setdefault(cid, []).append(tag_name) + + if import_ids: + for row in ( + self._db.query(CallImport) + .filter( + CallImport.organization_id == self._organization_id, + CallImport.id.in_(import_ids), + ) + .all() + ): + display_name, prefix = _call_import_label_parts( + row, + sorted(tag_names_by_import.get(row.id, [])), + ) + self._call_imports[row.id] = format_entity_label( + display_name, row.id, prefix + ) + + if comparison_ids: + for row in ( + self._db.query(TTSComparison) + .filter( + TTSComparison.organization_id == self._organization_id, + TTSComparison.id.in_(comparison_ids), + ) + .all() + ): + self._tts_comparisons[row.id] = _tts_comparison_display_name(row) + + if evaluator_result_ids: + for row in ( + self._db.query(EvaluatorResult) + .filter( + EvaluatorResult.organization_id == self._organization_id, + EvaluatorResult.id.in_(evaluator_result_ids), + ) + .all() + ): + if row.agent_id: + agent_ids.add(row.agent_id) + + missing_agent_ids = [uid for uid in agent_ids if uid not in self._agents] + if missing_agent_ids: + for row in ( + self._db.query(Agent) + .filter( + Agent.organization_id == self._organization_id, + Agent.id.in_(missing_agent_ids), + ) + .all() + ): + self._agents[row.id] = _agent_display_name(row) + + if row_ids: + for row in ( + self._db.query(CallImportRow) + .filter( + CallImportRow.organization_id == self._organization_id, + CallImportRow.id.in_(row_ids), + ) + .all() + ): + self._call_import_rows[row.id] = _clean_name( + row.conversation_id, "Conversation" + ) + + def evaluation_name(self, raw_id: str) -> str: + uid = parse_uuid(raw_id) + if uid and uid in self._evaluations: + return self._evaluations[uid] + if uid: + return format_entity_label(None, uid, "Evaluation") + return "Evaluation" + + def call_import_name(self, raw_id: str) -> str: + uid = parse_uuid(raw_id) + if uid and uid in self._call_imports: + return self._call_imports[uid] + if uid: + return format_entity_label(None, uid, "Import") + return "Import" + + def tts_comparison_name(self, raw_id: str) -> str: + uid = parse_uuid(raw_id) + if uid and uid in self._tts_comparisons: + return self._tts_comparisons[uid] + if uid: + return format_entity_label(None, uid, "Simulation") + return "Simulation" + + def agent_name(self, raw_id: str) -> str: + uid = parse_uuid(raw_id) + if uid and uid in self._agents: + return self._agents[uid] + if uid: + return format_entity_label(None, uid, "Agent") + return "Agent" + + def resource_name(self, raw_id: str, resource_type: Optional[str]) -> str: + if resource_type == "call_import_evaluation": + return self.evaluation_name(raw_id) + if resource_type == "call_import": + return self.call_import_name(raw_id) + if resource_type == "tts_comparison": + return self.tts_comparison_name(raw_id) + if resource_type == "agent": + return self.agent_name(raw_id) + uid = parse_uuid(raw_id) + if uid: + prefix = RESOURCE_TYPE_LABELS.get(resource_type or "", "Resource") + return format_entity_label(None, uid, prefix) + return RESOURCE_TYPE_LABELS.get(resource_type or "", "Unscoped") + + def call_import_row_name(self, raw_id: str) -> str: + uid = parse_uuid(raw_id) + if uid and uid in self._call_import_rows: + return self._call_import_rows[uid] + return "Conversation" + + +def parse_uuid(raw: Any) -> Optional[UUID]: + try: + return UUID(str(raw)) + except (ValueError, TypeError): + return None + + +def _clean_name(value: Optional[str], fallback: str) -> str: + text = (value or "").strip() + return text or fallback + + +def _call_import_title(row: CallImport) -> Optional[str]: + filename = (row.original_filename or "").strip() + if filename: + return filename + dataset = (row.dataset or "").strip() + if dataset: + return dataset + return None + + +def _call_import_meta_suffix(row: CallImport, tag_names: list[str]) -> str: + parts: list[str] = [] + dataset = (row.dataset or "").strip() + filename = (row.original_filename or "").strip() + if dataset and dataset != filename: + parts.append(dataset) + for name in tag_names: + if name and name not in parts: + parts.append(name) + if not parts: + return "" + return f" ({' · '.join(parts)})" + + +def _call_import_label_parts( + row: CallImport, + tag_names: list[str], +) -> tuple[Optional[str], str]: + base = _call_import_title(row) + suffix = _call_import_meta_suffix(row, tag_names) + if base: + return f"{base}{suffix}", "Import" + if suffix: + return suffix.strip(" ()"), "Import" + return None, "Import" + + +def _tts_comparison_display_name(row: TTSComparison) -> str: + name = (row.name or "").strip() or "Simulation" + sim = (row.simulation_id or "").strip() + if sim: + return f"{name} #{sim}" + return format_entity_label(name, row.id, "Simulation") + + +def _agent_display_name(row: Agent) -> str: + name = (row.name or "").strip() or "Agent" + short = (row.agent_id or "").strip() + if short: + return f"{name} #{short}" + return format_entity_label(name, row.id, "Agent") + + +def _normalize_context(raw: Any) -> Dict[str, str]: + if not raw or not isinstance(raw, dict): + return {} + return {str(k): str(v) for k, v in raw.items() if v is not None} + + +def _richest_context(contexts: list[Dict[str, str]]) -> Dict[str, str]: + if not contexts: + return {} + return max(contexts, key=lambda c: len(c)) + + +def build_usage_resource_label( + context: Optional[Dict[str, Any]], + resource_type: Optional[str], + resolver: UsageNameResolver, +) -> str: + """Hierarchical label: call import · evaluation · conversation.""" + ctx = _normalize_context(context) + parts: list[str] = [] + + call_import_id = ctx.get("call_import_id") + if call_import_id: + parts.append(resolver.call_import_name(call_import_id)) + elif resource_type == "call_import" and ctx.get("resource_id"): + parts.append(resolver.call_import_name(ctx["resource_id"])) + + evaluation_id = ctx.get("evaluation_id") + agent_id = ctx.get("agent_id") + resource_id = ctx.get("resource_id") + if evaluation_id: + parts.append(resolver.evaluation_name(evaluation_id)) + elif resource_type == "call_import_evaluation" and resource_id: + parts.append(resolver.evaluation_name(resource_id)) + elif resource_type == "tts_comparison" and resource_id: + parts.append(resolver.tts_comparison_name(resource_id)) + elif resource_type == "agent" and resource_id: + parts.append(resolver.agent_name(resource_id)) + elif agent_id: + parts.append(resolver.agent_name(agent_id)) + elif resource_type == "evaluator_result" and resource_id: + parts.append(resolver.resource_name(resource_id, resource_type)) + elif resource_type and resource_id: + parts.append(resolver.resource_name(resource_id, resource_type)) + + row_id = ctx.get("call_import_row_id") + if row_id: + parts.append(resolver.call_import_row_name(row_id)) + + if not parts: + rid = resource_id or agent_id + if rid and (resource_type == "agent" or agent_id): + return resolver.agent_name(rid) + if resource_type: + prefix = RESOURCE_TYPE_LABELS.get(resource_type, resource_type) + uid = parse_uuid(rid) + if uid: + return format_entity_label(None, uid, prefix) + return prefix + return "Unscoped" + + return " / ".join(parts) + + +def labels_for_resource_buckets( + buckets: list[tuple[Optional[str], Optional[str], list[Dict[str, Any]]]], + resolver: UsageNameResolver, +) -> Dict[str, str]: + """Map resource_id string -> label; buckets are (resource_id, resource_type, contexts).""" + labels: Dict[str, str] = {} + for raw_id, resource_type, contexts in buckets: + if not raw_id: + continue + ctx = _richest_context([_normalize_context(c) for c in contexts]) + merged = dict(ctx) + merged.setdefault("resource_id", raw_id) + if resource_type: + merged.setdefault("resource_type", resource_type) + labels[str(raw_id)] = build_usage_resource_label( + merged, resource_type, resolver + ) + return labels + + +def labels_for_call_import_ids( + import_ids: list[UUID], + resolver: UsageNameResolver, +) -> Dict[str, str]: + labels: Dict[str, str] = {} + for uid in import_ids: + labels[str(uid)] = resolver.call_import_name(str(uid)) + return labels + + +def collect_contexts_from_rows(rows: list[Any]) -> list[Dict[str, Any]]: + return [_normalize_context(r[0]) for r in rows if r and r[0]] diff --git a/app/services/usage/voice_usage_processor.py b/app/services/usage/voice_usage_processor.py new file mode 100644 index 00000000..21d82f14 --- /dev/null +++ b/app/services/usage/voice_usage_processor.py @@ -0,0 +1,110 @@ +"""Voice pipeline processor that records LLM/TTS usage from MetricsFrames.""" + +from __future__ import annotations + +import math +from typing import Optional +from uuid import UUID + +from loguru import logger + + +def create_llm_usage_recorder( + *, + organization_id: UUID | str | None, + workspace_id: UUID | str | None = None, + product_section: str = "playground", + resource_id: UUID | str | None = None, + resource_type: Optional[str] = None, +): + """Build a FrameProcessor that records LLM/TTS usage from MetricsFrames. + + Returns None when organization_id is missing or efficientai is unavailable. + """ + if not organization_id: + return None + + try: + from efficientai.frames.frames import Frame, MetricsFrame + from efficientai.metrics.metrics import ( + LLMUsageMetricsData, + ProcessingMetricsData, + TTSUsageMetricsData, + ) + from efficientai.processors.frame_processor import FrameDirection, FrameProcessor + + from app.services.usage.context import ( + LLMUsageContext, + LLMUsageProductSection, + ) + from app.services.usage.llm_usage import ( + record_llm_usage, + record_stt_usage, + record_tts_usage, + ) + from app.services.usage.normalize import UsageSnapshot, usage_snapshot_is_billable + except Exception as exc: + logger.debug("voice usage recorder unavailable: {}", exc) + return None + + try: + section = LLMUsageProductSection(product_section) + except ValueError: + section = LLMUsageProductSection.OTHER + + org_uuid = UUID(str(organization_id)) + ws_uuid = UUID(str(workspace_id)) if workspace_id else None + res_uuid = UUID(str(resource_id)) if resource_id else None + + usage_ctx = LLMUsageContext( + organization_id=org_uuid, + workspace_id=ws_uuid, + product_section=section, + resource_id=res_uuid, + resource_type=resource_type, + ) + + class LLMUsageRecorderProcessor(FrameProcessor): + async def process_frame(self, frame: Frame, direction: FrameDirection): + await super().process_frame(frame, direction) + if isinstance(frame, MetricsFrame): + for item in frame.data or []: + if isinstance(item, LLMUsageMetricsData) and item.value is not None: + tokens = item.value + snapshot = UsageSnapshot( + prompt_tokens=int(tokens.prompt_tokens or 0), + completion_tokens=int(tokens.completion_tokens or 0), + cache_read_tokens=int(tokens.cache_read_input_tokens or 0), + cache_creation_tokens=int( + tokens.cache_creation_input_tokens or 0 + ), + reasoning_tokens=int(tokens.reasoning_tokens or 0), + ) + if not usage_snapshot_is_billable(snapshot): + continue + record_llm_usage( + item.model or "unknown", + snapshot, + organization_id=org_uuid, + ctx=usage_ctx, + ) + elif isinstance(item, TTSUsageMetricsData): + record_tts_usage( + item.model or "unknown", + characters=int(item.value or 0), + organization_id=org_uuid, + ctx=usage_ctx, + ) + elif isinstance(item, ProcessingMetricsData): + processor_name = (item.processor or "").lower() + if "stt" in processor_name and item.value is not None: + seconds = max(1, int(math.ceil(float(item.value)))) + record_stt_usage( + item.model or "unknown", + audio_seconds=seconds, + organization_id=org_uuid, + ctx=usage_ctx, + ) + await self.push_frame(frame, direction) + + return LLMUsageRecorderProcessor() diff --git a/app/services/voice_agent/bot_fast_api.py b/app/services/voice_agent/bot_fast_api.py index 38fa2733..3bcd92a7 100644 --- a/app/services/voice_agent/bot_fast_api.py +++ b/app/services/voice_agent/bot_fast_api.py @@ -98,7 +98,7 @@ def _get_imports(): """ -async def run_bot(websocket_client, google_api_key: str, system_instruction: str = None, organization_id: str = None, agent_id: str = None, persona_id: str = None, scenario_id: str = None, evaluator_id: str = None, result_id: str = None, model_name: str = None, serializer=None, telephony_mode: bool = False, call_short_id: str = None, silence_hangup_secs: float | None = None): +async def run_bot(websocket_client, google_api_key: str, system_instruction: str = None, organization_id: str = None, agent_id: str = None, persona_id: str = None, scenario_id: str = None, evaluator_id: str = None, result_id: str = None, model_name: str = None, serializer=None, telephony_mode: bool = False, call_short_id: str = None, silence_hangup_secs: float | None = None, workspace_id: str = None): """ Run the voice agent bot with the provided Google API key. @@ -246,6 +246,17 @@ async def on_silence_hangup(): if user_transcript_processor: pipeline_processors.append(user_transcript_processor) pipeline_processors.append(llm) + from app.services.usage.voice_usage_processor import create_llm_usage_recorder + + usage_recorder = create_llm_usage_recorder( + organization_id=organization_id, + workspace_id=workspace_id, + product_section="agents" if agent_id else ("telephony" if telephony_mode else "playground"), + resource_id=agent_id, + resource_type="agent" if agent_id else None, + ) + if usage_recorder: + pipeline_processors.append(usage_recorder) if agent_transcript_processor: pipeline_processors.append(agent_transcript_processor) pipeline_processors.extend([ @@ -278,18 +289,32 @@ async def on_client_disconnected(transport, client): # RTVI events for efficientai client UI rtvi = imports["RTVIProcessor"](config=imports["RTVIConfig"](config=[])) - pipeline = imports["Pipeline"]( + from app.services.usage.voice_usage_processor import create_llm_usage_recorder + + usage_recorder = create_llm_usage_recorder( + organization_id=organization_id, + workspace_id=workspace_id, + product_section="agents" if agent_id else "playground", + resource_id=agent_id, + resource_type="agent" if agent_id else None, + ) + pipeline_steps = [ + ws_transport.input(), + user_recorder, + context_aggregator.user(), + rtvi, + llm, + ] + if usage_recorder: + pipeline_steps.append(usage_recorder) + pipeline_steps.extend( [ - ws_transport.input(), - user_recorder, - context_aggregator.user(), - rtvi, - llm, bot_recorder, ws_transport.output(), context_aggregator.assistant(), ] ) + pipeline = imports["Pipeline"](pipeline_steps) task = imports["PipelineTask"]( pipeline, diff --git a/app/services/voice_agent/voice_bundle.py b/app/services/voice_agent/voice_bundle.py index 296bf032..edebc4d0 100644 --- a/app/services/voice_agent/voice_bundle.py +++ b/app/services/voice_agent/voice_bundle.py @@ -510,6 +510,7 @@ async def run_voice_bundle_fastapi( websocket_client, system_instruction: str | None = None, organization_id: str | None = None, + workspace_id: str | None = None, agent_id: str | None = None, persona_id: str | None = None, scenario_id: str | None = None, @@ -835,8 +836,27 @@ async def on_silence_hangup(): if telephony_mode and call_short_id and agent_transcript_processor: pipeline_processors.append(agent_transcript_processor) + from app.services.usage.voice_usage_processor import create_llm_usage_recorder + + agent_usage_section = ( + "agents" + if agent_id + else ("telephony" if telephony_mode else "voice_playground") + ) + usage_recorder = create_llm_usage_recorder( + organization_id=organization_id, + workspace_id=workspace_id, + product_section=agent_usage_section, + resource_id=agent_id, + resource_type="agent" if agent_id else None, + ) + pipeline_processors.extend([ tts, + ]) + if usage_recorder: + pipeline_processors.append(usage_recorder) + pipeline_processors.extend([ bot_recorder if use_aligned_recorders else audio_buffer_output, ws_transport.output(), context_aggregator.assistant(), @@ -865,6 +885,7 @@ async def on_client_disconnected(transport, client): await task.cancel() else: rtvi = imports["RTVIProcessor"](config=imports["RTVIConfig"](config=[])) + rtvi_processors = [ws_transport.input()] if silence_hangup_processor: rtvi_processors.append(silence_hangup_processor) @@ -875,6 +896,17 @@ async def on_client_disconnected(transport, client): rtvi, llm, tts, + ]) + usage_recorder = create_llm_usage_recorder( + organization_id=organization_id, + workspace_id=workspace_id, + product_section="agents" if agent_id else "voice_playground", + resource_id=agent_id, + resource_type="agent" if agent_id else None, + ) + if usage_recorder: + rtvi_processors.append(usage_recorder) + rtvi_processors.extend([ audio_buffer_output, ws_transport.output(), context_aggregator.assistant(), diff --git a/app/services/webrtc_bridge/test_agent_processor.py b/app/services/webrtc_bridge/test_agent_processor.py index 6d5ea658..990fa6bd 100644 --- a/app/services/webrtc_bridge/test_agent_processor.py +++ b/app/services/webrtc_bridge/test_agent_processor.py @@ -9,8 +9,9 @@ import asyncio import io import os -from typing import Optional, Callable, Awaitable, List, Dict, Any +from typing import Optional, Callable, Awaitable, List, Dict, Any, Union from dataclasses import dataclass, field +from uuid import UUID from loguru import logger # TTS service imports @@ -90,6 +91,9 @@ class TestAgentConfig: response_delay_ms: int = 500 # Delay before responding (more natural) allow_interruptions: bool = False + organization_id: Optional[Union[UUID, str]] = None + workspace_id: Optional[Union[UUID, str]] = None + class TestAgentProcessor: """ @@ -366,7 +370,9 @@ async def _generate_llm_response(self) -> Optional[str]: max_tokens=self.config.llm_max_tokens if self.config.llm_max_tokens is not None else 150, temperature=self.config.llm_temperature if self.config.llm_temperature is not None else 0.7, ) - + + self._record_llm_usage(response=response) + return response.choices[0].message.content.strip() except Exception as e: @@ -392,6 +398,8 @@ async def _text_to_speech(self, text: str) -> Optional[bytes]: else: audio = await self._tts_cartesia(text) + if audio: + self._record_tts_usage(text=text) return audio except Exception as e: logger.error(f"[TestAgent] TTS ({provider}) error: {e}") @@ -400,6 +408,82 @@ async def _text_to_speech(self, text: str) -> Optional[bytes]: def _tts_settings(self) -> Dict[str, Any]: return dict(self.config.tts_config or {}) + def _record_tts_usage(self, *, text: str) -> None: + if not self.config.organization_id: + return + try: + from app.services.usage.context import ( + LLMUsageContext, + LLMUsageProductSection, + llm_usage_context, + ) + from app.services.usage.llm_usage import record_tts_usage + + model = self.config.tts_model or TTS_DEFAULT_MODELS.get( + self.config.tts_provider.lower(), "unknown" + ) + org_id = UUID(str(self.config.organization_id)) + ws_id = ( + UUID(str(self.config.workspace_id)) + if self.config.workspace_id + else None + ) + with llm_usage_context( + LLMUsageContext( + organization_id=org_id, + workspace_id=ws_id, + product_section=LLMUsageProductSection.TEST_AGENT, + ) + ): + record_tts_usage( + model, + characters=len(text or ""), + organization_id=org_id, + ) + except Exception as exc: + logger.debug("test agent tts usage record skipped: {}", exc) + + def _record_llm_usage(self, *, response: Any) -> None: + if not self.config.organization_id: + return + try: + from app.services.usage.context import ( + LLMUsageContext, + LLMUsageProductSection, + llm_usage_context, + ) + from app.services.usage.llm_usage import record_llm_usage + from app.services.usage.normalize import UsageSnapshot + + usage = getattr(response, "usage", None) + if usage is None: + return + org_id = UUID(str(self.config.organization_id)) + ws_id = ( + UUID(str(self.config.workspace_id)) + if self.config.workspace_id + else None + ) + model = self.config.llm_model or "unknown" + snapshot = UsageSnapshot( + prompt_tokens=int(getattr(usage, "prompt_tokens", 0) or 0), + completion_tokens=int(getattr(usage, "completion_tokens", 0) or 0), + ) + with llm_usage_context( + LLMUsageContext( + organization_id=org_id, + workspace_id=ws_id, + product_section=LLMUsageProductSection.TEST_AGENT, + ) + ): + record_llm_usage( + model, + snapshot, + organization_id=org_id, + ) + except Exception as exc: + logger.debug("test agent llm usage record skipped: {}", exc) + async def _tts_cartesia(self, text: str) -> Optional[bytes]: """Synthesize speech via Cartesia.""" import httpx diff --git a/app/workers/config.py b/app/workers/config.py index f4e44f36..c6862c92 100644 --- a/app/workers/config.py +++ b/app/workers/config.py @@ -53,6 +53,7 @@ os.environ.setdefault("TOKENIZERS_PARALLELISM", "false") from celery import Celery # noqa: E402 (env vars must be set first) +from celery.schedules import crontab # noqa: E402 from loguru import logger # noqa: E402 from app.config import settings, load_config_from_file # noqa: E402 @@ -74,6 +75,39 @@ # Queues consumed by the dedicated call-import / evaluation worker. IMPORTS_WORKER_QUEUES = "imports,diarization,eval-control,evaluations" EVAL_CONTROL_QUEUE = "eval-control" +USAGE_WORKER_QUEUE = "usage" +PLATFORM_WORKER_QUEUE = "platform" + + +def _usage_flush_beat_seconds() -> float: + raw = os.environ.get("USAGE_FLUSH_BEAT_SECONDS", "120") + try: + return max(30.0, float(raw)) + except (TypeError, ValueError): + return 120.0 + + +def _platform_beat_schedule() -> dict: + """Periodic platform tasks — run from dedicated ``celery beat`` (single replica).""" + return { + "flush-usage-counters": { + "task": "flush_usage_counters", + "schedule": _usage_flush_beat_seconds(), + }, + "evaluate-alerts": { + "task": "evaluate_alerts", + "schedule": crontab(minute="*/5"), + }, + "refresh-fx-rates": { + "task": "refresh_fx_rates", + "schedule": crontab(hour=6, minute=0), + }, + "prune-oss-usage-history": { + "task": "prune_oss_usage_history", + "schedule": crontab(hour=3, minute=0), + }, + } + # Create Celery app celery_app = Celery( @@ -106,6 +140,7 @@ # any leaked OMP threads, and HuggingFace tokenizer state, and prevents # slow drift from accumulating over hours of import processing. worker_max_tasks_per_child=20, + beat_schedule=_platform_beat_schedule(), ) # Route call-import recording fetch to the imports queue (preferred by workers @@ -138,4 +173,13 @@ "evaluate_studio_run_item": {"queue": "evaluations"}, "generate_agent_flowchart": {"queue": "celery"}, "map_agent_flowchart_prompt_sections": {"queue": "celery"}, + "flush_usage_counters": {"queue": USAGE_WORKER_QUEUE}, + "recompute_usage_costs": {"queue": USAGE_WORKER_QUEUE}, + "evaluate_alerts": {"queue": PLATFORM_WORKER_QUEUE}, + "refresh_fx_rates": {"queue": PLATFORM_WORKER_QUEUE}, + "prune_oss_usage_history": {"queue": PLATFORM_WORKER_QUEUE}, + "dispatch_cron_jobs": {"queue": "celery"}, + "run_cron_evaluator_job": {"queue": "celery"}, } + +from app.workers import cron_bootstrap # noqa: F401,E402 diff --git a/app/workers/cron_bootstrap.py b/app/workers/cron_bootstrap.py new file mode 100644 index 00000000..b0c1d62a --- /dev/null +++ b/app/workers/cron_bootstrap.py @@ -0,0 +1,21 @@ +"""Bootstrap cron dispatcher on worker startup.""" + +from __future__ import annotations + +from celery.signals import worker_ready +from loguru import logger + + +@worker_ready.connect +def _bootstrap_cron_dispatcher(sender, **kwargs): + try: + from app.services.cron.dispatcher_lock import try_acquire_dispatcher_leader + + if not try_acquire_dispatcher_leader(): + return + from app.workers.tasks.dispatch_cron_jobs import dispatch_cron_jobs_task + + dispatch_cron_jobs_task.apply_async(countdown=5) + logger.info("Cron dispatcher bootstrap enqueued") + except Exception as exc: + logger.warning("Cron dispatcher bootstrap skipped: {}", exc) diff --git a/app/workers/tasks/__init__.py b/app/workers/tasks/__init__.py index e60a9ad2..a8542175 100644 --- a/app/workers/tasks/__init__.py +++ b/app/workers/tasks/__init__.py @@ -23,6 +23,13 @@ from . import finalize_telephony_recording from . import evaluate_studio_run_item from . import call_import_bulk_ops +from . import flush_usage_counters +from . import recompute_usage_costs +from . import prune_oss_usage_history +from . import dispatch_cron_jobs +from . import evaluate_alerts +from . import refresh_fx_rates +from . import run_cron_evaluator_job from app.workers.concurrency import eval_dispatch from app.workers.concurrency import fair_dispatch from app.workers.concurrency import fair_diarization_dispatch @@ -120,3 +127,10 @@ materialize_call_import_evaluation_task = ( call_import_bulk_ops.materialize_call_import_evaluation_task ) +flush_usage_counters_task = flush_usage_counters.flush_usage_counters_task +recompute_usage_costs_task = recompute_usage_costs.recompute_usage_costs_task +prune_oss_usage_history_task = prune_oss_usage_history.prune_oss_usage_history_task +dispatch_cron_jobs_task = dispatch_cron_jobs.dispatch_cron_jobs_task +evaluate_alerts_task = evaluate_alerts.evaluate_alerts_task +refresh_fx_rates_task = refresh_fx_rates.refresh_fx_rates_task +run_cron_evaluator_job_task = run_cron_evaluator_job.run_cron_evaluator_job_task diff --git a/app/workers/tasks/agent_flowchart_jobs.py b/app/workers/tasks/agent_flowchart_jobs.py index 8704fd6c..02431794 100644 --- a/app/workers/tasks/agent_flowchart_jobs.py +++ b/app/workers/tasks/agent_flowchart_jobs.py @@ -52,13 +52,19 @@ def generate_agent_flowchart_task( ) return - graph, provider_enum, model_str = generate_agent_flowchart( - prompt_text=partial.content, - organization_id=partial.organization_id, - db=db, - provider=provider, - model=model, + from app.services.usage.context import ( + llm_usage_context, + usage_context_for_prompt_partial, ) + + with llm_usage_context(usage_context_for_prompt_partial(partial)): + graph, provider_enum, model_str = generate_agent_flowchart( + prompt_text=partial.content, + organization_id=partial.organization_id, + db=db, + provider=provider, + model=model, + ) partial.agent_flowchart = graph.model_dump(mode="json") if isinstance(partial.agent_flowchart, dict): generated_at = graph.generated_at @@ -130,18 +136,24 @@ def map_agent_flowchart_prompt_sections_task( ): raise ValueError("Generate a flowchart before mapping prompt sections") + from app.services.usage.context import ( + llm_usage_context, + usage_context_for_prompt_partial, + ) + graph = apply_prompt_hash_staleness( AgentFlowGraph.model_validate(partial.agent_flowchart), partial.content, ) - mapped_graph = map_all_flow_nodes_to_prompt( - prompt_text=partial.content, - graph=graph, - organization_id=partial.organization_id, - db=db, - provider=provider, - model=model, - ) + with llm_usage_context(usage_context_for_prompt_partial(partial)): + mapped_graph = map_all_flow_nodes_to_prompt( + prompt_text=partial.content, + graph=graph, + organization_id=partial.organization_id, + db=db, + provider=provider, + model=model, + ) partial.agent_flowchart = mapped_graph.model_dump(mode="json") if isinstance(partial.agent_flowchart, dict): if mapped_graph.generated_at is not None: diff --git a/app/workers/tasks/dispatch_cron_jobs.py b/app/workers/tasks/dispatch_cron_jobs.py new file mode 100644 index 00000000..3c658c66 --- /dev/null +++ b/app/workers/tasks/dispatch_cron_jobs.py @@ -0,0 +1,58 @@ +"""Self-scheduling cron job dispatcher (replaces Celery Beat for platform jobs).""" + +from __future__ import annotations + +import os + +from loguru import logger + +from app.database import SessionLocal +from app.workers.config import celery_app + + +def _dispatch_interval_seconds() -> int: + raw = os.environ.get("CRON_DISPATCH_INTERVAL_SECONDS", "30") + try: + return max(10, int(raw)) + except (TypeError, ValueError): + return 30 + + +@celery_app.task(name="dispatch_cron_jobs") +def dispatch_cron_jobs_task() -> dict: + from app.services.cron.dispatcher_lock import ( + acquire_dispatcher_run_lock, + refresh_dispatcher_leader, + release_dispatcher_run_lock, + ) + from app.services.cron.job_dispatch import ( + advance_cron_job, + enqueue_cron_job, + list_due_cron_jobs, + ) + + interval = _dispatch_interval_seconds() + refresh_dispatcher_leader() + + if not acquire_dispatcher_run_lock(): + dispatch_cron_jobs_task.apply_async(countdown=interval) + return {"skipped": "locked"} + + dispatched: list[dict] = [] + db = SessionLocal() + try: + for job in list_due_cron_jobs(db): + try: + meta = enqueue_cron_job(job) + advance_cron_job(db, job) + db.commit() + dispatched.append({"job_id": str(job.id), "job_type": job.job_type, **meta}) + except Exception as exc: + db.rollback() + logger.warning("cron dispatch failed for job {}: {}", job.id, exc) + finally: + db.close() + release_dispatcher_run_lock() + + dispatch_cron_jobs_task.apply_async(countdown=interval) + return {"dispatched": len(dispatched), "jobs": dispatched} diff --git a/app/workers/tasks/evaluate_alerts.py b/app/workers/tasks/evaluate_alerts.py new file mode 100644 index 00000000..5d541eda --- /dev/null +++ b/app/workers/tasks/evaluate_alerts.py @@ -0,0 +1,17 @@ +"""Celery task: evaluate all active alerts.""" + +from __future__ import annotations + +from app.database import SessionLocal +from app.workers.config import celery_app + + +@celery_app.task(name="evaluate_alerts") +def evaluate_alerts_task() -> dict: + from app.services.alerts.alert_evaluation_service import alert_evaluation_service + + db = SessionLocal() + try: + return alert_evaluation_service.evaluate_all_alerts(db) + finally: + db.close() diff --git a/app/workers/tasks/evaluate_call_import_row.py b/app/workers/tasks/evaluate_call_import_row.py index 29247d24..c4090f6c 100644 --- a/app/workers/tasks/evaluate_call_import_row.py +++ b/app/workers/tasks/evaluate_call_import_row.py @@ -332,6 +332,7 @@ def evaluate_call_import_row_task( slot_task_id = _eval_slot_task_id or self.request.id scoring_inputs: dict[str, Any] | None = None restricted_metric_uuids: list[UUID] | None = None + usage_ctx_token = None try: from app.db_sharding.row_ops import ( close_row_sessions, @@ -363,6 +364,20 @@ def evaluate_call_import_row_task( row_db.commit() return {"status": "failed", "reason": "evaluation_missing"} + from app.services.usage.call_import_context import ( + call_import_evaluation_usage_context, + ) + from app.services.usage.context import set_usage_context + + usage_ctx_token = set_usage_context( + call_import_evaluation_usage_context( + organization_id=evaluation.organization_id, + workspace_id=evaluation.workspace_id, + evaluation_id=evaluation.id, + call_import_id=evaluation.call_import_id, + ) + ) + previous_row_status = eval_row.status eval_row.status = "running" @@ -761,6 +776,10 @@ def evaluate_call_import_row_task( if row_db is not None: close_row_sessions(row_db, catalog_db) finally: + if usage_ctx_token is not None: + from app.services.usage.context import reset_usage_context + + reset_usage_context(usage_ctx_token) from app.workers.concurrency.fair_dispatch import ( finish_eval_work_and_redispatch, ) diff --git a/app/workers/tasks/flush_usage_counters.py b/app/workers/tasks/flush_usage_counters.py new file mode 100644 index 00000000..3dc158d0 --- /dev/null +++ b/app/workers/tasks/flush_usage_counters.py @@ -0,0 +1,18 @@ +"""Celery task: flush Redis LLM usage counters into catalog rollups.""" + +from __future__ import annotations + +from loguru import logger + +from app.database import SessionLocal +from app.workers.config import celery_app + + +@celery_app.task(name="flush_usage_counters") +def flush_usage_counters_task() -> dict: + from app.services.usage.llm_usage import flush_all_usage_to_catalog + + flushed = flush_all_usage_to_catalog(SessionLocal) + if flushed: + logger.info("Flushed {} LLM usage rollup buckets", flushed) + return {"flushed_buckets": flushed} diff --git a/app/workers/tasks/generate_evaluation_metric_clusters.py b/app/workers/tasks/generate_evaluation_metric_clusters.py index e91dc108..8cb59775 100644 --- a/app/workers/tasks/generate_evaluation_metric_clusters.py +++ b/app/workers/tasks/generate_evaluation_metric_clusters.py @@ -120,19 +120,23 @@ def on_progress(completed: int, total: int) -> None: flag_modified(evaluation, "metric_clusters") db.commit() - state = generate_metric_clusters( - db, - evaluation, - evaluation.organization_id, - provider_enum, - model_str, - completed_row_pairs=completed_pairs, - metrics=metrics, - policies=policies, - on_progress=on_progress, - max_llm_calls=max_llm_calls, - is_cancelled=_reload_cancelled, - ) + from app.services.usage.call_import_context import usage_context_for_evaluation + from app.services.usage.context import llm_usage_context + + with llm_usage_context(usage_context_for_evaluation(evaluation)): + state = generate_metric_clusters( + db, + evaluation, + evaluation.organization_id, + provider_enum, + model_str, + completed_row_pairs=completed_pairs, + metrics=metrics, + policies=policies, + on_progress=on_progress, + max_llm_calls=max_llm_calls, + is_cancelled=_reload_cancelled, + ) if _reload_cancelled(): return diff --git a/app/workers/tasks/generate_evaluation_prompt_improvements.py b/app/workers/tasks/generate_evaluation_prompt_improvements.py index 0ba282a2..fd22b1fb 100644 --- a/app/workers/tasks/generate_evaluation_prompt_improvements.py +++ b/app/workers/tasks/generate_evaluation_prompt_improvements.py @@ -107,25 +107,29 @@ def generate_evaluation_prompt_improvements_task( eval_rows, ) - state = generate_prompt_improvements( - evaluation=evaluation, - imported_agent=imported_agent, - clusters_state=clusters_state, - organization_id=evaluation.organization_id, - db=db, - provider=provider, - model=model, - credential_id=UUID(credential_id) if credential_id else None, - period_deltas=period_deltas, - ) - evaluation.prompt_improvements = prompt_improvements_state_to_db(state) - flag_modified(evaluation, "prompt_improvements") - db.commit() - logger.info( - "Prompt improvements completed for evaluation {} ({} suggestions)", - evaluation_id, - len(state.suggestions), - ) + from app.services.usage.call_import_context import usage_context_for_evaluation + from app.services.usage.context import llm_usage_context + + with llm_usage_context(usage_context_for_evaluation(evaluation)): + state = generate_prompt_improvements( + evaluation=evaluation, + imported_agent=imported_agent, + clusters_state=clusters_state, + organization_id=evaluation.organization_id, + db=db, + provider=provider, + model=model, + credential_id=UUID(credential_id) if credential_id else None, + period_deltas=period_deltas, + ) + evaluation.prompt_improvements = prompt_improvements_state_to_db(state) + flag_modified(evaluation, "prompt_improvements") + db.commit() + logger.info( + "Prompt improvements completed for evaluation {} ({} suggestions)", + evaluation_id, + len(state.suggestions), + ) except Exception as exc: logger.exception( "Prompt improvements failed for evaluation {}: {}", diff --git a/app/workers/tasks/generate_evaluation_tldr_insights.py b/app/workers/tasks/generate_evaluation_tldr_insights.py index c1e5de73..6f9e96a1 100644 --- a/app/workers/tasks/generate_evaluation_tldr_insights.py +++ b/app/workers/tasks/generate_evaluation_tldr_insights.py @@ -40,14 +40,27 @@ def generate_evaluation_tldr_insights_task( if evaluation is None: return {"error": "evaluation_not_found", "status_code": 404} + from app.services.usage.call_import_context import ( + call_import_evaluation_usage_context, + ) + from app.services.usage.context import llm_usage_context + try: - summary = _generate_and_persist_tldr_summary( - db, - evaluation, - organization_id=UUID(organization_id), - provider=provider, - model=model, - ) + with llm_usage_context( + call_import_evaluation_usage_context( + organization_id=evaluation.organization_id, + workspace_id=evaluation.workspace_id, + evaluation_id=evaluation.id, + call_import_id=evaluation.call_import_id, + ) + ): + summary = _generate_and_persist_tldr_summary( + db, + evaluation, + organization_id=UUID(organization_id), + provider=provider, + model=model, + ) except HTTPException as exc: return {"error": exc.detail, "status_code": exc.status_code} diff --git a/app/workers/tasks/generate_evaluation_user_insights.py b/app/workers/tasks/generate_evaluation_user_insights.py index 6029ce4a..71f666df 100644 --- a/app/workers/tasks/generate_evaluation_user_insights.py +++ b/app/workers/tasks/generate_evaluation_user_insights.py @@ -81,21 +81,25 @@ def on_progress(completed: int, total: int) -> None: flag_modified(evaluation, "user_insights") db.commit() - state = generate_user_insights( - db, - evaluation, - evaluation.organization_id, - provider_enum, - model_str, - completed_row_pairs=completed_pairs, - metrics=metrics, - aggregate=aggregate, - on_progress=on_progress, - max_llm_calls=max_llm_calls, - ) - evaluation.user_insights = user_insights_state_to_db(state) - flag_modified(evaluation, "user_insights") - db.commit() + from app.services.usage.call_import_context import usage_context_for_evaluation + from app.services.usage.context import llm_usage_context + + with llm_usage_context(usage_context_for_evaluation(evaluation)): + state = generate_user_insights( + db, + evaluation, + evaluation.organization_id, + provider_enum, + model_str, + completed_row_pairs=completed_pairs, + metrics=metrics, + aggregate=aggregate, + on_progress=on_progress, + max_llm_calls=max_llm_calls, + ) + evaluation.user_insights = user_insights_state_to_db(state) + flag_modified(evaluation, "user_insights") + db.commit() except Exception as exc: # noqa: BLE001 logger.exception( "User insights generation failed for evaluation {}: {}", diff --git a/app/workers/tasks/process_evaluator_result.py b/app/workers/tasks/process_evaluator_result.py index e19c061b..7a5c5c49 100644 --- a/app/workers/tasks/process_evaluator_result.py +++ b/app/workers/tasks/process_evaluator_result.py @@ -1,751 +1,952 @@ -"""Celery task: process evaluator result (transcribe and evaluate metrics).""" - -import time -import uuid as _uuid -from uuid import UUID - -from loguru import logger -from sqlalchemy import or_ - -from app.database import SessionLocal -from app.models.database import ModelProvider - -from app.workers.config import celery_app -from app.workers.tasks.helpers.constants import ( - REMOVED_EVALUATION_METRIC_NAMES, - AUDIO_ONLY_METRIC_NAMES, -) -from app.workers.tasks.helpers.score_utils import provider_matches, get_metric_type_value -from app.workers.tasks.helpers.audio_evaluation import ( - evaluate_audio_metrics, - handle_audio_evaluation_error, -) -from app.workers.tasks.helpers.llm_evaluation import ( - evaluate_with_llm, - handle_llm_evaluation_error, -) -from app.services.evaluators.evaluator_result_call_data import slim_call_data_for_evaluator_result - - -def _commit_evaluator_result(db, result) -> None: - from app.services.live_entity_storage import sync_evaluator_result - - sync_evaluator_result(db, result) - db.commit() - - -def _make_json_serializable(obj): - """Recursively convert non-JSON-native values (e.g., NumPy types).""" - try: - import numpy as np - except Exception: - np = None - - if isinstance(obj, dict): - return {k: _make_json_serializable(v) for k, v in obj.items()} - if isinstance(obj, list): - return [_make_json_serializable(item) for item in obj] - if isinstance(obj, tuple): - return tuple(_make_json_serializable(item) for item in obj) - - if np is not None: - if isinstance(obj, np.integer): - return int(obj) - if isinstance(obj, np.floating): - return float(obj) - if isinstance(obj, np.ndarray): - return obj.tolist() - if isinstance(obj, np.bool_): - return bool(obj) - - return obj - - -def _load_related_entities(db, result): - """Load evaluator, agent, persona, scenario from database.""" - from app.models.database import Evaluator, Agent, Persona, Scenario - - evaluator = None - if result.evaluator_id: - evaluator = db.query(Evaluator).filter(Evaluator.id == result.evaluator_id).first() - if not evaluator: - logger.warning( - f"[EvaluatorResult {result.result_id}] Evaluator {result.evaluator_id} not found, " - "continuing without evaluator" - ) - - agent = None - if result.agent_id: - agent = db.query(Agent).filter(Agent.id == result.agent_id).first() - - persona = None - if result.persona_id: - persona = db.query(Persona).filter(Persona.id == result.persona_id).first() - - scenario = None - if result.scenario_id: - scenario = db.query(Scenario).filter(Scenario.id == result.scenario_id).first() - - return evaluator, agent, persona, scenario - - -def _transcribe_audio(result, ai_providers, db): - """Transcribe audio file and return transcript with timing info.""" - from app.services.ai.transcription_service import transcription_service - - stt_provider = ModelProvider.OPENAI - stt_model = "whisper-1" - - openai_provider = next( - (p for p in ai_providers if provider_matches(p.provider, ModelProvider.OPENAI)), - None, - ) - if not openai_provider: - logger.warning( - f"[EvaluatorResult {result.result_id}] No OpenAI provider found, using default whisper-1" - ) - - transcription_start_time = time.time() - transcription_result = transcription_service.transcribe( - audio_file_key=result.audio_s3_key, - stt_provider=stt_provider, - stt_model=stt_model, - organization_id=result.organization_id, - db=db, - language=None, - enable_speaker_diarization=True, - ) - transcription_time = time.time() - transcription_start_time - - return ( - transcription_result.get("transcript", ""), - transcription_result.get("speaker_segments", []), - transcription_time, - ) - - -def _generate_call_analysis(transcription, ai_providers, organization_id, result_id, db, agent=None, scenario=None): - """Generate call analysis (summary, sentiment, success) using LLM.""" - from app.services.ai.llm_service import llm_service - - llm_provider = ModelProvider.OPENAI - llm_model = "gpt-4o-mini" - - chosen_provider = next( - (p for p in ai_providers if provider_matches(p.provider, llm_provider)), - None, - ) - if not chosen_provider: - llm_provider = ModelProvider.GOOGLE - llm_model = "gemini-2.5-flash" - chosen_provider = next( - (p for p in ai_providers if provider_matches(p.provider, llm_provider)), - None, - ) - if not chosen_provider: - logger.warning(f"[EvaluatorResult {result_id}] No LLM provider available for call analysis") - return None - - agent_context = "" - if agent and agent.description: - agent_context = f"\n\nAgent Description:\n{agent.description}" - if scenario: - scenario_name = getattr(scenario, 'name', '') - scenario_desc = getattr(scenario, 'description', '') - if scenario_name or scenario_desc: - agent_context += f"\n\nScenario: {scenario_name}" - if scenario_desc: - agent_context += f"\nScenario Description: {scenario_desc}" - - messages = [ - { - "role": "system", - "content": ( - "You are a call analysis expert. Analyze the following conversation transcript " - "and provide a structured analysis. Respond ONLY with valid JSON, no markdown." - ), - }, - { - "role": "user", - "content": f"""Analyze this conversation transcript and provide: -1. A concise summary of the call (2-3 sentences) -2. The user/caller's overall sentiment (one of: Positive, Negative, Neutral, Mixed) -3. Whether the call was successful in achieving its objective (true/false) -{agent_context} - -Transcript: -{transcription} - -Respond in this exact JSON format: -{{"call_summary": "...", "user_sentiment": "...", "call_successful": true/false}}""", - }, - ] - - try: - llm_result = llm_service.generate_response( - messages=messages, - llm_provider=llm_provider, - llm_model=llm_model, - organization_id=organization_id, - db=db, - temperature=0.3, - max_tokens=500, - ) - import json - import re - text = llm_result.get("text", "") - json_match = re.search(r'\{[^{}]*\}', text, re.DOTALL) - if json_match: - analysis = json.loads(json_match.group()) - required_keys = {"call_summary", "user_sentiment", "call_successful"} - if required_keys.issubset(analysis.keys()): - logger.info(f"[EvaluatorResult {result_id}] Call analysis generated successfully") - return analysis - logger.warning(f"[EvaluatorResult {result_id}] Could not parse call analysis from LLM response") - return None - except Exception as e: - logger.error(f"[EvaluatorResult {result_id}] Call analysis failed: {e}", exc_info=True) - return None - - -def _categorize_metrics(enabled_metrics, has_audio): - """Split metrics into LLM-evaluable and audio-only categories.""" - llm_metrics = [] - audio_metrics = [] - skipped_scores = {} - - for m in enabled_metrics: - if m.name.lower() in AUDIO_ONLY_METRIC_NAMES: - if has_audio: - audio_metrics.append(m) - else: - skipped_scores[str(m.id)] = { - "value": None, - "type": get_metric_type_value(m), - "metric_name": m.name, - "skipped": "audio_required", - } - else: - llm_metrics.append(m) - - return llm_metrics, audio_metrics, skipped_scores - - -def _evaluate_llm_metrics_grouped( - *, - transcription: str, - llm_metrics: list, - ai_providers, - organization_id, - result_id: str, - db, - evaluator, - agent, - persona, - scenario, -) -> tuple[dict, float | None]: - """Evaluate LLM metrics, grouping categorization children by parent.""" - from app.workers.tasks.evaluate_call_import_row_core import build_parent_groups - - parents_by_id, children_by_parent, standalone_metrics = build_parent_groups( - db, llm_metrics - ) - metric_scores: dict = {} - evaluation_time: float | None = None - - def _run_bucket(bucket, parent_metric=None): - nonlocal evaluation_time - scores, eval_time = evaluate_with_llm( - transcription=transcription, - llm_metrics=bucket, - ai_providers=ai_providers, - organization_id=organization_id, - result_id=result_id, - db=db, - evaluator=evaluator, - agent=agent, - persona=persona, - scenario=scenario, - parent_metric=parent_metric, - ) - metric_scores.update(scores) - if eval_time is not None: - evaluation_time = eval_time - - if standalone_metrics: - _run_bucket(standalone_metrics, parent_metric=None) - - for parent_id, children in children_by_parent.items(): - parent_metric = parents_by_id.get(parent_id) - if not parent_metric: - logger.warning( - f"[EvaluatorResult {result_id}] Parent metric {parent_id} not found; " - "evaluating children as flat metrics" - ) - _run_bucket(children, parent_metric=None) - continue - _run_bucket(children, parent_metric=parent_metric) - - return metric_scores, evaluation_time - - -def _normalize_platform(platform: object) -> str: - """Normalize provider platform enum/string into lowercase string.""" - if not platform: - return "" - if hasattr(platform, "value"): - return str(platform.value).lower() - return str(platform).lower() - - -def _extract_audio_url(call_data: dict, platform: str) -> str | None: - """Extract provider-specific audio URL from call data.""" - recording_urls = call_data.get("recording_urls", {}) if isinstance(call_data, dict) else {} - provider_payload = call_data.get("provider_payload", {}) if isinstance(call_data, dict) else {} - artifact = call_data.get("artifact", {}) if isinstance(call_data, dict) else {} - recording = artifact.get("recording", {}) if isinstance(artifact, dict) else {} - mono_recording = recording.get("mono", {}) if isinstance(recording, dict) else {} - if platform == "elevenlabs": - return recording_urls.get("conversation_audio") - if platform == "retell": - return call_data.get("recording_url") - if platform == "vapi": - return ( - call_data.get("recordingUrl") - or call_data.get("stereoRecordingUrl") - or artifact.get("recordingUrl") - or artifact.get("stereoRecordingUrl") - or mono_recording.get("combinedUrl") - or recording_urls.get("combined_url") - or recording_urls.get("stereo_url") - or call_data.get("recordingUrl") - or provider_payload.get("recordingUrl") - or provider_payload.get("stereoRecordingUrl") - ) - if platform == "smallest": - return ( - call_data.get("recording_url") - or call_data.get("recordingUrl") - or recording_urls.get("combined_url") - or recording_urls.get("conversation_audio") - ) - return None - - -def _recover_missing_audio_for_result(result, db, refresh_call_data: bool = True) -> bool: - """ - Attempt to recover missing audio from provider, upload to S3, and persist key. - - Returns True when a new S3 key is successfully stored. - """ - import requests as _http - - from app.core.encryption import decrypt_api_key - from app.models.database import Agent, Integration - from app.services.storage.s3_service import s3_service - from app.services.voice_providers import get_voice_provider - - platform = _normalize_platform(result.provider_platform) - if platform not in {"retell", "vapi", "elevenlabs", "smallest"}: - return False - if not result.provider_call_id: - return False - - agent = db.query(Agent).filter(Agent.id == result.agent_id).first() if result.agent_id else None - integration = None - decrypted_key = None - if agent and agent.voice_ai_integration_id: - integration = db.query(Integration).filter( - Integration.id == agent.voice_ai_integration_id, - Integration.organization_id == result.organization_id, - ).first() - if integration: - try: - decrypted_key = decrypt_api_key(integration.api_key) - except Exception as decrypt_err: - logger.warning( - f"[EvaluatorResult {result.result_id}] Unable to decrypt integration key " - f"for audio recovery: {decrypt_err}" - ) - - call_data = result.call_data or {} - if refresh_call_data and decrypted_key: - try: - provider_class = get_voice_provider(platform) - provider_kwargs = {"api_key": decrypted_key} - if platform == "vapi" and integration and getattr(integration, "public_key", None): - provider_kwargs["public_key"] = integration.public_key - provider = provider_class(**provider_kwargs) - if hasattr(provider, "retrieve_call_metrics"): - refreshed = provider.retrieve_call_metrics(result.provider_call_id) - if isinstance(refreshed, dict) and refreshed: - call_data = refreshed - result.call_data = refreshed - db.commit() - except Exception as refresh_err: - logger.warning( - f"[EvaluatorResult {result.result_id}] Audio recovery could not refresh provider metrics: " - f"{refresh_err}" - ) - - audio_url = _extract_audio_url(call_data, platform) - if not audio_url: - logger.warning( - f"[EvaluatorResult {result.result_id}] Audio recovery failed: no provider recording URL available" - ) - return False - - headers = {"xi-api-key": decrypted_key} if platform == "elevenlabs" and decrypted_key else None - try: - response = _http.get(audio_url, headers=headers, timeout=120) - except Exception as download_err: - logger.warning( - f"[EvaluatorResult {result.result_id}] Audio recovery download failed: {download_err}" - ) - return False - - if response.status_code != 200 or not response.content: - logger.warning( - f"[EvaluatorResult {result.result_id}] Audio recovery download returned " - f"status={response.status_code}" - ) - return False - - content_type = response.headers.get("content-type", "audio/mpeg") - ext = "wav" if "wav" in content_type else "mp3" - org_id = str(result.organization_id) - s3_key = ( - f"audio/organizations/{org_id}/evaluations/" - f"{result.provider_call_id}/{_uuid.uuid4()}.{ext}" - ) - - try: - s3_service.upload_file_by_key(response.content, s3_key, content_type=content_type) - except Exception as upload_err: - logger.warning( - f"[EvaluatorResult {result.result_id}] Audio recovery upload failed: {upload_err}" - ) - return False - - result.audio_s3_key = s3_key - _commit_evaluator_result(db, result) - db.refresh(result) - logger.info( - f"[EvaluatorResult {result.result_id}] Recovered missing audio and stored at {s3_key}" - ) - return True - - -def _all_audio_scores_download_failed(audio_scores: dict[str, dict[str, object]]) -> bool: - """Check whether every audio metric failed due to missing/unreadable S3 object.""" - if not audio_scores: - return False - return all( - isinstance(score, dict) and score.get("error") == "audio_download_failed" - for score in audio_scores.values() - ) - - -def _playground_call_recording(db, result): - from app.models.database import CallRecording, CallRecordingSource - - return ( - db.query(CallRecording) - .filter( - CallRecording.evaluator_result_id == result.id, - CallRecording.source == CallRecordingSource.PLAYGROUND, - ) - .first() - ) - - -@celery_app.task(name="process_evaluator_result", bind=True, max_retries=3) -def process_evaluator_result_task(self, result_id: str): - """ - Celery task to process an evaluator result: transcribe audio and evaluate metrics. - - Workflow: - 1. QUEUED -> Job is created and queued - 2. TRANSCRIBING -> Audio is being transcribed - 3. EVALUATING -> Transcription is being evaluated against metrics - 4. COMPLETED -> All processing is complete - 5. FAILED -> An error occurred - """ - db = SessionLocal() - task_start_time = time.time() - - try: - from app.models.database import ( - EvaluatorResult, - EvaluatorResultStatus, - Metric, - AIProvider, - ) - - result_uuid = UUID(result_id) - result = db.query(EvaluatorResult).filter(EvaluatorResult.id == result_uuid).first() - - if not result: - logger.error(f"[EvaluatorResult {result_id}] Job not found in database") - return {"error": "Evaluator result not found"} - - logger.info(f"[EvaluatorResult {result.result_id}] Starting processing task") - - result.celery_task_id = self.request.id - db.commit() - - try: - if not result.audio_s3_key: - _recover_missing_audio_for_result(result, db, refresh_call_data=True) - - has_existing_transcript = bool(result.transcription) - if not result.audio_s3_key and not has_existing_transcript: - raise ValueError("No audio S3 key or existing transcript found") - - evaluator, agent, persona, scenario = _load_related_entities(db, result) - is_custom_evaluator = evaluator and ( - bool(evaluator.custom_prompt) - or evaluator.agent_id is None - ) - - if not is_custom_evaluator and not agent: - raise ValueError("Agent not found and no custom prompt available") - - ai_providers = db.query(AIProvider).filter( - AIProvider.organization_id == result.organization_id, - AIProvider.is_active == True, - ).all() - - # Step 1: Transcription - if has_existing_transcript: - transcription = result.transcription - speaker_segments = result.speaker_segments or [] - transcription_time = 0.0 - else: - result.status = EvaluatorResultStatus.TRANSCRIBING.value - db.commit() - - transcription, speaker_segments, transcription_time = _transcribe_audio( - result, ai_providers, db - ) - result.transcription = transcription - # Avoid duplicating transcript structure when provider call_data already carries it. - if not result.call_data: - result.speaker_segments = speaker_segments if speaker_segments else None - db.commit() - - # Step 2: Load and categorize metrics - # Include metrics that have "agent" in their enabled_surfaces so users can - # restrict metrics to specific surfaces from the metrics page. Legacy rows - # (created before the surfaces column existed, or via fixtures that don't - # set the field) keep the original behavior: an enabled=True metric with - # an empty enabled_surfaces list is treated as agent-enabled. - enabled_metrics = db.query(Metric).filter( - Metric.organization_id == result.organization_id, - Metric.enabled == True, - or_(Metric.lifecycle.is_(None), Metric.lifecycle == "active"), - ).all() - enabled_metrics = [ - m for m in enabled_metrics - if (m.name or "").strip().lower() not in REMOVED_EVALUATION_METRIC_NAMES - and ( - "agent" in (m.enabled_surfaces or []) - or not (m.enabled_surfaces or []) # legacy/unset → default to agent - ) - ] - - # Explicit metric selection on the evaluator/suite. Categorization - # parents are stored as a single ID and expanded to child labels here. - from app.services.evaluators.evaluator_helpers import expand_metric_ids_for_evaluation - - selected_metric_ids = { - str(mid) for mid in (getattr(evaluator, "metric_ids", None) or []) - } - if selected_metric_ids: - expanded_ids = expand_metric_ids_for_evaluation( - db, - result.organization_id, - list(selected_metric_ids), - ) or set() - enabled_metrics = [ - m for m in enabled_metrics if str(m.id) in expanded_ids - ] - - has_audio = bool(result.audio_s3_key) - llm_metrics, audio_metrics, metric_scores = _categorize_metrics(enabled_metrics, has_audio) - selected_metric_count = len(llm_metrics) + len(audio_metrics) - - call_recording = _playground_call_recording(db, result) - if call_recording: - from app.services.billing.flexprice_service import ( - record_playground_call_evaluated, - ) - - evaluation_attempt_id = f"{result.id}:{self.request.id}" - record_playground_call_evaluated( - result.organization_id, - evaluation_attempt_id, - evaluator_result_id=result.id, - workspace_id=result.workspace_id, - call_short_id=call_recording.call_short_id, - metric_count=selected_metric_count, - ) - - evaluation_time = None - - # Step 3: Audio metrics evaluation - if audio_metrics and has_audio: - try: - audio_scores = evaluate_audio_metrics( - audio_s3_key=result.audio_s3_key, - audio_metrics=audio_metrics, - result_id=result.result_id, - ) - - if _all_audio_scores_download_failed(audio_scores): - logger.warning( - f"[EvaluatorResult {result.result_id}] Existing S3 audio unavailable; " - "attempting provider audio recovery" - ) - recovered = _recover_missing_audio_for_result(result, db, refresh_call_data=True) - if recovered and result.audio_s3_key: - audio_scores = evaluate_audio_metrics( - audio_s3_key=result.audio_s3_key, - audio_metrics=audio_metrics, - result_id=result.result_id, - ) - - metric_scores.update(audio_scores) - except Exception as audio_err: - logger.error( - f"[EvaluatorResult {result.result_id}] Audio analysis failed: {audio_err}", - exc_info=True, - ) - metric_scores.update(handle_audio_evaluation_error(audio_metrics, audio_err)) - - # Step 4: LLM metrics evaluation - if llm_metrics and transcription: - result.status = EvaluatorResultStatus.EVALUATING.value - db.commit() - - try: - llm_scores, evaluation_time = _evaluate_llm_metrics_grouped( - transcription=transcription, - llm_metrics=llm_metrics, - ai_providers=ai_providers, - organization_id=result.organization_id, - result_id=result.result_id, - db=db, - evaluator=evaluator, - agent=agent, - persona=persona, - scenario=scenario, - ) - metric_scores.update(llm_scores) - except Exception as llm_err: - error_msg = str(llm_err).replace("{", "{{").replace("}", "}}") - logger.error( - f"[EvaluatorResult {result.result_id}] ✗ LLM evaluation failed: {error_msg}", - exc_info=True, - ) - metric_scores.update(handle_llm_evaluation_error(llm_metrics, llm_err)) - else: - if not llm_metrics: - logger.warning( - f"[EvaluatorResult {result.result_id}] No LLM-evaluable metrics found " - "(audio-only metrics were skipped), skipping evaluation" - ) - if not transcription: - logger.warning( - f"[EvaluatorResult {result.result_id}] No transcription available, " - "skipping evaluation" - ) - - # Step 5: Call Analysis - if transcription and not (result.call_data and result.call_data.get("call_analysis")): - try: - call_analysis = _generate_call_analysis( - transcription=transcription, - ai_providers=ai_providers, - organization_id=result.organization_id, - result_id=result.result_id, - db=db, - agent=agent, - scenario=scenario, - ) - if call_analysis: - existing_call_data = dict(result.call_data) if isinstance(result.call_data, dict) else {} - existing_call_data["call_analysis"] = call_analysis - result.call_data = slim_call_data_for_evaluator_result(existing_call_data) - except Exception as analysis_err: - logger.warning( - f"[EvaluatorResult {result.result_id}] Call analysis failed (non-fatal): {analysis_err}" - ) - - # Step 6: Complete - from sqlalchemy.orm.attributes import flag_modified - - result.metric_scores = _make_json_serializable(metric_scores) - flag_modified(result, "metric_scores") - if isinstance(result.call_data, dict): - result.call_data = slim_call_data_for_evaluator_result(result.call_data) - if isinstance(result.call_data, (dict, list)): - result.call_data = _make_json_serializable(result.call_data) - flag_modified(result, "call_data") - result.status = EvaluatorResultStatus.COMPLETED.value - result.error_message = None - _commit_evaluator_result(db, result) - - from app.services.billing.flexprice_service import ( - record_playground_evaluation_completed, - ) - - call_recording = _playground_call_recording(db, result) - if call_recording: - record_playground_evaluation_completed( - result.organization_id, - f"{result.id}:{self.request.id}", - evaluator_result_id=result.id, - workspace_id=result.workspace_id, - call_short_id=call_recording.call_short_id, - duration_seconds=result.duration_seconds, - metric_count=len(metric_scores) or selected_metric_count, - ) - - total_time = time.time() - task_start_time - logger.info( - f"[EvaluatorResult {result.result_id}] Completed in {total_time:.2f}s, " - f"{len(metric_scores)} metrics evaluated" - ) - - return { - "result_id": result_id, - "status": "completed", - "transcription": transcription, - "metrics_evaluated": len(metric_scores), - "processing_time": total_time, - "transcription_time": transcription_time, - "evaluation_time": evaluation_time, - } - - except Exception as e: - db.rollback() - logger.error(f"[EvaluatorResult {result_id}] Processing failed: {e}", exc_info=True) - try: - failed_result = db.query(EvaluatorResult).filter(EvaluatorResult.id == result_uuid).first() - if failed_result: - failed_result.status = EvaluatorResultStatus.FAILED.value - failed_result.error_message = str(e) - db.commit() - except Exception as persist_err: - db.rollback() - logger.error( - f"[EvaluatorResult {result_id}] Failed to persist FAILED status: {persist_err}", - exc_info=True, - ) - raise - - except Exception as exc: - raise self.retry(exc=exc, countdown=60) - finally: - db.close() +"""Celery task: process evaluator result (transcribe and evaluate metrics).""" + +import math +import time +import uuid as _uuid +from uuid import UUID + +from loguru import logger +from sqlalchemy import or_ + +from app.database import SessionLocal +from app.models.database import ModelProvider + +from app.workers.config import celery_app +from app.workers.tasks.helpers.constants import ( + REMOVED_EVALUATION_METRIC_NAMES, + AUDIO_ONLY_METRIC_NAMES, +) +from app.workers.tasks.helpers.score_utils import provider_matches, get_metric_type_value +from app.workers.tasks.helpers.audio_evaluation import ( + evaluate_audio_metrics, + handle_audio_evaluation_error, +) +from app.workers.tasks.helpers.llm_evaluation import ( + evaluate_with_llm, + handle_llm_evaluation_error, +) +from app.services.evaluators.evaluator_result_call_data import slim_call_data_for_evaluator_result +from app.services.evaluators.call_data_transcript import extract_transcript_from_call_data + + +class EvaluatorInputUnavailableError(ValueError): + """Permanent failure: no audio or transcript available for evaluation.""" + + +def _commit_evaluator_result(db, result) -> None: + from app.services.live_entity_storage import sync_evaluator_result + + sync_evaluator_result(db, result) + db.commit() + + +def _make_json_serializable(obj): + """Recursively convert non-JSON-native values (e.g., NumPy types).""" + try: + import numpy as np + except Exception: + np = None + + if isinstance(obj, dict): + return {k: _make_json_serializable(v) for k, v in obj.items()} + if isinstance(obj, list): + return [_make_json_serializable(item) for item in obj] + if isinstance(obj, tuple): + return tuple(_make_json_serializable(item) for item in obj) + + if np is not None: + if isinstance(obj, np.integer): + return int(obj) + if isinstance(obj, np.floating): + return float(obj) + if isinstance(obj, np.ndarray): + return obj.tolist() + if isinstance(obj, np.bool_): + return bool(obj) + + return obj + + +def _load_related_entities(db, result): + """Load evaluator, agent, persona, scenario from database.""" + from app.models.database import Evaluator, Agent, Persona, Scenario + + evaluator = None + if result.evaluator_id: + evaluator = db.query(Evaluator).filter(Evaluator.id == result.evaluator_id).first() + if not evaluator: + logger.warning( + f"[EvaluatorResult {result.result_id}] Evaluator {result.evaluator_id} not found, " + "continuing without evaluator" + ) + + agent = None + if result.agent_id: + agent = db.query(Agent).filter(Agent.id == result.agent_id).first() + + persona = None + if result.persona_id: + persona = db.query(Persona).filter(Persona.id == result.persona_id).first() + + scenario = None + if result.scenario_id: + scenario = db.query(Scenario).filter(Scenario.id == result.scenario_id).first() + + return evaluator, agent, persona, scenario + + +def _transcribe_audio(result, ai_providers, db): + """Transcribe audio file and return transcript with timing info.""" + from app.services.ai.transcription_service import transcription_service + + stt_provider = ModelProvider.OPENAI + stt_model = "whisper-1" + + openai_provider = next( + (p for p in ai_providers if provider_matches(p.provider, ModelProvider.OPENAI)), + None, + ) + if not openai_provider: + logger.warning( + f"[EvaluatorResult {result.result_id}] No OpenAI provider found, using default whisper-1" + ) + + transcription_start_time = time.time() + transcription_result = transcription_service.transcribe( + audio_file_key=result.audio_s3_key, + stt_provider=stt_provider, + stt_model=stt_model, + organization_id=result.organization_id, + db=db, + language=None, + enable_speaker_diarization=True, + ) + transcription_time = time.time() - transcription_start_time + + return ( + transcription_result.get("transcript", ""), + transcription_result.get("speaker_segments", []), + transcription_time, + ) + + +def _generate_call_analysis(transcription, ai_providers, organization_id, result_id, db, agent=None, scenario=None): + """Generate call analysis (summary, sentiment, success) using LLM.""" + from app.services.ai.llm_service import llm_service + + llm_provider = ModelProvider.OPENAI + llm_model = "gpt-4o-mini" + + chosen_provider = next( + (p for p in ai_providers if provider_matches(p.provider, llm_provider)), + None, + ) + if not chosen_provider: + llm_provider = ModelProvider.GOOGLE + llm_model = "gemini-2.5-flash" + chosen_provider = next( + (p for p in ai_providers if provider_matches(p.provider, llm_provider)), + None, + ) + if not chosen_provider: + logger.warning(f"[EvaluatorResult {result_id}] No LLM provider available for call analysis") + return None + + agent_context = "" + if agent and agent.description: + agent_context = f"\n\nAgent Description:\n{agent.description}" + if scenario: + scenario_name = getattr(scenario, 'name', '') + scenario_desc = getattr(scenario, 'description', '') + if scenario_name or scenario_desc: + agent_context += f"\n\nScenario: {scenario_name}" + if scenario_desc: + agent_context += f"\nScenario Description: {scenario_desc}" + + messages = [ + { + "role": "system", + "content": ( + "You are a call analysis expert. Analyze the following conversation transcript " + "and provide a structured analysis. Respond ONLY with valid JSON, no markdown." + ), + }, + { + "role": "user", + "content": f"""Analyze this conversation transcript and provide: +1. A concise summary of the call (2-3 sentences) +2. The user/caller's overall sentiment (one of: Positive, Negative, Neutral, Mixed) +3. Whether the call was successful in achieving its objective (true/false) +{agent_context} + +Transcript: +{transcription} + +Respond in this exact JSON format: +{{"call_summary": "...", "user_sentiment": "...", "call_successful": true/false}}""", + }, + ] + + try: + llm_result = llm_service.generate_response( + messages=messages, + llm_provider=llm_provider, + llm_model=llm_model, + organization_id=organization_id, + db=db, + temperature=0.3, + max_tokens=500, + ) + import json + import re + text = llm_result.get("text", "") + json_match = re.search(r'\{[^{}]*\}', text, re.DOTALL) + if json_match: + analysis = json.loads(json_match.group()) + required_keys = {"call_summary", "user_sentiment", "call_successful"} + if required_keys.issubset(analysis.keys()): + logger.info(f"[EvaluatorResult {result_id}] Call analysis generated successfully") + return analysis + logger.warning(f"[EvaluatorResult {result_id}] Could not parse call analysis from LLM response") + return None + except Exception as e: + logger.error(f"[EvaluatorResult {result_id}] Call analysis failed: {e}", exc_info=True) + return None + + +def _categorize_metrics(enabled_metrics, has_audio): + """Split metrics into LLM-evaluable and audio-only categories.""" + llm_metrics = [] + audio_metrics = [] + skipped_scores = {} + + for m in enabled_metrics: + if m.name.lower() in AUDIO_ONLY_METRIC_NAMES: + if has_audio: + audio_metrics.append(m) + else: + skipped_scores[str(m.id)] = { + "value": None, + "type": get_metric_type_value(m), + "metric_name": m.name, + "skipped": "audio_required", + } + else: + llm_metrics.append(m) + + return llm_metrics, audio_metrics, skipped_scores + + +def _evaluate_llm_metrics_grouped( + *, + transcription: str, + llm_metrics: list, + ai_providers, + organization_id, + result_id: str, + db, + evaluator, + agent, + persona, + scenario, +) -> tuple[dict, float | None]: + """Evaluate LLM metrics, grouping categorization children by parent.""" + from app.workers.tasks.evaluate_call_import_row_core import build_parent_groups + + parents_by_id, children_by_parent, standalone_metrics = build_parent_groups( + db, llm_metrics + ) + metric_scores: dict = {} + evaluation_time: float | None = None + + def _run_bucket(bucket, parent_metric=None): + nonlocal evaluation_time + scores, eval_time = evaluate_with_llm( + transcription=transcription, + llm_metrics=bucket, + ai_providers=ai_providers, + organization_id=organization_id, + result_id=result_id, + db=db, + evaluator=evaluator, + agent=agent, + persona=persona, + scenario=scenario, + parent_metric=parent_metric, + ) + metric_scores.update(scores) + if eval_time is not None: + evaluation_time = eval_time + + if standalone_metrics: + _run_bucket(standalone_metrics, parent_metric=None) + + for parent_id, children in children_by_parent.items(): + parent_metric = parents_by_id.get(parent_id) + if not parent_metric: + logger.warning( + f"[EvaluatorResult {result_id}] Parent metric {parent_id} not found; " + "evaluating children as flat metrics" + ) + _run_bucket(children, parent_metric=None) + continue + _run_bucket(children, parent_metric=parent_metric) + + return metric_scores, evaluation_time + + +_PERMANENT_VAPI_ENDED_REASONS = frozenset( + { + "call.in-progress.error-assistant-did-not-receive-customer-audio", + "call.in-progress.error-assistant-did-not-receive-customer-media", + "customer-did-not-give-microphone-permission", + "customer-did-not-answer", + } +) + + +def _call_data_for_transcript_extraction(result, db) -> tuple[dict, str]: + """Resolve provider payload and platform for transcript extraction.""" + platform = _normalize_platform(result.provider_platform) + recording = _playground_call_recording(db, result) if db is not None else None + if recording: + if not result.provider_call_id and recording.provider_call_id: + result.provider_call_id = recording.provider_call_id + if not result.provider_platform and recording.provider_platform: + result.provider_platform = recording.provider_platform + platform = _normalize_platform(result.provider_platform or recording.provider_platform) + + result_data = result.call_data if isinstance(result.call_data, dict) else {} + recording_data = ( + recording.call_data + if recording and isinstance(recording.call_data, dict) + else {} + ) + + if not result_data and not recording_data: + return {}, platform + + merged = dict(recording_data) + merged.update(result_data) + for key in ("endedReason", "messages", "transcript", "transcript_object"): + if not merged.get(key) and recording_data.get(key): + merged[key] = recording_data[key] + return merged, platform + + +def _hydrate_transcript_from_call_data(result, db) -> bool: + """Copy transcript from provider call_data onto the result row when missing.""" + if (result.transcription or "").strip(): + return True + + call_data, platform = _call_data_for_transcript_extraction(result, db) + if not call_data or not platform: + return False + + transcript_text, speaker_segments = extract_transcript_from_call_data( + call_data, platform + ) + if not (transcript_text or "").strip(): + return False + + result.transcription = transcript_text.strip() + if speaker_segments: + result.speaker_segments = speaker_segments + + if isinstance(result.call_data, dict) and result.call_data: + result.call_data = slim_call_data_for_evaluator_result(result.call_data) + + _commit_evaluator_result(db, result) + logger.info( + f"[EvaluatorResult {result.result_id}] Hydrated transcript from provider call_data" + ) + return True + + +def _permanent_input_failure_message(result, db) -> str: + """Return a user-facing error when retrying cannot succeed.""" + call_data, platform = _call_data_for_transcript_extraction(result, db) + if not isinstance(call_data, dict): + call_data = {} + + ended_reason = str( + call_data.get("endedReason") or call_data.get("ended_reason") or "" + ).strip() + if not platform and ended_reason.startswith("call."): + platform = "vapi" + + if platform == "vapi": + if ended_reason in _PERMANENT_VAPI_ENDED_REASONS: + return ( + "Call ended before customer audio was available " + f"({ended_reason}). Check microphone permissions and try again." + ) + if ended_reason: + return f"Call ended without usable audio or transcript ({ended_reason})." + + return "No audio or transcript available for evaluation." + + +def _normalize_platform(platform: object) -> str: + """Normalize provider platform enum/string into lowercase string.""" + if not platform: + return "" + if hasattr(platform, "value"): + return str(platform.value).lower() + return str(platform).lower() + + +def _extract_audio_url(call_data: dict, platform: str) -> str | None: + """Extract provider-specific audio URL from call data.""" + recording_urls = call_data.get("recording_urls", {}) if isinstance(call_data, dict) else {} + provider_payload = call_data.get("provider_payload", {}) if isinstance(call_data, dict) else {} + artifact = call_data.get("artifact", {}) if isinstance(call_data, dict) else {} + recording = artifact.get("recording", {}) if isinstance(artifact, dict) else {} + mono_recording = recording.get("mono", {}) if isinstance(recording, dict) else {} + if platform == "elevenlabs": + return recording_urls.get("conversation_audio") + if platform == "retell": + return call_data.get("recording_url") + if platform == "vapi": + return ( + call_data.get("recordingUrl") + or call_data.get("stereoRecordingUrl") + or artifact.get("recordingUrl") + or artifact.get("stereoRecordingUrl") + or mono_recording.get("combinedUrl") + or recording_urls.get("combined_url") + or recording_urls.get("stereo_url") + or call_data.get("recordingUrl") + or provider_payload.get("recordingUrl") + or provider_payload.get("stereoRecordingUrl") + ) + if platform == "smallest": + return ( + call_data.get("recording_url") + or call_data.get("recordingUrl") + or recording_urls.get("combined_url") + or recording_urls.get("conversation_audio") + ) + return None + + +def _recover_missing_audio_for_result(result, db, refresh_call_data: bool = True) -> bool: + """ + Attempt to recover missing audio from provider, upload to S3, and persist key. + + Returns True when a new S3 key is successfully stored. + """ + import requests as _http + + from app.core.encryption import decrypt_api_key + from app.models.database import Agent, Integration + from app.services.storage.s3_service import s3_service + from app.services.voice_providers import get_voice_provider + + platform = _normalize_platform(result.provider_platform) + if platform not in {"retell", "vapi", "elevenlabs", "smallest"}: + return False + if not result.provider_call_id: + return False + + agent = db.query(Agent).filter(Agent.id == result.agent_id).first() if result.agent_id else None + integration = None + decrypted_key = None + if agent and agent.voice_ai_integration_id: + integration = db.query(Integration).filter( + Integration.id == agent.voice_ai_integration_id, + Integration.organization_id == result.organization_id, + ).first() + if integration: + try: + decrypted_key = decrypt_api_key(integration.api_key) + except Exception as decrypt_err: + logger.warning( + f"[EvaluatorResult {result.result_id}] Unable to decrypt integration key " + f"for audio recovery: {decrypt_err}" + ) + + call_data = result.call_data or {} + if refresh_call_data and decrypted_key: + try: + provider_class = get_voice_provider(platform) + provider_kwargs = {"api_key": decrypted_key} + if platform == "vapi" and integration and getattr(integration, "public_key", None): + provider_kwargs["public_key"] = integration.public_key + provider = provider_class(**provider_kwargs) + if hasattr(provider, "retrieve_call_metrics"): + refreshed = provider.retrieve_call_metrics(result.provider_call_id) + if isinstance(refreshed, dict) and refreshed: + call_data = refreshed + result.call_data = refreshed + db.commit() + except Exception as refresh_err: + logger.warning( + f"[EvaluatorResult {result.result_id}] Audio recovery could not refresh provider metrics: " + f"{refresh_err}" + ) + + audio_url = _extract_audio_url(call_data, platform) + if not audio_url: + logger.warning( + f"[EvaluatorResult {result.result_id}] Audio recovery failed: no provider recording URL available" + ) + return False + + headers = None + if platform == "elevenlabs" and decrypted_key: + headers = {"xi-api-key": decrypted_key} + elif platform == "vapi" and decrypted_key: + headers = {"Authorization": f"Bearer {decrypted_key}"} + try: + response = _http.get(audio_url, headers=headers, timeout=120) + except Exception as download_err: + logger.warning( + f"[EvaluatorResult {result.result_id}] Audio recovery download failed: {download_err}" + ) + return False + + if response.status_code != 200 or not response.content: + logger.warning( + f"[EvaluatorResult {result.result_id}] Audio recovery download returned " + f"status={response.status_code}" + ) + return False + + content_type = response.headers.get("content-type", "audio/mpeg") + ext = "wav" if "wav" in content_type else "mp3" + org_id = str(result.organization_id) + s3_key = ( + f"audio/organizations/{org_id}/evaluations/" + f"{result.provider_call_id}/{_uuid.uuid4()}.{ext}" + ) + + try: + s3_service.upload_file_by_key(response.content, s3_key, content_type=content_type) + except Exception as upload_err: + logger.warning( + f"[EvaluatorResult {result.result_id}] Audio recovery upload failed: {upload_err}" + ) + return False + + result.audio_s3_key = s3_key + _commit_evaluator_result(db, result) + db.refresh(result) + logger.info( + f"[EvaluatorResult {result.result_id}] Recovered missing audio and stored at {s3_key}" + ) + return True + + +def _all_audio_scores_download_failed(audio_scores: dict[str, dict[str, object]]) -> bool: + """Check whether every audio metric failed due to missing/unreadable S3 object.""" + if not audio_scores: + return False + return all( + isinstance(score, dict) and score.get("error") == "audio_download_failed" + for score in audio_scores.values() + ) + + +def _playground_call_recording(db, result): + from app.models.database import CallRecording, CallRecordingSource + + return ( + db.query(CallRecording) + .filter( + CallRecording.evaluator_result_id == result.id, + CallRecording.source == CallRecordingSource.PLAYGROUND, + ) + .first() + ) + + +_EXTERNAL_VOICE_PROVIDER_PLATFORMS = frozenset( + {"vapi", "retell", "elevenlabs", "smallest"} +) + + +def _should_record_external_agent_call_usage(result) -> bool: + """Only bill completed calls that ran on an external voice provider. + + Internal voice-bundle / WebSocket calls already emit LLM/STT/TTS usage + from the live pipeline; recording again here would double-count. + """ + if not result.agent_id: + return False + platform = (getattr(result, "provider_platform", None) or "").strip().lower() + return platform in _EXTERNAL_VOICE_PROVIDER_PLATFORMS + + +def _resolve_call_duration_seconds(result) -> int: + if result.duration_seconds: + try: + return max(0, int(math.ceil(float(result.duration_seconds)))) + except (TypeError, ValueError): + pass + call_data = result.call_data if isinstance(result.call_data, dict) else {} + for key in ("duration_seconds", "duration"): + raw = call_data.get(key) + if raw is not None: + try: + return max(0, int(math.ceil(float(raw)))) + except (TypeError, ValueError): + continue + started_at = call_data.get("startedAt") or call_data.get("start_timestamp") + ended_at = call_data.get("endedAt") or call_data.get("end_timestamp") + if started_at and ended_at: + try: + from datetime import datetime + + start = datetime.fromisoformat(str(started_at).replace("Z", "+00:00")) + end = datetime.fromisoformat(str(ended_at).replace("Z", "+00:00")) + return max(0, int(math.ceil((end - start).total_seconds()))) + except Exception: + return 0 + return 0 + + +def _record_agent_call_usage(result, *, usage_ctx) -> None: + if not _should_record_external_agent_call_usage(result): + return + try: + from app.services.usage.llm_usage import record_call_usage + + record_call_usage( + "voice-agent-call", + organization_id=result.organization_id, + ctx=usage_ctx, + audio_seconds=_resolve_call_duration_seconds(result), + ) + except Exception as exc: + logger.debug( + "[EvaluatorResult {}] agent call usage record skipped: {}", + result.result_id, + exc, + ) + + +@celery_app.task(name="process_evaluator_result", bind=True, max_retries=3) +def process_evaluator_result_task(self, result_id: str): + """ + Celery task to process an evaluator result: transcribe audio and evaluate metrics. + + Workflow: + 1. QUEUED -> Job is created and queued + 2. TRANSCRIBING -> Audio is being transcribed + 3. EVALUATING -> Transcription is being evaluated against metrics + 4. COMPLETED -> All processing is complete + 5. FAILED -> An error occurred + """ + db = SessionLocal() + task_start_time = time.time() + + try: + from app.models.database import ( + EvaluatorResult, + EvaluatorResultStatus, + Metric, + AIProvider, + ) + + result_uuid = UUID(result_id) + result = db.query(EvaluatorResult).filter(EvaluatorResult.id == result_uuid).first() + + if not result: + logger.error(f"[EvaluatorResult {result_id}] Job not found in database") + return {"error": "Evaluator result not found"} + + from app.services.usage.context import ( + llm_usage_context, + usage_context_for_evaluator_result, + ) + + usage_ctx = usage_context_for_evaluator_result(result) + with llm_usage_context(usage_ctx): + logger.info(f"[EvaluatorResult {result.result_id}] Starting processing task") + + result.celery_task_id = self.request.id + db.commit() + + try: + if not result.audio_s3_key: + _recover_missing_audio_for_result(result, db, refresh_call_data=True) + + _hydrate_transcript_from_call_data(result, db) + + has_existing_transcript = bool((result.transcription or "").strip()) + if not result.audio_s3_key and not has_existing_transcript: + raise EvaluatorInputUnavailableError( + _permanent_input_failure_message(result, db) + ) + + evaluator, agent, persona, scenario = _load_related_entities(db, result) + is_custom_evaluator = evaluator and ( + bool(evaluator.custom_prompt) + or evaluator.agent_id is None + ) + + if not is_custom_evaluator and not agent: + raise ValueError("Agent not found and no custom prompt available") + + ai_providers = db.query(AIProvider).filter( + AIProvider.organization_id == result.organization_id, + AIProvider.is_active == True, + ).all() + + # Step 1: Transcription + if has_existing_transcript: + transcription = result.transcription + speaker_segments = result.speaker_segments or [] + transcription_time = 0.0 + else: + result.status = EvaluatorResultStatus.TRANSCRIBING.value + db.commit() + + transcription, speaker_segments, transcription_time = _transcribe_audio( + result, ai_providers, db + ) + result.transcription = transcription + # Avoid duplicating transcript structure when provider call_data already carries it. + if not result.call_data: + result.speaker_segments = speaker_segments if speaker_segments else None + db.commit() + + # Step 2: Load and categorize metrics + # Include metrics that have "agent" in their enabled_surfaces so users can + # restrict metrics to specific surfaces from the metrics page. Legacy rows + # (created before the surfaces column existed, or via fixtures that don't + # set the field) keep the original behavior: an enabled=True metric with + # an empty enabled_surfaces list is treated as agent-enabled. + enabled_metrics = db.query(Metric).filter( + Metric.organization_id == result.organization_id, + Metric.enabled == True, + or_(Metric.lifecycle.is_(None), Metric.lifecycle == "active"), + ).all() + enabled_metrics = [ + m for m in enabled_metrics + if (m.name or "").strip().lower() not in REMOVED_EVALUATION_METRIC_NAMES + and ( + "agent" in (m.enabled_surfaces or []) + or not (m.enabled_surfaces or []) # legacy/unset → default to agent + ) + ] + + # Explicit metric selection on the evaluator/suite. Categorization + # parents are stored as a single ID and expanded to child labels here. + from app.services.evaluators.evaluator_helpers import expand_metric_ids_for_evaluation + + selected_metric_ids = { + str(mid) for mid in (getattr(evaluator, "metric_ids", None) or []) + } + if selected_metric_ids: + expanded_ids = expand_metric_ids_for_evaluation( + db, + result.organization_id, + list(selected_metric_ids), + ) or set() + enabled_metrics = [ + m for m in enabled_metrics if str(m.id) in expanded_ids + ] + + has_audio = bool(result.audio_s3_key) + llm_metrics, audio_metrics, metric_scores = _categorize_metrics(enabled_metrics, has_audio) + selected_metric_count = len(llm_metrics) + len(audio_metrics) + + call_recording = _playground_call_recording(db, result) + if call_recording: + from app.services.billing.flexprice_service import ( + record_playground_call_evaluated, + ) + + evaluation_attempt_id = f"{result.id}:{self.request.id}" + record_playground_call_evaluated( + result.organization_id, + evaluation_attempt_id, + evaluator_result_id=result.id, + workspace_id=result.workspace_id, + call_short_id=call_recording.call_short_id, + metric_count=selected_metric_count, + ) + + evaluation_time = None + + # Step 3: Audio metrics evaluation + if audio_metrics and has_audio: + try: + audio_scores = evaluate_audio_metrics( + audio_s3_key=result.audio_s3_key, + audio_metrics=audio_metrics, + result_id=result.result_id, + ) + + if _all_audio_scores_download_failed(audio_scores): + logger.warning( + f"[EvaluatorResult {result.result_id}] Existing S3 audio unavailable; " + "attempting provider audio recovery" + ) + recovered = _recover_missing_audio_for_result(result, db, refresh_call_data=True) + if recovered and result.audio_s3_key: + audio_scores = evaluate_audio_metrics( + audio_s3_key=result.audio_s3_key, + audio_metrics=audio_metrics, + result_id=result.result_id, + ) + + metric_scores.update(audio_scores) + except Exception as audio_err: + logger.error( + f"[EvaluatorResult {result.result_id}] Audio analysis failed: {audio_err}", + exc_info=True, + ) + metric_scores.update(handle_audio_evaluation_error(audio_metrics, audio_err)) + + # Step 4: LLM metrics evaluation + if llm_metrics and transcription: + result.status = EvaluatorResultStatus.EVALUATING.value + db.commit() + + try: + llm_scores, evaluation_time = _evaluate_llm_metrics_grouped( + transcription=transcription, + llm_metrics=llm_metrics, + ai_providers=ai_providers, + organization_id=result.organization_id, + result_id=result.result_id, + db=db, + evaluator=evaluator, + agent=agent, + persona=persona, + scenario=scenario, + ) + metric_scores.update(llm_scores) + except Exception as llm_err: + error_msg = str(llm_err).replace("{", "{{").replace("}", "}}") + logger.error( + f"[EvaluatorResult {result.result_id}] ✗ LLM evaluation failed: {error_msg}", + exc_info=True, + ) + metric_scores.update(handle_llm_evaluation_error(llm_metrics, llm_err)) + else: + if not llm_metrics: + logger.warning( + f"[EvaluatorResult {result.result_id}] No LLM-evaluable metrics found " + "(audio-only metrics were skipped), skipping evaluation" + ) + if not transcription: + logger.warning( + f"[EvaluatorResult {result.result_id}] No transcription available, " + "skipping evaluation" + ) + + # Step 5: Call Analysis + if transcription and not (result.call_data and result.call_data.get("call_analysis")): + try: + call_analysis = _generate_call_analysis( + transcription=transcription, + ai_providers=ai_providers, + organization_id=result.organization_id, + result_id=result.result_id, + db=db, + agent=agent, + scenario=scenario, + ) + if call_analysis: + existing_call_data = dict(result.call_data) if isinstance(result.call_data, dict) else {} + existing_call_data["call_analysis"] = call_analysis + result.call_data = slim_call_data_for_evaluator_result(existing_call_data) + except Exception as analysis_err: + logger.warning( + f"[EvaluatorResult {result.result_id}] Call analysis failed (non-fatal): {analysis_err}" + ) + + # Step 6: Complete + from sqlalchemy.orm.attributes import flag_modified + + result.metric_scores = _make_json_serializable(metric_scores) + flag_modified(result, "metric_scores") + if isinstance(result.call_data, dict): + result.call_data = slim_call_data_for_evaluator_result(result.call_data) + if isinstance(result.call_data, (dict, list)): + result.call_data = _make_json_serializable(result.call_data) + flag_modified(result, "call_data") + result.status = EvaluatorResultStatus.COMPLETED.value + result.error_message = None + _record_agent_call_usage(result, usage_ctx=usage_ctx) + _commit_evaluator_result(db, result) + + from app.services.billing.flexprice_service import ( + record_playground_evaluation_completed, + ) + + call_recording = _playground_call_recording(db, result) + if call_recording: + record_playground_evaluation_completed( + result.organization_id, + f"{result.id}:{self.request.id}", + evaluator_result_id=result.id, + workspace_id=result.workspace_id, + call_short_id=call_recording.call_short_id, + duration_seconds=result.duration_seconds, + metric_count=len(metric_scores) or selected_metric_count, + ) + + total_time = time.time() - task_start_time + logger.info( + f"[EvaluatorResult {result.result_id}] Completed in {total_time:.2f}s, " + f"{len(metric_scores)} metrics evaluated" + ) + + return { + "result_id": result_id, + "status": "completed", + "transcription": transcription, + "metrics_evaluated": len(metric_scores), + "processing_time": total_time, + "transcription_time": transcription_time, + "evaluation_time": evaluation_time, + } + + except EvaluatorInputUnavailableError: + db.rollback() + raise + except Exception as e: + db.rollback() + logger.error(f"[EvaluatorResult {result_id}] Processing failed: {e}", exc_info=True) + try: + failed_result = db.query(EvaluatorResult).filter(EvaluatorResult.id == result_uuid).first() + if failed_result: + failed_result.status = EvaluatorResultStatus.FAILED.value + failed_result.error_message = str(e) + db.commit() + except Exception as persist_err: + db.rollback() + logger.error( + f"[EvaluatorResult {result_id}] Failed to persist FAILED status: {persist_err}", + exc_info=True, + ) + raise + + except EvaluatorInputUnavailableError as exc: + logger.warning(f"[EvaluatorResult {result_id}] Input unavailable: {exc}") + try: + from app.models.database import EvaluatorResult, EvaluatorResultStatus + + failed_result = db.query(EvaluatorResult).filter( + EvaluatorResult.id == result_uuid + ).first() + if failed_result: + failed_result.status = EvaluatorResultStatus.FAILED.value + failed_result.error_message = str(exc) + db.commit() + except Exception as persist_err: + db.rollback() + logger.error( + f"[EvaluatorResult {result_id}] Failed to persist FAILED status: {persist_err}", + exc_info=True, + ) + return {"error": str(exc), "status": "failed"} + except Exception as exc: + raise self.retry(exc=exc, countdown=60) + finally: + db.close() diff --git a/app/workers/tasks/prune_oss_usage_history.py b/app/workers/tasks/prune_oss_usage_history.py new file mode 100644 index 00000000..a91b186c --- /dev/null +++ b/app/workers/tasks/prune_oss_usage_history.py @@ -0,0 +1,22 @@ +"""Celery task: prune OSS usage rollup rows beyond history window.""" + +from __future__ import annotations + +from loguru import logger + +from app.database import SessionLocal +from app.workers.config import celery_app + + +@celery_app.task(name="prune_oss_usage_history") +def prune_oss_usage_history_task() -> dict: + from app.services.usage.retention import prune_oss_usage_history + + db = SessionLocal() + try: + result = prune_oss_usage_history(db) + finally: + db.close() + if result.get("deleted"): + logger.info("prune_oss_usage_history_task {}", result) + return result diff --git a/app/workers/tasks/recompute_usage_costs.py b/app/workers/tasks/recompute_usage_costs.py new file mode 100644 index 00000000..d1a2aea7 --- /dev/null +++ b/app/workers/tasks/recompute_usage_costs.py @@ -0,0 +1,83 @@ +"""Celery task: recompute stored usage costs for rollup rows.""" + +from __future__ import annotations + +from datetime import date +from typing import Optional +from uuid import UUID + +from loguru import logger + +from app.database import SessionLocal +from app.workers.config import celery_app + + +@celery_app.task(name="recompute_usage_costs") +def recompute_usage_costs_task( + job_id: Optional[str] = None, + organization_id: Optional[str] = None, + model: Optional[str] = None, + usage_kind: Optional[str] = None, + start_date: Optional[str] = None, + end_date: Optional[str] = None, +) -> dict: + from app.models.database import UsageCostRecomputeJob + from app.services.usage.pricing import recompute_usage_costs + from app.services.usage.pricing_jobs import ( + mark_job_completed, + mark_job_failed, + mark_job_running, + update_job_progress, + ) + + db = SessionLocal() + try: + if job_id: + job = ( + db.query(UsageCostRecomputeJob) + .filter(UsageCostRecomputeJob.id == UUID(job_id)) + .first() + ) + if job is None: + return {"updated_rows": 0, "error": "job not found"} + mark_job_running(db, job.id) + org_uuid = job.organization_id + model = job.model + usage_kind = job.usage_kind + start = job.start_date + end = job.end_date + progress_job_id = job.id + else: + org_uuid = UUID(organization_id) if organization_id else None + start = date.fromisoformat(start_date) if start_date else None + end = date.fromisoformat(end_date) if end_date else None + progress_job_id = None + + def _on_progress(updated_rows: int) -> None: + if progress_job_id is not None: + update_job_progress(db, progress_job_id, updated_rows) + + updated = recompute_usage_costs( + db, + organization_id=org_uuid, + model=model, + usage_kind=usage_kind, + start_date=start, + end_date=end, + on_progress=_on_progress if progress_job_id else None, + ) + if progress_job_id is not None: + mark_job_completed(db, progress_job_id, updated) + if updated: + logger.info("Recomputed usage costs for {} rollup row(s)", updated) + if org_uuid is not None and updated: + from app.services.usage.read_cache import invalidate_org_usage_read_cache + + invalidate_org_usage_read_cache(org_uuid) + return {"updated_rows": updated, "job_id": job_id} + except Exception as exc: + if job_id: + mark_job_failed(db, UUID(job_id), str(exc)) + raise + finally: + db.close() diff --git a/app/workers/tasks/refresh_fx_rates.py b/app/workers/tasks/refresh_fx_rates.py new file mode 100644 index 00000000..170e9103 --- /dev/null +++ b/app/workers/tasks/refresh_fx_rates.py @@ -0,0 +1,12 @@ +"""Celery task: refresh USD/INR display FX rate.""" + +from __future__ import annotations + +from app.workers.config import celery_app + + +@celery_app.task(name="refresh_fx_rates") +def refresh_fx_rates_task() -> dict: + from app.services.usage.fx_rates import refresh_usd_inr_rate + + return refresh_usd_inr_rate() diff --git a/app/workers/tasks/run_cron_evaluator_job.py b/app/workers/tasks/run_cron_evaluator_job.py new file mode 100644 index 00000000..508ba1b2 --- /dev/null +++ b/app/workers/tasks/run_cron_evaluator_job.py @@ -0,0 +1,23 @@ +"""Celery task: run evaluator_ids for an org cron job.""" + +from __future__ import annotations + +from uuid import UUID + +from app.database import SessionLocal +from app.models.database import CronJob +from app.workers.config import celery_app + + +@celery_app.task(name="run_cron_evaluator_job") +def run_cron_evaluator_job_task(job_id: str) -> dict: + from app.services.cron.job_dispatch import run_evaluator_cron_job + + db = SessionLocal() + try: + job = db.query(CronJob).filter(CronJob.id == UUID(job_id)).first() + if job is None: + return {"error": "job not found", "job_id": job_id} + return run_evaluator_cron_job(db, job) + finally: + db.close() diff --git a/app/workers/tasks/run_judge_alignment.py b/app/workers/tasks/run_judge_alignment.py index 6ffbb655..053546cc 100644 --- a/app/workers/tasks/run_judge_alignment.py +++ b/app/workers/tasks/run_judge_alignment.py @@ -79,7 +79,13 @@ def run_judge_alignment_task( return {"error": "No samples"} try: - metrics = run_judge(run, dataset, evaluator, samples, db) + from app.services.usage.context import ( + llm_usage_context, + usage_context_for_judge_run, + ) + + with llm_usage_context(usage_context_for_judge_run(run)): + metrics = run_judge(run, dataset, evaluator, samples, db) except Exception as exc: logger.error( f"[JudgeAlignment] Run {judge_run_id} crashed: {exc}", diff --git a/app/workers/tasks/run_prompt_optimization.py b/app/workers/tasks/run_prompt_optimization.py index dbd9c869..b7efb247 100644 --- a/app/workers/tasks/run_prompt_optimization.py +++ b/app/workers/tasks/run_prompt_optimization.py @@ -106,19 +106,24 @@ def run_prompt_optimization_task(self, optimization_run_id: str): ) from app.services.optimization import run_optimization - - result = run_optimization( - agent=agent, - evaluator=evaluator, - voice_bundle=voice_bundle, - training_data=training_data, - metrics=enabled_metrics, - ai_providers=ai_providers, - organization_id=run.organization_id, - db=db, - config=run.config, + from app.services.usage.context import ( + llm_usage_context, + usage_context_for_prompt_optimization_run, ) + with llm_usage_context(usage_context_for_prompt_optimization_run(run)): + result = run_optimization( + agent=agent, + evaluator=evaluator, + voice_bundle=voice_bundle, + training_data=training_data, + metrics=enabled_metrics, + ai_providers=ai_providers, + organization_id=run.organization_id, + db=db, + config=run.config, + ) + run.best_prompt = result["best_candidate"] run.best_score = result["best_score"] run.metric_history = result["metric_history"] diff --git a/app/workers/tasks/transcribe_call_import_row.py b/app/workers/tasks/transcribe_call_import_row.py index 68368b93..525a7a48 100644 --- a/app/workers/tasks/transcribe_call_import_row.py +++ b/app/workers/tasks/transcribe_call_import_row.py @@ -391,6 +391,8 @@ def _persist_diarization_failure( def _run_diarization_pipeline(ctx: dict[str, Any]) -> dict[str, Any]: """STT / S3 / LLM diarisation without a long-lived DB session.""" from app.models.enums import ModelProvider + from app.services.usage.call_import_context import call_import_row_usage_context + from app.services.usage.context import llm_usage_context from app.workers.tasks.helpers.llm_diarisation import ( LLMDiarisationError, diarize_audio_with_llm, @@ -406,6 +408,52 @@ def _run_diarization_pipeline(ctx: dict[str, Any]) -> dict[str, Any]: llm_credential_uuid = ctx["llm_credential_uuid"] effective_prompt = ctx["effective_prompt"] + evaluation_id = ctx.get("evaluation_id") + call_import_id = ctx.get("call_import_id") + if not call_import_id: + raise ValueError("call_import_id missing from transcribe pipeline context") + + with llm_usage_context( + call_import_row_usage_context( + organization_id=organization_id, + workspace_id=ctx.get("workspace_id"), + call_import_id=call_import_id, + evaluation_id=evaluation_id, + ) + ): + return _run_diarization_pipeline_inner( + ctx, + row_id=row_id, + normalised_mode=normalised_mode, + recording_key=recording_key, + organization_id=organization_id, + llm_provider_value=llm_provider_value, + llm_model_value=llm_model_value, + llm_credential_uuid=llm_credential_uuid, + effective_prompt=effective_prompt, + ModelProvider=ModelProvider, + LLMDiarisationError=LLMDiarisationError, + diarize_audio_with_llm=diarize_audio_with_llm, + diarize_transcript_with_llm=diarize_transcript_with_llm, + ) + + +def _run_diarization_pipeline_inner( + ctx: dict[str, Any], + *, + row_id, + normalised_mode, + recording_key, + organization_id, + llm_provider_value, + llm_model_value, + llm_credential_uuid, + effective_prompt, + ModelProvider, + LLMDiarisationError, + diarize_audio_with_llm, + diarize_transcript_with_llm, +) -> dict[str, Any]: plain_text: Optional[str] = None raw_turns: Optional[List[Dict[str, Any]]] = None @@ -923,6 +971,16 @@ def transcribe_call_import_row_task( "normalised_mode": normalised_mode, "recording_key": recording_key, "organization_id": row.organization_id, + "workspace_id": getattr(row, "workspace_id", None), + "call_import_id": row.call_import_id, + "evaluation_id": ( + UUID(evaluation_id_for_dispatch) + if evaluation_id_for_dispatch + else None + ), + "evaluation_row_id": ( + UUID(run_eval_row_id) if run_eval_row_id else None + ), "stt_provider": provider_enum.value if provider_enum else None, "stt_model": stt_model, "credential_uuid": credential_uuid, diff --git a/app/workers/tasks/tts_comparison.py b/app/workers/tasks/tts_comparison.py index 8ae42ea9..2c8968ac 100644 --- a/app/workers/tasks/tts_comparison.py +++ b/app/workers/tasks/tts_comparison.py @@ -180,6 +180,7 @@ def generate_tts_comparison_task(self, comparison_id: str): ) from app.services.ai.tts_service import tts_service, get_audio_file_extension from app.services.storage.s3_service import s3_service + from app.services.usage.context import LLMUsageContext, LLMUsageProductSection, llm_usage_context db = SessionLocal() try: @@ -258,15 +259,24 @@ def _resolve_voice_meta(sample_obj): f"provider={sample.provider} voice={sample.voice_id} " f"sample_rate_hz={sample_rate_hz} config={tts_config}" ) - audio_bytes, latency_ms, ttfb_ms = tts_service.synthesize_timed( - text=sample.text, - tts_provider=provider_enum, - tts_model=sample.model, - organization_id=comp.organization_id, - db=db, - voice=sample.voice_id, - config=tts_config or None, - ) + with llm_usage_context( + LLMUsageContext( + organization_id=comp.organization_id, + workspace_id=comp.workspace_id, + product_section=LLMUsageProductSection.VOICE_PLAYGROUND, + resource_id=comp.id, + resource_type="tts_comparison", + ) + ): + audio_bytes, latency_ms, ttfb_ms = tts_service.synthesize_timed( + text=sample.text, + tts_provider=provider_enum, + tts_model=sample.model, + organization_id=comp.organization_id, + db=db, + voice=sample.voice_id, + config=tts_config or None, + ) audio_ext = get_audio_file_extension( sample.provider, int(sample_rate_hz) if sample_rate_hz else None @@ -514,6 +524,20 @@ def evaluate_tts_comparison_task(self, comparison_id: str): db.commit() return {"evaluated": 0} + from app.services.usage.context import ( + LLMUsageContext, + LLMUsageProductSection, + llm_usage_context, + ) + + usage_ctx = LLMUsageContext( + organization_id=comp.organization_id, + workspace_id=comp.workspace_id, + product_section=LLMUsageProductSection.VOICE_PLAYGROUND, + resource_id=comp.id, + resource_type="tts_comparison", + ) + stt_provider_str, stt_model = _resolve_stt_config(comp, db) stt_available = bool(stt_provider_str and stt_model) if stt_available: @@ -535,86 +559,86 @@ def evaluate_tts_comparison_task(self, comparison_id: str): ) evaluated = 0 - for sample in samples: - tmp_path = None - try: - audio_bytes = s3_service.download_file_by_key(sample.audio_s3_key) - if not audio_bytes: - continue - - ext = ".mp3" - if sample.audio_s3_key: - key_ext = os.path.splitext(sample.audio_s3_key)[1].lower() - if key_ext in {".wav", ".mp3", ".flac", ".ogg", ".m4a"}: - ext = key_ext - tmp_fd, tmp_path = tempfile.mkstemp(suffix=ext) - os.close(tmp_fd) - with open(tmp_path, "wb") as f: - f.write(audio_bytes) - - if any_qualitative_enabled: - metrics = qualitative_voice_service.calculate_all_metrics(tmp_path) - metrics = _filter_qualitative_metrics(metrics, enabled_voice_metric_names) - else: - metrics = {} - - if stt_available and sample.text: - selected_stt_provider = ModelProvider(stt_provider_str) - selected_language = None - if selected_stt_provider == ModelProvider.SARVAM: - # Sarvam saarika models accept language_code; saaras models do not. - if "saarika" in (stt_model or "").lower(): - selected_language = "hi-IN" - asr_transcript = transcription_service.transcribe_text_only( - audio_file_path=tmp_path, - stt_provider=selected_stt_provider, - stt_model=stt_model, - organization_id=comp.organization_id, - db=db, - language=selected_language, - ) - if asr_transcript: - score_bundle = _compute_wer_cer(sample.text, asr_transcript) - metrics["WER Raw"] = score_bundle.get("raw_wer") - metrics["CER Raw"] = score_bundle.get("raw_cer") - metrics["WER Normalized"] = score_bundle.get("normalized_wer") - metrics["CER Normalized"] = score_bundle.get("normalized_cer") - metrics["WER"] = ( - score_bundle.get("normalized_wer") - if score_bundle.get("normalized_wer") is not None - else score_bundle.get("raw_wer") - ) - metrics["CER"] = ( - score_bundle.get("normalized_cer") - if score_bundle.get("normalized_cer") is not None - else score_bundle.get("raw_cer") - ) - metrics["ASR Transcript"] = asr_transcript + with llm_usage_context(usage_ctx): + for sample in samples: + tmp_path = None + try: + audio_bytes = s3_service.download_file_by_key(sample.audio_s3_key) + if not audio_bytes: + continue + + ext = ".mp3" + if sample.audio_s3_key: + key_ext = os.path.splitext(sample.audio_s3_key)[1].lower() + if key_ext in {".wav", ".mp3", ".flac", ".ogg", ".m4a"}: + ext = key_ext + tmp_fd, tmp_path = tempfile.mkstemp(suffix=ext) + os.close(tmp_fd) + with open(tmp_path, "wb") as f: + f.write(audio_bytes) + + if any_qualitative_enabled: + metrics = qualitative_voice_service.calculate_all_metrics(tmp_path) + metrics = _filter_qualitative_metrics(metrics, enabled_voice_metric_names) else: - metrics["WER"] = None - metrics["CER"] = None - metrics["WER Raw"] = None - metrics["CER Raw"] = None - metrics["WER Normalized"] = None - metrics["CER Normalized"] = None - metrics["ASR Transcript"] = None - - _evaluate_custom_voice_metrics(sample, metrics, comp, db) - - sample.evaluation_metrics = _sanitize_metrics(metrics) - db.commit() - evaluated += 1 - - logger.info(f"[TTS Eval] Sample {sample.id} metrics: {metrics}") - - except Exception as e: - logger.warning("[TTS Eval] Sample {} eval failed: {}", sample.id, e) - finally: - if tmp_path and os.path.exists(tmp_path): - try: - os.unlink(tmp_path) - except Exception: - pass + metrics = {} + + if stt_available and sample.text: + selected_stt_provider = ModelProvider(stt_provider_str) + selected_language = None + if selected_stt_provider == ModelProvider.SARVAM: + if "saarika" in (stt_model or "").lower(): + selected_language = "hi-IN" + asr_transcript = transcription_service.transcribe_text_only( + audio_file_path=tmp_path, + stt_provider=selected_stt_provider, + stt_model=stt_model, + organization_id=comp.organization_id, + db=db, + language=selected_language, + ) + if asr_transcript: + score_bundle = _compute_wer_cer(sample.text, asr_transcript) + metrics["WER Raw"] = score_bundle.get("raw_wer") + metrics["CER Raw"] = score_bundle.get("raw_cer") + metrics["WER Normalized"] = score_bundle.get("normalized_wer") + metrics["CER Normalized"] = score_bundle.get("normalized_cer") + metrics["WER"] = ( + score_bundle.get("normalized_wer") + if score_bundle.get("normalized_wer") is not None + else score_bundle.get("raw_wer") + ) + metrics["CER"] = ( + score_bundle.get("normalized_cer") + if score_bundle.get("normalized_cer") is not None + else score_bundle.get("raw_cer") + ) + metrics["ASR Transcript"] = asr_transcript + else: + metrics["WER"] = None + metrics["CER"] = None + metrics["WER Raw"] = None + metrics["CER Raw"] = None + metrics["WER Normalized"] = None + metrics["CER Normalized"] = None + metrics["ASR Transcript"] = None + + _evaluate_custom_voice_metrics(sample, metrics, comp, db) + + sample.evaluation_metrics = _sanitize_metrics(metrics) + db.commit() + evaluated += 1 + + logger.info(f"[TTS Eval] Sample {sample.id} metrics: {metrics}") + + except Exception as e: + logger.warning("[TTS Eval] Sample {} eval failed: {}", sample.id, e) + finally: + if tmp_path and os.path.exists(tmp_path): + try: + os.unlink(tmp_path) + except Exception: + pass from app.api.v1.routes.voice_playground import _recompute_summary diff --git a/docker-compose.observability.yml b/docker-compose.observability.yml index 402ac2e0..35b84190 100644 --- a/docker-compose.observability.yml +++ b/docker-compose.observability.yml @@ -62,6 +62,17 @@ services: loki-max-backoff: "800ms" loki-external-labels: "service=worker-imports,environment=docker" + worker-usage: + depends_on: + - loki + logging: + driver: loki + options: + loki-url: "http://localhost:3100/loki/api/v1/push" + loki-retries: "5" + loki-max-backoff: "800ms" + loki-external-labels: "service=worker-usage,environment=docker" + # Pin image tags instead of :latest to reduce exposure to stale third-party # binaries (e.g. Go stdlib CVEs in observability tooling). loki: diff --git a/docker-compose.yml b/docker-compose.yml index 552884f3..1d6f0d0b 100644 --- a/docker-compose.yml +++ b/docker-compose.yml @@ -154,12 +154,45 @@ services: # - ./secrets/gcp-sa.json:/app/secrets/gcp-sa.json:ro command: eai worker --config /app/config.yml --loglevel info --queues celery,audio-metrics --concurrency 8 + # Celery Beat — platform schedules + platform task worker (single replica; do not scale). + # Beat enqueues tasks; co-located worker runs evaluate_alerts, refresh_fx_rates, prune_oss_usage_history. + # flush_usage_counters still runs on worker-usage (usage queue). + beat: + image: ghcr.io/efficientai-tech/efficientai-worker:${EFFICIENTAI_VERSION:-latest} + build: + context: . + dockerfile: docker/Dockerfile.worker + args: + INSTALL_EXTRAS: "" + container_name: efficientai_beat + env_file: + - .env + environment: + DATABASE_URL: postgresql://${POSTGRES_USER:-efficientai}:${POSTGRES_PASSWORD:-password}@db:5432/${POSTGRES_DB:-efficientai} + REDIS_URL: redis://redis:6379/0 + CELERY_BROKER_URL: redis://redis:6379/0 + CELERY_RESULT_BACKEND: redis://redis:6379/0 + ENCRYPTION_KEY: ${ENCRYPTION_KEY:-} + depends_on: + db: + condition: service_healthy + redis: + condition: service_healthy + volumes: + - ./uploads:/app/uploads + - ./.data:/app/.data + - ./config.docker.yml:/app/config.yml:ro + command: > + sh -c "celery -A app.workers.celery_app worker -Q platform --pool threads --concurrency 2 --loglevel=info & + exec celery -A app.workers.celery_app beat --loglevel=info" + # Dedicated worker for call-import + evaluation queues. Celery drains # ``imports`` (recording fetch) before ``diarization`` (manual diarise), # then ``eval-control`` (cancel/retry/materialize), then ``evaluations`` # (fair dispatch + LLM scoring). - # The default ``worker`` service handles ``celery`` (legacy) and - # ``audio-metrics`` (Praat/UTMOS audio metric tasks). + # The default ``worker`` service handles ``celery`` (evaluator cron dispatch only) + # and ``audio-metrics`` (Praat/UTMOS audio metric tasks). + # Platform schedules run via ``beat`` (scheduler + ``platform`` queue worker). worker-imports: image: ghcr.io/efficientai-tech/efficientai-worker:${EFFICIENTAI_VERSION:-latest} build: @@ -193,5 +226,33 @@ services: # (each task may hold catalog + shard connections for tens of seconds). command: eai worker --config /app/config.yml --loglevel info --queues imports,diarization,eval-control,evaluations --pool threads --concurrency 12 + # Low-priority usage pricing: Redis flush + cost recompute/backfill. + worker-usage: + image: ghcr.io/efficientai-tech/efficientai-worker:${EFFICIENTAI_VERSION:-latest} + build: + context: . + dockerfile: docker/Dockerfile.worker + args: + INSTALL_EXTRAS: "" + container_name: efficientai_worker_usage + env_file: + - .env + environment: + DATABASE_URL: postgresql://${POSTGRES_USER:-efficientai}:${POSTGRES_PASSWORD:-password}@db:5432/${POSTGRES_DB:-efficientai} + REDIS_URL: redis://redis:6379/0 + CELERY_BROKER_URL: redis://redis:6379/0 + CELERY_RESULT_BACKEND: redis://redis:6379/0 + ENCRYPTION_KEY: ${ENCRYPTION_KEY:-} + depends_on: + db: + condition: service_healthy + redis: + condition: service_healthy + volumes: + - ./uploads:/app/uploads + - ./.data:/app/.data + - ./config.docker.yml:/app/config.yml:ro + command: eai worker --config /app/config.yml --loglevel info --queues usage --pool threads --concurrency 4 + volumes: postgres_data: diff --git a/docs-fumadocs/content/docs/intro.mdx b/docs-fumadocs/content/docs/intro.mdx index 36d2e905..f95f6f49 100644 --- a/docs-fumadocs/content/docs/intro.mdx +++ b/docs-fumadocs/content/docs/intro.mdx @@ -33,6 +33,7 @@ EfficientAI gives you an end-to-end loop for voice AI quality: - **Metric-driven evaluation**: built-in and custom metrics with categorization labels, surface targeting, and org-wide or workspace scope. - **Cloud storage**: store audio in Amazon S3, S3-compatible services, or Google Cloud Storage via the Data Sources UI. - **Prompt optimization workflows**: run optimization loops, compare candidates, accept winners, and push selected prompts to providers. +- **Usage tracking**: org-wide LLM, STT, and TTS consumption with cost estimates and drill-down by workspace and product area. ## Platform model @@ -48,6 +49,7 @@ This structure keeps testing reproducible while still reflecting real-world voic ## Key guides +- [Usage](/docs/monitoring/usage/) — LLM/STT/TTS consumption, cost estimates, and drill-down by workspace and product area - [Workspaces](/docs/getting-started/workspaces/) — project isolation within your organization, plus workspace roles (Viewer / Editor / Workspace Admin) and how they interact with org roles - [Cloud Storage](/docs/getting-started/cloud-storage/) — S3 and GCS configuration - [Metrics](/docs/products/metrics/) — custom rubrics, surfaces, and evaluation scope diff --git a/docs-fumadocs/content/docs/monitoring/meta.json b/docs-fumadocs/content/docs/monitoring/meta.json index 2ace091c..e0a482ad 100644 --- a/docs-fumadocs/content/docs/monitoring/meta.json +++ b/docs-fumadocs/content/docs/monitoring/meta.json @@ -1,6 +1,7 @@ { "title": "Monitoring", "pages": [ + "usage", "calls", "alerting", "cron-jobs" diff --git a/docs-fumadocs/content/docs/monitoring/usage.mdx b/docs-fumadocs/content/docs/monitoring/usage.mdx new file mode 100644 index 00000000..1c2cb701 --- /dev/null +++ b/docs-fumadocs/content/docs/monitoring/usage.mdx @@ -0,0 +1,198 @@ +--- +id: usage +title: Usage +sidebar_position: 1 +--- + +# Usage + +**Usage** is org-scoped analytics for LLM, STT, and TTS consumption with estimated costs. It helps you see how much each workspace, product area, call import, and model is consuming — and what that usage likely costs. + +Usage is **not** a quota or billing portal. Optional Flexprice event metering is a separate system and is not shown in the Usage UI. + +Models and providers tracked here come from your enabled [integrations](/docs/getting-started/integrations/). Usage is attributed per [workspace](/docs/getting-started/workspaces/) when the underlying workflow is workspace-scoped. + +--- + +## Opening the Usage page + +In the sidebar, go to **Usage → Overview** (`/usage`). + +The page has two tabs when your org is licensed for enterprise features and you are an org admin: + +| Tab | URL | Who can access | +|-----|-----|----------------| +| **Overview** | `/usage` | All org members | +| **Pricing overrides** | `/usage?tab=pricing` | Org admins with enterprise license | + +--- + +## Summary metrics + +The top of the Overview tab shows rollup cards for the selected date range and filters: + +| Metric | Description | +|--------|-------------| +| **Input tokens** | Prompt or input tokens sent to LLM providers | +| **Output tokens** | Completion or output tokens returned by LLMs | +| **Total tokens** | Sum of input and output tokens | +| **LLM calls** | Number of LLM API calls | +| **STT audio** | Speech-to-text audio duration (shown when non-zero) | +| **TTS characters** | Text-to-speech characters synthesized (shown when non-zero) | +| **Cache read** | Tokens read from provider prompt cache (shown when non-zero) | +| **Cache write** | Tokens written to provider prompt cache (shown when non-zero) | +| **Reasoning** | Reasoning tokens billed separately by some providers (shown when non-zero) | +| **Estimated cost** | Total estimated cost for the filtered range | + +Click **Cost breakdown** to open a modal with line items: + +- Input, output, cache read, cache write, reasoning +- Audio (STT), TTS +- Total estimated cost + +If any usage in the range has no matching catalog rate, the breakdown shows an **unpriced usage** warning. Those rows still appear in token/volume metrics but do not contribute to cost totals. + +### Currency display + +Toggle between **USD** and **INR** in the filter bar. INR amounts use a live USD→INR rate from Frankfurter when available, with a fallback estimate when the FX service is unreachable. + +--- + +## Filters + +Filters narrow the summary cards and drill-down table. All filter state is stored in the URL, so you can bookmark or share a specific view. + +| Filter | Description | +|--------|-------------| +| **Date range** | Start and end dates, interpreted in your browser's IANA timezone | +| **Workspace** | Limit to one workspace | +| **Call import** | Limit to one call import batch | +| **Dataset** | Filter by dataset name on call import rows | +| **Tag** | Filter by call import tag | +| **Evaluation run** | Filter by evaluation resource | +| **Usage kind** | `LLM`, `STT`, or `TTS` | +| **Model** | Provider model identifier | +| **Source / product section** | Product area (playground, evaluators, call imports, etc.) | + +--- + +## Drill-down navigation + +Click rows in the breakdown table to drill deeper. Breadcrumbs at the top show your current path; click a breadcrumb to go back up. + +```mermaid +flowchart TD + Org[Organization] --> Workspace[Workspace] + Workspace --> Composite[Call imports and product areas] + Composite --> CallImport[Call import batch] + CallImport --> EvalRun[Evaluation run] + EvalRun --> Model[Model] + Model --> Kind[Usage kind] + Composite --> ProductSection[Product section] + ProductSection --> Model +``` + +At the organization level, the table groups by **workspace**. Inside a workspace you see a composite view: + +- **Call import batches** — CSV uploads or manual audio recordings +- **Product areas** — usage from other parts of the platform (not tied to a single call import) + +From a call import batch you can drill into evaluation runs, then model, then usage kind. From a product area you drill into model, then usage kind. + +Each drill level returns at most **100 rows**. If more exist, results are truncated at that level. + +### Product sections + +| Section | What it tracks | +|---------|----------------| +| **Call imports** | Call import batch processing | +| **Call import evaluations** | Evaluations run on imported calls | +| **Playground** | Text playground and experiments | +| **Voice playground** | Voice agent playground — LLM, STT, and TTS | +| **Chat** | Chat conversations | +| **Telephony** | Telephony and live calls | +| **Evaluators** | Evaluator definitions and runs | +| **Metrics** | Metrics and scoring | +| **Judge alignment** | Judge alignment workflows | +| **Prompt optimization** | Prompt optimization jobs | +| **Personas** | Persona generation | +| **Agents** | Agent configuration | +| **Prompt partials** | Prompt partials | +| **Conversation evaluations** | Conversation evaluations | +| **Test agent** | Test agent sessions | +| **Other** | Usage not attributed to a named product area | + +--- + +## Data freshness and history + +### Freshness + +Usage counters are buffered in Redis and flushed to Postgres by the `worker-usage` service on a Celery Beat schedule (default: every **2 minutes**). The UI reads Postgres only. + +The page shows **Updated** with a `last_updated_at` timestamp when available. Expect roughly **2 minutes** of lag between new API usage and what appears on this page. + +For fresh data in self-hosted deployments, ensure `beat`, `worker-usage`, and the default `worker` are running (or use `eai start-all`). + +### History limits + +| Deployment | History window | Pricing overrides tab | +|------------|----------------|----------------------| +| **OSS** (no license) | Last **7 days** | Hidden | +| **Enterprise** (`EFFICIENTAI_LICENSE`) | Unlimited | Org admins only | + +On OSS deployments, an amber banner explains the 7-day cap and points to `EFFICIENTAI_LICENSE`. If you pick a wider date range, it is automatically clamped to the allowed window. + +See [Configuration](/docs/reference/configuration/) for license setup. + +--- + +## Pricing overrides + +Enterprise org admins can open **Pricing overrides** (`/usage?tab=pricing`) to set per-model rates that override the built-in catalog for cost estimation. + +### Override fields by usage kind + +| Usage kind | Rate fields (USD) | +|------------|-------------------| +| **LLM** | Input / 1M tokens, output / 1M tokens, cache read / 1M, cache write / 1M, reasoning / 1M, audio / minute | +| **STT** | Audio / minute | +| **TTS** | Characters / 1M | + +For each override you set: + +- **Provider credential** — models come from enabled integrations +- **Usage kind** — LLM, STT, or TTS +- **Model** — provider model identifier +- **Effective from** — date the override starts applying +- **Rates** — USD values for the fields above + +You can prefill rates from the catalog or an existing override. Saving creates or updates the override; deleting removes it. + +### How overrides affect costs + +Overrides apply to **new usage** recorded on or after `effective_from`. Costs already stamped on daily rollup rows are **not** changed automatically. + +To backfill historical costs after a catalog or override change, use the CLI or API recompute workflow. The Usage UI does not expose recompute jobs today — see [CLI Commands](/docs/reference/cli-commands/#usage-pricing). + +--- + +## Operations + +Self-hosted operators manage pricing catalogs and cost backfills outside the UI. + +| Requirement | Purpose | +|-------------|---------| +| `beat` | Schedules usage flush, FX refresh, OSS history prune | +| `worker-usage` | Flushes Redis counters and runs cost recompute jobs | +| Default `worker` | Evaluator cron dispatch (indirectly drives much platform usage) | + +Common tuning variables (see `env.example`): + +| Variable | Default | Purpose | +|----------|---------|---------| +| `USAGE_FLUSH_BEAT_SECONDS` | `120` | Flush interval (~2 min UI lag) | +| `USAGE_READ_CACHE_TTL_SECONDS` | `90` | Redis cache TTL for summary/breakdown/filters | +| `USAGE_FLUSH_MAX_BATCHES_PER_RUN` | `30` | Batches per flush tick | + +For seeding rates, diffing catalogs, and recomputing stored costs, see [CLI Commands — Usage pricing](/docs/reference/cli-commands/#usage-pricing). diff --git a/docs-fumadocs/content/docs/reference/cli-commands.mdx b/docs-fumadocs/content/docs/reference/cli-commands.mdx index 6e09c83a..c0ce1658 100644 --- a/docs-fumadocs/content/docs/reference/cli-commands.mdx +++ b/docs-fumadocs/content/docs/reference/cli-commands.mdx @@ -107,3 +107,50 @@ eai migrate --verbose ``` **Note**: Migrations run automatically on application startup. You only need to run them manually if you want to apply migrations before starting the server. + +## Usage pricing + +Manage model pricing rates and backfill stored usage costs on `llm_usage_daily` rollups. Requires `beat`, `worker-usage`, and the default `worker` (or `eai start-all`). + +```bash +# Upsert model_pricing_rates from app/config/models.json +eai usage seed-rates --config config.yml + +# Compare models.json pricing vs Postgres +eai usage diff-rates --config config.yml + +# Backfill costs in-process (all orgs; use after migrate or catalog change) +eai usage recompute --config config.yml --sync + +# Async recompute via usage queue (requires --organization-id) +eai usage recompute --config config.yml --organization-id + +# Optional scopes: --model, --usage-kind, --start-date, --end-date + +# Optional: fetch LiteLLM prices into pricing_catalog.json +eai usage sync-litellm --local +eai usage sync-litellm --local --write-models +``` + +**After migrations or catalog changes:** + +```bash +eai migrate +eai usage seed-rates --config config.yml +eai usage recompute --config config.yml --sync +``` + +**Flush / Usage UI tuning** — set in `.env` (see `env.example`): + +| Variable | Default | Purpose | +|----------|---------|---------| +| `USAGE_FLUSH_BUCKET_BATCH_SIZE` | `500` | Buckets per DB transaction | +| `USAGE_FLUSH_MAX_BATCHES_PER_RUN` | `30` | Batches per flush tick (≤ **15,000** buckets/run) | +| `USAGE_FLUSH_BEAT_SECONDS` | `120` | Celery Beat flush interval (~2 min lag vs Redis) | +| `USAGE_FLUSH_LOCK_TTL_SECONDS` | `300` | Per-org flush lock TTL | +| `USAGE_READ_CACHE_TTL_SECONDS` | `90` | Redis cache TTL for usage summary/breakdown/filters | +| `CRON_DISPATCH_INTERVAL_SECONDS` | `30` | Evaluator cron dispatcher tick (default worker) | + +The Usage UI reads Postgres only (summary/breakdown/filters). Redis counters flush on the Celery Beat schedule (~2 min eventual consistency). If Redis backlog grows, lower `USAGE_FLUSH_BEAT_SECONDS` or raise `USAGE_FLUSH_MAX_BATCHES_PER_RUN`. + +See [Usage](/docs/monitoring/usage/) for the end-user guide. diff --git a/docs-fumadocs/content/feature-contributors.json b/docs-fumadocs/content/feature-contributors.json index 28ca6227..e7990cc8 100644 --- a/docs-fumadocs/content/feature-contributors.json +++ b/docs-fumadocs/content/feature-contributors.json @@ -287,6 +287,23 @@ ], "lastReviewed": "2026-06-15" }, + { + "featureId": "monitoring/usage", + "docPath": "docs-fumadocs/content/docs/monitoring/usage.mdx", + "historyPath": "docs-fumadocs/content/docs/monitoring/usage.mdx", + "owners": [ + "aadhar-EAI", + "Tejas Narayan" + ], + "contributors": [ + { + "name": "aadhar-EAI", + "email": "aadhar@efficientai.cloud", + "commits": 1 + } + ], + "lastReviewed": "2026-08-18" + }, { "featureId": "monitoring/cron-jobs", "docPath": "docs-fumadocs/content/docs/monitoring/cron-jobs.mdx", diff --git a/docs-fumadocs/public/search-index.json b/docs-fumadocs/public/search-index.json index e2d517c7..2d7daadc 100644 --- a/docs-fumadocs/public/search-index.json +++ b/docs-fumadocs/public/search-index.json @@ -1,19 +1,28 @@ { - "generatedAt": "2026-06-13T19:35:24.170Z", - "count": 30, + "generatedAt": "2026-08-18T13:39:50.745Z", + "count": 35, "records": [ { "id": "advanced/architecture", - "url": "/docs/advanced/architecture", + "url": "/docs/advanced/architecture/", "title": "Architecture", "breadcrumbs": [ "Advanced" ], "content": "Architecture Simple Overview EfficientAI is built like a modern web application. The Brain (API Server) : Controls everything. The Worker : Does the heavy lifting in the background, like processing audio files so the website stays fast. The Interface (Frontend) : The website you see and click on. The Memory (Database) : Where we store all your agents, test results, and user data. Technical Deep Dive EfficientAI is built as a modular, containerized application designed for scalability and extensibility. (See original Architecture documentation below) System Components The platform consists of four primary components: 1. API Server (FastAPI) : The central control plane. 2. Worker (Celery) : Handles asynchronous background tasks (transcription, evaluation). 3. Frontend (React/Vite) : The user interface. 4. Data Stores : PostgreSQL (State) and Redis (Queue/Cache). Core Services 1. API Server ( ) Built with FastAPI, it provides REST endpoints for: Resource Management : CRUD for Agents, Personas, Scenarios. Orchestration : Real time control of test conversations. Analysis : Serving evaluation results and dashboards. 2. Asynchronous Workers ( ) Powered by Celery and Redis, the workers handle long running operations: Transcription : Processing audio files (using Whisper, Deepgram, etc.). Evaluation : Running metric calculations (WER/CER) on completed conversations. 3. Test Agent Service ( ) This is the heart of the testing engine. It: Manages the state of the conversation. Generates accurate system prompts for the Persona. Handles the latency sensitive loop of: ." }, + { + "id": "advanced/call-import-sharding", + "url": "/docs/advanced/call-import-sharding/", + "title": "Call Import Sharding", + "breadcrumbs": [ + "Advanced" + ], + "content": "Call Import Sharding For large batches (10k+ rows), call import row data can be spread across multiple PostgreSQL data shards with a catalog database for metadata, routing, and parent counters. Architecture Catalog DB — , , shard slice registry, dispatch metadata Data shards — , (partitioned by consistent hash on ) Scatter/gather reads — API and workers query each shard and merge results Fair dispatch — import/eval workers respect per shard pending scans and Redis progress keys Enable sharding in under . See and for profiles. Live telephony / evaluator results use a separate payload sharding path ( , ) keyed by workspace — see in the repo. Operations Connection pools When , use smaller per process pools ( 3–5, 5–10) so API + workers × shards stay under Postgres . PgBouncer (optional) Point , , and each shard at PgBouncer ( ) in transaction pooling mode. Keep SQLAlchemy enabled. Observability Row Celery tasks log on the import worker hot path. Watch Redis eval/import progress keys ( , ) alongside catalog parent counters. Rebalance Dry run registry updates: Use only after pausing the import. The rebalance tool copies rows to the target shard, updates the catalog registry, then removes copies from the source shard. Troubleshooting stalled evaluations If evaluation runs stall after recordings import (rows show import but no diarization/scoring): 1. Restart API and after deploying fixes. 2. Retry evaluation from the UI (use Overwrite existing transcripts if diarization previously failed). 3. Ensure the imports worker consumes ." + }, { "id": "advanced/database", - "url": "/docs/advanced/database", + "url": "/docs/advanced/database/", "title": "Database", "breadcrumbs": [ "Advanced" @@ -22,7 +31,7 @@ }, { "id": "advanced/development", - "url": "/docs/advanced/development", + "url": "/docs/advanced/development/", "title": "Development & Troubleshooting", "breadcrumbs": [ "Advanced" @@ -31,25 +40,25 @@ }, { "id": "getting-started/authentication", - "url": "/docs/getting-started/authentication", + "url": "/docs/getting-started/authentication/", "title": "Authentication", "breadcrumbs": [ "Getting Started" ], - "content": "🔐 Authentication EfficientAI ships with a pluggable authentication system that scales from a single operator OSS install to an enterprise deployment behind your existing identity provider. You pick the providers you want in (or via in ) and the API/frontend adapt automatically. Deployment models | Model | Providers | License needed | | | | | | OSS self hosted (default) | , | None | | Enterprise SSO (BYO IdP) | , | | — the header, always available, for programmatic access (CI pipelines, SDKs, scripts). — email + password, verified against the local users table, returns an app signed HS256 Bearer token. Enabled by default. — license gated. Verifies a Bearer JWT issued by your OIDC compliant IdP (Okta, Azure AD / Entra ID, Google Workspace, AWS Cognito, Auth0, Ping, JumpCloud, OneLogin, …) against the issuer's JWKS. :::info Why no bundled IdP? In practice every enterprise already runs one. Shipping our own Keycloak alongside the app just added another thing for you to operate and lock down. talks to whatever you already have. ::: Self hosted (OSS) This is the default after or . No license, no IdP, no external dependencies. The equivalent environment variables (for Docker Compose / ): First time bootstrap 1. Start the stack. 2. Open and click Create account on the login screen. The first user you create becomes the admin of a fresh organization. 3. Mint an API key from Profile → API Keys (or via ) for programmatic access. Password login (email + password) When is enabled, the login screen shows a Sign in and (if ) a Create account tab. Signing up provisions a new user and a new organization, and makes that user the of it. If you leave the organization name blank, the server derives one from the email's local part. Once signed in, the SPA holds a short lived Bearer token and silently re authenticates before it expires. You can change the token lifetime with — the default is 12 hours. Linking a password to an API key only account If you bootstrapped with , the backend provisions a placeholder user behind that key (its email ends in ). You can upgrade this identity to a real email + password login so you can sign in interactively with the same user. Do it from Profile → Sign in Password while signed in via the API key — the page detects the placeholder email and prompts you to pick a real one and a password. After saving, the same user can log in either with the original API key (for machines) or with email + password (for humans). Rules the UI enforces: If the user already has a password, the form asks for the current one before accepting a new one. You can only set the email from that screen while it's still the placeholder address; \"real\" users change their email from the main profile edit flow. Hardening before you expose it to the internet Turn off self service signup once your team is in: Rotate to invalidate existing sessions. Put the app behind a reverse proxy (Nginx, Caddy, Cloudflare) that terminates TLS and enforces HSTS. Restrict to the exact domain(s) serving the SPA. Team management: invitations & organizations EfficientAI is multi tenant from the ground up. Every piece of data is scoped to an organization , a user can be a member of more than one organization, and each membership has a role that controls what they can do. Roles | Role | Can do | | | | | | Read only access to everything in the org. | | | Everything a reader can + create/update/delete most resources. | | | Everything a writer can + manage users, invitations, roles, API keys, and org settings. | The role is stored per membership, so the same user can be an in one org and a in another. Inviting a teammate Admins invite teammates from Settings → Team . An invitation captures an email and a role and stays valid for 7 days. From the same page, admins can also: See the current members of the org and change their role (with a guard so you can't demote the last admin). Remove a member from the organization. Revoke a pending invitation. Delivery. The backend cre" + "content": "🔐 Authentication EfficientAI ships with a pluggable authentication system that scales from a single operator OSS install to an enterprise deployment behind your existing identity provider. You pick the providers you want in (or via in ) and the API/frontend adapt automatically. Deployment models | Model | Providers | License needed | | | | | | OSS self hosted (default) | , | None | | Enterprise SSO (BYO IdP) | , | | — the header, always available, for programmatic access (CI pipelines, SDKs, scripts). — email + password, verified against the local users table, returns an app signed HS256 Bearer token. Enabled by default. — license gated. Verifies a Bearer JWT issued by your OIDC compliant IdP (Okta, Azure AD / Entra ID, Google Workspace, AWS Cognito, Auth0, Ping, JumpCloud, OneLogin, …) against the issuer's JWKS. :::info Why no bundled IdP? In practice every enterprise already runs one. Shipping our own Keycloak alongside the app just added another thing for you to operate and lock down. talks to whatever you already have. ::: Self hosted (OSS) This is the default after or . No license, no IdP, no external dependencies. The equivalent environment variables (for Docker Compose / ): First time bootstrap 1. Start the stack. 2. Open and click Create account on the login screen. The first user you create becomes the admin of a fresh organization. 3. Mint an API key from Profile → API Keys (or via ) for programmatic access. Password login (email + password) When is enabled, the login screen shows a Sign in and (if ) a Create account tab. Signing up provisions a new user and a new organization, and makes that user the of it. If you leave the organization name blank, the server derives one from the email's local part. Once signed in, the SPA holds a short lived access token (15 minutes by default) plus a refresh token. The client silently refreshes the access token before it expires. You can change lifetimes with and . Logout revokes the refresh token and blacklists the current access token server side. Linking a password to an API key only account If you bootstrapped with , the backend provisions a placeholder user behind that key (its email ends in ). You can upgrade this identity to a real email + password login so you can sign in interactively with the same user. Do it from Profile → Sign in Password while signed in via the API key — the page detects the placeholder email and prompts you to pick a real one and a password. After saving, the same user can log in either with the original API key (for machines) or with email + password (for humans). Rules the UI enforces: If the user already has a password, the form asks for the current one before accepting a new one. You can only set the email from that screen while it's still the placeholder address; \"real\" users change their email from the main profile edit flow. Hardening before you expose it to the internet Turn off self service signup once your team is in: Rotate to invalidate existing sessions. Put the app behind a reverse proxy (Nginx, Caddy, Cloudflare) that terminates TLS and enforces HSTS. The bundled FastAPI server sends baseline security headers on all responses ( , , , , and by default). If you terminate traffic at an external reverse proxy, keep those headers (or stricter CSP ) enabled there too. After reviewing CSP violation reports, set to enforce the policy. Pin third party observability container images to fixed tags and rebuild external reverse proxy images on patched runtimes. If a scanner reports a Go stdlib CVE in a binary this repo does not build, identify the flagged container or proxy artifact and upgrade it separately. Restrict to the exact domain(s) serving the SPA. Set in production so , , and are not served on the public hostname. Keep so and return 404 to anonymous public clients (including vulnerability scanners hitting your ALB hostname). AWS ALB target health checks connect directly from VPC addresses (no ); include your VPC/LB CIDRs in . Full migrat" }, { "id": "getting-started/cloud-storage", - "url": "/docs/getting-started/cloud-storage", + "url": "/docs/getting-started/cloud-storage/", "title": "Cloud Storage", "breadcrumbs": [ "Getting Started" ], - "content": "Cloud Storage Overview EfficientAI can store audio files and recordings in the cloud using one active blob backend at a time: Amazon S3 (or any S3 compatible service such as MinIO, DigitalOcean Spaces, Cloudflare R2) Google Cloud Storage (GCS) Cloud storage is useful for: Storing large audio files outside your application server Scaling storage independently from compute Integrating with existing AWS or GCP infrastructure Browsing, uploading, and managing audio from the Data Sources page in the UI Select the backend with in ( or ), then enable and configure the matching block below. Storage block | Option | Description | | | | | | Local fallback directory for uploads when cloud storage is disabled | | | Maximum upload size in megabytes | | | Active cloud backend: or | | | File extensions accepted for audio uploads | Amazon S3 and S3 compatible storage Configuration Set and enable the block: | Option | Required | Description | | | | | | | Yes | Set to to enable S3 storage | | | Yes | Name of your S3 bucket | | | Yes | AWS region (e.g., , ) | | | Yes | AWS Access Key ID | | | Yes | AWS Secret Access Key | | | No | Custom endpoint for S3 compatible services | | | No | Folder prefix for uploaded files (default: ) | S3 compatible examples MinIO DigitalOcean Spaces Cloudflare R2 Google Cloud Storage Configuration Set and enable the block: | Option | Required | Description | | | | | | | Yes | Set to to enable GCS storage | | | Yes | Name of your GCS bucket | | | Yes | GCP project ID | | | No | Path to a service account JSON key file | | | No | Folder prefix for uploaded files (default: ) | Authentication EfficientAI resolves GCS credentials in this order: 1. in (relative paths resolve from the working directory) 2. environment variable 3. Application Default Credentials (ADC) on GCE, GKE, or Cloud Run The client is included in the standard EfficientAI install ( ); no extra Python package step is required for GCS. GCP setup 1. Create a GCS bucket in your project (Console or ). 2. Create a service account with object read/write access (e.g., on the bucket, or a tighter custom role). 3. Download a JSON key and set , or mount the key and set . 4. Set , , and fill in and . Object layout under matches S3 (organization scoped paths), so you can migrate between providers without changing application logic. Data Sources UI Once cloud storage is configured and connected, open Configuration → Data Sources to: Browse folders and audio files in your bucket Upload new audio files Preview playback in the browser Delete files Test the connection (labels show Amazon S3 or Google Cloud Storage based on ) Verifying connection Restart the application after changing storage settings: Use Test Connection on the Data Sources page, or upload a test file and confirm it appears under the configured in your bucket. Troubleshooting | Issue | Solution | | | | | (S3) | Check IAM permissions and bucket policy | | (S3) | Verify bucket name and region | | (S3) | Check for S3 compatible services | | (S3) | Verify and | | GCS | Confirm service account has on the bucket | | GCS bucket not found | Verify , , and that the bucket exists | | GCS auth failure | Set or ; on GCP VMs you can use ADC | | Wrong provider in UI | Ensure matches the enabled block ( vs ) |" + "content": "Cloud Storage Overview EfficientAI can store audio files and recordings in the cloud using one active blob backend at a time: Amazon S3 (or any S3 compatible service such as MinIO, DigitalOcean Spaces, Cloudflare R2) Google Cloud Storage (GCS) Azure Blob Storage Cloud storage is useful for: Storing large audio files outside your application server Scaling storage independently from compute Integrating with existing AWS, GCP, or Azure infrastructure Browsing, uploading, and managing audio from the Data Sources page in the UI Select the backend with in ( , , or ), then enable and configure the matching block below. Storage block | Option | Description | | | | | | Local fallback directory for uploads when cloud storage is disabled | | | Maximum upload size in megabytes | | | Active cloud backend: , , or | | | File extensions accepted for audio uploads | Amazon S3 and S3 compatible storage Configuration Set and enable the block: | Option | Required | Description | | | | | | | Yes | Set to to enable S3 storage | | | Yes | Name of your S3 bucket | | | Yes | AWS region (e.g., , ) | | | Yes | AWS Access Key ID | | | Yes | AWS Secret Access Key | | | No | Custom endpoint for S3 compatible services | | | No | Folder prefix for uploaded files (default: ) | S3 compatible examples MinIO DigitalOcean Spaces Cloudflare R2 Google Cloud Storage Configuration Set and enable the block: | Option | Required | Description | | | | | | | Yes | Set to to enable GCS storage | | | Yes | Name of your GCS bucket | | | Yes | GCP project ID | | | No | Path to a service account JSON key file | | | No | Override service account email for signed URL generation when ADC does not expose it | | | No | Folder prefix for uploaded files (default: ) | Authentication EfficientAI resolves GCS credentials in this order: 1. in (relative paths resolve from the working directory) 2. environment variable 3. Application Default Credentials (ADC) on GCE, GKE, or Cloud Run Uploads and server side downloads work with ADC alone (including GKE Workload Identity ). Signed URLs for browser playback require either a service account JSON key with a private key, or Workload Identity plus IAM signBlob (see below). The client is included in the standard EfficientAI install ( ); no extra Python package step is required for GCS. GCP setup 1. Create a GCS bucket in your project (Console or ). 2. Create a service account with object read/write access (e.g., on the bucket, or a tighter custom role). 3. Authenticate the app using one of: GKE Workload Identity (recommended): bind your Kubernetes service account to the GCP service account. Enable the and grant the GCP service account on itself so EfficientAI can sign playback URLs without a JSON key. Service account JSON key: download a key and set , or mount the key and set . 4. Set , , and fill in and . Example IAM binding for Workload Identity signed URLs: If ADC does not expose the service account email (some federation setups), set to the workload GCP service account email. Object layout under matches S3 (organization scoped paths), so you can migrate between providers without changing application logic. Azure Blob Storage Configuration Set and enable the block: | Option | Required | Description | | | | | | | Yes | Set to to enable Azure Blob Storage | | | Yes | Name of your Azure storage container | | | Yes | Azure storage account name ( not required if is set) | | | Yes | Storage account access key ( not required if is set) | | | No | Full Azure storage connection string (overrides + ) | | | No | Folder prefix for uploaded files (default: ) | Authentication EfficientAI resolves Azure credentials in this order: 1. in 2. + Managed Identity is not supported in this release; use a connection string or account key. SAS URLs (temporary download links in the UI) require an account key. Connection strings include the key automatically. Azure setup 1. Create a Storage Account in the Azure Portal (or via ). 2. Create a container (e.g., ) under tha" }, { "id": "getting-started/installation", - "url": "/docs/getting-started/installation", + "url": "/docs/getting-started/installation/", "title": "Installation", "breadcrumbs": [ "Getting Started" @@ -58,41 +67,50 @@ }, { "id": "getting-started/integrations", - "url": "/docs/getting-started/integrations", + "url": "/docs/getting-started/integrations/", "title": "Integrations", "breadcrumbs": [ "Getting Started" ], - "content": "Integrations Integrations in EfficientAI are not limited to external voice agent platforms. You can connect integrations across three layers of the stack: 1. Voice agent platforms (agent/runtime side) 2. AI providers (LLM/STT/TTS model side) 3. Telephony providers (PSTN/number/routing side) This lets you test and evaluate the complete call path from model behavior to phone network delivery. Integrations UI 1) Voice platform integrations (agent side) Voice platform integrations connect EfficientAI to externally hosted voice agents.