import numpy as np import faiss from typing import List, Dict, Optional from sentence_transformers import SentenceTransformer from models.schemas import RAGSearchResult from config.settings import settings from core.multilingual_manager import MultilingualManager class EnhancedRAGSystem: def __init__(self): self.documents: List[str] = [] self.metadatas: List[Dict] = [] self.embeddings: Optional[np.ndarray] = None self.index: Optional[faiss.Index] = None # Multilingual support self.multilingual_manager = MultilingualManager() self.current_dimension = settings.EMBEDDING_DIMENSION self._initialize_sample_data() # SỬA TÊN HÀM def _initialize_sample_data(self): # SỬA TÊN HÀM """Khởi tạo dữ liệu mẫu""" # Vietnamese sample data vietnamese_data = [ "Rau xanh cung cấp nhiều vitamin và chất xơ tốt cho sức khỏe", "Trái cây tươi chứa nhiều vitamin C và chất chống oxy hóa", "Cá hồi giàu omega-3 tốt cho tim mạch và trí não", "Nước rất quan trọng cho cơ thể, nên uống ít nhất 2 lít mỗi ngày", "Hà Nội là thủ đô của Việt Nam, nằm ở miền Bắc", "Thành phố Hồ Chí Minh là thành phố lớn nhất Việt Nam", "Việt Nam có khí hậu nhiệt đới gió mùa với 4 mùa rõ rệt" ] # English sample data english_data = [ "Green vegetables provide many vitamins and fiber that are good for health", "Fresh fruits contain lots of vitamin C and antioxidants", "Salmon is rich in omega-3 which is good for heart and brain", "Water is very important for the body, should drink at least 2 liters per day", "London is the capital of England and the United Kingdom", "New York City is the most populous city in the United States", "The United States has diverse climate zones from tropical to arctic" ] # Vietnamese metadata - SỬA LỖI SYNTAX vietnamese_metadatas = [ {"type": "nutrition", "source": "sample", "language": "vi"}, {"type": "nutrition", "source": "sample", "language": "vi"}, {"type": "nutrition", "source": "sample", "language": "vi"}, {"type": "health", "source": "sample", "language": "vi"}, {"type": "geography", "source": "sample", "language": "vi"}, {"type": "geography", "source": "sample", "language": "vi"}, {"type": "geography", "source": "sample", "language": "vi"} ] # English metadata - SỬA LỖI SYNTAX english_metadatas = [ {"type": "nutrition", "source": "sample", "language": "en"}, {"type": "nutrition", "source": "sample", "language": "en"}, {"type": "nutrition", "source": "sample", "language": "en"}, {"type": "health", "source": "sample", "language": "en"}, {"type": "geography", "source": "sample", "language": "en"}, {"type": "geography", "source": "sample", "language": "en"}, {"type": "geography", "source": "sample", "language": "en"} ] # Add documents with language metadata self.add_documents(vietnamese_data, vietnamese_metadatas) self.add_documents(english_data, english_metadatas) def add_documents(self, documents: List[str], metadatas: List[Dict] = None): """Thêm documents vào database với embedding phù hợp""" if not documents: return # Ensure metadatas has the same length as documents if metadatas is None: metadatas = [{} for _ in documents] elif len(metadatas) != len(documents): # Extend or truncate metadatas to match documents length if len(metadatas) < len(documents): metadatas = metadatas + [{} for _ in range(len(documents) - len(metadatas))] else: metadatas = metadatas[:len(documents)] # Detect language for each document and create embeddings accordingly new_embeddings_list = [] valid_documents = [] valid_metadatas = [] for i, doc in enumerate(documents): if not doc or len(doc.strip()) == 0: continue language = metadatas[i].get('language', 'vi') embedding_model = self.multilingual_manager.get_embedding_model(language) if embedding_model is not None: try: # Create embedding for this document doc_embedding = embedding_model.encode([doc]) new_embeddings_list.append(doc_embedding[0]) valid_documents.append(doc) valid_metadatas.append(metadatas[i]) except Exception as e: print(f"❌ Lỗi tạo embedding cho document {i}: {e}") if not valid_documents: return # Convert list of embeddings to numpy array new_embeddings = np.array(new_embeddings_list) # Handle dimension mismatch if self.embeddings is not None and self.embeddings.shape[1] != new_embeddings.shape[1]: print(f"⚠️ Phát hiện dimension mismatch ({self.embeddings.shape[1]} vs {new_embeddings.shape[1]}), tạo index mới...") self.embeddings = None self.index = None # Update embeddings if self.embeddings is None: self.embeddings = new_embeddings self.current_dimension = new_embeddings.shape[1] else: self.embeddings = np.vstack([self.embeddings, new_embeddings]) # Update FAISS index self._update_faiss_index() self.documents.extend(valid_documents) self.metadatas.extend(valid_metadatas) print(f"✅ Đã thêm {len(valid_documents)} documents vào RAG database") def _update_faiss_index(self): """Cập nhật FAISS index với embeddings hiện tại""" if self.embeddings is None or len(self.embeddings) == 0: return try: dimension = self.embeddings.shape[1] self.index = faiss.IndexFlatIP(dimension) # Normalize embeddings for cosine similarity faiss.normalize_L2(self.embeddings) self.index.add(self.embeddings.astype(np.float32)) print(f"✅ Đã cập nhật FAISS index với dimension {dimension}") except Exception as e: print(f"❌ Lỗi cập nhật FAISS index: {e}") def semantic_search(self, query: str, top_k: int = None) -> List[RAGSearchResult]: """Tìm kiếm ngữ nghĩa với model phù hợp theo ngôn ngữ""" if top_k is None: top_k = settings.TOP_K_RESULTS if not self.documents or self.index is None: return self._fallback_keyword_search(query, top_k) # Detect query language and get appropriate model query_language = self.multilingual_manager.detect_language(query) embedding_model = self.multilingual_manager.get_embedding_model(query_language) if embedding_model is None: return self._fallback_keyword_search(query, top_k) try: # Encode query with appropriate model query_embedding = embedding_model.encode([query]) # Normalize query embedding for cosine similarity faiss.normalize_L2(query_embedding) # Search in FAISS index similarities, indices = self.index.search( query_embedding.astype(np.float32), min(top_k, len(self.documents)) ) results = [] for i, (similarity, idx) in enumerate(zip(similarities[0], indices[0])): if idx < len(self.documents): results.append(RAGSearchResult( id=str(idx), text=self.documents[idx], similarity=float(similarity), metadata=self.metadatas[idx] if idx < len(self.metadatas) else {} )) # Filter results by language relevance filtered_results = self._filter_by_language_relevance(results, query_language) print(f"🔍 Tìm kiếm '{query[:50]}...' (ngôn ngữ: {query_language}) - Tìm thấy {len(filtered_results)} kết quả") return filtered_results except Exception as e: print(f"❌ Lỗi tìm kiếm ngữ nghĩa: {e}") return self._fallback_keyword_search(query, top_k) def _filter_by_language_relevance(self, results: List[RAGSearchResult], query_language: str) -> List[RAGSearchResult]: """Lọc kết quả theo độ liên quan ngôn ngữ""" if not results: return results # Boost scores for documents in the same language for result in results: doc_language = result.metadata.get('language', 'vi') if doc_language == query_language: # Boost similarity score for same language documents result.similarity = min(result.similarity * 1.2, 1.0) # Re-sort by updated similarity scores results.sort(key=lambda x: x.similarity, reverse=True) return results def _fallback_keyword_search(self, query: str, top_k: int) -> List[RAGSearchResult]: """Tìm kiếm dự phòng dựa trên từ khóa""" query_lower = query.lower() results = [] for i, doc in enumerate(self.documents): score = 0 doc_language = self.metadatas[i].get('language', 'vi') if i < len(self.metadatas) else 'vi' query_language = self.multilingual_manager.detect_language(query) # Language matching bonus if doc_language == query_language: score += 0.5 # Keyword matching for word in query_lower.split(): if len(word) > 2 and word in doc.lower(): score += 1 if score > 0: results.append(RAGSearchResult( id=str(i), text=doc, similarity=min(score / 5, 1.0), metadata=self.metadatas[i] if i < len(self.metadatas) else {} )) results.sort(key=lambda x: x.similarity, reverse=True) return results[:top_k] def get_collection_stats(self) -> Dict: """Lấy thống kê collection với thông tin đa ngôn ngữ""" language_stats = {} for metadata in self.metadatas: lang = metadata.get('language', 'unknown') language_stats[lang] = language_stats.get(lang, 0) + 1 return { 'total_documents': len(self.documents), 'embedding_count': len(self.embeddings) if self.embeddings is not None else 0, 'embedding_dimension': self.current_dimension, 'language_distribution': language_stats, 'name': 'multilingual_rag_system', 'status': 'active', 'has_embeddings': self.embeddings is not None }