diff --git a/src/api_gateway/db/models/inference_log.py b/src/api_gateway/db/models/inference_log.py index 55ebd1b7..fa6632ab 100644 --- a/src/api_gateway/db/models/inference_log.py +++ b/src/api_gateway/db/models/inference_log.py @@ -16,7 +16,8 @@ class InferenceLog(Base): ) user_id = Column(String, index=True, nullable=False) # User ID or API key context ip_address = Column(String, index=True, nullable=True) # Client IP supplied by caller/proxy - + request_id = Column(String, index=True, nullable=True) # X-Request-ID; null on rows logged before it was stored + # Request Info request_payload = Column(JSON, nullable=True) # Full request body (optional to save space) model = Column(String, index=True, nullable=False) diff --git a/src/api_gateway/gateway/router.py b/src/api_gateway/gateway/router.py index 80db5249..426a167a 100644 --- a/src/api_gateway/gateway/router.py +++ b/src/api_gateway/gateway/router.py @@ -155,6 +155,7 @@ async def _persist_log_background(log_data: InferenceLogCreate, log_id: str): deployment_id=log_data.deployment_id, user_id=log_data.user_id, ip_address=log_data.ip_address, + request_id=log_data.request_id, model=log_data.model, request_payload=request_payload, latency_ms=log_data.latency_ms, diff --git a/src/api_gateway/management/insights.py b/src/api_gateway/management/insights.py index a5afbf05..7ab0a6ba 100644 --- a/src/api_gateway/management/insights.py +++ b/src/api_gateway/management/insights.py @@ -79,6 +79,7 @@ def _build_filters( ip_address: str | None, status: InsightsStatusFilter, deployment_type: InsightsDeploymentTypeFilter, + request_id: str | None = None, ) -> List[Any]: conditions: List[Any] = [ DBDeployment.org_id == org_id, @@ -101,6 +102,9 @@ def _build_filters( if ip_address: conditions.append(DBInferenceLog.ip_address == ip_address) + if request_id: + conditions.append(DBInferenceLog.request_id == request_id) + if status == "success": conditions.append(DBInferenceLog.status_code < 400) elif status == "error": @@ -529,6 +533,7 @@ async def get_insights_logs( deployment_id: str | None = Query(None), model: str | None = Query(None), ip_address: str | None = Query(None), + request_id: str | None = Query(None, max_length=128), status: InsightsStatusFilter = Query("all"), deployment_type: InsightsDeploymentTypeFilter = Query("all"), limit: int = Query(50, ge=1, le=200), @@ -553,6 +558,7 @@ async def get_insights_logs( ip_address, status, deployment_type, + request_id=request_id, ) clamped_limit = min(max(limit, 1), 200) diff --git a/src/api_gateway/schemas/logging.py b/src/api_gateway/schemas/logging.py index 9f7bedad..69a43b11 100644 --- a/src/api_gateway/schemas/logging.py +++ b/src/api_gateway/schemas/logging.py @@ -9,6 +9,7 @@ class InferenceLogCreate(BaseModel): user_id: str model: str ip_address: Optional[str] = None + request_id: Optional[str] = None request_payload: Optional[Dict[str, Any]] = None latency_ms: Optional[int] = None ttft_ms: Optional[int] = None @@ -30,6 +31,7 @@ class InferenceLogResponse(BaseModel): user_id: str model: str ip_address: Optional[str] = None + request_id: Optional[str] = None request_payload: Optional[Dict[str, Any]] = None latency_ms: Optional[int] = None ttft_ms: Optional[int] = None diff --git a/src/inference/client.py b/src/inference/client.py index 851a8b07..9c61a386 100644 --- a/src/inference/client.py +++ b/src/inference/client.py @@ -221,6 +221,8 @@ async def _send_log(): "is_streaming": is_streaming, "applied_policies": applied_policies, "ip_address": ip_address, + # create_task copied the request's context, so this is its id. + "request_id": request_id_ctx.get() or None, "request_type": request_type, } diff --git a/src/infra/schema/global_schema.sql b/src/infra/schema/global_schema.sql index 392480d7..516fa313 100644 --- a/src/infra/schema/global_schema.sql +++ b/src/infra/schema/global_schema.sql @@ -525,7 +525,8 @@ CREATE TABLE inference_logs ( deployment_id UUID NOT NULL, user_id VARCHAR NOT NULL, ip_address VARCHAR, - request_payload JSON, + request_id VARCHAR, + request_payload JSON, model VARCHAR NOT NULL, latency_ms INTEGER, ttft_ms INTEGER, @@ -543,6 +544,7 @@ CREATE TABLE inference_logs ( ); CREATE INDEX ix_inference_logs_user_id ON inference_logs (user_id); CREATE INDEX ix_inference_logs_ip_address ON inference_logs (ip_address); +CREATE INDEX ix_inference_logs_request_id ON inference_logs (request_id); CREATE INDEX ix_inference_logs_id ON inference_logs (id); CREATE INDEX ix_inference_logs_model ON inference_logs (model); CREATE INDEX ix_inference_logs_created_at ON inference_logs (created_at); diff --git a/src/infra/schema/migrations/20261026_inference_logs_request_id.sql b/src/infra/schema/migrations/20261026_inference_logs_request_id.sql new file mode 100644 index 00000000..b1a02bea --- /dev/null +++ b/src/infra/schema/migrations/20261026_inference_logs_request_id.sql @@ -0,0 +1,14 @@ +-- The inference app gives every request an X-Request-ID (#442). Storing it on +-- the log row lets a request seen elsewhere, in a client's error or a log +-- line, be found here. Rows logged before this keep it null. +-- +-- inference_logs is often the largest table, so the index is built +-- CONCURRENTLY and does not block inserts. That cannot run inside a +-- transaction, so the runner executes the part after @SPLIT@ on its own. +ALTER TABLE inference_logs + ADD COLUMN IF NOT EXISTS request_id VARCHAR; + +-- @SPLIT@ + +CREATE INDEX CONCURRENTLY IF NOT EXISTS ix_inference_logs_request_id + ON inference_logs (request_id); diff --git a/src/tests/api_gateway/test_inference_log_background.py b/src/tests/api_gateway/test_inference_log_background.py index 2679f45a..6fc8a67c 100644 --- a/src/tests/api_gateway/test_inference_log_background.py +++ b/src/tests/api_gateway/test_inference_log_background.py @@ -177,3 +177,20 @@ async def test_endpoint_returns_log_id(self): assert result["status"] == "ok" assert "log_id" in result assert len(result["log_id"]) == 36 # UUID length + + +@pytest.mark.asyncio +async def test_the_request_id_is_stored_on_the_row(): + from api_gateway.gateway.router import _persist_log_background + + mock_session = AsyncMock() + mock_session.__aenter__ = AsyncMock(return_value=mock_session) + mock_session.__aexit__ = AsyncMock(return_value=None) + added = [] + mock_session.add = added.append + + with patch("api_gateway.db.database.AsyncSessionLocal", return_value=mock_session): + await _persist_log_background(_make_log_data(request_id="trace-42"), "log-rid") + await _persist_log_background(_make_log_data(), "log-no-rid") + + assert [row.request_id for row in added] == ["trace-42", None] diff --git a/src/tests/api_gateway/test_management_insights.py b/src/tests/api_gateway/test_management_insights.py index 0288b74b..b22ae301 100644 --- a/src/tests/api_gateway/test_management_insights.py +++ b/src/tests/api_gateway/test_management_insights.py @@ -284,6 +284,37 @@ async def test_logs_pagination_and_limit_clamping(): assert len(response.items) == 2 +@pytest.mark.asyncio +async def test_logs_can_be_found_by_request_id(): + request = _make_request() + db = AsyncMock() + count_result = MagicMock() + count_result.scalar.return_value = 1 + logs_result = MagicMock() + logs_result.scalars.return_value.all.return_value = [_make_log("l1")] + db.execute.side_effect = [count_result, logs_result] + + await get_insights_logs( + request=request, + start_time=_now() - timedelta(days=1), + end_time=_now(), + deployment_id=None, + model=None, + ip_address=None, + request_id="trace-42", + status="all", + deployment_type="all", + limit=50, + offset=0, + db=db, + ) + + for call in db.execute.call_args_list: + stmt = call.args[0] + assert "inference_logs.request_id" in str(stmt) + assert "trace-42" in stmt.compile().params.values() + + @pytest.mark.asyncio async def test_empty_dataset_returns_zeroed_shapes(): request = _make_request() diff --git a/src/tests/inference/test_request_id.py b/src/tests/inference/test_request_id.py index aad8ced3..2855cce0 100644 --- a/src/tests/inference/test_request_id.py +++ b/src/tests/inference/test_request_id.py @@ -82,3 +82,25 @@ def test_calls_to_the_gateway_forward_the_id(): assert client._get_headers()["X-Request-ID"] == "trace-7" finally: request_id_ctx.reset(token) + + +@pytest.mark.asyncio +async def test_the_log_carries_the_id_of_the_request_being_logged(monkeypatch): + import asyncio + + client = ApiGatewayClient() + sent = [] + + class _Http: + async def post(self, path, json, headers): + sent.append(json) + + monkeypatch.setattr(client, "_get_client", lambda: _Http()) + token = request_id_ctx.set("trace-9") + try: + await client.log_inference(deployment_id="d1", user_id="apikey:k1", model="m") + finally: + request_id_ctx.reset(token) # the send runs later, in the context it was created in + await asyncio.sleep(0) + await asyncio.sleep(0) + assert sent[0]["request_id"] == "trace-9"