From 5f3a731db69b75ca6e0708ed7398cba59ee66299 Mon Sep 17 00:00:00 2001 From: Nico Loesch Date: Wed, 22 Jul 2026 04:09:31 +0000 Subject: [PATCH 1/2] Remove config calls from internal functions and expose only at interface boundary --- src/omop_emb/backends/__init__.py | 3 +- src/omop_emb/backends/base_backend.py | 47 ++++--- src/omop_emb/cli/cli_diagnostics.py | 7 +- src/omop_emb/cli/cli_embeddings.py | 59 ++++---- src/omop_emb/cli/cli_legacy.py | 7 +- src/omop_emb/cli/cli_maintenance.py | 24 ++-- src/omop_emb/config.py | 29 +++- src/omop_emb/embeddings/embedding_client.py | 82 +++++------- tests/test_embedding_client.py | 141 ++++++++++---------- 9 files changed, 220 insertions(+), 179 deletions(-) diff --git a/src/omop_emb/backends/__init__.py b/src/omop_emb/backends/__init__.py index 65b752c..662962e 100644 --- a/src/omop_emb/backends/__init__.py +++ b/src/omop_emb/backends/__init__.py @@ -2,7 +2,7 @@ from typing import TYPE_CHECKING -from .base_backend import EmbeddingBackend, resolve_backend +from .base_backend import EmbeddingBackend, resolve_backend, resolve_backend_from_config from .sqlitevec import SQLiteVecEmbeddingBackend if TYPE_CHECKING: @@ -11,6 +11,7 @@ __all__ = [ "EmbeddingBackend", "resolve_backend", + "resolve_backend_from_config", "SQLiteVecEmbeddingBackend", "PGVectorEmbeddingBackend", ] diff --git a/src/omop_emb/backends/base_backend.py b/src/omop_emb/backends/base_backend.py index 67c43cf..a5d4fc9 100644 --- a/src/omop_emb/backends/base_backend.py +++ b/src/omop_emb/backends/base_backend.py @@ -993,46 +993,45 @@ def validate_embeddings_and_records( def resolve_backend( - backend_type: Optional[Union[str, BackendType]] = None, + backend_type: str | BackendType, + *, + sqlite_path: Optional[str] = None, ) -> EmbeddingBackend: - """Return the configured embedding backend via oa-configurator. + """Return the embedding backend for *backend_type*. - When backend_type is omitted, reads from the active oa-configurator config. Connection details are resolved via the oa-configurator Resolver: - - ``sqlitevec``: uses sqlite_path from config (required; must be set explicitly). + - ``sqlitevec``: uses *sqlite_path* (required; must be set explicitly). - ``pgvector``: uses the ``emb_db`` resource from oa-configurator. + + Callers that have a loaded ``OmopEmbConfig`` should prefer + :func:`resolve_backend_from_config` over calling this directly. """ - cfg = OmopEmbConfig.get_config() - if backend_type is None: - backend_str = cfg.backend - else: - backend_str = ( - backend_type if isinstance(backend_type, str) else backend_type.value - ) + backend_str = ( + backend_type if isinstance(backend_type, str) else backend_type.value + ) try: - resolved_backend = BackendType(backend_str.lower()) + resolved_backend_type = BackendType(backend_str.lower()) except ValueError: raise RuntimeError( f"Unknown backend {backend_str!r}. " f"Supported: {[b.value for b in BackendType]}." ) - if resolved_backend == BackendType.SQLITEVEC: + if resolved_backend_type == BackendType.SQLITEVEC: from omop_emb.backends.sqlitevec import SQLiteVecEmbeddingBackend - if not cfg.sqlite_path: + if not sqlite_path: raise RuntimeError( "sqlitevec backend requires 'sqlite_path' to be configured. " "Set it via `omop-config configure omop-emb`. " "To use an ephemeral in-memory store intentionally, set sqlite_path = ':memory:'." ) - path = cfg.sqlite_path - logger.info(f"Using SQLiteVec backend with database file: {path}") - return SQLiteVecEmbeddingBackend.from_path(path) + logger.info(f"Using SQLiteVec backend with database file: {sqlite_path}") + return SQLiteVecEmbeddingBackend.from_path(sqlite_path) - if resolved_backend == BackendType.PGVECTOR: + if resolved_backend_type == BackendType.PGVECTOR: engine = resolve_omop_emb_engine() if engine.dialect.name != "postgresql": raise RuntimeError( @@ -1049,4 +1048,14 @@ def resolve_backend( logger.info(f"Using pgvector backend with engine: {engine.url}") return PGVectorEmbeddingBackend(emb_engine=engine) - raise RuntimeError(f"Implementation for {resolved_backend.value} is not available.") + raise RuntimeError(f"Implementation for {resolved_backend_type.value} is not available.") + + +def resolve_backend_from_config(cfg: OmopEmbConfig) -> EmbeddingBackend: + """Resolve the embedding backend using an already-loaded config object. + + The one place config is read for backend selection; call sites that + already have *cfg* in scope should prefer this over calling + :func:`resolve_backend` directly. + """ + return resolve_backend(cfg.backend, sqlite_path=cfg.sqlite_path) diff --git a/src/omop_emb/cli/cli_diagnostics.py b/src/omop_emb/cli/cli_diagnostics.py index dc72803..00af6fd 100644 --- a/src/omop_emb/cli/cli_diagnostics.py +++ b/src/omop_emb/cli/cli_diagnostics.py @@ -5,8 +5,8 @@ import sqlalchemy as sa import typer -from omop_emb.backends import resolve_backend -from omop_emb.config import MetricType, resolve_omop_cdm_engine +from omop_emb.backends import resolve_backend_from_config +from omop_emb.config import MetricType, load_omop_emb_config, resolve_omop_cdm_engine from omop_emb.interface import list_registered_models logger = logging.getLogger(__name__) @@ -17,7 +17,8 @@ name="health-check", help="Verify backend connectivity and list registered models." ) def health_check(): - backend = resolve_backend() + cfg = load_omop_emb_config() + backend = resolve_backend_from_config(cfg) typer.echo(f"Backend: {backend.backend_type.value} | connected.") # CDM connectivity is optional for the health check diff --git a/src/omop_emb/cli/cli_embeddings.py b/src/omop_emb/cli/cli_embeddings.py index 47e2eea..c8ba91b 100644 --- a/src/omop_emb/cli/cli_embeddings.py +++ b/src/omop_emb/cli/cli_embeddings.py @@ -9,12 +9,13 @@ from omop_emb.utils.cdm import check_concept_cdm from omop_emb.backends.index_config import index_config_from_index_type -from omop_emb.backends import resolve_backend +from omop_emb.backends import resolve_backend_from_config from omop_emb.config import ( IndexType, MetricType, OmopEmbConfig, ProviderType, + load_omop_emb_config, provider_type_examples, resolve_omop_cdm_engine, ) @@ -53,17 +54,26 @@ def consolidate_queries( raise ValueError("No queries provided.") -def _get_config() -> OmopEmbConfig: - """Load and return the active OmopEmbConfig, converting a missing-file - error into an actionable setup message. - """ - try: - return OmopEmbConfig.get_config() - except FileNotFoundError: - raise RuntimeError( - "No omop-emb configuration file found. " - "Run `omop-config configure omop-emb` to set it up." - ) +def _build_embedding_client( + cfg: OmopEmbConfig, + *, + model: str, + api_base: str, + api_key: str, + provider_type: ProviderType, + embedding_batch_size: int = 32, +) -> EmbeddingClient: + """Construct an EmbeddingClient, applying config-derived dim/prefixes.""" + return EmbeddingClient( + model=model, + api_base=api_base, + api_key=api_key, + embedding_batch_size=embedding_batch_size, + provider_type=provider_type, + embedding_dim=cfg.embedding_dim, + document_embedding_prefix=cfg.document_embedding_prefix, + query_embedding_prefix=cfg.query_embedding_prefix, + ) def _render_search_results( @@ -167,21 +177,22 @@ def add_embeddings( ``create-index`` afterwards to upgrade to an HNSW approximate index. """ - cfg = _get_config() + cfg = load_omop_emb_config() resolved_api_base = api_base or cfg.api_base resolved_api_key = api_key or cfg.api_key resolved_provider = provider or cfg.provider_type resolved_model = model or cfg.embedding_model - backend = resolve_backend() + backend = resolve_backend_from_config(cfg) omop_cdm_engine = resolve_omop_cdm_engine() - embedding_client = EmbeddingClient( + embedding_client = _build_embedding_client( + cfg, model=resolved_model, api_base=resolved_api_base, api_key=resolved_api_key, - embedding_batch_size=batch_size, provider_type=resolved_provider, + embedding_batch_size=batch_size, ) # FLAT registration: metric_type=COSINE is used only for upsert validation; # FLAT accepts any backend-supported metric, so COSINE is always valid here. @@ -322,14 +333,15 @@ def create_index( locked in and all subsequent queries must use the same metric. """ - cfg = _get_config() + cfg = load_omop_emb_config() resolved_api_base = api_base or cfg.api_base resolved_api_key = api_key or cfg.api_key resolved_provider = provider or cfg.provider_type resolved_model = model or cfg.embedding_model - backend = resolve_backend() - embedding_client = EmbeddingClient( + backend = resolve_backend_from_config(cfg) + embedding_client = _build_embedding_client( + cfg, model=resolved_model, api_base=resolved_api_base, api_key=resolved_api_key, @@ -611,7 +623,7 @@ def search( ] = None, ): - cfg = _get_config() + cfg = load_omop_emb_config() resolved_api_base = api_base or cfg.api_base resolved_api_key = api_key or cfg.api_key resolved_provider = provider or cfg.provider_type @@ -619,7 +631,7 @@ def search( resolved_model = model or cfg.embedding_model queries_generator = consolidate_queries(queries=queries, queries_file=queries_file) - backend = resolve_backend() + backend = resolve_backend_from_config(cfg) # CDM enrichment is optional for search try: @@ -630,12 +642,13 @@ def search( "CDM engine not configured; concept names will not be enriched in results." ) - embedding_client = EmbeddingClient( + embedding_client = _build_embedding_client( + cfg, model=resolved_model, api_base=resolved_api_base, api_key=resolved_api_key, - embedding_batch_size=batch_size, provider_type=resolved_provider, + embedding_batch_size=batch_size, ) embedding_reader = EmbeddingReaderInterface( model=embedding_client.canonical_model_name, diff --git a/src/omop_emb/cli/cli_legacy.py b/src/omop_emb/cli/cli_legacy.py index a2510b5..976783c 100644 --- a/src/omop_emb/cli/cli_legacy.py +++ b/src/omop_emb/cli/cli_legacy.py @@ -12,9 +12,9 @@ import typer from tqdm import tqdm -from omop_emb.backends import resolve_backend +from omop_emb.backends import resolve_backend_from_config from omop_emb.backends.index_config import IndexConfig, index_config_from_index_type -from omop_emb.config import IndexType, MetricType, ProviderType +from omop_emb.config import IndexType, MetricType, ProviderType, load_omop_emb_config from omop_emb.storage import embedding_bundle if TYPE_CHECKING: @@ -249,7 +249,8 @@ def import_legacy_faiss_cache( err=True, ) - backend = resolve_backend() + cfg = load_omop_emb_config() + backend = resolve_backend_from_config(cfg) try: from omop_emb.storage.faiss import FAISSCache except ImportError as e: diff --git a/src/omop_emb/cli/cli_maintenance.py b/src/omop_emb/cli/cli_maintenance.py index 1e57faf..aa02557 100644 --- a/src/omop_emb/cli/cli_maintenance.py +++ b/src/omop_emb/cli/cli_maintenance.py @@ -4,12 +4,13 @@ from typing import Annotated, Optional import typer -from omop_emb.backends import resolve_backend +from omop_emb.backends import resolve_backend_from_config from omop_emb.backends.index_config import index_config_from_index_type from omop_emb.config import ( IndexType, MetricType, ProviderType, + load_omop_emb_config, ) from omop_emb.embeddings.embedding_providers import get_provider_from_provider_type from omop_emb.interface import list_registered_models @@ -38,7 +39,8 @@ def list_models( ] = None, ): - backend = resolve_backend() + cfg = load_omop_emb_config() + backend = resolve_backend_from_config(cfg) records = list_registered_models( backend=backend, provider_type=provider_type, @@ -128,7 +130,8 @@ def rebuild_index( embedding_provider = get_provider_from_provider_type(provider_type) model = embedding_provider.canonical_model_name(model) - backend = resolve_backend() + cfg = load_omop_emb_config() + backend = resolve_backend_from_config(cfg) record = backend.get_registered_model(model_name=model) if record is None: @@ -196,7 +199,8 @@ def delete_model( abort=True, ) - backend = resolve_backend() + cfg = load_omop_emb_config() + backend = resolve_backend_from_config(cfg) record = backend.get_registered_model(model_name=model) if record is None: typer.echo( @@ -254,7 +258,8 @@ def export_bundle_cmd( embedding_provider = get_provider_from_provider_type(provider_type) model = embedding_provider.canonical_model_name(model) - backend = resolve_backend() + cfg = load_omop_emb_config() + backend = resolve_backend_from_config(cfg) meta, h5_path = embedding_bundle.export_bundle( backend=backend, @@ -339,7 +344,8 @@ def build_faiss_cache( embedding_provider = get_provider_from_provider_type(provider_type) model = embedding_provider.canonical_model_name(model) - backend = resolve_backend() + cfg = load_omop_emb_config() + backend = resolve_backend_from_config(cfg) index_config = index_config_from_index_type( index_type, @@ -407,7 +413,8 @@ def check_faiss_cache( embedding_provider = get_provider_from_provider_type(provider_type) model = embedding_provider.canonical_model_name(model) - backend = resolve_backend() + cfg = load_omop_emb_config() + backend = resolve_backend_from_config(cfg) record = backend.get_registered_model(model_name=model) if record is None: typer.echo( @@ -475,7 +482,8 @@ def import_bundle_cmd( ] = False, ): - backend = resolve_backend() + cfg = load_omop_emb_config() + backend = resolve_backend_from_config(cfg) try: imported = embedding_bundle.import_bundle( backend=backend, diff --git a/src/omop_emb/config.py b/src/omop_emb/config.py index c27eabd..503687e 100644 --- a/src/omop_emb/config.py +++ b/src/omop_emb/config.py @@ -7,7 +7,7 @@ from pydantic import Field from sqlalchemy import Engine -from oa_configurator import DatabaseConfig, PackageConfigBase, ResourceSpec +from oa_configurator import ConfigurationError, DatabaseConfig, PackageConfigBase, ResourceSpec from omop_alchemy.config import OmopAlchemyConfig @@ -124,6 +124,33 @@ class OmopEmbConfig(PackageConfigBase): ) +def load_omop_emb_config() -> OmopEmbConfig: + """Load the active OmopEmbConfig, the one place CLI entry points should + read config from. + + Converts config errors into an actionable RuntimeError instead of + letting them propagate raw. + + Raises + ------ + RuntimeError + If no config file exists, or the config exists but is missing a + required resource (e.g. the CDM database). + """ + try: + return OmopEmbConfig.get_config() + except FileNotFoundError: + raise RuntimeError( + "No omop-emb configuration file found. " + "Run `omop-config configure omop-emb` to set it up." + ) + except ConfigurationError as exc: + raise RuntimeError( + f"omop-emb configuration is incomplete: {exc} " + "Run `omop-config configure omop-emb` to fix it." + ) from exc + + def resolve_omop_cdm_engine() -> Engine: """Resolve CDM engine via oa-configurator, used read-only.""" return OmopEmbConfig.get_engine(OmopAlchemyConfig.CDM_DB.semantic_name) diff --git a/src/omop_emb/embeddings/embedding_client.py b/src/omop_emb/embeddings/embedding_client.py index 99535df..43bedfc 100644 --- a/src/omop_emb/embeddings/embedding_client.py +++ b/src/omop_emb/embeddings/embedding_client.py @@ -16,7 +16,7 @@ from openai import OpenAI from .embedding_providers import EmbeddingProvider, get_provider_from_provider_type -from omop_emb.config import OmopEmbConfig, ProviderType +from omop_emb.config import ProviderType logger = logging.getLogger(__name__) @@ -56,6 +56,16 @@ class EmbeddingClient: when both are supplied. provider_type : ProviderType, optional Used to construct a provider when *provider* is not supplied. + embedding_dim : int, optional + Known embedding dimensionality, e.g. resolved by the caller from + config. When omitted, resolved lazily on first access to + ``embedding_dim`` (provider API discovery, then a live probe call). + document_embedding_prefix : str, optional + Text prefix prepended to texts embedded with ``EmbeddingRole.DOCUMENT``. + Defaults to ``""`` (no prefix). + query_embedding_prefix : str, optional + Text prefix prepended to texts embedded with ``EmbeddingRole.QUERY``. + Defaults to ``""`` (no prefix). Raises ------ @@ -71,6 +81,9 @@ def __init__( embedding_batch_size: int = 32, provider: Optional[EmbeddingProvider] = None, provider_type: Optional[ProviderType] = None, + embedding_dim: Optional[int] = None, + document_embedding_prefix: str = "", + query_embedding_prefix: str = "", ) -> None: if provider is not None and provider_type is not None: logger.warning( @@ -85,14 +98,26 @@ def __init__( raise ValueError("Must supply either provider or provider_type.") self._model = self._provider.canonical_model_name(model) self._embedding_batch_size = embedding_batch_size - self._embedding_dim: Optional[int] = None + self._embedding_dim: Optional[int] = embedding_dim self._base_client = OpenAI(base_url=api_base, api_key=api_key) - doc_prefix, query_prefix = self.load_embedding_prefixes() self._embedding_prefixes = { - EmbeddingRole.DOCUMENT: doc_prefix, - EmbeddingRole.QUERY: query_prefix, + EmbeddingRole.DOCUMENT: document_embedding_prefix, + EmbeddingRole.QUERY: query_embedding_prefix, } + for role, prefix in self._embedding_prefixes.items(): + if prefix: + logger.info( + f"{role.value.capitalize()} embedding prefix set: {prefix!r}. " + f"All {role.value} texts will be prepended with this prefix." + ) + else: + logger.warning( + f"{role.value.capitalize()} embedding prefix is not set. " + f"This is fine for symmetric models. For asymmetric models (e.g. nomic-embed-text, " + f"E5, BGE), set {role.value.lower()}_embedding_prefix via " + f"'omop-config configure omop_emb'." + ) logger.info( f"{EmbeddingClient.__name__} initialised for model={self._model!r}.\n" @@ -132,7 +157,7 @@ def embedding_dim(self) -> int: """Embedding vector dimension, resolved on first access and cached. Resolution order: - 1. ``OMOP_EMB_EMBEDDING_DIM`` environment variable (explicit override). + 1. ``embedding_dim`` constructor argument (explicit override). 2. Provider API discovery (e.g. Ollama ``/api/show``). 3. Live probe: embed the string ``"test"`` and read the returned shape. One extra API call, but works for any OpenAI-compatible endpoint @@ -141,15 +166,6 @@ def embedding_dim(self) -> int: if self._embedding_dim is not None: return self._embedding_dim - try: - cfg_dim = OmopEmbConfig.get_config().embedding_dim - except FileNotFoundError: - cfg_dim = None - if cfg_dim is not None: - self._embedding_dim = cfg_dim - logger.debug(f"Embedding dimension set from config: {cfg_dim}.") - return cfg_dim - provider_dim = self._provider.get_embedding_dim( model=self._model, api_base=self.api_base ) @@ -271,42 +287,6 @@ def euclidean_distance( b = self.embeddings(text2, embedding_role=text2_role) return float(np.linalg.norm(a - b)) - @staticmethod - def load_embedding_prefixes() -> Tuple[str, str]: - """Load embedding prefixes for document and query roles from the OA_Configurator config. - - Returns - ------- - Tuple[str, str] - A tuple containing the document embedding prefix and the query embedding prefix. - """ - try: - cfg = OmopEmbConfig.get_config() - document_embedding_prefix = cfg.document_embedding_prefix - query_embedding_prefix = cfg.query_embedding_prefix - except Exception: - document_embedding_prefix = "" - query_embedding_prefix = "" - - for role, prefix in [ - (EmbeddingRole.DOCUMENT, document_embedding_prefix), - (EmbeddingRole.QUERY, query_embedding_prefix), - ]: - if prefix: - logger.info( - f"{role.value.capitalize()} embedding prefix loaded from config: {prefix!r}. " - f"All {role.value} texts will be prepended with this prefix." - ) - else: - logger.warning( - f"{role.value.capitalize()} embedding prefix is not set in config. " - f"This is fine for symmetric models. For asymmetric models (e.g. nomic-embed-text, " - f"E5, BGE), set document_embedding_prefix / query_embedding_prefix via " - f"'omop-config configure omop_emb'." - ) - - return document_embedding_prefix, query_embedding_prefix - def _apply_embedding_prefix( self, texts: str | Tuple[str, ...] | List[str], diff --git a/tests/test_embedding_client.py b/tests/test_embedding_client.py index b665626..f618265 100644 --- a/tests/test_embedding_client.py +++ b/tests/test_embedding_client.py @@ -6,24 +6,17 @@ from __future__ import annotations +import logging from unittest.mock import MagicMock, Mock, patch import numpy as np import pytest -from omop_emb.config import OmopEmbConfig, ProviderType +from omop_emb.config import ProviderType from omop_emb.embeddings import EmbeddingClient, EmbeddingRole, OllamaProvider from omop_emb.embeddings.embedding_client import EmbeddingClientError -def _make_emb_config(doc_prefix: str = "", query_prefix: str = "") -> OmopEmbConfig: - """Build a minimal OmopEmbConfig with the requested embedding prefixes.""" - return OmopEmbConfig( - document_embedding_prefix=doc_prefix, - query_embedding_prefix=query_prefix, - ) - - OLLAMA_BASE = "http://localhost:11434/v1" OPENAI_BASE = "https://api.openai.com/v1" OLLAMA_MODEL = "nomic-embed-text:v1.5" @@ -164,6 +157,19 @@ def test_cached_after_first_access(self, mock_openai): _ = client.embedding_dim provider.get_embedding_dim.assert_called_once() + def test_embedding_dim_from_constructor_short_circuits_provider( + self, mock_openai + ): + provider = self._mock_provider(768) + client = EmbeddingClient( + model=OLLAMA_MODEL, + api_base=OLLAMA_BASE, + provider=provider, + embedding_dim=512, + ) + assert client.embedding_dim == 512 + provider.get_embedding_dim.assert_not_called() + # --------------------------------------------------------------------------- # embeddings(): batching, shapes, input coercions @@ -470,82 +476,80 @@ def test_preserves_message(self): # --------------------------------------------------------------------------- -# load_embedding_prefixes(): env var loading and startup logging +# Embedding prefix resolution and startup logging (constructor kwargs) # --------------------------------------------------------------------------- @pytest.mark.unit -class TestLoadEmbeddingPrefixes: - def test_returns_empty_strings_when_not_configured(self, monkeypatch): - monkeypatch.setattr(OmopEmbConfig, "get_config", lambda: _make_emb_config()) - doc_prefix, query_prefix = EmbeddingClient.load_embedding_prefixes() - assert doc_prefix == "" - assert query_prefix == "" - - def test_returns_configured_prefixes(self, monkeypatch): - monkeypatch.setattr( - OmopEmbConfig, - "get_config", - lambda: _make_emb_config("search_document: ", "search_query: "), - ) - doc_prefix, query_prefix = EmbeddingClient.load_embedding_prefixes() - assert doc_prefix == "search_document: " - assert query_prefix == "search_query: " - - def test_logs_info_when_document_prefix_is_set(self, monkeypatch, caplog): - monkeypatch.setattr( - OmopEmbConfig, - "get_config", - lambda: _make_emb_config(doc_prefix="passage: "), - ) - import logging +class TestEmbeddingPrefixLogging: + def test_prefixes_default_to_empty_string(self, mock_openai): + client = EmbeddingClient( + model=OLLAMA_MODEL, api_base=OLLAMA_BASE, provider=OllamaProvider() + ) + prefixes = client.embedding_role_prefixes() + assert prefixes[EmbeddingRole.DOCUMENT] == "" + assert prefixes[EmbeddingRole.QUERY] == "" + def test_prefixes_stored_from_constructor_args(self, mock_openai): + client = EmbeddingClient( + model=OLLAMA_MODEL, + api_base=OLLAMA_BASE, + provider=OllamaProvider(), + document_embedding_prefix="search_document: ", + query_embedding_prefix="search_query: ", + ) + prefixes = client.embedding_role_prefixes() + assert prefixes[EmbeddingRole.DOCUMENT] == "search_document: " + assert prefixes[EmbeddingRole.QUERY] == "search_query: " + + def test_logs_info_when_document_prefix_is_set(self, mock_openai, caplog): with caplog.at_level( logging.INFO, logger="omop_emb.embeddings.embedding_client" ): - EmbeddingClient.load_embedding_prefixes() + EmbeddingClient( + model=OLLAMA_MODEL, + api_base=OLLAMA_BASE, + provider=OllamaProvider(), + document_embedding_prefix="passage: ", + ) assert any( "passage: " in r.message and r.levelname == "INFO" for r in caplog.records ) - def test_logs_info_when_query_prefix_is_set(self, monkeypatch, caplog): - monkeypatch.setattr( - OmopEmbConfig, - "get_config", - lambda: _make_emb_config(query_prefix="search_query: "), - ) - import logging - + def test_logs_info_when_query_prefix_is_set(self, mock_openai, caplog): with caplog.at_level( logging.INFO, logger="omop_emb.embeddings.embedding_client" ): - EmbeddingClient.load_embedding_prefixes() + EmbeddingClient( + model=OLLAMA_MODEL, + api_base=OLLAMA_BASE, + provider=OllamaProvider(), + query_embedding_prefix="search_query: ", + ) assert any( "search_query: " in r.message and r.levelname == "INFO" for r in caplog.records ) - def test_logs_warning_when_document_prefix_not_set(self, monkeypatch, caplog): - monkeypatch.setattr(OmopEmbConfig, "get_config", lambda: _make_emb_config()) - import logging - + def test_logs_warning_when_document_prefix_not_set(self, mock_openai, caplog): with caplog.at_level( logging.WARNING, logger="omop_emb.embeddings.embedding_client" ): - EmbeddingClient.load_embedding_prefixes() + EmbeddingClient( + model=OLLAMA_MODEL, api_base=OLLAMA_BASE, provider=OllamaProvider() + ) warning_messages = [ r.message for r in caplog.records if r.levelname == "WARNING" ] assert any("omop-config configure omop_emb" in m for m in warning_messages) - def test_logs_warning_when_query_prefix_not_set(self, monkeypatch, caplog): - monkeypatch.setattr(OmopEmbConfig, "get_config", lambda: _make_emb_config()) - import logging - + def test_logs_warning_when_query_prefix_not_set(self, mock_openai, caplog): with caplog.at_level( logging.WARNING, logger="omop_emb.embeddings.embedding_client" ): - EmbeddingClient.load_embedding_prefixes() + EmbeddingClient( + model=OLLAMA_MODEL, api_base=OLLAMA_BASE, provider=OllamaProvider() + ) warning_messages = [ r.message for r in caplog.records if r.levelname == "WARNING" ] @@ -560,17 +564,17 @@ def test_logs_warning_when_query_prefix_not_set(self, monkeypatch, caplog): @pytest.mark.unit class TestApplyEmbeddingPrefix: @pytest.fixture - def client_with_prefixes(self, monkeypatch, mock_openai): - monkeypatch.setattr( - OmopEmbConfig, "get_config", lambda: _make_emb_config("doc: ", "query: ") - ) + def client_with_prefixes(self, mock_openai): return EmbeddingClient( - model=OLLAMA_MODEL, api_base=OLLAMA_BASE, provider=OllamaProvider() + model=OLLAMA_MODEL, + api_base=OLLAMA_BASE, + provider=OllamaProvider(), + document_embedding_prefix="doc: ", + query_embedding_prefix="query: ", ) @pytest.fixture - def client_no_prefixes(self, monkeypatch, mock_openai): - monkeypatch.setattr(OmopEmbConfig, "get_config", lambda: _make_emb_config()) + def client_no_prefixes(self, mock_openai): return EmbeddingClient( model=OLLAMA_MODEL, api_base=OLLAMA_BASE, provider=OllamaProvider() ) @@ -621,25 +625,22 @@ def test_no_prefix_returns_texts_unchanged_list(self, client_no_prefixes): result = c._apply_embedding_prefix(["a", "b"], text_role=EmbeddingRole.QUERY) assert result == ["a", "b"] - def test_prefix_reflected_in_api_call(self, monkeypatch, mock_openai): + def test_prefix_reflected_in_api_call(self, mock_openai): """Verify the prefixed text reaches the OpenAI API call.""" _, oi = mock_openai - monkeypatch.setattr( - OmopEmbConfig, - "get_config", - lambda: _make_emb_config(doc_prefix="passage: "), - ) c = EmbeddingClient( - model=OLLAMA_MODEL, api_base=OLLAMA_BASE, provider=OllamaProvider() + model=OLLAMA_MODEL, + api_base=OLLAMA_BASE, + provider=OllamaProvider(), + document_embedding_prefix="passage: ", ) oi.embeddings.create.return_value = _make_embedding_response([[0.1, 0.2]]) c.embeddings("diabetes", embedding_role=EmbeddingRole.DOCUMENT) call_input = oi.embeddings.create.call_args.kwargs["input"] assert call_input == ("passage: diabetes",) - def test_no_prefix_passes_text_verbatim_to_api(self, monkeypatch, mock_openai): + def test_no_prefix_passes_text_verbatim_to_api(self, mock_openai): _, oi = mock_openai - monkeypatch.setattr(OmopEmbConfig, "get_config", lambda: _make_emb_config()) c = EmbeddingClient( model=OLLAMA_MODEL, api_base=OLLAMA_BASE, provider=OllamaProvider() ) From c2b17e605525c3fce7f5c2c8772588dd0ca92b46 Mon Sep 17 00:00:00 2001 From: Nico Loesch Date: Wed, 22 Jul 2026 23:06:37 +0000 Subject: [PATCH 2/2] Remove unused import --- src/omop_emb/backends/base_backend.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/omop_emb/backends/base_backend.py b/src/omop_emb/backends/base_backend.py index a5d4fc9..be5e7f4 100644 --- a/src/omop_emb/backends/base_backend.py +++ b/src/omop_emb/backends/base_backend.py @@ -4,7 +4,7 @@ from functools import wraps import logging from datetime import datetime -from typing import Any, Callable, Generic, Iterable, Mapping, Optional, Sequence, Tuple, TypeVar, Union +from typing import Any, Callable, Generic, Iterable, Mapping, Optional, Sequence, Tuple, TypeVar from numpy import ndarray from sqlalchemy import Engine from sqlalchemy.orm import sessionmaker