LLM Client

machine/translation/llm/__init__.py

@@ -0,0 +1,4 @@
+from .llm_completion_client import Completion, LlmCompletionClient, TokenLogprob
+from .openrouter_llm_completion_client import OpenRouterLlmCompletionClient
+
+__all__ = ["Completion", "LlmCompletionClient", "OpenRouterLlmCompletionClient", "TokenLogprob"]

machine/translation/llm/llm_completion_client.py

@@ -0,0 +1,45 @@
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from dataclasses import dataclass, field
+from types import TracebackType
+from typing import Dict, List, Optional, Type
+
+
+@dataclass(frozen=True)
+class TokenLogprob:
+    token: str
+    logprob: float
+
+
+@dataclass(frozen=True)
+class Completion:
+    """A reply from the model, with whatever the provider reported alongside it."""
+
+    text: str
+    cost: float
+    token_logprobs: List[TokenLogprob] = field(default_factory=list)
+    input_token_count: int = 0
+    output_token_count: int = 0
+
+
+class LlmCompletionClient(ABC):
+    @abstractmethod
+    def complete(self, messages: List[Dict[str, str]], logprobs: bool = False) -> Completion: ...
+
+    def supports_logprobs(self) -> bool:
+        return False
+
+    def count_tokens(self, text: str) -> Optional[int]:
+        return None
+
+    def __enter__(self) -> LlmCompletionClient:
+        return self
+
+    def __exit__(
+        self,
+        exc_type: Optional[Type[BaseException]],
+        exc_value: Optional[BaseException],
+        traceback: Optional[TracebackType],
+    ) -> None:
+        pass

machine/translation/llm/openrouter_llm_completion_client.py

@@ -0,0 +1,114 @@
+import logging
+from types import TracebackType
+from typing import Any, Dict, List, Optional, Type
+
+from .llm_completion_client import Completion, LlmCompletionClient, TokenLogprob
+
+logger = logging.getLogger(__name__)
+
+
+class OpenRouterLlmCompletionClient(LlmCompletionClient):
+    def __init__(
+        self,
+        model: str,
+        temperature: float = 0.0,
+        max_tokens: Optional[int] = None,
+        api_key: Optional[str] = None,
+    ) -> None:
+        self._model = model
+        self._temperature = temperature
+        self._max_tokens = max_tokens
+        self._logprobs_supported: Optional[bool] = None
+        self._closed = False
+        try:
+            from openrouter import OpenRouter
+        except ImportError as error:
+            raise ImportError(
+                "OpenRouter completions require the 'openrouter' package. Install it with `uv sync --extra llm`."
+            ) from error
+        self._client = OpenRouter(api_key=api_key)
+
+    def complete(self, messages: List[Dict[str, str]], logprobs: bool = False) -> Completion:
+        if self._closed:
+            raise RuntimeError("OpenRouter completion client is closed.")
+        try:
+            from openrouter.components import (
+                ChatAssistantMessage,
+                ChatAssistantMessageContent,
+                ChatContentText,
+                ChatResult,
+                ChatUsage,
+            )
+        except ImportError as error:
+            raise ImportError(
+                "OpenRouter completions require the 'openrouter' package. Install it with `uv sync --extra llm`."
+            ) from error
+        request = {
+            "model": self._model,
+            "messages": messages,
+            "temperature": self._temperature,
+            "logprobs": logprobs,
+            "stream": False,
+        }
+        if self._max_tokens is not None:
+            request["max_completion_tokens"] = self._max_tokens
+        response: ChatResult = self._client.chat.send(**request)
+        choice = response.choices[0] if response.choices else None
+        if choice is not None and getattr(choice, "finish_reason", None) == "length" and self._max_tokens is not None:
+            logger.warning("Completion was finished because the max tokens ceiling %d was hit", self._max_tokens)
+        message: Optional[ChatAssistantMessage] = getattr(choice, "message", None)
+        content: Optional[ChatAssistantMessageContent] = getattr(message, "content", None)
+        if isinstance(content, str):
+            text = content
+        elif isinstance(content, list):
+            text = "".join(p.text for p in content if isinstance(p, ChatContentText))
+        else:
+            text = ""
+        token_logprobs = self._get_logprobs(choice) if logprobs else []
+        usage: Optional[ChatUsage] = getattr(response, "usage", None)
+        cost: Optional[float] = getattr(usage, "cost", None)
+        prompt_tokens: int = getattr(usage, "prompt_tokens", 0)
+        completion_tokens: int = getattr(usage, "completion_tokens", 0)
+        return Completion(
+            text=text,
+            cost=cost or 0.0,
+            token_logprobs=token_logprobs,
+            input_token_count=prompt_tokens,
+            output_token_count=completion_tokens,
+        )
+
+    def supports_logprobs(self) -> bool:
+        if self._closed:
+            raise RuntimeError("OpenRouter completion client is closed.")
+        if self._logprobs_supported is None:
+            try:
+                author, separator, slug = self._model.partition("/")
+                if not separator or not author or not slug:
+                    return False
+                model = self._client.models.get(author=author, slug=slug).data
+                self._logprobs_supported = "logprobs" in model.supported_parameters
+            except Exception:
+                logger.warning(
+                    "Could not determine whether OpenRouter model '%s' supports logprobs.", self._model, exc_info=True
+                )
+        return self._logprobs_supported or False
+
+    def __enter__(self) -> "OpenRouterLlmCompletionClient":
+        if self._closed:
+            raise RuntimeError("OpenRouter completion client is closed.")
+        self._client.__enter__()
+        return self
+
+    def __exit__(
+        self,
+        exc_type: Optional[Type[BaseException]],
+        exc_value: Optional[BaseException],
+        traceback: Optional[TracebackType],
+    ) -> None:
+        self._client.__exit__(exc_type, exc_value, traceback)
+        self._closed = True
+
+    def _get_logprobs(self, choice: Any) -> List[TokenLogprob]:
+        logprobs = getattr(choice, "logprobs", None)
+        entries = getattr(logprobs, "content", None) or []
+        return [TokenLogprob(str(entry.token), float(entry.logprob)) for entry in entries]

machine/translation/translation_engine.py

@@ -1,13 +1,13 @@
 from __future__ import annotations
 
-from abc import abstractmethod
+from abc import ABC, abstractmethod
 from types import TracebackType
 from typing import Optional, Sequence, Type, Union
 
 from .translation_result import TranslationResult
 
 
-class TranslationEngine:
+class TranslationEngine(ABC):
     @abstractmethod
     def translate(self, segment: Union[str, Sequence[str]]) -> TranslationResult: ...
 

pyproject.toml

@@ -54,6 +54,7 @@ jobs = [
     "dynaconf>=3.2.5,<4.0",
     "pydantic>=2.13.4,<3.0",
 ]
+llm = ["openrouter>=1.3.95"]
 
 [build-system]
 requires = ["uv_build>=0.12.17,<0.13"]

tests/translation/llm/test_openrouter_llm_completion_client.py

@@ -0,0 +1,160 @@
+from contextlib import contextmanager
+from types import SimpleNamespace
+from unittest.mock import patch
+
+import pytest
+
+from machine.translation.llm import OpenRouterLlmCompletionClient
+from machine.translation.llm.llm_completion_client import Completion, TokenLogprob
+
+
+class _FakeChat:
+    def __init__(self, response):
+        self.response = response
+        self.calls = []
+
+    def send(self, **kwargs):
+        self.calls.append(kwargs)
+        return self.response
+
+
+class _FakeModels:
+    def __init__(self, supported_parameters, error=None):
+        self.supported_parameters = supported_parameters
+        self.error = error
+        self.calls = []
+
+    def get(self, *, author, slug):
+        self.calls.append((author, slug))
+        if self.error is not None:
+            raise self.error
+        return SimpleNamespace(data=SimpleNamespace(supported_parameters=self.supported_parameters))
+
+
+class _FakeOpenRouter:
+    def __init__(self, response, supported_parameters=None, model_error=None):
+        self.chat = _FakeChat(response)
+        self.models = _FakeModels(supported_parameters or [], model_error)
+        self.enter_calls = 0
+        self.exit_calls = 0
+
+    def __enter__(self):
+        self.enter_calls += 1
+        return self
+
+    def __exit__(self, exc_type, exc_value, traceback):
+        self.exit_calls += 1
+
+
+@contextmanager
+def _client(response, supported_parameters=None, model_error=None):
+    sdk_client = _FakeOpenRouter(response, supported_parameters, model_error)
+    with patch("openrouter.OpenRouter", return_value=sdk_client):
+        yield sdk_client
+
+
+def test_openrouter_client_sends_chat_request_and_maps_completion_fields():
+    response = SimpleNamespace(
+        choices=[SimpleNamespace(message=SimpleNamespace(content="hola"))],
+        usage=SimpleNamespace(cost=0.012, prompt_tokens=8, completion_tokens=2),
+    )
+    with _client(response) as sdk_client:
+        with OpenRouterLlmCompletionClient(
+            "openai/gpt-4o", temperature=0.2, max_tokens=50, api_key="test-key"
+        ) as completion_client:
+            result = completion_client.complete([{"role": "user", "content": "hello"}])
+
+    assert result == Completion("hola", 0.012, input_token_count=8, output_token_count=2)
+    assert sdk_client.chat.calls == [
+        {
+            "model": "openai/gpt-4o",
+            "messages": [{"role": "user", "content": "hello"}],
+            "temperature": 0.2,
+            "max_completion_tokens": 50,
+            "logprobs": False,
+            "stream": False,
+        }
+    ]
+    assert sdk_client.enter_calls == 1
+    assert sdk_client.exit_calls == 1
+
+
+def test_openrouter_client_max_tokens_is_none():
+    response = SimpleNamespace(choices=[], usage=None)
+    with _client(response) as sdk_client:
+        with OpenRouterLlmCompletionClient("openai/gpt-4o", api_key="test-key") as client:
+            client.complete([])
+
+    assert sdk_client.chat.calls == [
+        {
+            "model": "openai/gpt-4o",
+            "messages": [],
+            "temperature": 0.0,
+            "logprobs": False,
+            "stream": False,
+        }
+    ]
+
+
+def test_openrouter_client_fails_after_client_is_closed():
+    response = SimpleNamespace(choices=[], usage=None)
+    with _client(response):
+        client = OpenRouterLlmCompletionClient("openai/gpt-4o", api_key="test-key")
+        with client:
+            pass
+
+        with pytest.raises(RuntimeError, match="closed"):
+            client.complete([])
+
+        with pytest.raises(RuntimeError, match="closed"):
+            with client:
+                client.complete([])
+
+
+def test_openrouter_client_maps_requested_token_logprobs():
+    response = SimpleNamespace(
+        choices=[
+            SimpleNamespace(
+                message=SimpleNamespace(content="hola"),
+                logprobs=SimpleNamespace(content=[SimpleNamespace(token="hola", logprob=-0.2)]),
+            )
+        ],
+        usage=None,
+    )
+    with _client(response):
+        with OpenRouterLlmCompletionClient("openai/gpt-4o", api_key="test-key") as client:
+            result = client.complete([], logprobs=True)
+
+    assert result == Completion("hola", 0.0, [TokenLogprob("hola", -0.2)])
+
+
+def test_openrouter_client_reports_logprobs_supported_by_its_model_and_caches_lookup():
+    response = SimpleNamespace(choices=[], usage=None)
+    with _client(response, ["temperature", "logprobs"]) as sdk_client:
+        with OpenRouterLlmCompletionClient("openai/gpt-4o", api_key="test-key") as client:
+            assert client.supports_logprobs()
+            assert client.supports_logprobs()
+
+    assert sdk_client.models.calls == [("openai", "gpt-4o")]
+
+
+def test_openrouter_client_reports_logprobs_unsupported_by_its_model():
+    response = SimpleNamespace(choices=[], usage=None)
+    with _client(response, ["temperature"]):
+        with OpenRouterLlmCompletionClient("openai/gpt-4o", api_key="test-key") as client:
+            assert not client.supports_logprobs()
+
+
+def test_openrouter_client_reports_logprobs_unsupported_when_model_lookup_fails():
+    response = SimpleNamespace(choices=[], usage=None)
+    with _client(response, model_error=RuntimeError("metadata unavailable")):
+        with OpenRouterLlmCompletionClient("openai/gpt-4o", api_key="test-key") as client:
+            assert not client.supports_logprobs()
+
+
+def test_openrouter_client_reports_logprobs_unsupported_for_unqualified_model_ids():
+    response = SimpleNamespace(choices=[], usage=None)
+    with _client(response, ["logprobs"]) as sdk_client:
+        with OpenRouterLlmCompletionClient("gpt-4o", api_key="test-key") as client:
+            assert not client.supports_logprobs()
+            assert sdk_client.models.calls == []