TICKET-611: Add openai-remote provider for the ai-telemetry-gateway

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")

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)

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,

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