dotfiles/python/ebook_search/llm_interface.py

"""LLM provider HTTP adapters."""

from __future__ import annotations

import logging
from typing import TYPE_CHECKING

import httpx

if TYPE_CHECKING:
    from collections.abc import Sequence

    from python.ebook_search.config import EbookSearchConfig, RerankConfig

logger = logging.getLogger(__name__)


def auth_headers(api_key: str) -> dict[str, str]:
    """Build authorization headers when an API key is configured."""
    if api_key == "not-needed":
        return {}
    return {"Authorization": f"Bearer {api_key}"}


async def request_embeddings(
    client: httpx.AsyncClient,
    texts: Sequence[str],
    config: EbookSearchConfig,
) -> list[list[float]]:
    """Request embeddings from the configured OpenAI-compatible endpoint.

    Args:
        client (httpx.AsyncClient): Shared async client for LLM calls.
        texts (Sequence[str]): Texts to embed.
        config (EbookSearchConfig): Runtime settings supplying the endpoint, model, and auth.

    Returns:
        list[list[float]]: One embedding vector per input text.

    Raises:
        RuntimeError: If the request fails or the response cannot be parsed.
    """
    try:
        response = await client.post(
            f"{config.embedding_base_url.rstrip('/')}/embeddings",
            headers=auth_headers(config.embedding_api_key),
            json={"model": config.embedding_model, "input": list(texts)},
            timeout=config.embedding_timeout_seconds,
        )
        response.raise_for_status()
        return embedding_vectors_from_response(response.json())
    except (httpx.HTTPError, ValueError, KeyError, TypeError) as error:
        logger.exception(
            f"ebook_embed_request_failed {config.embedding_base_url=} {config.embedding_model=} count={len(texts)}"
        )
        msg = f"Embedding request failed. base_url={config.embedding_base_url} model={config.embedding_model}"
        raise RuntimeError(msg) from error


async def check_embedding_endpoint(
    client: httpx.AsyncClient,
    config: EbookSearchConfig,
    *,
    timeout_seconds: float = 5.0,
) -> bool:
    """Return whether the configured embedding endpoint answers a model listing."""
    try:
        response = await client.get(
            f"{config.embedding_base_url.rstrip('/')}/models",
            headers=auth_headers(config.embedding_api_key),
            timeout=timeout_seconds,
        )
        response.raise_for_status()
    except httpx.HTTPError as error:
        logger.warning(f"ebook_embedding_endpoint_unreachable {config.embedding_base_url=} {error=}")
        return False
    return True


async def check_chat_endpoint(
    client: httpx.AsyncClient,
    config: EbookSearchConfig,
    *,
    timeout_seconds: float = 5.0,
) -> bool:
    """Return whether the configured chat (answering) endpoint answers a model listing."""
    try:
        response = await client.get(
            f"{config.vllm_base_url.rstrip('/')}/models",
            headers=auth_headers(config.vllm_api_key),
            timeout=timeout_seconds,
        )
        response.raise_for_status()
    except httpx.HTTPError as error:
        logger.warning(f"ebook_chat_endpoint_unreachable {config.vllm_base_url=} {error=}")
        return False
    return True


def embedding_vectors_from_response(body: object) -> list[list[float]]:
    """Extract embedding vectors from an OpenAI-compatible embedding response."""
    if not isinstance(body, dict):
        msg = "Embedding response is not an object"
        raise TypeError(msg)

    data = body["data"]
    if not isinstance(data, list):
        msg = "Embedding response data is not a list"
        raise TypeError(msg)

    vectors: list[list[float]] = []
    for item in data:
        if not isinstance(item, dict):
            msg = "Embedding item is not an object"
            raise TypeError(msg)
        embedding = item["embedding"]
        if not isinstance(embedding, list):
            msg = "Embedding value is not a list"
            raise TypeError(msg)
        vectors.append([float(value) for value in embedding])
    return vectors


async def request_rerank(
    client: httpx.AsyncClient,
    query: str,
    documents: Sequence[str],
    config: RerankConfig,
) -> object | None:
    """Request rerank scores from the configured vLLM endpoint.

    Args:
        client (httpx.AsyncClient): Shared async client for LLM calls.
        query (str): Query the documents are scored against.
        documents (Sequence[str]): Candidate documents to score.
        config (RerankConfig): Rerank endpoint settings.

    Returns:
        object | None: The decoded response body, or ``None`` when it is not valid JSON.
    """
    payload = {
        "model": config.model,
        "query": query,
        "documents": list(documents),
    }
    response = await client.post(
        f"{config.base_url.rstrip('/')}/rerank",
        json=payload,
        timeout=config.timeout_seconds,
    )
    response.raise_for_status()
    try:
        return response.json()
    except ValueError:
        logger.debug("ebook_rerank_response_invalid_json", extra={"response": response.text})
        return None


async def request_chat_completion(
    client: httpx.AsyncClient,
    config: EbookSearchConfig,
    messages: Sequence[dict[str, str]],
) -> str:
    """Request a chat completion over a shared async client.

    Args:
        client (httpx.AsyncClient): Shared async client whose connection pool bounds concurrency.
        config (EbookSearchConfig): Runtime settings supplying the endpoint, model, and auth.
        messages (Sequence[dict[str, str]]): OpenAI-style chat messages.

    Returns:
        str: The assistant message text.

    Raises:
        RuntimeError: If the request fails or the response cannot be parsed.
    """
    try:
        response = await client.post(
            f"{config.vllm_base_url.rstrip('/')}/chat/completions",
            headers=auth_headers(config.vllm_api_key),
            json={
                "model": config.chat_model,
                "messages": list(messages),
                "temperature": 0,
            },
            timeout=config.chat_timeout_seconds,
        )
        response.raise_for_status()
        return chat_content_from_response(response.json())
    except (httpx.HTTPError, ValueError, KeyError, TypeError) as error:
        msg = f"Chat request failed. base_url={config.vllm_base_url} model={config.chat_model}"
        raise RuntimeError(msg) from error


def chat_content_from_response(body: object) -> str:
    """Extract text content from an OpenAI-compatible chat response."""
    if not isinstance(body, dict):
        msg = "Chat response is not an object"
        raise TypeError(msg)

    choices = body["choices"]
    if not isinstance(choices, list) or not choices:
        msg = "Chat response has no choices"
        raise ValueError(msg)

    first = choices[0]
    if not isinstance(first, dict):
        msg = "Chat choice is not an object"
        raise TypeError(msg)

    message = first["message"]
    if not isinstance(message, dict):
        msg = "Chat message is not an object"
        raise TypeError(msg)

    content = message.get("content") or ""
    if not isinstance(content, str):
        msg = "Chat content is not text"
        raise TypeError(msg)
    return content