Update core/rag_system.py

core/rag_system.py

@@ -2,11 +2,9 @@ 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
+from models.schemas import RAGSearchResult
+from config.settings import settings
+from core.multilingual_manager import MultilingualManager
 
 class EnhancedRAGSystem:
     def __init__(self):
@@ -15,12 +13,13 @@ class EnhancedRAGSystem:
         self.embeddings: Optional[np.ndarray] = None
         self.index: Optional[faiss.Index] = None
         
-        # Multilingual support - simplified for now
-        self.current_dimension = 384  # Default dimension
+        # Multilingual support
+        self.multilingual_manager = MultilingualManager()
+        self.current_dimension = settings.EMBEDDING_DIMENSION
         
-        self._initialize_sample_data()
+        self._initialize_sample_data()  # SỬA TÊN HÀM
     
-    def _initialize_sample_data(self):
+    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 = [
@@ -44,7 +43,7 @@ class EnhancedRAGSystem:
             "The United States has diverse climate zones from tropical to arctic"
         ]
         
-        # Vietnamese metadata
+        # Vietnamese metadata - SỬA LỖI SYNTAX
         vietnamese_metadatas = [
             {"type": "nutrition", "source": "sample", "language": "vi"},
             {"type": "nutrition", "source": "sample", "language": "vi"},
@@ -55,7 +54,7 @@ class EnhancedRAGSystem:
             {"type": "geography", "source": "sample", "language": "vi"}
         ]
         
-        # English metadata
+        # English metadata - SỬA LỖI SYNTAX
         english_metadatas = [
             {"type": "nutrition", "source": "sample", "language": "en"},
             {"type": "nutrition", "source": "sample", "language": "en"},
@@ -70,15 +69,6 @@ class EnhancedRAGSystem:
         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...")
@@ -107,11 +97,11 @@ class EnhancedRAGSystem:
         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
+            if doc and isinstance(doc, str) and len(doc.strip()) > 5:  # At least 5 characters
                 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"⚠️ Bỏ qua document {i}: không hợp lệ")
         
         print(f"📊 Documents hợp lệ: {len(valid_documents)}/{len(documents)}")
         
@@ -120,36 +110,30 @@ class EnhancedRAGSystem:
             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
+                language = valid_metadatas[i].get('language', 'vi')
+                embedding_model = self.multilingual_manager.get_embedding_model(language)
+                
+                if embedding_model is None:
+                    print(f"⚠️ Không có embedding model cho document {i}")
+                    continue
+                    
+                # Create embedding
                 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)
+            print("❌ Không tạo được embeddings nào")
             return
         
         # Convert to numpy array
@@ -158,7 +142,6 @@ class EnhancedRAGSystem:
             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
@@ -176,26 +159,19 @@ class EnhancedRAGSystem:
                 # 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ó")
+                    print("🔄 Tạo system mới do dimension không khớp")
+                    self.embeddings = new_embeddings
+                    self.documents = valid_documents
+                    self.metadatas = valid_metadatas
+                else:
+                    # 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
@@ -205,16 +181,9 @@ class EnhancedRAGSystem:
         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:
@@ -229,20 +198,23 @@ class EnhancedRAGSystem:
         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"""
+    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 = 5
+            top_k = settings.TOP_K_RESULTS
             
         if not self.documents or self.index is None:
             return self._fallback_keyword_search(query, top_k)
         
-        embedding_model = self._get_embedding_model()
+        # 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
+            # Encode query with appropriate model
             query_embedding = embedding_model.encode([query])
             
             # Normalize query embedding for cosine similarity
@@ -257,27 +229,52 @@ class EnhancedRAGSystem:
             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 {}
-                    })
+                    results.append(RAGSearchResult(
+                        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
+            # 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 _fallback_keyword_search(self, query: str, top_k: int) -> List[Dict]:
+    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():
@@ -285,18 +282,18 @@ class EnhancedRAGSystem:
                     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.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)
+        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"""
+        """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')
@@ -310,4 +307,4 @@ class EnhancedRAGSystem:
             'name': 'multilingual_rag_system',
             'status': 'active',
             'has_embeddings': self.embeddings is not None
-        }
+        }