Spaces:
Runtime error
Runtime error
File size: 3,631 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 | import logging
from typing import List, Dict
from .knowledge_base_manager import KnowledgeBaseManager
try:
from .vector_rag_manager import VectorRAGManager
VECTOR_AVAILABLE = True
except ImportError:
VECTOR_AVAILABLE = False
logger = logging.getLogger(__name__)
class RAGHelper:
def __init__(self, knowledge_base_path: str = "knowledge_base", use_vector_search: bool = True):
self.kb_manager = KnowledgeBaseManager(knowledge_base_path)
self.kb_manager.scan_documents()
self.use_vector_search = use_vector_search and VECTOR_AVAILABLE
self.vector_rag = None
if self.use_vector_search:
try:
self.vector_rag = VectorRAGManager(knowledge_base_path, vector_db_path="/app/data/vector_db")
logger.info("Vector RAG initialized")
except Exception as e:
logger.warning(f"Vector RAG failed: {e}")
self.use_vector_search = False
self.max_context_docs = 3
self.similarity_threshold = 0.30
def _route_query_to_categories(self, query: str) -> List[str]:
ql = query.lower()
cat_keywords = {
'policies': ['return','refund','warranty','policy','shipping'],
'faqs': ['how','what','why','help','troubleshoot','problem'],
'product_manuals': ['spec','manual','guide','connect','cable','hdmi','usb']
}
scores = {}
for cat, kw in cat_keywords.items():
score = sum(1 for k in kw if k in ql)
if score:
scores[cat] = score
if scores:
return sorted(scores, key=scores.get, reverse=True)
return ['policies','faqs','product_manuals']
def get_relevant_context(self, query: str) -> str:
categories = self._route_query_to_categories(query)
if self.use_vector_search and self.vector_rag:
try:
results = self.vector_rag.retrieve_and_rerank_filtered(
query, target_categories=categories, initial_k=20, final_k=3,
similarity_threshold=self.similarity_threshold
)
return self._build_context_from_results(results)
except Exception as e:
logger.warning(f"Vector search failed: {e}, falling back to keyword")
return self._get_keyword_context(query)
def _get_keyword_context(self, query: str) -> str:
results = self.kb_manager.search_documents(query)
context = ""
for r in results[:self.max_context_docs]:
content = self.kb_manager.load_document_content(r['path'])
if content:
context += f"\n--- {r['title']} (Category: {r['category']}) ---\n{content[:1500]}\n"
return context
def _build_context_from_results(self, results: List[Dict]) -> str:
if not results:
return ""
chunks = []
for res in results:
doc = res['document']
meta = res['metadata']
chunks.append(f"--- {meta['document_title']} (Category: {meta['category']}) ---\n{doc[:1200]}")
return "\n".join(chunks[:3])
def get_knowledge_base_stats(self) -> Dict:
stats = self.kb_manager.get_stats() if hasattr(self.kb_manager, 'get_stats') else {'total_documents': 0}
stats['vector_search_available'] = self.use_vector_search
return stats
def ensure_vector_index(self, force_reindex: bool = False):
if self.vector_rag:
return self.vector_rag.index_documents(force_reindex=force_reindex)
return {"error": "Vector not available"}
|