Zeri00/Cogni-Chat-document-reader-v2

1import os2from dotenv import load_dotenv3from operator import itemgetter4from langchain_groq import ChatGroq5from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder6from langchain_core.runnables import RunnableParallel, RunnablePassthrough7from langchain_core.output_parsers import StrOutputParser8from langchain_core.runnables.history import RunnableWithMessageHistory9from langchain_core.documents import Document10from query_expansion import expand_query_simple11from typing import List, Optional12import time13 14class GroqAPIKeyManager:15    def __init__(self, api_keys: List[str]):16        self.api_keys = [key for key in api_keys if key and key != "your_groq_api_key_here"]17        if not self.api_keys:18            raise ValueError("No valid API keys provided!")19        20        self.current_index = 021        self.failed_keys = set()22        self.success_count = {key: 0 for key in self.api_keys}23        self.failure_count = {key: 0 for key in self.api_keys}24        25        print(f"API Key Manager: Loaded {len(self.api_keys)} API keys")26    27    def get_current_key(self) -> str:28        return self.api_keys[self.current_index]29    30    def mark_success(self, api_key: str):31        if api_key in self.success_count:32            self.success_count[api_key] += 133            if api_key in self.failed_keys:34                self.failed_keys.remove(api_key)35                print(f"API Key #{self.api_keys.index(api_key) + 1} recovered!")36    37    def mark_failure(self, api_key: str):38        if api_key in self.failure_count:39            self.failure_count[api_key] += 140            self.failed_keys.add(api_key)41    42    def rotate_to_next_key(self) -> bool:43        initial_index = self.current_index44        attempts = 045        46        while attempts < len(self.api_keys):47            self.current_index = (self.current_index + 1) % len(self.api_keys)48            attempts += 149            50            current_key = self.api_keys[self.current_index]51            if attempts >= len(self.api_keys):52                print(f"All keys attempted, retrying with key #{self.current_index + 1}")53                return True54            if current_key not in self.failed_keys:55                print(f"Switching to API Key #{self.current_index + 1}")56                return True57        58        return False59    60    def get_statistics(self) -> str:61        stats = []62        for i, key in enumerate(self.api_keys):63            success = self.success_count[key]64            failure = self.failure_count[key]65            status = "FAILED" if key in self.failed_keys else "ACTIVE"66            masked_key = key[:8] + "..." + key[-4:] if len(key) > 12 else "***"67            stats.append(f"   Key #{i+1} ({masked_key}): {success} success, {failure} failures [{status}]")68        return "\n".join(stats)69 70 71def load_api_keys_from_hf_secrets() -> List[str]:72    api_keys = []73    secret_names = ['GROQ_API_KEY_1', 'GROQ_API_KEY_2', 'GROQ_API_KEY_3', 'GROQ_API_KEY_4']74    75    print("Loading API keys from Hugging Face Secrets...")76    77    for secret_name in secret_names:78        try:79            api_key = os.getenv(secret_name)80            81            if api_key and api_key.strip() and api_key != "your_groq_api_key_here":82                api_keys.append(api_key.strip())83                print(f" Loaded: {secret_name}")84            else:85                print(f" Not found or empty: {secret_name}")86        except Exception as e:87            print(f" Error loading {secret_name}: {str(e)}")88    return api_keys89 90 91def create_llm_with_fallback(92    api_key_manager: GroqAPIKeyManager,93    model_name: str,94    temperature: float,95    max_retries: int = 396) -> ChatGroq:97    for attempt in range(max_retries):98        current_key = api_key_manager.get_current_key()99        100        try:101            llm = ChatGroq(102                model_name=model_name,103                api_key=current_key,104                temperature=temperature105            )106            test_result = llm.invoke("test")107            api_key_manager.mark_success(current_key)108            return llm109            110        except Exception as e:111            error_msg = str(e).lower()112            api_key_manager.mark_failure(current_key)113            if "rate" in error_msg or "limit" in error_msg:114                print(f"  Rate limit hit on API Key #{api_key_manager.current_index + 1}")115            elif "auth" in error_msg or "api" in error_msg:116                print(f"  Authentication failed on API Key #{api_key_manager.current_index + 1}")117            else:118                print(f"  Error with API Key #{api_key_manager.current_index + 1}: {str(e)[:50]}")119            120            if attempt < max_retries - 1:121                if api_key_manager.rotate_to_next_key():122                    print(f" Retrying with next API key (Attempt {attempt + 2}/{max_retries})...")123                    time.sleep(1)124                else:125                    raise ValueError("All API keys failed!")126            else:127                raise ValueError(f"Failed to initialize LLM after {max_retries} attempts")128    129    raise ValueError("Failed to create LLM with any available API key")130 131 132def create_multi_query_retriever(base_retriever, llm, strategy: str = "balanced"):133    def multi_query_retrieve(query: str) -> List[Document]:134        query_variations = expand_query_simple(query, strategy=strategy, llm=llm)135        all_docs = []136        seen_content = set()137        for i, query_var in enumerate(query_variations):138            try:139                docs = base_retriever.invoke(query_var)140                for doc in docs:141                    content_hash = hash(doc.page_content)142                    if content_hash not in seen_content:143                        seen_content.add(content_hash)144                        all_docs.append(doc)145            except Exception as e:146                print(f" Query Expansion Error (Query {i+1}): {str(e)[:50]}")147                continue148        print(f" Query Expansion: Retrieved {len(all_docs)} unique documents.")149        return all_docs150    return multi_query_retrieve151 152 153def get_system_prompt(temperature: float) -> str:154    if temperature <= 0.4:155        return """You are CogniChat, an expert document analysis assistant specializing in comprehensive and well-structured answers.156 157RESPONSE GUIDELINES:158 159**Structure & Formatting:**160- Start with a direct answer to the question161- Use **bold** for key terms, important concepts, and technical terminology162- Use bullet points (•) for lists, features, or multiple items163- Use numbered lists (1., 2., 3.) for steps, procedures, or sequential information164- Use ### Headers to organize different sections or topics165- Add blank lines between sections for readability166 167**Source Citation:**168- Always cite information using: [Source: filename, Page: X] and cite it at the end of the entire answer only169- Place citations at the end of your final answer only 170- Do not cite sources within the body of your answer171- Multiple sources: [Source: doc1.pdf, Page: 3; doc2.pdf, Page: 7]172 173**Completeness:**174- Provide thorough, detailed answers using ALL relevant information from context175- Summarize and properly elaborate each point for increased clarity176- If the question has multiple parts, address each part clearly177 178**Accuracy:**179- ONLY use information from the provided context documents below180- If information is incomplete, state what IS available and what ISN'T181- If the answer isn't in the context, clearly state: "I cannot find this information in the uploaded documents"182- Never make assumptions or add information not in the context183 184---185 186{context}187 188---189 190Now answer the following question comprehensively using the context above:"""191    192    elif temperature <= 0.8:193        return """You are CogniChat, an intelligent document analysis assistant that combines accuracy with engaging communication.194 195RESPONSE GUIDELINES:196 197**Communication Style:**198- Present information in a clear, engaging manner199- Use **bold** for emphasis on important concepts200- Balance structure with natural flow201- Make complex topics accessible and interesting202 203**Content Approach:**204- Ground your response firmly in the provided context205- Add helpful explanations and connections between concepts206- Use analogies or examples when they help clarify ideas (but keep them brief)207- Organize information logically with headers (###) and lists where appropriate208 209**Source Attribution:**210- Cite sources at the end: [Source: filename, Page: X]211- Be transparent about what the documents do and don't contain212 213**Accuracy:**214- Base your answer on the context documents provided215- If information is partial, explain what's available216- Acknowledge gaps: "The documents don't cover this aspect"217 218---219 220{context}221 222---223 224Now answer the following question in an engaging yet accurate way:"""225    226    else:  # temperature > 0.8227        # Creative BUT CLEAR prompt - REVISED VERSION228        return """You are CogniChat, a creative document analyst who makes complex information clear, memorable, and engaging.229 230 YOUR CORE MISSION: **CLARITY FIRST, CREATIVITY SECOND**231 232Make information easier to understand, not harder. Your creativity should illuminate, not obscure.233 234**CREATIVE CLARITY PRINCIPLES:**235 2361. **Simplify, Don't Complicate**237   - Break down complex concepts into simple, digestible parts238   - Use everyday language alongside technical terms239   - Explain jargon immediately in plain English240   - Short sentences for complex ideas, varied length for rhythm241 2422. **Smart Use of Examples & Analogies** (Use Sparingly!)243   - Only use analogies when they genuinely make something clearer244   - Keep analogies simple and relatable (everyday objects/experiences)245   - Never use metaphors that require explanation themselves246   - If you can explain it directly in simple terms, do that instead247 2483. **Engaging Structure**249   - Start with the core answer in one clear sentence250   - Use **bold** to highlight key takeaways251   - Break information into logical chunks with ### headers252   - Use bullet points for clarity, not decoration253   - Add brief transition phrases to connect ideas smoothly254 2554. **Conversational Yet Precise**256   - Write like you're explaining to a smart friend257   - Use "you" and active voice to engage readers258   - Ask occasional rhetorical questions only if they aid understanding259   - Vary sentence length to maintain interest260   - Use emojis sparingly (1-2 max) and only where they add clarity261 2625. **Visual Clarity**263   - Strategic use of formatting: **bold** for key terms, *italics* for emphasis264   - White space between sections for easy scanning265   - Progressive disclosure: simple concepts first, details after266   - Numbered lists for sequences, bullets for related items267 268**WHAT TO AVOID:**269-  Flowery or overly descriptive language270-  Complex metaphors that need their own explanation271-  Long narrative storytelling that buries the facts272-  Multiple rhetorical questions in a row273-  Overuse of emojis or exclamation points274-  Making simple things sound complicated275 276**ACCURACY BOUNDARIES:**277-  Creative explanation and presentation of facts278-  Simple, helpful examples from common knowledge279-  Reorganizing information for better understanding280-  Never invent facts not in the documents281-  Don't contradict source material282-  If info is missing, say so clearly and briefly283 284**Source Attribution:**285- End with: [Source: filename, Page: X]286- Keep it simple and clear287 288---289 290{context}291 292---293 294Now, explain the answer clearly and engagingly. Remember: if your grandmother couldn't understand it, simplify more:"""295    296 297 298def create_rag_chain(299    retriever,300    get_session_history_func,301    enable_query_expansion=True,302    expansion_strategy="balanced",303    model_name: str = "moonshotai/kimi-k2-instruct",304    temperature: float = 0.2,305    api_keys: Optional[List[str]] = None306):307    if api_keys is None:308        api_keys = load_api_keys_from_hf_secrets()309    310    if not api_keys:311        raise ValueError(312            "No valid API keys found! Please set GROQ_API_KEY or GROQ_API_KEY_1, "313            "GROQ_API_KEY_2, GROQ_API_KEY_3, GROQ_API_KEY_4 in your .env file"314        )315    316    api_key_manager = GroqAPIKeyManager(api_keys)317    318    print(f" RAG: Initializing LLM - Model: {model_name}, Temp: {temperature}")319    320    if temperature <= 0.4:321        creativity_mode = "FACTUAL & STRUCTURED"322    elif temperature <= 0.8:323        creativity_mode = "BALANCED & ENGAGING"324    else:325        creativity_mode = "CREATIVE & STORYTELLING"326    print(f"Creativity Mode: {creativity_mode}")327    328    llm = create_llm_with_fallback(api_key_manager, model_name, temperature)329    print(f"LLM initialized with API Key #{api_key_manager.current_index + 1}")330 331    if enable_query_expansion:332        print(f"RAG: Query Expansion ENABLED (Strategy: {expansion_strategy})")333        enhanced_retriever = create_multi_query_retriever(334            base_retriever=retriever,335            llm=llm,336            strategy=expansion_strategy337        )338    else:339        enhanced_retriever = retriever340 341    rewrite_template = """You are an expert at optimizing search queries for document retrieval.342 343Given the conversation history and a follow-up question, create a comprehensive standalone question that:3441. Incorporates all relevant context from the chat history3452. Expands abbreviations and resolves all pronouns (it, they, this, that, etc.)3463. Includes key technical terms and concepts that would help find relevant documents3474. Maintains the original intent, specificity, and detail level3485. If the question asks for comparison or multiple items, ensure all items are in the query349 350Chat History:351{chat_history}352 353Follow-up Question: {question}354 355Optimized Standalone Question:"""356    rewrite_prompt = ChatPromptTemplate.from_messages([357        ("system", rewrite_template),358        MessagesPlaceholder(variable_name="chat_history"),359        ("human", "{question}")360    ])361    query_rewriter = rewrite_prompt | llm | StrOutputParser()362 363    def format_docs(docs):364        if not docs:365            return "No relevant documents found in the knowledge base."366        367        formatted_parts = []368        for i, doc in enumerate(docs, 1):369            source = doc.metadata.get('source', 'Unknown Document')370            page = doc.metadata.get('page', 'N/A')371            rerank_score = doc.metadata.get('rerank_score')372            content = doc.page_content.strip()373            374            doc_header = f"{'='*60}\nDOCUMENT {i}\n{'='*60}"375            metadata_line = f"Source: {source} | Page: {page}"376            if rerank_score:377                metadata_line += f" | Relevance: {rerank_score:.3f}"378            379            formatted_parts.append(380                f"{doc_header}\n"381                f"{metadata_line}\n"382                f"{'-'*60}\n"383                f"{content}\n"384            )385        return f"RETRIEVED CONTEXT ({len(docs)} documents):\n\n" + "\n".join(formatted_parts)386 387    rag_template = get_system_prompt(temperature)388    389    rag_prompt = ChatPromptTemplate.from_messages([390        ("system", rag_template),391        MessagesPlaceholder(variable_name="chat_history"),392        ("human", "{question}"),393    ])394 395    rewriter_input = RunnableParallel({396        "question": itemgetter("question"),397        "chat_history": itemgetter("chat_history"),398    })399 400    retrieval_chain = rewriter_input | query_rewriter | enhanced_retriever | format_docs401 402    conversational_rag_chain = RunnableParallel({403        "context": retrieval_chain,404        "question": itemgetter("question"),405        "chat_history": itemgetter("chat_history"),406    }) | rag_prompt | llm | StrOutputParser()407 408    chain_with_memory = RunnableWithMessageHistory(409        conversational_rag_chain,410        get_session_history_func,411        input_messages_key="question",412        history_messages_key="chat_history",413    )414    415    print("RAG: Chain created successfully.")416    print("\n" + api_key_manager.get_statistics())417    418    return chain_with_memory, api_key_manager