feat: add openai api support and chunking

pyproject.toml

@@ -11,6 +11,7 @@ dependencies = [
     "google-cloud-core>=2.4.3",
     "google-genai>=1.59.0",
     "json-repair~=0.40.0",
+    "openai>=2.44.0",
     "pydantic>=2.11.7",
     "rapidfuzz>=3.13.0",
     "requests>=2.32.4",

src/genai_utils/chunking.py

@@ -0,0 +1,76 @@
+"""
+We frequently end up chunking input sentences when we do things with Gen AI.
+This has become replicated across multiple projects,
+so we made some laughably simple helper functions here.
+"""
+
+import math
+from typing import Iterator
+
+
+def calculate_n_chunks(n_sentences: int, chunk_size: int, overlap_size: int) -> int:
+    """
+    Calculates the number of chunks that ``make_chunks`` will produce.
+
+    Parameters
+    ----------
+    n_sentences: int
+        Total number of sentences.
+    chunk_size: int
+        Number of sentences per chunk.
+    overlap_size: int
+        Number of overlapping sentences between consecutive chunks.
+
+    Returns
+    -------
+    int
+        The number of chunks.
+    """
+    if n_sentences == 0:
+        return 0
+    # The first chunk that reaches the end of the input is the last one we
+    # emit; any further chunk would be a subset of it (see ``make_chunks``).
+    if n_sentences <= chunk_size:
+        return 1
+    step = chunk_size - overlap_size
+    return math.ceil((n_sentences - chunk_size) / step) + 1
+
+
+def make_chunks(
+    sentences: list[str], chunk_size: int, overlap_size: int
+) -> Iterator[list[str]]:
+    """
+    Takes a list of sentences, and yields smaller chunks of sentences,
+    with overlaps between chunks.
+
+    Parameters
+    ----------
+    text: list[str]
+        List of sentences to be chunked.
+    chunk_size: int
+        Number of sentences to be in each chunk.
+    overlap_size: int
+        Size of overlap of consecutive chunks.
+
+    Returns
+    -------
+    Iterator[list[str]]
+        Yields one chunk of sentences at a time.
+
+    Raises
+    ------
+    ValueError:
+        if the overlap is negative or the overlap is larger than the chunk size.
+    """
+    if not (0 <= overlap_size < chunk_size):
+        raise ValueError(
+            f"overlap_size ({overlap_size}) must be >= 0 and "
+            f"< chunk_size ({chunk_size})"
+        )
+    step = chunk_size - overlap_size
+    for i in range(0, len(sentences), step):
+        yield sentences[i : i + chunk_size]
+        # Once a chunk reaches the end, stop: any later chunk would start
+        # inside this one's overlap and be fully contained in it.
+        if i + chunk_size >= len(sentences):
+            break

src/genai_utils/openai_llm.py

@@ -0,0 +1,148 @@
+"""
+This file contains basic functionality for making requests to a hosted
+OpenAI API compatible LLM.
+"""
+
+import logging
+import time
+
+import openai
+from openai import OpenAI
+from tenacity import (
+    before_sleep_log,
+    retry,
+    retry_if_exception,
+    stop_after_attempt,
+    wait_exponential,
+)
+
+_logger = logging.getLogger(__name__)
+
+
+def _is_model_loading(exception: BaseException) -> bool:
+    # Servers (vLLM, llama.cpp) return 503 while up but not yet ready to
+    # serve — e.g. still loading the model, or a freshly-started replica.
+    # Retrying rides out that window; stop_after_attempt bounds the wait.
+    return (
+        isinstance(exception, openai.InternalServerError)
+        and exception.status_code == 503
+    )
+
+
+def log_tokens_per_second(response: object, elapsed: float) -> None:
+    """
+    Logs generation throughput (tokens/second) at info level.
+
+    Prefers the server's own decode timing when available (llama.cpp reports a
+    ``timings`` object with ``predicted_per_second``), otherwise falls back to
+    ``completion_tokens`` over the measured wall-clock time, which works for any
+    OpenAI-compatible server such as vLLM.
+    """
+    if not _logger.isEnabledFor(logging.INFO):
+        return
+
+    # llama.cpp: authoritative decode rate from the server, no network overhead.
+    # The openai client stashes unknown response fields on ``model_extra``.
+    extra = getattr(response, "model_extra", None) or {}
+    timings = extra.get("timings") if isinstance(extra, dict) else None
+    if isinstance(timings, dict) and timings.get("predicted_per_second") is not None:
+        _logger.info(
+            "Generation throughput: %.1f tok/s (%d tokens in %.2fs, server timings)",
+            timings["predicted_per_second"],
+            timings.get("predicted_n", 0),
+            timings.get("predicted_ms", 0) / 1000,
+        )
+        return
+
+    # vLLM / generic fallback: completion tokens over wall-clock time.
+    usage = getattr(response, "usage", None)
+    completion_tokens = getattr(usage, "completion_tokens", None)
+    if completion_tokens and elapsed > 0:
+        _logger.info(
+            "Generation throughput: %.1f tok/s (%d completion tokens in %.2fs)",
+            completion_tokens / elapsed,
+            completion_tokens,
+            elapsed,
+        )
+
+
+class NoOutputError(Exception):
+    pass
+
+
+class LLM:
+    """
+    An LLM using the Open AI API.
+    This might be something served via VLLM or Llama.cpp.
+    Mostly just wraps the methods for sending requests, etc.
+    If the model needs an API key, you will need to provide one.
+    """
+
+    def __init__(
+        self,
+        model_url: str,
+        model_name: str = "local-model",
+        api_key: str = "not-needed",
+    ):
+        """
+        Args:
+            model_url:
+                The url where the model is hosted.
+            model_name:
+                The name of the model.
+                This will have been set when you deployed the model.
+                The previous hardcoded name was "local-model"
+                so that behaviour remains as default.
+            api_key:
+                The API key for the model, if you need one.
+        """
+        self.model_url = model_url
+        self.model_name = model_name
+        self.client = OpenAI(
+            base_url=f"{model_url}/v1",
+            api_key=api_key,
+        )
+
+    @retry(
+        retry=retry_if_exception(_is_model_loading),
+        wait=wait_exponential(multiplier=2, min=5, max=30),
+        stop=stop_after_attempt(5),
+        before_sleep=before_sleep_log(_logger, logging.WARNING),
+        reraise=True,
+    )
+    def run_message(self, message_content: str, use_thinking: bool) -> str:
+        """
+        Sends the message to the client in the correct formatting.
+        Returns the result.
+        Will use thinking if you ask it to.
+        """
+        extra_model_config = (
+            {
+                "chat_template_kwargs": {"enable_thinking": False},
+                "reasoning_budget": 0,
+            }
+            if not use_thinking
+            else {
+                "chat_template_kwargs": {"enable_thinking": True},
+            }
+        )
+        # Ask llama.cpp to report per-request timings. Ignored by servers that
+        # don't support it (e.g. vLLM), so it's safe to always send.
+        extra_model_config["timings_per_token"] = True
+
+        start = time.perf_counter()
+        response = self.client.chat.completions.create(
+            model=self.model_name,
+            messages=[
+                {
+                    "role": "user",
+                    "content": str(message_content),
+                },
+            ],
+            extra_body=extra_model_config,
+        )
+        log_tokens_per_second(response, time.perf_counter() - start)
+        response_text = response.choices[0].message.content
+        if response_text is None:
+            raise NoOutputError()
+        return response_text

tests/genai_utils/test_chunking.py

@@ -0,0 +1,98 @@
+from pytest import mark, param, raises
+
+from genai_utils.chunking import (
+    calculate_n_chunks,
+    make_chunks,
+)
+
+
+@mark.parametrize(
+    "n_sentences,chunk_size,overlap_size,expected_chunk_count",
+    [
+        param(50, 100, 10, 1, id="short transcript with overlap"),
+        param(350, 100, 10, 4, id="long transcript with overlap"),
+        param(50, 100, 0, 1, id="short transcript without overlap"),
+        param(350, 100, 0, 4, id="long transcript without overlap"),
+        param(350, 100, 50, 6, id="long transcript with big overlap"),
+        param(10, 100, 15, 1, id="text shorter than overlap"),
+        param(50, 50, 0, 1, id="sentences same as chunk size"),
+        param(20, 50, 0, 1, id="sentences shorter than chunk size"),
+        param(20, 19, 0, 2, id="sentences one longer than chunk size"),
+        param(0, 100, 0, 0, id="no sentences"),
+    ],
+)
+def test_calculate_n_chunks(
+    n_sentences, chunk_size, overlap_size, expected_chunk_count
+) -> None:
+    assert (
+        calculate_n_chunks(n_sentences, chunk_size, overlap_size)
+        == expected_chunk_count
+    )
+
+
+@mark.parametrize(
+    "n_sentences,chunk_size,overlap_size",
+    [
+        param(50, 100, 10, id="short transcript with overlap"),
+        param(350, 100, 10, id="long transcript with overlap"),
+        param(350, 100, 50, id="long transcript with big overlap"),
+        param(20, 19, 0, id="sentences one longer than chunk size"),
+        param(0, 100, 0, id="no sentences"),
+    ],
+)
+def test_make_chunks_count_matches_calculation(
+    n_sentences, chunk_size, overlap_size
+) -> None:
+    """``make_chunks`` should yield exactly as many chunks as predicted."""
+    sentences = [f"sentence {i}" for i in range(n_sentences)]
+    chunks = list(make_chunks(sentences, chunk_size, overlap_size))
+    assert len(chunks) == calculate_n_chunks(n_sentences, chunk_size, overlap_size)
+
+
+def test_make_chunks_overlap_content() -> None:
+    """Consecutive chunks should share ``overlap_size`` sentences."""
+    sentences = [f"s{i}" for i in range(10)]
+    chunks = list(make_chunks(sentences, chunk_size=4, overlap_size=2))
+
+    # first chunk is the start of the input
+    assert chunks[0] == ["s0", "s1", "s2", "s3"]
+    # the tail of one chunk is the head of the next
+    assert chunks[0][-2:] == chunks[1][:2]
+    assert chunks[1] == ["s2", "s3", "s4", "s5"]
+
+
+def test_make_chunks_no_overlap_partitions_input() -> None:
+    """With no overlap the chunks should join back into the original list."""
+    sentences = [f"s{i}" for i in range(10)]
+    chunks = list(make_chunks(sentences, chunk_size=3, overlap_size=0))
+    flattened = [sentence for chunk in chunks for sentence in chunk]
+    assert flattened == sentences
+
+
+def test_make_chunks_no_redundant_trailing_chunk() -> None:
+    """A final chunk fully contained in the previous one should not be emitted."""
+    sentences = [f"s{i}" for i in range(10)]
+    # chunk_size=4, overlap=2, step=2 would step to index 8 and yield a final
+    # ["s8", "s9"] that is already wholly inside the previous ["s6"..."s9"] chunk.
+    chunks = list(make_chunks(sentences, chunk_size=4, overlap_size=2))
+    assert chunks[-1] == ["s6", "s7", "s8", "s9"]
+    # every chunk carries new content the previous one didn't have
+    assert all(chunks[i] != chunks[i - 1] for i in range(1, len(chunks)))
+
+
+def test_make_chunks_empty_input() -> None:
+    assert list(make_chunks([], chunk_size=5, overlap_size=2)) == []
+
+
+@mark.parametrize(
+    "chunk_size,overlap_size",
+    [
+        param(10, 15, id="overlap bigger than chunk"),
+        param(10, 10, id="overlap same as chunk"),
+        param(100, -1, id="negative overlap"),
+    ],
+)
+def test_make_chunks_invalid_overlap_raises(chunk_size, overlap_size) -> None:
+    with raises(ValueError):
+        # wrap in list to force the generator to run
+        list(make_chunks(["x"] * 50, chunk_size, overlap_size))