Claude/kg retrieval logs merge m4tvyg

shepherd_server/base_routes.py

@@ -23,12 +23,17 @@
     get_callback_query_id,
     get_logs,
     get_message,
+    get_query_log_level,
     get_query_state,
     remove_callback_id,
     save_logs,
     save_message,
 )
-from shepherd_utils.logger import attach_query_handler, setup_logging
+from shepherd_utils.logger import (
+    attach_query_handler,
+    resolve_log_level,
+    setup_logging,
+)
 from shepherd_utils.otel import setup_tracer
 
 setup_logging()
@@ -86,8 +91,11 @@ async def run_query(
     query_id = str(uuid.uuid4())[:8]
     response_id = str(uuid.uuid4())[:8]
     # Set up logger
-    log_level = query.get("log_level") or settings.log_level
-    level_number = logging._nameToLevel[log_level]
+    # Same resolver the callback handler and the merge use, so an unparseable
+    # level from a client falls back to the default instead of failing intake.
+    level_number = resolve_log_level(
+        query.get("log_level"), resolve_log_level(settings.log_level)
+    )
     logger = logging.getLogger(f"shepherd.{query_id}")
     logger.setLevel(level_number)
     attach_query_handler(logger)
@@ -346,10 +354,22 @@ async def callback(
             content={"detail": "Invalid request body"},
             status_code=422,
         )
-    # Now that the body is parsed, apply the query's requested log level.
-    log_level = response.get("log_level") or "INFO"
-    level_number = logging._nameToLevel[log_level]
+    # get associated query id for this callback. Resolved before anything else
+    # is logged: the level to log it at is a property of the query, and this
+    # mapping is the only route back to it.
+    original_query = await get_callback_query_id(callback_id, logger)
+    if original_query is None:
+        # No callback->query mapping, so there's no response_id to persist these
+        # logs under; surface it in the console/collector at least.
+        logger.warning(f"Callback {callback_id}: couldn't find original query.")
+        return Response("Couldn't find original query.", 500)
+    # Apply the level the client asked for. It lives in the stored query -- a
+    # TRAPI response has no log_level field, so the body we were just posted
+    # can't tell us (it used to be read from there, which quietly meant INFO for
+    # every callback and dropped a DEBUG query's logs from here on).
+    level_number = await get_query_log_level(original_query[0], logger)
     logger.setLevel(level_number)
+    logger.debug(f"Got original query: {original_query}")
     # logger.info(response)
     results = response["message"].get("results")
     if results is None:
@@ -367,14 +387,6 @@ async def callback(
     logger.debug(
         f"[{callback_id}] for query graph: {response['message'].get('query_graph')}"
     )
-    # get associated query id for this callback
-    original_query = await get_callback_query_id(callback_id, logger)
-    logger.debug(f"Got original query: {original_query}")
-    if original_query is None:
-        # No callback->query mapping, so there's no response_id to persist these
-        # logs under; surface it in the console/collector at least.
-        logger.warning(f"Callback {callback_id}: couldn't find original query.")
-        return Response("Couldn't find original query.", 500)
     # if len(response["message"]["results"]) > 0:
     #     with open(
     #         f"shepherd_server/debug/{query_id}_{callback_id}_response.json",
@@ -445,7 +457,7 @@ async def get_query_response(
     query_id: str,
 ):
     """Get a query response."""
-    level_number = logging._nameToLevel["INFO"]
+    level_number = logging.INFO
     logger = logging.getLogger("shepherd.get_query")
     logger.setLevel(level_number)
     attach_query_handler(logger)

shepherd_utils/db.py

@@ -14,7 +14,7 @@
 from psycopg_pool import AsyncConnectionPool
 
 from .config import settings
-from .logger import get_query_handler
+from .logger import get_query_handler, resolve_log_level
 
 PG_RETRIES = 5
 
@@ -410,6 +410,33 @@ async def get_message(
     return message
 
 
+async def get_query_log_level(
+    query_id: str,
+    logger: logging.Logger,
+    default: Union[int, None] = None,
+) -> int:
+    """The log level the client asked for, read back from the stored query.
+
+    The stored query is the only record of the requested level once a request
+    has been handed off. A TRAPI *response* has no ``log_level`` field, so
+    nothing a subservice posts back to ``/callback`` carries it -- everything
+    hanging off a callback (the handler's own logs, the merge task it enqueues,
+    the retrieval logs that merge folds into the query's log list) has to come
+    back here for it.
+
+    Falls back to the server default when the query didn't ask for a level, or
+    when it can no longer be read.
+    """
+    if default is None:
+        default = resolve_log_level(settings.log_level)
+    try:
+        query = await get_message(query_id, logger)
+    except Exception as e:
+        logger.warning(f"Couldn't read the log level for query {query_id}: {e}")
+        return default
+    return resolve_log_level(query.get("log_level"), default)
+
+
 # ---------------------------------------------------------------------------
 # Per-query "ready callback" index
 #

shepherd_utils/logger.py

@@ -96,6 +96,30 @@ def log_handler(self):
         return self._log_handler
 
 
+# TRAPI's LogLevel enum (plus CRITICAL), mapped to the ``logging`` levels the
+# pipeline filters against. The one place level names are understood, so the
+# server, the workers and the log entries subservices send back all agree.
+LOG_LEVELS = {
+    "DEBUG": logging.DEBUG,
+    "INFO": logging.INFO,
+    "WARNING": logging.WARNING,
+    "ERROR": logging.ERROR,
+    "CRITICAL": logging.CRITICAL,
+}
+
+
+def resolve_log_level(name, default: int = logging.INFO) -> int:
+    """Turn a TRAPI log level name into a ``logging`` level number.
+
+    Anything unusable -- absent, empty, or a name we don't know -- falls back to
+    ``default`` rather than raising: a query asking for a level we can't parse
+    should still run.
+    """
+    if not name:
+        return default
+    return LOG_LEVELS.get(str(name).upper(), default)
+
+
 def get_query_handler(logger: logging.Logger):
     """Return the query log handler attached to ``logger``, or None."""
     return next(

shepherd_utils/shared.py

@@ -25,7 +25,7 @@
 from .config import settings
 from .db import initialize_db, save_logs
 from .heartbeat import Heartbeat
-from .logger import attach_query_handler, setup_logging
+from .logger import attach_query_handler, resolve_log_level, setup_logging
 from .reclaim import reclaim_orphaned
 
 # Cap each per-stream duration queue so a stopped monitor can't OOM the broker.
@@ -482,7 +482,7 @@ async def get_tasks(
     it ``None`` to fall back to ``PER_STREAM_MIN_IDLE_SEC`` / settings.
     """
     # Set up logger
-    level_number = logging._nameToLevel[settings.log_level]
+    level_number = resolve_log_level(settings.log_level)
     worker_logger = logging.getLogger(f"shepherd.{stream}.{consumer}")
     worker_logger.setLevel(level_number)
     attach_query_handler(worker_logger)

tests/unit/test_callback_log_level.py

@@ -0,0 +1,136 @@
+"""Tests for the log level the /callback endpoint runs at.
+
+The level a query asked for lives in the query itself. A TRAPI *response* has
+no ``log_level`` field, so the body a subservice posts to ``/callback`` can't
+carry it -- the handler has to read it back from the stored query. Getting this
+wrong is quiet: everything downstream of the callback (the handler's own logs,
+the merge task it enqueues, and the retrieval logs merge folds into the query's
+log list) silently runs at INFO and a DEBUG query loses its logs.
+"""
+
+import logging
+
+import orjson
+import pytest
+from starlette.requests import Request
+
+from shepherd_server import base_routes
+from shepherd_server.base_routes import ARATargetEnum, callback
+from shepherd_utils.db import save_message
+
+logger = logging.getLogger(__name__)
+
+
+def _make_request(body: bytes) -> Request:
+    """A Starlette Request that streams ``body`` in one chunk."""
+    scope = {
+        "type": "http",
+        "method": "POST",
+        "path": "/aragorn/callback/cb-1",
+        "headers": [(b"content-length", str(len(body)).encode())],
+    }
+    sent = False
+
+    async def receive():
+        nonlocal sent
+        if sent:
+            return {"type": "http.disconnect"}
+        sent = True
+        return {"type": "http.request", "body": body, "more_body": False}
+
+    return Request(scope, receive)
+
+
+def _patch_callback_deps(monkeypatch):
+    """Stub the postgres-backed lookups and capture the enqueued task."""
+    tasks = []
+
+    async def _get_callback_query_id(callback_id, logger):
+        return ("q-1", "{}")
+
+    async def _get_query_state(query_id, logger):
+        # response_id lives at index 7 of the shepherd_brain row.
+        return [None, None, None, None, None, None, None, "resp-1"]
+
+    async def _add_ready_callback(response_id, callback_id, logger):
+        return None
+
+    async def _add_task(stream, fields, logger):
+        tasks.append((stream, fields))
+
+    async def _save_logs(response_id, logger):
+        return None
+
+    monkeypatch.setattr(base_routes, "get_callback_query_id", _get_callback_query_id)
+    monkeypatch.setattr(base_routes, "get_query_state", _get_query_state)
+    monkeypatch.setattr(base_routes, "add_ready_callback", _add_ready_callback)
+    monkeypatch.setattr(base_routes, "add_task", _add_task)
+    monkeypatch.setattr(base_routes, "save_logs", _save_logs)
+    return tasks
+
+
+def _callback_body(**extra):
+    body = {"message": {"results": [], "knowledge_graph": {"nodes": {}, "edges": {}}}}
+    body.update(extra)
+    return orjson.dumps(body)
+
+
+@pytest.mark.asyncio
+async def test_callback_takes_its_level_from_the_stored_query(redis_mock, monkeypatch):
+    """A DEBUG query keeps logging at DEBUG once its callbacks come back."""
+    tasks = _patch_callback_deps(monkeypatch)
+    await save_message("q-1", {"log_level": "DEBUG", "message": {}}, logger)
+
+    response = await callback(
+        ARATargetEnum.ARAGORN, "cb-1", _make_request(_callback_body())
+    )
+
+    assert response.status_code == 200
+    assert logging.getLogger("shepherd.cb-1").level == logging.DEBUG
+    # ...and the level rides along to the merge, which filters the retrieval's
+    # own log entries against it.
+    assert tasks[0][0] == "merge_message"
+    assert tasks[0][1]["log_level"] == logging.DEBUG
+
+
+@pytest.mark.asyncio
+async def test_callback_ignores_a_level_claimed_by_the_body(redis_mock, monkeypatch):
+    """The regression: the level used to be read off the posted body. Nothing a
+    subservice sends back gets to lower (or raise) what the client asked for."""
+    tasks = _patch_callback_deps(monkeypatch)
+    await save_message("q-1", {"log_level": "DEBUG", "message": {}}, logger)
+
+    await callback(
+        ARATargetEnum.ARAGORN, "cb-2", _make_request(_callback_body(log_level="ERROR"))
+    )
+
+    assert logging.getLogger("shepherd.cb-2").level == logging.DEBUG
+    assert tasks[0][1]["log_level"] == logging.DEBUG
+
+
+@pytest.mark.asyncio
+async def test_callback_falls_back_to_the_server_default(redis_mock, monkeypatch):
+    """A query that didn't ask for a level gets the configured default."""
+    tasks = _patch_callback_deps(monkeypatch)
+    await save_message("q-1", {"message": {}}, logger)
+
+    await callback(ARATargetEnum.ARAGORN, "cb-3", _make_request(_callback_body()))
+
+    assert logging.getLogger("shepherd.cb-3").level == logging.INFO
+    assert tasks[0][1]["log_level"] == logging.INFO
+
+
+@pytest.mark.asyncio
+async def test_callback_survives_an_unreadable_query(redis_mock, monkeypatch):
+    """The stored query can have expired out from under a late callback; that
+    shouldn't fail the callback, just fall back to the default level."""
+    tasks = _patch_callback_deps(monkeypatch)
+    # "q-1" intentionally not stored.
+
+    response = await callback(
+        ARATargetEnum.ARAGORN, "cb-4", _make_request(_callback_body())
+    )
+
+    assert response.status_code == 200
+    assert logging.getLogger("shepherd.cb-4").level == logging.INFO
+    assert tasks[0][1]["log_level"] == logging.INFO

tests/unit/test_logger.py

@@ -10,6 +10,7 @@
     get_logging_config,
     get_query_handler,
     get_worker_logger,
+    resolve_log_level,
 )
 
 
@@ -69,6 +70,24 @@ def test_query_logger_handler_named_query_log_handler():
     assert handler.name == "query_log_handler"
 
 
+def test_resolve_log_level_maps_trapi_level_names():
+    """The one place level names are understood -- the requested level on a
+    query, and the level on a log entry a subservice sends back."""
+    assert resolve_log_level("DEBUG") == logging.DEBUG
+    assert resolve_log_level("debug") == logging.DEBUG
+    assert resolve_log_level("WARNING") == logging.WARNING
+    assert resolve_log_level("ERROR") == logging.ERROR
+
+
+def test_resolve_log_level_falls_back_for_anything_unusable():
+    """A query naming a level we can't parse should still run."""
+    assert resolve_log_level(None) == logging.INFO
+    assert resolve_log_level("") == logging.INFO
+    assert resolve_log_level("LOUD") == logging.INFO
+    assert resolve_log_level("LOUD", logging.WARNING) == logging.WARNING
+    assert resolve_log_level(None, logging.DEBUG) == logging.DEBUG
+
+
 def test_drain_empties_the_queue_and_returns_oldest_first():
     """Reading is destructive: the same records must not be handed out twice."""
     handler = QueryLogger().log_handler