core/rag_pipeline.py

import logging
import google.generativeai as genai
from core.embedding_model import get_embedding_model
from config import GEMINI_API_KEY, HUMAN_PROMPT_TEMPLATE, SYSTEM_PROMPT, TOP_K_RESULTS, TEMPERATURE, MAX_OUTPUT_TOKENS
import os
import re

logger = logging.getLogger(__name__)

# Cấu hình Gemini API
genai.configure(api_key=GEMINI_API_KEY)

class RAGPipeline:
    def __init__(self):
        # Khởi tạo RAG Pipeline với embedding model
        logger.info("Đang khởi tạo RAG Pipeline")
        
        self.embedding_model = get_embedding_model()
        self.gemini_model = genai.GenerativeModel('gemini-2.0-flash')
        
        logger.info("RAG Pipeline đã sẵn sàng hoạt động")
    
    def generate_response(self, query, age=1):
        # Tạo phản hồi cho câu hỏi của người dùng sử dụng RAG
        try:
            logger.info(f"Bắt đầu tạo phản hồi cho câu hỏi: {query[:50]}... (tuổi: {age})")
            
            # Tìm kiếm thông tin liên quan trong ChromaDB
            logger.info("Đang tìm kiếm thông tin liên quan trong cơ sở dữ liệu")
            search_results = self.embedding_model.search(query, top_k=TOP_K_RESULTS, age_filter=age)
            # search_results = self.embedding_model.search(query, top_k=TOP_K_RESULTS)
            
            if not search_results or len(search_results) == 0:
                logger.warning("Không tìm thấy thông tin liên quan trong cơ sở dữ liệu")
                return {
                    "success": True,
                    "response": "Xin lỗi, tôi không tìm thấy thông tin liên quan đến câu hỏi của bạn trong tài liệu.",
                    "sources": []
                }
            
            # Chuẩn bị ngữ cảnh từ kết quả tìm kiếm
            contexts = []
            sources = []
            
            for result in search_results:
                metadata = result.get('metadata', {})
                content = result.get('document', '')
                
                # Thêm nội dung vào ngữ cảnh
                contexts.append({
                    "content": content,
                    "metadata": metadata
                })
                
                #Tạo thông tin nguồn tài liệu với chunk ID
                chunk_id = metadata.get('chunk_id', 'unknown_chunk')
                chapter = metadata.get('chapter', '')
                original_title = metadata.get('title', '')
                
                # Tạo tên tài liệu dựa trên chapter
                if 'bai1' in chapter:
                    document_name = "Bài 1: Dinh dưỡng theo lứa tuổi học sinh"
                elif 'bai2' in chapter:
                    document_name = "Bài 2: An toàn thực phẩm"
                elif 'bai3' in chapter:
                    document_name = "Bài 3: Vệ sinh dinh dưỡng"
                elif 'bai4' in chapter:
                    document_name = "Bài 4: Giáo dục dinh dưỡng"
                elif 'phuluc' in chapter:
                    document_name = "Phụ lục"
                elif chapter.startswith('bosung'):
                    document_name = metadata.get('document_title') or f"Tài liệu bổ sung {chapter}"
                else:
                    document_name = metadata.get('document_title') or "Tài liệu dinh dưỡng"
                
                # ✅ Source title với format: "Tên tài liệu - ID: chunk_id"
                if original_title and original_title.strip() and len(original_title.strip()) > 3:
                    # Kiểm tra title không phải là generic title
                    generic_titles = ['chunk', 'section', 'part', 'mục', 'phần']
                    if not any(generic in original_title.lower() for generic in generic_titles):
                        source_title = f"{document_name} - {original_title} - ID: {chunk_id}"
                    else:
                        source_title = f"{document_name} - ID: {chunk_id}"
                else:
                    source_title = f"{document_name} - ID: {chunk_id}"
                
                source_info = {
                    "title": source_title,
                    "pages": metadata.get('pages'),
                    "content_type": metadata.get('content_type', 'text'),
                    "chunk_id": chunk_id,
                    "chapter": chapter
                }
                
                # Kiểm tra trùng lặp dựa trên chunk_id
                if not any(s.get('chunk_id') == chunk_id for s in sources):
                    sources.append(source_info)
            
            # Định dạng ngữ cảnh cho prompt
            formatted_contexts = self._format_contexts(contexts)
            
            # Tạo prompt với ngữ cảnh độ tuổi
            full_prompt = self._create_prompt_with_age_context(query, age, formatted_contexts)
            
            # Tạo phản hồi với Gemini AI
            logger.info("Đang tạo phản hồi với Gemini AI")
            response = self.gemini_model.generate_content(
                full_prompt,
                generation_config=genai.types.GenerationConfig(
                    temperature=TEMPERATURE,
                    max_output_tokens=MAX_OUTPUT_TOKENS
                )
            )
            
            if not response or not response.text:
                logger.error("Gemini AI không trả về phản hồi")
                return {
                    "success": False,
                    "error": "Không thể tạo phản hồi"
                }
            
            response_text = response.text.strip()
            
            # Xử lý các đường dẫn hình ảnh trong phản hồi
            response_text = self._process_image_links(response_text)
            
            logger.info("Đã tạo phản hồi thành công")
            
            return {
                "success": True,
                "response": response_text,
                "sources": sources[:3]  # ✅ Giới hạn 3 sources để không quá dài
            }
            
        except Exception as e:
            logger.error(f"Lỗi khi tạo phản hồi: {str(e)}")
            return {
                "success": False,
                "error": f"Lỗi tạo phản hồi: {str(e)}"
            }
    
    def _format_contexts(self, contexts):
        # Định dạng ngữ cảnh với tên tài liệu thực tế + chunk ID
        formatted = []
        
        for i, context in enumerate(contexts, 1):
            content = context['content']
            metadata = context['metadata']
            
            # Lấy thông tin để tạo context với số thứ tự
            chunk_id = metadata.get('chunk_id', f'chunk_{i}')
            chapter = metadata.get('chapter', '')
            original_title = metadata.get('title', '')
            
            # Tạo tên tài liệu dựa trên chapter
            if 'bai1' in chapter:
                document_name = "Bài 1: Dinh dưỡng theo lứa tuổi học sinh"
            elif 'bai2' in chapter:
                document_name = "Bài 2: An toàn thực phẩm"
            elif 'bai3' in chapter:
                document_name = "Bài 3: Vệ sinh dinh dưỡng"
            elif 'bai4' in chapter:
                document_name = "Bài 4: Giáo dục dinh dưỡng"
            elif 'phuluc' in chapter:
                document_name = "Phụ lục"
            else:
                document_name = metadata.get('document_title') or "Tài liệu dinh dưỡng"
            
            # Format: [{tên bài} - ID: {chunk_id}]
            context_str = f"[{document_name} - ID: {chunk_id}]"
            
            # Thêm title của chunk nếu có và có ý nghĩa
            if original_title and original_title.strip() and len(original_title.strip()) > 3:
                generic_titles = ['chunk', 'section', 'part', 'mục', 'phần']
                if not any(generic in original_title.lower() for generic in generic_titles):
                    context_str += f" - {original_title}"
            
            context_str += f"\n{content}\n"
            formatted.append(context_str)
        
        return "\n".join(formatted)
    
    def _create_prompt_with_age_context(self, query, age, contexts):        
        # Xác định hướng dẫn theo nhóm tuổi
        if age <= 3:
            age_guidance = "Sử dụng ngôn ngữ đơn giản, dễ hiểu cho phụ huynh có con nhỏ."
        elif age <= 6:
            age_guidance = "Tập trung vào dinh dưỡng cho trẻ mầm non, ngôn ngữ phù hợp với phụ huynh."
        elif age <= 12:
            age_guidance = "Nội dung phù hợp cho trẻ tiểu học, có thể giải thích đơn giản cho trẻ hiểu."
        elif age <= 15:
            age_guidance = "Thông tin chi tiết hơn, phù hợp cho học sinh trung học cơ sở."
        else:
            age_guidance = "Thông tin đầy đủ, chi tiết cho học sinh trung học phổ thông."
        
        # Tạo system prompt có tính đến độ tuổi (giữ nguyên SYSTEM_PROMPT gốc)
        age_aware_system_prompt = f"""{SYSTEM_PROMPT}

QUAN TRỌNG - Hướng dẫn theo độ tuổi:
Người dùng hiện tại {age} tuổi. {age_guidance}
- Điều chỉnh ngôn ngữ và nội dung cho phù hợp
- Đưa ra lời khuyên cụ thể cho độ tuổi này
- Tránh thông tin quá phức tạp hoặc không phù hợp
"""
        
        # Tạo human prompt từ template với hướng dẫn trích dẫn
        human_prompt = f"""
Câu hỏi: {query}

Độ tuổi người dùng: {age} tuổi

Tài liệu tham khảo:
{contexts}

Dựa vào thông tin trong tài liệu tham khảo và ngữ cảnh cuộc trò chuyện (nếu có), hãy trả lời câu hỏi một cách chi tiết, dễ hiểu và phù hợp với độ tuổi của người dùng.

Nếu trong tài liệu có bảng biểu hoặc hình ảnh, hãy giữ nguyên và đưa vào câu trả lời. 

Nếu câu hỏi hiện tại có liên quan đến các cuộc trò chuyện trước đó, hãy cố gắng duy trì sự nhất quán trong câu trả lời.
"""
        
        return f"{age_aware_system_prompt}\n\n{human_prompt}"
    
    def _process_image_links(self, response_text):
        # Xử lý và chuyển đổi các đường dẫn hình ảnh trong phản hồi
        try:
            import re
            
            # Tìm các pattern markdown: ![alt](path)
            image_pattern = r'!\[([^\]]*)\]\(([^)]+)\)'
            
            def replace_image_path(match):
                alt_text = match.group(1)
                image_path = match.group(2)
                
                # Xử lý đường dẫn local (Windows/Linux)
                if '\\' in image_path or image_path.startswith('/') or ':' in image_path:
                    # Trích xuất tên file từ đường dẫn local
                    filename = image_path.split('\\')[-1].split('/')[-1]
                    
                    # Tìm bai_id từ tên file (format: baiX_filename)
                    bai_match = re.match(r'^(bai\d+)_', filename)
                    if bai_match:
                        bai_id = bai_match.group(1)
                        # Tạo URL API
                        api_url = f"/api/figures/{bai_id}/{filename}"
                        return f"![{alt_text}]({api_url})"
                
                # Nếu đã là đường dẫn API, giữ nguyên
                elif image_path.startswith('/api/figures/'):
                    return match.group(0)
                
                # Xử lý đường dẫn tương đối
                elif '../figures/' in image_path:
                    filename = image_path.split('../figures/')[-1]
                    bai_match = re.match(r'^(bai\d+)_', filename)
                    if bai_match:
                        bai_id = bai_match.group(1)
                        api_url = f"/api/figures/{bai_id}/{filename}"
                        return f"![{alt_text}]({api_url})"

                return match.group(0)
            
            # Thay thế tất cả các liên kết hình ảnh
            processed_text = re.sub(image_pattern, replace_image_path, response_text)
            
            image_count = len(re.findall(image_pattern, response_text))
            if image_count > 0:
                logger.info(f"Đã xử lý {image_count} liên kết hình ảnh")
            
            return processed_text
            
        except Exception as e:
            logger.error(f"Lỗi khi xử lý liên kết hình ảnh: {e}")
            return response_text

    def generate_follow_up_questions(self, query, answer, age=1):
        # Tạo câu hỏi gợi ý dựa trên cuộc hội thoại hiện tại
        try:
            logger.info("Đang tạo câu hỏi gợi ý")
            
            follow_up_prompt = f"""
Dựa trên cuộc hội thoại sau, hãy tạo 3-5 câu hỏi gợi ý phù hợp cho người dùng {age} tuổi về chủ đề dinh dưỡng:

Câu hỏi gốc: {query}
Câu trả lời: {answer}

Hãy tạo các câu hỏi:
1. Liên quan trực tiếp đến chủ đề
2. Phù hợp với độ tuổi {age}
3. Thực tế và hữu ích
4. Ngắn gọn, dễ hiểu

Trả về danh sách câu hỏi, mỗi câu một dòng, không đánh số.
"""
            
            response = self.gemini_model.generate_content(
                follow_up_prompt,
                generation_config=genai.types.GenerationConfig(
                    temperature=0.7,
                    max_output_tokens=500
                )
            )
            
            if not response or not response.text:
                return {
                    "success": False,
                    "error": "Không thể tạo câu hỏi gợi ý"
                }
            
            # Chuyển đổi phản hồi thành danh sách câu hỏi
            questions = []
            lines = response.text.strip().split('\n')
            
            for line in lines:
                line = line.strip()
                # Lọc các dòng hợp lệ (không rỗng, không phải comment, đủ dài)
                if line and not line.startswith('#') and len(line) > 10:
                    # Loại bỏ số thứ tự nếu có (1. 2. hoặc 1) 2))
                    line = re.sub(r'^\d+[\.\)]\s*', '', line)
                    questions.append(line)
            
            # Giới hạn tối đa 5 câu hỏi
            questions = questions[:5]
            
            logger.info(f"Đã tạo {len(questions)} câu hỏi gợi ý")
            
            return {
                "success": True,
                "questions": questions
            }
            
        except Exception as e:
            logger.error(f"Lỗi khi tạo câu hỏi gợi ý: {str(e)}")
            return {
                "success": False,
                "error": f"Lỗi tạo câu hỏi gợi ý: {str(e)}"
            }