feat(embeddings): opt-in Ollama embedding backend for GPU acceleration - #982
feat(embeddings): opt-in Ollama embedding backend for GPU acceleration#982felipetruman wants to merge 2 commits into
Conversation
…ation
Adds `get_embedding_function()` in `backends.chroma` that returns an
`OllamaEmbeddingFunction` when `EMBEDDING_PROVIDER=ollama`, and `None`
otherwise (preserving the current `DefaultEmbeddingFunction` behavior).
The embedding function is threaded through every ChromaDB collection
accessor — `ChromaBackend.get_collection`, `get_or_create_collection`,
`create_collection`, and `mcp_server._get_collection` — so the same EF
is applied on both write and read paths, as ChromaDB requires.
Environment variables (all optional; safe defaults):
EMBEDDING_PROVIDER set to "ollama" to enable (default: unset → Chroma default)
OLLAMA_URL base URL (default: http://localhost:11434)
OLLAMA_EMBED_MODEL model tag (default: nomic-embed-text)
OLLAMA_EMBED_TIMEOUT seconds (default: 60)
Motivation: on machines with a local GPU (e.g. AMD Radeon via ROCm), a
persistent Ollama server is dramatically faster than the bundled ONNX
MiniLM that runs on CPU — ~14 ms/embed warm vs. hundreds of ms/embed,
and embedding-heavy operations like `mempalace mine` become minutes
instead of hours.
Trade-off: nomic-embed-text produces 768-dim vectors vs. MiniLM's
384-dim, so switching providers on an existing palace is not
backward-compatible — users must `mempalace nuke` and remine. This is
opt-in and gated behind an env var, so existing palaces are unaffected.
Verified end-to-end:
- `get_embedding_function()` returns OllamaEmbeddingFunction with the
correct base URL / model / timeout
- A 20-file mine into a fresh palace produced 425 embeddings, all
768-dim (confirmed via `SELECT dimension FROM collections` on
chroma.sqlite3)
- `mempalace search` round-trips against the new collection and
returns semantically ranked results
- VRAM usage on AMD Radeon RX 9060 XT: 2.33 GB (nomic loaded)
Tested on chromadb 1.5.7, Python 3.14, Linux (ROCm).
There was a problem hiding this comment.
Pull request overview
Adds an environment-variable gated embedding backend selection so MemPalace can route ChromaDB embedding generation through a local Ollama server (typically GPU-accelerated) while keeping the existing Chroma default embedding behavior unchanged unless explicitly enabled.
Changes:
- Introduces
get_embedding_function()factory in the Chroma backend to select Ollama vs default based on env vars. - Passes the resolved embedding function through Chroma collection creation/access paths in
ChromaBackend. - Ensures
mcp_server._get_collection()also uses the same embedding function for consistency with the backend.
Reviewed changes
Copilot reviewed 2 out of 2 changed files in this pull request and generated 7 comments.
| File | Description |
|---|---|
mempalace/backends/chroma.py |
Adds env-gated embedding function factory and forwards EF into Chroma collection APIs. |
mempalace/mcp_server.py |
Updates MCP server’s direct Chroma collection access to use the same embedding function. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
| ef = get_embedding_function() | ||
| if create: | ||
| _collection_cache = ChromaCollection( | ||
| client.get_or_create_collection( | ||
| _config.collection_name, metadata={"hnsw:space": "cosine"} | ||
| _config.collection_name, | ||
| metadata={"hnsw:space": "cosine"}, | ||
| embedding_function=ef, # type: ignore[arg-type] | ||
| ) |
There was a problem hiding this comment.
The new # type: ignore[arg-type] on embedding_function=ef is masking a real typing mismatch introduced by allowing ef to be None. Prefer avoiding the ignore by only passing embedding_function when ef is not None (e.g., build a kwargs dict), or by adjusting get_embedding_function()'s return type / local typing so the call sites stay type-safe.
| client = _get_client() | ||
| ef = get_embedding_function() | ||
| if create: |
There was a problem hiding this comment.
Because _get_collection() wraps the whole block in except Exception: return None, any misconfiguration in get_embedding_function() (e.g., invalid OLLAMA_EMBED_TIMEOUT, missing Ollama EF in older chromadb installs) will be silently swallowed and later surfaces as "no palace" behavior. Consider at least logging the exception (or narrowing the exception handling) so embedding-provider configuration errors are debuggable.
| client = self._client(palace_path) | ||
| ef = get_embedding_function() | ||
| if create: | ||
| collection = client.get_or_create_collection( | ||
| collection_name, metadata={"hnsw:space": "cosine"} | ||
| collection_name, | ||
| metadata={"hnsw:space": "cosine"}, | ||
| embedding_function=ef, | ||
| ) | ||
| else: | ||
| collection = client.get_collection(collection_name) | ||
| collection = client.get_collection( | ||
| collection_name, embedding_function=ef | ||
| ) |
There was a problem hiding this comment.
ef = get_embedding_function() can be None, but embedding_function=ef is always passed into get_or_create_collection() / get_collection(). To avoid type mismatches (and potential runtime issues if chromadb doesn't accept None here), consider only including the embedding_function kwarg when ef is not None and otherwise letting Chroma use its default behavior.
| collection = self._client(palace_path).create_collection( | ||
| collection_name, metadata={"hnsw:space": hnsw_space} | ||
| collection_name, | ||
| metadata={"hnsw:space": hnsw_space}, | ||
| embedding_function=get_embedding_function(), | ||
| ) |
There was a problem hiding this comment.
create_collection(..., embedding_function=get_embedding_function()) will pass None through when Ollama is not enabled. If the intent is "use Chroma's default embedding function", consider conditionally omitting the kwarg when get_embedding_function() returns None for consistency with older/newer chromadb versions and to keep typing clean.
| def get_embedding_function(): | ||
| """Return an embedding function based on env, or ``None`` for Chroma's default. | ||
|
|
||
| Set ``EMBEDDING_PROVIDER=ollama`` to route embeddings through a local | ||
| Ollama server (GPU-accelerated). Tunables: | ||
|
|
||
| - ``OLLAMA_URL`` (default ``http://localhost:11434``) - base URL | ||
| - ``OLLAMA_EMBED_MODEL`` (default ``nomic-embed-text``) | ||
| - ``OLLAMA_EMBED_TIMEOUT`` seconds (default ``60``) | ||
|
|
||
| Returning ``None`` keeps ChromaDB's ``DefaultEmbeddingFunction`` (ONNX | ||
| MiniLM, 384 dims, CPU) for backward compatibility with existing palaces. | ||
| """ | ||
| provider = os.environ.get("EMBEDDING_PROVIDER", "").lower() | ||
| if provider != "ollama": | ||
| return None | ||
| from chromadb.utils.embedding_functions import OllamaEmbeddingFunction | ||
|
|
||
| return OllamaEmbeddingFunction( | ||
| url=os.environ.get("OLLAMA_URL") or "http://localhost:11434", | ||
| model_name=os.environ.get("OLLAMA_EMBED_MODEL") or "nomic-embed-text", | ||
| timeout=int(os.environ.get("OLLAMA_EMBED_TIMEOUT") or "60"), | ||
| ) |
There was a problem hiding this comment.
get_embedding_function() is newly introduced but currently has no automated coverage. Consider adding unit tests that (1) default to None when EMBEDDING_PROVIDER is unset/other values, and (2) with EMBEDDING_PROVIDER=ollama reads URL/model/timeout defaults + overrides from env so regressions are caught without requiring a live Ollama server (can assert the returned object type/attributes or that it is callable).
| return OllamaEmbeddingFunction( | ||
| url=os.environ.get("OLLAMA_URL") or "http://localhost:11434", | ||
| model_name=os.environ.get("OLLAMA_EMBED_MODEL") or "nomic-embed-text", | ||
| timeout=int(os.environ.get("OLLAMA_EMBED_TIMEOUT") or "60"), | ||
| ) |
There was a problem hiding this comment.
timeout=int(os.environ.get("OLLAMA_EMBED_TIMEOUT") or "60") will raise ValueError for a non-integer env value, and that error will be fairly opaque to users. Consider validating/parsing with a clearer exception message (e.g., "OLLAMA_EMBED_TIMEOUT must be an integer number of seconds") or falling back to the default when parsing fails.
| Returning ``None`` keeps ChromaDB's ``DefaultEmbeddingFunction`` (ONNX | ||
| MiniLM, 384 dims, CPU) for backward compatibility with existing palaces. | ||
| """ | ||
| provider = os.environ.get("EMBEDDING_PROVIDER", "").lower() |
There was a problem hiding this comment.
provider = os.environ.get("EMBEDDING_PROVIDER", "").lower() doesn't handle leading/trailing whitespace (e.g., " ollama"), which can make the opt-in flag unexpectedly fail. Consider normalizing with .strip().lower() before comparing.
| provider = os.environ.get("EMBEDDING_PROVIDER", "").lower() | |
| provider = os.environ.get("EMBEDDING_PROVIDER", "").strip().lower() |
Addresses MemPalace#982 (review) 1. EMBEDDING_PROVIDER now normalized with .strip().lower() so leading/ trailing whitespace no longer silently disables the opt-in flag. 2. OLLAMA_EMBED_TIMEOUT parsing moved into a dedicated _parse_timeout() helper that raises ValueError with a clear message naming the env var instead of a bare ValueError from int(), and treats empty/whitespace as "fall back to 60" rather than crashing. 3. New _ef_kwargs() helper returns {"embedding_function": ef} when an EF is configured and {} otherwise. All ChromaBackend + mcp_server call sites now spread that dict via ** so: - the embedding_function kwarg is never present on the default path (was passed as None before) - the previous # type: ignore[arg-type] annotations are gone - older chromadb builds that reject embedding_function=None keep working without change Unrelated Pyright diagnostics that pre-date this PR are left alone. 4. _get_collection()'s broad except now logs via logger.debug(..., exc_info=True) so embedding-provider misconfigurations (bad OLLAMA_EMBED_TIMEOUT, chromadb without OllamaEmbeddingFunction, server unreachable, ...) are diagnosable instead of silently surfacing as "no palace" later. The broad except is intentional — palace-not-found must stay non-fatal. 5. New tests/test_embeddings_ollama.py (32 assertions across 25 tests), all passing, covering: - provider-disabled returns None (unset, empty, whitespace, other) - provider-enabled case/whitespace normalization - URL/model/timeout defaults + env overrides with whitespace trim - bad OLLAMA_EMBED_TIMEOUT raises helpful ValueError - _parse_timeout valid/invalid parametrised paths - _ef_kwargs empty vs populated + spread-into-dict invariant - no cached module-level state across toggles No Ollama server is required; tests only inspect the EF object's attributes and the kwargs dict.
|
Thanks for the review @copilot-pull-request-reviewer. All 7 comments addressed in 06956aa. Item-by-item:
Classification: all 7 fell into "quality / design" — no blocking critical bugs. One design pattern change (the CI: no checks report back on this PR — let me know if there's something I should wire up locally. |
igorls
added
enhancement
|
Hi, get_embedding_function() selects the env var value before stripping, so whitespace-only OLLAMA_URL or OLLAMA_EMBED_MODEL becomes an empty string and is passed into OllamaEmbeddingFunction, causing embedding requests to fail at runtime. Severity: action required | Category: correctness How to fix: Validate stripped URL/model strings Agent prompt to fix - you can give this to your LLM of choice:
We noticed a couple of other issues in this PR as well - happy to share if helpful. Spotted by Qodo code review - free for open-source projects. |
|
Hi, thanks for the contribution. This PR has merge conflicts with Could you rebase onto If this change is no longer relevant, feel free to close the PR. (This message is part of a periodic backlog pass, sent to all open PRs that match this state.) |
igorls
added
the
needs-rebase
|
+1 from a heavy multilingual user — landing this would close a real gap. I've been monkey-patching the same hook locally since v3.0.0 (FR/EN corpus, 68k drawers). The architectural cleanup in 3.3.4 ( One suggestion: consider documenting Tested today against v3.3.4 with this PR's approach (re-implemented locally as a Happy to share concrete metrics on a 68k-drawer FR/EN palace (latency, recall@k, hybrid score distributions) if it would help reviewers, just ping me. |
Summary
Adds an opt-in
OllamaEmbeddingFunctionpath to the ChromaDB backend. WhenEMBEDDING_PROVIDER=ollamais set, MemPalace routes all embedding generation through a local Ollama server — typically GPU-accelerated on machines with a supported GPU (AMD via ROCm, NVIDIA via CUDA, Apple Silicon).Default behavior is unchanged: without the env var, ChromaDB's bundled
DefaultEmbeddingFunction(ONNX MiniLM, 384-dim, CPU) is used.Why
On a machine with a local GPU, the default ONNX MiniLM running on CPU becomes the bottleneck for
mempalace mineon large projects — on my workstation (AMD Radeon RX 9060 XT, 16 GB VRAM), mining~/freedomdigitalhubwith MiniLM/CPU was still running after 2 hours. A persistent Ollama+nomic-embed-text server on the same GPU generates embeddings at ~14 ms/embed warm, turning hour-scale mines into minute-scale ones.What changed
get_embedding_function()inmempalace/backends/chroma.py— pure env-var gated factory.ChromaBackend.get_collection,get_or_create_collection, andcreate_collectionnow forward the embedding function to ChromaDB.mcp_server._get_collection(which bypasses ChromaBackend for its own inode-based cache) also forwards the EF, keeping writes and reads consistent.Environment variables (all optional)
EMBEDDING_PROVIDEROLLAMA_URLOLLAMA_EMBED_MODELOLLAMA_EMBED_TIMEOUTCompatibility
This is opt-in. Existing palaces are not touched unless the user sets
EMBEDDING_PROVIDER=ollama.Because
nomic-embed-textproduces 768-dim vectors and MiniLM produces 384-dim, switching providers on an existing palace requiresmempalace nuke+ remine. The docs should make this explicit — happy to add a note in a follow-up commit or in this PR, whichever reviewers prefer.Verification
End-to-end on the reporter's workstation (Python 3.14, chromadb 1.5.7, ROCm, AMD RX 9060 XT):
Test plan
🤖 Generated with Claude Code