NSAF MCP Server

from typing import Dict, Any, List, Optional
import json
from pathlib import Path
from .base_skill import BaseSkill
from ..utils.validation import validate_params
from langchain.document_loaders import DirectoryLoader, TextLoader
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain.embeddings import OpenAIEmbeddings
from langchain.vectorstores import FAISS
from langchain.chains import RetrievalQA
from langchain.chat_models import ChatOpenAI

class RAGSkill(BaseSkill):
    """Skill for handling Retrieval-Augmented Generation (RAG) content."""

    def __init__(self):
        super().__init__()
        self.name = "rag"
        self.config = self._load_config()
        self.vector_store = None
        self.qa_chain = None

    def _load_config(self) -> Dict[str, Any]:
        """Load RAG skill configuration."""
        config_path = Path(__file__).parent / "configs" / "rag.json"
        with open(config_path, "r") as f:
            return json.load(f)

    def _initialize_rag(self, data_path: str) -> None:
        """Initialize RAG components with user data."""
        # Load documents
        loader = DirectoryLoader(data_path, glob="**/*.txt", loader_cls=TextLoader)
        documents = loader.load()

        # Split documents
        text_splitter = RecursiveCharacterTextSplitter(
            chunk_size=1000,
            chunk_overlap=200,
            length_function=len,
        )
        texts = text_splitter.split_documents(documents)

        # Create vector store
        embeddings = OpenAIEmbeddings()
        self.vector_store = FAISS.from_documents(texts, embeddings)

        # Initialize QA chain
        llm = ChatOpenAI(temperature=0.7)
        self.qa_chain = RetrievalQA.from_chain_type(
            llm=llm,
            chain_type="stuff",
            retriever=self.vector_store.as_retriever()
        )

    def _generate_content(self, query: str, context: Dict[str, Any]) -> str:
        """Generate content using RAG."""
        if not self.qa_chain:
            raise ValueError("RAG system not initialized. Please upload data first.")

        # Enhance query with context
        enhanced_query = f"""
        Generate content for: {query}
        Goal: {context.get('goal', 'general')}
        Tone: {context.get('tone', 'neutral')}
        Target Audience: {context.get('target_audience', 'general')}
        Additional Context: {context.get('additional_context', '')}
        """

        # Generate response
        response = self.qa_chain.run(enhanced_query)
        return response

    def execute(self, params: Dict[str, Any]) -> Dict[str, Any]:
        """Execute RAG skill operations."""
        try:
            validate_params(params, self.config["params_schema"])
            action = params["action"]

            if action == "initialize":
                if "data_path" not in params:
                    raise ValueError("data_path is required for initialization")
                self._initialize_rag(params["data_path"])
                return {
                    "success": True,
                    "message": "RAG system initialized successfully"
                }

            elif action == "generate":
                if not all(k in params for k in ["query", "context"]):
                    raise ValueError("query and context are required for generation")
                content = self._generate_content(params["query"], params["context"])
                return {
                    "success": True,
                    "content": content
                }

            else:
                raise ValueError(f"Unknown action: {action}")

        except Exception as e:
            return {
                "success": False,
                "errors": [str(e)]
            }

    def cleanup(self) -> None:
        """Clean up resources."""
        if self.vector_store:
            # Clean up vector store resources if needed
            self.vector_store = None
        if self.qa_chain:
            # Clean up QA chain resources if needed
            self.qa_chain = None