chat_3.py

import os
import uuid
from dotenv import load_dotenv

from utils.chat_prompts import (
    NON_RAG_PROMPT,
    RAG_CHAT_PROMPT_ENG,
    RAG_CHAT_PROMPT_TH,
    RAG_CHAT_PROMPT_KOREAN,
    QUERY_REWRITING_PROMPT_OBJ
)
from get_retriever_2 import final_retrievers # Retriever ที่อาจผ่านการ rerank
from input_classifier import classify_input_type, detect_language

from langchain_openai import ChatOpenAI
from langchain_core.messages import HumanMessage, AIMessage
from langchain_google_genai import ChatGoogleGenerativeAI

from langfuse.callback import CallbackHandler

# Load environment variables from .env file
load_dotenv()

langfuse_handler = CallbackHandler(
  secret_key=os.environ['LANGFUSE_SECRET_KEY'],
  public_key=os.environ['LANGFUSE_PUBLIC_KEY'],
  host="https://us.cloud.langfuse.com"
)

class Chat:
    def __init__(self, model_name_llm="jai-chat-1-3-2", temperature=0):
        self.session_id = str(uuid.uuid4())[:8]
        self.model_name_llm = model_name_llm # เก็บชื่อโมเดลไว้เผื่อใช้

        # --- LLM Initialization ---
        # เราจะสร้าง LLM สำหรับการตอบคำถามหลัก และอาจจะใช้ตัวเดียวกันหรือตัวที่เล็กกว่าสำหรับ Query Rewriting
        if model_name_llm == "jai-chat-1-3-2":
            self.llm_main = ChatOpenAI(
                model=model_name_llm,
                api_key=os.getenv("JAI_API_KEY"),
                base_url=os.getenv("CHAT_BASE_URL"),
                temperature=temperature,
                max_tokens=2048,
                max_retries=2,
                seed=13
            )
            # สำหรับ query rewriting, ถ้า JAI มี model ที่เร็วกว่า/ถูกกว่า ก็สามารถใช้ตัวนั้นได้
            # หรือใช้ตัวเดียวกันไปก่อน
            self.llm_rewriter = self.llm_main
        elif model_name_llm == "gemini-2.0-flash":
            GEMINI_API_KEY = os.getenv("GOOGLE_API_KEY")
            if not GEMINI_API_KEY:
                raise ValueError("GOOGLE_API_KEY (for Gemini) not found in environment variables.")
            
            common_gemini_config = {
                "google_api_key": GEMINI_API_KEY,
                "temperature": temperature,
                "max_output_tokens": 2048, # ใช้ max_output_tokens สำหรับ Gemini
                "convert_system_message_to_human": True, # สำคัญสำหรับ Gemini
                # "top_p": 0.95,
                # "top_k": 40,
            }
            self.llm_main = ChatGoogleGenerativeAI(
                model="gemini-1.5-flash-latest", # e.g., "gemini-1.5-flash-latest" or "gemini-1.5-pro-latest"
                **common_gemini_config
            )
            # Gemini มี "gemini-1.5-flash-latest" ที่เหมาะกับงานเร็วๆ เช่น rewriting
            # ถ้า llm_main เป็น pro, อาจจะใช้ flash สำหรับ rewriter
            # ในที่นี้ถ้า llm_main เป็น flash อยู่แล้ว ก็ใช้ตัวเดียวกัน
            self.llm_rewriter = ChatGoogleGenerativeAI(
                model="gemini-2.0-flash", # หรือ model_name_llm ถ้ามันคือ flash อยู่แล้ว
                **common_gemini_config
            )

        else:
            raise ValueError(f"Unsupported LLM model '{model_name_llm}'.")

        self.history = [] # Store Langchain Message objects

    def append_history(self, message: [HumanMessage, AIMessage]):
        self.history.append(message)

    def get_formatted_history_for_llm(self, n_turns: int = 3) -> list:
        """Returns the last n_turns of history as a list of Message objects."""
        return self.history[-(n_turns * 2):]

    def get_stringified_history_for_rewrite(self, n_turns: int = 2) -> str:
        """
        Formats the last n_turns of history (excluding the current un-added user input)
        as a string for the query rewriter prompt.
        """
        history_to_format = self.history[-(n_turns * 2):]
        if not history_to_format:
            return "No history available."

        history_str_parts = []
        for msg in history_to_format:
            role = "User" if isinstance(msg, HumanMessage) else "AI"
            history_str_parts.append(f"{role}: {msg.content}")
        return "\n".join(history_str_parts)

    def classify_input(self, user_input: str) -> str:
        history_content_list = [msg.content for msg in self.history] # เอาเฉพาะ content
        return classify_input_type(user_input, history=history_content_list)

    def format_docs(self, docs: list) -> str:
        return "\n\n".join(doc.page_content for doc in docs)

    def get_retriever_and_prompt(self, lang_code: str):
        """
        Returns the appropriate retriever and RAG prompt based on the language.
        Handles potential errors if retriever or prompt is not found.
        """
        retriever = final_retrievers.get(lang_code)
        
        if lang_code == "Thai":
            prompt_template = RAG_CHAT_PROMPT_TH
        elif lang_code == "Korean":
            prompt_template = RAG_CHAT_PROMPT_KOREAN
        elif lang_code == "English":
            prompt_template = RAG_CHAT_PROMPT_ENG
        else:
            print(f"Warning: Unsupported language '{lang_code}' for RAG. Defaulting to English.")
            retriever = final_retrievers.get('English') # Default to English
            prompt_template = RAG_CHAT_PROMPT_ENG

        if not retriever:
            # Attempt to get any available retriever if the specific or default English one is missing
            available_langs = list(final_retrievers.keys())
            if available_langs:
                fallback_lang = available_langs[0]
                retriever = final_retrievers[fallback_lang]
                print(f"Warning: No retriever for '{lang_code}' or 'English'. Using first available: '{fallback_lang}'.")
                # Match prompt to this fallback retriever if possible
                if fallback_lang == "Thai": prompt_template = RAG_CHAT_PROMPT_TH
                elif fallback_lang == "Korean": prompt_template = RAG_CHAT_PROMPT_KOREAN
                else: prompt_template = RAG_CHAT_PROMPT_ENG # Default to English prompt anyway
            else:
                raise ValueError("CRITICAL: No retrievers configured at all.")

        if not prompt_template: # Should not happen if logic above is correct
             raise ValueError(f"CRITICAL: No RAG prompt template found for language '{lang_code}' or effective fallback.")
            
        return retriever, prompt_template

    def _rewrite_query_if_needed(self, user_input: str, input_lang: str) -> str:
        """
        Internal method to rewrite the user query using chat history if there is history.
        """
        if not self.history: # No history, no need to rewrite
            return user_input

        chat_history_str = self.get_stringified_history_for_rewrite(n_turns=2)

        try:
            rewrite_prompt_messages = QUERY_REWRITING_PROMPT_OBJ.format_messages(
                chat_history=chat_history_str,
                question=user_input
            )
            
            response = self.llm_rewriter.invoke(rewrite_prompt_messages)
            rewritten_query = response.content.strip()

            # Basic validation of rewritten query
            if rewritten_query and len(rewritten_query) < (len(user_input) + 250) and len(rewritten_query) > 0: # Avoid empty or overly long
                print(f"Original query: '{user_input}', Rewritten query for retriever: '{rewritten_query}'")
                return rewritten_query
            else:
                print(f"Rewritten query validation failed or empty. Using original: '{user_input}'")
                return user_input
        except Exception as e:
            print(f"Error during query rewriting: {e}. Using original query.")
            return user_input

    def call_rag(self, user_input: str, input_lang: str) -> str:
        try:
            retriever, selected_rag_prompt = self.get_retriever_and_prompt(input_lang)
        except ValueError as e:
            print(f"Error in RAG setup: {e}")
            return f"Sorry, I encountered a configuration issue for {input_lang} RAG. Please contact support."

        # --- Query Rewriting Step ---
        # user_input is the current question. self.history does not yet contain it.
        query_for_retriever = self._rewrite_query_if_needed(user_input, input_lang)

        print(f"Retrieving documents for query: '{query_for_retriever}' (lang: {input_lang})")
        try:
            context_docs = retriever.invoke(query_for_retriever)
        except Exception as e:
            print(f"Error during document retrieval: {e}")
            return "Sorry, I had trouble finding relevant information for your query."
        
        print(f"Retrieved {len(context_docs)} documents. (Max possible after rerank: {os.getenv('FINAL_TOP_K_RERANK', 'N/A')})")
        # for i, doc in enumerate(context_docs):
        # print(f"  Doc {i+1} (Score: {doc.metadata.get('rerank_score', 'N/A')}): {doc.page_content[:100]}...")

        context_str = self.format_docs(context_docs)
        
        # History for the RAG prompt (LLM context)
        history_for_llm_prompt = self.get_formatted_history_for_llm(n_turns=3)

        rag_input_data = {
            "question": user_input, # Use the original user_input for the question to LLM
            "context": context_str,
            "history": history_for_llm_prompt # Pass Langchain Message objects
        }

        try:
            prompt_messages = selected_rag_prompt.format_messages(**rag_input_data)
            # print(f"DEBUG: RAG Prompt Messages: {prompt_messages}")
            # response = self.llm_main.invoke(prompt_messages)
            response = self.llm_main.invoke(prompt_messages, config={"callbacks": [langfuse_handler]})
            return response.content.strip()
        except Exception as e:
            print(f"Error during RAG LLM call: {e}")
            return "Sorry, I encountered an error while generating the response."

    def call_non_rag(self, user_input: str, input_lang: str) -> str:
        # Ensure NON_RAG_PROMPT is a ChatPromptTemplate or a string
        try:
            if hasattr(NON_RAG_PROMPT, "format_messages"): # Is a ChatPromptTemplate
                prompt_messages = NON_RAG_PROMPT.format(user_input=user_input, input_lang=input_lang)
            elif isinstance(NON_RAG_PROMPT, str): # Is a plain string template
                formatted_prompt_str = NON_RAG_PROMPT.format(user_input=user_input, input_lang=input_lang)
                prompt_messages = [HumanMessage(content=formatted_prompt_str)]
            else:
                raise TypeError("NON_RAG_PROMPT is of an unsupported type.")
            
            # response = self.llm_main.invoke(prompt_messages)
            response = self.llm_main.invoke(prompt_messages, config={"callbacks": [langfuse_handler]})

            return response.content.strip()

        except Exception as e:
            print(f"Error during Non-RAG LLM call: {e}")
            return "Sorry, I had trouble processing your general request."

    def chat(self, user_input: str) -> str:
        print(f"\n\n-- USER INPUT: {user_input} --")
        
        # 1. Detect language of the current user input
        # This should ideally be robust and return one of "Thai", "Korean", "English"
        # or a defined set that get_retriever_and_prompt can handle.
        try:
            input_lang_detected = detect_language(user_input)
            print(f"Language detected: {input_lang_detected}")
        except Exception as e:
            print(f"Error detecting language: {e}. Defaulting to Thai.") # Or your most common language
            input_lang_detected = "Thai"

        history_before_current_input = self.history[:] # Make a copy of history *before* appending current input

        # Append current user input to history
        self.append_history(HumanMessage(content=user_input))
        
        # 3. Classify input type based on current input and now-updated history
        try:
            input_type = self.classify_input(user_input) # classify_input now sees history with current input
        except Exception as e:
            print(f"Error classifying input type: {e}. Defaulting to Non-RAG.")
            input_type = "Non-RAG"

        ai_response_content = ""
        if input_type == "RAG":
            print("[RAG FLOW]")
            # Pass the original user_input and the detected language.
            # call_rag will handle query rewriting internally using history_before_current_input
            ai_response_content = self.call_rag_v2(user_input, input_lang_detected, history_before_current_input)
        else: # "Non-RAG" or any other classification
            print(f"[{input_type} FLOW (Treated as NON-RAG)]")
            ai_response_content = self.call_non_rag(user_input, input_lang_detected)

        # 4. Append AI response to history
        self.append_history(AIMessage(content=ai_response_content))
        
        print(f"AI:::: {ai_response_content}")
        return ai_response_content


    # New call_rag version that accepts history_before_current_input
    def call_rag_v2(self, user_input: str, input_lang: str, history_for_rewrite: list) -> str:
        try:
            retriever, selected_rag_prompt = self.get_retriever_and_prompt(input_lang)
        except ValueError as e:
            print(f"Error in RAG setup: {e}")
            return f"Sorry, I encountered a configuration issue for {input_lang} RAG. Please contact support."

        # --- Query Rewriting Step ---
        query_for_retriever = self._rewrite_query_if_needed_v2(user_input, history_for_rewrite)


        print(f"Retrieving documents for query: '{query_for_retriever}' (lang: {input_lang})")
        try:
            context_docs = retriever.invoke(query_for_retriever)
        except Exception as e:
            print(f"Error during document retrieval: {e}")
            return "Sorry, I had trouble finding relevant information for your query."
        
        print(f"Retrieved {len(context_docs)} documents.")

        context_str = self.format_docs(context_docs)
        print(f"\n----> CONTEXT DOCS (from call_rag_v2)\n{context_str}")
        
        # History for the RAG prompt (LLM context) - this should be the full history *including* current user_input
        history_for_llm_prompt = self.get_formatted_history_for_llm(n_turns=3)


        rag_input_data = {
            "question": user_input,
            "context": context_str,
            "history": history_for_llm_prompt
        }

        try:
            prompt_messages = selected_rag_prompt.format_messages(**rag_input_data)
            # response = self.llm_main.invoke(prompt_messages)
            response = self.llm_main.invoke(prompt_messages, config={"callbacks": [langfuse_handler]})

            return response.content.strip()

        except Exception as e:
            print(f"Error during RAG LLM call: {e}")
            return "Sorry, I encountered an error while generating the response."

    # New rewrite_query version that accepts history
    def _rewrite_query_if_needed_v2(self, user_input: str, history_list: list) -> str:
        if not history_list:
            return user_input

        # Format the passed history_list (which is before current user_input)
        history_str_parts = []
        for msg in history_list[-(2*2):]: # Take last 2 turns from the provided history
             role = "User" if isinstance(msg, HumanMessage) else "AI"
             history_str_parts.append(f"{role}: {msg.content}")
        chat_history_str = "\n".join(history_str_parts) if history_str_parts else "No relevant history."

        try:
            rewrite_prompt_messages = QUERY_REWRITING_PROMPT_OBJ.format_messages(
                chat_history=chat_history_str,
                question=user_input
            )
            response = self.llm_rewriter.invoke(rewrite_prompt_messages)
            rewritten_query = response.content.strip()

            if rewritten_query and len(rewritten_query) < (len(user_input) + 250) and len(rewritten_query) > 0:
                print(f"Original query: '{user_input}', Rewritten query for retriever: '{rewritten_query}'")
                return rewritten_query
            else:
                print(f"Rewritten query validation failed. Using original: '{user_input}'")
                return user_input
        except Exception as e:
            print(f"Error during query rewriting: {e}. Using original query.")
            return user_input