core/rag_system.py

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
        }