master/ui_pack/app.py

import gradio as gr
import time
from datetime import datetime
import json

from ..rag_pipeline import ChatAssistant, get_embeddings, vretrieve
from ..rag_pipeline import determine_action, judge_prompt, request_retrieve_prompt
from ..utils import load_local

# DEVELOPER: Add or remove models here.
AVAILABLE_MODELS = {
    # "mistral large (mistral)": ("mistral-large-2", "mistral"),
    "mistral medium (mistral)": ("mistral-medium", "mistral"),
    "mistral small (mistral)": ("mistral-small", "mistral"),
    "llama3 8B" : ("llama3:8b", "ollama"),
    "llama3.1 8B": ("llama3.1:8b", "ollama"),
    "gpt-oss 20B": ("gpt-oss-20b", "ollama"),
    "gemma3 12B": ("gemma3:12b", "ollama"),
    # "gpt 4o mini": ("gpt-4o-mini", "openai"),
    # "gpt 4o": ("gpt-4o", "openai"),
}
DEFAULT_MODEL_KEY = "mistral medium (mistral)"

EMBEDDING_MODEL_ID = "alibaba-nlp/gte-multilingual-base"
VECTORSTORE_PATH = "master/knowledge/vectorstore_law"
LOG_FILE_PATH = "log.txt"
MAX_HISTORY_CONVERSATION = 1

CUSTOMER = """
{"name":"Đảng Cộng sản Việt Nam",
"birth":1930,
"current_position":"Đảng cầm quyền và là chính đảng duy nhất được phép hoạt động tại Việt Nam",
"email":"ngjabach@example.com",
"relations":[]}
"""
# System prompt for the medical assistant
sys = """
**Persona và Vai trò:**
Bạn là một Trợ lý AI Phân tích Pháp lý, hoạt động như một chuyên gia trong dự án VNPT AI Agent. Vai trò của bạn là cung cấp các phân tích khách quan, chính xác và dựa trên bằng chứng, giúp bảo vệ khách hàng khỏi các rủi ro thông tin trên không gian mạng. Giọng văn của bạn phải luôn chuyên nghiệp, có thẩm quyền và hoàn toàn trung lập.

**Nhiệm vụ cốt lõi:**
Phân tích và đối chiếu thông tin từ các nguồn được chỉ định (bài viết, tin tức, v.v.) với một cơ sở tri thức pháp lý đã được xác thực. Dựa trên sự đối chiếu này, bạn sẽ đưa ra những đánh giá và khuyến nghị cụ thể.

**Các Nguyên tắc Bắt buộc:**

1.  **Tính toàn vẹn của Nguồn:** Mọi kết luận, nhận định và phân tích của bạn phải xuất phát **DUY NHẤT** từ "Tài liệu Nguồn" được cung cấp trong mỗi yêu cầu. **Nghiêm cấm** việc sử dụng kiến thức bên ngoài, thông tin chưa được xác thực, hoặc các giả định cá nhân.
2.  **Tổng hợp và Trích xuất:** Nhiệm vụ của bạn là tổng hợp và chắt lọc thông tin liên quan từ tài liệu nguồn để trả lời truy vấn một cách chính xác. Không diễn giải hoặc sáng tạo thông tin ngoài phạm vi tài liệu.
3.  **Xử lý Thông tin Không Đầy đủ:** Nếu tài liệu được cung cấp không chứa thông tin cần thiết để trả lời một câu hỏi, bạn phải tuyên bố rõ ràng: *"Dựa trên các tài liệu được cung cấp, không có đủ thông tin để đưa ra nhận định về vấn đề này."* Tuyệt đối không được phỏng đoán.
4.  **Duy trì Tính Khách quan:** Luôn giữ một lập trường phân tích, không thiên vị. Tránh mọi hình thức suy đoán, ý kiến cá nhân, hoặc sử dụng ngôn ngữ mang tính cảm tính. Mọi lập luận phải được củng cố bằng các trích dẫn hoặc tham chiếu trực tiếp từ tài liệu.
5.  **Tập trung vào Từng Nhiệm vụ:** Chỉ sử dụng lịch sử trò chuyện để làm rõ bối cảnh của câu hỏi hiện tại. Mỗi yêu cầu phân tích là một nhiệm vụ độc lập và phải được xử lý dựa trên bộ tài liệu mới được cung cấp cho yêu cầu đó."""

# --- Initial Setup (runs once) ---
vectorstore, docs = None, None
print("Initializing models and data...")
embedding_model = get_embeddings(EMBEDDING_MODEL_ID, show_progress=False)
vectorstore, docs = load_local(VECTORSTORE_PATH, embedding_model)
print("Initialization complete.")

# --- Helper Functions ---
def log(log_txt: str):
    """Appends a log entry to the log file."""
    with open(LOG_FILE_PATH, "a", encoding="utf-8") as log_file:
        log_file.write(log_txt + "\n")

def retrieve(chat_assistant: ChatAssistant, message: str, history: list):
    history = history[-MAX_HISTORY_CONVERSATION:]
    conversation = "".join(f"User: {user_msg}\nBot: {bot_msg}\n" for user_msg, bot_msg in history)
    query_for_rag = conversation + f"User: {message}\nBot:"

    rag_query = chat_assistant.get_response(request_retrieve_prompt.format(input=query_for_rag))
    rag_query = rag_query[rag_query.lower().rfind("[") + 1: rag_query.rfind("]")]

    if "NO NEED" not in rag_query:
        retrieve_results = vretrieve(rag_query, vectorstore, docs, k=4, metric="mmr", threshold=0.7)
    else:
        retrieve_results = []

    retrieved_docs = "\n".join([f"Document {i+1}:\n" + doc.page_content for i, doc in enumerate(retrieve_results)])
    log(f"%% RAG query %%: {rag_query}")
    log(f"%% Retrieved documents %%:\n{retrieved_docs}")
    return retrieved_docs


# --- Core Chatbot Logic ---
def chatbot_logic(message: str, history: list, selected_model_key: str, customer_data: str):
    """
    Handles the main logic for receiving a message, performing RAG, and generating a response.
    """
    model_id, model_provider = AVAILABLE_MODELS[selected_model_key]

    log(f"%% Current time %%: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}")
    log(f"%% User message %%: {message}")
    log(f"%% Using Model %%: {model_id} ({model_provider})")

    try:
        chat_assistant = ChatAssistant(model_id, model_provider)
    except Exception as e:
        yield f"Error: Could not initialize the model. Please check the ID and provider. Details: {e}"
        return
    
    action = chat_assistant.get_response(determine_action.format(input=message), sys)
    action = action[action.lower().rfind("[") + 1: action.rfind("]")]
    log(f"%% Action %%: {action}")

    if "RAG" in action:
        # customer_data = json.loads(CUSTOMER)
        retrieved_docs = retrieve(chat_assistant, message, history)
        final_prompt = judge_prompt.format(customer=customer_data, law_docs=retrieved_docs, post=message)
    elif "CHAT" in action:
        final_prompt = "".join(f"User: {user_msg}\nBot: {bot_msg}\n" for user_msg, bot_msg in history)

    response = ""
    for token in chat_assistant.get_streaming_response(final_prompt, sys):
        response += token
        yield response
    log(f"%% Bot response %%: {response}")
    log("=" * 50 + "\n\n")

# --- UI Helper Function ---
def start_new_chat():
    """Clears the chatbot and input box to start a new conversation."""
    return None, ""







customer_data = json.loads(CUSTOMER)

# Define a professional, classic theme
professional_theme = gr.themes.Base(
    primary_hue=gr.themes.colors.blue,
    secondary_hue=gr.themes.colors.sky,
).set(
    button_primary_background_fill_hover='*primary_600',
)

with gr.Blocks(theme=professional_theme) as chatbot_ui:
    # Use gr.State to manage the customer data throughout the session
    customer_state = gr.State(customer_data)

    gr.Markdown("# VNPT Guardian")
    gr.Markdown("### Your Guardian AI Assistant")

    with gr.Row(equal_height=True):
        # --- Control Panel (Left Column) ---
        with gr.Column(scale=1, min_width=300): # Increased min_width for new elements
            model_selector = gr.Dropdown(
                label="Select Model",
                choices=list(AVAILABLE_MODELS.keys()),
                value=DEFAULT_MODEL_KEY,
            )
            
            new_chat_btn = gr.Button("New Chat", variant="secondary")

            # --- NEW: Customer Information Feature ---
            edit_info_btn = gr.Button("Edit Customer Info", variant="secondary")

            with gr.Accordion("Customer Details", open=False) as edit_form:
                customer_name_input = gr.Textbox(label="Name")
                customer_birth_input = gr.Number(label="Birth/Establish Year", precision=0)
                customer_pos_input = gr.Textbox(label="Description", lines=3)
                customer_email_input = gr.Textbox(label="Email")
                save_info_btn = gr.Button("Save Changes", variant="primary")
            # --- End of Feature ---
            
            gr.Markdown(
                "--- \n"
                "**Note:** Your conversations are saved for quality assurance."
            )

        # --- Chat Interface (Right Column) ---
        with gr.Column(scale=4):
            chatbot = gr.Chatbot(
                label="Chat Window", 
                height=600, 
                bubble_full_width=False
            )
            
            msg_input = gr.Textbox(
                label="Your Message", 
                placeholder="Type your question here and press Enter...",
            )

    # --- Event Handler Functions ---
    def respond(message, chat_history, selected_model_key):
        """Wrapper function to connect chatbot_logic with Gradio's state."""
        chat_history = chat_history or []
        bot_message_stream = chatbot_logic(message, chat_history, selected_model_key, customer_data)
        chat_history.append([message, ""])
        for token in bot_message_stream:
            chat_history[-1][1] = token
            yield chat_history
            
    # --- NEW: Functions for Customer Info Feature ---
    def show_edit_form(current_data):
        """Populates the form with current data and makes it visible."""
        return {
            edit_form: gr.update(open=True),
            customer_name_input: current_data['name'],
            customer_birth_input: current_data['birth'],
            customer_pos_input: current_data['current_position'],
            customer_email_input: current_data['email'],
        }

    def save_customer_info(current_data, name, birth, position, email):
        """Updates the customer data state and hides the form."""
        current_data['name'] = name
        current_data['birth'] = int(birth)
        current_data['current_position'] = position
        current_data['email'] = email
        
        gr.Info("Customer information saved successfully!")
        
        # Return the updated data to the state and close the accordion
        return {
            customer_state: current_data,
            edit_form: gr.update(open=False)
        }

    # --- Event Listeners ---
    # Event handler for submitting a message
    msg_input.submit(
        respond,
        [msg_input, chatbot, model_selector],
        [chatbot]
    ).then(
        lambda: gr.update(value=""), None, [msg_input], queue=False
    )

    # Event handler for the "New Chat" button
    new_chat_btn.click(
        start_new_chat,
        inputs=None,
        outputs=[chatbot, msg_input],
        queue=False
    )

    # --- NEW: Event handlers for the customer info feature ---
    edit_info_btn.click(
        fn=show_edit_form,
        inputs=[customer_state],
        outputs=[
            edit_form,
            customer_name_input, 
            customer_birth_input, 
            customer_pos_input, 
            customer_email_input
        ],
        queue=False
    )
    
    save_info_btn.click(
        fn=save_customer_info,
        inputs=[
            customer_state,
            customer_name_input,
            customer_birth_input,
            customer_pos_input,
            customer_email_input
        ],
        outputs=[customer_state, edit_form],
        queue=False
    )


# --- Launch the App ---
if __name__ == "__main__":
    chatbot_ui.launch(debug=True, share=True)