Spaces:
Runtime error
Runtime error
File size: 4,984 Bytes
a5778f2 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 | import os
import logging
from pathlib import Path
from typing import List, Dict
import chromadb
from chromadb.config import Settings
from sentence_transformers import SentenceTransformer
from .knowledge_base_manager import KnowledgeBaseManager
os.environ['ANONYMIZED_TELEMETRY'] = 'False'
os.environ['CHROMA_TELEMETRY_ENABLED'] = 'false'
logger = logging.getLogger(__name__)
class VectorRAGManager:
def __init__(self, knowledge_base_path: str, vector_db_path: str = "vector_db", embedding_model: str = "all-MiniLM-L6-v2"):
self.kb_path = Path(knowledge_base_path)
self.vector_db_path = Path(vector_db_path)
self.vector_db_path.mkdir(exist_ok=True, parents=True)
self.kb_manager = KnowledgeBaseManager(knowledge_base_path)
self.chroma_client = chromadb.PersistentClient(path=str(self.vector_db_path), settings=Settings(anonymized_telemetry=False))
self.embedding_model = SentenceTransformer(embedding_model)
self.collection_name = "tmc_documents"
self._init_collection()
self.chunk_size = 800
self.chunk_overlap = 120
def _init_collection(self):
try:
self.collection = self.chroma_client.get_collection(self.collection_name)
except:
self.collection = self.chroma_client.create_collection(self.collection_name)
def chunk_text(self, text: str) -> List[str]:
if len(text) <= self.chunk_size:
return [text]
chunks = []
start = 0
while start < len(text):
end = min(start + self.chunk_size, len(text))
if end < len(text):
for i in range(end, max(start, end-200), -1):
if text[i] in '.!?\n':
end = i+1
break
chunks.append(text[start:end].strip())
start = end - self.chunk_overlap
return [c for c in chunks if len(c) > 50]
def generate_embedding(self, text: str) -> List[float]:
return self.embedding_model.encode(text).tolist()
def index_documents(self, force_reindex: bool = False) -> Dict:
if force_reindex and self.collection.count() > 0:
self.chroma_client.delete_collection(self.collection_name)
self._init_collection()
self.kb_manager.scan_documents()
stats = {"documents": 0, "chunks": 0}
for category, cat_info in self.kb_manager.documents.items():
if category == 'development':
logger.info(f"Skipping development folder")
continue
for doc in cat_info['documents']:
content = self.kb_manager.load_document_content(doc['path'])
if not content:
continue
chunks = self.chunk_text(content)
if not chunks:
continue
ids = []
embeds = []
texts = []
metas = []
for i, chunk in enumerate(chunks):
cid = f"{doc['filename']}_{i}"
ids.append(cid)
texts.append(chunk)
embeds.append(self.generate_embedding(chunk))
metas.append({
"document_title": doc['title'],
"document_path": doc['path'],
"category": category,
"chunk_index": i
})
self.collection.add(ids=ids, embeddings=embeds, documents=texts, metadatas=metas)
stats["documents"] += 1
stats["chunks"] += len(chunks)
logger.info(f"Indexed {stats['documents']} docs, {stats['chunks']} chunks")
return stats
def semantic_search(self, query: str, n_results: int = 5) -> List[Dict]:
if self.collection.count() == 0:
return []
q_emb = self.generate_embedding(query)
results = self.collection.query(query_embeddings=[q_emb], n_results=n_results, include=["documents","metadatas","distances"])
formatted = []
if results['documents'] and results['documents'][0]:
for doc, meta, dist in zip(results['documents'][0], results['metadatas'][0], results['distances'][0]):
similarity = max(0.0, 1.0 - (dist / 2.0))
formatted.append({"document": doc, "metadata": meta, "similarity": similarity})
return formatted
def retrieve_and_rerank_filtered(self, query: str, target_categories: List[str], initial_k: int = 20, final_k: int = 3, similarity_threshold: float = 0.30) -> List[Dict]:
candidates = self.semantic_search(query, n_results=initial_k)
filtered = [c for c in candidates if c['similarity'] >= similarity_threshold and c['metadata'].get('category') in target_categories]
return filtered[:final_k]
def get_collection_stats(self) -> Dict:
return {"total_chunks": self.collection.count()}
|