feat(umapyai): Simple web interface

5 files changed, 198 insertions, 6 deletions

diff --git a/pyproject.toml b/pyproject.toml
index c831093..49fde4f 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -10,6 +10,8 @@ dependencies = [
     "psutil>=7.0.0",
     "loguru>=0.7.3",
     "beautifulsoup4>=4.13.4",
+    "flask>=3.1.1",
+    "flask-cors>=6.0.1",
 ]
 readme = "README.md"
 requires-python = ">= 3.8"
diff --git a/requirements-dev.lock b/requirements-dev.lock
index bc67d2b..c1671dc 100644
--- a/requirements-dev.lock
+++ b/requirements-dev.lock
@@ -24,6 +24,8 @@ bcrypt==4.3.0
     # via chromadb
 beautifulsoup4==4.13.4
     # via umapyai
+blinker==1.9.0
+    # via flask
 build==1.2.2.post1
     # via chromadb
 cachetools==5.5.2
@@ -38,6 +40,7 @@ charset-normalizer==3.4.2
 chromadb==1.0.15
     # via umapyai
 click==8.2.1
+    # via flask
     # via typer
     # via uvicorn
 coloredlogs==15.0.1
@@ -50,6 +53,11 @@ filelock==3.18.0
     # via huggingface-hub
     # via torch
     # via transformers
+flask==3.1.1
+    # via flask-cors
+    # via umapyai
+flask-cors==6.0.1
+    # via umapyai
 flatbuffers==25.2.10
     # via onnxruntime
 fsspec==2025.7.0
@@ -87,7 +95,10 @@ importlib-metadata==8.7.0
     # via opentelemetry-api
 importlib-resources==6.5.2
     # via chromadb
+itsdangerous==2.2.0
+    # via flask
 jinja2==3.1.6
+    # via flask
     # via torch
 joblib==1.5.1
     # via scikit-learn
@@ -102,7 +113,9 @@ loguru==0.7.3
 markdown-it-py==3.0.0
     # via rich
 markupsafe==3.0.2
+    # via flask
     # via jinja2
+    # via werkzeug
 mdurl==0.1.2
     # via markdown-it-py
 mmh3==5.1.0
@@ -282,6 +295,9 @@ websocket-client==1.8.0
     # via kubernetes
 websockets==15.0.1
     # via uvicorn
+werkzeug==3.1.3
+    # via flask
+    # via flask-cors
 yapf==0.43.0
 zipp==3.23.0
     # via importlib-metadata
diff --git a/requirements.lock b/requirements.lock
index eff78fa..ef68198 100644
--- a/requirements.lock
+++ b/requirements.lock
@@ -24,6 +24,8 @@ bcrypt==4.3.0
     # via chromadb
 beautifulsoup4==4.13.4
     # via umapyai
+blinker==1.9.0
+    # via flask
 build==1.2.2.post1
     # via chromadb
 cachetools==5.5.2
@@ -38,6 +40,7 @@ charset-normalizer==3.4.2
 chromadb==1.0.15
     # via umapyai
 click==8.2.1
+    # via flask
     # via typer
     # via uvicorn
 coloredlogs==15.0.1
@@ -50,6 +53,11 @@ filelock==3.18.0
     # via huggingface-hub
     # via torch
     # via transformers
+flask==3.1.1
+    # via flask-cors
+    # via umapyai
+flask-cors==6.0.1
+    # via umapyai
 flatbuffers==25.2.10
     # via onnxruntime
 fsspec==2025.7.0
@@ -87,7 +95,10 @@ importlib-metadata==8.7.0
     # via opentelemetry-api
 importlib-resources==6.5.2
     # via chromadb
+itsdangerous==2.2.0
+    # via flask
 jinja2==3.1.6
+    # via flask
     # via torch
 joblib==1.5.1
     # via scikit-learn
@@ -102,7 +113,9 @@ loguru==0.7.3
 markdown-it-py==3.0.0
     # via rich
 markupsafe==3.0.2
+    # via flask
     # via jinja2
+    # via werkzeug
 mdurl==0.1.2
     # via markdown-it-py
 mmh3==5.1.0
@@ -279,5 +292,8 @@ websocket-client==1.8.0
     # via kubernetes
 websockets==15.0.1
     # via uvicorn
+werkzeug==3.1.3
+    # via flask
+    # via flask-cors
 zipp==3.23.0
     # via importlib-metadata
diff --git a/src/umapyai/__init__.py b/src/umapyai/__init__.py
index feca194..342b41a 100644
--- a/src/umapyai/__init__.py
+++ b/src/umapyai/__init__.py
@@ -4,6 +4,9 @@ import chromadb
 from sentence_transformers import SentenceTransformer
 import requests
 from loguru import logger
+from threading import Thread
+from flask import Flask, request, jsonify, send_file
+from flask_cors import CORS
 from .constants import (ARTICLES_DIRECTORY, CHROMA_DIRECTORY, CHROMA_COLLECTION,
                         CHUNK_SIZE, EMBEDDING_MODEL, OLLAMA_MODEL, TOP_K,
                         OLLAMA_URL)
@@ -17,6 +20,39 @@ logger.add(
     format="<green>{time:YYYY-MM-DD HH:mm:ss}</green> | <level>{level: <8}</level> | <cyan>{function}</cyan> | <level>{message}</level>"
 )
 
+app = Flask(__name__)
+
+CORS(app)
+
+
+def prompt(context, user_query):
+  return ("You are an expert Uma Musume: Pretty Derby build guide advisor.\n"
+          "Answer the user's question using ONLY the following context. "
+          "If the answer isn't in the context, say you don't know.\n\n"
+          f"Context:\n{context}\n\n"
+          f"Question: {user_query}\nAnswer:")
+
+
+def start_flask(find_relevant_chunks, query_ollama):
+
+  @app.route("/")
+  def index():
+    return send_file("chat.html")
+
+  @app.route("/api/ask", methods=["POST"])
+  def api_ask():
+    data = request.get_json()
+    user_query = data.get("question", "")
+    top_chunks = find_relevant_chunks(user_query)
+    context = "\n\n".join([c[0] for c in top_chunks])
+    full_prompt = prompt(context, user_query)
+    answer = query_ollama(full_prompt)
+    sources = ", ".join(sorted(set(meta['source'] for _, meta in top_chunks)))
+
+    return jsonify({"answer": answer, "sources": sources})
+
+  app.run(host="0.0.0.0", port=5000, debug=False, use_reloader=False)
+
 
 def main():
   ollama_process = None
@@ -107,8 +143,15 @@ def main():
       except Exception as error:
         return f"Error communicating with Ollama: {error}"
 
+    flask_thread = Thread(
+        target=start_flask,
+        args=(find_relevant_chunks, query_ollama),
+        daemon=True)
+
+    flask_thread.start()
     logger.success(
         "Ready! Ask your Uma Musume build questions (type 'exit' to quit).")
+    logger.info("Web chat available at http://localhost:5000/")
 
     while True:
       user_query = input("\n> ")
@@ -118,12 +161,7 @@ def main():
 
       top_chunks = find_relevant_chunks(user_query)
       context = "\n\n".join([c[0] for c in top_chunks])
-      full_prompt = (
-          "You are an expert Uma Musume: Pretty Derby build guide advisor.\n"
-          "Answer the user's question using ONLY the following context. "
-          "If the answer isn't in the context, say you don't know.\n\n"
-          f"Context:\n{context}\n\n"
-          f"Question: {user_query}\nAnswer:")
+      full_prompt = prompt(context, user_query)
       answer = query_ollama(full_prompt)
 
       print("\n", answer)
diff --git a/src/umapyai/chat.html b/src/umapyai/chat.html
new file mode 100644
index 0000000..a4cccfb
--- /dev/null
+++ b/src/umapyai/chat.html
@@ -0,0 +1,120 @@
+<!doctype html>
+<html>
+  <head>
+    <title>umapyai</title>
+
+    <style>
+      body {
+        font-family: Arial, sans-serif;
+        max-width: 66vw;
+        margin: 1vh auto;
+        background: #131313;
+        color: #ddd;
+      }
+
+      #chatbox {
+        height: 66vh;
+        overflow-y: scroll;
+        background: #1e1d1e;
+        border: 1px solid #404040;
+        margin-bottom: 1vh;
+        padding: 1vh;
+      }
+
+      .user {
+        font-weight: bold;
+      }
+
+      .ai {
+        white-space: pre-wrap;
+      }
+
+      .source {
+        font-size: 0.9em;
+      }
+
+      #prompt {
+        width: 33vw;
+        font-size: 1.1em;
+        background: #363636;
+        color: #ddd;
+      }
+
+      #send-button {
+        font-size: 1.1em;
+        background: #363636;
+        color: #ddd;
+      }
+    </style>
+  </head>
+
+  <body>
+    <div id="chatbox"></div>
+
+    <input
+      id="prompt"
+      placeholder="Ask your Uma Musume build question ..."
+      autocomplete="off"
+    />
+
+    <button id="send-button">Send</button>
+
+    <script>
+      window.onload = function () {
+        let chatbox = document.getElementById("chatbox");
+        let prompt = document.getElementById("prompt");
+        let sendButton = document.getElementById("send-button");
+        let chat = [];
+
+        const render = () => {
+          chatbox.innerHTML = "";
+
+          for (let message of chat) {
+            if (message.user)
+              chatbox.innerHTML +=
+                "<div class='user'>> " + message.text + "</div>";
+            else
+              chatbox.innerHTML +=
+                "<div class='ai'>" +
+                message.text +
+                "<div class='source'>" +
+                message.sources +
+                "</div></div>";
+          }
+
+          chatbox.scrollTop = chatbox.scrollHeight;
+        };
+
+        sendButton.onclick = async () => {
+          let query = prompt.value.trim();
+
+          if (!query) return;
+
+          chat.push({ user: 1, text: query });
+          render();
+
+          prompt.value = "";
+
+          let response = await fetch("/api/ask", {
+            method: "POST",
+            headers: { "Content-Type": "application/json" },
+            body: JSON.stringify({ question: query }),
+          });
+          let responseData = await response.json();
+
+          chat.push({
+            user: 0,
+            text: responseData.answer,
+            sources: responseData.sources,
+          });
+          render();
+        };
+
+        prompt.addEventListener("keydown", (event) => {
+          if (event.key == "Enter") sendButton.onclick();
+        });
+        render();
+      };
+    </script>
+  </body>
+</html>