feat(storage): add ChromaDB and LanceDB vector backends#50
Conversation
Implements two new embedded vector backends following the existing StorageBackend interface: - ChromaBackend: supports EphemeralClient (in-memory), PersistentClient (local disk), or HttpClient (remote server). Install with `pip install 'vektori[chroma]'`. - LanceDBBackend: serverless columnar vector DB backed by Apache Arrow. Supports local paths and cloud URIs (S3/GCS/Azure). Install with `pip install 'vektori[lancedb]'`. Both backends are wired into the factory via `storage_backend="chroma"` and `storage_backend="lancedb"`. Adds 23 tests (10 unit + 13 Chroma integration); LanceDB integration tests auto-skip when not installed. Closes #45, closes #46
|
Warning Rate limit exceeded
Your organization is not enrolled in usage-based pricing. Contact your admin to enable usage-based pricing to continue reviews beyond the rate limit, or try again in 56 minutes and 55 seconds. ⌛ How to resolve this issue?After the wait time has elapsed, a review can be triggered using the We recommend that you space out your commits to avoid hitting the rate limit. 🚦 How do rate limits work?CodeRabbit enforces hourly rate limits for each developer per organization. Our paid plans have higher rate limits than the trial, open-source and free plans. In all cases, we re-allow further reviews after a brief timeout. Please see our FAQ for further information. ℹ️ Review info⚙️ Run configurationConfiguration used: defaults Review profile: CHILL Plan: Pro Run ID: 📒 Files selected for processing (1)
📝 WalkthroughWalkthroughAdds two new async storage backends (Chroma and LanceDB), wires them into the storage factory and config, updates optional dependencies, and adds unit and integration tests covering sentences, facts, episodes, sessions, relationships, supersession chains, and GDPR-style user deletion. Changes
Sequence Diagram(s)sequenceDiagram
participant App as Application
participant Factory as create_storage
participant Backend as StorageBackend (Chroma/LanceDB)
participant VDB as Vector DB
App->>Factory: create_storage(config with storage_backend)
activate Factory
Factory-->>Backend: instantiate ChromaBackend/LanceDBBackend
deactivate Factory
App->>Backend: search_facts(embedding, user_id, ...)
activate Backend
Backend->>VDB: vector similarity query + metadata filters
activate VDB
VDB-->>Backend: raw results (+ distances)
deactivate VDB
Backend->>Backend: normalize rows, apply post-filters, resolve supersession chain
Backend-->>App: results list
deactivate Backend
sequenceDiagram
participant App as Application
participant Backend as StorageBackend (Chroma/LanceDB)
participant VDB as Vector DB
App->>Backend: upsert_sentences(sentences, embeddings, user_id)
activate Backend
Backend->>VDB: upsert/merge vectors + metadata
activate VDB
VDB-->>Backend: success / ids
deactivate VDB
Backend-->>App: count / ids inserted
deactivate Backend
Estimated code review effort🎯 4 (Complex) | ⏱️ ~70 minutes Possibly related PRs
Poem
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
There was a problem hiding this comment.
Actionable comments posted: 3
🧹 Nitpick comments (3)
vektori/storage/factory.py (1)
12-15: Update docstring to include new backends.The docstring at line 13 lists only
"sqlite", "postgres", "memory", "neo4j", "qdrant"but doesn't mention the newly added"chroma"and"lancedb"backends.📝 Suggested fix
"""Resolve and initialize the correct storage backend from config. Backend selection priority: - 1. config.storage_backend key ("sqlite", "postgres", "memory", "neo4j", "qdrant") + 1. config.storage_backend key ("sqlite", "postgres", "memory", "neo4j", "qdrant", "chroma", "lancedb") 2. URL prefix heuristic (postgresql://, bolt://, neo4j://, http://localhost:6333) """🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@vektori/storage/factory.py` around lines 12 - 15, Update the module docstring in vektori/storage/factory.py (the Backend selection priority docstring) to list the newly supported backends "chroma" and "lancedb" alongside the existing entries ("sqlite", "postgres", "memory", "neo4j", "qdrant") so the documentation matches the implemented backends.vektori/storage/lancedb_backend.py (1)
250-256:find_sentences_by_similarityreturns empty list.Same as ChromaBackend — this method is stubbed to return
[]. If this is intentional as a fallback strategy, consider adding a comment for clarity.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@vektori/storage/lancedb_backend.py` around lines 250 - 256, The method find_sentences_by_similarity currently returns an empty list as a stub; either implement its similarity search logic to mirror ChromaBackend (use the same embedding/nearest-neighbor query flow and the session_id/threshold parameters) or, if returning an empty list is intentional as a fallback, replace the stub with a clear docstring/comment stating that this backend does not support sentence similarity and why. Update the function find_sentences_by_similarity in lancedb_backend.py to either call the implemented LanceDB query routine (using the same input symbols as other retrieval methods) or add the explanatory comment so future readers know this is deliberate.vektori/storage/chroma_backend.py (1)
214-214: Consider replacing deprecateddatetime.utcnow().
datetime.utcnow()is deprecated as of Python 3.12 and will emitDeprecationWarning. While the project currently targets Python 3.10+, consider using the timezone-aware alternative for future compatibility.This same pattern appears at line 336, 703, and 825.
📝 Suggested fix (one example)
+from datetime import datetime, timezone ... - "created_at": datetime.utcnow().isoformat(), + "created_at": datetime.now(timezone.utc).isoformat(),🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@vektori/storage/chroma_backend.py` at line 214, Replace uses of the deprecated naive timestamp creation "datetime.utcnow().isoformat()" in chroma_backend.py (the metadata "created_at" assignments) with a timezone-aware call like "datetime.now(timezone.utc).isoformat()" and import timezone from datetime; update every occurrence (the ones currently at the "created_at" assignments around the reported spots) to ensure timestamps are timezone-aware and avoid DeprecationWarning.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@tests/unit/test_chroma_lancedb_factory.py`:
- Around line 3-12: The import block in
tests/unit/test_chroma_lancedb_factory.py is mis-ordered and includes an unused
pytest import; remove the unused "import pytest" and reorder the imports into
canonical groups (future -> stdlib -> 3rd-party -> local) so they are sorted
(e.g., keep "from __future__ import annotations" first, then "from unittest.mock
import AsyncMock, patch", then third-party imports like pytest if present — but
here remove it — then local vektori imports: VektoriConfig, ChromaBackend,
create_storage, LanceDBBackend); run your formatter/isort to ensure the import
order is compliant with I001.
In `@vektori/storage/chroma_backend.py`:
- Around line 258-264: The find_sentences_by_similarity method in
chroma_backend.py is an intentional no-op and should include a clarifying
comment; update the async def find_sentences_by_similarity(self, quotes,
session_id, threshold=0.75) implementation to document that it is superseded by
search_sentences_in_session (embedding-based) and then return an empty list,
i.e., add a one-line comment like "Superseded by search_sentences_in_session
(embedding-based)." immediately before the existing return [] to match other
backend implementations.
In `@vektori/storage/lancedb_backend.py`:
- Around line 360-363: The date-filtering logic that appends to the local
variable where when before_date/after_date are set should exclude empty-string
event_time values to avoid incorrect string comparisons; update the branches
that currently append "AND event_time <= '...'" and "AND event_time >= '...'" to
also require a non-empty event_time (e.g., append "AND event_time != '' AND
event_time <= '...'" and "AND event_time != '' AND event_time >= '...'") or
otherwise add a single "AND event_time != ''" conjunct to the composed where
clause so facts with empty event_time are not matched; locate the code that
references the where variable and the before_date/after_date parameters in
lancedb_backend.py and make the change there.
---
Nitpick comments:
In `@vektori/storage/chroma_backend.py`:
- Line 214: Replace uses of the deprecated naive timestamp creation
"datetime.utcnow().isoformat()" in chroma_backend.py (the metadata "created_at"
assignments) with a timezone-aware call like
"datetime.now(timezone.utc).isoformat()" and import timezone from datetime;
update every occurrence (the ones currently at the "created_at" assignments
around the reported spots) to ensure timestamps are timezone-aware and avoid
DeprecationWarning.
In `@vektori/storage/factory.py`:
- Around line 12-15: Update the module docstring in vektori/storage/factory.py
(the Backend selection priority docstring) to list the newly supported backends
"chroma" and "lancedb" alongside the existing entries ("sqlite", "postgres",
"memory", "neo4j", "qdrant") so the documentation matches the implemented
backends.
In `@vektori/storage/lancedb_backend.py`:
- Around line 250-256: The method find_sentences_by_similarity currently returns
an empty list as a stub; either implement its similarity search logic to mirror
ChromaBackend (use the same embedding/nearest-neighbor query flow and the
session_id/threshold parameters) or, if returning an empty list is intentional
as a fallback, replace the stub with a clear docstring/comment stating that this
backend does not support sentence similarity and why. Update the function
find_sentences_by_similarity in lancedb_backend.py to either call the
implemented LanceDB query routine (using the same input symbols as other
retrieval methods) or add the explanatory comment so future readers know this is
deliberate.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro
Run ID: 236de056-b5fb-4e34-a48f-5ef6509ed6ac
📒 Files selected for processing (8)
pyproject.tomltests/integration/test_chroma_backend.pytests/integration/test_lancedb_backend.pytests/unit/test_chroma_lancedb_factory.pyvektori/config.pyvektori/storage/chroma_backend.pyvektori/storage/factory.pyvektori/storage/lancedb_backend.py
| async def find_sentences_by_similarity( | ||
| self, | ||
| quotes: list[str], | ||
| session_id: str, | ||
| threshold: float = 0.75, | ||
| ) -> list[str]: | ||
| return [] |
There was a problem hiding this comment.
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
# Check how find_sentences_by_similarity is used in the codebase
rg -n "find_sentences_by_similarity" --type py -C 2Repository: vektori-ai/vektori
Length of output: 2402
🏁 Script executed:
# Check implementations in other backends
echo "=== QDRANT Backend ===" && sed -n '259,275p' vektori/storage/qdrant_backend.py
echo ""
echo "=== SQLITE Backend ===" && sed -n '229,245p' vektori/storage/sqlite.py
echo ""
echo "=== MEMORY Backend ===" && sed -n '96,112p' vektori/storage/memory.pyRepository: vektori-ai/vektori
Length of output: 1661
🏁 Script executed:
# Check the chroma backend implementation in full context
echo "=== CHROMA Backend Full Implementation ===" && sed -n '256,270p' vektori/storage/chroma_backend.pyRepository: vektori-ai/vektori
Length of output: 435
🏁 Script executed:
# Check if find_sentences_by_similarity is called anywhere in the codebase (usage)
rg -n "\.find_sentences_by_similarity\(" --type py -B 2 -A 2Repository: vektori-ai/vektori
Length of output: 44
Add clarifying comment to document intentional stub.
The method stub is consistent with other backend implementations (SQLite, Memory) where find_sentences_by_similarity is superseded by search_sentences_in_session (embedding-based). However, unlike those backends, the chroma backend lacks the explanatory comment. Add:
# Superseded by search_sentences_in_session (embedding-based).
return []🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@vektori/storage/chroma_backend.py` around lines 258 - 264, The
find_sentences_by_similarity method in chroma_backend.py is an intentional no-op
and should include a clarifying comment; update the async def
find_sentences_by_similarity(self, quotes, session_id, threshold=0.75)
implementation to document that it is superseded by search_sentences_in_session
(embedding-based) and then return an empty list, i.e., add a one-line comment
like "Superseded by search_sentences_in_session (embedding-based)." immediately
before the existing return [] to match other backend implementations.
| if before_date is not None: | ||
| where += f" AND event_time <= '{before_date.isoformat()}'" | ||
| if after_date is not None: | ||
| where += f" AND event_time >= '{after_date.isoformat()}'" |
There was a problem hiding this comment.
Date filtering may miss facts with empty event_time.
When before_date or after_date is specified, the SQL filter compares against event_time directly. Facts with empty event_time ("") will be included or excluded based on string comparison, which may produce unexpected results.
Consider adding an explicit check for non-empty event_time:
🔧 Suggested fix
if before_date is not None:
- where += f" AND event_time <= '{before_date.isoformat()}'"
+ where += f" AND event_time != '' AND event_time <= '{before_date.isoformat()}'"
if after_date is not None:
- where += f" AND event_time >= '{after_date.isoformat()}'"
+ where += f" AND event_time != '' AND event_time >= '{after_date.isoformat()}'"📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| if before_date is not None: | |
| where += f" AND event_time <= '{before_date.isoformat()}'" | |
| if after_date is not None: | |
| where += f" AND event_time >= '{after_date.isoformat()}'" | |
| if before_date is not None: | |
| where += f" AND event_time != '' AND event_time <= '{before_date.isoformat()}'" | |
| if after_date is not None: | |
| where += f" AND event_time != '' AND event_time >= '{after_date.isoformat()}'" |
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@vektori/storage/lancedb_backend.py` around lines 360 - 363, The
date-filtering logic that appends to the local variable where when
before_date/after_date are set should exclude empty-string event_time values to
avoid incorrect string comparisons; update the branches that currently append
"AND event_time <= '...'" and "AND event_time >= '...'" to also require a
non-empty event_time (e.g., append "AND event_time != '' AND event_time <=
'...'" and "AND event_time != '' AND event_time >= '...'") or otherwise add a
single "AND event_time != ''" conjunct to the composed where clause so facts
with empty event_time are not matched; locate the code that references the where
variable and the before_date/after_date parameters in lancedb_backend.py and
make the change there.
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@tests/integration/test_chroma_backend.py`:
- Around line 30-33: The teardown currently yields backend then calls await
backend.delete_user("test-user") and await
backend.delete_user("delete-test-user") before await backend.close(), which
means backend.close() may be skipped if a delete_user raises; wrap the
post-yield cleanup in a try/finally (or use an async contextmanager) so that
await backend.close() is always executed in the finally block; ensure
delete_user calls run inside the try and any exceptions are caught/logged (or
suppressed) so they don't prevent the finally from calling backend.close().
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro
Run ID: ebdf3569-6a17-49e1-a608-2391bc9fe117
📒 Files selected for processing (3)
tests/integration/test_chroma_backend.pytests/integration/test_lancedb_backend.pyvektori/config.py
✅ Files skipped from review due to trivial changes (1)
- vektori/config.py
🚧 Files skipped from review as they are similar to previous changes (1)
- tests/integration/test_lancedb_backend.py
| yield backend | ||
| await backend.delete_user("test-user") | ||
| await backend.delete_user("delete-test-user") | ||
| await backend.close() |
There was a problem hiding this comment.
Ensure backend close() always runs during teardown.
If either cleanup call fails at Line 31 or Line 32, Line 33 won’t run, which can leak resources and cause cascading test flakiness.
Proposed teardown hardening
`@pytest.fixture`
async def chroma_backend():
@@
backend = ChromaBackend(prefix="test", embedding_dim=4)
await backend.initialize()
- yield backend
- await backend.delete_user("test-user")
- await backend.delete_user("delete-test-user")
- await backend.close()
+ try:
+ yield backend
+ finally:
+ try:
+ await backend.delete_user("test-user")
+ await backend.delete_user("delete-test-user")
+ finally:
+ await backend.close()📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| yield backend | |
| await backend.delete_user("test-user") | |
| await backend.delete_user("delete-test-user") | |
| await backend.close() | |
| try: | |
| yield backend | |
| finally: | |
| try: | |
| await backend.delete_user("test-user") | |
| await backend.delete_user("delete-test-user") | |
| finally: | |
| await backend.close() |
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@tests/integration/test_chroma_backend.py` around lines 30 - 33, The teardown
currently yields backend then calls await backend.delete_user("test-user") and
await backend.delete_user("delete-test-user") before await backend.close(),
which means backend.close() may be skipped if a delete_user raises; wrap the
post-yield cleanup in a try/finally (or use an async contextmanager) so that
await backend.close() is always executed in the finally block; ensure
delete_user calls run inside the try and any exceptions are caught/logged (or
suppressed) so they don't prevent the finally from calling backend.close().
Summary
ChromaBackend— embedded/HTTP ChromaDB backend (closes feat(storage): Chroma vector backend #45)LanceDBBackend— serverless columnar LanceDB backend (closes feat(storage): LanceDB vector backend #46)StorageBackendinterface: sentences, facts, episodes, sessions, GDPR deletestorage_backend="chroma"/storage_backend="lancedb"pip install 'vektori[chroma]'andpip install 'vektori[lancedb]'Chroma modes
ChromaBackend()— no path, no hostChromaBackend(path="/my/db")ordatabase_url="/my/db"ChromaBackend(host="localhost", port=8000)ordatabase_url="http://host:8000"LanceDB modes
LanceDBBackend(uri="/my/lancedb")LanceDBBackend(uri="s3://bucket/path")Test plan
tests/unit/test_chroma_lancedb_factory.pyEphemeralClient(no server required) — all passtmp_pathfixture — auto-skip when not installedruff check+ruff formatclean🤖 Generated with Claude Code
Summary by CodeRabbit