rag.py

from langchain_openai import OpenAIEmbeddings
from langchain_openai import AzureOpenAIEmbeddings
from langchain_community.vectorstores import FAISS
from langchain_community.document_loaders import WebBaseLoader
from langchain_text_splitters import RecursiveCharacterTextSplitter
from langchain_openai import ChatOpenAI
from langchain_core.messages import SystemMessage, HumanMessage, AIMessage
from langchain_community.callbacks import get_openai_callback
from typing import List, Dict, Any
import json
import os
import re
import shutil

# Load configuration from JSON file
def load_config():
    config_path = os.path.join(os.path.dirname(__file__), 'config.json')
    with open(config_path, 'r', encoding='utf-8') as f:
        return json.load(f)

# Load configuration
CONFIG = load_config()
CUBIX_DOCS = CONFIG['cubix_docs']
MOB_MAPPINGS = CONFIG['mob_mappings']
COMMAND_SPECS = CONFIG['command_specs']
SYSTEM_PROMPT_TEMPLATE = CONFIG['system_prompt_template']

class RAGSystem:
    def __init__(self, openai_api_key: str):
        self.openai_api_key = openai_api_key
        self.embeddings = OpenAIEmbeddings(
            openai_api_key=openai_api_key, 
            model="text-embedding-3-large",
            dimensions=1536,  # Explicitly setting dimensions for consistency
            show_progress_bar=True
        )
        self.document_store = None
        self.user_conversations = {}
        self.model = ChatOpenAI(
            openai_api_key=openai_api_key,
            model_name="gpt-4o-mini",
            temperature=0.1
        )
        self.documents_loaded = False  # Add this line to track document loading
        self.initialize_knowledge_base()
        
    def initialize_knowledge_base(self):
        """Initialize or load the knowledge base if it exists."""
        if self.documents_loaded:  # Check if documents are already loaded
            print("Documents have already been loaded.")
            return
        
        try:
            # First check if index exists
            if not os.path.exists("faiss_index"):
                print("No existing knowledge base found. Creating new one...")
                self.create_new_knowledge_base()
                self.documents_loaded = True  # Set flag to True after loading
                return

            try:
                # Try to load the existing index
                self.document_store = FAISS.load_local(
                    "faiss_index", 
                    self.embeddings, 
                    allow_dangerous_deserialization=True
                )
                
                # Verify by running a test query
                self.document_store.similarity_search_with_score("test query", k=1)
                print("Loaded existing knowledge base.")
            except AssertionError:
                # Dimension mismatch detected
                print("Embedding dimensions mismatch detected. Rebuilding knowledge base...")
                self.create_new_knowledge_base(backup_old=True)
            except Exception as e:
                # Other loading errors
                print(f"Error loading knowledge base: {e}")
                self.create_new_knowledge_base(backup_old=True)
                
            self.documents_loaded = True  # Set flag to True after loading
            
        except Exception as e:
            print(f"Error during knowledge base initialization: {e}")
            self.create_new_knowledge_base()
    
    def create_new_knowledge_base(self, backup_old=False):
        """Create a new knowledge base from scratch"""
        # Backup old index if needed
        if backup_old and os.path.exists("faiss_index"):
            backup_dir = f"faiss_index_backup"
            shutil.move("faiss_index", backup_dir)
            print(f"Old index backed up to {backup_dir}")
        
        # Create directory if needed
        os.makedirs("faiss_index", exist_ok=True)
        
        # Load documents from web
        docs = []
        for url in CUBIX_DOCS:
            try:
                loader = WebBaseLoader(url)
                web_docs = loader.load()
                
                # Clean and preprocess documents
                for doc in web_docs:
                    # Clean HTML artifacts and normalize whitespace
                    content = doc.page_content
                    content = re.sub(r'\s+', ' ', content)  # Normalize whitespace
                    content = re.sub(r'<[^>]+>', ' ', content)  # Remove HTML tags
                    content = re.sub(r'\[.*?\]', '', content)  # Remove square brackets content
                    
                    # Update document content
                    doc.page_content = content.strip()
                    
                    # Add source URL to metadata if not present
                    if 'source' not in doc.metadata:
                        doc.metadata['source'] = url
                
                docs.extend(web_docs)
                print(f"Loaded document from {url}")
            except Exception as e:
                print(f"Error loading document from {url}: {e}")
                
        if not docs:
            print("No documents could be loaded.")
            # Initialize empty vector store
            self.document_store = FAISS.from_texts(["placeholder"], self.embeddings)
            return
            
        # Split documents
        text_splitter = RecursiveCharacterTextSplitter(
            chunk_size=800,  # Increased from 500 for better context
            chunk_overlap=100,  # Increased overlap to prevent breaking important context
            separators=["\n\n", "\n", ". ", " ", ""],  # More intelligent splitting
            length_function=len
        )
        split_docs = text_splitter.split_documents(docs)
        
        # Create new vector store from scratch
        self.document_store = FAISS.from_documents(split_docs, self.embeddings)
        
        try:
            # Save the new index
            self.document_store.save_local("faiss_index")
            print(f"Successfully saved new knowledge base to faiss_index directory with {len(split_docs)} chunks")
        except Exception as e:
            print(f"Error saving knowledge base: {e}")
            
    def add_web_documents(self, urls: List[str]):
        """Add web documents to the knowledge base"""
        if not urls:
            return False
            
        # Load documents from web
        docs = []
        for url in urls:
            try:
                print(f"Loading document from {url}")
                loader = WebBaseLoader(url)
                web_docs = loader.load()
                
                # Clean and preprocess documents
                for doc in web_docs:
                    # Clean HTML artifacts and normalize whitespace
                    content = doc.page_content
                    content = re.sub(r'\s+', ' ', content)  # Normalize whitespace
                    content = re.sub(r'<[^>]+>', ' ', content)  # Remove HTML tags
                    content = re.sub(r'\[.*?\]', '', content)  # Remove square brackets content
                    
                    # Update document content
                    doc.page_content = content.strip()
                    
                    # Add source URL to metadata if not present
                    if 'source' not in doc.metadata:
                        doc.metadata['source'] = url
                        
                docs.extend(web_docs)
                print(f"Successfully loaded document from {url}")
            except Exception as e:
                print(f"Error loading document from {url}: {e}")
                
        if not docs:
            print("No documents could be loaded.")
            return False
            
        # Split documents
        text_splitter = RecursiveCharacterTextSplitter(
            chunk_size=500,  # Increased from 150 for better context
            chunk_overlap=50,  # Added overlap to prevent breaking important context
            separators=["\n\n", "\n", " ", ""],  # More intelligent splitting
            length_function=len
        )
        split_docs = text_splitter.split_documents(docs)
        
        # Create or update vector store
        if self.document_store is None:
            self.document_store = FAISS.from_documents(split_docs, self.embeddings)
        else:
            # Add documents to existing store
            self.document_store.add_documents(split_docs)
            
        try:
            # Save the updated index
            self.document_store.save_local("faiss_index")
            print(f"Added {len(split_docs)} document chunks to the knowledge base")
            return True
        except Exception as e:
            print(f"Error saving index: {e}")
            return False
        
    def _filter_response_fields(self, response_dict: dict) -> dict:
        """Filter out fields that are not in required_fields for the command type."""
        if not isinstance(response_dict, dict) or 'type' not in response_dict:
            return response_dict

        command_type = response_dict['type']
        if command_type not in COMMAND_SPECS:
            return response_dict

        required_fields = COMMAND_SPECS[command_type]['required_fields']
        return {k: v for k, v in response_dict.items() if k in required_fields}

    def generate_response(self, user_id: str, message: str) -> Dict[str, Any]:
        """Generate a response for a user message."""
        if user_id not in self.user_conversations:
            self.user_conversations[user_id] = []
            
        # Clean and normalize the input message
        message = message.strip()
        
        # Add the new message using proper LangChain message object
        self.user_conversations[user_id].append(HumanMessage(content=message))
        
        # Extract relevant context with improved search parameters
        relevant_context = ""
        context_docs = []
        if self.document_store:
            try:
                # Enhanced retrieval strategy - hybrid search approach
                # First try with higher relevance score threshold for more accurate results
                results = self.document_store.similarity_search_with_score(
                    message, 
                    k=5,  # Increased from 4 for better coverage
                    score_threshold=0.6  # Relaxed from 0.75 for better recall with Russian queries
                )
                
                if results:
                    # Sort by score and take top results
                    results.sort(key=lambda x: x[1])
                    
                    # Format with scores and metadata for better context
                    formatted_results = []
                    for doc, score in results:
                        if score < 1.5:  # Only include relevance context (lower score is better in FAISS)
                            # Extract source for better attribution
                            source = doc.metadata.get('source', 'Unknown source')
                            if isinstance(source, str) and source.startswith('http'):
                                source = source.split('/')[-1] if '/' in source else source
                                
                            # Apply intelligent content trimming - keep key information while reducing token usage
                            content = doc.page_content
                            
                            # Advanced cleaning: remove redundant whitespace and normalize
                            content = re.sub(r'\s+', ' ', content).strip()
                            
                            # Limit content length based on relevance score - higher relevance gets more tokens
                            max_length = int(min(1000, 1500))  # Increased length for better context
                            if len(content) > max_length:
                                content = content[:max_length] + "..."
                                
                            context = {
                                'content': content,
                                'score': score,
                                'source': source,
                                'metadata': doc.metadata if hasattr(doc, 'metadata') else {}
                            }
                            context_docs.append(context)
                            formatted_results.append(f"SOURCE: {source}\nRELEVANCE: {score:.2f}\n{content}")
                    
                    if formatted_results:
                        relevant_context = "\n\n---\n\n".join(formatted_results)
                
                # If no high-relevance results found, try broader search
                if not relevant_context:
                    results = self.document_store.similarity_search(
                        message,
                        k=4  # Increased from 3 for better coverage
                    )
                    if results:
                        formatted_results = []
                        for doc in results:
                            # Extract source for better attribution
                            source = doc.metadata.get('source', 'Unknown source')
                            if isinstance(source, str) and source.startswith('http'):
                                source = source.split('/')[-1] if '/' in source else source
                                
                            # Truncate content if too long (token optimization)
                            content = doc.page_content
                            if len(content) > 1000:  # Increased from 800 for better context
                                content = content[:1000] + "..."
                                
                            formatted_results.append(f"SOURCE: {source}\n{content}")
                            
                        relevant_context = "\n\n---\n\n".join(formatted_results)
            except Exception as e:
                print(f"Error during context retrieval: {e}")
                relevant_context = f"Error retrieving context: {str(e)}"
        
        # Build command specifications for the system prompt
        command_specs_text = "# Available Commands and Required Fields\n\n"
        
        for cmd_name, cmd_spec in COMMAND_SPECS.items():
            command_specs_text += f"## {cmd_name}\n"
            command_specs_text += f"Description: {cmd_spec['description']}\n"
            command_specs_text += "Required fields:\n"
            
            for field in cmd_spec['required_fields']:
                command_specs_text += f"- {field}\n"
            
            command_specs_text += "Examples:\n"
            for example in cmd_spec['examples']:
                command_specs_text += f"Query: \"{example['input']}\"\n"
                command_specs_text += f"Response: {json.dumps(example['output'], ensure_ascii=False)}\n\n"
        
        # Add mob mappings information
        mob_mappings_text = "# Entity IDs for Minecraft Mobs\n\nWhen processing kill_mob commands, use these entity IDs in the 'mob' field:\n\n"
        for mob_name, entity_id in MOB_MAPPINGS.items():
            mob_mappings_text += f"- {mob_name}: {entity_id}\n"
        mob_mappings_text += "\nWhen a player mentions a mob, always map it to the corresponding entity ID in your response.\n\n"
        
        # Build system prompt with new structure
        system_prompt = ""
        
        # Add intro
        if 'intro' in SYSTEM_PROMPT_TEMPLATE:
            system_prompt += SYSTEM_PROMPT_TEMPLATE['intro'] + "\n\n"
        
        # Add all sections from the template in order
        sections = [
            'command_rules',
            'validation_checklist',
            'examples',
            'error_handlers'
        ]
        
        for section in sections:
            if section in SYSTEM_PROMPT_TEMPLATE:
                for item in SYSTEM_PROMPT_TEMPLATE[section]:
                    system_prompt += item + "\n"
                system_prompt += "\n"
        
        # Add command specs and mob mappings from our generated text
        system_prompt += command_specs_text + "\n"
        system_prompt += mob_mappings_text + "\n"
        
        # Add relevant context if available
        if relevant_context:
            system_prompt += f"""
## Relevant Information from CubixWorld Documentation
Use this information to guide your response. If the user is asking about game updates, features, or information that appears in this documentation, provide that information in your response:

{relevant_context}

IMPORTANT: If the user is asking about information contained in the documentation above, you SHOULD provide that information in your response. Do not say you cannot answer questions about game updates or features if the information is available in the documentation.
            """
        
        # Create messages list with proper LangChain message objects
        messages = [SystemMessage(content=system_prompt)]
        
        # Get the last 5 conversation messages (reduced from 10 for more focused context)
        history = self.user_conversations[user_id][-5:] if len(self.user_conversations[user_id]) > 0 else []
        
        # Add conversation history to the messages
        messages.extend(history)
        
        # Add context summary to help model
        if context_docs:
            # Enhanced context summary with confidence signals
            context_summary = "\nContext relevance summary (sorted by relevance):\n"
            
            # Sort context by relevance score
            sorted_contexts = sorted(context_docs, key=lambda x: x['score'])
            
            for i, ctx in enumerate(sorted_contexts):
                # Add confidence classification (lower score is better in FAISS)
                confidence = "High" if ctx['score'] < 0.3 else "Medium" if ctx['score'] < 0.6 else "Moderate"
                # Add a brief preview of content with source
                source = ctx['source'].split('/')[-1] if '/' in ctx['source'] else ctx['source']
                preview = ctx['content'][:100].replace('\n', ' ') + "..."
                context_summary += f"{i+1}. [{confidence} confidence, score {ctx['score']:.2f}] From {source}: {preview}\n"
                
            messages.append(SystemMessage(content=context_summary))
        
        # Add special instruction for information questions
        if any(keyword in message.lower() for keyword in ["что", "какие", "когда", "обновление", "новое", "осеннее", "autumn"]):
            info_instruction = """
IMPORTANT INSTRUCTION: The user is asking about game information or updates. If you have relevant information in the context provided, 
you MUST share that information in your response using the message command type. Do not refuse to answer questions about game updates 
or features if the information is available in the context.
"""
            messages.append(SystemMessage(content=info_instruction))
        
        # Initialize token usage variables
        total_tokens = 0
        prompt_tokens = 0
        completion_tokens = 0
        total_cost = 0.0
        
        # Get response from the model with token tracking
        with get_openai_callback() as cb:
            response = self.model.predict_messages(messages)
            response_content = response.content
            
            # Store token usage metrics
            total_tokens = cb.total_tokens
            prompt_tokens = cb.prompt_tokens
            completion_tokens = cb.completion_tokens
            total_cost = cb.total_cost
            
        try:
            # Simple JSON extraction
            if "```json" in response_content:
                json_start = response_content.find("```json") + 7
                json_end = response_content.find("```", json_start)
                response_content = response_content[json_start:json_end].strip()
            elif "```" in response_content:
                json_start = response_content.find("```") + 3
                json_end = response_content.find("```", json_start)
                response_content = response_content[json_start:json_end].strip()
                
            parsed_response = json.loads(response_content)
            filtered_response = self._filter_response_fields(parsed_response)
            json_response = {'response': filtered_response}
            
            # Add token usage information to the response
            json_response["token_usage"] = {
                "total_tokens": total_tokens,
                "prompt_tokens": prompt_tokens,
                "completion_tokens": completion_tokens,
                "total_cost_usd": total_cost
            }
            
            # Store the bot's response in conversation history as proper AIMessage
            self.user_conversations[user_id].append(AIMessage(content=json.dumps(json_response)))
            
            return json_response
            
        except json.JSONDecodeError:
            # Simple fallback if JSON parsing fails
            fallback_response = {
               'type': 'message',
               'message': response_content
            }
            
            self.user_conversations[user_id].append(AIMessage(content=json.dumps(fallback_response)))
            
            return fallback_response