"""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.""" return await _check_endpoint( client, base_url=config.embedding_base_url, api_key=config.embedding_api_key, timeout_seconds=timeout_seconds, unavailable_log=f"ebook_embedding_endpoint_unreachable {config.embedding_base_url=}", ) 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.""" return await _check_endpoint( client, base_url=config.vllm_base_url, api_key=config.vllm_api_key, timeout_seconds=timeout_seconds, unavailable_log=f"ebook_chat_endpoint_unreachable {config.vllm_base_url=}", ) async def _check_endpoint( client: httpx.AsyncClient, *, base_url: str, api_key: str, timeout_seconds: float, unavailable_log: str, ) -> bool: """Return whether an OpenAI-compatible endpoint answers a model listing.""" try: response = await client.get( f"{base_url.rstrip('/')}/models", headers=auth_headers(api_key), timeout=timeout_seconds, ) response.raise_for_status() except httpx.HTTPError as error: logger.warning(f"{unavailable_log} {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]], *, response_format: dict[str, object] | None = None, ) -> 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. response_format (dict[str, object] | None): Optional OpenAI-compatible structured output constraint. 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} | ({"response_format": response_format} if response_format is not None else {}), 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