Expanded LLM/STT/TTS usage tracking across pre-prod flows and some fixes

.github/workflows/backend-tests-postgres.yml

@@ -41,6 +41,8 @@ jobs:
       REDIS_URL: redis://localhost:6379/0
       CELERY_BROKER_URL: redis://localhost:6379/0
       CELERY_RESULT_BACKEND: redis://localhost:6379/0
+      EFFICIENTAI_PYTEST: "1"
+      FLEXPRICE_ENABLED: "false"
 
     steps:
       - name: Checkout repository
@@ -83,6 +85,8 @@ jobs:
       SHARD_DATABASE_URL_02: postgresql://postgres:postgres@localhost:5432/efficientai_data_02
       SHARDING_INTEGRATION_TEST: "1"
       REDIS_URL: redis://localhost:6379/0
+      EFFICIENTAI_PYTEST: "1"
+      FLEXPRICE_ENABLED: "false"
 
     steps:
       - name: Checkout repository

app/api/v1/routes/agents.py

@@ -2,16 +2,19 @@
 Agents API Routes
 Complete CRUD operations for test agents
 """
-from fastapi import APIRouter, Depends, HTTPException, status, Query
+from uuid import uuid4
+
+from fastapi import APIRouter, BackgroundTasks, Depends, HTTPException, status, Query
 from fastapi.responses import JSONResponse, Response
 from sqlalchemy.orm import Session
 from typing import List, Optional
-from uuid import UUID
+from uuid import UUID, uuid4
 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.services.billing.flexprice_service import record_agent_test_setup_generated
 from app.models.database import (
     Agent, ConversationEvaluation, TestAgentConversation, VoiceBundle,
     AIProvider, Integration, IntegrationPlatform, CallMediumEnum,
@@ -254,6 +257,7 @@ def _scenario_draft_responses(scenarios) -> list[GeneratedScenarioDraftResponse]
 @router.post("/generate-test-prompt", response_model=GenerateTestPromptResponse)
 async def generate_test_prompt(
     data: GenerateTestPromptRequest,
+    background_tasks: BackgroundTasks,
     organization_id: UUID = Depends(get_organization_id),
     workspace_id: UUID = Depends(get_workspace_id),
     api_key: str = Depends(get_api_key),
@@ -297,6 +301,14 @@ async def generate_test_prompt(
                 llm_config=data.llm_config,
                 credential_id=data.credential_id,
             )
+        background_tasks.add_task(
+            record_agent_test_setup_generated,
+            organization_id,
+            uuid4(),
+            workspace_id=workspace_id,
+            purpose="test_prompt",
+            model=result.model,
+        )
         return GenerateTestPromptResponse(
             sections=_test_prompt_section_responses(result.sections),
             test_agent_prompt=result.test_agent_prompt,
@@ -313,6 +325,7 @@ async def generate_test_prompt(
 @router.post("/generate-scenarios-from-prompt", response_model=GenerateScenariosFromPromptResponse)
 async def generate_scenarios_from_prompt(
     data: GenerateScenariosFromPromptRequest,
+    background_tasks: BackgroundTasks,
     organization_id: UUID = Depends(get_organization_id),
     workspace_id: UUID = Depends(get_workspace_id),
     api_key: str = Depends(get_api_key),
@@ -357,6 +370,15 @@ async def generate_scenarios_from_prompt(
                 llm_config=data.llm_config,
                 credential_id=data.credential_id,
             )
+        background_tasks.add_task(
+            record_agent_test_setup_generated,
+            organization_id,
+            uuid4(),
+            workspace_id=workspace_id,
+            purpose="scenarios",
+            model=result.model,
+            scenario_count=len(result.scenarios),
+        )
         return GenerateScenariosFromPromptResponse(
             scenarios=_scenario_draft_responses(result.scenarios),
             provider=result.provider,
@@ -372,6 +394,7 @@ async def generate_scenarios_from_prompt(
 @router.post("/generate-test-setup", response_model=GenerateTestSetupResponse)
 async def generate_test_setup(
     data: GenerateTestSetupRequest,
+    background_tasks: BackgroundTasks,
     organization_id: UUID = Depends(get_organization_id),
     workspace_id: UUID = Depends(get_workspace_id),
     api_key: str = Depends(get_api_key),
@@ -430,6 +453,15 @@ async def generate_test_setup(
                 llm_config=data.llm_config,
                 credential_id=data.credential_id,
             )
+        background_tasks.add_task(
+            record_agent_test_setup_generated,
+            organization_id,
+            uuid4(),
+            workspace_id=workspace_id,
+            purpose="full_setup",
+            model=scenario_result.model,
+            scenario_count=len(scenario_result.scenarios),
+        )
         return GenerateTestSetupResponse(
             sections=_test_prompt_section_responses(prompt_result.sections),
             test_agent_prompt=prompt_result.test_agent_prompt,

app/api/v1/routes/auth.py

@@ -287,6 +287,18 @@ def _extract_bearer(authorization: Optional[str]) -> Optional[str]:
     return token.strip()
 
 
+def _revoke_local_password_access_token(bearer: str) -> None:
+    try:
+        claims = decode_access_token(bearer)
+        jti = claims.get("jti")
+        exp = claims.get("exp")
+        if jti and exp:
+            ttl = max(int(exp) - int(datetime.now(timezone.utc).timestamp()), 1)
+            revoke_access_jti(jti, ttl)
+    except JWTError:
+        pass
+
+
 def _issue_session_tokens(
     db: Session,
     *,
@@ -611,15 +623,7 @@ def logout(
     """Revoke the current session's refresh token and blacklist the access token."""
     bearer = _extract_bearer(authorization)
     if bearer and principal.auth_method == AuthMethod.LOCAL_PASSWORD:
-        try:
-            claims = decode_access_token(bearer)
-            jti = claims.get("jti")
-            exp = claims.get("exp")
-            if jti and exp:
-                ttl = max(int(exp) - int(datetime.now(timezone.utc).timestamp()), 1)
-                revoke_access_jti(jti, ttl)
-        except JWTError:
-            pass
+        _revoke_local_password_access_token(bearer)
 
     if payload and payload.refresh_token:
         revoke_refresh_token(db, payload.refresh_token)
@@ -697,12 +701,14 @@ def refresh_session(payload: RefreshRequest, db: Session = Depends(get_db)) -> T
 
 class SwitchOrgRequest(BaseModel):
     organization_id: str
+    refresh_token: Optional[str] = None
 
 
 @router.post("/switch-org", response_model=TokenResponse)
 def switch_organization(
     payload: SwitchOrgRequest,
     principal: Principal = Depends(get_principal),
+    authorization: Optional[str] = Header(None, alias="Authorization"),
     db: Session = Depends(get_db),
 ) -> TokenResponse:
     """
@@ -766,6 +772,13 @@ def switch_organization(
             detail="User is no longer active.",
         )
 
+    if principal.auth_method == AuthMethod.LOCAL_PASSWORD:
+        bearer = _extract_bearer(authorization)
+        if bearer:
+            _revoke_local_password_access_token(bearer)
+        if payload.refresh_token:
+            revoke_refresh_token(db, payload.refresh_token)
+
     user.last_login_at = datetime.now(timezone.utc)
     db.commit()
     db.refresh(user)

app/api/v1/routes/call_import_evaluations.py

@@ -1281,6 +1281,20 @@ def _name_for_source(source: str) -> Optional[str]:
                 transcribe_overwrite=payload.transcribe_overwrite,
             )
 
+    from app.services.billing.flexprice_service import (
+        record_call_import_evaluation_started,
+    )
+
+    for evaluation in created_evaluations:
+        record_call_import_evaluation_started(
+            organization_id,
+            evaluation.id,
+            workspace_id=evaluation.workspace_id,
+            call_import_id=call_import.id,
+            total_rows=int(evaluation.total_rows or 0),
+            metric_count=len(leaf_metric_ids),
+        )
+
     for evaluation in created_evaluations:
         db.refresh(evaluation)
 
@@ -3980,6 +3994,18 @@ async def generate_call_import_evaluation_pdf_report(
             detail="Failed to store PDF report due to a concurrent duplicate request.",
         ) from None
     db.refresh(pdf_report)
+    from app.services.billing.flexprice_service import (
+        record_call_import_pdf_report_generated,
+    )
+
+    record_call_import_pdf_report_generated(
+        organization_id,
+        pdf_report.id,
+        workspace_id=pdf_report.workspace_id,
+        evaluation_id=evaluation.id,
+        call_import_id=call_import_id,
+        report_type=pdf_report.report_type,
+    )
     return _pdf_report_response_from_row(pdf_report)
 
 
@@ -6926,6 +6952,16 @@ def _resolve_alias(alias_map: Dict[str, str], key: str) -> str:
 # in ``app/workers/tasks/helpers/llm_evaluation.py`` — kept local here to
 # avoid a worker import cycle from the routes module.
 DISCOVERED_METRICS_KEY = "__discovered_metrics__"
+METRIC_SCORES_META_KEYS = frozenset({DISCOVERED_METRICS_KEY, "_billing"})
+
+
+def _is_metric_scores_meta_key(key: str) -> bool:
+    normalized = str(key or "").strip().lower()
+    if not normalized:
+        return True
+    if normalized in {item.lower() for item in METRIC_SCORES_META_KEYS}:
+        return True
+    return normalized.endswith("__discovered")
 
 # Allowed values for an LLM-suggested top-level metric type. Kept in
 # sync with ``DiscoveredMetricSuggestedType`` in
@@ -8674,7 +8710,8 @@ def _reset_eval_row_for_retry(
         eval_row.metric_scores = {
             key: value
             for key, value in existing.items()
-            if str(key).lower() not in target_keys
+            if _is_metric_scores_meta_key(str(key))
+            or str(key).lower() not in target_keys
         }
     else:
         eval_row.metric_scores = {}
@@ -9181,12 +9218,15 @@ async def retry_call_import_evaluation(
         payload.metric_ids if payload else None
     )
     if metric_ids is not None:
+        metric_ids = [
+            mid for mid in metric_ids if not _is_metric_scores_meta_key(str(mid))
+        ]
         if not metric_ids:
             raise HTTPException(
                 status_code=400,
                 detail=(
-                    "metric_ids must be a non-empty list. Omit the "
-                    "field to re-run all metrics."
+                    "metric_ids must be a non-empty list of metric UUIDs. "
+                    "Omit the field to re-run all metrics."
                 ),
             )
 

app/api/v1/routes/chat.py

@@ -28,6 +28,7 @@ class ChatRequest(BaseModel):
     temperature: Optional[float] = 0.7
     max_tokens: Optional[int] = None
     llm_config: Optional[Dict[str, Any]] = None
+    usage_purpose: Optional[str] = "scenario_description"
 
 
 class ChatResponse(BaseModel):
@@ -80,6 +81,7 @@ async def chat_completion(
             uuid4(),
             workspace_id=workspace_id,
             model=result.get("model", request.model),
+            purpose=request.usage_purpose or "scenario_description",
         )
 
         return ChatResponse(

app/api/v1/routes/evaluator_results.py

@@ -864,15 +864,17 @@ def re_evaluate_result(
             detail="Cannot re-evaluate: this result has no transcription. It must be transcribed first."
         )
 
-    if not result.evaluator_id:
+    if not result.evaluator_id and not result.agent_id:
         raise HTTPException(
             status_code=400,
-            detail="Cannot re-evaluate: this result is not linked to an evaluator."
+            detail="Cannot re-evaluate: this result is not linked to an agent or evaluator."
         )
 
-    evaluator = db.query(Evaluator).filter(Evaluator.id == result.evaluator_id).first()
-    if not evaluator:
-        raise HTTPException(status_code=404, detail="Linked evaluator no longer exists")
+    evaluator = None
+    if result.evaluator_id:
+        evaluator = db.query(Evaluator).filter(Evaluator.id == result.evaluator_id).first()
+        if not evaluator:
+            raise HTTPException(status_code=404, detail="Linked evaluator no longer exists")
 
     # ------------------------------------------------------------------
     # If no audio in S3 yet, try to download from the voice provider

app/api/v1/routes/metric_studio.py

@@ -25,6 +25,7 @@
     MetricStudioRunRetryRequest,
 )
 from app.services.metric_studio.metric_selection import expand_studio_metric_selection
+from app.services.metric_studio.run_rollup import rollup_metric_studio_run
 from app.services.metric_studio.source_resolver import resolve_source
 
 router = APIRouter(prefix="/metric-studio", tags=["metric-studio"])
@@ -117,26 +118,7 @@ def _serialize_result(
 
 
 def _rollup_run_status(db: Session, run: MetricStudioRun) -> None:
-    results = (
-        db.query(MetricStudioRunResult)
-        .filter(MetricStudioRunResult.run_id == run.id)
-        .all()
-    )
-    completed = sum(1 for r in results if r.status == "completed")
-    failed = sum(1 for r in results if r.status == "failed")
-    pending = sum(1 for r in results if r.status in {"pending", "running"})
-    run.completed_items = completed
-    run.failed_items = failed
-    if pending:
-        run.status = "running"
-    elif failed and completed:
-        run.status = "partial"
-    elif failed:
-        run.status = "failed"
-    else:
-        run.status = "completed"
-        run.finished_at = datetime.now(timezone.utc)
-    db.flush()
+    rollup_metric_studio_run(db, run, emit_flexprice=True, commit=False)
 
 
 @router.post(
@@ -297,6 +279,8 @@ def get_metric_studio_run(
     )
     if not run:
         raise HTTPException(status_code=404, detail="Studio run not found.")
+    rollup_metric_studio_run(db, run, emit_flexprice=True, commit=True)
+    db.refresh(run)
     return _serialize_run(run)
 
 

app/api/v1/routes/observability.py

@@ -5,15 +5,11 @@
 from typing import Any, Dict, List, Optional, Union
 from uuid import UUID
 
-from fastapi import APIRouter, BackgroundTasks, Depends, HTTPException, status
+from fastapi import APIRouter, Depends, HTTPException, status
 from pydantic import BaseModel, ConfigDict
 from sqlalchemy.orm import Session
 
 from app.dependencies import get_api_key, get_db, get_organization_id, get_workspace_id
-from app.services.billing.flexprice_service import (
-    record_observability_call_evaluated,
-    record_observability_call_ingested,
-)
 from app.models.database import (
     Agent, APIKey, CallRecording, CallRecordingStatus, CallRecordingSource,
     Evaluator, EvaluatorResult, EvaluatorResultStatus, Scenario, Workspace,
@@ -243,13 +239,6 @@ def _upsert_call_recording(
 
     response = _serialize_call_recording(call_recording, include_data=True, agent=agent_obj)
     response["action"] = action
-    if action == "created":
-        record_observability_call_ingested(
-            organization_id,
-            call_recording.call_short_id,
-            workspace_id=workspace_id,
-            provider=provider_platform,
-        )
     return response
 
 
@@ -765,7 +754,6 @@ def _messages_to_speaker_segments(messages: List[Dict[str, Any]]) -> List[Dict[s
 async def evaluate_call(
     call_short_id: str,
     payload: EvaluateCallPayload,
-    background_tasks: BackgroundTasks,
     organization_id: UUID = Depends(get_organization_id),
     workspace_id: UUID = Depends(get_workspace_id),
     api_key: str = Depends(get_api_key),
@@ -877,13 +865,6 @@ async def evaluate_call(
     except Exception:
         pass
 
-    background_tasks.add_task(
-        record_observability_call_evaluated,
-        organization_id,
-        call_short_id,
-        workspace_id=workspace_id,
-    )
-
     return {
         "evaluator_result_id": str(evaluator_result.id),
         "result_id": evaluator_result.result_id,