Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion CLAUDE.md
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down
17 changes: 17 additions & 0 deletions EMBEDDINGS.md
Original file line number Diff line number Diff line change
Expand Up @@ -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)

Expand Down Expand Up @@ -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.
Expand Down
10 changes: 8 additions & 2 deletions src/cocoindex_code/embedder_params.py
Original file line number Diff line number Diff line change
Expand Up @@ -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",
Expand Down Expand Up @@ -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:
Expand Down
30 changes: 30 additions & 0 deletions src/cocoindex_code/settings.py
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand All @@ -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),
Expand Down Expand Up @@ -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:
Expand Down Expand Up @@ -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:
Expand Down Expand Up @@ -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 = (
Expand Down
68 changes: 66 additions & 2 deletions src/cocoindex_code/shared.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand All @@ -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

Expand Down
132 changes: 132 additions & 0 deletions tests/test_custom_embedder.py
Original file line number Diff line number Diff line change
@@ -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