Deploy daa0049

app/pipeline/nodes/log_eval.py

@@ -7,7 +7,7 @@ from datetime import datetime, timezone
 from typing import Callable
 
 from app.models.pipeline import PipelineState
-from app.services.loki_sink import ship_to_loki
+from app.services.axiom_sink import ship_to_axiom
 
 logger = logging.getLogger(__name__)
 _PENDING_TASKS: set[asyncio.Task[None]] = set()
@@ -73,7 +73,7 @@ def _source_hit_proxy(state: PipelineState) -> int:
 def make_log_eval_node(db_path: str) -> Callable[[PipelineState], dict]:
     """
     Writes interaction to SQLite synchronously (<5ms) inside the request lifespan.
-    Ships a sanitised observability record to Grafana Loki asynchronously.
+    Ships a sanitised observability record to Axiom asynchronously.
 
     The `path` field tags which pipeline branch produced the answer:
       "cache_hit"    — served from semantic cache, no LLM called.
@@ -141,7 +141,7 @@ def make_log_eval_node(db_path: str) -> Callable[[PipelineState], dict]:
             )
             return cursor.lastrowid  # type: ignore[return-value]
 
-    def _build_loki_record(state: PipelineState) -> dict:
+    def _build_axiom_record(state: PipelineState) -> dict:
         reranked_chunks = state.get("reranked_chunks", [])
         return {
             "timestamp": datetime.now(tz=timezone.utc).isoformat(),
@@ -183,13 +183,13 @@ def make_log_eval_node(db_path: str) -> Callable[[PipelineState], dict]:
             row_id = _write_to_sqlite(state)
             try:
                 loop = asyncio.get_running_loop()
-                task = loop.create_task(ship_to_loki(_build_loki_record(state)))
+                task = loop.create_task(ship_to_axiom(_build_axiom_record(state)))
                 _PENDING_TASKS.add(task)
                 task.add_done_callback(_PENDING_TASKS.discard)
             except RuntimeError:
                 # Called outside an event loop (for example, synchronous unit tests).
-                # SQLite write still succeeds; Loki shipping is skipped.
-                logger.debug("No running event loop; skipping async Loki ship for this call.")
+                # SQLite write still succeeds; Axiom shipping is skipped.
+                logger.debug("No running event loop; skipping async Axiom ship for this call.")
 
             return {"interaction_id": row_id}
         except Exception as e:

app/services/axiom_sink.py

@@ -0,0 +1,106 @@
+from __future__ import annotations
+
+import hashlib
+import json
+import logging
+import os
+from typing import Any
+
+import httpx
+
+logger = logging.getLogger(__name__)
+
+_AXIOM_TOKEN = os.environ.get("AXIOM_TOKEN", "").strip()
+_AXIOM_DATASET = os.environ.get("AXIOM_DATASET", "").strip()
+_AXIOM_ORG_ID = os.environ.get("AXIOM_ORG_ID", "").strip()
+_AXIOM_INGEST_URL = os.environ.get("AXIOM_INGEST_URL", "").strip()
+
+if _AXIOM_INGEST_URL:
+    _AXIOM_URL = _AXIOM_INGEST_URL
+elif _AXIOM_DATASET:
+    _AXIOM_URL = f"https://api.axiom.co/v1/datasets/{_AXIOM_DATASET}/ingest"
+else:
+    _AXIOM_URL = ""
+
+_AXIOM_ENABLED = bool(_AXIOM_URL and _AXIOM_TOKEN)
+if not _AXIOM_ENABLED:
+    logger.warning("Axiom is not configured; Axiom shipping is disabled.")
+
+
+def _sha_prefix(value: str, prefix_len: int) -> str:
+    return hashlib.sha256(value.encode("utf-8")).hexdigest()[:prefix_len]
+
+
+def _to_float_or_none(value: Any) -> float | None:
+    if isinstance(value, (int, float)):
+        return float(value)
+    return None
+
+
+def _resolve_source_hit_proxy(record: dict[str, Any], top_rerank_score: float | None, chunk_count: int) -> int:
+    explicit = record.get("source_hit_proxy")
+    if explicit is not None:
+        return int(explicit)
+    return int(top_rerank_score is not None and top_rerank_score > -1.5 and chunk_count >= 2)
+
+
+def _build_sanitized_record(record: dict[str, Any]) -> dict[str, Any]:
+    query = str(record.get("query", ""))
+    session_id = str(record.get("session_id", ""))
+
+    rerank_scores = record.get("rerank_scores") or []
+    source_types_used = record.get("source_types_used") or []
+    chunk_count = int(record.get("chunk_count", 0) or 0)
+    top_chunk_doc_id = str(record.get("top_chunk_doc_id", "") or "")
+
+    top_rerank_score = _to_float_or_none(record.get("top_rerank_score"))
+    source_hit_proxy = _resolve_source_hit_proxy(record, top_rerank_score, chunk_count)
+
+    return {
+        "timestamp": str(record.get("timestamp", "")),
+        "session_id": _sha_prefix(session_id, 12) if session_id else "",
+        "path": str(record.get("path", "rag") or "rag"),
+        "cached": bool(record.get("cached", False)),
+        "latency_ms": int(record.get("latency_ms", 0) or 0),
+        "retrieval_attempts": int(record.get("retrieval_attempts", 0) or 0),
+        "top_rerank_score": top_rerank_score,
+        "focused_source_type": str(record.get("focused_source_type", "") or ""),
+        "sibling_expansion_count": int(record.get("sibling_expansion_count", 0) or 0),
+        "critic_groundedness": record.get("critic_groundedness"),
+        "critic_completeness": record.get("critic_completeness"),
+        "critic_specificity": record.get("critic_specificity"),
+        "critic_quality": str(record.get("critic_quality", "") or ""),
+        "is_enumeration_query": bool(record.get("is_enumeration_query", False)),
+        "guard_passed": bool(record.get("guard_passed", False)),
+        "query_complexity": str(record.get("query_complexity", "") or ""),
+        "is_followup": bool(record.get("is_followup", False)),
+        "is_audio_mode": bool(record.get("is_audio_mode", False)),
+        "query_hash": _sha_prefix(query, 16) if query else "",
+        "chunk_count": chunk_count,
+        "top_chunk_doc_id": top_chunk_doc_id,
+        "source_types_used": sorted(str(source_type) for source_type in source_types_used if str(source_type)),
+        "follow_up_count": len(record.get("follow_ups") or []),
+        "rerank_scores": [float(score) for score in rerank_scores if isinstance(score, (int, float))],
+        "source_hit_proxy": source_hit_proxy,
+    }
+
+
+async def ship_to_axiom(record: dict[str, Any]) -> None:
+    if not _AXIOM_ENABLED:
+        return
+
+    try:
+        sanitized = _build_sanitized_record(record)
+        headers = {
+            "Authorization": f"Bearer {_AXIOM_TOKEN}",
+            "Content-Type": "application/json",
+        }
+        if _AXIOM_ORG_ID:
+            headers["X-Axiom-Org-Id"] = _AXIOM_ORG_ID
+
+        timeout = httpx.Timeout(3.0)
+        async with httpx.AsyncClient(timeout=timeout, headers=headers) as client:
+            response = await client.post(_AXIOM_URL, json=[sanitized])
+            response.raise_for_status()
+    except Exception as exc:
+        logger.warning("Axiom ship failed: %s", exc)

tests/test_axiom_sink.py

@@ -0,0 +1,166 @@
+import asyncio
+import json
+import sqlite3
+
+import httpx
+import pytest
+
+from app.pipeline.nodes.log_eval import make_log_eval_node
+from app.services import axiom_sink
+
+
+class _FakeResponse:
+    def raise_for_status(self) -> None:
+        return None
+
+
+@pytest.mark.asyncio
+async def test_ship_to_axiom_sanitises_pii(monkeypatch) -> None:
+    captured: dict = {}
+
+    class _FakeAsyncClient:
+        def __init__(self, *args, **kwargs):
+            self._args = args
+            self._kwargs = kwargs
+
+        async def __aenter__(self):
+            return self
+
+        async def __aexit__(self, exc_type, exc, tb):
+            return None
+
+        async def post(self, url, json):
+            await asyncio.sleep(0)
+            captured["url"] = url
+            captured["payload"] = json
+            return _FakeResponse()
+
+    monkeypatch.setattr(axiom_sink, "_AXIOM_ENABLED", True)
+    monkeypatch.setattr(axiom_sink, "_AXIOM_URL", "https://api.axiom.co/v1/datasets/personabot/ingest")
+    monkeypatch.setattr(axiom_sink, "_AXIOM_TOKEN", "xaat-test-token")
+    monkeypatch.setattr(axiom_sink, "_AXIOM_ORG_ID", "")
+    monkeypatch.setattr(axiom_sink.httpx, "AsyncClient", _FakeAsyncClient)
+
+    await axiom_sink.ship_to_axiom(
+        {
+            "timestamp": "2026-04-21T00:00:00Z",
+            "session_id": "session-123",
+            "query": "my private query",
+            "answer": "long text",
+            "path": "rag",
+            "cached": False,
+            "latency_ms": 100,
+            "retrieval_attempts": 1,
+            "top_rerank_score": -1.0,
+            "focused_source_type": "resume",
+            "sibling_expansion_count": 0,
+            "critic_groundedness": 3,
+            "critic_completeness": 3,
+            "critic_specificity": 3,
+            "critic_quality": "high",
+            "is_enumeration_query": False,
+            "guard_passed": True,
+            "query_complexity": "simple",
+            "is_followup": False,
+            "is_audio_mode": False,
+            "follow_ups": [],
+            "chunk_count": 0,
+            "top_chunk_doc_id": "",
+            "source_types_used": [],
+            "rerank_scores": [],
+        }
+    )
+
+    serialized = json.dumps(captured["payload"])
+    assert "my private query" not in serialized
+    assert "long text" not in serialized
+
+    payload_record = captured["payload"][0]
+    assert payload_record.get("query_hash")
+
+
+@pytest.mark.asyncio
+async def test_ship_to_axiom_noop_when_not_configured(monkeypatch) -> None:
+    called = {"post": False}
+
+    class _NeverCalledClient:
+        def __init__(self, *args, **kwargs):
+            self._args = args
+            self._kwargs = kwargs
+
+        async def __aenter__(self):
+            return self
+
+        async def __aexit__(self, exc_type, exc, tb):
+            return None
+
+        async def post(self, url, json):
+            await asyncio.sleep(0)
+            called["post"] = True
+            return _FakeResponse()
+
+    monkeypatch.setattr(axiom_sink, "_AXIOM_ENABLED", False)
+    monkeypatch.setattr(axiom_sink.httpx, "AsyncClient", _NeverCalledClient)
+
+    await axiom_sink.ship_to_axiom({"query": "x", "answer": "y"})
+
+    assert called["post"] is False
+
+
+@pytest.mark.asyncio
+async def test_ship_to_axiom_swallows_http_error(monkeypatch) -> None:
+    class _FailingClient:
+        def __init__(self, *args, **kwargs):
+            self._args = args
+            self._kwargs = kwargs
+
+        async def __aenter__(self):
+            return self
+
+        async def __aexit__(self, exc_type, exc, tb):
+            return None
+
+        async def post(self, url, json):
+            await asyncio.sleep(0)
+            raise httpx.ConnectError("connection failed")
+
+    monkeypatch.setattr(axiom_sink, "_AXIOM_ENABLED", True)
+    monkeypatch.setattr(axiom_sink, "_AXIOM_URL", "https://api.axiom.co/v1/datasets/personabot/ingest")
+    monkeypatch.setattr(axiom_sink, "_AXIOM_TOKEN", "xaat-test-token")
+    monkeypatch.setattr(axiom_sink, "_AXIOM_ORG_ID", "")
+    monkeypatch.setattr(axiom_sink.httpx, "AsyncClient", _FailingClient)
+
+    await axiom_sink.ship_to_axiom({"query": "q", "answer": "a", "path": "rag", "chunk_count": 0})
+
+
+@pytest.mark.asyncio
+async def test_source_hit_proxy_logged_to_sqlite(tmp_path) -> None:
+    db_path = str(tmp_path / "interactions.db")
+    node = make_log_eval_node(db_path)
+
+    state = {
+        "session_id": "s1",
+        "query": "What work experience does Darshan have?",
+        "answer": "He worked at VK Live.",
+        "reranked_chunks": [
+            {"text": "a", "metadata": {"doc_id": "d1", "source_type": "resume", "rerank_score": -1.0}},
+            {"text": "b", "metadata": {"doc_id": "d2", "source_type": "resume", "rerank_score": -1.2}},
+            {"text": "c", "metadata": {"doc_id": "d3", "source_type": "resume", "rerank_score": -1.3}},
+        ],
+        "latency_ms": 123,
+        "cached": False,
+        "path": "rag",
+        "is_enumeration_query": False,
+        "top_rerank_score": -1.0,
+        "retrieval_attempts": 1,
+        "follow_ups": [],
+    }
+
+    result = node(state)
+    assert "interaction_id" in result
+
+    with sqlite3.connect(db_path) as conn:
+        row = conn.execute("SELECT source_hit_proxy FROM interactions LIMIT 1").fetchone()
+
+    assert row is not None
+    assert row[0] == 1

tests/test_log_eval_privacy.py

@@ -46,16 +46,16 @@ def test_log_eval_stores_chunk_metadata_without_text(tmp_path) -> None:
 
 
 @pytest.mark.asyncio
-async def test_log_eval_sends_sanitized_loki_payload(monkeypatch, tmp_path) -> None:
+async def test_log_eval_sends_sanitized_axiom_payload(monkeypatch, tmp_path) -> None:
     db_path = str(tmp_path / "interactions.db")
     node = make_log_eval_node(db_path)
     captured: dict = {}
 
-    async def _fake_ship_to_loki(record: dict) -> None:
+    async def _fake_ship_to_axiom(record: dict) -> None:
         await asyncio.sleep(0)
         captured["record"] = record
 
-    monkeypatch.setattr("app.pipeline.nodes.log_eval.ship_to_loki", _fake_ship_to_loki)
+    monkeypatch.setattr("app.pipeline.nodes.log_eval.ship_to_axiom", _fake_ship_to_axiom)
 
     node(
         {