ed/manual_slop
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

adjustments to rag engine

Browse Source
This commit is contained in:
Edward R. Gonzalez
2026-05-13 06:32:26 -04:00
parent 1a529ed750
commit 8e9725792f
1 changed files with 226 additions and 217 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/rag_engine.py
+216 -207
View File
@@ -8,248 +8,257 @@ from chromadb.config import Settings
from src import models from src import models
from src import mcp_client from src import mcp_client
try: _SENTENCE_TRANSFORMERS = None
from sentence_transformers import SentenceTransformer _GOOGLE_GENAI = None
except ImportError:
SentenceTransformer = None
from google import genai def _get_sentence_transformers():
from google.genai import types global _SENTENCE_TRANSFORMERS
from src import ai_client if _SENTENCE_TRANSFORMERS is None:
from sentence_transformers import SentenceTransformer
_SENTENCE_TRANSFORMERS = SentenceTransformer
return _SENTENCE_TRANSFORMERS
def _get_google_genai():
global _GOOGLE_GENAI
if _GOOGLE_GENAI is None:
from google import genai
from google.genai import types
_GOOGLE_GENAI = (genai, types)
return _GOOGLE_GENAI
class BaseEmbeddingProvider: class BaseEmbeddingProvider:
def embed(self, texts: List[str]) -> List[List[float]]: def embed(self, texts: List[str]) -> List[List[float]]:
raise NotImplementedError() raise NotImplementedError()
class LocalEmbeddingProvider(BaseEmbeddingProvider): class LocalEmbeddingProvider(BaseEmbeddingProvider):
def __init__(self, model_name: str = 'all-MiniLM-L6-v2'): def __init__(self, model_name: str = 'all-MiniLM-L6-v2'):
if SentenceTransformer is None: ST = _get_sentence_transformers()
raise ImportError("sentence-transformers is not installed") if ST is None:
self.model = SentenceTransformer(model_name) raise ImportError("sentence-transformers is not installed")
self.model = ST(model_name)
def embed(self, texts: List[str]) -> List[List[float]]: def embed(self, texts: List[str]) -> List[List[float]]:
embeddings = self.model.encode(texts) embeddings = self.model.encode(texts)
return embeddings.tolist() return embeddings.tolist()
class GeminiEmbeddingProvider(BaseEmbeddingProvider): class GeminiEmbeddingProvider(BaseEmbeddingProvider):
def __init__(self, model_name: str = 'text-embedding-004'): def __init__(self, model_name: str = 'text-embedding-004'):
self.model_name = model_name self.model_name = model_name
def embed(self, texts: List[str]) -> List[List[float]]: def embed(self, texts: List[str]) -> List[List[float]]:
ai_client._ensure_gemini_client() google_module = _get_google_genai()
client = ai_client._gemini_client if google_module is None:
if not client: raise ImportError("google-genai is not installed")
raise ValueError("Gemini client not initialized") genai_pkg, types = google_module
from src import ai_client
# For text-embedding-004, we can embed a batch ai_client._ensure_gemini_client()
res = client.models.embed_content( client = ai_client._gemini_client
model=self.model_name, if not client:
contents=texts, raise ValueError("Gemini client not initialized")
config=types.EmbedContentConfig(task_type="RETRIEVAL_DOCUMENT") res = client.models.embed_content(
) model=self.model_name,
return [e.values for e in res.embeddings] contents=texts,
config=types.EmbedContentConfig(task_type="RETRIEVAL_DOCUMENT")
)
return [e.values for e in res.embeddings]
class RAGEngine: class RAGEngine:
def __init__(self, config: models.RAGConfig, base_dir: str = "."): def __init__(self, config: models.RAGConfig, base_dir: str = "."):
self.config = config self.config = config
self.base_dir = base_dir self.base_dir = base_dir
self.client = None self.client = None
self.collection = None self.collection = None
self.embedding_provider = None self.embedding_provider = None
if not self.config.enabled: if not self.config.enabled:
return return
self._init_embedding_provider() self._init_embedding_provider()
self._init_vector_store() self._init_vector_store()
def _init_embedding_provider(self): def _init_embedding_provider(self):
if self.config.embedding_provider == 'gemini': if self.config.embedding_provider == 'gemini':
self.embedding_provider = GeminiEmbeddingProvider() self.embedding_provider = GeminiEmbeddingProvider()
elif self.config.embedding_provider == 'local': elif self.config.embedding_provider == 'local':
self.embedding_provider = LocalEmbeddingProvider() self.embedding_provider = LocalEmbeddingProvider()
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(self):
vs_config = self.config.vector_store vs_config = self.config.vector_store
if vs_config.provider == 'chroma': if vs_config.provider == 'chroma':
db_path = os.path.join(self.base_dir, ".slop_cache", "chroma_db") db_path = os.path.join(self.base_dir, ".slop_cache", "chroma_db")
os.makedirs(db_path, exist_ok=True) os.makedirs(db_path, exist_ok=True)
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)
elif vs_config.provider == 'mock': elif vs_config.provider == 'mock':
self.client = "mock" self.client = "mock"
self.collection = "mock" self.collection = "mock"
else: else:
raise ValueError(f"Unknown vector store provider: {vs_config.provider}") raise ValueError(f"Unknown vector store provider: {vs_config.provider}")
def is_empty(self) -> bool: def is_empty(self) -> bool:
if not self.config.enabled: if not self.config.enabled:
return True return True
if self.config.vector_store.provider == 'mock' or self.collection == "mock": if self.config.vector_store.provider == 'mock' or self.collection == "mock":
return True return True
if self.collection is None: if self.collection is None:
return True return True
return self.collection.count() == 0 return self.collection.count() == 0
def add_documents(self, ids: List[str], texts: List[str], metadatas: Optional[List[Dict[str, Any]]] = None): def add_documents(self, ids: List[str], texts: List[str], metadatas: Optional[List[Dict[str, Any]]] = None):
""" """
[C: tests/test_rag_engine.py:test_rag_engine_chroma] [C: tests/test_rag_engine.py:test_rag_engine_chroma]
""" """
if not self.config.enabled or self.collection == "mock": if not self.config.enabled or self.collection == "mock":
return return
embeddings = self.embedding_provider.embed(texts) embeddings = self.embedding_provider.embed(texts)
self.collection.upsert( self.collection.upsert(
ids=ids, ids=ids,
embeddings=embeddings, embeddings=embeddings,
documents=texts, documents=texts,
metadatas=metadatas metadatas=metadatas
) )
def _chunk_text(self, content: str) -> List[str]: def _chunk_text(self, content: str) -> List[str]:
"""Character-based chunking with overlap.""" """Character-based chunking with overlap."""
chunks = [] chunks = []
if not content: if not content:
return chunks return chunks
chunk_size = self.config.chunk_size chunk_size = self.config.chunk_size
overlap = self.config.chunk_overlap overlap = self.config.chunk_overlap
start = 0 start = 0
while start < len(content): while start < len(content):
end = start + chunk_size end = start + chunk_size
chunks.append(content[start:end]) chunks.append(content[start:end])
if end >= len(content): if end >= len(content):
break break
start += (chunk_size - overlap) start += (chunk_size - overlap)
return chunks return chunks
def _chunk_code(self, content: str, file_path: str) -> List[str]: def _chunk_code(self, content: str, file_path: str) -> List[str]:
"""AST-aware chunking for Python code.""" """AST-aware chunking for Python code."""
try: try:
from src.file_cache import ASTParser from src.file_cache import ASTParser
parser = ASTParser("python") parser = ASTParser("python")
tree = parser.parse(content) tree = parser.parse(content)
chunks = [] chunks = []
# Capture classes and top-level functions for node in tree.root_node.children:
for node in tree.root_node.children: if node.type in ("function_definition", "class_definition"):
if node.type in ("function_definition", "class_definition"): chunks.append(content[node.start_byte:node.end_byte])
chunks.append(content[node.start_byte:node.end_byte])
# Fallback if no structural chunks found or if file is small if not chunks or len(content) < self.config.chunk_size:
if not chunks or len(content) < self.config.chunk_size: return self._chunk_text(content)
return self._chunk_text(content) return chunks
return chunks except Exception:
except Exception: return self._chunk_text(content)
return self._chunk_text(content)
def index_file(self, file_path: str): def index_file(self, file_path: str):
"""Reads, chunks, and indexes a file into the vector store.""" """Reads, chunks, and indexes a file into the vector store."""
if not self.config.enabled or self.collection == "mock": if not self.config.enabled or self.collection == "mock":
return return
full_path = os.path.join(self.base_dir, file_path) full_path = os.path.join(self.base_dir, file_path)
if not os.path.exists(full_path): if not os.path.exists(full_path):
return return
try: try:
mtime = os.path.getmtime(full_path) mtime = os.path.getmtime(full_path)
except Exception: except Exception:
return return
# Incremental check: see if we already have this file with the same mtime try:
try: res = self.collection.get(where={"path": file_path}, limit=1, include=["metadatas"])
res = self.collection.get(where={"path": file_path}, limit=1, include=["metadatas"]) if res and res["metadatas"] and res["metadatas"][0]:
if res and res["metadatas"] and res["metadatas"][0]: if res["metadatas"][0].get("mtime") == mtime:
if res["metadatas"][0].get("mtime") == mtime: return
return except Exception:
except Exception: pass
pass
try: try:
with open(full_path, "r", encoding="utf-8", errors="ignore") as f: with open(full_path, "r", encoding="utf-8", errors="ignore") as f:
content = f.read() content = f.read()
except Exception: except Exception:
return return
# Remove old entries for this file self.collection.delete(where={"path": file_path})
self.collection.delete(where={"path": file_path})
if file_path.lower().endswith(".py"): if file_path.lower().endswith(".py"):
chunks = self._chunk_code(content, file_path) chunks = self._chunk_code(content, file_path)
else: else:
chunks = self._chunk_text(content) chunks = self._chunk_text(content)
if not chunks: if not chunks:
return return
ids = [f"{file_path}_{i}" for i in range(len(chunks))] ids = [f"{file_path}_{i}" for i in range(len(chunks))]
metadatas = [{"path": file_path, "chunk": i, "mtime": mtime} for i in range(len(chunks))] metadatas = [{"path": file_path, "chunk": i, "mtime": mtime} for i in range(len(chunks))]
self.add_documents(ids, chunks, metadatas) self.add_documents(ids, chunks, metadatas)
def _search_mcp(self, query: str, top_k: int = 5) -> List[Dict[str, Any]]: def _search_mcp(self, query: str, top_k: int = 5) -> List[Dict[str, Any]]:
async def _async_search_mcp(): async def _async_search_mcp():
tool_name = self.config.vector_store.mcp_tool or "rag_search" tool_name = self.config.vector_store.mcp_tool or "rag_search"
args = {"query": query, "top_k": top_k} args = {"query": query, "top_k": top_k}
res_str = await mcp_client.async_dispatch(tool_name, args) res_str = await mcp_client.async_dispatch(tool_name, args)
try: try:
data = json.loads(res_str) data = json.loads(res_str)
if isinstance(data, list): if isinstance(data, list):
return data return data
elif isinstance(data, dict) and "results" in data: elif isinstance(data, dict) and "results" in data:
return data["results"] return data["results"]
return [] return []
except: except:
return [] return []
return asyncio.run(_async_search_mcp()) return asyncio.run(_async_search_mcp())
def search(self, query: str, top_k: int = 5) -> List[Dict[str, Any]]: def search(self, query: str, top_k: int = 5) -> List[Dict[str, Any]]:
""" """
[C: tests/mock_concurrent_mma.py:main, tests/test_rag_engine.py:test_rag_engine_chroma] [C: tests/mock_concurrent_mma.py:main, tests/test_rag_engine.py:test_rag_engine_chroma]
""" """
if not self.config.enabled: if not self.config.enabled:
return [] return []
if self.config.vector_store.provider == 'mcp': if self.config.vector_store.provider == 'mcp':
return self._search_mcp(query, top_k) return self._search_mcp(query, top_k)
if self.collection == "mock": if self.collection == "mock":
return [] return []
query_embedding = self.embedding_provider.embed([query])[0] query_embedding = self.embedding_provider.embed([query])[0]
results = self.collection.query( results = self.collection.query(
query_embeddings=[query_embedding], query_embeddings=[query_embedding],
n_results=top_k n_results=top_k
) )
ret = [] ret = []
if results and results["ids"] and results["ids"][0]: if results and results["ids"] and results["ids"][0]:
for i in range(len(results["ids"][0])): for i in range(len(results["ids"][0])):
ret.append({ ret.append({
"id": results["ids"][0][i], "id": results["ids"][0][i],
"document": results["documents"][0][i], "document": results["documents"][0][i],
"metadata": results["metadatas"][0][i] if results["metadatas"] else {}, "metadata": results["metadatas"][0][i] if results["metadatas"] else {},
"distance": results["distances"][0][i] if "distances" in results and results["distances"] else 0.0 "distance": results["distances"][0][i] if "distances" in results and results["distances"] else 0.0
}) })
return ret return ret
def delete_documents(self, ids: List[str]): def delete_documents(self, ids: List[str]):
""" """
[C: tests/test_rag_engine.py:test_rag_engine_chroma] [C: tests/test_rag_engine.py:test_rag_engine_chroma]
""" """
if not self.config.enabled or self.collection == "mock": if not self.config.enabled or self.collection == "mock":
return return
self.collection.delete(ids=ids) self.collection.delete(ids=ids)
def get_all_indexed_paths(self) -> List[str]: def get_all_indexed_paths(self) -> List[str]:
if not self.config.enabled or self.collection == "mock": if not self.config.enabled or self.collection == "mock":
return [] return []
res = self.collection.get(include=["metadatas"]) res = self.collection.get(include=["metadatas"])
if not res or not res["metadatas"]: if not res or not res["metadatas"]:
return [] return []
return list(set(m.get("path") for m in res["metadatas"] if m.get("path"))) return list(set(m.get("path") for m in res["metadatas"] if m.get("path")))
def delete_documents_by_path(self, file_paths: List[str]): def delete_documents_by_path(self, file_paths: List[str]):
if not self.config.enabled or self.collection == "mock": if not self.config.enabled or self.collection == "mock":
return return
for path in file_paths: for path in file_paths:
self.collection.delete(where={"path": path}) self.collection.delete(where={"path": path})