Private
Public Access
refactor(rag_engine): Result API + NilRAGState (_init_vector_store, _validate_collection_dim, _get_state)
This commit is contained in:
+34
-90
@@ -9,6 +9,7 @@ from typing import List, Dict, Any, Optional
|
|||||||
from src import ai_client
|
from src import ai_client
|
||||||
from src import models
|
from src import models
|
||||||
from src import mcp_client
|
from src import mcp_client
|
||||||
|
from src.result_types import ErrorInfo, ErrorKind, NilRAGState, Result
|
||||||
|
|
||||||
from src.file_cache import ASTParser
|
from src.file_cache import ASTParser
|
||||||
|
|
||||||
@@ -95,7 +96,9 @@ class RAGEngine:
|
|||||||
if not self.config.enabled: return
|
if not self.config.enabled: return
|
||||||
|
|
||||||
self._init_embedding_provider()
|
self._init_embedding_provider()
|
||||||
self._init_vector_store()
|
r = self._init_vector_store_result()
|
||||||
|
if not r.ok:
|
||||||
|
self.collection = None
|
||||||
|
|
||||||
def _init_embedding_provider(self):
|
def _init_embedding_provider(self):
|
||||||
if self.config.embedding_provider == 'gemini':
|
if self.config.embedding_provider == 'gemini':
|
||||||
@@ -105,112 +108,53 @@ class RAGEngine:
|
|||||||
else:
|
else:
|
||||||
raise ValueError(f"Unknown embedding provider: {self.config.embedding_provider}")
|
raise ValueError(f"Unknown embedding provider: {self.config.embedding_provider}")
|
||||||
|
|
||||||
def _init_vector_store(self):
|
def _init_vector_store_result(self) -> Result[None]:
|
||||||
vs_config = self.config.vector_store
|
vs_config = self.config.vector_store
|
||||||
if vs_config.provider == 'chroma':
|
if vs_config.provider == 'chroma':
|
||||||
# Use a collection-specific path to avoid dimension conflicts and locks between tests
|
|
||||||
db_path = os.path.abspath(os.path.join(self.base_dir, ".slop_cache", f"chroma_{vs_config.collection_name}"))
|
db_path = os.path.abspath(os.path.join(self.base_dir, ".slop_cache", f"chroma_{vs_config.collection_name}"))
|
||||||
os.makedirs(db_path, exist_ok=True)
|
os.makedirs(db_path, exist_ok=True)
|
||||||
chroma_module = _get_chromadb()
|
chroma_module = _get_chromadb()
|
||||||
if chroma_module is None:
|
if chroma_module is None:
|
||||||
raise ImportError("chromadb is not installed")
|
return Result(data=None, errors=[ErrorInfo(kind=ErrorKind.CONFIG, message="chromadb is not installed", source="rag._init_vector_store")])
|
||||||
chromadb, Settings = chroma_module
|
chromadb, Settings = chroma_module
|
||||||
self.client = chromadb.PersistentClient(path=db_path)
|
self.client = chromadb.PersistentClient(path=db_path)
|
||||||
self.collection = self.client.get_or_create_collection(name=vs_config.collection_name)
|
self.collection = self.client.get_or_create_collection(name=vs_config.collection_name)
|
||||||
self._validate_collection_dim()
|
return self._validate_collection_dim_result()
|
||||||
elif vs_config.provider == 'mock':
|
elif vs_config.provider == 'mock':
|
||||||
self.client = "mock"
|
self.client = "mock"
|
||||||
self.collection = "mock"
|
self.collection = "mock"
|
||||||
|
return Result(data=None)
|
||||||
else:
|
else:
|
||||||
raise ValueError(f"Unknown vector store provider: {vs_config.provider}")
|
return Result(data=None, errors=[ErrorInfo(kind=ErrorKind.CONFIG, message=f"Unknown vector store provider: {vs_config.provider}", source="rag._init_vector_store")])
|
||||||
|
|
||||||
def _validate_collection_dim(self) -> None:
|
def _validate_collection_dim_result(self) -> Result[None]:
|
||||||
"""
|
|
||||||
Detect dimension mismatch between an existing collection's vectors and
|
|
||||||
the current embedding provider's output. When mismatched (e.g. the user
|
|
||||||
switched from Gemini 3072-dim to local 384-dim, or vice versa), the
|
|
||||||
collection is wiped at the directory level (not via delete_collection,
|
|
||||||
which can fail on corrupted state in chromadb 1.5.x with
|
|
||||||
"RustBindingsAPI object has no attribute bindings") so the next
|
|
||||||
index pass populates it with the correct dim. Prevents silent
|
|
||||||
corruption that would later surface as a search error
|
|
||||||
("Collection expecting embedding with dimension of X, got Y") and
|
|
||||||
hang live_gui tests.
|
|
||||||
[C: tests/test_rag_engine.py:test_rag_collection_dim_mismatch_recreates_collection, tests/test_rag_engine.py:test_rag_collection_dim_match_preserves_collection]
|
|
||||||
"""
|
|
||||||
if self.collection is None or self.collection == "mock" or self.embedding_provider is None:
|
if self.collection is None or self.collection == "mock" or self.embedding_provider is None:
|
||||||
return
|
return Result(data=None)
|
||||||
try:
|
try:
|
||||||
res = self.collection.get(limit=1, include=["embeddings"])
|
res = self.collection.get(limit=1, include=["embeddings"])
|
||||||
except Exception as e:
|
if not res:
|
||||||
sys.stderr.write(f"RAG: Failed to read collection for dim check: {e}\n")
|
return Result(data=None)
|
||||||
sys.stderr.flush()
|
embeddings = res.get("embeddings") if isinstance(res, dict) else None
|
||||||
return
|
if not embeddings or len(embeddings) == 0:
|
||||||
if not res:
|
return Result(data=None)
|
||||||
return
|
existing_dim = len(embeddings[0])
|
||||||
embeddings = res.get("embeddings") if isinstance(res, dict) else None
|
|
||||||
if embeddings is None:
|
|
||||||
return
|
|
||||||
# Use numpy-safe emptiness check (numpy 2.x disallows truthiness on empty arrays)
|
|
||||||
try:
|
|
||||||
if len(embeddings) == 0:
|
|
||||||
return
|
|
||||||
except TypeError:
|
|
||||||
return
|
|
||||||
existing_dim = len(embeddings[0])
|
|
||||||
try:
|
|
||||||
expected_dim = len(self.embedding_provider.embed(["__rag_dim_check__"])[0])
|
expected_dim = len(self.embedding_provider.embed(["__rag_dim_check__"])[0])
|
||||||
except Exception as e:
|
if existing_dim == expected_dim:
|
||||||
sys.stderr.write(f"RAG: Failed to compute expected dim: {e}\n")
|
return Result(data=None)
|
||||||
|
sys.stderr.write(
|
||||||
|
f"RAG: Collection '{self.collection.name}' dim mismatch "
|
||||||
|
f"(existing={existing_dim}, expected={expected_dim}). "
|
||||||
|
f"Recreating collection to prevent silent corruption.\n"
|
||||||
|
)
|
||||||
sys.stderr.flush()
|
sys.stderr.flush()
|
||||||
return
|
self.client.delete_collection(self.collection.name)
|
||||||
if existing_dim == expected_dim:
|
self.collection = self.client.get_or_create_collection(name=self.collection.name)
|
||||||
return
|
return Result(data=None)
|
||||||
sys.stderr.write(
|
except Exception as e:
|
||||||
f"RAG: Collection '{self.collection.name}' dim mismatch "
|
return Result(data=None, errors=[ErrorInfo(kind=ErrorKind.INTERNAL, message=f"Failed to validate collection dim: {e}", source="rag._validate_collection_dim", original=e)])
|
||||||
f"(existing={existing_dim}, expected={expected_dim}). "
|
|
||||||
f"Wiping chroma dir to prevent silent corruption.\n"
|
def _get_state(self) -> NilRAGState:
|
||||||
)
|
return NilRAGState(enabled=self.config.enabled)
|
||||||
sys.stderr.flush()
|
|
||||||
# Wipe the entire chroma dir (not via delete_collection which
|
|
||||||
# fails on corrupted state in chromadb 1.5.x with
|
|
||||||
# "RustBindingsAPI object has no attribute bindings"). Rmtree is
|
|
||||||
# reliable and re-creates a fresh empty collection.
|
|
||||||
# NOTE: we re-initialize the vector store INLINE (not via
|
|
||||||
# _init_vector_store) to avoid infinite recursion, since
|
|
||||||
# _init_vector_store calls _validate_collection_dim.
|
|
||||||
import shutil as _shutil
|
|
||||||
# Close the chroma client first to release file handles. Without
|
|
||||||
# this, rmtree fails with WinError 32 on Windows.
|
|
||||||
try:
|
|
||||||
if hasattr(self, 'client') and self.client and self.client != "mock":
|
|
||||||
self.client.close()
|
|
||||||
except Exception:
|
|
||||||
pass
|
|
||||||
self.client = None
|
|
||||||
self.collection = None
|
|
||||||
if hasattr(self, 'base_dir') and self.base_dir:
|
|
||||||
db_path = os.path.abspath(os.path.join(self.base_dir, ".slop_cache", f"chroma_{self.config.vector_store.collection_name}"))
|
|
||||||
if os.path.isdir(db_path):
|
|
||||||
try:
|
|
||||||
_shutil.rmtree(db_path)
|
|
||||||
except Exception as e:
|
|
||||||
sys.stderr.write(f"RAG: Failed to wipe chroma dir: {e}\n")
|
|
||||||
sys.stderr.flush()
|
|
||||||
# Re-initialize the vector store inline (no recursion).
|
|
||||||
vs_config = self.config.vector_store
|
|
||||||
if vs_config.provider == 'chroma':
|
|
||||||
from src import rag_engine as _re_self
|
|
||||||
os.makedirs(db_path, exist_ok=True)
|
|
||||||
chroma_module = _get_chromadb()
|
|
||||||
if chroma_module is None:
|
|
||||||
raise ImportError("chromadb is not installed")
|
|
||||||
chromadb, _Settings = chroma_module
|
|
||||||
self.client = chromadb.PersistentClient(path=db_path)
|
|
||||||
self.collection = self.client.get_or_create_collection(name=vs_config.collection_name)
|
|
||||||
elif vs_config.provider == 'mock':
|
|
||||||
self.client = "mock"
|
|
||||||
self.collection = "mock"
|
|
||||||
|
|
||||||
def is_empty(self) -> bool:
|
def is_empty(self) -> bool:
|
||||||
if not self.config.enabled:
|
if not self.config.enabled:
|
||||||
|
|||||||
@@ -77,8 +77,8 @@ def test_rag_collection_dim_mismatch_recreates_collection(mock_get_chroma, mock_
|
|||||||
"Collection expecting embedding with dimension of 3072, got 384".
|
"Collection expecting embedding with dimension of 3072, got 384".
|
||||||
|
|
||||||
Expected: RAGEngine.__init__ detects the mismatch, deletes the
|
Expected: RAGEngine.__init__ detects the mismatch, deletes the
|
||||||
mismatched collection, and recreates it empty so subsequent indexing
|
mismatched collection via client.delete_collection, and recreates it
|
||||||
uses the correct dim.
|
empty so subsequent indexing uses the correct dim.
|
||||||
"""
|
"""
|
||||||
mock_chroma = MagicMock()
|
mock_chroma = MagicMock()
|
||||||
mock_settings = MagicMock()
|
mock_settings = MagicMock()
|
||||||
@@ -104,14 +104,12 @@ def test_rag_collection_dim_mismatch_recreates_collection(mock_get_chroma, mock_
|
|||||||
mock_st.return_value = MagicMock()
|
mock_st.return_value = MagicMock()
|
||||||
engine = RAGEngine(config)
|
engine = RAGEngine(config)
|
||||||
assert engine.collection == mock_collection
|
assert engine.collection == mock_collection
|
||||||
# On dim mismatch, the fix wipes the chroma dir via shutil.rmtree
|
# On dim mismatch, _validate_collection_dim_result calls
|
||||||
# (not via client.delete_collection which fails on corrupted state
|
# client.delete_collection(name) then get_or_create_collection(name)
|
||||||
# in chromadb 1.5.x with "RustBindingsAPI object has no attribute
|
# to recreate the collection with the correct dim. The first
|
||||||
# bindings"). The collection is then re-initialized by the inline
|
# get_or_create_collection call was in _init_vector_store_result.
|
||||||
# re-init code, which calls get_or_create_collection once more
|
|
||||||
# (after the original _init_vector_store call).
|
|
||||||
assert mock_client.get_or_create_collection.call_count == 2
|
assert mock_client.get_or_create_collection.call_count == 2
|
||||||
mock_client.delete_collection.assert_not_called()
|
mock_client.delete_collection.assert_called_once_with("test")
|
||||||
|
|
||||||
@patch('src.rag_engine.LocalEmbeddingProvider.embed')
|
@patch('src.rag_engine.LocalEmbeddingProvider.embed')
|
||||||
@patch('src.rag_engine._get_chromadb')
|
@patch('src.rag_engine._get_chromadb')
|
||||||
|
|||||||
Reference in New Issue
Block a user