[pull] main from onyx-dot-app:main

.vscode/extensions.json

@@ -0,0 +1,5 @@
+{
+    "recommendations": [
+        "anysphere.remote-containers"
+    ]
+}

backend/ee/onyx/server/log_export/api.py

@@ -26,6 +26,8 @@
 from onyx import __version__
 from onyx.auth.permissions import require_permission
 from onyx.background.celery.versioned_apps.client import app as client_app
+from onyx.cache.interface import CacheBackendType
+from onyx.configs.app_configs import CACHE_BACKEND
 from onyx.configs.constants import (
     OnyxCeleryPriority,
     OnyxCeleryQueues,
@@ -137,9 +139,9 @@ def start_log_export(
     Starts an export: fans out one collector task per worker type, collects the
     api_server's logs inline, and returns the export ID to poll.
 
-    Fan-out failures (e.g. deployments with no broker or workers, like the
-    onyx-lite overlay) degrade the export to just the api_server's logs instead
-    of failing.
+    The fan-out is skipped on deployments that have no celery broker (the
+    onyx-lite overlay); a failing broker degrades the export to the workers
+    already enqueued instead of failing it.
     """
     if MULTI_TENANT:
         raise OnyxError(
@@ -165,33 +167,44 @@ def start_log_export(
 
         # Fan out before the inline collection below so workers get the full
         # window before ``expires=`` discards their tasks, and their collection
-        # overlaps the api_server's.
+        # overlaps the api_server's. When redis is absent by design (the
+        # onyx-lite overlay), there is no broker behind ``send_task`` and no
+        # workers to collect from, so the fan-out is skipped outright
+        # (``maybe_schedule_license_reclaim`` applies the same rule); a broker
+        # that exists but is down degrades per-send below instead.
         enqueued_worker_names: list[str] = []
-        for worker_name, queue in WORKER_COLLECT_QUEUES.items():
-            try:
-                client_app.send_task(
-                    OnyxCeleryTask.EXPORT_LOGS_COLLECT_TASK,
-                    priority=OnyxCeleryPriority.HIGHEST,
-                    queue=queue,
-                    expires=deadline,
-                    kwargs={
-                        "export_id": export_id,
-                        "worker_name": worker_name,
-                    },
-                )
-            except Exception as e:
-                # All sends share one broker, so the first failure means the
-                # rest would fail too. Only the workers already enqueued are
-                # awaited.
-                logger.warning(
-                    "Log export fan-out failed while enqueueing %s; continuing "
-                    "with %s: %s",
-                    worker_name,
-                    enqueued_worker_names,
-                    e,
-                )
-                break
-            enqueued_worker_names.append(worker_name)
+        if CACHE_BACKEND != CacheBackendType.REDIS:
+            logger.info(
+                "Log export fan-out skipped: this deployment has no celery "
+                "broker (CACHE_BACKEND=%s).",
+                CACHE_BACKEND.value,
+            )
+        else:
+            for worker_name, queue in WORKER_COLLECT_QUEUES.items():
+                try:
+                    client_app.send_task(
+                        OnyxCeleryTask.EXPORT_LOGS_COLLECT_TASK,
+                        priority=OnyxCeleryPriority.HIGHEST,
+                        queue=queue,
+                        expires=deadline,
+                        kwargs={
+                            "export_id": export_id,
+                            "worker_name": worker_name,
+                        },
+                    )
+                except Exception as e:
+                    # All sends share one broker, so the first failure means the
+                    # rest would fail too. Only the workers already enqueued are
+                    # awaited.
+                    logger.warning(
+                        "Log export fan-out failed while enqueueing %s; "
+                        "continuing with %s: %s",
+                        worker_name,
+                        enqueued_worker_names,
+                        e,
+                    )
+                    break
+                enqueued_worker_names.append(worker_name)
 
         manifest = LogExportManifest(
             export_id=export_id,

backend/ee/onyx/server/reporting/usage_export_generation.py

@@ -2,6 +2,7 @@
 import tempfile
 import uuid
 import zipfile
+from collections.abc import Iterable
 from datetime import datetime, timedelta, timezone
 
 from fastapi_users_db_sqlalchemy import UUID_ID
@@ -20,33 +21,37 @@
 )
 from onyx.configs.constants import FileOrigin
 from onyx.db.models import User
+from onyx.db.user_usage import UsageExportRow, iter_usage_export
 from onyx.db.users import get_all_users
 from onyx.file_store.constants import MAX_IN_MEMORY_SIZE
 from onyx.file_store.file_store import FileStore, get_default_file_store
 from onyx.utils.csv_utils import sanitize_csv_cell_or_none
+from onyx.utils.logger import setup_logger
+
+logger = setup_logger()
+
+
+def _normalize_period(
+    period: tuple[datetime, datetime] | None,
+) -> tuple[datetime, datetime]:
+    if period is None:
+        return (
+            datetime.fromtimestamp(0, tz=timezone.utc),
+            datetime.now(tz=timezone.utc),
+        )
+    # time-picker sends a time which is at the beginning of the day
+    # so we need to add one day to the end time to make it inclusive
+    return (period[0], period[1] + timedelta(days=1))
 
 
 def generate_chat_messages_report(
     db_session: Session,
     file_store: FileStore,
     report_id: str,
-    period: tuple[datetime, datetime] | None,
+    period: tuple[datetime, datetime],
 ) -> str:
     file_name = f"{report_id}_chat_sessions"
 
-    if period is None:
-        period = (
-            datetime.fromtimestamp(0, tz=timezone.utc),
-            datetime.now(tz=timezone.utc),
-        )
-    else:
-        # time-picker sends a time which is at the beginning of the day
-        # so we need to add one day to the end time to make it inclusive
-        period = (
-            period[0],
-            period[1] + timedelta(days=1),
-        )
-
     with tempfile.SpooledTemporaryFile(
         max_size=MAX_IN_MEMORY_SIZE, mode="w+"
     ) as temp_file:
@@ -128,6 +133,57 @@ def generate_user_report(
     return file_id
 
 
+def generate_usage_breakdown_report(
+    file_store: FileStore,
+    report_id: str,
+    rows: Iterable[UsageExportRow],
+) -> str:
+    file_name = f"{report_id}_usage_by_user"
+
+    with tempfile.SpooledTemporaryFile(
+        max_size=MAX_IN_MEMORY_SIZE, mode="w+"
+    ) as temp_file:
+        csvwriter = csv.writer(temp_file, delimiter=",")
+        csvwriter.writerow(
+            [
+                "user_email",
+                "day",
+                "model",
+                "flow",
+                "provider",
+                "input_tokens",
+                "output_tokens",
+                "cache_read_tokens",
+                "cost_cents",
+            ]
+        )
+        for row in rows:
+            # User-controlled strings: formula-injection guard.
+            csvwriter.writerow(
+                [
+                    sanitize_csv_cell_or_none(row.email),
+                    row.day,
+                    sanitize_csv_cell_or_none(row.model),
+                    sanitize_csv_cell_or_none(row.flow),
+                    sanitize_csv_cell_or_none(row.provider),
+                    row.input_tokens,
+                    row.output_tokens,
+                    row.cache_read_tokens,
+                    row.cost_cents,
+                ]
+            )
+
+        temp_file.seek(0)
+        file_id = file_store.save_file(
+            content=temp_file,
+            display_name=file_name,
+            file_origin=FileOrigin.GENERATED_REPORT,
+            file_type="text/csv",
+        )
+
+    return file_id
+
+
 def create_new_usage_report(
     db_session: Session,
     user_id: UUID_ID | None,  # None = auto-generated
@@ -136,47 +192,72 @@ def create_new_usage_report(
 ) -> UsageReportMetadata:
     report_id = report_id or str(uuid.uuid4())
     file_store = get_default_file_store()
+    normalized_period = _normalize_period(period)
 
-    messages_file_id = generate_chat_messages_report(
-        db_session, file_store, report_id, period
-    )
-    users_file_id = generate_user_report(db_session, file_store, report_id)
-
-    # Re-check just before writing the final report: the API-level check
-    # happens before this (async) task runs, so a second request with the
-    # same client-supplied report_id can slip past it while this task is
-    # still generating the first report.
-    if usage_report_id_in_use(db_session, uuid.UUID(report_id)):
-        raise ValueError(f"report_id {report_id} is already in use")
-
-    with tempfile.SpooledTemporaryFile(max_size=MAX_IN_MEMORY_SIZE) as zip_buffer:
-        with zipfile.ZipFile(zip_buffer, "a", zipfile.ZIP_DEFLATED) as zip_file:
-            # write messages
-            chat_messages_tmpfile = file_store.read_file(
-                messages_file_id, mode="b", use_tempfile=True
-            )
-            zip_file.writestr(
-                "chat_messages.csv",
-                chat_messages_tmpfile.read(),
-            )
+    intermediate_file_ids: list[str] = []
+    try:
+        messages_file_id = generate_chat_messages_report(
+            db_session, file_store, report_id, normalized_period
+        )
+        intermediate_file_ids.append(messages_file_id)
+        users_file_id = generate_user_report(db_session, file_store, report_id)
+        intermediate_file_ids.append(users_file_id)
 
-            # write users
-            users_tmpfile = file_store.read_file(
-                users_file_id, mode="b", use_tempfile=True
-            )
-            zip_file.writestr("users.csv", users_tmpfile.read())
+        query_start, query_end = normalized_period
+        usage_rows = iter_usage_export(db_session, query_start, query_end)
+        usage_breakdown_file_id = generate_usage_breakdown_report(
+            file_store, report_id, usage_rows
+        )
+        intermediate_file_ids.append(usage_breakdown_file_id)
 
-        zip_buffer.seek(0)
+        # Re-check just before writing the final report: the API-level check
+        # happens before this (async) task runs, so a second request with the
+        # same client-supplied report_id can slip past it while this task is
+        # still generating the first report.
+        if usage_report_id_in_use(db_session, uuid.UUID(report_id)):
+            raise ValueError(f"report_id {report_id} is already in use")
 
-        # store zip blob to file_store
-        report_name = f"{datetime.now(tz=timezone.utc).strftime('%Y-%m-%d')}_{report_id}_usage_report.zip"
-        file_store.save_file(
-            content=zip_buffer,
-            display_name=report_name,
-            file_origin=FileOrigin.GENERATED_REPORT,
-            file_type="application/zip",
-            file_id=report_name,
-        )
+        with tempfile.SpooledTemporaryFile(max_size=MAX_IN_MEMORY_SIZE) as zip_buffer:
+            with zipfile.ZipFile(zip_buffer, "a", zipfile.ZIP_DEFLATED) as zip_file:
+                # write messages
+                chat_messages_tmpfile = file_store.read_file(
+                    messages_file_id, mode="b", use_tempfile=True
+                )
+                zip_file.writestr(
+                    "chat_messages.csv",
+                    chat_messages_tmpfile.read(),
+                )
+
+                # write users
+                users_tmpfile = file_store.read_file(
+                    users_file_id, mode="b", use_tempfile=True
+                )
+                zip_file.writestr("users.csv", users_tmpfile.read())
+
+                usage_breakdown_tmpfile = file_store.read_file(
+                    usage_breakdown_file_id, mode="b", use_tempfile=True
+                )
+                zip_file.writestr("usage_by_user.csv", usage_breakdown_tmpfile.read())
+
+            zip_buffer.seek(0)
+
+            # store zip blob to file_store
+            report_name = f"{datetime.now(tz=timezone.utc).strftime('%Y-%m-%d')}_{report_id}_usage_report.zip"
+            file_store.save_file(
+                content=zip_buffer,
+                display_name=report_name,
+                file_origin=FileOrigin.GENERATED_REPORT,
+                file_type="application/zip",
+                file_id=report_name,
+            )
+    finally:
+        for file_id in intermediate_file_ids:
+            try:
+                file_store.delete_file(file_id, error_on_missing=False)
+            except Exception:
+                logger.exception(
+                    "Failed to delete temporary usage report file %s", file_id
+                )
 
     # add report after zip file is written
     new_report = write_usage_report(db_session, report_name, user_id, period)

backend/onyx/db/user_usage.py

@@ -4,7 +4,7 @@
 model, flow, provider), not an append-only per-call ledger."""
 
 from collections import defaultdict
-from collections.abc import Sequence
+from collections.abc import Iterator, Sequence
 from datetime import datetime, timedelta
 from math import ceil
 from typing import Any, cast
@@ -210,12 +210,11 @@ def get_user_usage_by_day_and_model(
     ]
 
 
-def get_usage_export(
-    db_session: Session,
+def _get_usage_export_query(
     start: datetime,
     end: datetime,
     model: str | None = None,
-) -> list[UsageExportRow]:
+) -> Any:
     utc_day = func.date(func.timezone("UTC", UserUsage.window_start))
     # Deleted users/API keys leave user_id NULL but keep their spend. An inner
     # join would hide that spend here while the tenant-wide totals still count
@@ -254,10 +253,22 @@ def get_usage_export(
     if model is not None:
         query = query.where(UserUsage.model == model)
 
-    rows = db_session.execute(query).all()
+    return query
 
-    return [
-        UsageExportRow(
+
+def iter_usage_export(
+    db_session: Session,
+    start: datetime,
+    end: datetime,
+    model: str | None = None,
+) -> Iterator[UsageExportRow]:
+    result = db_session.execute(
+        _get_usage_export_query(start, end, model).execution_options(
+            stream_results=True
+        )
+    ).yield_per(1000)
+    for email, mdl, flow, provider, day, in_tok, out_tok, cache_tok, cost in result:
+        yield UsageExportRow(
             email=str(email),
             model=mdl,
             flow=flow,
@@ -268,8 +279,15 @@ def get_usage_export(
             cache_read_tokens=int(cache_tok or 0),
             cost_cents=float(cost or 0.0),
         )
-        for email, mdl, flow, provider, day, in_tok, out_tok, cache_tok, cost in rows
-    ]
+
+
+def get_usage_export(
+    db_session: Session,
+    start: datetime,
+    end: datetime,
+    model: str | None = None,
+) -> list[UsageExportRow]:
+    return list(iter_usage_export(db_session, start, end, model))
 
 
 def get_usage_reset_window_start(