feat: Integrate VecLite as a native vector storage backend - #51
Merged
Conversation
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
📝 WalkthroughWalkthroughAdds a new VecLite storage backend implementation and tests, registers it in the storage factory, declares a Changes
Sequence Diagram(s)sequenceDiagram
participant Client
participant Factory
participant VecLiteBackend as "VecLiteBackend\n(vektori/storage)"
participant VecLiteDB as "veclite DB files"
Client->>Factory: create_storage(config)
Factory->>VecLiteBackend: instantiate(path)
VecLiteBackend->>VecLiteDB: initialize() (open/create DB files)
Client->>VecLiteBackend: upsert_sentences / insert_fact / search_*
VecLiteBackend->>VecLiteDB: write embeddings/records / similarity query
VecLiteDB-->>VecLiteBackend: results
VecLiteBackend-->>Client: result objects
Estimated code review effort🎯 4 (Complex) | ⏱️ ~45 minutes Possibly related issues
Possibly related PRs
Poem
🚥 Pre-merge checks | ✅ 2 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (2 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 |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 9
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@pyproject.toml`:
- Around line 33-39: Remove test-only packages "pytest" and "pytest-asyncio" and
the optional backend "veclite-db" from the core dependencies list (the top-level
dependencies array) so they are not installed for every user; ensure "pytest"
and "pytest-asyncio" remain only in the dev extras and "veclite-db" remains
under the [project.optional-dependencies] -> veclite extra. If you intend an
"all" extra to install every backend, add "veclite-db>=1.0.9" to that extra as
well so backends are optional unless explicitly requested.
In `@tests/unit/test_veclite_backend.py`:
- Around line 1-3: Remove the unused stdlib imports and stray whitespace-only
blank lines at the top of the test module: delete the unused imports (datetime,
timezone, os) from the import block and keep only the used import(s) like
pytest, and remove any blank lines that contain only spaces so there are no
trailing whitespace-only lines (also apply the same cleanup to the similar
import/blank-line region around lines 15-29).
In `@vektori/storage/factory.py`:
- Around line 78-86: The code currently checks URL heuristics before honoring an
explicit backend_key, so storage_backend="veclite" can be ignored when
database_url matches an earlier scheme; update the logic so that an explicit
backend_key wins: either move the branch that tests backend_key == "veclite" to
run before any database_url scheme checks, or split the condition so that if
backend_key == "veclite" you always instantiate VecLiteBackend(path) regardless
of database_url; reference backend_key, database_url and VecLiteBackend when
making the change so the explicit-key path is evaluated first.
In `@vektori/storage/veclite_backend.py`:
- Around line 105-111: The current stubbed method find_sentences_by_similarity
returns an empty list which silently disables paraphrase matching; replace that
behavior by either implementing a real fallback or raising a clear error: if
VecLite backend cannot perform paraphrase/similarity lookups yet, change
find_sentences_by_similarity to raise NotImplementedError (or a custom
BackendNotSupportedError) with a concise message like "VecLite backend does not
support paraphrase similarity lookups"; otherwise implement the fallback inside
find_sentences_by_similarity by computing embeddings for the input quotes,
querying the VecLite similarity/index API for matches above the threshold, and
returning the matched sentence strings — reference the
find_sentences_by_similarity method to locate where to make this change.
- Around line 438-445: delete_user currently only removes in-memory dict entries
(_sentences, _facts, _episodes, _sessions) but leaves persisted VecLite vectors
and related mappings (_edges, _fact_sources, _episode_facts) on disk; update
delete_user to collect all vector ids associated with the user, remove those ids
from the VecLite indices (call the VecLite index client/delete API or mark them
deleted), and delete entries from _edges, _fact_sources, and _episode_facts that
reference those ids; after removals, persist or flush the index and any
file-backed stores so on-disk state matches the in-memory deletions and return
the total removed count.
- Around line 30-53: The backend currently keeps metadata only in memory
(_sentences, _facts, _edges, _fact_sources, _sessions, _episodes,
_episode_facts) and never persists or reloads it in initialize() / close(),
causing lost joins across restarts; fix by serializing these structures to disk
(e.g., JSON files like sentences.json, facts.json, edges.json,
fact_sources.json, sessions.json, episodes.json, episode_facts.json inside
self.base_path) whenever closing or mutating, and load/validate them in
initialize() before using self._vec_sentences / self._vec_facts /
self._vec_episodes so IDs line up; ensure atomic writes (tmp file + rename) and
graceful handling of missing/corrupt files (recreate empty structures) and keep
file names and keys matching the in-memory dict/list shapes used by methods that
reference _sentences, _facts, _edges, _fact_sources, _sessions, _episodes, and
_episode_facts.
- Around line 89-103: The code currently calls _vec_sentences.search and then
post-filters results by user_id/agent_id/session_id/subject (using
_sentences.get and checks) which can drop most matches if the shared index's
top-N is dominated by other tenants; update the search logic in this method (and
the similar blocks at the other ranges you noted) to either (a) perform
per-scope partitioned searches when possible (e.g., call the vector index with a
tenant/agent/session-specific partition or namespace) or (b) implement an
iterative-expand-and-rerank strategy: repeatedly call _vec_sentences.search with
increasing candidate limits (or fetch all candidates) and accumulate filtered
matches from _sentences until you reach the desired limit or exhaust the index;
ensure you reference and preserve the existing distance computation (1.0 - sim)
and break when results length >= limit to avoid returning incomplete scoped
results.
- Around line 217-229: get_active_facts currently ignores the optional agent_id,
causing cross-agent reads; update get_active_facts (and the other fact-reading
method get_all_facts) to filter facts by agent when agent_id is provided by
adding a condition like (agent_id is None or f.get("agent_id") == agent_id) to
the list comprehension (while keeping the existing user_id and is_active checks)
so slicing/limit/offset behavior remains unchanged.
- Around line 57-65: In upsert_sentences, add a pre-loop check that sentences
and embeddings have the same length and fail fast if they don't: compare
len(sentences) and len(embeddings) before the for loop and raise a clear
exception (e.g., ValueError) including both lengths and any identifying context
(user_id/agent_id) so mismatched responses are not silently truncated by zip;
keep the rest of the logic (count, loop over sent/emb) unchanged.
🪄 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: 76b83555-1b5e-4892-ba76-30511c3f74f9
⛔ Files ignored due to path filters (1)
uv.lockis excluded by!**/*.lock
📒 Files selected for processing (5)
pyproject.tomltests/unit/test_new_backends_factory.pytests/unit/test_veclite_backend.pyvektori/storage/factory.pyvektori/storage/veclite_backend.py
There was a problem hiding this comment.
Actionable comments posted: 2
🧹 Nitpick comments (2)
vektori/storage/veclite_backend.py (2)
454-457:insert_fact_sourcelacks deduplication, unlikeinsert_episode_fact.This method always appends without checking for existing links, while the similar
insert_episode_factmethod (lines 403-409) properly deduplicates. This inconsistency could lead to duplicate fact-source entries on retries.♻️ Proposed fix for consistency with insert_episode_fact
async def insert_fact_source(self, fact_id: str, sentence_id: str) -> None: + for link in self._fact_sources: + if link["fact_id"] == fact_id and link["sentence_id"] == sentence_id: + return link = {"fact_id": fact_id, "sentence_id": sentence_id} self._fact_sources.append(link) self._sync_relation(f"fs:{uuid.uuid4()}", link)🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@vektori/storage/veclite_backend.py` around lines 454 - 457, The insert_fact_source method currently always appends a new link and can create duplicates on retries; update insert_fact_source to mirror insert_episode_fact by building the link {"fact_id": fact_id, "sentence_id": sentence_id"}, checking if that link already exists in self._fact_sources before appending, and only call self._sync_relation(f"fs:{uuid.uuid4()}", link) when you actually added the link to avoid duplicate entries and duplicate sync calls. Ensure you reference insert_fact_source and insert_episode_fact so the deduplication logic is consistent between them.
346-350:insert_edgeslacks deduplication.The base interface contract specifies "ON CONFLICT DO NOTHING" but this implementation always appends, potentially creating duplicate edges. While edges are typically inserted once per session, repeated processing or retries could accumulate duplicates.
♻️ Proposed fix to add deduplication
async def insert_edges(self, edges: list[dict[str, Any]]) -> int: + existing = {(e.get("source"), e.get("target"), e.get("type")) for e in self._edges} + count = 0 for edge in edges: + key = (edge.get("source"), edge.get("target"), edge.get("type")) + if key in existing: + continue self._edges.append(edge) + existing.add(key) self._sync_relation(f"edge:{uuid.uuid4()}", edge) - return len(edges) + count += 1 + return count🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@vektori/storage/veclite_backend.py` around lines 346 - 350, insert_edges currently always appends every edge (and calls _sync_relation with a fresh uuid), causing duplicate edges; change it to deduplicate before appending by computing a canonical representation (e.g., json.dumps(edge, sort_keys=True) or equivalent) and tracking seen representations in a set, only appending and calling _sync_relation for edges whose representation is not already present in self._edges; also use a stable/deterministic key for _sync_relation (derived from the canonical representation, e.g., a hash or uuid5) instead of uuid.uuid4(), and return the number of edges actually inserted (unique adds) rather than the input length.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@vektori/storage/veclite_backend.py`:
- Around line 561-579: delete_user currently filters in-memory lists (_edges,
_fact_sources, _episode_facts) but does not remove corresponding persisted
entries in _vec_relations, leaving orphaned relation rows; fix by, after
collecting sentence_ids, fact_ids, episode_ids and after the existing in-memory
filters in delete_user, iterate through self._vec_relations.get_all() (guarded
by hasattr(self._vec_relations, "delete")), parse each meta JSON, and call
self._vec_relations.delete(rid) for relation ids whose prefix is "edge:", "fs:"
or "ef:" and whose parsed fields (source/target for edges, fact_id/sentence_id
for fs, episode_id/fact_id for ef) match the deleted id sets so persisted
relations are removed when their referenced entities were deleted.
- Around line 29-31: The test imports VecLiteBackend even when the optional
veclite-db package isn't installed; update the test to call
pytest.importorskip("veclite") before importing VecLiteBackend so the test is
skipped when the veclite module is missing, ensuring the import of
vektori.storage.veclite_backend.VecLiteBackend only happens after the
skip-check.
---
Nitpick comments:
In `@vektori/storage/veclite_backend.py`:
- Around line 454-457: The insert_fact_source method currently always appends a
new link and can create duplicates on retries; update insert_fact_source to
mirror insert_episode_fact by building the link {"fact_id": fact_id,
"sentence_id": sentence_id"}, checking if that link already exists in
self._fact_sources before appending, and only call
self._sync_relation(f"fs:{uuid.uuid4()}", link) when you actually added the link
to avoid duplicate entries and duplicate sync calls. Ensure you reference
insert_fact_source and insert_episode_fact so the deduplication logic is
consistent between them.
- Around line 346-350: insert_edges currently always appends every edge (and
calls _sync_relation with a fresh uuid), causing duplicate edges; change it to
deduplicate before appending by computing a canonical representation (e.g.,
json.dumps(edge, sort_keys=True) or equivalent) and tracking seen
representations in a set, only appending and calling _sync_relation for edges
whose representation is not already present in self._edges; also use a
stable/deterministic key for _sync_relation (derived from the canonical
representation, e.g., a hash or uuid5) instead of uuid.uuid4(), and return the
number of edges actually inserted (unique adds) rather than the input length.
🪄 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: b09253bd-0808-4422-a784-3fd5b91b4df6
⛔ Files ignored due to path filters (1)
uv.lockis excluded by!**/*.lock
📒 Files selected for processing (4)
pyproject.tomltests/unit/test_veclite_backend.pyvektori/storage/factory.pyvektori/storage/veclite_backend.py
✅ Files skipped from review due to trivial changes (1)
- pyproject.toml
🚧 Files skipped from review as they are similar to previous changes (2)
- vektori/storage/factory.py
- tests/unit/test_veclite_backend.py
Contributor
|
Hey @rithulkamesh, thanks so much for this contribution! The code looks really solid and the integration with the factory pattern and StorageBackend interface is super clean. there is one minor issues that should be fixed before merging: |
There was a problem hiding this comment.
Actionable comments posted: 2
♻️ Duplicate comments (1)
vektori/storage/veclite_backend.py (1)
319-328:⚠️ Potential issue | 🟠 Major
find_fact_by_textstill ignoresagent_id.
get_active_factswas updated at line 300 to filter byagent_id, but this method — which is used for write-time dedup — was not. Two agents writing the same fact text for the same user will dedupe against each other, breaking(user_id, agent_id)isolation and causing one agent to silently adopt the other agent's fact row.🔐 Suggested fix
async def find_fact_by_text( self, user_id: str, text: str, agent_id: str | None = None, ) -> dict[str, Any] | None: for f in self._facts.values(): - if f.get("user_id") == user_id and f.get("text") == text and f.get("is_active", True): + if ( + f.get("user_id") == user_id + and f.get("text") == text + and f.get("is_active", True) + and (agent_id is None or f.get("agent_id") == agent_id) + ): return f return None🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@vektori/storage/veclite_backend.py` around lines 319 - 328, The find_fact_by_text method currently ignores agent_id and thus can dedupe facts across agents; update the loop in find_fact_by_text to also compare f.get("agent_id") (treating None appropriately) against the agent_id parameter and only return when user_id, text, is_active, and agent_id match; ensure the function signature (find_fact_by_text) and callers that pass agent_id still work and that agent_id=None behavior mirrors get_active_facts' filtering logic.
🧹 Nitpick comments (1)
tests/unit/test_veclite_backend.py (1)
1-44: Add tests for the two fixes this PR advertises.The PR description highlights (1) durability via VecLite-backed metadata and (2) the ghost-data fix where relation deletes are now swept from
relations.dbondelete_user. Neither is exercised here — the only test covers the happy-path ofupsert_sentences+search_sentences. Without regression coverage, a future refactor can silently reintroduce the exact bugs this PR closed.Worth adding, at minimum:
- A restart-persistence test: open
VecLiteBackend(path), upsert sentences/facts/edges,close(), open a fresh instance at the same path, callinitialize(), assert_sentences/_facts/_edges/_fact_sources/_episode_factsare repopulated andsearch_sentencesstill returns the record.- A
delete_usercascade test: seed sentences + facts + an edge + a fact_source + an episode_fact, calldelete_user, then reopen the backend and assert there are no ghostedge:/fs:/ef:/sess:rows in_vec_relations.get_all().Want me to draft those two tests?
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@tests/unit/test_veclite_backend.py` around lines 1 - 44, Add two new async tests using VecLiteBackend: (1) restart-persistence: create VecLiteBackend(path=temp_dir), call initialize(), upsert sentences/facts/edges/fact_sources/episode_facts using existing methods (e.g. upsert_sentences and whatever upsert_* helpers exist), close() then create a fresh VecLiteBackend at same path, call initialize(), and assert internal maps _sentences/_facts/_edges/_fact_sources/_episode_facts are repopulated and search_sentences returns the previously inserted sentence; (2) delete_user cascade: seed sentences, facts, an edge, a fact_source and an episode_fact for a user, call delete_user(user_id), then reopen backend (initialize()) and assert _vec_relations.get_all() contains no rows starting with the relation prefixes "edge:", "fs:", "ef:", or "sess:" (and that corresponding _edges/_fact_sources/_episode_facts/_sentences entries are removed). Ensure tests use the same temp_dir fixture and close() between reopenings.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@vektori/storage/veclite_backend.py`:
- Around line 587-608: The loop mutates self._vec_relations while iterating
self._vec_relations.get_all(), which can invalidate the iterator; instead, first
build a list/set of relation IDs to remove by iterating get_all() and checking
the same conditions (use sentence_ids, fact_ids, episode_ids ->
sent_set/fact_set/ep_set and inspect rid, meta_str, data), then after the
iteration complete call self._vec_relations.delete(rid) for each collected id;
keep the same rid prefixes ("edge:", "fs:", "ef:") and JSON parsing logic but
perform deletions only after collecting matching rids.
---
Duplicate comments:
In `@vektori/storage/veclite_backend.py`:
- Around line 319-328: The find_fact_by_text method currently ignores agent_id
and thus can dedupe facts across agents; update the loop in find_fact_by_text to
also compare f.get("agent_id") (treating None appropriately) against the
agent_id parameter and only return when user_id, text, is_active, and agent_id
match; ensure the function signature (find_fact_by_text) and callers that pass
agent_id still work and that agent_id=None behavior mirrors get_active_facts'
filtering logic.
---
Nitpick comments:
In `@tests/unit/test_veclite_backend.py`:
- Around line 1-44: Add two new async tests using VecLiteBackend: (1)
restart-persistence: create VecLiteBackend(path=temp_dir), call initialize(),
upsert sentences/facts/edges/fact_sources/episode_facts using existing methods
(e.g. upsert_sentences and whatever upsert_* helpers exist), close() then create
a fresh VecLiteBackend at same path, call initialize(), and assert internal maps
_sentences/_facts/_edges/_fact_sources/_episode_facts are repopulated and
search_sentences returns the previously inserted sentence; (2) delete_user
cascade: seed sentences, facts, an edge, a fact_source and an episode_fact for a
user, call delete_user(user_id), then reopen backend (initialize()) and assert
_vec_relations.get_all() contains no rows starting with the relation prefixes
"edge:", "fs:", "ef:", or "sess:" (and that corresponding
_edges/_fact_sources/_episode_facts/_sentences entries are removed). Ensure
tests use the same temp_dir fixture and close() between reopenings.
🪄 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: fdaf656b-e6b2-477a-988a-8c35a2c11d90
📒 Files selected for processing (2)
tests/unit/test_veclite_backend.pyvektori/storage/veclite_backend.py
…d persistence tests
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
This PR introduces VecLite (via the
veclite-dbPython package) as a native Database Storage Provider for Vektori. This enables users to leverage VecLite's fast, embedded vector search capabilities seamlessly within the Vektori ecosystem.Key Changes
veclite_backend.py): Implemented theVecLiteBackendwhich satisfies theStorageBackendinterface. It utilizesveclite.VecLite()instances for fast vector similarity search (sentences.db,facts.db,episodes.db) while maintaining metadata state.factory.py): Updated the storage factory to automatically resolve and initialize the VecLite backend when the configuration specifiesstorage_backend="veclite"or uses aveclite://<path>connection string.veclite-db>=1.0.9as an optional dependency inpyproject.toml.tests/unit/test_veclite_backend.pyto verify vector insertion and cosine similarity search behavior.tests/unit/test_new_backends_factory.pyto validate correct routing and path parsing for VecLite URLs.How to test
You can now initialize a Vektori client using VecLite by passing the appropriate config:
Summary by CodeRabbit
New Features
Tests
Chores