app.py

import logging
import gradio as gr
import configs.config as config
import services.scraper
import stores.chroma
from llm_setup.llm_setup import LLMService
from caching.lfu import LFUCache
from configs.config import MODEL_REGISTRY, DEFAULT_PROVIDER
import time 

logger = logging.getLogger()  # Create a logger object
logger.setLevel(logging.INFO)  # Set the logging level to INFO

config.set_envs()  # Set environment variables using the config module

store = stores.chroma.ChromaDB(config.EMBEDDINGS)
service = services.scraper.Service(store)

# Scrape data and get the store vector retriever
service.scrape_and_get_store_vector_retriever(config.URLS)

# Initialize the LLMService with logger, prompt, and store vector retriever
llm_svc = LLMService(
    logger=logger,
    system_prompt=config.SYSTEM_PROMPT,
    web_retriever=store.get_chroma_instance().as_retriever(),
    provider=config.DEFAULT_PROVIDER,
    llm_model_name=config.LLM_MODEL_NAME,
)

def respond(user_input,session_hash):
    if user_input == "clear_chat_history_aisdb_override":
        llm_svc.store={}
        return "Memory Cache cleared"
    response = llm_svc.conversational_rag_chain().invoke(
        {"input": user_input},
        config={"configurable": {"session_id": session_hash}},
    )["answer"]

    return response

def echo(text, chat_history, request: gr.Request):
    if request:
        session_hash = request.session_hash
        resp = respond(text, session_hash)
        for i in range(len(resp)):
            time.sleep(0.01)
            yield resp[: i + 1]
    else:
        return "No request object received."


def on_reset_button_click():
    llm_svc.store=LFUCache(capacity=50)

def on_apply_model(provider, model_name, api_key):
    key = api_key.strip() if api_key and api_key.strip() else None
    try:
        llm_svc.update_llm(provider, model_name, key)
        return f"Switched to {provider} / {model_name}"
    except Exception as e:
        return f"Error: {str(e)}"

def on_provider_change(provider):
    models = MODEL_REGISTRY.get(provider, [])
    return gr.update(choices=models, value=models[0] if models else None)

# --- Maritime Theme ---
maritime_blue = gr.themes.Color(
    c50="#f0f9ff", c100="#e0f2fe", c200="#b9e6fe", c300="#7dd4fc",
    c400="#38bdf8", c500="#0ea5e9", c600="#0284c7", c700="#0369a1",
    c800="#075985", c900="#0c4a6e", c950="#082f49",
    name="maritime-blue",
)

teal_accent = gr.themes.Color(
    c50="#f0fdfa", c100="#ccfbf1", c200="#99f6e4", c300="#5eead4",
    c400="#2dd4bf", c500="#14b8a6", c600="#0d9488", c700="#0f766e",
    c800="#115e59", c900="#134e4a", c950="#042f2e",
    name="teal-accent",
)

try:
    stormy_theme = gr.themes.Ocean(
        primary_hue=maritime_blue,
        secondary_hue=teal_accent,
        neutral_hue="slate",
        spacing_size=gr.themes.sizes.spacing_md,
        radius_size=gr.themes.sizes.radius_lg,
        text_size=gr.themes.sizes.text_md,
        font=(gr.themes.GoogleFont("Inter"), "ui-sans-serif", "system-ui", "sans-serif"),
        font_mono=(gr.themes.GoogleFont("JetBrains Mono"), "ui-monospace", "Consolas", "monospace"),
    )
except AttributeError:
    stormy_theme = gr.themes.Soft(
        primary_hue=maritime_blue,
        secondary_hue=teal_accent,
        neutral_hue="slate",
        spacing_size=gr.themes.sizes.spacing_md,
        radius_size=gr.themes.sizes.radius_lg,
        text_size=gr.themes.sizes.text_md,
        font=(gr.themes.GoogleFont("Inter"), "ui-sans-serif", "system-ui", "sans-serif"),
        font_mono=(gr.themes.GoogleFont("JetBrains Mono"), "ui-monospace", "Consolas", "monospace"),
    )

stormy_theme = stormy_theme.set(
    body_background_fill="#f0f9ff",
    body_background_fill_dark="#0c1929",
    body_text_color="#0c4a6e",
    body_text_color_dark="#e0f2fe",
    block_background_fill="#ffffff",
    block_background_fill_dark="#0f2942",
    block_border_color="#b9e6fe",
    block_border_color_dark="#0369a1",
    button_primary_background_fill="linear-gradient(135deg, #0ea5e9, #0d9488)",
    button_primary_background_fill_hover="linear-gradient(135deg, #38bdf8, #14b8a6)",
    button_primary_background_fill_dark="linear-gradient(135deg, #0369a1, #0f766e)",
    button_primary_text_color="#ffffff",
    button_secondary_background_fill="#e0f2fe",
    button_secondary_background_fill_hover="#b9e6fe",
    button_secondary_background_fill_dark="#0f2942",
    button_secondary_text_color="#0c4a6e",
    button_secondary_text_color_dark="#7dd4fc",
    input_background_fill="#f8fafc",
    input_background_fill_dark="#0f2942",
    input_border_color="#b9e6fe",
    input_border_color_focus="#0ea5e9",
    input_border_color_dark="#0369a1",
    shadow_drop="0 2px 8px rgba(14, 165, 233, 0.08)",
    shadow_drop_lg="0 4px 16px rgba(14, 165, 233, 0.12)",
)

custom_css = """
.stormy-header {
    text-align: center;
    padding: 1.5rem 1rem 1rem 1rem;
    background: linear-gradient(135deg, #0c4a6e 0%, #0ea5e9 50%, #0d9488 100%);
    border-radius: 12px;
    margin-bottom: 0.5rem;
    color: white;
}
.stormy-header h1 {
    font-size: 1.8rem;
    margin: 0 0 0.25rem 0;
    font-weight: 700;
    color: #ffffff !important;
}
.stormy-header p {
    font-size: 0.95rem;
    margin: 0;
    color: #e0f2fe !important;
    opacity: 0.9;
}
.reset-btn {
    max-width: 200px !important;
}
.stormy-footer {
    text-align: center;
    font-size: 0.8rem;
    color: #64748b;
    padding-top: 0.5rem;
}
"""

if __name__ == '__main__':
    logging.info("Starting AIVIz Bot")

    with gr.Blocks(theme=stormy_theme, css=custom_css, title="Stormy - AISdb Assistant") as demo:

        # Branding Header
        gr.Markdown(
            """
            <div class="stormy-header">
                <h1>Stormy - AISdb Assistant</h1>
                <p>Your maritime data companion. Ask about AIS vessel tracking, data processing, machine learning, and more.</p>
            </div>
            """,
            elem_id="header",
        )

        # Chat Interface
        chatbot = gr.Chatbot(
            placeholder=(
                "<strong>Welcome aboard!</strong><br>"
                "I'm Stormy, your AISdb documentation assistant.<br>"
                "Ask me about vessel tracking, data queries, or machine learning with AIS data."
            ),
            height=500,
            type="messages",
            show_copy_button=True,
        )

        gr.ChatInterface(
            fn=echo,
            type="messages",
            chatbot=chatbot,
            textbox=gr.Textbox(
                placeholder="Ask Stormy about AISdb...",
                container=False,
                scale=7,
            ),
            examples=[
                "How do I get started with AISdb?",
                "How can I query vessel tracks by MMSI?",
                "What machine learning models work with AIS data?",
                "How do I visualize ship trajectories on a map?",
            ],
        )

        # Action Bar
        with gr.Row():
            with gr.Column(scale=3):
                with gr.Accordion("About Stormy & AISdb", open=False):
                    gr.Markdown(
                        """
                        **Stormy** is an AI assistant built on the AISdb (Automatic Identification System Database)
                        documentation. It can help you with:

                        - **Data Access**: Loading AIS data, creating databases, CSV export
                        - **Querying**: SQL queries, filtering by MMSI, time ranges, geographic areas
                        - **Processing**: Data cleaning, track interpolation, decimation
                        - **Visualization**: Plotting vessel trajectories, hexagon discretization
                        - **Machine Learning**: Seq2Seq models, autoencoders for AIS data
                        - **Geospatial**: Haversine distance, shore distance, bathymetric data

                        Powered by AISdb documentation from [aisviz.gitbook.io](https://aisviz.gitbook.io/documentation)
                        and [MAPS Lab](https://mapslab.tech/).
                        """
                    )
            with gr.Column(scale=1, min_width=200):
                reset_button = gr.Button(
                    "Reset Chat Memory",
                    variant="secondary",
                    size="sm",
                    elem_classes=["reset-btn"],
                )
                reset_button.click(on_reset_button_click)

        # Model Settings Panel
        with gr.Accordion("Model Settings", open=False):
            with gr.Row():
                provider_dropdown = gr.Dropdown(
                    choices=list(MODEL_REGISTRY.keys()),
                    value=DEFAULT_PROVIDER,
                    label="Provider",
                    interactive=True,
                    scale=1,
                )
                model_dropdown = gr.Dropdown(
                    choices=MODEL_REGISTRY[DEFAULT_PROVIDER],
                    value=config.LLM_MODEL_NAME,
                    label="Model",
                    interactive=True,
                    scale=1,
                )
            with gr.Row():
                api_key_input = gr.Textbox(
                    label="API Key (optional override)",
                    placeholder="Leave blank to use environment variable",
                    type="password",
                    scale=3,
                )
                apply_button = gr.Button(
                    "Apply",
                    variant="primary",
                    size="sm",
                    scale=1,
                )
            status_text = gr.Textbox(
                label="Status",
                interactive=False,
                value=f"Active: {DEFAULT_PROVIDER} / {config.LLM_MODEL_NAME}",
                max_lines=1,
            )

            provider_dropdown.change(
                fn=on_provider_change,
                inputs=[provider_dropdown],
                outputs=[model_dropdown],
            )
            apply_button.click(
                fn=on_apply_model,
                inputs=[provider_dropdown, model_dropdown, api_key_input],
                outputs=[status_text],
            )

        # Footer
        gr.Markdown(
            '<div class="stormy-footer">Built with Gradio & LangChain | AISdb Documentation Assistant</div>'
        )

    demo.launch()