app.py

from asyncio.log import logger
import yaml
from pathlib import Path
import gradio as gr
import asyncio
from langchain_mcp_client import lc_mcp_exec
from dotenv import load_dotenv
import os
import base64
from memory_store import MemoryStore
import logging


# Load environment variables
load_dotenv()
VERSION = "0.0.1"

# ======================================= Load DB configs
def load_db_configs():
    """Load database configurations from configs.yaml"""
    configs_path = Path("configs.yaml")

    if not configs_path.exists():
        raise FileNotFoundError("configs.yaml not found")

    with open(configs_path) as f:
        configs = yaml.safe_load(f)

    return configs["db_configs"]


def image_to_base64_html(image_path, alt_text="Customer Status", width=600):
    with open(image_path, "rb") as f:
        encoded = base64.b64encode(f.read()).decode("utf-8")
        return f'<img src="data:image/png;base64,{encoded}" alt="{alt_text}" style="width:{width}px; max-width:100%; display:block; margin:0 auto 24px auto; border-radius:12px; box-shadow:0 2px 12px rgba(0,0,0,0.08);" />'


# ====================================== Async-compatible wrapper
async def run_agent(request, history=None):
    try:
        # Process request using existing memory
        response, messages = await lc_mcp_exec(request)
        
        # Handle image processing
        image_path = ""
        load_dotenv()
        PANDAS_EXPORTS_PATH = os.environ["PANDAS_EXPORTS_PATH"]
        
        # Ensure the exports directory exists
        os.makedirs(PANDAS_EXPORTS_PATH, exist_ok=True)
        
        # Check for generated charts
        generated_files = [f for f in os.listdir(PANDAS_EXPORTS_PATH) 
                         if f.startswith("temp_chart_") and f.endswith(".png")]
        
        if generated_files:
            image_path = os.path.join(PANDAS_EXPORTS_PATH, generated_files[0])
            try:
                image_markdown = image_to_base64_html(image_path)
                output = f"{image_markdown}\n\n{response}"
                os.remove(image_path)  # Clean up the image file
            except Exception as e:
                logger.error(f"Error processing image: {e}")
                output = response
        else:
            output = response

        return output

    except Exception as e:
        logger.error(f"Error in run_agent: {str(e)}", exc_info=True)
        return f"Error: {str(e)}"


# ====================================== Gradio UI with history
LOGO_PATH = "resources/pialogo.png"

# CSS customizations
custom_css = """
@import url('https://fonts.googleapis.com/css2?family=Inter:wght@400;600&display=swap');
body, .markdown-content, .container, .chat-container, .message-container {
    font-family: 'Inter', Arial, sans-serif !important;
}
.container {
    max-width: 2200px !important;
    margin: auto;
    padding: 20px;
}
.chat-container {
    height: 1000px !important;
    min-height: 1000px !important;
    overflow-y: auto;
}
.message-container {
    padding: 15px;
    border-radius: 10px;
    margin: 10px 0;
}
.markdown-content {
    font-size: 18px;
    line-height: 1.7;
} 

button.svelte-lixn6qd {
    display: none !important;
}
#component-16 {
    display: none !important;
}
""" 

xtheme = gr.themes.Soft() 

with gr.Blocks(css=custom_css, theme=xtheme) as demo:
    with gr.Row(elem_classes="container"):
        # with gr.Column(scale=0.5):
        #     gr.Image(value=LOGO_PATH, height=100, show_label=False, show_download_button=False, show_fullscreen_button=False)
        with gr.Column(scale=5):
            gr.Markdown(
                """
                <h1 style='text-align: center; margin-bottom: 1rem'>Talk to Your Data</h1>
                <p style='text-align: center'>Ask questions about your database, analyze and visualize data.</p>
                <p style='text-align: center; color: #666; font-size: 0.9em; margin-top: -0.5rem'>Version: {}</p>
                """.format(VERSION)
            )
    with gr.Row(elem_classes="container"):
        with gr.Column(scale=3):
            chat = gr.ChatInterface(
                fn=run_agent,
                chatbot=gr.Chatbot(
                    height=1000,
                    show_label=False,
                    elem_classes="chat-container",
                    render_markdown=True,
                    type="messages"
                ),
                textbox=gr.Textbox(
                    placeholder="Type your questions here...",
                    container=False,
                    scale=4
                ),
                examples=[
                    "Describe the database",
                    "List all tables in the database",
                    "Show me the top 10 customers by ticket count and visualize it bar chart with different colors",
                    "Show distribution of tickets by status, use pie chart",
                    "List all tables with columns and data types",
                    "How many comments are there per ticket channel (email, chat, portal)? Also Visualize it as a pie chart",
                    "Delete ticket with id BR-90001",
                    "Show top 5 products with highest number of tickets in a bar chart, add legand, print product names vertically",
                    "How many customers are in each industry?",
                    "List the 5 most active agents by ticket count in 2025.",
                    "How many tickets were reopened at least once?",
                    "In which of the records, have the customer provided positive feedback as a result of the service they got?",
                    "Show the number of comments per ticket, grouped by channel (email, chat, portal), for tickets with comments. Visualize using box plot.",
                    "Display the number of tickets and average reopen count per customer, colored by industry, for customers with resolved tickets. Visualize using Scatter Plot."
                ],
                save_history=True,
                type="messages"
            )
        with gr.Column(scale=1):
            with gr.Accordion("Example Questions", open=True):
                gr.Markdown("""
                - 📋 Describe the database
                - 📊 List all tables in database
                - 👥 Total number of customers
                - 📈 Top 10 customers by ticket count
                - 📋 Ticket count by status and visualize
                - 📆 Average ticket reopened
                - 🧹 Clear memory : `/clear-cache`
                """)


# TODO: maybe we can add a mcp tool to validate the results (those converted to DataFrame) to make sure the valid type is passed to the visualization tool by ReAct agent


if __name__ == "__main__":
    demo.launch(
        server_name="0.0.0.0",
        # server_port=7860,
        #share=True,
        debug=True
    )