app.py

from __future__ import annotations

import html
import inspect
import os
from functools import lru_cache

import gradio as gr
from langchain_core.messages import AIMessage, BaseMessage, HumanMessage

from memory_agent.agent import UniversalMemoryAgent
from memory_agent.config import AppConfig
from memory_agent.errors import RATE_LIMIT_MESSAGE, is_rate_limit_error

NAMESPACE = "default_user"

HERO_HTML = """
<section class="hero">
  <h1>Universal Memory Agent</h1>
  <p>Learns from interactions, remembers your domain, and answers from stored knowledge.</p>
</section>
"""

CSS = """
:root {
    --bg: #04070d;
    --panel: #0c1324;
    --panel-2: #121b31;
    --text: #e9f2ff;
    --muted: #9aa9c3;
    --yellow: #ffd84d;
    --blue: #2d7cff;
    --blue-soft: #4bc4ff;
}

body, .gradio-container {
    color: var(--text);
    background:
        radial-gradient(1200px 720px at 8% -12%, rgba(36, 128, 255, 0.42), transparent 56%),
        radial-gradient(980px 620px at 96% 3%, rgba(185, 230, 255, 0.24), transparent 58%),
        radial-gradient(920px 580px at 98% 2%, rgba(255, 216, 77, 0.16), transparent 62%),
        linear-gradient(180deg, #050911 0%, var(--bg) 100%) !important;
}

.app-shell {
    max-width: 920px;
    margin: 0 auto;
    padding: 1.6rem 0.4rem 1.2rem 0.4rem;
}

.hero {
    position: relative;
    overflow: hidden;
    border-radius: 22px;
    border: 1px solid rgba(75, 196, 255, 0.26);
    background: linear-gradient(140deg, rgba(12, 19, 36, 0.95), rgba(6, 10, 19, 0.98));
    padding: 1.25rem 1.25rem 1rem 1.25rem;
    box-shadow: 0 18px 45px rgba(0, 0, 0, 0.42);
    margin-bottom: 0.9rem;
}

.hero::after {
    content: "";
    position: absolute;
    inset: 0;
    background:
        radial-gradient(560px 200px at -5% 0%, rgba(255, 216, 77, 0.22), transparent 72%),
        radial-gradient(520px 240px at 105% 0%, rgba(42, 142, 255, 0.5), transparent 76%),
        radial-gradient(430px 190px at 82% 18%, rgba(215, 244, 255, 0.2), transparent 74%);
    pointer-events: none;
}

.hero h1 {
    margin: 0;
    color: #f5fbff;
    font-size: clamp(2rem, 4vw, 3.05rem);
    letter-spacing: 0.25px;
    line-height: 1.1;
}

.hero p {
    margin: 0.52rem 0 0 0;
    color: var(--muted);
    font-size: 1.03rem;
}

#clear-btn {
    border-radius: 999px;
    border: 1px solid rgba(255, 216, 77, 0.55);
    background: linear-gradient(110deg, rgba(255, 216, 77, 0.14), rgba(45, 124, 255, 0.17));
    color: var(--text);
    font-weight: 650;
}

#chatbot {
    min-height: 460px;
    border: 1px solid rgba(45, 124, 255, 0.24);
    border-radius: 18px;
    background: linear-gradient(155deg, rgba(15, 24, 45, 0.82), rgba(9, 14, 26, 0.9));
    box-shadow: 0 10px 25px rgba(0, 0, 0, 0.34);
}

#chatbot .message,
#chatbot .message-row .message {
    border-radius: 16px !important;
    border: 1px solid rgba(45, 124, 255, 0.22) !important;
    background: linear-gradient(155deg, rgba(15, 24, 45, 0.92), rgba(9, 14, 26, 0.92)) !important;
    color: #eaf3ff !important;
}

#composer {
    margin-top: 0.7rem;
}

#user-input {
    border-radius: 14px;
    border: 1px solid rgba(255, 216, 77, 0.46);
    background: rgba(6, 11, 21, 0.94);
}

#user-input textarea {
    color: #f3f8ff !important;
    font-size: 1.02rem !important;
}

#user-input textarea::placeholder {
    color: #9fb1cf !important;
}

#send-btn {
    border-radius: 12px;
    border: 1px solid rgba(255, 216, 77, 0.55);
    background: linear-gradient(110deg, rgba(255, 216, 77, 0.14), rgba(45, 124, 255, 0.17));
    color: var(--text);
    font-weight: 700;
    min-width: 56px;
}

#error-box {
    margin-top: 0.9rem;
    border-radius: 12px;
    border: 1px solid rgba(255, 119, 119, 0.35);
    background: rgba(80, 25, 25, 0.35);
    color: #ffdede;
    padding: 0.8rem 0.9rem;
    font-size: 0.96rem;
}

@media (max-width: 768px) {
    .app-shell {
        padding: 1.1rem 0.2rem 1rem 0.2rem;
    }
    .hero {
        padding: 1rem;
    }
}
"""


@lru_cache(maxsize=1)
def build_agent() -> UniversalMemoryAgent:
    config = AppConfig.from_env()
    return UniversalMemoryAgent(config=config)


def _history_to_langchain(history: list[dict[str, str]]) -> list[BaseMessage]:
    messages: list[BaseMessage] = []
    for item in history:
        role = item.get("role", "")
        content = item.get("content", "")
        if role == "user":
            messages.append(HumanMessage(content=content))
        elif role == "assistant":
            messages.append(AIMessage(content=content))
    return messages


def _escape(text: str) -> str:
    return html.escape(text, quote=False)


def _build_error_message(error: Exception) -> str:
    return (
        "Failed to initialize agent. Configure API_KEY (or MISTRAL_API_KEY).\n"
        f"Details: {_escape(str(error))}"
    )


def handle_chat(user_input: str, history: list[dict[str, str]] | None) -> tuple[str, list[dict[str, str]], list[dict[str, str]]]:
    chat_history = list(history or [])
    text = (user_input or "").strip()
    if not text:
        return "", chat_history, chat_history

    try:
        agent = build_agent()
        prior_messages = _history_to_langchain(chat_history)
        assistant_text = agent.run(
            user_input=text,
            chat_history=prior_messages,
            namespace=NAMESPACE,
        )
    except Exception as error:
        if is_rate_limit_error(error):
            assistant_text = RATE_LIMIT_MESSAGE
        else:
            assistant_text = f"Temporary error: {error}"

    chat_history.append({"role": "user", "content": text})
    chat_history.append({"role": "assistant", "content": assistant_text})
    return "", chat_history, chat_history


def clear_chat() -> tuple[list[dict[str, str]], list[dict[str, str]]]:
    return [], []


def build_ui() -> gr.Blocks:
    error_message = ""
    try:
        build_agent()
    except Exception as error:
        error_message = _build_error_message(error)

    with gr.Blocks(title="Universal Memory Agent") as demo:
        with gr.Column(elem_classes=["app-shell"]):
            gr.HTML(HERO_HTML)
            clear_button = gr.Button("Clear chat history", elem_id="clear-btn")
            chatbot = gr.Chatbot(
                [],
                show_label=False,
                elem_id="chatbot",
                height=500,
            )
            with gr.Row(elem_id="composer"):
                user_input = gr.Textbox(
                    placeholder="Share a fact or ask a question...",
                    show_label=False,
                    container=True,
                    lines=1,
                    max_lines=1,
                    elem_id="user-input",
                    scale=20,
                    autofocus=True,
                )
                send_button = gr.Button("↑", elem_id="send-btn", scale=1, min_width=56)

            if error_message:
                gr.HTML(f"<div id='error-box'>{error_message}</div>")

            session_history = gr.State([])

            user_input.submit(
                fn=handle_chat,
                inputs=[user_input, session_history],
                outputs=[user_input, chatbot, session_history],
            )
            send_button.click(
                fn=handle_chat,
                inputs=[user_input, session_history],
                outputs=[user_input, chatbot, session_history],
            )
            clear_button.click(
                fn=clear_chat,
                inputs=None,
                outputs=[chatbot, session_history],
            )

    return demo


def _resolve_host() -> str:
    return os.getenv("APP_HOST", "0.0.0.0")


def _resolve_port() -> int | None:
    # Hugging Face Spaces sets PORT. Local/dev can use APP_PORT.
    env_port = os.getenv("PORT") or os.getenv("APP_PORT")
    if not env_port:
        return None
    try:
        return int(env_port)
    except ValueError:
        return None


demo = build_ui()


if __name__ == "__main__":
    launch_kwargs: dict[str, object] = {
        "server_name": _resolve_host(),
        "css": CSS,
    }
    resolved_port = _resolve_port()
    if resolved_port is not None:
        launch_kwargs["server_port"] = resolved_port
    if "show_api" in inspect.signature(demo.launch).parameters:
        launch_kwargs["show_api"] = False
    # HF Spaces may run Python 3.13 where Gradio SSR can trigger asyncio loop cleanup warnings.
    if "ssr_mode" in inspect.signature(demo.launch).parameters:
        launch_kwargs["ssr_mode"] = False
    try:
        demo.launch(**launch_kwargs)
    except OSError as error:
        # If requested port is busy, retry with Gradio automatic port selection.
        if "Cannot find empty port" not in str(error) or "server_port" not in launch_kwargs:
            raise
        launch_kwargs.pop("server_port", None)
        demo.launch(**launch_kwargs)