ComputerAssignments/CA6_systematic_generalization/integrations/advanced_rag_system.py

"""
سیستم RAG پیشرفته با LangChain و Gemini API
Advanced RAG System with LangChain and Gemini API
"""

import os
import logging
from typing import List, Dict, Any, Optional, Tuple
from dataclasses import dataclass
import json
import pickle
from pathlib import Path

# LangChain imports
from langchain.text_splitter import RecursiveCharacterTextSplitter, TokenTextSplitter
from langchain.document_loaders import TextLoader, PyPDFLoader, DirectoryLoader
from langchain.embeddings import HuggingFaceEmbeddings
from langchain.vectorstores import FAISS, Chroma, Pinecone
from langchain.chains import RetrievalQA, ConversationalRetrievalChain
from langchain.memory import ConversationBufferMemory, ConversationSummaryMemory
from langchain.prompts import PromptTemplate, ChatPromptTemplate
from langchain.schema import Document
from langchain.callbacks import StreamingStdOutCallbackHandler

# Gemini integration
from langchain_google_genai import ChatGoogleGenerativeAI, GoogleGenerativeAIEmbeddings
from langchain_community.embeddings import HuggingFaceEmbeddings

# Other utilities
import numpy as np
from sentence_transformers import SentenceTransformer
import chromadb
from chromadb.config import Settings

# Local imports
from gemini_config import GeminiConfig, GeminiSystematicGeneralization

logger = logging.getLogger(__name__)


@dataclass
class RAGConfig:
    """پیکربندی سیستم RAG"""

    # Embedding settings
    embedding_model: str = "sentence-transformers/all-MiniLM-L6-v2"
    embedding_dimension: int = 384

    # Text splitting
    chunk_size: int = 1000
    chunk_overlap: int = 200

    # Vector store
    vector_store_type: str = "faiss"  # faiss, chroma, pinecone
    persist_directory: str = "./vector_store"

    # Retrieval
    top_k: int = 5
    similarity_threshold: float = 0.7

    # Generation
    max_tokens: int = 2048
    temperature: float = 0.7


class AdvancedRAGSystem:
    """سیستم RAG پیشرفته با قابلیت‌های Systematic Generalization"""

    def __init__(self, rag_config: RAGConfig, gemini_config: GeminiConfig):
        self.rag_config = rag_config
        self.gemini_config = gemini_config
        self.vector_store = None
        self.retriever = None
        self.llm = None
        self.embeddings = None
        self.memory = None

        self._setup_components()

    def _setup_components(self):
        """تنظیم اجزای سیستم RAG"""
        try:
            # تنظیم embeddings
            self._setup_embeddings()

            # تنظیم LLM (Gemini)
            self._setup_llm()

            # تنظیم memory
            self._setup_memory()

            logger.info("✅ اجزای سیستم RAG با موفقیت تنظیم شدند")

        except Exception as e:
            logger.error(f"❌ خطا در تنظیم اجزای RAG: {e}")
            raise

    def _setup_embeddings(self):
        """تنظیم مدل embeddings"""
        try:
            # استفاده از HuggingFace embeddings
            self.embeddings = HuggingFaceEmbeddings(
                model_name=self.rag_config.embedding_model,
                model_kwargs={"device": "cpu"},
            )
            logger.info(f"✅ Embeddings تنظیم شد: {self.rag_config.embedding_model}")

        except Exception as e:
            logger.error(f"❌ خطا در تنظیم embeddings: {e}")
            raise

    def _setup_llm(self):
        """تنظیم مدل زبان Gemini"""
        try:
            self.llm = ChatGoogleGenerativeAI(
                model=self.gemini_config.model_name,
                google_api_key=self.gemini_config.api_key,
                temperature=self.gemini_config.temperature,
                max_output_tokens=self.gemini_config.max_tokens,
                convert_system_message_to_human=True,
            )
            logger.info(f"✅ LLM Gemini تنظیم شد: {self.gemini_config.model_name}")

        except Exception as e:
            logger.error(f"❌ خطا در تنظیم LLM: {e}")
            raise

    def _setup_memory(self):
        """تنظیم حافظه مکالمه"""
        self.memory = ConversationBufferMemory(
            memory_key="chat_history", return_messages=True, output_key="answer"
        )
        logger.info("✅ حافظه مکالمه تنظیم شد")

    def load_documents(
        self, file_paths: List[str], file_types: Optional[List[str]] = None
    ) -> List[Document]:
        """
        بارگذاری اسناد از فایل‌های مختلف

        Args:
            file_paths: مسیرهای فایل‌ها
            file_types: انواع فایل (txt, pdf, md)

        Returns:
            لیست اسناد بارگذاری شده
        """
        documents = []

        for i, file_path in enumerate(file_paths):
            try:
                file_type = (
                    file_types[i] if file_types and i < len(file_types) else "txt"
                )

                if file_type.lower() == "pdf":
                    loader = PyPDFLoader(file_path)
                else:
                    loader = TextLoader(file_path, encoding="utf-8")

                docs = loader.load()
                documents.extend(docs)

                logger.info(f"✅ فایل بارگذاری شد: {file_path}")

            except Exception as e:
                logger.error(f"❌ خطا در بارگذاری فایل {file_path}: {e}")
                continue

        logger.info(f"✅ مجموعاً {len(documents)} سند بارگذاری شد")
        return documents

    def split_documents(self, documents: List[Document]) -> List[Document]:
        """تقسیم اسناد به قطعات کوچک‌تر"""
        try:
            # استفاده از RecursiveCharacterTextSplitter
            text_splitter = RecursiveCharacterTextSplitter(
                chunk_size=self.rag_config.chunk_size,
                chunk_overlap=self.rag_config.chunk_overlap,
                length_function=len,
                separators=["\n\n", "\n", " ", ""],
            )

            split_docs = text_splitter.split_documents(documents)

            logger.info(f"✅ اسناد به {len(split_docs)} قطعه تقسیم شدند")
            return split_docs

        except Exception as e:
            logger.error(f"❌ خطا در تقسیم اسناد: {e}")
            raise

    def create_vector_store(
        self, documents: List[Document], save_path: Optional[str] = None
    ):
        """ایجاد و ذخیره vector store"""
        try:
            if self.rag_config.vector_store_type.lower() == "faiss":
                self.vector_store = FAISS.from_documents(documents, self.embeddings)

                if save_path:
                    self.vector_store.save_local(save_path)

            elif self.rag_config.vector_store_type.lower() == "chroma":
                self.vector_store = Chroma.from_documents(
                    documents,
                    self.embeddings,
                    persist_directory=self.rag_config.persist_directory,
                )

            else:
                raise ValueError(
                    f"نوع vector store پشتیبانی نمی‌شود: {self.rag_config.vector_store_type}"
                )

            # ایجاد retriever
            self.retriever = self.vector_store.as_retriever(
                search_kwargs={"k": self.rag_config.top_k}
            )

            logger.info(
                f"✅ Vector store ایجاد شد: {self.rag_config.vector_store_type}"
            )

        except Exception as e:
            logger.error(f"❌ خطا در ایجاد vector store: {e}")
            raise

    def load_vector_store(self, load_path: str):
        """بارگذاری vector store از فایل"""
        try:
            if self.rag_config.vector_store_type.lower() == "faiss":
                self.vector_store = FAISS.load_local(
                    load_path, self.embeddings, allow_dangerous_deserialization=True
                )

            elif self.rag_config.vector_store_type.lower() == "chroma":
                self.vector_store = Chroma(
                    persist_directory=load_path, embedding_function=self.embeddings
                )

            self.retriever = self.vector_store.as_retriever(
                search_kwargs={"k": self.rag_config.top_k}
            )

            logger.info(f"✅ Vector store بارگذاری شد از: {load_path}")

        except Exception as e:
            logger.error(f"❌ خطا در بارگذاری vector store: {e}")
            raise

    def create_systematic_prompt_template(self) -> PromptTemplate:
        """ایجاد template برای Systematic Generalization"""

        template = """
شما یک سیستم هوش مصنوعی پیشرفته هستید که در Systematic Generalization تخصص دارید.

اصول Systematic Generalization:
1. ترکیب‌پذیری (Compositionality): معنی عبارات پیچیده از ترکیب معانی اجزای آن‌ها حاصل می‌شود
2. تولیدی بودن (Productivity): توانایی تولید ترکیبات نامحدود از اجزای محدود  
3. سیستماتیک بودن (Systematicity): اگر سیستم بتواند یک ترکیب را پردازش کند، باید ترکیبات مرتبط را نیز پردازش کند

زمینه (Context):
{context}

تاریخچه مکالمه:
{chat_history}

سوال: {question}

لطفاً پاسخ خود را بر اساس اصول Systematic Generalization ارائه دهید و از دانش بازیابی شده استفاده کنید.

پاسخ سیستماتیک:
"""

        return PromptTemplate(
            template=template, input_variables=["context", "chat_history", "question"]
        )

    def create_rag_chain(self) -> ConversationalRetrievalChain:
        """ایجاد زنجیره RAG برای مکالمه"""
        try:
            prompt_template = self.create_systematic_prompt_template()

            chain = ConversationalRetrievalChain.from_llm(
                llm=self.llm,
                retriever=self.retriever,
                memory=self.memory,
                combine_docs_chain_kwargs={"prompt": prompt_template},
                return_source_documents=True,
                verbose=True,
            )

            logger.info("✅ زنجیره RAG مکالمه ایجاد شد")
            return chain

        except Exception as e:
            logger.error(f"❌ خطا در ایجاد زنجیره RAG: {e}")
            raise

    def query_systematic(self, question: str) -> Dict[str, Any]:
        """پرسش سیستماتیک با RAG"""
        try:
            if not self.retriever:
                raise ValueError(
                    "Retriever تنظیم نشده است. ابتدا vector store را ایجاد کنید."
                )

            # بازیابی اسناد مرتبط
            relevant_docs = self.retriever.get_relevant_documents(question)

            # ایجاد context از اسناد بازیابی شده
            context = "\n\n".join([doc.page_content for doc in relevant_docs])

            # تولید پاسخ با Gemini
            gemini_system = GeminiSystematicGeneralization(self.gemini_config)
            response = gemini_system.generate_systematic_response(
                prompt=question, context=context
            )

            return {
                "question": question,
                "answer": response["response"],
                "source_documents": [
                    {
                        "content": doc.page_content,
                        "metadata": doc.metadata,
                        "score": getattr(doc, "score", 0.0),
                    }
                    for doc in relevant_docs
                ],
                "context_length": len(context),
                "success": response["success"],
            }

        except Exception as e:
            logger.error(f"❌ خطا در پرسش سیستماتیک: {e}")
            return {
                "question": question,
                "answer": f"خطا در پردازش سوال: {str(e)}",
                "success": False,
                "error": str(e),
            }

    def evaluate_retrieval_quality(
        self, questions: List[str], expected_answers: List[str]
    ) -> Dict[str, Any]:
        """ارزیابی کیفیت بازیابی"""
        results = {
            "total_questions": len(questions),
            "successful_retrievals": 0,
            "average_context_length": 0,
            "detailed_results": [],
        }

        total_context_length = 0

        for i, question in enumerate(questions):
            try:
                response = self.query_systematic(question)

                if response["success"]:
                    results["successful_retrievals"] += 1
                    total_context_length += response["context_length"]

                results["detailed_results"].append(
                    {
                        "question": question,
                        "expected_answer": (
                            expected_answers[i] if i < len(expected_answers) else ""
                        ),
                        "retrieved_answer": response["answer"],
                        "context_length": response["context_length"],
                        "success": response["success"],
                    }
                )

            except Exception as e:
                logger.error(f"❌ خطا در ارزیابی سوال {i}: {e}")
                results["detailed_results"].append(
                    {"question": question, "error": str(e), "success": False}
                )

        results["retrieval_success_rate"] = (
            results["successful_retrievals"] / results["total_questions"]
        )
        results["average_context_length"] = total_context_length / max(
            results["successful_retrievals"], 1
        )

        return results

    def get_system_info(self) -> Dict[str, Any]:
        """اطلاعات سیستم"""
        return {
            "rag_config": {
                "embedding_model": self.rag_config.embedding_model,
                "chunk_size": self.rag_config.chunk_size,
                "vector_store_type": self.rag_config.vector_store_type,
                "top_k": self.rag_config.top_k,
            },
            "gemini_config": {
                "model_name": self.gemini_config.model_name,
                "temperature": self.gemini_config.temperature,
                "max_tokens": self.gemini_config.max_tokens,
            },
            "components_status": {
                "embeddings": self.embeddings is not None,
                "llm": self.llm is not None,
                "vector_store": self.vector_store is not None,
                "retriever": self.retriever is not None,
                "memory": self.memory is not None,
            },
        }


class SystematicGeneralizationRAG:
    """کلاس اصلی برای Systematic Generalization با RAG"""

    def __init__(self, rag_config: RAGConfig, gemini_config: GeminiConfig):
        self.rag_system = AdvancedRAGSystem(rag_config, gemini_config)
        self.gemini_system = GeminiSystematicGeneralization(gemini_config)

    def setup_knowledge_base(
        self,
        documents_path: str,
        file_types: Optional[List[str]] = None,
        save_vector_store: bool = True,
    ) -> bool:
        """تنظیم پایگاه دانش"""
        try:
            # بارگذاری اسناد
            documents = self.rag_system.load_documents([documents_path], file_types)

            # تقسیم اسناد
            split_docs = self.rag_system.split_documents(documents)

            # ایجاد vector store
            save_path = (
                self.rag_system.rag_config.persist_directory
                if save_vector_store
                else None
            )
            self.rag_system.create_vector_store(split_docs, save_path)

            logger.info("✅ پایگاه دانش با موفقیت تنظیم شد")
            return True

        except Exception as e:
            logger.error(f"❌ خطا در تنظیم پایگاه دانش: {e}")
            return False

    def test_systematic_generalization(
        self, training_examples: List[Dict], test_cases: List[Dict]
    ) -> Dict[str, Any]:
        """تست Systematic Generalization"""

        results = {
            "training_examples": len(training_examples),
            "test_cases": len(test_cases),
            "rag_results": [],
            "gemini_results": [],
            "comparison": {},
        }

        # تست با RAG
        for test_case in test_cases:
            rag_response = self.rag_system.query_systematic(test_case["input"])
            results["rag_results"].append(rag_response)

        # تست با Gemini خالص
        gemini_results = self.gemini_system.test_compositional_generalization(
            training_examples, test_cases
        )
        results["gemini_results"] = gemini_results

        # مقایسه نتایج
        results["comparison"] = {
            "rag_success_rate": len([r for r in results["rag_results"] if r["success"]])
            / len(test_cases),
            "gemini_accuracy": gemini_results["accuracy"],
            "rag_vs_gemini": (
                "RAG"
                if results["comparison"]["rag_success_rate"]
                > gemini_results["accuracy"]
                else "Gemini"
            ),
        }

        return results


# مثال استفاده
if __name__ == "__main__":
    # تنظیم پیکربندی‌ها
    rag_config = RAGConfig(
        embedding_model="sentence-transformers/all-MiniLM-L6-v2",
        chunk_size=1000,
        vector_store_type="faiss",
    )

    gemini_config = GeminiConfig(
        api_key="your-gemini-api-key-here", model_name="gemini-1.5-pro", temperature=0.7
    )

    # ایجاد سیستم
    systematic_rag = SystematicGeneralizationRAG(rag_config, gemini_config)

    # تنظیم پایگاه دانش (مثال)
    # systematic_rag.setup_knowledge_base("./documents", ["txt", "pdf"])

    # تست ساده
    test_question = "چگونه می‌توانم مفهوم ترکیب‌پذیری را در هوش مصنوعی پیاده‌سازی کنم؟"
    response = systematic_rag.rag_system.query_systematic(test_question)

    print("پاسخ RAG:")
    print(response["answer"])