diff --git a/tests/_fake_openai_endpoint_server.py b/tests/_fake_openai_endpoint_server.py index caf3fb5ba2a..ac83e74b66a 100644 --- a/tests/_fake_openai_endpoint_server.py +++ b/tests/_fake_openai_endpoint_server.py @@ -8,11 +8,11 @@ those jobs failed with ``404 Application not found`` even though nothing in the PR was broken. This process is the local stand-in. A model points its ``api_base`` here and -gets back a well-formed chat/text/embedding response with realistic ``usage`` so -cost tracking and spend accounting still exercise their real code paths. The one -behavioral special case mirrors the old hosted mock: a request whose ``model`` -is ``429`` returns HTTP 429 so rate-limit and cooldown tests still have -something to trip on. +gets back a well-formed chat/text/embedding/moderation response with realistic +``usage`` so cost tracking and spend accounting still exercise their real code +paths. The one behavioral special case mirrors the old hosted mock: a request +whose ``model`` is ``429`` returns HTTP 429 so rate-limit and cooldown tests +still have something to trip on. """ from __future__ import annotations @@ -35,6 +35,21 @@ _SLOW_MODEL: Final = "slow-endpoint" _SLOW_RESPONSE_SECONDS: Final = 3.0 _PROMPT_TOKENS: Final = 20 _COMPLETION_TOKENS: Final = 20 +_MODERATION_CATEGORIES: Final = ( + "harassment", + "harassment/threatening", + "hate", + "hate/threatening", + "illicit", + "illicit/violent", + "self-harm", + "self-harm/instructions", + "self-harm/intent", + "sexual", + "sexual/minors", + "violence", + "violence/graphic", +) def _usage() -> dict[str, int]: @@ -220,6 +235,28 @@ async def triton_embeddings(_request: Request) -> Response: ) +def _moderation_result() -> dict[str, object]: + return { + "flagged": False, + "categories": {category: False for category in _MODERATION_CATEGORIES}, + "category_scores": {category: 0.0 for category in _MODERATION_CATEGORIES}, + "category_applied_input_types": {category: ["text"] for category in _MODERATION_CATEGORIES}, + } + + +async def moderations(request: Request) -> Response: + body: Final = await _parse_body(request) + raw_input: Final = body.get("input", "") + count: Final = len(raw_input) if isinstance(raw_input, list) else 1 + return JSONResponse( + { + "id": f"modr-{uuid.uuid4().hex[:24]}", + "model": _requested_model(body), + "results": [_moderation_result() for _ in range(max(count, 1))], + } + ) + + async def list_models(_request: Request) -> Response: return JSONResponse( { @@ -247,6 +284,8 @@ app = Starlette( Route("/embeddings", embeddings, methods=["POST"]), Route("/v1/embeddings", embeddings, methods=["POST"]), Route("/triton/embeddings", triton_embeddings, methods=["POST"]), + Route("/moderations", moderations, methods=["POST"]), + Route("/v1/moderations", moderations, methods=["POST"]), Route("/models", list_models, methods=["GET"]), Route("/v1/models", list_models, methods=["GET"]), ] diff --git a/tests/local_testing/test_fake_openai_endpoint.py b/tests/local_testing/test_fake_openai_endpoint.py index 79d8b4f97e3..d5236d3de1b 100644 --- a/tests/local_testing/test_fake_openai_endpoint.py +++ b/tests/local_testing/test_fake_openai_endpoint.py @@ -13,9 +13,11 @@ from __future__ import annotations import re from pathlib import Path +from typing import Final import httpx import pytest +from openai.types import ModerationCreateResponse from tests.fake_openai_endpoint import ( _LOCAL_DEFAULT, @@ -56,6 +58,20 @@ def test_chat_completion_shape(): assert body["usage"]["total_tokens"] == 40 +def test_moderations_route_parses_as_an_openai_response(): + base: Final = ensure_fake_openai_endpoint() + response: Final = httpx.post( + f"{base}/v1/moderations", + json={"input": ["I want to harm someone", "hello"], "model": "omni-moderation-latest"}, + timeout=10, + ) + assert response.status_code == 200 + parsed: Final = ModerationCreateResponse.model_validate(response.json()) + assert parsed.model == "omni-moderation-latest" + assert len(parsed.results) == 2 + assert parsed.results[0].categories.violence is False + + def test_triton_embeddings_route(): base = ensure_fake_openai_endpoint() response = httpx.post(f"{base}/triton/embeddings", json={"inputs": []}, timeout=10)