fix requirements

app.py

@@ -5,107 +5,331 @@ import gradio as gr
 import uuid
 from utils.chat import ChatLaborLaw
 
+# ==============================================================================
+# 1. GLOBAL INITIALIZATION (ทำครั้งเดียวตอนแอปเริ่มทำงาน)
+# ==============================================================================
+# --- Langfuse Handler ---
+LANGFUSE_HANDLER = CallbackHandler()
 
-# Function to initialize a new session and create chatbot instance for that session
-async def initialize_session():
-    session_id = str(uuid.uuid4())[:8]
-    chatbot = ChatLaborLaw()
-    # chatbot = Chat("gemini-2.0-flash")
-    history = []
-    return "", session_id, chatbot, history
+# --- LLM Initialization ---
+# (ปรับแก้ส่วนนี้เพื่อเลือกว่าจะใช้โมเดลไหนเป็น default)
+MODEL_NAME_LLM = "jai-chat-1-3-2" 
+TEMPERATURE = 0
 
+if MODEL_NAME_LLM == "jai-chat-1-3-2":
+    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
+    )
+elif MODEL_NAME_LLM == "gemini-2.0-flash":
+    LLM_MAIN = ChatGoogleGenerativeAI(
+        model="gemini-1.5-flash", 
+        google_api_key=os.getenv("GOOGLE_API_KEY"),
+        temperature=TEMPERATURE,
+        max_output_tokens=2048, 
+        convert_system_message_to_human=True, 
+    )
+else:
+    raise ValueError(f"Unsupported LLM model '{MODEL_NAME_LLM}'.")
 
-# Function to handle user input and chatbot response
-async def chat_function(prompt, history, session_id, chatbot):
-    if chatbot is None:
-        return history, "", session_id, chatbot # Skip if chatbot not ready
-    
-    # Append the user's input to the message history
-    history.append({"role": "user", "content": prompt})
+# --- Database and Retriever Initialization ---
+MONGO_CONNECTION_STR = os.getenv('MONGO_CONNECTION_STRING')
+MONGO_DATABASE = os.getenv('MONGO_DATABASE')
+MONGO_COLLECTION = os.getenv('MONGO_COLLECTION')
+
+MONGO_CLIENT = MongoClient(MONGO_CONNECTION_STR)
+DB = MONGO_CLIENT[MONGO_DATABASE]
+MONGO_COLLECTION_INSTANCE = DB[MONGO_COLLECTION]
+RETRIEVER = RerankRetriever()
+
+print("Global objects initialized successfully.")
+
+# ==============================================================================
+# 2. HELPER FUNCTIONS (แปลงมาจากเมธอดใน Class)
+# ==============================================================================
+
+def format_main_context(list_of_documents):
+    formatted_docs = []
+    for i, doc in enumerate(list_of_documents):
+        metadata = doc.metadata
+        formatted = f"Doc{i}\n{metadata.get('law_name', '-')}\nมาตรา\t{metadata.get('section_number', '-')}\n{doc.page_content}\nประกาศ\t{metadata.get('publication_date', '-')}\nเริ่มใช้\t{metadata.get('effective_date', '-')}"
+        formatted_docs.append(formatted)
+    return "\n\n".join(formatted_docs)
 
-    # Get the response from the chatbot
-    response = await chatbot.chat(prompt)  # ใช้ await ได้แล้ว
+def format_ref_context(list_of_docs):
+    formatted_ref_docs = []
+    for i, doc in enumerate(list_of_docs):
+        formatted = f"{doc.get('law_name', '-')}\nมาตรา\t{doc.get('section_number', '-')}\n{doc.get('text', '-')}"
+        formatted_ref_docs.append(formatted)
+    return "\n\n".join(formatted_ref_docs)
+
+def get_main_context(user_query, **kwargs):
+    compression_retriever = RETRIEVER.get_compression_retriever(**kwargs)
+    return compression_retriever.invoke(user_query)
+
+def get_ref_context(main_context_docs):
+    all_reference_ids = set()
+    for context in main_context_docs:
+        references_list = context.metadata.get('references', [])
+        if isinstance(references_list, list):
+            for ref_str in references_list:
+                all_reference_ids.add(ref_str.replace("มาตรา", "").strip())
     
-    # Append the assistant's response to the message history
-    history.append({"role": "assistant", "content": response})
+    if not all_reference_ids:
+        return []
+
+    mongo_query = {"law_type": "summary", "section_number": {"$in": list(all_reference_ids)}}
+    projection = {"text": 1, "law_name": 1, "section_number": 1}
+    return list(MONGO_COLLECTION_INSTANCE.find(mongo_query, projection))
+
+# ==============================================================================
+# 3. CORE LOGIC FUNCTIONS (RAG / Non-RAG)
+# ==============================================================================
+
+async def call_rag(user_input: str, langchain_history: list) -> str:
+    context_docs = get_main_context(user_input, law_type="summary")
+    main_context_str = format_main_context(context_docs)
+
+    ref_context_docs = get_ref_context(context_docs)
+    ref_context_str = format_ref_context(ref_context_docs) if ref_context_docs else "-"
+
+    rag_input_data = {
+        "question": user_input,
+        "main_context": main_context_str,
+        "ref_context": ref_context_str,
+        "history": langchain_history
+    }
     
-    return history, "", session_id, chatbot
+    try:
+        prompt_messages = RAG_CHAT_PROMPT.format_messages(**rag_input_data)
+        response = await LLM_MAIN.ainvoke(prompt_messages, config={"callbacks": [LANGFUSE_HANDLER]})
+        clean_response = re.sub(r"<[^>]+>|#+", "", response.content).strip()
+        return clean_response
+    except Exception as e:
+        print(f"Error during RAG LLM call: {e}")
+        return "ขออภัย ระบบขัดข้องขณะประมวลผลคำตอบ"
 
+async def call_non_rag(user_input: str) -> str:
+    prompt_messages = NON_RAG_PROMPT.format(user_input=user_input)
+    response = await LLM_MAIN.ainvoke(prompt_messages, config={"callbacks": [LANGFUSE_HANDLER]})
+    return response.content.strip() if response and response.content else "ขออภัย ระบบไม่สามารถตอบคำถามได้ในขณะนี้"
 
-# Function to save feedback with chat history
-async def send_feedback(feedback, history, session_id, chatbot):
-    os.makedirs("app/feedback", exist_ok=True)
-    filename = f"app/feedback/feedback_{session_id}.txt"
+# ==============================================================================
+# 4. GRADIO EVENT HANDLERS
+# ==============================================================================
+
+def initialize_session():
+    """รีเซ็ต State ทั้งหมดสำหรับ Session ใหม่"""
+    session_id = str(uuid.uuid4())[:8]
+    return "", session_id, [], []  # user_input, session_id, ui_history, langchain_history
+
+async def chat_orchestrator(prompt: str, ui_history: list, langchain_history: list):
+    """
+    ฟังก์ชันหลักที่จัดการการสนทนาทั้งหมด
+    """
+    if not prompt.strip():
+        return ui_history, langchain_history, ""
+
+    # 1. อัปเดต Langchain History ด้วยข้อความใหม่
+    langchain_history.append(HumanMessage(content=prompt))
+
+    # 2. จำแนกประเภทของ Input
+    try:
+        history_content_list = [msg.content for msg in langchain_history]
+        input_type = classify_input_type(prompt, history=history_content_list)
+    except Exception as e:
+        print(f"Error classifying input type: {e}. Defaulting to Non-RAG.")
+        input_type = "Non-RAG"
+
+    # 3. เรียกใช้ Flow ที่เหมาะสม
+    if input_type == "RAG":
+        ai_response = await call_rag(prompt, langchain_history)
+    else:
+        ai_response = await call_non_rag(prompt)
+
+    # 4. อัปเดต History ทั้งสองรูปแบบ
+    langchain_history.append(AIMessage(content=ai_response))
+    ui_history.append((prompt, ai_response))
+    
+    # 5. ส่งค่ากลับไปอัปเดต UI และ State
+    return ui_history, langchain_history, "" # ui_history, langchain_history, user_input (ให้เป็นค่าว่าง)
+
+def send_feedback(feedback: str, history: list, session_id: str):
+    """บันทึก Feedback"""
+    if not feedback.strip(): return ""
+    os.makedirs("feedback", exist_ok=True)
+    filename = f"feedback/feedback_{session_id}.txt"
     with open(filename, "a", encoding="utf-8") as f:
-        f.write("=== Feedback Received ===\n")
-        f.write(f"Session ID: {session_id}\n")
-        f.write(f"Feedback: {feedback}\n")
-        f.write("Chat History:\n")
-        for msg in history:
-            f.write(f"{msg['role']}: {msg['content']}\n")
+        f.write(f"=== Feedback Received ===\nSession ID: {session_id}\nFeedback: {feedback}\nChat History:\n")
+        for user_msg, assistant_msg in history:
+            f.write(f"User: {user_msg}\nAssistant: {assistant_msg}\n")
         f.write("\n--------------------------\n\n")
-    return "" # Clear feedback input
+    gr.Info("ขอบคุณสำหรับข้อเสนอแนะ!")
+    return ""
 
+# ==============================================================================
+# 5. GRADIO UI DEFINITION
+# ==============================================================================
 
-# Create the Gradio interface
 with gr.Blocks(theme=gr.themes.Soft(primary_hue="amber")) as demo:
     gr.Markdown("# สอบถามเรื่องกฎหมายแรงงาน")
 
-    # Initialize State
-    session_state = gr.State()
-    chatbot_instance = gr.State()
-    chatbot_history = gr.State([])
+    # --- States ---
+    # session_id_state: เก็บ ID ของ session ปัจจุบัน
+    # langchain_history_state: เก็บประวัติการสนทนาในรูปแบบ Langchain Message (HumanMessage, AIMessage)
+    session_id_state = gr.State()
+    langchain_history_state = gr.State([])
 
-    # Chat UI
-    chatbot_interface = gr.Chatbot(type="messages", label="Chat History")
-    user_input = gr.Textbox(placeholder="Type your message here...", elem_id="user_input", lines=1)
-
-    submit_button = gr.Button("Send")
-    clear_button = gr.Button("Delete Chat History")
+    # --- UI Components ---
+    chatbot_interface = gr.Chatbot(label="ประวัติการสนทนา", height=550, bubble_styling=False, show_copy_button=True)
+    user_input = gr.Textbox(placeholder="พิมพ์คำถามของคุณที่นี่...", label="คำถาม", lines=2)
+    with gr.Row():
+        submit_button = gr.Button("ส่ง", variant="primary", scale=4)
+        clear_button = gr.Button("เริ่มการสนทนาใหม่", scale=1)
 
-    # Submit actions
+    # --- Event Wiring ---
     submit_button.click(
-        fn=chat_function,
-        inputs=[user_input, chatbot_history, session_state, chatbot_instance],
-        outputs=[chatbot_interface, user_input, session_state, chatbot_instance]
+        fn=chat_orchestrator,
+        inputs=[user_input, chatbot_interface, langchain_history_state],
+        outputs=[chatbot_interface, langchain_history_state, user_input]
     )
-
     user_input.submit(
-        fn=chat_function,
-        inputs=[user_input, chatbot_history, session_state, chatbot_instance],
-        outputs=[chatbot_interface, user_input, session_state, chatbot_instance]
+        fn=chat_orchestrator,
+        inputs=[user_input, chatbot_interface, langchain_history_state],
+        outputs=[chatbot_interface, langchain_history_state, user_input]
     )
-
-    # # Clear history
-    # clear_button.click(lambda: [], outputs=chatbot_interface)
     clear_button.click(
         fn=initialize_session,
         inputs=[],
-        outputs=[user_input, session_state, chatbot_instance, chatbot_history]
-    ).then(
-        fn=lambda: gr.update(value=[]),
-        inputs=[],
-        outputs=chatbot_interface
+        outputs=[user_input, session_id_state, chatbot_interface, langchain_history_state],
+        queue=False
     )
-
-
-    # Feedback section
-    with gr.Row():
-        feedback_input = gr.Textbox(placeholder="Send us feedback...", label="Feedback")
-        send_feedback_button = gr.Button("Send Feedback")
-
+    
+    with gr.Accordion("ส่งข้อเสนอแนะ (Feedback)", open=False):
+        feedback_input = gr.Textbox(placeholder="ความคิดเห็นของคุณมีความสำคัญต่อการพัฒนาของเรา...", label="Feedback", lines=2, scale=4)
+        send_feedback_button = gr.Button("ส่ง Feedback")
+    
     send_feedback_button.click(
         fn=send_feedback,
-        inputs=[feedback_input, chatbot_history, session_state, chatbot_instance],
-        outputs=[feedback_input]
+        inputs=[feedback_input, chatbot_interface, session_id_state],
+        outputs=[feedback_input],
+        queue=False
     )
 
-    # Initialize session on load
     demo.load(
         fn=initialize_session,
         inputs=[],
-        outputs=[user_input, session_state, chatbot_instance, chatbot_history]
+        outputs=[user_input, session_id_state, chatbot_interface, langchain_history_state]
     )
 
-# Launch
-demo.launch(share=True)
+demo.queue().launch()
+
+
+# # Function to initialize a new session and create chatbot instance for that session
+# async def initialize_session():
+#     session_id = str(uuid.uuid4())[:8]
+#     chatbot = ChatLaborLaw()
+#     # chatbot = Chat("gemini-2.0-flash")
+#     history = []
+#     return "", session_id, chatbot, history
+
+
+# # Function to handle user input and chatbot response
+# async def chat_function(prompt, history, session_id, chatbot):
+#     if chatbot is None:
+#         return history, "", session_id, chatbot # Skip if chatbot not ready
+    
+#     # Append the user's input to the message history
+#     history.append({"role": "user", "content": prompt})
+
+#     # Get the response from the chatbot
+#     response = await chatbot.chat(prompt)  # ใช้ await ได้แล้ว
+    
+#     # Append the assistant's response to the message history
+#     history.append({"role": "assistant", "content": response})
+    
+#     return history, "", session_id, chatbot
+
+
+# # Function to save feedback with chat history
+# async def send_feedback(feedback, history, session_id, chatbot):
+#     os.makedirs("app/feedback", exist_ok=True)
+#     filename = f"app/feedback/feedback_{session_id}.txt"
+#     with open(filename, "a", encoding="utf-8") as f:
+#         f.write("=== Feedback Received ===\n")
+#         f.write(f"Session ID: {session_id}\n")
+#         f.write(f"Feedback: {feedback}\n")
+#         f.write("Chat History:\n")
+#         for msg in history:
+#             f.write(f"{msg['role']}: {msg['content']}\n")
+#         f.write("\n--------------------------\n\n")
+#     return "" # Clear feedback input
+
+
+# # Create the Gradio interface
+# with gr.Blocks(theme=gr.themes.Soft(primary_hue="amber")) as demo:
+#     gr.Markdown("# สอบถามเรื่องกฎหมายแรงงาน")
+
+#     # Initialize State
+#     session_state = gr.State()
+#     chatbot_instance = gr.State()
+#     chatbot_history = gr.State([])
+
+#     # Chat UI
+#     chatbot_interface = gr.Chatbot(type="messages", label="Chat History")
+#     user_input = gr.Textbox(placeholder="Type your message here...", elem_id="user_input", lines=1)
+
+#     submit_button = gr.Button("Send")
+#     clear_button = gr.Button("Delete Chat History")
+
+#     # Submit actions
+#     submit_button.click(
+#         fn=chat_function,
+#         inputs=[user_input, chatbot_history, session_state, chatbot_instance],
+#         outputs=[chatbot_interface, user_input, session_state, chatbot_instance]
+#     )
+
+#     user_input.submit(
+#         fn=chat_function,
+#         inputs=[user_input, chatbot_history, session_state, chatbot_instance],
+#         outputs=[chatbot_interface, user_input, session_state, chatbot_instance]
+#     )
+
+#     # # Clear history
+#     # clear_button.click(lambda: [], outputs=chatbot_interface)
+#     clear_button.click(
+#         fn=initialize_session,
+#         inputs=[],
+#         outputs=[user_input, session_state, chatbot_instance, chatbot_history]
+#     ).then(
+#         fn=lambda: gr.update(value=[]),
+#         inputs=[],
+#         outputs=chatbot_interface
+#     )
+
+
+#     # Feedback section
+#     with gr.Row():
+#         feedback_input = gr.Textbox(placeholder="Send us feedback...", label="Feedback")
+#         send_feedback_button = gr.Button("Send Feedback")
+
+#     send_feedback_button.click(
+#         fn=send_feedback,
+#         inputs=[feedback_input, chatbot_history, session_state, chatbot_instance],
+#         outputs=[feedback_input]
+#     )
+
+#     # Initialize session on load
+#     demo.load(
+#         fn=initialize_session,
+#         inputs=[],
+#         outputs=[user_input, session_state, chatbot_instance, chatbot_history]
+#     )
+
+# # Launch
+# demo.launch(share=True)