diff --git a/src/common/extensions.py b/src/common/extensions.py index e5ec3516..c9abf48f 100644 --- a/src/common/extensions.py +++ b/src/common/extensions.py @@ -1,10 +1,10 @@ """Extensions installed alongside InferiaLLM, found through the ``inferia.extensions`` entry point. Each entry point names a callable that returns an extension object. The gateway -calls its ``mount(app)`` at startup and the migrate step its ``migrate(dsn)``; an -extension may provide either. With nothing installed nothing changes, and an -extension that fails is logged and skipped, never raised: InferiaLLM must start -without it. +calls its ``mount(app)`` at startup, the migrate step its ``migrate(dsn)``, and the +inference app its ``before_request(...)`` on each chat request; an extension may +provide any of them. With nothing installed nothing changes, and an extension +that fails is logged and skipped, never raised: InferiaLLM must work without it. """ from __future__ import annotations @@ -12,6 +12,8 @@ from functools import lru_cache from importlib.metadata import entry_points +from starlette.exceptions import HTTPException + log = logging.getLogger(__name__) GROUP = "inferia.extensions" @@ -46,3 +48,23 @@ async def migrate_all(dsn: str) -> None: await ext.migrate(dsn) except Exception: log.exception("extension %r failed to migrate; skipped", ext) + + +async def before_request(**request) -> None: + """Let extensions check a request before the engine is called. + + ``request`` holds request_id, user_id, org_id, deployment_id, model, + request_type and body (the JSON body, never file uploads). An extension + refuses by raising HTTPException, which reaches the client like core's own + refusals; any other error is logged and the request goes on. + """ + for ext in installed(): + hook = getattr(ext, "before_request", None) + if hook is None: + continue + try: + await hook(**request) + except HTTPException: + raise + except Exception: + log.exception("extension %r failed in before_request; skipped", ext) diff --git a/src/inference/core/handlers/completion.py b/src/inference/core/handlers/completion.py index b1a41468..e6392d36 100644 --- a/src/inference/core/handlers/completion.py +++ b/src/inference/core/handlers/completion.py @@ -11,6 +11,8 @@ from fastapi import BackgroundTasks, HTTPException from fastapi.responses import StreamingResponse +from common import extensions +from common.http_client import request_id_ctx from inference.client import api_gateway_client from inference.config import settings from ..providers import resolve_upstream @@ -79,6 +81,18 @@ async def handle( applied_policies.append("quota") await api_gateway_client.check_quota(user_context_id, model) + # 3b. Installed extensions; one may refuse with an HTTPException. + # /v1/messages comes through here too, already in this format. + await extensions.before_request( + request_id=request_id_ctx.get() or "", + user_id=user_context_id, + org_id=org_id, + deployment_id=str(deployment_id) if deployment_id else None, + model=model, + request_type="llm", + body=body, + ) + # 4. Prepare Provider Request endpoint_url = deployment.get("endpoint") if not endpoint_url: diff --git a/src/tests/common/test_extensions.py b/src/tests/common/test_extensions.py index 6026f9a5..d222c6d8 100644 --- a/src/tests/common/test_extensions.py +++ b/src/tests/common/test_extensions.py @@ -93,3 +93,38 @@ async def ext(dsn): monkeypatch.setattr(init, "migrate_all", ext) await init._apply_migrations("postgresql://x") assert order == [("core", "postgresql://x"), ("extensions", "postgresql://x")] + + +class _Checker: + def __init__(self, error=None): + self.error, self.seen = error, [] + + async def before_request(self, **request): + self.seen.append(request) + if self.error: + raise self.error + + +async def test_before_request_reaches_every_extension_that_has_it(install): + first, second = _Checker(), _Checker() + install(_EP("first", first), _EP("no_hook", _Ext()), _EP("second", second)) + await extensions.before_request(request_id="r1", model="m", body={"messages": []}) + assert first.seen == second.seen == [{"request_id": "r1", "model": "m", "body": {"messages": []}}] + + +async def test_a_refusal_reaches_the_caller(install): + from fastapi import HTTPException + + later = _Checker() + install(_EP("refuses", _Checker(HTTPException(status_code=400, detail="no"))), _EP("later", later)) + with pytest.raises(HTTPException) as e: + await extensions.before_request(request_id="r1") + assert e.value.status_code == 400 + assert later.seen == [] # refused requests stop at the first refusal + + +async def test_a_broken_check_is_skipped_and_the_request_goes_on(install): + good = _Checker() + install(_EP("broken", _Checker(RuntimeError("bug"))), _EP("good", good)) + await extensions.before_request(request_id="r1") + assert good.seen == [{"request_id": "r1"}] diff --git a/src/tests/inference/test_extension_request_check.py b/src/tests/inference/test_extension_request_check.py new file mode 100644 index 00000000..3ac919c9 --- /dev/null +++ b/src/tests/inference/test_extension_request_check.py @@ -0,0 +1,77 @@ +"""Chat requests go past installed extensions after the quota check, before the engine.""" +from unittest.mock import AsyncMock, MagicMock + +import pytest +from fastapi import HTTPException + +from common.http_client import request_id_ctx +from inference.core.handlers import completion as ch + +pytestmark = pytest.mark.asyncio + +BODY = {"model": "qwen", "messages": [{"role": "user", "content": "hi"}]} + + +@pytest.fixture +def handler(monkeypatch): + order = [] + context = { + "deployment": {"id": "dep-1", "endpoint": "http://engine.test", "engine": "vllm"}, + "user_id_context": "apikey:k1", + "org_id": "org-1", + "rate_limit_config": None, + "log_payloads": True, + } + monkeypatch.setattr(ch.GatewayService, "resolve_context", AsyncMock(return_value=context)) + + async def check_quota(user_id, model): + order.append("quota") + + monkeypatch.setattr(ch.api_gateway_client, "check_quota", check_quota) + + def resolve_upstream(*_a, **_k): + order.append("engine") + raise RuntimeError("stop before the engine call") + + monkeypatch.setattr(ch, "resolve_upstream", resolve_upstream) + seen = [] + + async def before_request(**request): + order.append("extensions") + seen.append(request) + + monkeypatch.setattr(ch.extensions, "before_request", before_request) + return order, seen, monkeypatch + + +async def _handle(): + token = request_id_ctx.set("trace-1") + try: + return await ch.CompletionHandler.handle("sk-key", dict(BODY), MagicMock()) + finally: + request_id_ctx.reset(token) + + +async def test_extensions_see_the_request_after_the_quota_and_before_the_engine(handler): + order, seen, _ = handler + with pytest.raises(RuntimeError, match="stop before the engine"): + await _handle() + assert order == ["quota", "extensions", "engine"] + assert seen == [{ + "request_id": "trace-1", "user_id": "apikey:k1", "org_id": "org-1", + "deployment_id": "dep-1", "model": "qwen", "request_type": "llm", "body": BODY, + }] + + +async def test_a_refusal_stops_the_request_before_the_engine(handler): + order, _, monkeypatch = handler + + async def refuse(**request): + order.append("extensions") + raise HTTPException(status_code=400, detail="refused") + + monkeypatch.setattr(ch.extensions, "before_request", refuse) + with pytest.raises(HTTPException) as e: + await _handle() + assert e.value.status_code == 400 + assert order == ["quota", "extensions"]