diff --git a/README.md b/README.md index ce44612..3722d5c 100644 --- a/README.md +++ b/README.md @@ -39,6 +39,30 @@ omop-emb embeddings search --api-base http://localhost:11434/v1 --api-key ollama --standard-only --domain Condition --k 5 ``` +**Ingest concepts (OpenAI-hosted model):** + +Configure everything once via oa-configurator (see +[Configuration via oa-configurator](#configuration-via-oa-configurator) below) +rather than exporting connection details and passing an API key on the command +line: + +```bash +omop-config configure omop_alchemy +# CDM database connection + +omop-config configure omop_emb +# backend: sqlitevec +# sqlite_path: /data/omop_emb.db +# provider_type: openai +# api_base: https://api.openai.com/v1 +# api_key: +# embedding_model: text-embedding-3-large +``` + +```bash +omop-emb embeddings add-embeddings # backend/CDM/provider/model/api_base/api_key all resolved from config +``` + **pgvector with HNSW index:** ```bash diff --git a/docs/usage/cli.md b/docs/usage/cli.md index ec80b67..521ae7f 100644 --- a/docs/usage/cli.md +++ b/docs/usage/cli.md @@ -97,7 +97,7 @@ omop-emb embeddings add-embeddings --api-base --provider [OPTIONS] |---|---|---|---| | `--api-base` | | **required** | Base URL of the embedding API. | | `--api-key` | | `ollama` | API key for the embedding API. | -| `--provider` | | from config (default `ollama`) | Embedding provider type (e.g. `ollama`). | +| `--provider` | | from config (default `ollama`) | Embedding provider type (`ollama` or `openai`). | | `--model` | `-m` | `text-embedding-3-small` | Embedding model name. | | `--batch-size` | `-b` | `100` | Concepts per API batch. | @@ -160,7 +160,7 @@ omop-emb embeddings create-index --api-base --provider --model --provider --query "hypertens |---|---|---|---| | `--api-base` | | **required** | Base URL of the embedding API. | | `--api-key` | | `ollama` | API key. | -| `--provider` | | from config (default `ollama`) | Embedding provider type (e.g. `ollama`). | +| `--provider` | | from config (default `ollama`) | Embedding provider type (`ollama` or `openai`). | | `--model` | `-m` | `text-embedding-3-small` | Embedding model name. | | `--batch-size` | `-b` | `100` | Batch size for embedding generation. | diff --git a/docs/usage/interface-guide.md b/docs/usage/interface-guide.md index 41c7428..e5391f9 100644 --- a/docs/usage/interface-guide.md +++ b/docs/usage/interface-guide.md @@ -220,12 +220,12 @@ primary backend — no CDM round-trip at query time. ## EmbeddingClient and providers -!!! note - Currently, only `OllamaProvider` is supported. - `EmbeddingClient` wraps any OpenAI-compatible endpoint. It canonicalises the model name at construction time and exposes `canonical_model_name` as the stable -identifier used in the registry. +identifier used in the registry. Two providers are supported: `OllamaProvider` +for self-hosted models served via Ollama, and `OpenAIProvider` for OpenAI-hosted +models (or any OpenAI-compatible API reachable with an API key, without an +Ollama compatibility layer in front of it). ```python from omop_emb import EmbeddingClient @@ -242,6 +242,19 @@ print(client.canonical_model_name) # "nomic-embed-text:v1.5" print(client.embedding_dim) # auto-discovered via Ollama /api/show ``` +```python +# OpenAI — hosted model, authenticated via API key +client = EmbeddingClient( + model="text-embedding-3-large", + api_base="https://api.openai.com/v1", + api_key="sk-...", + provider_type=ProviderType.OPENAI, +) + +print(client.canonical_model_name) # "text-embedding-3-large" +print(client.embedding_dim) # discovered via a live probe call (no discovery endpoint) +``` + --- ## Model name validation @@ -254,6 +267,11 @@ print(client.embedding_dim) # auto-discovered via Ollama /api/show - ✅ `llama3:8b` - Any name with an explicit, immutable tag +**OpenAI:** + +- ✅ `text-embedding-3-large` +- Any name — no tag normalisation is required or applied + ### Invalid names (raise `ValueError`) **Ollama:** @@ -265,7 +283,9 @@ print(client.embedding_dim) # auto-discovered via Ollama /api/show **Why the strictness?** In long-term healthcare data storage, `:latest` is a moving target. Running `ollama pull llama3` silently changes which model version `:latest` points to, breaking consistency between stored embeddings - and new query embeddings. + and new query embeddings. OpenAI-hosted model identifiers do not have this + problem: a given name (e.g. `text-embedding-3-large`) refers to a fixed + model version, so no equivalent tag requirement applies. --- diff --git a/src/omop_emb/__init__.py b/src/omop_emb/__init__.py index e9609b4..1ef2deb 100644 --- a/src/omop_emb/__init__.py +++ b/src/omop_emb/__init__.py @@ -6,6 +6,7 @@ EmbeddingClient, EmbeddingProvider, OllamaProvider, + OpenAIProvider, ) from omop_emb.config import ( BackendType, @@ -43,6 +44,7 @@ "EmbeddingClient", "EmbeddingProvider", "OllamaProvider", + "OpenAIProvider", "BackendType", "IndexType", "MetricType", diff --git a/src/omop_emb/cli/cli_embeddings.py b/src/omop_emb/cli/cli_embeddings.py index f65b4e8..47e2eea 100644 --- a/src/omop_emb/cli/cli_embeddings.py +++ b/src/omop_emb/cli/cli_embeddings.py @@ -15,6 +15,7 @@ MetricType, OmopEmbConfig, ProviderType, + provider_type_examples, resolve_omop_cdm_engine, ) from omop_emb.embeddings import EmbeddingClient @@ -104,7 +105,7 @@ def add_embeddings( Optional[ProviderType], typer.Option( "--provider", - help="Embedding provider type (e.g. 'ollama'). Defaults to the value configured via omop-config.", + help=f"Embedding provider type (e.g. {provider_type_examples()}). Defaults to the value configured via omop-config.", rich_help_panel="Embedding API Options", ), ] = None, @@ -260,7 +261,7 @@ def create_index( Optional[ProviderType], typer.Option( "--provider", - help="Embedding provider type (e.g. 'ollama'). Defaults to the value configured via omop-config.", + help=f"Embedding provider type (e.g. {provider_type_examples()}). Defaults to the value configured via omop-config.", rich_help_panel="Embedding API Options", ), ] = None, @@ -376,7 +377,7 @@ def add_embeddings_with_index( Optional[ProviderType], typer.Option( "--provider", - help="Embedding provider type (e.g. 'ollama'). Defaults to the value configured via omop-config.", + help=f"Embedding provider type (e.g. {provider_type_examples()}). Defaults to the value configured via omop-config.", rich_help_panel="Embedding API Options", ), ] = None, @@ -524,7 +525,7 @@ def search( Optional[ProviderType], typer.Option( "--provider", - help="Embedding provider type (e.g. 'ollama'). Defaults to the value configured via omop-config.", + help=f"Embedding provider type (e.g. {provider_type_examples()}). Defaults to the value configured via omop-config.", rich_help_panel="Embedding API Options", ), ] = None, diff --git a/src/omop_emb/config.py b/src/omop_emb/config.py index 4026632..c27eabd 100644 --- a/src/omop_emb/config.py +++ b/src/omop_emb/config.py @@ -19,9 +19,25 @@ class ProviderType(StrEnum): ------- OLLAMA Self-hosted models served via the Ollama runtime. + OPENAI + OpenAI-hosted models (or any OpenAI-compatible API), authenticated via + an API key rather than a local, unauthenticated endpoint. """ OLLAMA = "ollama" + OPENAI = "openai" + + +def provider_type_examples() -> str: + """Return a human-readable list of ProviderType values for help text. + + e.g. ``"'ollama' or 'openai'"`` for two values, or + ``"'ollama', 'openai', or 'anthropic'"`` once a third is added. + """ + values = [f"'{p.value}'" for p in ProviderType] + if len(values) <= 2: + return " or ".join(values) + return f"{', '.join(values[:-1])}, or {values[-1]}" class OmopEmbConfig(PackageConfigBase): @@ -100,7 +116,7 @@ class OmopEmbConfig(PackageConfigBase): ) provider_type: ProviderType = Field( default=ProviderType.OLLAMA, - description="Embedding provider type (e.g. 'ollama').", + description=f"Embedding provider type (e.g. {provider_type_examples()}).", ) embedding_model: str = Field( default="qwen3-embedding:0.6b", diff --git a/src/omop_emb/embeddings/__init__.py b/src/omop_emb/embeddings/__init__.py index 0bc2f1f..40e8deb 100644 --- a/src/omop_emb/embeddings/__init__.py +++ b/src/omop_emb/embeddings/__init__.py @@ -7,6 +7,7 @@ EmbeddingProvider, get_provider_from_provider_type, OllamaProvider, + OpenAIProvider, ) __all__ = [ @@ -16,4 +17,5 @@ "EmbeddingProvider", "get_provider_from_provider_type", "OllamaProvider", + "OpenAIProvider", ] diff --git a/src/omop_emb/embeddings/embedding_providers.py b/src/omop_emb/embeddings/embedding_providers.py index f4e072c..fe8c2a9 100644 --- a/src/omop_emb/embeddings/embedding_providers.py +++ b/src/omop_emb/embeddings/embedding_providers.py @@ -58,7 +58,14 @@ def canonical_model_name(self, name: str) -> str: str Canonical model name, e.g. ``'llama3:8b'`` or ``'text-embedding-3-small'``. + + Raises + ------ + ValueError + If *name* is empty or whitespace-only. """ + if not name.strip(): + raise ValueError("model name must not be empty.") canonical_model_name = self._canonical_model_name_impl(name) logger.info( f"Set canonical model name for provider {self.provider_type.value!r}: {canonical_model_name!r}" @@ -67,7 +74,12 @@ def canonical_model_name(self, name: str) -> str: @abstractmethod def _canonical_model_name_impl(self, name: str) -> str: - """Provider-specific implementation of canonical model name resolution.""" + """Provider-specific implementation of canonical model name resolution. + + *name* is guaranteed non-empty and non-whitespace-only by the caller + (:meth:`canonical_model_name`) — implementations do not need to + re-check for that. + """ ... def get_embedding_dim(self, model: str, api_base: URL) -> Optional[int]: @@ -181,8 +193,42 @@ def get_embedding_dim(self, model: str, api_base: URL) -> int: ) +class OpenAIProvider(EmbeddingProvider): + """Provider for OpenAI-hosted models (or any OpenAI-compatible API). + + Model names require no tag normalisation, unlike Ollama's ``name:tag`` + requirement. Embedding dimensions have no discovery endpoint equivalent to + Ollama's ``POST /api/show``, so :meth:`get_embedding_dim` uses the base + class's default (``None``) — :attr:`EmbeddingClient.embedding_dim` already + falls back to a live probe (embedding a test string and reading the + resulting shape) whenever a provider cannot discover the dimension via + its API. + """ + + @property + def provider_type(self) -> ProviderType: + return ProviderType.OPENAI + + def _canonical_model_name_impl(self, name: str) -> str: + """Return *name* unchanged (no tag normalisation required). + + Parameters + ---------- + name : str + Model name, e.g. ``'text-embedding-3-large'``. + + Returns + ------- + str + The same model name, stripped of surrounding whitespace. + """ + return name.strip() + + def get_provider_from_provider_type(provider_type: ProviderType) -> EmbeddingProvider: """Map a ProviderType enum to an EmbeddingProvider instance.""" if provider_type == ProviderType.OLLAMA: return OllamaProvider() + if provider_type == ProviderType.OPENAI: + return OpenAIProvider() raise ValueError(f"Unsupported provider type: {provider_type}") diff --git a/tests/test_providers.py b/tests/test_providers.py index dbb7fdc..e28422a 100644 --- a/tests/test_providers.py +++ b/tests/test_providers.py @@ -5,6 +5,7 @@ from omop_emb.config import ProviderType from omop_emb.embeddings import ( OllamaProvider, + OpenAIProvider, get_provider_from_provider_type, ) @@ -49,6 +50,59 @@ def test_strips_whitespace_before_validation(self): def test_strips_whitespace_with_explicit_tag(self): assert OllamaProvider().canonical_model_name(" llama3:8b ") == "llama3:8b" + def test_raises_for_empty_name(self): + with pytest.raises(ValueError, match="must not be empty"): + OllamaProvider().canonical_model_name("") + + def test_raises_for_whitespace_only_name(self): + with pytest.raises(ValueError, match="must not be empty"): + OllamaProvider().canonical_model_name(" ") + + +class TestOpenAIProviderCanonicalModelName: + """OpenAIProvider.canonical_model_name applies no tag normalisation.""" + + def test_returns_name_unchanged(self): + assert ( + OpenAIProvider().canonical_model_name("text-embedding-3-large") + == "text-embedding-3-large" + ) + + def test_strips_whitespace(self): + assert ( + OpenAIProvider().canonical_model_name(" text-embedding-3-large ") + == "text-embedding-3-large" + ) + + def test_idempotent(self): + provider = OpenAIProvider() + canonical = provider.canonical_model_name("text-embedding-3-large") + assert provider.canonical_model_name(canonical) == canonical + + def test_untagged_name_is_not_rejected(self): + """Unlike Ollama, a bare name with no ':tag' is perfectly valid.""" + assert OpenAIProvider().canonical_model_name("text-embedding-3-large") + + def test_raises_for_empty_name(self): + with pytest.raises(ValueError, match="must not be empty"): + OpenAIProvider().canonical_model_name("") + + def test_raises_for_whitespace_only_name(self): + with pytest.raises(ValueError, match="must not be empty"): + OpenAIProvider().canonical_model_name(" ") + + +@pytest.mark.unit +class TestOpenAIProviderGetEmbeddingDim: + def test_returns_none(self): + """No discovery endpoint exists; EmbeddingClient falls back to a live probe.""" + assert ( + OpenAIProvider().get_embedding_dim( + "text-embedding-3-large", "https://api.openai.com/v1" + ) + is None + ) + @pytest.mark.unit class TestGetProviderFromProviderType: @@ -62,6 +116,16 @@ def test_ollama_result_has_correct_provider_type(self): == ProviderType.OLLAMA ) + def test_openai_type_returns_openai_provider(self): + provider = get_provider_from_provider_type(ProviderType.OPENAI) + assert isinstance(provider, OpenAIProvider) + + def test_openai_result_has_correct_provider_type(self): + assert ( + get_provider_from_provider_type(ProviderType.OPENAI).provider_type + == ProviderType.OPENAI + ) + def test_each_call_returns_a_fresh_instance(self): """Provider instances must not be shared/cached between calls.""" a = get_provider_from_provider_type(ProviderType.OLLAMA)