diff --git a/ai_feedback/__main__.py b/ai_feedback/__main__.py index 68c5507..a5c00fb 100644 --- a/ai_feedback/__main__.py +++ b/ai_feedback/__main__.py @@ -9,7 +9,13 @@ from . import code_processing, image_processing, text_processing from .helpers import arg_options from .helpers.constants import HELP_MESSAGES -from .models import ModelFactory +from .models import GatewayError, ModelFactory + +# Scopes not listed here are handled as code. +_PROCESSOR_BY_SCOPE = { + "image": image_processing.process_image, + "text": text_processing.process_text, +} _TYPE_BY_EXTENSION = { '.c': 'C', @@ -337,19 +343,18 @@ def main() -> int: print(f"Error: {e}") sys.exit(1) + prompt = prompt_content if args.scope == "image": prompt = {"prompt_content": prompt_content} - request, response = image_processing.process_image( - model, args, prompt, system_instructions, marking_instructions - ) - elif args.scope == "text": - request, response = text_processing.process_text( - model, args, prompt_content, system_instructions, marking_instructions - ) - else: - request, response = code_processing.process_code( - model, args, prompt_content, system_instructions, marking_instructions - ) + + process = _PROCESSOR_BY_SCOPE.get(args.scope, code_processing.process_code) + try: + request, response = process(model, args, prompt, system_instructions, marking_instructions) + except GatewayError as e: + # stderr, because the autotester reports a failed run from the subprocess's + # stderr; on stdout this would be swallowed and shown as "exit status 1". + print(f"Error: {e}", file=sys.stderr) + sys.exit(1) markdown_template = load_markdown_template(args.output_template) timestamp = datetime.now().strftime("%Y%m%d_%H%M%S") diff --git a/ai_feedback/models/OpenAIRemoteModel.py b/ai_feedback/models/OpenAIRemoteModel.py new file mode 100644 index 0000000..7a8eddd --- /dev/null +++ b/ai_feedback/models/OpenAIRemoteModel.py @@ -0,0 +1,135 @@ +import json +import os +from typing import Optional + +import openai +from ollama import Message + +from .OpenAIModel import OpenAIModel + + +class GatewayError(RuntimeError): + """The gateway would not serve the call. + + Raised instead of letting ``openai.APIError`` escape, so the reason (budget + exhausted, kill switch, missing attribution, gateway unreachable) reaches the + instructor as one readable line rather than the last line of a stack trace. + """ + + +def _failure_reason(error: openai.APIError) -> str: + """The gateway's own message, or the SDK's summary when the body is not ours. + + Our hooks answer with ``{"error": {"message": ...}}``, which the OpenAI SDK + parses into ``error.body``. Anything else — a proxy error page, a connection + failure carrying no body — falls back to the SDK's rendering. + """ + body = getattr(error, "body", None) + if isinstance(body, dict) and body.get("message"): + return str(body["message"]) + return str(error) + + +class OpenAIRemoteModel(OpenAIModel): + """An OpenAI-compatible model served through the MarkUs LiteLLM gateway. + + This is the sibling of :class:`RemoteModel`. ``RemoteModel`` talks to the + ``markus-ai-server`` proxy ("polymouth") with a custom payload and an + ``X-API-KEY`` header. ``OpenAIRemoteModel`` instead speaks the standard + OpenAI chat-completion contract (``Authorization: Bearer`` + OpenAI + request/response schema), which is what the self-hosted LiteLLM proxy + expects. The only differences from :class:`OpenAIModel` are the endpoint + (the LiteLLM gateway rather than ``api.openai.com``) and the per-call + attribution header. + + Attribution metadata (instance, course_id, assignment_id, group_id, + batch_id, category) is read from the ``LITELLM_SPEND_METADATA`` environment + variable and forwarded verbatim as the ``x-litellm-spend-logs-metadata`` + header. The autotester's AI tester sets that variable per invocation. The + gateway's pre-call hook reads the header to attribute spend to the right + course and to enforce the gatekeeper budget. See the ai-telemetry-gateway + project for the receiving side. + """ + + #: Header LiteLLM reads to persist arbitrary metadata on each spend-log row. + METADATA_HEADER = "x-litellm-spend-logs-metadata" + + #: Reply-size cap sent when the caller does not pick one. The gateway + #: rejects calls that omit max_tokens, and the autotester exposes no + #: model-options field, so this default is what keeps a plain config working. + DEFAULT_MAX_TOKENS = 1024 + + def __init__( + self, + remote_url: str = "http://localhost:4000/v1", + model_name: str = "gpt-4o-mini", + ) -> None: + """Initialize a client pointed at the LiteLLM gateway. + + Args: + remote_url: Base URL of the LiteLLM proxy's OpenAI-compatible API + (the ``/v1`` root). Supplied by the autotester via ``--remote_url``. + model_name: The model to request, e.g. ``gpt-4o-mini``. Must be one + of the models the gateway is configured to allow. + """ + # Bypass OpenAIModel.__init__ on purpose: it builds a client against + # api.openai.com using OPENAI_API_KEY, which is not how we authenticate + # to the gateway. We build our own client below. + super(OpenAIModel, self).__init__(model_name) + self.client = openai.OpenAI( + base_url=remote_url, + api_key=self._require_api_key(), + default_headers=self._attribution_headers(), + ) + + def _call_openai( + self, prompt: str, system_instructions: str, model_options: Optional[dict] = None, schema: Optional[dict] = None + ) -> str: + """Delegate to OpenAIModel with max_tokens defaulted; the gateway rejects calls without it.""" + model_options = dict(model_options or {}) + model_options.setdefault("max_tokens", self.DEFAULT_MAX_TOKENS) + try: + return super()._call_openai(prompt, system_instructions, model_options, schema) + except openai.APIError as exc: + raise GatewayError(_failure_reason(exc)) from exc + + def process_image(self, message: Message, args) -> str: + """Delegate to OpenAIModel, reporting gateway refusals the same way as text calls.""" + try: + return super().process_image(message, args) + except openai.APIError as exc: + raise GatewayError(_failure_reason(exc)) from exc + + @staticmethod + def _require_api_key() -> str: + """The LiteLLM virtual key sent as 'Authorization: Bearer'.""" + api_key = os.getenv("LITELLM_API_KEY") + if not api_key: + raise RuntimeError( + "LITELLM_API_KEY is not set. The gateway authenticates callers " + "with a LiteLLM virtual key sent as 'Authorization: Bearer'." + ) + return api_key + + @classmethod + def _attribution_headers(cls) -> dict: + """The x-litellm-spend-logs-metadata header, or {} when unset. + + Forwarded as-is — the autotester produces the JSON — but we fail loud on + malformed JSON rather than ship a broken header. + """ + metadata = os.getenv("LITELLM_SPEND_METADATA") + if not metadata: + return {} + try: + json.loads(metadata) + except (json.JSONDecodeError, TypeError) as exc: + raise RuntimeError( + "LITELLM_SPEND_METADATA is not valid JSON; refusing to send a malformed attribution header." + ) from exc + return {cls.METADATA_HEADER: metadata} + + @property + def spend_logs_metadata(self) -> Optional[str]: + """The attribution header value sent on every call, or None if unset.""" + return self.client.default_headers.get(self.METADATA_HEADER) diff --git a/ai_feedback/models/__init__.py b/ai_feedback/models/__init__.py index 476e9ab..0070d7d 100644 --- a/ai_feedback/models/__init__.py +++ b/ai_feedback/models/__init__.py @@ -6,6 +6,7 @@ from .Model import Model from .OpenAIModel import OpenAIModel from .OpenAIModelVector import OpenAIModelVector +from .OpenAIRemoteModel import GatewayError, OpenAIRemoteModel from .RemoteModel import RemoteModel @@ -14,6 +15,7 @@ class ModelFactory: _registry: Dict[str, Type[Model]] = { "remote": RemoteModel, + "openai-remote": OpenAIRemoteModel, "claude": ClaudeModel, "openai": OpenAIModel, "codellama": CodeLlamaModel, diff --git a/tests/test_main.py b/tests/test_main.py new file mode 100644 index 0000000..587621c --- /dev/null +++ b/tests/test_main.py @@ -0,0 +1,54 @@ +"""Tests for the CLI's handling of a gateway refusal. + +The autotester runs this package as a subprocess and reports a failed run from +its stderr, so a refused call must exit non-zero with the reason on stderr. +""" + +import sys + +import pytest + +from ai_feedback import code_processing +from ai_feedback.__main__ import main +from ai_feedback.models import GatewayError, ModelFactory + +REASON = "Course budget exhausted for course_id=1: spent CAD 0.01 of CAD 0.01." + + +@pytest.fixture +def cli(monkeypatch, tmp_path): + """Run the CLI against a stub model, returning a runner for the given processor.""" + submission = tmp_path / "submission.py" + submission.write_text("print('hi')\n") + monkeypatch.setattr(ModelFactory, "create", lambda *args, **kwargs: object()) + monkeypatch.setattr( + sys, + "argv", + [ + "ai_feedback", + "--scope", + "code", + "--submission", + str(submission), + "--provider", + "openai-remote", + "--prompt_text", + "Review this code.", + ], + ) + return main + + +def test_gateway_refusal_exits_with_the_reason_on_stderr(monkeypatch, capsys, cli): + def _refuse(*args, **kwargs): + raise GatewayError(REASON) + + monkeypatch.setattr(code_processing, "process_code", _refuse) + + with pytest.raises(SystemExit) as exited: + cli() + + assert exited.value.code == 1 + captured = capsys.readouterr() + assert REASON in captured.err + assert REASON not in captured.out diff --git a/tests/test_openai_remote_model.py b/tests/test_openai_remote_model.py new file mode 100644 index 0000000..8c7683e --- /dev/null +++ b/tests/test_openai_remote_model.py @@ -0,0 +1,194 @@ +"""Tests for OpenAIRemoteModel — the LiteLLM-gateway-backed OpenAI model. + +These tests never touch the network. They patch ``openai.OpenAI`` with a fake +client that records construction kwargs and returns a canned chat completion, +so we can assert the gateway endpoint, auth, and the attribution header without +a running proxy. +""" + +import json +import types + +import httpx +import openai +import pytest +from ollama import Message + +from ai_feedback.models import GatewayError, ModelFactory, OpenAIRemoteModel + +METADATA = { + "instance": "markus.cs.toronto.edu", + "course_id": 12, + "assignment_id": 34, + "group_id": 56, + "batch_id": None, + "category": "student", +} + + +class _FakeCompletions: + def __init__(self, content): + self._content = content + self.calls = [] + self.error = None # set to make create() fail the way the OpenAI SDK does + + def create(self, **kwargs): + if self.error: + raise self.error + self.calls.append(kwargs) + message = types.SimpleNamespace(content=self._content) + choice = types.SimpleNamespace(message=message) + return types.SimpleNamespace(choices=[choice]) + + +class _FakeClient: + """Stand-in for openai.OpenAI that records how it was built.""" + + last = None + + def __init__(self, *, base_url=None, api_key=None, default_headers=None, content="feedback"): + self.base_url = base_url + self.api_key = api_key + self.default_headers = default_headers or {} + self.completions = _FakeCompletions(content) + self.chat = types.SimpleNamespace(completions=self.completions) + _FakeClient.last = self + + +def _status_error(body): + """A real openai.APIStatusError carrying ``body``, as the SDK would raise it.""" + request = httpx.Request("POST", "http://gateway:4000/v1/chat/completions") + return openai.APIStatusError("Error code: 400", response=httpx.Response(400, request=request), body=body) + + +@pytest.fixture +def fake_openai(monkeypatch): + monkeypatch.setattr(openai, "OpenAI", _FakeClient) + return _FakeClient + + +@pytest.fixture(autouse=True) +def gateway_env(monkeypatch): + monkeypatch.setenv("LITELLM_API_KEY", "sk-test-virtual-key") + monkeypatch.delenv("LITELLM_SPEND_METADATA", raising=False) + + +def test_provider_is_registered(): + assert ModelFactory.is_registered("openai-remote") + assert ModelFactory.get_model_class("openai-remote") is OpenAIRemoteModel + + +def test_client_targets_gateway_with_bearer_auth(fake_openai): + OpenAIRemoteModel(remote_url="http://gateway:4000/v1", model_name="gpt-4o-mini") + assert fake_openai.last.base_url == "http://gateway:4000/v1" + assert fake_openai.last.api_key == "sk-test-virtual-key" + + +def test_attaches_metadata_header_when_set(monkeypatch, fake_openai): + monkeypatch.setenv("LITELLM_SPEND_METADATA", json.dumps(METADATA)) + model = OpenAIRemoteModel() + header = fake_openai.last.default_headers[OpenAIRemoteModel.METADATA_HEADER] + assert json.loads(header) == METADATA + assert model.spend_logs_metadata == header + + +def test_no_metadata_header_when_unset(fake_openai): + model = OpenAIRemoteModel() + assert OpenAIRemoteModel.METADATA_HEADER not in fake_openai.last.default_headers + assert model.spend_logs_metadata is None + + +def test_missing_api_key_fails_loud(monkeypatch, fake_openai): + monkeypatch.delenv("LITELLM_API_KEY", raising=False) + with pytest.raises(RuntimeError, match="LITELLM_API_KEY"): + OpenAIRemoteModel() + + +def test_malformed_metadata_fails_loud(monkeypatch, fake_openai): + monkeypatch.setenv("LITELLM_SPEND_METADATA", "{not valid json") + with pytest.raises(RuntimeError, match="not valid JSON"): + OpenAIRemoteModel() + + +def test_generate_response_speaks_openai_contract(fake_openai): + model = OpenAIRemoteModel(model_name="gpt-4o-mini") + prompt, response = model.generate_response( + prompt="Review this code.", + submission_file=None, + system_instructions="You are a TA.", + model_options={}, + ) + assert response == "feedback" + assert prompt == "Review this code." + sent = fake_openai.last.completions.calls[0] + assert sent["model"] == "gpt-4o-mini" + roles = [m["role"] for m in sent["messages"]] + assert roles == ["system", "user"] + + +def test_max_tokens_defaults_when_caller_omits_it(fake_openai): + model = OpenAIRemoteModel() + model.generate_response( + prompt="Review this code.", + submission_file=None, + system_instructions="You are a TA.", + model_options={}, + ) + sent = fake_openai.last.completions.calls[0] + assert sent["max_tokens"] == OpenAIRemoteModel.DEFAULT_MAX_TOKENS + + +def test_caller_max_tokens_wins_over_default(fake_openai): + model = OpenAIRemoteModel() + model.generate_response( + prompt="Review this code.", + submission_file=None, + system_instructions="You are a TA.", + model_options={"max_tokens": 2048}, + ) + sent = fake_openai.last.completions.calls[0] + assert sent["max_tokens"] == 2048 + + +def _failing_with(sdk_error): + """Make the gateway raise ``sdk_error``, and return the GatewayError it becomes.""" + model = OpenAIRemoteModel() + model.client.completions.error = sdk_error + with pytest.raises(GatewayError) as raised: + model.generate_response( + prompt="Review this code.", + submission_file=None, + system_instructions="You are a TA.", + model_options={}, + ) + return raised.value + + +def test_budget_rejection_surfaces_the_gateway_message(fake_openai): + """The instructor-facing reason, not a stack trace ending in BadRequestError.""" + reason = "Course budget exhausted for course_id=1: spent CAD 0.01 of CAD 0.01." + assert str(_failing_with(_status_error({"message": reason}))) == reason + + +def test_failure_without_our_body_falls_back_to_the_sdk_message(fake_openai): + """A proxy error page has no 'message' key; we must still say something.""" + assert "Error code: 400" in str(_failing_with(_status_error("502 Bad Gateway"))) + + +def test_unreachable_gateway_is_reported_the_same_way(fake_openai): + """A gateway restart mid-batch must not print a stack trace either.""" + request = httpx.Request("POST", "http://gateway:4000/v1/chat/completions") + assert "Connection error" in str(_failing_with(openai.APIConnectionError(request=request))) + + +def test_image_failure_is_reported_as_a_gateway_error(fake_openai): + """Image feedback goes through process_image, not _call_openai — same treatment.""" + model = OpenAIRemoteModel() + model.client.completions.error = _status_error({"message": "Course budget exhausted for course_id=1."}) + with pytest.raises(GatewayError, match="Course budget exhausted"): + model.process_image(Message(role="user", content="Describe this plot.", images=[]), args=None) + + +def test_failure_keeps_the_original_error_for_debugging(fake_openai): + error = _failing_with(_status_error({"message": "Upstream API key is disabled."})) + assert isinstance(error.__cause__, openai.APIStatusError)