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

ai_feedback/models/OpenAIRemoteModel.py

@@ -0,0 +1,102 @@
+import json
+import os
+from typing import Optional
+
+import openai
+
+from .OpenAIModel import OpenAIModel
+
+
+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)
+        return super()._call_openai(prompt, system_instructions, model_options, schema)
+
+    @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 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_openai_remote_model.py

@@ -0,0 +1,139 @@
+"""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 openai
+import pytest
+
+from ai_feedback.models import 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 = []
+
+    def create(self, **kwargs):
+        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
+
+
+@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