feat(umapyai): Improve tool handling and web UI integration

1 files changed, 100 insertions, 64 deletions

diff --git a/src/umapyai/__init__.py b/src/umapyai/__init__.py
index 8618862..93ec61d 100644
--- a/src/umapyai/__init__.py
+++ b/src/umapyai/__init__.py
@@ -14,7 +14,8 @@ from .constants import (ARTICLES_DIRECTORY, CHROMA_DIRECTORY, CHROMA_COLLECTION,
 from .ollama_server import start_ollama_server, is_ollama_live, ensure_model_pulled, kill_ollama
 from collections import defaultdict
 from .language import clean_for_match, get_query_phrases
-from .tools import tools, call_tool
+from .tools import tools as UMA_TOOLS, call_tool
+import copy
 
 logger.remove()
 logger.add(
@@ -30,23 +31,110 @@ socket = Sock(app)
 CORS(app)
 
 
+def query_ollama(prompt, ollama_history):
+  messages: list[dict] = []
+
+  if ollama_history:
+    messages.extend(copy.deepcopy(ollama_history))
+
+  messages.append({"role": "user", "content": prompt})
+
+  while True:
+    answer_buffer = ""
+    tool_called = False
+
+    try:
+      for chunk in ollama.chat(
+          model=OLLAMA_MODEL,
+          messages=messages,
+          tools=UMA_TOOLS,
+          stream=True,
+      ):
+        if chunk.get("done"):
+          messages.append({"role": "assistant", "content": answer_buffer})
+
+          yield {"type": "history", "data": messages}
+
+          break
+
+        message = chunk.get("message", {})
+
+        if not message:
+          continue
+
+        if "tool_calls" in message:
+          tool_call_dicts = []
+
+          for i, raw_call in enumerate(message["tool_calls"]):
+            if isinstance(raw_call, dict):
+              call_as_dict = raw_call
+            else:
+              call_as_dict = (
+                  raw_call.dict() if hasattr(raw_call, "dict") else json.loads(
+                      json.dumps(raw_call, default=str)))
+
+            tool_call_dicts.append(call_as_dict)
+
+          messages.append({
+              "role": "assistant",
+              "content": None,
+              "tool_calls": tool_call_dicts,
+          })
+
+          for i, tool_call in enumerate(tool_call_dicts):
+            tool_call_id = tool_call.get("id", f"tool_call_{i}")
+
+            try:
+              tool_result = call_tool(tool_call)
+            except Exception as error:
+              tool_result = {"error": str(error)}
+
+            messages.append({
+                "role": "tool",
+                "tool_call_id": tool_call_id,
+                "content": json.dumps(tool_result),
+            })
+
+          tool_called = True
+
+          break
+
+        if "content" in message and message["content"]:
+          text_chunk = message["content"]
+          answer_buffer += text_chunk
+
+          yield {"type": "answer_chunk", "data": text_chunk}
+
+    except Exception as error:
+      yield {
+          "type": "error",
+          "data": f"Error communicating with Ollama: {error}"
+      }
+
+      break
+
+    if tool_called:
+      continue
+    else:
+      break
+
+
 def prompt(rag_context, user_query, chat_history):
   system_prompt = (
-      'You are a friendly and expert "Uma Musume: Pretty Derby" build guide advisor. '
-      'Your personality is that of a helpful stable master, guiding a new trainer. You are not anyone in particular. '
-      'Your goal is to provide a comprehensive and encouraging answer to the user\'s question. '
-      'You are a fan-made guide and are not affiliated with Cygames or any other entity.\n\n'
-      'IMPORTANT: If the user asks a question that could be answered by a tool (API), ALWAYS call the tool, even if you believe you know the answer from context or memory. '
-      'NEVER rely solely on the provided documentation or context for character data, builds, live news, or support card info—always call the relevant tool endpoint if one exists.'
+      'You are a friendly and expert "Uma Musume: Pretty Derby" trainer and build-guide advisor. '
+      'Speak in an encouraging, knowledgeable tone—like a helpful stable master coaching a newer trainer. You are not anyone in particular. '
+      'You are fan-made (not affiliated with Cygames or any other entity) and you already *know* the game\'s characters, stats, and meta. '
+      'When absolutely necessary—e.g., for time-sensitive data such as current banners, news, or a character\'s most recent balance patch—you may *silently* use the available API tools. '
+      'Never mention or apologise for using tools; simply answer as if you had the knowledge at hand.'
   )
   instruction = (
       "Carefully analyse the user's question, the provided context, and the chat history. "
       "Synthesise the information from the context and the chat history to form a coherent answer. "
       "Connect different pieces of information and draw logical conclusions. "
-      "If the context does not contain enough information to give a complete answer, "
-      "say so, but still provide any relevant information you did find.")
-  history_string = "\n".join(
-      [f"{message['role']}: {message['content']}" for message in chat_history])
+      "If something is uncertain, briefly explain what is known and, if helpful, suggest what additional info would clarify it."
+  )
+  history_string = "\n".join(f"{msg['role'].capitalize()}: {msg['content']}"
+                             for msg in chat_history[-6:])
 
   return (
       f"{system_prompt}\n\n"
@@ -81,7 +169,7 @@ def start_flask(find_relevant_chunks, query_ollama):
       webSocket.send(json.dumps({"type": "sources", "data": sources}))
 
       for ollama_response in query_ollama(full_prompt, ollama_history):
-        webSocket.send(json.dumps(ollama_response))
+        webSocket.send(json.dumps(ollama_response, default=str))
 
   app.run(host="0.0.0.0", port=5000, debug=False, use_reloader=False)
 
@@ -186,58 +274,6 @@ def main():
 
       return merged[:top_k]
 
-    def query_ollama(prompt, context=None):
-      try:
-        messages = [{"role": "user", "content": prompt}]
-
-        if context:
-          messages = context + messages
-
-        while True:
-          tool_called = False
-
-          for chunk in ollama.chat(
-              model=OLLAMA_MODEL, messages=messages, stream=True, tools=tools):
-            message = chunk.get("message", {})
-
-            if "tool_calls" in message:
-              for tool_call in message["tool_calls"]:
-                tool_result = call_tool(tool_call)
-
-                messages.append({
-                    "role": "assistant",
-                    "tool_call_id": tool_call.get("id", "tool_call_1"),
-                    "name": tool_call.get("name", "unknown_tool"),
-                    "content": None,
-                    "tool_calls": [tool_call],
-                })
-                messages.append({
-                    "role": "tool",
-                    "tool_call_id": tool_call.get("id", "tool_call_1"),
-                    "content": str(tool_result),
-                })
-
-                tool_called = True
-
-              break
-            elif "content" in message and message["content"]:
-              yield {"type": "answer_chunk", "data": message["content"]}
-
-            if chunk.get("done"):
-              yield {
-                  "type": "history",
-                  "data": messages + ([message] if message else [])
-              }
-
-          if not tool_called:
-            break
-      except Exception as error:
-        error_message = f"Error communicating with Ollama: {error}"
-
-        logger.error(error_message)
-
-        yield {"type": "error", "data": error_message}
-
     flask_thread = Thread(
         target=start_flask,
         args=(find_relevant_chunks, query_ollama),