litellm/litellm/integrations/focus/database.py

"""Database access helpers for Focus export."""

from __future__ import annotations

from datetime import datetime
from typing import Final

import polars as pl


class FocusLiteLLMDatabase:
    """Retrieves LiteLLM usage data for Focus export workflows."""

    def _ensure_prisma_client(self):
        from litellm.proxy.proxy_server import prisma_client

        if prisma_client is None:
            raise RuntimeError(
                "Database not connected. Connect a database to your proxy - "
                "https://docs.litellm.ai/docs/simple_proxy#managing-auth---virtual-keys"
            )
        return prisma_client

    async def get_usage_data(
        self,
        *,
        limit: int | None = None,
        start_time_utc: datetime | None = None,
        end_time_utc: datetime | None = None,
    ) -> pl.DataFrame:
        """Return usage data for the requested window."""
        client: Final = self._ensure_prisma_client()

        where_clauses: Final[list[str]] = []
        query_params: Final[list[datetime | int]] = []
        placeholder_index = 1
        if start_time_utc:
            where_clauses.append(f"dus.updated_at >= ${placeholder_index}::timestamptz")
            query_params.append(start_time_utc)
            placeholder_index += 1
        if end_time_utc:
            where_clauses.append(f"dus.updated_at <= ${placeholder_index}::timestamptz")
            query_params.append(end_time_utc)
            placeholder_index += 1

        where_clause = ""
        if where_clauses:
            where_clause = "WHERE " + " AND ".join(where_clauses)

        limit_clause = ""
        if limit is not None:
            try:
                limit_value: Final = int(limit)
            except (TypeError, ValueError) as exc:  # pragma: no cover - defensive guard
                raise ValueError("limit must be an integer") from exc
            if limit_value < 0:
                raise ValueError("limit must be non-negative")
            limit_clause = f" LIMIT ${placeholder_index}"
            query_params.append(limit_value)

        query: Final = f"""
        SELECT
            dus.id,
            dus.date,
            dus.user_id,
            dus.api_key,
            dus.model,
            dus.model_group,
            dus.custom_llm_provider,
            dus.prompt_tokens,
            dus.completion_tokens,
            dus.spend,
            dus.api_requests,
            dus.successful_requests,
            dus.failed_requests,
            dus.cache_creation_input_tokens,
            dus.cache_read_input_tokens,
            dus.created_at,
            dus.updated_at,
            vt.team_id,
            vt.key_alias as api_key_alias,
            tt.team_alias,
            ut.user_email as user_email,
            COALESCE(vt.organization_id, tt.organization_id) as organization_id,
            ot.organization_alias as organization_alias
        FROM "LiteLLM_DailyUserSpend" dus
        LEFT JOIN "LiteLLM_VerificationToken" vt ON dus.api_key = vt.token
        LEFT JOIN "LiteLLM_TeamTable" tt ON vt.team_id = tt.team_id
        LEFT JOIN "LiteLLM_UserTable" ut ON dus.user_id = ut.user_id
        LEFT JOIN "LiteLLM_OrganizationTable" ot
            ON ot.organization_id = COALESCE(vt.organization_id, tt.organization_id)
        {where_clause}
        ORDER BY dus.date DESC, dus.created_at DESC
        {limit_clause}
        """

        try:
            db_response: Final = await client.db.query_raw(query, *query_params)
            from litellm.proxy.spend_tracking.key_metadata_recovery import (
                fill_missing_api_key_aliases,
            )

            usage_rows: Final = (
                db_response.to_dicts()
                if isinstance(db_response, pl.DataFrame)
                else db_response
                if isinstance(db_response, list)
                else []
            )
            recovered_rows: Final = await fill_missing_api_key_aliases(client, usage_rows)
            return pl.DataFrame([dict(row) for row in recovered_rows], infer_schema_length=None)
        except Exception as exc:
            raise RuntimeError(f"Error retrieving usage data: {exc}") from exc

    async def get_table_info(self) -> dict[str, object]:
        """Return metadata about the spend table for diagnostics."""
        client: Final = self._ensure_prisma_client()

        info_query: Final = """
        SELECT column_name, data_type, is_nullable
        FROM information_schema.columns
        WHERE table_name = 'LiteLLM_DailyUserSpend'
        ORDER BY ordinal_position;
        """
        try:
            columns_response: Final = await client.db.query_raw(info_query)
            return {"columns": columns_response, "table_name": "LiteLLM_DailyUserSpend"}
        except Exception as exc:
            raise RuntimeError(f"Error getting table info: {exc}") from exc