mirror of
https://github.com/headroomlabs-ai/headroom.git
synced 2026-08-27 14:17:10 -04:00
1725cd1f83
## Description
The process-wide embedder cache can hand a caller an embedder bound to
the wrong Ollama server.
`_create_embedder` caches by `(backend, model)`:
```python
key = (
config.embedder_backend.value if hasattr(...) else str(...),
config.embedder_model or "",
)
```
But the Ollama branch constructs the embedder with the server URL:
```python
embedder = OllamaEmbedder(base_url=config.ollama_base_url, model_name=config.embedder_model)
```
So two configs in the same process that share a backend and model but
point at different Ollama servers (for example a per-project storage
router, or a fail-over host) collide on the same cache key. The first
call builds and caches an `OllamaEmbedder` bound to server A; the second
call, asking for server B, gets server A's embedder back and silently
embeds against the wrong host.
The code already reasoned about the analogous `openai_api_key` omission
and worked around it with an up-front validation guard (see the comment
above the key), but `ollama_base_url` has no such guard, so it just
resolves to the wrong server.
## Fix
Add `config.ollama_base_url` to the cache key. Same server still hits
the cache (one model load); a different server gets its own embedder.
Non-Ollama backends are unaffected (the URL just becomes an extra,
constant key component).
Closes #
## Type of Change
- [x] Bug fix (non-breaking change that fixes an issue)
- [ ] New feature (non-breaking change that adds functionality)
- [ ] Breaking change (fix or feature that would cause existing
functionality to change)
- [ ] Documentation update
- [ ] Performance improvement
- [ ] Code refactoring (no functional changes)
## Changes Made
- `headroom/memory/factory.py`: include `config.ollama_base_url` in the
embedder cache key, with a comment explaining why.
- `tests/test_memory/test_factory_embedder_cache.py`: new file with
`test_ollama_embedder_cache_keys_on_base_url` (different servers get
different embedders) and
`test_ollama_embedder_cache_reuses_same_base_url` (same server still
caches). Kept out of `test_factory.py` because that module skips
wholesale without `hnswlib`, which these cases don't need.
- `CHANGELOG.md`: Bug Fixes entry.
## Testing
- [ ] Unit tests pass (`pytest`)
- [x] Linting passes (`ruff check .`)
- [ ] Type checking passes (`mypy headroom`)
- [x] New tests added for new functionality
- [ ] Manual testing performed
### Test Output
```text
$ uvx ruff@0.15.17 check headroom/memory/factory.py tests/test_memory/test_factory_embedder_cache.py
All checks passed!
$ python -m py_compile headroom/memory/factory.py tests/test_memory/test_factory_embedder_cache.py
OK
```
## Real Behavior Proof
- Environment: Windows 11, Python 3.12, `uvx ruff@0.15.17`. Importing
`headroom` pulls in the torch/transformers stack and a full `pytest`
gets OOM-killed on this box, so I verified the cache-key behavior with a
dependency-free script that models the `(backend, model)` vs `(backend,
model, base_url)` keys against a simulated cache, and left the full
pytest to CI.
- Exact command / steps: created two configs with the same backend and
model but `ollama_base_url` of `http://gpu1:11434` and
`http://gpu2:11434`, and resolved each through the old key and the new
key against a shared cache.
- Observed result: the old key serves the same embedder object for both,
and the config asking for `gpu2` is handed the `gpu1`-bound embedder;
the new key gives each config its own embedder bound to its own server.
The new tests assert distinct embedders with the right `_base_url` for
different servers, and cache reuse for the same server.
- Not tested: a live Ollama round-trip (`OllamaEmbedder` construction is
offline — it stores the URL and lazily creates its client); full local
`pytest` deferred to CI (OOM, per above).
## Review Readiness
- [x] I have performed a self-review
- [x] This PR is ready for human review
## Checklist
- [x] My code follows the project's style guidelines
- [x] I have performed a self-review of my code
- [x] I have commented my code, particularly in hard-to-understand areas
- [ ] I have made corresponding changes to the documentation
- [x] My changes generate no new warnings
- [x] I have added tests that prove my fix is effective or that my
feature works
- [ ] New and existing unit tests pass locally with my changes
- [x] I have updated the CHANGELOG.md if applicable
## Additional Notes
The "unit tests pass locally" and "type checking" boxes are unchecked
because the full suite imports the ML stack, which I can't run in this
environment; the change adds one component to a cache-key tuple in a
pure function, verified by the standalone proof and the two new tests
for CI. The tests construct only the lightweight (offline) Ollama
embedder, so they don't need a running server or the vector-index deps.
Co-authored-by: Tejas Chopra <chopratejas@gmail.com>
52 lines
1.8 KiB
Python
52 lines
1.8 KiB
Python
"""The embedder cache must not serve an embedder bound to the wrong server.
|
|
|
|
Kept out of ``test_factory.py`` (which skips wholesale without hnswlib) because
|
|
these cases only construct the lightweight Ollama embedder and need no vector
|
|
index.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
from headroom.memory.config import EmbedderBackend, MemoryConfig
|
|
from headroom.memory.factory import _create_embedder, _reset_embedder_cache_for_tests
|
|
|
|
|
|
def test_ollama_embedder_cache_keys_on_base_url():
|
|
"""Two configs that share backend + model but differ in ollama_base_url must
|
|
not share a cached embedder — the second would otherwise get an embedder
|
|
bound to the first server."""
|
|
_reset_embedder_cache_for_tests()
|
|
try:
|
|
cfg1 = MemoryConfig(
|
|
embedder_backend=EmbedderBackend.OLLAMA,
|
|
embedder_model="nomic-embed-text",
|
|
ollama_base_url="http://gpu1:11434",
|
|
)
|
|
cfg2 = MemoryConfig(
|
|
embedder_backend=EmbedderBackend.OLLAMA,
|
|
embedder_model="nomic-embed-text",
|
|
ollama_base_url="http://gpu2:11434",
|
|
)
|
|
|
|
e1 = _create_embedder(cfg1)
|
|
e2 = _create_embedder(cfg2)
|
|
|
|
assert e1 is not e2
|
|
assert e1._base_url == "http://gpu1:11434"
|
|
assert e2._base_url == "http://gpu2:11434"
|
|
finally:
|
|
_reset_embedder_cache_for_tests()
|
|
|
|
|
|
def test_ollama_embedder_cache_reuses_same_base_url():
|
|
"""Same backend + model + base_url still hits the cache (one model load)."""
|
|
_reset_embedder_cache_for_tests()
|
|
try:
|
|
cfg = MemoryConfig(
|
|
embedder_backend=EmbedderBackend.OLLAMA,
|
|
embedder_model="nomic-embed-text",
|
|
ollama_base_url="http://gpu1:11434",
|
|
)
|
|
assert _create_embedder(cfg) is _create_embedder(cfg)
|
|
finally:
|
|
_reset_embedder_cache_for_tests()
|