rag_processor.py

import os
from dotenv import load_dotenv
from operator import itemgetter
from langchain_groq import ChatGroq
from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
from langchain_core.runnables import RunnableParallel, RunnablePassthrough
from langchain_core.output_parsers import StrOutputParser
from langchain_core.runnables.history import RunnableWithMessageHistory
from langchain_core.documents import Document
from query_expansion import expand_query_simple
from typing import List, Optional
import time

class GroqAPIKeyManager:
    """Manages multiple Groq API keys with automatic rotation and fallback."""
    
    def __init__(self, api_keys: List[str]):
        """
        Initialize with a list of API keys.
        
        Args:
            api_keys: List of Groq API keys to use
        """
        self.api_keys = [key for key in api_keys if key and key != "your_groq_api_key_here"]
        if not self.api_keys:
            raise ValueError("No valid API keys provided!")
        
        self.current_index = 0
        self.failed_keys = set()
        self.success_count = {key: 0 for key in self.api_keys}
        self.failure_count = {key: 0 for key in self.api_keys}
        
        print(f"🔑 API Key Manager: Loaded {len(self.api_keys)} API keys")
    
    def get_current_key(self) -> str:
        """Get the current API key."""
        return self.api_keys[self.current_index]
    
    def mark_success(self, api_key: str):
        """Mark an API key as successful."""
        if api_key in self.success_count:
            self.success_count[api_key] += 1
            # Remove from failed keys if it was there
            if api_key in self.failed_keys:
                self.failed_keys.remove(api_key)
                print(f"   ✅ API Key #{self.api_keys.index(api_key) + 1} recovered!")
    
    def mark_failure(self, api_key: str):
        """Mark an API key as failed."""
        if api_key in self.failure_count:
            self.failure_count[api_key] += 1
            self.failed_keys.add(api_key)
    
    def rotate_to_next_key(self) -> bool:
        """
        Rotate to the next available API key.
        
        Returns:
            True if a new key is available, False if all keys failed
        """
        initial_index = self.current_index
        attempts = 0
        
        while attempts < len(self.api_keys):
            self.current_index = (self.current_index + 1) % len(self.api_keys)
            attempts += 1
            
            current_key = self.api_keys[self.current_index]
            
            # If we've tried all keys, allow retry even failed ones
            if attempts >= len(self.api_keys):
                print(f"   ⚠️ All keys attempted, retrying with key #{self.current_index + 1}")
                return True
            
            # Skip recently failed keys unless it's been a while
            if current_key not in self.failed_keys:
                print(f"   🔄 Switching to API Key #{self.current_index + 1}")
                return True
        
        return False
    
    def get_statistics(self) -> str:
        """Get statistics about API key usage."""
        stats = []
        for i, key in enumerate(self.api_keys):
            success = self.success_count[key]
            failure = self.failure_count[key]
            status = "❌ FAILED" if key in self.failed_keys else "✅ ACTIVE"
            masked_key = key[:8] + "..." + key[-4:] if len(key) > 12 else "***"
            stats.append(f"   Key #{i+1} ({masked_key}): {success} success, {failure} failures [{status}]")
        return "\n".join(stats)


def load_api_keys_from_hf_secrets() -> List[str]:
    """
    Load API keys from Hugging Face Spaces Secrets.
    
    In your Hugging Face Space settings, add these secrets:
    - GROQ_API_KEY_1
    - GROQ_API_KEY_2
    - GROQ_API_KEY_3
    - GROQ_API_KEY_4
    
    Returns:
        List of API keys retrieved from HF secrets
    """
    api_keys = []
    secret_names = ['GROQ_API_KEY_1', 'GROQ_API_KEY_2', 'GROQ_API_KEY_3', 'GROQ_API_KEY_4']
    
    print("🔐 Loading API keys from Hugging Face Secrets...")
    
    for secret_name in secret_names:
        try:
            # HF Spaces secrets are available as environment variables
            api_key = os.getenv(secret_name)
            
            if api_key and api_key.strip() and api_key != "your_groq_api_key_here":
                api_keys.append(api_key.strip())
                print(f"   ✅ Loaded: {secret_name}")
            else:
                print(f"   ⚠️ Not found or empty: {secret_name}")
        except Exception as e:
            print(f"   ❌ Error loading {secret_name}: {str(e)}")
    
    # ADD THIS RETURN STATEMENT - this was missing!
    return api_keys


def create_llm_with_fallback(
    api_key_manager: GroqAPIKeyManager,
    model_name: str,
    temperature: float,
    max_retries: int = 3
) -> ChatGroq:
    """
    Create a ChatGroq LLM with automatic API key fallback.
    
    Args:
        api_key_manager: Manager handling multiple API keys
        model_name: Name of the model to use
        temperature: Temperature setting
        max_retries: Maximum number of retry attempts
    
    Returns:
        ChatGroq instance
    """
    for attempt in range(max_retries):
        current_key = api_key_manager.get_current_key()
        
        try:
            llm = ChatGroq(
                model_name=model_name,
                api_key=current_key,
                temperature=temperature
            )
            # Test the connection with a simple call
            test_result = llm.invoke("test")
            api_key_manager.mark_success(current_key)
            return llm
            
        except Exception as e:
            error_msg = str(e).lower()
            api_key_manager.mark_failure(current_key)
            
            # Check if it's a rate limit or auth error
            if "rate" in error_msg or "limit" in error_msg:
                print(f"   ⚠️ Rate limit hit on API Key #{api_key_manager.current_index + 1}")
            elif "auth" in error_msg or "api" in error_msg:
                print(f"   ❌ Authentication failed on API Key #{api_key_manager.current_index + 1}")
            else:
                print(f"   ❌ Error with API Key #{api_key_manager.current_index + 1}: {str(e)[:50]}")
            
            # Try next key if available
            if attempt < max_retries - 1:
                if api_key_manager.rotate_to_next_key():
                    print(f"   🔄 Retrying with next API key (Attempt {attempt + 2}/{max_retries})...")
                    time.sleep(1)  # Brief pause before retry
                else:
                    raise ValueError("All API keys failed!")
            else:
                raise ValueError(f"Failed to initialize LLM after {max_retries} attempts")
    
    raise ValueError("Failed to create LLM with any available API key")


def create_multi_query_retriever(base_retriever, llm, strategy: str = "balanced"):
    """Wraps a base retriever with query expansion capabilities."""
    def multi_query_retrieve(query: str) -> List[Document]:
        """Retrieves documents using expanded query variations."""
        query_variations = expand_query_simple(query, strategy=strategy, llm=llm)
        all_docs = []
        seen_content = set()
        for i, query_var in enumerate(query_variations):
            try:
                docs = base_retriever.invoke(query_var)
                for doc in docs:
                    content_hash = hash(doc.page_content)
                    if content_hash not in seen_content:
                        seen_content.add(content_hash)
                        all_docs.append(doc)
            except Exception as e:
                print(f"   ✗ Query Expansion Error (Query {i+1}): {str(e)[:50]}")
                continue
        print(f"   📊 Query Expansion: Retrieved {len(all_docs)} unique documents.")
        return all_docs
    return multi_query_retrieve


def get_system_prompt(temperature: float) -> str:
    """
    Returns a system prompt dynamically based on temperature setting.
    
    Temperature ranges:
    - 0.0-0.4: Highly factual, structured, conservative
    - 0.4-0.8: Balanced approach with moderate creativity
    - 0.8-1.0: Creative, engaging, storytelling mode
    """
    
    if temperature <= 0.4:
        # Conservative, structured prompt
        return """You are CogniChat, an expert document analysis assistant specializing in comprehensive and well-structured answers.

RESPONSE GUIDELINES:

**Structure & Formatting:**
- Start with a direct answer to the question
- Use **bold** for key terms, important concepts, and technical terminology
- Use bullet points (•) for lists, features, or multiple items
- Use numbered lists (1., 2., 3.) for steps, procedures, or sequential information
- Use ### Headers to organize different sections or topics
- Add blank lines between sections for readability

**Source Citation:**
- Always cite information using: [Source: filename, Page: X] and cite it at the end of the entire answer only
- Place citations at the end of your final answer only 
- Do not cite sources within the body of your answer
- Multiple sources: [Source: doc1.pdf, Page: 3; doc2.pdf, Page: 7]

**Completeness:**
- Provide thorough, detailed answers using ALL relevant information from context
- Summarize and properly elaborate each point for increased clarity
- If the question has multiple parts, address each part clearly

**Accuracy:**
- ONLY use information from the provided context documents below
- If information is incomplete, state what IS available and what ISN'T
- If the answer isn't in the context, clearly state: "I cannot find this information in the uploaded documents"
- Never make assumptions or add information not in the context

---

{context}

---

Now answer the following question comprehensively using the context above:"""
    
    elif temperature <= 0.8:
        # Balanced prompt
        return """You are CogniChat, an intelligent document analysis assistant that combines accuracy with engaging communication.

RESPONSE GUIDELINES:

**Communication Style:**
- Present information in a clear, engaging manner
- Use **bold** for emphasis on important concepts
- Balance structure with natural flow
- Make complex topics accessible and interesting

**Content Approach:**
- Ground your response firmly in the provided context
- Add helpful explanations and connections between concepts
- Use analogies or examples when they help clarify ideas (but keep them brief)
- Organize information logically with headers (###) and lists where appropriate

**Source Attribution:**
- Cite sources at the end: [Source: filename, Page: X]
- Be transparent about what the documents do and don't contain

**Accuracy:**
- Base your answer on the context documents provided
- If information is partial, explain what's available
- Acknowledge gaps: "The documents don't cover this aspect"

---

{context}

---

Now answer the following question in an engaging yet accurate way:"""
    
    else:  # temperature > 0.8
        # Creative BUT CLEAR prompt - REVISED VERSION
        return """You are CogniChat, a creative document analyst who makes complex information clear, memorable, and engaging.

🎯 YOUR CORE MISSION: **CLARITY FIRST, CREATIVITY SECOND**

Make information easier to understand, not harder. Your creativity should illuminate, not obscure.

**CREATIVE CLARITY PRINCIPLES:**

1. **Simplify, Don't Complicate**
   - Break down complex concepts into simple, digestible parts
   - Use everyday language alongside technical terms
   - Explain jargon immediately in plain English
   - Short sentences for complex ideas, varied length for rhythm

2. **Smart Use of Examples & Analogies** (Use Sparingly!)
   - Only use analogies when they genuinely make something clearer
   - Keep analogies simple and relatable (everyday objects/experiences)
   - Never use metaphors that require explanation themselves
   - If you can explain it directly in simple terms, do that instead

3. **Engaging Structure**
   - Start with the core answer in one clear sentence
   - Use **bold** to highlight key takeaways
   - Break information into logical chunks with ### headers
   - Use bullet points for clarity, not decoration
   - Add brief transition phrases to connect ideas smoothly

4. **Conversational Yet Precise**
   - Write like you're explaining to a smart friend
   - Use "you" and active voice to engage readers
   - Ask occasional rhetorical questions only if they aid understanding
   - Vary sentence length to maintain interest
   - Use emojis sparingly (1-2 max) and only where they add clarity

5. **Visual Clarity**
   - Strategic use of formatting: **bold** for key terms, *italics* for emphasis
   - White space between sections for easy scanning
   - Progressive disclosure: simple concepts first, details after
   - Numbered lists for sequences, bullets for related items

**WHAT TO AVOID:**
- ❌ Flowery or overly descriptive language
- ❌ Complex metaphors that need their own explanation
- ❌ Long narrative storytelling that buries the facts
- ❌ Multiple rhetorical questions in a row
- ❌ Overuse of emojis or exclamation points
- ❌ Making simple things sound complicated

**ACCURACY BOUNDARIES:**
- ✅ Creative explanation and presentation of facts
- ✅ Simple, helpful examples from common knowledge
- ✅ Reorganizing information for better understanding
- ❌ Never invent facts not in the documents
- ❌ Don't contradict source material
- ❌ If info is missing, say so clearly and briefly

**Source Attribution:**
- End with: [Source: filename, Page: X]
- Keep it simple and clear

---

{context}

---

Now, explain the answer clearly and engagingly. Remember: if your grandmother couldn't understand it, simplify more:"""
    


def create_rag_chain(
    retriever,
    get_session_history_func,
    enable_query_expansion=True,
    expansion_strategy="balanced",
    model_name: str = "moonshotai/kimi-k2-instruct",
    temperature: float = 0.2,
    api_keys: Optional[List[str]] = None
):
    """
    Creates an advanced RAG chain with temperature-adaptive prompting and API key rotation.
    
    Args:
        retriever: Document retriever
        get_session_history_func: Function to get session history
        enable_query_expansion: Whether to enable query expansion
        expansion_strategy: Strategy for query expansion
        model_name: Name of the LLM model
        temperature: Temperature setting (0.0-1.0)
        api_keys: Optional list of API keys. If None, loads from environment
    """
    
    # Load API keys from HF Secrets
    if api_keys is None:
        api_keys = load_api_keys_from_hf_secrets()
    
    if not api_keys:
        raise ValueError(
            "No valid API keys found! Please set GROQ_API_KEY or GROQ_API_KEY_1, "
            "GROQ_API_KEY_2, GROQ_API_KEY_3, GROQ_API_KEY_4 in your .env file"
        )
    
    # Initialize API key manager
    api_key_manager = GroqAPIKeyManager(api_keys)
    
    print(f"⚙️ RAG: Initializing LLM - Model: {model_name}, Temp: {temperature}")
    
    # Display creativity mode based on temperature
    if temperature <= 0.4:
        creativity_mode = "FACTUAL & STRUCTURED"
    elif temperature <= 0.8:
        creativity_mode = "BALANCED & ENGAGING"
    else:
        creativity_mode = "CREATIVE & STORYTELLING"
    print(f"🎭 Creativity Mode: {creativity_mode}")
    
    # Create LLM with fallback
    llm = create_llm_with_fallback(api_key_manager, model_name, temperature)
    print(f"✅ LLM initialized with API Key #{api_key_manager.current_index + 1}")

    if enable_query_expansion:
        print(f"✨ RAG: Query Expansion ENABLED (Strategy: {expansion_strategy})")
        enhanced_retriever = create_multi_query_retriever(
            base_retriever=retriever,
            llm=llm,
            strategy=expansion_strategy
        )
    else:
        enhanced_retriever = retriever

    rewrite_template = """You are an expert at optimizing search queries for document retrieval.

Given the conversation history and a follow-up question, create a comprehensive standalone question that:
1. Incorporates all relevant context from the chat history
2. Expands abbreviations and resolves all pronouns (it, they, this, that, etc.)
3. Includes key technical terms and concepts that would help find relevant documents
4. Maintains the original intent, specificity, and detail level
5. If the question asks for comparison or multiple items, ensure all items are in the query

Chat History:
{chat_history}

Follow-up Question: {question}

Optimized Standalone Question:"""
    rewrite_prompt = ChatPromptTemplate.from_messages([
        ("system", rewrite_template),
        MessagesPlaceholder(variable_name="chat_history"),
        ("human", "{question}")
    ])
    query_rewriter = rewrite_prompt | llm | StrOutputParser()

    def format_docs(docs):
        """Format retrieved documents with clear structure and metadata."""
        if not docs:
            return "No relevant documents found in the knowledge base."
        
        formatted_parts = []
        for i, doc in enumerate(docs, 1):
            source = doc.metadata.get('source', 'Unknown Document')
            page = doc.metadata.get('page', 'N/A')
            rerank_score = doc.metadata.get('rerank_score')
            content = doc.page_content.strip()
            
            doc_header = f"{'='*60}\nDOCUMENT {i}\n{'='*60}"
            metadata_line = f"Source: {source} | Page: {page}"
            if rerank_score:
                metadata_line += f" | Relevance: {rerank_score:.3f}"
            
            formatted_parts.append(
                f"{doc_header}\n"
                f"{metadata_line}\n"
                f"{'-'*60}\n"
                f"{content}\n"
            )
        return f"RETRIEVED CONTEXT ({len(docs)} documents):\n\n" + "\n".join(formatted_parts)

    # Get temperature-adaptive system prompt
    rag_template = get_system_prompt(temperature)
    
    rag_prompt = ChatPromptTemplate.from_messages([
        ("system", rag_template),
        MessagesPlaceholder(variable_name="chat_history"),
        ("human", "{question}"),
    ])

    # Rewriter input construction
    rewriter_input = RunnableParallel({
        "question": itemgetter("question"),
        "chat_history": itemgetter("chat_history"),
    })

    # Main retrieval pipeline
    retrieval_chain = rewriter_input | query_rewriter | enhanced_retriever | format_docs

    # Final conversational RAG chain
    conversational_rag_chain = RunnableParallel({
        "context": retrieval_chain,
        "question": itemgetter("question"),
        "chat_history": itemgetter("chat_history"),
    }) | rag_prompt | llm | StrOutputParser()

    chain_with_memory = RunnableWithMessageHistory(
        conversational_rag_chain,
        get_session_history_func,
        input_messages_key="question",
        history_messages_key="chat_history",
    )
    
    print("✅ RAG: Chain created successfully.")
    print("\n" + api_key_manager.get_statistics())
    
    return chain_with_memory, api_key_manager  # Return manager for statistics