diff --git a/app/core/router.py b/app/core/router.py index c0d44b0..6bef649 100644 --- a/app/core/router.py +++ b/app/core/router.py @@ -13,6 +13,16 @@ # SERVICE_MAP entry: this gateway proxies to it, it does not implement # any ServiceNow logic itself. "servicenow": "http://servicenow:8097", + # Public API v1 ("M1": GET /v1/usage): the gateway proxies this one + # read-only call into omnibioai-billing's existing + # GET /billing/organizations/{id}/subscription/usage-limits, reusing + # the same build_upstream_headers()-forwarded bearer token every + # other SERVICE_MAP entry relies on -- billing-service independently + # verifies it against the same shared-platform JWT secret (see its + # app/core/iam.py::_verify_caller), exactly like omnibioai-rag does + # for the "rag" entry above. See SERVICE_PERMISSION_MAP below for the + # permission this requires. + "billing": "http://billing-service:8005", } @@ -51,6 +61,13 @@ def resolve_service(service: str) -> str | None: # finer check" pattern omnibioai-rag's dataset.read/app/api/iam.py # already established, not a new authorization model. "servicenow": "servicenow_incident.read", + # Public API v1 (GET /v1/usage): omnibioai-auth's Permission Registry + # already has "usage.read" registered as "reserved -- not yet + # enforced by any route" -- the exact same state dataset.read/ + # model.use/workflow.execute were in before this gateway's IAM + # Foundation integration made them real. This is usage.read's first + # real consumer. + "billing": "usage.read", } diff --git a/app/routes/v1.py b/app/routes/v1.py index 82782fd..eaba63c 100644 --- a/app/routes/v1.py +++ b/app/routes/v1.py @@ -27,13 +27,20 @@ from app.core.config import Config from app.core.router import resolve_service from app.routes.gateway import build_upstream_headers, proxy -from app.services.literature_contract import UnsupportedRequestError, build_public_answer, build_rag_query +from app.services.literature_contract import ( + UnsupportedRequestError, + build_public_answer, + build_public_search, + build_rag_query, + build_rag_search_query, +) from app.services.v1_store import V1Store, request_fingerprint router = APIRouter(prefix="/v1") store = V1Store(Config.V1_REDIS_URL, Config.USAGE_REDIS_URL) ANSWER_RESOURCE = "literature.answer" +SEARCH_RESOURCE = "literature.search" _IDEMPOTENCY_KEY = re.compile(r"^[\x21-\x7e]{1,255}$") @@ -70,13 +77,17 @@ async def _rate_limited(request: Request, subject: str, request_id: str): return headers, None -async def _forward(request: Request, method: str, path: str, body=None): - url = f"{resolve_service('rag')}/{path}" +async def _forward_to(service: str, request: Request, method: str, path: str, body=None): + url = f"{resolve_service(service)}/{path}" if request.url.query: url = f"{url}?{request.url.query}" return await proxy.forward(url=url, method=method, headers=build_upstream_headers(request), body=body) +async def _forward(request: Request, method: str, path: str, body=None): + return await _forward_to("rag", request, method, path, body) + + def _upstream_error(status: int, response, request_id: str, headers: dict): if status >= 500: return _error(502 if status == 500 else status, "upstream_error", @@ -85,8 +96,18 @@ def _upstream_error(status: int, response, request_id: str, headers: dict): "The literature service rejected the request.", request_id, headers, detail=response) -@router.post("/literature/answers") -async def literature_answers(request: Request): +async def _handle_billable_literature_call( + request: Request, *, resource: str, build_rag_body, build_public_response, quota_exceeded_message: str, +): + """Shared lifecycle for every billable /v1/literature/* call + (currently /answers and /search): auth'd-org check, contract + translation, rate limit, idempotency replay, quota check, forward to + RAG, response translation, usage emission, quota consumption. The two + callers differ only in which resource they bill, how they translate + their request/response, and their quota-exceeded wording -- every + other step (in particular the idempotency/quota/usage sequencing) + must stay identical between them, so it lives here once rather than + as two copies that could silently drift apart.""" request_id = getattr(request.state, "trace_id", "") subject, org_id, user_id = _caller(request) if not org_id: @@ -99,7 +120,7 @@ async def literature_answers(request: Request): return _error(400, "invalid_request", "Request body must be JSON.", request_id) try: - rag_body = build_rag_query(body) + rag_body = build_rag_body(body) except UnsupportedRequestError as exc: return _error(400, "unsupported_request", exc.message, request_id, detail={"field": exc.field}) @@ -127,13 +148,11 @@ async def literature_answers(request: Request): return _error(409, "idempotency_in_progress", "A request with this Idempotency-Key is still running.", request_id, headers) - remaining = await store.quota_remaining(org_id, ANSWER_RESOURCE) + remaining = await store.quota_remaining(org_id, resource) if remaining is not None and remaining <= 0: if idempotency_key is not None: await store.idempotency_finish(subject, idempotency_key, fingerprint, 402, None) - return _error(402, "quota_exceeded", - "Your organization has used its included answers. Add a payment method or upgrade the plan.", - request_id, headers) + return _error(402, "quota_exceeded", quota_exceeded_message, request_id, headers) started = time.monotonic() status, response = await _forward(request, "POST", "v1/query", rag_body) @@ -144,7 +163,7 @@ async def literature_answers(request: Request): await store.idempotency_finish(subject, idempotency_key, fingerprint, status, response) return _upstream_error(status, response, request_id, headers) - public_response = build_public_answer( + public_response = build_public_response( response, domain=body.get("domain"), request_id=request_id, latency_ms=latency_ms, ) if idempotency_key is not None: @@ -154,7 +173,7 @@ async def literature_answers(request: Request): await store.emit_usage( org_id=org_id, user_id=user_id, - resource=ANSWER_RESOURCE, + resource=resource, trace_id=request_id, dedup_key=f"{subject}:{idempotency_key}" if idempotency_key else request_id, metadata={ @@ -164,10 +183,38 @@ async def literature_answers(request: Request): "idempotency_key_sha256": request_fingerprint(idempotency_key) if idempotency_key else None, }, ) - await store.consume_quota(org_id, ANSWER_RESOURCE) + await store.consume_quota(org_id, resource) return JSONResponse(public_response, status_code=status, headers={**headers, "X-Request-Id": request_id}) +@router.post("/literature/answers") +async def literature_answers(request: Request): + return await _handle_billable_literature_call( + request, + resource=ANSWER_RESOURCE, + build_rag_body=build_rag_query, + build_public_response=build_public_answer, + quota_exceeded_message="Your organization has used its included answers. " + "Add a payment method or upgrade the plan.", + ) + + +@router.post("/literature/search") +async def literature_search(request: Request): + """Billable unit: 1 search (see the design doc's pricing table -- + priced around 1/10th of an answer). Retrieval only: never invokes an + LLM, via RAG's mode="search" (app/services/literature_contract.py's + build_rag_search_query sets it).""" + return await _handle_billable_literature_call( + request, + resource=SEARCH_RESOURCE, + build_rag_body=build_rag_search_query, + build_public_response=build_public_search, + quota_exceeded_message="Your organization has used its included searches. " + "Add a payment method or upgrade the plan.", + ) + + @router.get("/literature/studies") async def literature_studies(request: Request): """Free: the queryable studies/domains, from omnibioai-rag GET /v1/studies. @@ -181,3 +228,110 @@ async def literature_studies(request: Request): if not 200 <= status < 300: return _upstream_error(status, response, request_id, headers) return JSONResponse(response, status_code=status, headers={**headers, "X-Request-Id": request_id}) + + +@router.get("/literature/domains") +async def literature_domains(request: Request): + """Free: the queryable research domains, under the frozen public + name the design doc uses ("domain", not RAG's internal "study"). + Same underlying data as /literature/studies above (kept as-is for + backward compatibility) via the same omnibioai-rag GET /v1/studies + call, reshaped to the public contract. Rate-limited like every /v1 + call, never billed.""" + request_id = getattr(request.state, "trace_id", "") + subject, _, _ = _caller(request) + headers, limited = await _rate_limited(request, subject, request_id) + if limited: + return limited + status, response = await _forward(request, "GET", "v1/studies") + if not 200 <= status < 300: + return _upstream_error(status, response, request_id, headers) + domains = [ + {"name": s.get("name"), "abstract_count": s.get("abstract_count")} + for s in (response.get("studies") or []) + ] + return JSONResponse({"domains": domains}, status_code=status, headers={**headers, "X-Request-Id": request_id}) + + +@router.get("/usage") +async def literature_usage(request: Request): + """Free: the caller's organization's included/used/remaining units + for the current billing period, from omnibioai-billing's existing + GET /billing/organizations/{id}/subscription/usage-limits -- the + gateway's first synchronous call into billing-service (see + app/core/router.py's SERVICE_MAP/SERVICE_PERMISSION_MAP entries, + gated on usage.read). Rate-limited like every /v1 call, never + billed. + + Estimated charge in dollars (also mentioned in the design doc) is + deliberately omitted: that needs billing's cost-summary endpoint and + its own start_date/end_date period math, which this does not yet do. + Reporting a wrong number would be worse than omitting it -- the same + principle app/services/literature_contract.py applies to token + counts. + """ + request_id = getattr(request.state, "trace_id", "") + subject, org_id, _ = _caller(request) + if not org_id: + return _error(403, "organization_required", + "This API is billed to an organization; your account has none.", request_id) + + headers, limited = await _rate_limited(request, subject, request_id) + if limited: + return limited + + status, response = await _forward_to( + "billing", request, "GET", f"billing/organizations/{org_id}/subscription/usage-limits", + ) + if status == 404: + return _error(404, "no_active_plan", "Your organization has no active billing plan.", request_id, headers) + if not 200 <= status < 300: + return _error(502 if status >= 500 else status, "upstream_error", + "The billing service failed to report usage.", request_id, headers, + detail=response if status < 500 else None) + + usage = [ + { + "resource": item.get("resource"), + "unit": item.get("unit"), + "period": item.get("period"), + "included": item.get("included"), + "used": item.get("used"), + "remaining": item.get("remaining"), + } + for item in (response.get("limits") or []) + ] + return JSONResponse( + {"plan": response.get("plan_name"), "as_of": response.get("as_of"), "usage": usage}, + status_code=200, headers={**headers, "X-Request-Id": request_id}, + ) + + +@router.get("/models") +async def literature_models(request: Request): + """Free: the model catalog -- answer/embedding models currently + served, with source and price. Rate-limited like every /v1 call, + never billed. No upstream call: there is exactly one model path + today, RAG's own GPU-hosted default (see the design's "largest + gaps" #4 -- Claude/OpenAI routing and bring-your-own-key do not + exist yet, and build_rag_query already rejects any request that + would need one). + + price is null, not a placeholder dollar figure: the design doc's own + pricing section says to measure real GPU cost per answer first + (milestone M0, 1,000 representative questions) before setting + prices, and that measurement has not been run. The model actually + used for a given answer is already reported per-call in + /v1/literature/answers' response (its `model` field is the real + value RAG used; "default" here is the stable identifier a caller + passes back as this endpoint's own `model` request field). + """ + request_id = getattr(request.state, "trace_id", "") + subject, _, _ = _caller(request) + headers, limited = await _rate_limited(request, subject, request_id) + if limited: + return limited + models = [ + {"model": "default", "source": "omnibioai_gpu", "billed_by": "query", "price": None}, + ] + return JSONResponse({"models": models}, status_code=200, headers={**headers, "X-Request-Id": request_id}) diff --git a/app/services/literature_contract.py b/app/services/literature_contract.py index 6ea254b..183c66d 100644 --- a/app/services/literature_contract.py +++ b/app/services/literature_contract.py @@ -61,6 +61,62 @@ def build_rag_query(body: dict) -> dict: return query +def build_rag_search_query(body: dict) -> dict: + """Translate a public /v1/literature/search request body into RAG's + POST /v1/query body (with mode="search" set by the caller, not here -- + that's the one field this function doesn't own, since it's not part of + the public contract). Raises UnsupportedRequestError for a missing/ + blank question, the only field this endpoint requires. + + Unlike build_rag_query above, model/use_own_key/stream are not + validated here at all: search is retrieval-only regardless of what a + caller sends for them, so there is no "unsupported" case for fields + that were never going to change this call's behavior. + """ + if not isinstance(body, dict): + raise UnsupportedRequestError("body", "Request body must be a JSON object.") + + question = body.get("question") + if not isinstance(question, str) or not question.strip(): + raise UnsupportedRequestError("question", "\"question\" is required and must be a non-empty string.") + + query: dict = {"query": question, "study": body.get("domain") or "default", "mode": "search"} + max_results = body.get("max_results") + if isinstance(max_results, (int, float)) and not isinstance(max_results, bool) and max_results > 0: + query["top_k"] = int(max_results) + return query + + +def build_public_search(rag_result: dict, *, domain, request_id: str, latency_ms: int) -> dict: + """Translate RAG's /v1/query (mode="search") response into the frozen + public /v1/literature/search response shape: ranked documents, no + generated answer.""" + documents = rag_result.get("documents") or [] + results = [ + { + "pmid": doc.get("pmid"), + "title": doc.get("title"), + "year": doc.get("year"), + "score": doc.get("citation_confidence", doc.get("similarity_score")), + # The design doc calls this a "snippet"; RAG only has the full + # abstract text to offer, not a separately-generated excerpt, + # so that's what's returned under this name rather than + # fabricating a truncation. + "snippet": doc.get("abstract"), + } + for doc in documents + ] + return { + "id": f"srch_{request_id or uuid.uuid4().hex}", + "results": results, + "domain": rag_result.get("study", domain), + "usage": { + "searches": 1, + "latency_ms": latency_ms, + }, + } + + def build_public_answer(rag_result: dict, *, domain, request_id: str, latency_ms: int) -> dict: """Translate RAG's /v1/query response into the frozen public /v1/literature/answers response shape.""" diff --git a/tests/test_literature_contract.py b/tests/test_literature_contract.py index 8baaf7f..29dd6a7 100644 --- a/tests/test_literature_contract.py +++ b/tests/test_literature_contract.py @@ -4,7 +4,13 @@ """ import pytest -from app.services.literature_contract import UnsupportedRequestError, build_public_answer, build_rag_query +from app.services.literature_contract import ( + UnsupportedRequestError, + build_public_answer, + build_public_search, + build_rag_query, + build_rag_search_query, +) # --------------------------------------------------------------------------- @@ -122,3 +128,92 @@ def test_domain_falls_back_to_caller_supplied_value_when_rag_result_has_no_study def test_missing_request_id_still_produces_an_id(): answer = build_public_answer({}, domain=None, request_id="", latency_ms=1) assert answer["id"].startswith("ans_") and len(answer["id"]) > len("ans_") + + +# --------------------------------------------------------------------------- +# build_rag_search_query +# --------------------------------------------------------------------------- + +def test_search_minimal_request_sets_mode_search(): + assert build_rag_search_query({"question": "What is TP53?"}) == { + "query": "What is TP53?", "study": "default", "mode": "search", + } + + +def test_search_domain_maps_to_study(): + assert build_rag_search_query({"question": "q", "domain": "Oncology"})["study"] == "Oncology" + + +@pytest.mark.parametrize("max_results", [1, 10, 25.0]) +def test_search_max_results_maps_to_top_k(max_results): + assert build_rag_search_query({"question": "q", "max_results": max_results})["top_k"] == int(max_results) + + +@pytest.mark.parametrize("max_results", [0, -1, "10", None, True]) +def test_search_invalid_or_absent_max_results_omits_top_k(max_results): + body = {"question": "q"} + if max_results is not None: + body["max_results"] = max_results + assert "top_k" not in build_rag_search_query(body) + + +@pytest.mark.parametrize("body", [{}, {"question": ""}, {"question": " "}, {"question": 5}]) +def test_search_missing_or_blank_question_is_rejected(body): + with pytest.raises(UnsupportedRequestError) as exc: + build_rag_search_query(body) + assert exc.value.field == "question" + + +def test_search_non_dict_body_is_rejected(): + with pytest.raises(UnsupportedRequestError) as exc: + build_rag_search_query(["not", "a", "dict"]) + assert exc.value.field == "body" + + +@pytest.mark.parametrize("field", ["model", "use_own_key", "stream"]) +def test_search_does_not_reject_answer_only_fields(field): + # Unlike build_rag_query, search never calls an LLM regardless of + # these fields, so there is nothing to reject them for. + build_rag_search_query({"question": "q", field: "anything-truthy-or-not"}) # must not raise + + +# --------------------------------------------------------------------------- +# build_public_search +# --------------------------------------------------------------------------- + +def test_search_builds_full_public_shape_from_rag_result(): + rag_result = { + "study": "Oncology", "mode": "search", "summary": None, + "documents": [{"pmid": "123", "title": "TP53 review", "year": 2021, + "citation_confidence": 0.9, "abstract": "TP53 is a tumor suppressor."}], + } + result = build_public_search(rag_result, domain="Oncology", request_id="req-1", latency_ms=7) + assert result == { + "id": "srch_req-1", + "results": [{"pmid": "123", "title": "TP53 review", "year": 2021, "score": 0.9, + "snippet": "TP53 is a tumor suppressor."}], + "domain": "Oncology", + "usage": {"searches": 1, "latency_ms": 7}, + } + + +def test_search_falls_back_to_similarity_score_when_citation_confidence_absent(): + rag_result = {"documents": [{"pmid": "1", "similarity_score": 0.5}]} + result = build_public_search(rag_result, domain=None, request_id="r", latency_ms=1) + assert result["results"] == [{"pmid": "1", "title": None, "year": None, "score": 0.5, "snippet": None}] + + +def test_search_empty_documents_yields_empty_results(): + result = build_public_search({}, domain=None, request_id="r", latency_ms=1) + assert result["results"] == [] + assert "answer" not in result + + +def test_search_domain_falls_back_to_caller_supplied_value_when_rag_result_has_no_study(): + result = build_public_search({}, domain="Oncology", request_id="r", latency_ms=1) + assert result["domain"] == "Oncology" + + +def test_search_missing_request_id_still_produces_an_id(): + result = build_public_search({}, domain=None, request_id="", latency_ms=1) + assert result["id"].startswith("srch_") and len(result["id"]) > len("srch_") diff --git a/tests/test_v1_literature.py b/tests/test_v1_literature.py index e665cf2..80fa1d9 100644 --- a/tests/test_v1_literature.py +++ b/tests/test_v1_literature.py @@ -310,6 +310,118 @@ def test_studies_is_free_and_rate_limited(client, redis, upstream, monkeypatch): assert client.get("/v1/literature/studies", headers={"Authorization": f"Bearer {KEY}"}).status_code == 429 +def test_domains_is_free_and_rate_limited(client, redis, upstream, monkeypatch): + upstream.return_value = (200, {"studies": [{"name": "oncology", "abstract_count": 12}]}) + resp = client.get("/v1/literature/domains", headers={"Authorization": f"Bearer {KEY}"}) + assert resp.status_code == 200 + assert resp.json() == {"domains": [{"name": "oncology", "abstract_count": 12}]} + assert upstream.call_args.kwargs["url"] == "http://rag:8096/v1/studies" + assert _usage(redis) == [] + + upstream.return_value = (500, {}) + assert client.get("/v1/literature/domains", headers={"Authorization": f"Bearer {KEY}"}).status_code == 502 + + monkeypatch.setattr(Config, "V1_RATE_LIMIT_PER_MINUTE", 0) + assert client.get("/v1/literature/domains", headers={"Authorization": f"Bearer {KEY}"}).status_code == 429 + + +def test_search_forwards_search_mode_and_bills_search_resource(client, redis, upstream): + """/v1/literature/search uses the same billable-call lifecycle as + /v1/literature/answers, but bills "literature.search" (not + "literature.answer") and tells RAG mode="search" so no LLM is + invoked -- the response has no generated answer.""" + upstream.return_value = (200, { + "study": "default", "mode": "search", "summary": None, + "documents": [{"pmid": "1", "title": "TP53 review", "year": 2021, "citation_confidence": 0.9, + "abstract": "TP53 is a tumor suppressor."}], + }) + resp = client.post("/v1/literature/search", json={"question": "What does TP53 do?"}, + headers={"Authorization": f"Bearer {KEY}"}) + assert resp.status_code == 200 + body = resp.json() + assert body["id"].startswith("srch_") + assert body["results"] == [{"pmid": "1", "title": "TP53 review", "year": 2021, "score": 0.9, + "snippet": "TP53 is a tumor suppressor."}] + assert body["domain"] == "default" + assert body["usage"]["searches"] == 1 + assert "answer" not in body + + kwargs = upstream.call_args.kwargs + assert kwargs["method"] == "POST" + assert kwargs["body"] == {"query": "What does TP53 do?", "study": "default", "mode": "search"} + + (event,) = _usage(redis) + assert event["resource"] == "literature.search" + + +def test_search_rejects_missing_question(client, redis, upstream): + resp = client.post("/v1/literature/search", json={}, headers={"Authorization": f"Bearer {KEY}"}) + assert resp.status_code == 400 + assert resp.json()["error"]["detail"]["field"] == "question" + upstream.assert_not_called() + + +def test_search_quota_exceeded(client, redis, upstream): + with patch.object(v1.store, "quota_remaining", AsyncMock(return_value=0)): + resp = client.post("/v1/literature/search", json={"question": "q"}, + headers={"Authorization": f"Bearer {KEY}"}) + assert resp.status_code == 402 + assert resp.json()["error"]["type"] == "quota_exceeded" + upstream.assert_not_called() + + +def test_usage_requires_an_organization(client, redis, upstream): + with patch.object(_main_mod.iam, "validate_api_key", AsyncMock(return_value={**USER, "org_id": None})): + resp = client.get("/v1/usage", headers={"Authorization": f"Bearer {KEY}"}) + assert resp.status_code == 403 + assert resp.json()["error"]["type"] == "organization_required" + upstream.assert_not_called() + + +def test_usage_translates_billing_response_and_is_free(client, redis, upstream): + upstream.return_value = (200, { + "organization_id": 42, "billing_plan_id": 1, "plan_name": "Free", "as_of": "2026-10-03", + "limits": [ + {"service": "api", "action": "answer", "resource": "literature.answer", "unit": "requests", + "period": "monthly", "included": 100, "used": 12, "remaining": 88, "percentage_used": 12.0}, + ], + }) + resp = client.get("/v1/usage", headers={"Authorization": f"Bearer {KEY}"}) + assert resp.status_code == 200 + body = resp.json() + assert body["plan"] == "Free" + assert body["as_of"] == "2026-10-03" + assert body["usage"] == [ + {"resource": "literature.answer", "unit": "requests", "period": "monthly", + "included": 100, "used": 12, "remaining": 88}, + ] + assert upstream.call_args.kwargs["url"] == "http://billing-service:8005/billing/organizations/42/subscription/usage-limits" + assert _usage(redis) == [] + + +def test_usage_no_active_plan_returns_404(client, redis, upstream): + upstream.return_value = (404, {"detail": "No active subscription"}) + resp = client.get("/v1/usage", headers={"Authorization": f"Bearer {KEY}"}) + assert resp.status_code == 404 + assert resp.json()["error"]["type"] == "no_active_plan" + + +def test_usage_upstream_5xx_is_502(client, redis, upstream): + upstream.return_value = (500, {}) + resp = client.get("/v1/usage", headers={"Authorization": f"Bearer {KEY}"}) + assert resp.status_code == 502 + assert resp.json()["error"]["type"] == "upstream_error" + + +def test_models_is_free_static_and_never_calls_upstream(client, redis, upstream): + resp = client.get("/v1/models", headers={"Authorization": f"Bearer {KEY}"}) + assert resp.status_code == 200 + assert resp.json() == {"models": [{"model": "default", "source": "omnibioai_gpu", + "billed_by": "query", "price": None}]} + upstream.assert_not_called() + assert _usage(redis) == [] + + def test_unauthenticated_v1_is_rejected(client): assert client.post("/v1/literature/answers", json=BODY).status_code == 401