gradio_hf.py

import requests
import pandas as pd
import gradio as gr
import time
import re
import os
from dotenv import load_dotenv
load_dotenv()

API_BASE_URL = os.getenv("API_BASE_URL")


def clear_chat():
    """Clear all conversation history and reset all states via API"""
    try:
        requests.post(f"{API_BASE_URL}/clear_memory",
                      json={"reset_cache": True, "reset_model": False})
    except Exception as e:
        print(f"Warning: clear_memory failed: {e}")

    # Return empty values for all components that need to be reset
    # chatbot, state, raw_documents, chatbot_state, textbox
    return [], {}, {}, [], ""


def respond_to_chat(message, history, state, selected_model, debug_mode):
    """Handle chat responses with image support"""
    start = time.perf_counter()
    # Handle multimodal input - message is a dict with 'text' and 'files' keys
    if isinstance(message, dict):
        text_message = message.get('text', '')
        files = message.get('files', [])

        # Process any uploaded images
        image_path = None
        if files:
            image_path = files[0]
            print(image_path)

        user_query = text_message
    else:
        user_query = message
        image_path = None

    # call API
    try:
        if image_path:
            with open(image_path, 'rb') as f:
                files = {"image": f}
                data = {"message": user_query, "debug": debug_mode}
                resp = requests.post(
                    f"{API_BASE_URL}/chat_with_image", files=files, data=data, timeout=120)
        else:
            payload = {"message": user_query, "debug": debug_mode}
            resp = requests.post(f"{API_BASE_URL}/chat",
                                 json=payload, timeout=120)

        if resp.status_code == 200:
            j = resp.json()
            response = j.get("response", "")
            specs_advantages = j.get("specs_advantages")
            raw_documents = j.get("raw_documents")
            outputs = j.get("outputs")
        else:
            response = f"Error: API status {resp.status_code}"
            specs_advantages, raw_documents, outputs = None, None, None
    except Exception as e:
        response = f"Error calling API: {e}"
        specs_advantages, raw_documents, outputs = None, None, None

    end = time.perf_counter()

    if state is None:
        state = {}

    # specs_advantages is now always a dict for both products and solutions
    if isinstance(specs_advantages, dict):
        state["specs_advantages"] = specs_advantages
    state["raw_documents"] = raw_documents
    state["outputs"] = outputs

    return response


def show_specs(state, history=None):
    specs_map = state.get("specs_advantages", {})
    columns = ["Thông số"]
    raw_data = []
    if history is None:
        history = []

    if not specs_map:
        df = pd.DataFrame({}, columns=columns)
        markdown_table = df.to_markdown(index=False)
        # Fix: Use proper ChatInterface message format
        history.append(
            {"role": "assistant", "content": f"📄 Thông số kỹ thuật\n{markdown_table}"})
        return history

    # print(specs_map)
    for prod_id, data in specs_map.items():
        spec = data.get("specification", None)
        model = data.get("model", "")
        url = data.get("url", "")

        # Handle both products and solution packages
        if url:
            full_name = f"**[{data['name']} {model}]({url})**"
        else:
            full_name = f"**[{data['name']} {model}]**"

        if full_name not in columns:
            columns.append(full_name)

        if spec:
            # Check if this is a solution package (contains markdown table)
            if "### 📦" in spec:
                # For solution packages, parse the markdown table properly
                lines = spec.split('\n')
                in_table = False
                headers = []

                for line in lines:
                    line = line.strip()
                    if '|' in line and '---' not in line and line.startswith('|') and line.endswith('|'):
                        cells = [cell.strip()
                                 for cell in line.split('|')[1:-1]]

                        if not in_table:
                            # This is the header row
                            headers = cells
                            in_table = True
                            continue

                        # This is a data row
                        if len(cells) >= len(headers):
                            for i, header in enumerate(headers):
                                if i < len(cells):
                                    param_name = header
                                    param_value = cells[i]

                                    existing_row = None
                                    for row in raw_data:
                                        if row["Thông số"] == param_name:
                                            existing_row = row
                                            break

                                    if existing_row:
                                        existing_row[full_name] = param_value
                                    else:
                                        new_row = {"Thông số": param_name}
                                        for col in columns[1:]:
                                            new_row[col] = ""
                                        new_row[full_name] = param_value
                                        raw_data.append(new_row)
                    elif in_table and (not line or not line.startswith('|')):
                        in_table = False
            else:
                # For products, parse specification items
                items = re.split(r';|\n', spec)
                for item in items:
                    if ":" in item:
                        key, value = item.split(':', 1)
                        spec_key = key.strip().capitalize()
                        if spec_key == "Vậtl iệu":
                            spec_key = "Vật liệu"

                        existing_row = None
                        for row in raw_data:
                            if row["Thông số"] == spec_key:
                                existing_row = row
                                break

                        if existing_row:
                            existing_row[full_name] = value.strip(
                            ) if value else ""
                        else:
                            new_row = {"Thông số": spec_key}
                            for col in columns[1:]:
                                new_row[col] = ""
                            new_row[full_name] = value.strip() if value else ""
                            raw_data.append(new_row)

    if raw_data:
        df = pd.DataFrame(raw_data, columns=columns)
        df = df.fillna("").replace("None", "").replace("nan", "")
    else:
        df = pd.DataFrame(
            [["Không có thông số kỹ thuật", "", ""]], columns=columns)

    markdown_table = df.to_markdown(index=False)
    # Fix: Use proper ChatInterface message format
    history.append(
        {"role": "assistant", "content": f"📄 Thông số kỹ thuật\n{markdown_table}"})
    return history


def show_advantages(state, history=None):
    specs_map = state.get("specs_advantages", {})
    columns = ["Tên", "Ưu điểm nổi trội"]
    table_data = []
    if history is None:
        history = []

    if not specs_map:
        df = pd.DataFrame([["Không có ưu điểm", ""]], columns=columns)
        markdown_table = df.to_markdown(index=False)
        # Fix: Use proper ChatInterface message format
        history.append(
            {"role": "assistant", "content": f"💡 Ưu điểm nổi trội\n{markdown_table}"})
        return history

    for prod_id, data in specs_map.items():
        adv = data.get("advantages", "Không có ưu điểm")
        model = data.get("model", "")
        url = data.get("url", "")

        # Handle both products and solution packages
        if url:
            full_name = f"**[{data['name']} {model}]({url})**"
        else:
            full_name = f"**[{data['name']} {model}]**"

        if adv not in ["Không có ưu điểm", "", None]:
            table_data.append([full_name, adv])

    if table_data:
        df = pd.DataFrame(table_data, columns=columns)
    else:
        df = pd.DataFrame([["Không có ưu điểm", ""]], columns=columns)

    markdown_table = df.to_markdown(index=False)
    # Fix: Use proper ChatInterface message format
    history.append(
        {"role": "assistant", "content": f"💡 Ưu điểm nổi trội\n{markdown_table}"})
    return history


def show_solution_packages(state, history=None):
    """Show solution packages in a structured format"""
    specs_map = state.get("specs_advantages", {})
    if history is None:
        history = []

    if not specs_map:
        # Fix: Use proper ChatInterface message format
        history.append(
            {"role": "assistant", "content": "📦 Không có gói sản phẩm nào"})
        return history

    # Filter only solution packages
    solution_packages = {}
    for key, data in specs_map.items():
        if data.get("model") == "Gói giải pháp":
            solution_packages[key] = data

    if not solution_packages:
        # Fix: Use proper ChatInterface message format
        history.append(
            {"role": "assistant", "content": "📦 Không có gói sản phẩm nào"})
        return history

    # Build markdown content for each package
    markdown_content = "## 📦 Gói sản phẩm\n\n"

    for key, data in solution_packages.items():
        spec_content = data.get("specification", "")
        markdown_content += spec_content + "\n\n"

    # Fix: Use proper ChatInterface message format
    history.append({"role": "assistant", "content": markdown_content})
    return history


css = """

    .icon-button-wrapper.top-panel.hide-top-corner,

    .icon-button-wrapper.top-panel button[aria-label="Clear"],

    .icon-button-wrapper.top-panel button[title="Clear"],

    .chat-interface .clear-btn,

    .chat-interface button[aria-label="Clear"],

    .chat-interface button[title="Clear"],

    .chat-interface .svelte-1aq8tno button:last-child,

    .chat-interface .action-buttons button:last-child {

        display: none !important;

        visibility: hidden !important;

    }

    .svelte-vuh1yp .prose.svelte-lag733 .md.svelte-7ddecg.prose h1,

    .svelte-vuh1yp h1,

    .prose.svelte-lag733 h1 {

        visibility: hidden !important;

    }

    .top-right-clear {

        position: absolute;

        top: 10px;

        right: 10px;

        z-index: 1000;

    }

    .custom-button-row {

        margin-top: -10px;

        gap: 8px;

        justify-content: flex-start;

    }

    .custom-clear-btn {

        background: linear-gradient(135deg, #ff6b6b, #ee5a24) !important;

        color: white !important;

        border: none !important;

        border-radius: 6px !important;

        font-weight: 500 !important;

        transition: all 0.3s ease !important;

        min-width: 80px !important;

    }

    .custom-clear-btn:hover {

        background: linear-gradient(135deg, #ee5a24, #ff6b6b) !important;

        transform: translateY(-1px) !important;

        box-shadow: 0 4px 12px rgba(238, 90, 36, 0.3) !important;

    }

    .spec-adv-btn {

        background: linear-gradient(135deg, #4834d4, #686de0) !important;

        color: white !important;

        border: none !important;

        border-radius: 6px !important;

        font-weight: 500 !important;

        transition: all 0.3s ease !important;

    }

    .spec-adv-btn:hover {

        background: linear-gradient(135deg, #686de0, #4834d4) !important;

        transform: translateY(-1px) !important;

        box-shadow: 0 4px 12px rgba(72, 52, 212, 0.3) !important;

    }

    .gradio-container {

        position: relative;

    }

    .model-dropdown, .debug-dropdown {

        background: linear-gradient(135deg, #667eea, #764ba2) !important;

        color: white !important;

        border: none !important;

        border-radius: 6px !important;

        font-weight: 500 !important;

        font-size: 13px !important;

        transition: all 0.3s ease !important;

        min-height: 32px !important;

        box-shadow: 0 2px 4px rgba(0,0,0,0.15) !important;

    }

    .model-dropdown select, .debug-dropdown select {

        color: white !important;

        background: transparent !important;

        cursor: pointer !important;

    }

    .model-dropdown select:focus, .debug-dropdown select:focus {

        outline: none;

    }

    .model-dropdown select:not(:focus), .debug-dropdown select:not(:focus) {

        pointer-events: auto;

    }

"""

with gr.Blocks(fill_height=True, css=css) as demo:
    state = gr.State(value={})
    raw_documents = gr.State(value={})

    with gr.Row(elem_classes="top-right-clear"):
        with gr.Column(scale=20):
            gr.Markdown("### 🛍️ RangDong Sales Agent")

        with gr.Column(scale=1, min_width=180):
            model_dropdown = gr.Dropdown(
                choices=["Gemini 2.0 Flash",
                         "Gemini 2.5 Flash Lite",
                         "Gemini 2.0 Flash Lite"],
                value="Gemini 2.0 Flash",
                label="Model",
                show_label=False,
                elem_classes="model-dropdown",
                interactive=True,
                container=False
            )

        with gr.Column(scale=1, min_width=120):
            debug_dropdown = gr.Dropdown(
                choices=["Normal", "Debug"],
                value="Normal",
                label="Mode",
                show_label=False,
                elem_classes="debug-dropdown",
                interactive=True,
                container=False
            )

        with gr.Column(scale=1, min_width=30):
            clear_btn = gr.Button(
                "✨ Xoá", variant="secondary", size="sm", elem_classes="custom-clear-btn")

    chatbot = gr.Chatbot(
        avatar_images=[
            "https://cdn-icons-png.flaticon.com/512/219/219983.png",
            "assets/agent.png"
        ],
        type="messages",
        height=800
    )

    chat = gr.ChatInterface(
        fn=respond_to_chat,
        multimodal=True,
        type="messages",
        examples=[
            ["Tìm sản phẩm bình giữ nhiệt dung tích dưới 2 lít."],
            ["Tìm sản phẩm ổ cắm thông minh"],
            ["Tư vấn cho tôi đèn học chống cận cho con gái của tôi học lớp 6"]
        ],
        title="_",
        additional_inputs=[state, model_dropdown, debug_dropdown],
        chatbot=chatbot
    )

    with gr.Row(elem_classes="custom-button-row") as button_row:
        btn_spec = gr.Button("📄 Thông số kỹ thuật",
                             variant="secondary", elem_classes="spec-adv-btn")
        btn_adv = gr.Button("💡 Ưu điểm nổi trội",
                            variant="secondary", elem_classes="spec-adv-btn")
        btn_pckg = gr.Button(
            "📦 Gói sản phẩm", variant="secondary", elem_classes="spec-adv-btn")

    # Event handlers
    btn_spec.click(fn=show_specs, inputs=[state, chatbot], outputs=[chatbot])
    btn_adv.click(fn=show_advantages,
                  inputs=[state, chatbot],
                  outputs=[chatbot])

    # Gói sản phẩm button shows solution packages in structured format
    btn_pckg.click(fn=show_solution_packages,
                   inputs=[state, chatbot],
                   outputs=[chatbot])

    # Model change handler -> call API
    def _on_model_change(model_name):
        try:
            requests.post(f"{API_BASE_URL}/set_model",
                          json={"model_name": model_name})
        except Exception as e:
            print(f"Warning: set_model failed: {e}")

    model_dropdown.change(fn=_on_model_change,
                          inputs=[model_dropdown],
                          outputs=[]
                          )  # type: ignore[attr-defined]

    # Clear button - clear all ChatInterface internal components
    clear_btn.click(
        fn=clear_chat,
        inputs=[],
        outputs=[chatbot, state, raw_documents,
                 chat.chatbot_state, chat.textbox]
    )

if __name__ == "__main__":
    demo.launch()