radiology_rag/embedding.py

"""Embedding utilities (OpenAI-compatible, API-first)."""

from __future__ import annotations

from dataclasses import dataclass

from langchain_openai import OpenAIEmbeddings


@dataclass(frozen=True)
class EmbeddingConfig:
    base_url: str
    api_key: str
    model_name: str
    batch_size: int = 32


class EmbeddingClient:
    """Thin wrapper over LangChain OpenAIEmbeddings."""

    def __init__(self, cfg: EmbeddingConfig):
        self.cfg = cfg
        self._emb = OpenAIEmbeddings(
            base_url=cfg.base_url,
            api_key=cfg.api_key,
            model=cfg.model_name,
            chunk_size=int(cfg.batch_size or 32),
        )

    def embed_query(self, text: str) -> list[float]:
        return self._emb.embed_query(text)

    def embed_documents(self, texts: list[str]) -> list[list[float]]:
        return self._emb.embed_documents(texts)

    @property
    def langchain_embeddings(self) -> OpenAIEmbeddings:
        """Expose the underlying LangChain embeddings for Chroma."""
        return self._emb