|
| 1 | +"""Always-on, provider-agnostic per-query usage tracking. |
| 2 | +
|
| 3 | +This is deliberately independent of the optional LLM conversational-memory |
| 4 | +feature (``api/memory/graphiti_tool.py``). Memory writes are opt-in |
| 5 | +(``use_memory``), gated to OpenAI/Azure providers, and lazily created — so |
| 6 | +they cannot be used to measure adoption. This module records *every* query, |
| 7 | +regardless of provider or the ``use_memory`` flag, onto the central |
| 8 | +``Organizations`` graph (which already holds ``User``/``Identity``/``Token``). |
| 9 | +
|
| 10 | +For each query we maintain, fire-and-forget: |
| 11 | +
|
| 12 | +* Denormalized counters + activity timestamps on the ``User`` node |
| 13 | + (``query_count``/``success_count``/``error_count``/``last_active``/ |
| 14 | + ``first_query_at``) for cheap reads. |
| 15 | +* A per-query ``(:UsageEvent)`` node linked ``(User)-[:PERFORMED]->`` carrying |
| 16 | + ``graph_id``/``is_demo``/``success``/``timestamp`` for time-series, per-DB |
| 17 | + and success-rate analytics. |
| 18 | +
|
| 19 | +Writes never block or fail a request: they run as background tasks whose |
| 20 | +exceptions are logged and swallowed, mirroring |
| 21 | +``api.core.pipeline.save_memory_background``. |
| 22 | +""" |
| 23 | + |
| 24 | +import asyncio |
| 25 | +import base64 |
| 26 | +import binascii |
| 27 | +import hashlib |
| 28 | +import logging |
| 29 | +from typing import Optional |
| 30 | + |
| 31 | +from api.config import ORGANIZATIONS_GRAPH |
| 32 | +from api.core.db_resolver import resolve_db |
| 33 | +from api.core.pipeline import background_tasks_var, is_general_graph |
| 34 | + |
| 35 | +# Single round-trip: bump the User counters/timestamps and append a UsageEvent. |
| 36 | +# Uses MATCH (not MERGE) on User so an unknown email is a silent no-op rather |
| 37 | +# than creating a phantom user from the query path. ``timestamp()`` is FalkorDB |
| 38 | +# epoch-millis, matching every other timestamp in the Organizations graph. |
| 39 | +_RECORD_USAGE_CYPHER = """ |
| 40 | +MATCH (u:User {email: $email}) |
| 41 | +SET u.query_count = coalesce(u.query_count, 0) + 1, |
| 42 | + u.success_count = coalesce(u.success_count, 0) + (CASE WHEN $success THEN 1 ELSE 0 END), |
| 43 | + u.error_count = coalesce(u.error_count, 0) + (CASE WHEN $success THEN 0 ELSE 1 END), |
| 44 | + u.last_active = timestamp(), |
| 45 | + u.first_query_at = coalesce(u.first_query_at, timestamp()) |
| 46 | +CREATE (u)-[:PERFORMED]->(e:UsageEvent { |
| 47 | + graph_id: $graph_id, |
| 48 | + is_demo: $is_demo, |
| 49 | + success: $success, |
| 50 | + timestamp: timestamp() |
| 51 | +}) |
| 52 | +""" |
| 53 | + |
| 54 | + |
| 55 | +def _decode_email(user_id: str) -> Optional[str]: |
| 56 | + """Recover the user's email from the base64 ``user_id``. |
| 57 | +
|
| 58 | + Inverse of ``base64.b64encode(email.encode())`` in |
| 59 | + ``api/auth/user_management.py``. Returns ``None`` on malformed input so the |
| 60 | + caller can skip tracking instead of raising. |
| 61 | + """ |
| 62 | + if not user_id: |
| 63 | + return None |
| 64 | + try: |
| 65 | + email = base64.b64decode(user_id, validate=True).decode("utf-8") |
| 66 | + except (binascii.Error, ValueError, UnicodeDecodeError): |
| 67 | + logging.warning("Usage tracking: could not decode user_id to email") |
| 68 | + return None |
| 69 | + # b64decode is lenient about padding/length; require an email-shaped result |
| 70 | + # so a malformed id can't trigger a phantom DB write (matches the docstring). |
| 71 | + if "@" not in email: |
| 72 | + logging.warning("Usage tracking: decoded user_id is not a valid email") |
| 73 | + return None |
| 74 | + return email |
| 75 | + |
| 76 | + |
| 77 | +async def _write_usage(email: str, graph_id: str, is_demo: bool, success: bool, db) -> None: |
| 78 | + """Perform the single Cypher write against the Organizations graph.""" |
| 79 | + organizations_graph = resolve_db(db).select_graph(ORGANIZATIONS_GRAPH) |
| 80 | + await organizations_graph.query( |
| 81 | + _RECORD_USAGE_CYPHER, |
| 82 | + { |
| 83 | + "email": email, |
| 84 | + "graph_id": graph_id, |
| 85 | + "is_demo": is_demo, |
| 86 | + "success": success, |
| 87 | + }, |
| 88 | + ) |
| 89 | + # Structured-ish log line so usage is visible to log aggregators even |
| 90 | + # before any read API exists. graph_id is the namespaced name |
| 91 | + # ({base64(email)}_{db}) and base64 email is reversible, so log a short |
| 92 | + # stable hash instead of the raw value — this keeps user identity out of |
| 93 | + # logs and also neutralizes the CodeQL log-injection vector. |
| 94 | + graph_ref = hashlib.sha256(graph_id.encode()).hexdigest()[:12] |
| 95 | + logging.info( |
| 96 | + "usage_event graph=%s is_demo=%s success=%s", |
| 97 | + graph_ref, is_demo, success, |
| 98 | + ) |
| 99 | + |
| 100 | + |
| 101 | +def record_query_usage_background( |
| 102 | + user_id: str, |
| 103 | + namespaced: str, |
| 104 | + success: bool, |
| 105 | + *, |
| 106 | + db=None, |
| 107 | + task_sink: Optional[set] = None, |
| 108 | +) -> None: |
| 109 | + """Schedule fire-and-forget usage tracking for one query. |
| 110 | +
|
| 111 | + Returns immediately. The write runs as a background task whose failure is |
| 112 | + logged but never propagated, so tracking can never break or delay a query |
| 113 | + response. Called unconditionally at pipeline completion — independent of |
| 114 | + ``use_memory`` and the LLM provider. |
| 115 | +
|
| 116 | + Args: |
| 117 | + user_id: Base64-encoded email (the namespacing id used by the routes). |
| 118 | + namespaced: The fully-namespaced graph name the query ran against; |
| 119 | + already demo-aware, so it doubles as the recorded ``graph_id``. |
| 120 | + success: Whether SQL execution succeeded (no execution error). |
| 121 | + db: Optional FalkorDB handle; resolves to the server singleton when None. |
| 122 | + task_sink: Optional set the scheduled task is added to (and auto-removed |
| 123 | + from on completion) so callers can await any in-flight tracking |
| 124 | + writes before shutdown. |
| 125 | + """ |
| 126 | + email = _decode_email(user_id) |
| 127 | + if email is None: |
| 128 | + return |
| 129 | + |
| 130 | + is_demo = is_general_graph(namespaced) |
| 131 | + sink = task_sink if task_sink is not None else background_tasks_var.get() |
| 132 | + |
| 133 | + task = asyncio.create_task( |
| 134 | + _write_usage(email, namespaced, is_demo, success, db) |
| 135 | + ) |
| 136 | + |
| 137 | + if sink is not None: |
| 138 | + sink.add(task) |
| 139 | + task.add_done_callback(sink.discard) |
| 140 | + |
| 141 | + def _log_done(t: "asyncio.Task") -> None: |
| 142 | + if t.cancelled(): |
| 143 | + return |
| 144 | + exc = t.exception() |
| 145 | + if exc is not None: |
| 146 | + logging.error("Usage tracking save failed: %s", exc) # nosemgrep |
| 147 | + |
| 148 | + task.add_done_callback(_log_done) |
0 commit comments