app.py

import base64
import os
from pathlib import Path
from typing import Dict, List, Tuple

import gradio as gr
import openai


# --- Configuration -------------------------------------------------------
BASE_URL = "https://bai-ap.jts.co.th:10627/v1"
MODEL_NAME = "jai-mini-v2"
API_KEY = os.getenv("API_KEY", "NO_KEY_NEEDED")
SYSTEM_PROMPT = (
    "You are JAI(Jasmine Artificial Intelligence), built by JTS(Jasmine Technology Solution). Jasmine Technology Solution Public Company Limited (JTS) is a prominent company in Thailand, operating within the telecommunications and IT sectors. They are known for providing integrated ICT solutions, which encompass a wide range of services.\n\nHistorically, JTS has been a leader in designing and installing comprehensive telecommunication and IT systems as a System Integrator. More recently, they have expanded their focus significantly into cutting-edge technologies. This includes:\n\n* **Hyperscale Data Centers:** Providing robust infrastructure for data storage and processing.\n* **Digital Asset Businesses:** Notably, they are building one of the largest Bitcoin mining farms in Thailand, with aspirations to be the largest in Southeast Asia.\n* **Advanced AI Solutions:** JTS is a key player in developing Generative AI platforms and Thai Large Language Models (LLMs), aiming to drive digital transformation with intelligent, secure, and reliable AI. They leverage platforms like JAI Studio and high-performance GPU infrastructure for this development.\n\nJTS aims to be a leading \"Technology Enabler\" in ASEAN, with a mission that emphasizes continuous development of AI solutions, efficient network and cloud infrastructure, and leadership in blockchain technology with a focus on clean energy. They are committed to fostering innovation and contributing to the digital economy for sustainable value."
)

LOGO_PATH = Path(__file__).parent / "logo.png"


def _load_logo_data_uri() -> str:
    try:
        encoded = base64.b64encode(LOGO_PATH.read_bytes()).decode("ascii")
        return f"data:image/png;base64,{encoded}"
    except FileNotFoundError:
        return ""


LOGO_DATA_URI = _load_logo_data_uri()
HERO_LOGO_HTML = (
    f'<div class="hero-logo"><img src="{LOGO_DATA_URI}" alt="JAI Mini logo" /></div>'
    if LOGO_DATA_URI
    else ""
)

QUICK_PROMPTS: List[str] = [
    "Who is JTS?",
    "What is the benefit of small LLM?",
    "คุณเป็นใคร",
    "Ai ทำงานยังไง",
]

HERO_SECTION = f"""
<div class="hero-banner">
  <div class="hero-content">
    {HERO_LOGO_HTML}
    <div class="hero-copy">
      <span class="badge">Introducing JAI Mini</span>
      <p>Partner with JAI to deliver exceptional Thai language experiences, manage long-form context, and power reliable RAG-ready workflows. Purpose-built by the JTS research team.</p>
      <div class="hero-highlights">
        <span>⚡ Optimal for Thai and English</span>
        <span>🧠 Resilient in long-context RAG</span>
        <span>🧩 Compact and efficient</span>
      </div>
    </div>
  </div>
</div>
"""

INSIGHTS_HTML = """
<div class="surface-card">
  <h3>Why JAI?</h3>
  <ul class="feature-grid">
    <li>
      <span class="feature-icon">🧡</span>
      <div>
        <h4>Better Thai and English text comprehension</h4>
      </div>
    </li>
    <li>
      <span class="feature-icon">🪄</span>
      <div>
        <h4>RAG-focused</h4>
        <p>Stronger understanding of long-context inputs.</p>
      </div>
    </li>
    <li>
      <span class="feature-icon">🛡️</span>
      <div>
        <h4>Robust & lightweight</h4>
        <p>Compact 14B model designed for efficiency with a strong performance.</p>
      </div>
    </li>
  </ul>
</div>
<div class="surface-card secondary">
  <h3>Learn more at</h3>
  <ul class="moments-list">
    <li><strong>Blog : </li>
    <li><strong>Huggingface model : </li>
  </ul>
</div>
"""

FOOTER_HTML = """
<div class="footer-note">
  Crafted for Thai innovators by <strong>JTS</strong>. Share your ideas with us at <a href="mailto:jts.ai.team@gmail.com">jts.ai.team@gmail.com</a>.
</div>
"""


CUSTOM_CSS = """
body {
  margin: 0;
  background: radial-gradient(circle at 20% 20%, rgba(255, 173, 96, 0.22), transparent 45%),
              radial-gradient(circle at 85% 10%, rgba(255, 113, 91, 0.2), transparent 40%),
              #1f130b;
  color: #fff3e6;
  font-family: 'Inter', 'Segoe UI', system-ui, -apple-system, sans-serif;
}

.gradio-container {
  max-width: 1200px;
  margin: 0 auto;
  padding: 40px 28px 80px !important;
  background: transparent;
}

.hero-banner {
  position: relative;
  padding: 52px 48px 56px;
  border-radius: 28px;
  background: linear-gradient(135deg, rgba(255, 159, 67, 0.28), rgba(255, 98, 0, 0.32));
  backdrop-filter: blur(18px);
  border: 1px solid rgba(255, 240, 225, 0.24);
  overflow: hidden;
  box-shadow: 0 36px 120px rgba(110, 48, 0, 0.45);
}

.hero-banner::after {
  content: "";
  position: absolute;
  inset: 0;
  background: radial-gradient(circle at 18% 22%, rgba(255, 255, 255, 0.18), transparent 55%);
  pointer-events: none;
}

.hero-content {
  position: relative;
  display: flex;
  align-items: center;
  gap: 32px;
}

.hero-logo {
  flex: 0 0 auto;
  display: flex;
  align-items: center;
  justify-content: center;
  padding: 18px;
  border-radius: 24px;
  background: rgba(255, 246, 236, 0.08);
  border: 1px solid rgba(255, 232, 205, 0.28);
  box-shadow: 0 18px 50px rgba(88, 44, 9, 0.32);
}

.hero-logo img {
  width: 120px;
  height: auto;
  display: block;
}

.hero-copy {
  flex: 1 1 auto;
  display: flex;
  flex-direction: column;
  gap: 16px;
  min-width: 0;
}

.hero-banner h1 {
  font-size: 48px;
  line-height: 1.05;
  margin: 18px 0 18px;
  color: #fff1e3;
}

.hero-banner p {
  max-width: 640px;
  font-size: 19px;
  color: rgba(255, 244, 235, 0.88);
  margin-bottom: 26px;
}

.badge {
  display: inline-flex;
  align-items: center;
  gap: 8px;
  background: rgba(255, 255, 255, 0.16);
  border-radius: 999px;
  padding: 6px 16px;
  font-size: 14px;
  letter-spacing: 0.4px;
  text-transform: uppercase;
  color: #ffe2c3;
}

.hero-highlights {
  display: flex;
  flex-wrap: wrap;
  gap: 12px;
}

.hero-highlights span {
  display: inline-flex;
  align-items: center;
  gap: 10px;
  padding: 10px 18px;
  border-radius: 999px;
  background: rgba(92, 45, 7, 0.6);
  border: 1px solid rgba(255, 205, 169, 0.32);
  font-size: 15px;
  color: #ffddbe;
}

.content-row {
  margin-top: 42px;
  gap: 24px;
}

.chat-column {
  gap: 16px;
}

.chat-panel {
  border-radius: 24px !important;
  background: rgba(43, 21, 7, 0.86) !important;
  border: 1px solid rgba(255, 204, 152, 0.2) !important;
  box-shadow: 0 28px 80px rgba(46, 20, 3, 0.55);
}

.chat-panel .message.bot {
  background: rgba(255, 199, 150, 0.08) !important;
  border-radius: 18px;
}

.chat-panel .message.user {
  background: rgba(255, 142, 78, 0.22) !important;
  border-radius: 18px;
}

.compose-panel {
  padding: 20px;
  border-radius: 20px;
  background: rgba(40, 18, 5, 0.78);
  border: 1px solid rgba(255, 214, 176, 0.18);
  box-shadow: inset 0 0 0 1px rgba(255, 211, 173, 0.08);
  gap: 16px;
}

.compose-panel textarea {
  background: rgba(36, 18, 9, 0.72);
  border: 1px solid rgba(255, 198, 150, 0.18);
  border-radius: 16px;
  color: #ffeede;
  font-size: 16px;
  min-height: 94px !important;
}

.compose-panel textarea:focus {
  border-color: rgba(255, 183, 92, 0.65);
  box-shadow: 0 0 0 3px rgba(255, 162, 53, 0.22);
}

.actions-row {
  display: flex;
  justify-content: space-between;
  gap: 14px;
}

.primary-btn {
  background: linear-gradient(135deg, #ff9933, #ff6a2b) !important;
  border: none !important;
  color: #2a1204 !important;
  font-weight: 600 !important;
  padding: 12px 26px !important;
  border-radius: 14px !important;
}

.primary-btn:hover {
  filter: brightness(1.08);
}

.ghost-btn {
  background: rgba(255, 255, 255, 0.08) !important;
  border: 1px solid rgba(255, 214, 176, 0.24) !important;
  color: #ffe5ce !important;
  font-weight: 500 !important;
  padding: 12px 26px !important;
  border-radius: 14px !important;
}

.quick-prompts {
  margin-top: 42px;
}

.quick-prompts .label {
  font-size: 17px;
  color: rgba(255, 234, 214, 0.82);
  margin-bottom: 14px;
}

.quick-grid {
  display: grid;
  grid-template-columns: repeat(auto-fit, minmax(200px, 1fr));
  gap: 12px;
}

.quick-chip {
  justify-content: flex-start !important;
  border-radius: 14px !important;
  background: rgba(48, 20, 6, 0.75) !important;
  border: 1px solid rgba(255, 197, 150, 0.18) !important;
  color: #ffd9b7 !important;
  font-weight: 500 !important;
  min-height: 54px !important;
}

.quick-chip:hover {
  border-color: rgba(255, 170, 92, 0.55) !important;
  transform: translateY(-1px);
}

.info-column {
  display: flex;
  flex-direction: column;
  gap: 18px;
}

.surface-card {
  padding: 28px;
  border-radius: 24px;
  background: rgba(54, 24, 6, 0.78);
  border: 1px solid rgba(255, 202, 149, 0.22);
  box-shadow: 0 24px 70px rgba(48, 20, 4, 0.5);
}

.surface-card.secondary {
  background: rgba(64, 26, 8, 0.8);
}

.surface-card h3 {
  margin-top: 0;
  font-size: 22px;
  color: #fff1df;
}

.feature-grid {
  list-style: none;
  padding: 0;
  margin: 18px 0 0;
  display: flex;
  flex-direction: column;
  gap: 18px;
}

.feature-grid li {
  display: flex;
  gap: 18px;
  align-items: flex-start;
}

.feature-icon {
  font-size: 26px;
  line-height: 1;
}

.feature-grid h4 {
  margin: 0 0 4px;
  font-size: 18px;
}

.feature-grid p {
  margin: 0;
  color: rgba(255, 232, 210, 0.85);
  line-height: 1.55;
}

.moments-list {
  list-style: none;
  padding: 0;
  margin: 16px 0 0;
  display: grid;
  gap: 12px;
}

.moments-list li {
  color: rgba(255, 222, 191, 0.88);
  line-height: 1.55;
}

.footer-note {
  margin-top: 48px;
  text-align: center;
  color: rgba(255, 229, 199, 0.72);
  font-size: 15px;
}

.footer-note a {
  color: #ffb97a;
  font-weight: 500;
  text-decoration: none;
}

@media (max-width: 992px) {
  .hero-banner {
    padding: 36px 28px;
  }
  .hero-banner h1 {
    font-size: 38px;
  }
  .hero-content {
    flex-direction: column;
    align-items: flex-start;
    gap: 24px;
  }
  .hero-logo {
    padding: 16px;
  }
  .content-row {
    flex-direction: column;
  }
}

@media (max-width: 640px) {
  .hero-banner {
    padding: 28px 24px;
  }
  .hero-content {
    gap: 18px;
  }
  .hero-highlights {
    flex-direction: column;
  }
}
"""


# --- Initialize OpenAI Client -------------------------------------------
try:
    client = openai.OpenAI(
        base_url=BASE_URL,
        api_key=API_KEY,
    )
except Exception as exc:  # pragma: no cover - configuration feedback
    print(f"Error initializing OpenAI client: {exc}")
    client = None


# --- Chat Logic ----------------------------------------------------------
def stream_chat(history: List[Tuple[str, str]], user_message: str):
    history = history or []
    if not user_message or not user_message.strip():
        yield history, history
        return

    working_history = history + [(user_message, "")]
    yield working_history, working_history

    if not client:
        working_history[-1] = (user_message, "❌ Error: AI client not configured.")
        yield working_history, working_history
        return

    messages: List[Dict[str, str]] = [{"role": "system", "content": SYSTEM_PROMPT}]
    for previous_user, previous_bot in history:
        messages.append({"role": "user", "content": previous_user})
        messages.append({"role": "assistant", "content": previous_bot})
    messages.append({"role": "user", "content": user_message})

    assistant_text = ""
    try:
        stream = client.chat.completions.create(
            model=MODEL_NAME,
            messages=messages,
            extra_body={"chat_template_kwargs": {"enable_thinking": False}},
            stream=True,
        )
        for chunk in stream:
            delta = chunk.choices[0].delta.content
            if not delta:
                continue
            assistant_text += delta
            working_history[-1] = (user_message, assistant_text)
            yield working_history, working_history
    except Exception as error:  # pragma: no cover - runtime feedback path
        working_history[-1] = (user_message, f"⚠️ Error: {error}")
        yield working_history, working_history
        return

    if assistant_text:
        working_history[-1] = (user_message, assistant_text)
        yield working_history, working_history


def reset_conversation():
    return [], [], ""


# --- Interface -----------------------------------------------------------
with gr.Blocks(
    theme=gr.themes.Soft(primary_hue="orange", neutral_hue="gray"),
    title="Jai Mini · Modern AI Companion",
    css=CUSTOM_CSS,
) as demo:
    gr.HTML(HERO_SECTION)

    with gr.Row(elem_classes="content-row"):
        with gr.Column(scale=7, elem_classes="chat-column"):
            chatbot = gr.Chatbot(
                label="Conversation history",
                bubble_full_width=False,
                height=520,
                show_copy_button=True,
                elem_classes="chat-panel",
            )

            with gr.Group(elem_classes="compose-panel"):
                prompt_input = gr.Textbox(
                    label="Talk to JAI",
                    placeholder="Type your message here...",
                    lines=4,
                    autofocus=True,
                )
                with gr.Row(elem_classes="actions-row"):
                    submit_button = gr.Button("Send", elem_classes="primary-btn")
                    clear_button = gr.Button(
                        "Reset chat", elem_classes="ghost-btn", variant="secondary"
                    )

        with gr.Column(scale=5, min_width=320, elem_classes="info-column"):
            gr.HTML(INSIGHTS_HTML)

    with gr.Group(elem_classes="quick-prompts"):
        gr.HTML("<div class='label'>Suggest questions:</div>")
        with gr.Row(elem_classes="quick-grid"):
            quick_buttons = []
            for prompt in QUICK_PROMPTS:
                button = gr.Button(prompt, elem_classes="quick-chip", variant="secondary")
                quick_buttons.append((button, prompt))

    gr.HTML(FOOTER_HTML)

    history_state = gr.State([])

    submit_event = submit_button.click(
        fn=stream_chat,
        inputs=[history_state, prompt_input],
        outputs=[chatbot, history_state],
    )
    submit_event.then(fn=lambda: "", inputs=None, outputs=prompt_input, queue=False)

    submit_input_event = prompt_input.submit(
        fn=stream_chat,
        inputs=[history_state, prompt_input],
        outputs=[chatbot, history_state],
    )
    submit_input_event.then(fn=lambda: "", inputs=None, outputs=prompt_input, queue=False)

    clear_button.click(
        fn=reset_conversation,
        inputs=None,
        outputs=[chatbot, history_state, prompt_input],
        queue=False,
    )

    for button, prompt in quick_buttons:
        button.click(fn=lambda p=prompt: p, inputs=None, outputs=prompt_input, queue=False)


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