core/rag_system.py

import numpy as np
import faiss
from typing import List, Dict, Optional
from sentence_transformers import SentenceTransformer
import os
import json
import pandas as pd
from typing import List
import traceback

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 - simplified for now
        self.current_dimension = 384  # Default dimension
        
        self._initialize_sample_data()
    
    def _initialize_sample_data(self):
        """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
        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
        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 _get_embedding_model(self):
        """Lấy embedding model - simplified version"""
        try:
            # Sử dụng model nhỏ để tiết kiệm bộ nhớ
            return SentenceTransformer('sentence-transformers/all-MiniLM-L6-v2')
        except Exception as e:
            print(f"❌ Lỗi load embedding model: {e}")
            return None
    
    def add_documents(self, documents: List[str], metadatas: List[Dict] = None):
        """Thêm documents vào database - ĐÃ SỬA LỖI"""
        print(f"🔄 RAG System: Bắt đầu thêm {len(documents)} documents...")
        
        if not documents:
            print("❌ RAG System: Không có documents để thêm")
            return
        
        # Ensure metadatas has the same length as documents
        if metadatas is None:
            metadatas = [{} for _ in documents]
            print("📝 Tạo metadata mặc định")
        elif len(metadatas) != len(documents):
            print(f"⚠️ Metadata length mismatch: {len(metadatas)} vs {len(documents)}")
            # Fix metadata length
            new_metadatas = []
            for i in range(len(documents)):
                if i < len(metadatas):
                    new_metadatas.append(metadatas[i])
                else:
                    new_metadatas.append({"source": "upload", "language": "vi"})
            metadatas = new_metadatas
        
        # Filter valid documents
        valid_documents = []
        valid_metadatas = []
        
        for i, doc in enumerate(documents):
            if doc and isinstance(doc, str) and len(doc.strip()) > 3:  # Giảm độ dài tối thiểu
                valid_documents.append(doc.strip())
                valid_metadatas.append(metadatas[i] if i < len(metadatas) else {})
            else:
                print(f"⚠️ Bỏ qua document {i}: không hợp lệ - '{doc}'")
        
        print(f"📊 Documents hợp lệ: {len(valid_documents)}/{len(documents)}")
        
        if not valid_documents:
            print("❌ Không có documents hợp lệ để thêm")
            return
        
        # Create embeddings
        embedding_model = self._get_embedding_model()
        if embedding_model is None:
            print("❌ Không thể tạo embedding model")
            # Vẫn thêm documents không có embedding
            self._add_documents_without_embeddings(valid_documents, valid_metadatas)
            return
        
        new_embeddings_list = []
        successful_embeddings = 0
        
        for i, doc in enumerate(valid_documents):
            try:
                # Create embedding - sử dụng model duy nhất
                doc_embedding = embedding_model.encode([doc])
                new_embeddings_list.append(doc_embedding[0])
                successful_embeddings += 1
                
                if i % 10 == 0:  # Log tiến độ
                    print(f"📊 Đã embedding {i+1}/{len(valid_documents)} documents")
                
            except Exception as e:
                print(f"❌ Lỗi embedding document {i}: {e}")
                # Thêm document không có embedding
                new_embeddings_list.append(np.zeros(self.current_dimension))
        
        print(f"📊 Embeddings thành công: {successful_embeddings}/{len(valid_documents)}")
        
        if not new_embeddings_list:
            print("❌ Không tạo được embeddings nào, thêm documents không embedding")
            self._add_documents_without_embeddings(valid_documents, valid_metadatas)
            return
        
        # Convert to numpy array
        try:
            new_embeddings = np.array(new_embeddings_list)
            print(f"✅ Embedding matrix shape: {new_embeddings.shape}")
        except Exception as e:
            print(f"❌ Lỗi tạo embedding matrix: {e}")
            self._add_documents_without_embeddings(valid_documents, valid_metadatas)
            return
        
        # Handle existing embeddings
        old_doc_count = len(self.documents)
        
        if self.embeddings is None:
            # First time initialization
            self.embeddings = new_embeddings
            self.documents = valid_documents
            self.metadatas = valid_metadatas
            print("✅ Khởi tạo RAG system lần đầu")
        else:
            # Append to existing
            try:
                # Check dimension compatibility
                if self.embeddings.shape[1] != new_embeddings.shape[1]:
                    print(f"⚠️ Dimension mismatch: {self.embeddings.shape[1]} vs {new_embeddings.shape[1]}")
                    # Resize embeddings để phù hợp
                    if self.embeddings.shape[1] < new_embeddings.shape[1]:
                        # Pad existing embeddings
                        pad_width = new_embeddings.shape[1] - self.embeddings.shape[1]
                        self.embeddings = np.pad(self.embeddings, ((0,0), (0,pad_width)))
                    else:
                        # Truncate new embeddings
                        new_embeddings = new_embeddings[:, :self.embeddings.shape[1]]
                    
                    print("🔄 Đã điều chỉnh dimension")
                
                # Compatible dimensions, append
                self.embeddings = np.vstack([self.embeddings, new_embeddings])
                self.documents.extend(valid_documents)
                self.metadatas.extend(valid_metadatas)
                print("✅ Đã thêm vào system hiện có")
                    
            except Exception as e:
                print(f"❌ Lỗi khi thêm vào system: {e}")
                self._add_documents_without_embeddings(valid_documents, valid_metadatas)
                return
        
        # Update FAISS index
        self._update_faiss_index()
        
        new_doc_count = len(self.documents)
        print(f"🎉 THÀNH CÔNG: Đã thêm {new_doc_count - old_doc_count} documents mới")
        print(f"📊 Tổng documents: {new_doc_count}")
    
    def _add_documents_without_embeddings(self, documents: List[str], metadatas: List[Dict]):
        """Thêm documents không có embeddings (fallback)"""
        self.documents.extend(documents)
        self.metadatas.extend(metadatas)
        print(f"✅ Đã thêm {len(documents)} documents không có embeddings")
    
    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:
            print("⚠️ Không có embeddings để cập nhật index")
            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 = 5) -> List[Dict]:
        """Tìm kiếm ngữ nghĩa - simplified version"""
        if top_k is None:
            top_k = 5
            
        if not self.documents or self.index is None:
            return self._fallback_keyword_search(query, top_k)
        
        embedding_model = self._get_embedding_model()
        if embedding_model is None:
            return self._fallback_keyword_search(query, top_k)
            
        try:
            # Encode query
            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({
                        "id": str(idx),
                        "text": self.documents[idx],
                        "similarity": float(similarity),
                        "metadata": self.metadatas[idx] if idx < len(self.metadatas) else {}
                    })
            
            print(f"🔍 Tìm kiếm '{query[:50]}...' - Tìm thấy {len(results)} kết quả")
            return 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 _fallback_keyword_search(self, query: str, top_k: int) -> List[Dict]:
        """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
            
            # Keyword matching
            for word in query_lower.split():
                if len(word) > 2 and word in doc.lower():
                    score += 1
            
            if score > 0:
                results.append({
                    "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"""
        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
        }