Spaces:
Runtime error
Runtime error
| 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"} | |