diff --git a/CLAUDE.md b/CLAUDE.md index e0f85db..069ff4e 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -59,7 +59,7 @@ The daemon idle-exits after `daemon.idle_timeout_minutes` (default 180, 0 = neve The embedder is created once at daemon startup (`create_embedder()` in `shared.py`) and stored in the `ProjectRegistry`. Two param dicts (`indexing_params`, `query_params`) are threaded through CocoIndex context keys so asymmetric models (Cohere, Voyage, nomic-embed) can use different kwargs for document vs. query embedding. -Two install flavors: `[full]` bundles `sentence-transformers` (local inference); slim uses only LiteLLM (cloud APIs). The `PacedLiteLLMEmbedder` in `litellm_embedder.py` adds rate-limiting between requests. +Two install flavors: `[full]` bundles `sentence-transformers` (local inference); slim uses only LiteLLM (cloud APIs). The `PacedLiteLLMEmbedder` in `litellm_embedder.py` adds rate-limiting between requests. A third path, `provider: custom`, loads a user factory named by `embedding.embedder` (`"module.path:callable"`) through `load_custom_embedder_factory()` in `shared.py`; its params are passed through unvalidated and the result must satisfy `VectorSchemaProvider`. ### Settings and Project Discovery diff --git a/EMBEDDINGS.md b/EMBEDDINGS.md index 71d7fb4..3c0a52b 100644 --- a/EMBEDDINGS.md +++ b/EMBEDDINGS.md @@ -12,6 +12,7 @@ - **[Local Sentence-Transformers](#sentence-transformers-local)** - **[LiteLLM Remote (Cloud Providers)](#litellm-remote-cloud-providers)** - **[LiteLLM Local](#local-litellm-providers)** +- [Custom Embedder (`provider: custom`)](#custom-embedder-provider-custom) - [Choosing Based on Your Content](#choosing-based-on-your-content) - [Pacing & Rate Limits](#pacing--rate-limits) @@ -207,6 +208,22 @@ envs: --- +## Custom Embedder (`provider: custom`) + +When neither path fits (a provider litellm does not cover, a thinner dependency tree, a house embedding service), point the daemon at your own factory. The factory is an importable `"module.path:callable"` that cocoindex-code calls as `factory(settings, indexing_params)` once at daemon start; it must return an object with an async `embed(text, **params)` method returning a NumPy vector and an async `__coco_vector_schema__()` returning `cocoindex.resources.schema.VectorSchema` (cocoindex's `VectorSchemaProvider`). Give it a `__coco_memo_key__()` too, so cached embeddings are invalidated when the model or its request options change. + +```yaml +embedding: + provider: custom + embedder: my_embedders.voyage:make_embedder # module.path:callable + model: voyage-4-large # passed to your factory in settings.model + min_interval_ms: 20 # also available to the factory + indexing_params: {input_type: document} + query_params: {input_type: query} +``` + +`indexing_params` and `query_params` are not validated for a custom embedder: whatever keys you put there are passed to your `embed()` verbatim, the indexing side also reaching the factory as its second argument so defaults can be baked in. The package that provides the factory has to be importable by the daemon's interpreter, for example installed into the same environment (`uv tool install cocoindex-code --with my-embedders`). `ccc init` does not offer this provider; write the settings file by hand. A malformed `embedder` value, a module that does not import, or a factory that returns something without the contract fails at daemon start with the spec named in the error. + ## Choosing Based on Your Content - **Heavy Source Code, CPU-only**: Use **LateOn-Code** (Micro/Small) or **CodeSearch-ModernBERT-Crow-Plus** (Medium Encoder). Both are encoder-based and score >0.85 on code search benchmarks with good CPU throughput. diff --git a/src/cocoindex_code/embedder_params.py b/src/cocoindex_code/embedder_params.py index 01efaf9..c4c0d7e 100644 --- a/src/cocoindex_code/embedder_params.py +++ b/src/cocoindex_code/embedder_params.py @@ -10,7 +10,7 @@ from typing import Any, NamedTuple from .embedder_defaults import LEGACY_QUERY_PROMPT_MODELS -from .settings import EmbeddingSettings +from .settings import CUSTOM_PROVIDER, EmbeddingSettings __all__ = [ "EmbedderParams", @@ -48,7 +48,13 @@ def validate_params( indexing_params: dict[str, Any] | None, query_params: dict[str, Any] | None, ) -> None: - """Raise ``ValueError`` if either dict contains keys not accepted by *provider*.""" + """Raise ``ValueError`` if either dict contains keys not accepted by *provider*. + + A ``custom`` embedder owns its own kwargs: whatever its ``embed()`` accepts + is passed through verbatim, so nothing is validated here. + """ + if provider == CUSTOM_PROVIDER: + return accepted = accepted_kwargs_for(provider) for side, params in (("indexing_params", indexing_params), ("query_params", query_params)): if not params: diff --git a/src/cocoindex_code/settings.py b/src/cocoindex_code/settings.py index 0aa1faf..58c8e41 100644 --- a/src/cocoindex_code/settings.py +++ b/src/cocoindex_code/settings.py @@ -94,10 +94,18 @@ # --------------------------------------------------------------------------- +CUSTOM_PROVIDER = "custom" + + @dataclass class EmbeddingSettings: model: str provider: str = "litellm" + # ``provider: custom`` only: ``"module.path:callable"`` naming a factory that + # returns the embedder. The factory is called as ``factory(settings, + # indexing_params)`` and must return an object with an async ``embed`` + # method and ``__coco_vector_schema__`` (cocoindex's ``VectorSchemaProvider``). + embedder: str | None = None device: str | None = None min_interval_ms: int | None = None # PyTorch MPS allocator limits used by CocoIndex's isolated GPU runner. @@ -110,6 +118,17 @@ class EmbeddingSettings: query_params: dict[str, Any] | None = None def __post_init__(self) -> None: + if self.provider == CUSTOM_PROVIDER: + if not self.embedder or ":" not in self.embedder: + raise ValueError( + 'embedding.provider "custom" requires embedding.embedder set to ' + '"module.path:callable"' + ) + elif self.embedder is not None: + raise ValueError( + f"embedding.embedder is only used with provider {CUSTOM_PROVIDER!r} " + f"(got provider {self.provider!r})" + ) ratios = ( ("mps_low_watermark_ratio", self.mps_low_watermark_ratio), ("mps_high_watermark_ratio", self.mps_high_watermark_ratio), @@ -498,6 +517,8 @@ def _embedding_settings_to_dict(embedding: EmbeddingSettings) -> dict[str, Any]: "provider": embedding.provider, "model": embedding.model, } + if embedding.embedder is not None: + d["embedder"] = embedding.embedder if embedding.device is not None: d["device"] = embedding.device if embedding.min_interval_ms is not None: @@ -537,6 +558,8 @@ def _user_settings_from_dict(d: dict[str, Any]) -> UserSettings: emb_kwargs: dict[str, Any] = {"model": emb_dict["model"]} if "provider" in emb_dict: emb_kwargs["provider"] = emb_dict["provider"] + if "embedder" in emb_dict: + emb_kwargs["embedder"] = emb_dict["embedder"] if "device" in emb_dict: emb_kwargs["device"] = emb_dict["device"] if "min_interval_ms" in emb_dict: @@ -712,6 +735,13 @@ def save_user_settings(settings: UserSettings) -> Path: " # indexing_params: {}\n" " # query_params: {}\n" ), + "custom": ( + " #\n" + " # Extra kwargs passed to the embedder, whatever your factory's embed()\n" + " # accepts; cocoindex-code does not validate them for a custom embedder.\n" + " # indexing_params: {}\n" + " # query_params: {}\n" + ), } _MPS_SAFETY_COMMENT = ( diff --git a/src/cocoindex_code/shared.py b/src/cocoindex_code/shared.py index 2061aab..d718bb5 100644 --- a/src/cocoindex_code/shared.py +++ b/src/cocoindex_code/shared.py @@ -2,35 +2,54 @@ from __future__ import annotations +import importlib import importlib.util import logging import os import pathlib import sys import traceback as _tb +from collections.abc import Awaitable from dataclasses import dataclass -from typing import TYPE_CHECKING, Annotated, Any, NamedTuple, Union +from typing import TYPE_CHECKING, Annotated, Any, NamedTuple, Protocol, Union, runtime_checkable import cocoindex as coco import numpy as np import numpy.typing as npt from cocoindex.connectors import sqlite +from cocoindex.resources.schema import VectorSchema if TYPE_CHECKING: from cocoindex.ops.litellm import LiteLLMEmbedder from cocoindex.ops.sentence_transformers import SentenceTransformerEmbedder -from .settings import EmbeddingSettings +from .settings import CUSTOM_PROVIDER, EmbeddingSettings logger = logging.getLogger(__name__) SBERT_PREFIX = "sbert/" DEFAULT_LITELLM_MIN_INTERVAL_MS = 5 + +@runtime_checkable +class CustomEmbedder(Protocol): + """What a ``provider: custom`` factory must return. + + An async ``embed`` (one text in, one float32 vector out; extra kwargs are + the user's ``indexing_params`` / ``query_params``) and cocoindex's + ``VectorSchemaProvider`` method, so the vector column can be sized. + """ + + def embed(self, text: str, /, **kwargs: Any) -> Awaitable[npt.NDArray[np.float32]]: ... + + def __coco_vector_schema__(self) -> Awaitable[VectorSchema]: ... + + # Type alias Embedder = Union[ "SentenceTransformerEmbedder", "LiteLLMEmbedder", + CustomEmbedder, ] # Context keys @@ -126,12 +145,55 @@ async def check_embedding( return EmbeddingCheckResult(dim=None, error=msg, traceback=_tb.format_exc()) +def load_custom_embedder_factory(spec: str) -> Any: + """Import the ``"module.path:callable"`` named by ``embedding.embedder``. + + Raises ``ValueError`` with the spec in the message for a malformed spec, a + module that does not import, an attribute that does not exist, or one that + is not callable, so a settings typo fails at daemon start with a clear line + rather than deep inside the first embed call. + """ + module_path, sep, attr = spec.partition(":") + if not sep or not module_path or not attr: + raise ValueError(f'embedding.embedder must be "module.path:callable", got {spec!r}') + try: + module = importlib.import_module(module_path) + except ImportError as e: + raise ValueError(f"embedding.embedder {spec!r}: cannot import {module_path!r}: {e}") from e + factory = getattr(module, attr, None) + if factory is None: + raise ValueError(f"embedding.embedder {spec!r}: {module_path!r} has no attribute {attr!r}") + if not callable(factory): + raise ValueError(f"embedding.embedder {spec!r}: {attr!r} is not callable") + return factory + + +def _create_custom_embedder( + settings: EmbeddingSettings, + indexing_params: dict[str, Any] | None, +) -> CustomEmbedder: + assert settings.embedder is not None # enforced by EmbeddingSettings.__post_init__ + factory = load_custom_embedder_factory(settings.embedder) + instance = factory(settings, dict(indexing_params) if indexing_params else {}) + if not isinstance(instance, CustomEmbedder): + raise ValueError( + f"embedding.embedder {settings.embedder!r} returned {type(instance).__name__}, " + "which lacks an async embed() method and/or __coco_vector_schema__()" + ) + logger.info("Embedding model (custom %s): %s", settings.embedder, settings.model) + return instance + + def create_embedder( settings: EmbeddingSettings, indexing_params: dict[str, Any] | None = None, ) -> Embedder: """Create and return an embedder instance based on settings. + ``provider: custom`` loads the factory named by ``settings.embedder`` and + calls it as ``factory(settings, indexing_params)``; the result must satisfy + cocoindex's ``VectorSchemaProvider`` and expose an async ``embed()``. + For LiteLLM embedders, *indexing_params* (e.g. ``{"input_type": "passage"}``) are passed to the constructor as default kwargs forwarded into every ``litellm.aembedding`` call — including paths that don't go through @@ -147,6 +209,8 @@ def create_embedder( :data:`INDEXING_EMBED_PARAMS`. """ instance: Embedder + if settings.provider == CUSTOM_PROVIDER: + return _create_custom_embedder(settings, indexing_params) if settings.provider == "sentence-transformers": from cocoindex.ops.sentence_transformers import SentenceTransformerEmbedder diff --git a/tests/test_custom_embedder.py b/tests/test_custom_embedder.py new file mode 100644 index 0000000..a653823 --- /dev/null +++ b/tests/test_custom_embedder.py @@ -0,0 +1,132 @@ +"""Tests for ``provider: custom``: the embedder hook that loads a user factory.""" + +from __future__ import annotations + +import sys +import textwrap +from pathlib import Path +from typing import Any + +import numpy as np +import pytest +from cocoindex.resources.schema import VectorSchema + +from cocoindex_code.embedder_params import validate_params +from cocoindex_code.settings import ( + CUSTOM_PROVIDER, + EmbeddingSettings, + _embedding_settings_to_dict, + _user_settings_from_dict, +) +from cocoindex_code.shared import CustomEmbedder, create_embedder, load_custom_embedder_factory + +FACTORY_MODULE = textwrap.dedent( + """ + import numpy as np + from cocoindex.resources.schema import VectorSchema + + class FakeEmbedder: + def __init__(self, settings, indexing_params): + self.settings = settings + self.indexing_params = indexing_params + + async def embed(self, text, **kwargs): + return np.zeros(3, dtype=np.float32) + + async def __coco_vector_schema__(self): + return VectorSchema(dtype=np.dtype(np.float32), size=3) + + def make(settings, indexing_params): + return FakeEmbedder(settings, indexing_params) + + def make_wrong(settings, indexing_params): + return object() + + not_callable = 42 + """ +) + + +@pytest.fixture +def factory_module(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> str: + name = "custom_embedder_fixture_mod" + (tmp_path / f"{name}.py").write_text(FACTORY_MODULE) + monkeypatch.syspath_prepend(str(tmp_path)) + sys.modules.pop(name, None) + yield name + sys.modules.pop(name, None) + + +def test_settings_require_an_embedder_for_custom() -> None: + with pytest.raises(ValueError, match="requires embedding.embedder"): + EmbeddingSettings(model="m", provider=CUSTOM_PROVIDER) + with pytest.raises(ValueError, match="requires embedding.embedder"): + EmbeddingSettings(model="m", provider=CUSTOM_PROVIDER, embedder="no_colon") + + +def test_settings_reject_embedder_on_other_providers() -> None: + with pytest.raises(ValueError, match="only used with provider 'custom'"): + EmbeddingSettings(model="m", provider="litellm", embedder="pkg.mod:make") + + +def test_settings_round_trip_keeps_embedder() -> None: + settings = EmbeddingSettings( + model="voyage-4-large", + provider=CUSTOM_PROVIDER, + embedder="pkg.mod:make", + indexing_params={"input_type": "document"}, + ) + d = _embedding_settings_to_dict(settings) + assert d["provider"] == "custom" and d["embedder"] == "pkg.mod:make" + loaded = _user_settings_from_dict({"embedding": d}) + assert loaded.embedding == settings + + +def test_custom_params_are_not_validated() -> None: + # anything goes: the custom embedder owns its kwargs + validate_params(CUSTOM_PROVIDER, {"anything": 1}, {"else": 2}) + with pytest.raises(ValueError, match="unknown key"): + validate_params("litellm", {"anything": 1}, None) + + +def test_create_embedder_loads_the_factory(factory_module: str) -> None: + settings = EmbeddingSettings( + model="fake-model", provider=CUSTOM_PROVIDER, embedder=f"{factory_module}:make" + ) + embedder: Any = create_embedder(settings, {"input_type": "document"}) + assert isinstance(embedder, CustomEmbedder) + assert embedder.settings is settings + assert embedder.indexing_params == {"input_type": "document"} + + +def test_create_embedder_rejects_a_factory_without_the_contract(factory_module: str) -> None: + settings = EmbeddingSettings( + model="m", provider=CUSTOM_PROVIDER, embedder=f"{factory_module}:make_wrong" + ) + with pytest.raises(ValueError, match="lacks an async embed"): + create_embedder(settings) + + +@pytest.mark.parametrize( + ("spec", "message"), + [ + ("nope_module_xyz:make", "cannot import"), + ("{mod}:missing_attr", "has no attribute"), + ("{mod}:not_callable", "is not callable"), + ("{mod}", 'must be "module.path:callable"'), + ], +) +def test_loader_errors_name_the_spec(factory_module: str, spec: str, message: str) -> None: + with pytest.raises(ValueError, match=message): + load_custom_embedder_factory(spec.format(mod=factory_module)) + + +async def test_loaded_embedder_embeds_and_reports_schema(factory_module: str) -> None: + settings = EmbeddingSettings( + model="m", provider=CUSTOM_PROVIDER, embedder=f"{factory_module}:make" + ) + embedder: Any = create_embedder(settings) + vec = await embedder.embed("hello") + assert vec.shape == (3,) and vec.dtype == np.float32 + schema: VectorSchema = await embedder.__coco_vector_schema__() + assert schema.size == 3