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 ContentBlockError, OpenRouterLlmCompletionClient
+
+__all__ = ["Completion", "ContentBlockError", "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,140 @@
+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 ContentBlockError(Exception):
+    """This exception is raised when the completion was terminated due to a content block."""
+
+    def __init__(self, messages: List[Dict[str, str]]) -> None:
+        super().__init__("Completion response terminated because of content filtering.")
+        self.messages = messages
+
+
+class OpenRouterLlmCompletionClient(LlmCompletionClient):
+    def __init__(
+        self,
+        model: str,
+        max_tokens: Optional[int] = None,
+        api_key: Optional[str] = None,
+        **default_kwargs: Any,
+    ) -> None:
+        self._model = model
+        self._max_tokens = max_tokens
+        self._default_kwargs = default_kwargs
+        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, **kwargs: Any) -> Completion:
+        if self._closed:
+            raise RuntimeError("OpenRouter completion client is closed.")
+        try:
+            from openrouter.components import (
+                ChatAssistantMessage,
+                ChatAssistantMessageContent,
+                ChatChoice,
+                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 = dict(self._default_kwargs)
+        request.update(kwargs)
+        request.update(
+            {
+                "model": self._model,
+                "messages": messages,
+                "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)
+        choices = getattr(response, "choices", None)
+        if not choices:
+            raise ValueError("Malformed OpenRouter completion response: No choices were returned.")
+        choice: Optional[ChatChoice] = choices[0]
+        if choice is None:
+            raise ValueError("Malformed OpenRouter completion response: The first choice is null.")
+        if 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)
+        elif getattr(choice, "finish_reason", None) == "content_filter":
+            raise ContentBlockError(messages)
+        message: Optional[ChatAssistantMessage] = getattr(choice, "message", None)
+        if message is None:
+            raise ValueError("Malformed OpenRouter completion response: The first choice has no message.")
+        content: Optional[ChatAssistantMessageContent] = getattr(message, "content", None)
+        if content is None:
+            raise ValueError("Malformed OpenRouter completion response: The message has no content.")
+        if isinstance(content, str):
+            text = content
+        elif isinstance(content, list):
+            text_parts = [p.text for p in content if isinstance(p, ChatContentText) and isinstance(p.text, str)]
+            if not text_parts:
+                raise ValueError("Malformed OpenRouter completion response: The message has no text content.")
+            text = "".join(text_parts)
+        else:
+            raise ValueError("Malformed OpenRouter completion response: The message content is not 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"]