feat(umapyai): Improve prompt and context use

2 files changed, 65 insertions, 17 deletions

diff --git a/src/umapyai/__init__.py b/src/umapyai/__init__.py
index 342b41a..ce2099f 100644
--- a/src/umapyai/__init__.py
+++ b/src/umapyai/__init__.py
@@ -25,12 +25,35 @@ 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 prompt(rag_context, user_query, is_first_turn=True):
+  if is_first_turn:
+    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. '
+        'Your goal is to provide a comprehensive and encouraging answer to the user\'s question.'
+    )
+    instruction = (
+        "Carefully analyse the user's question and the provided context. "
+        "Synthesise the information from the context to form a coherent answer. "
+        "Connect different pieces of information and draw logical conclusions from the text. "
+        "If the context does not contain enough information to give a complete answer, "
+        "say so, but still provide any relevant information you did find.")
+
+    return (
+        f"{system_prompt}\n\n"
+        f"## Instruction\n{instruction}\n\n"
+        "## Context Provided\n"
+        f"--- START OF CONTEXT ---\n{rag_context}\n--- END OF CONTEXT---\n\n"
+        f"## User's Question\n{user_query}\n\n"
+        "## Your Answer")
+  else:
+    return (
+        "Here is some new information that might be relevant to your follow-up question. "
+        "Please synthesise it with our ongoing conversation to provide your answer.\n\n"
+        "## Additional Context\n"
+        f"--- START OF CONTEXT ---\n{rag_context}\n--- END OF CONTEXT---\n\n"
+        f"## User's Question\n{user_query}\n\n"
+        "## Your Answer")
 
 
 def start_flask(find_relevant_chunks, query_ollama):
@@ -43,13 +66,19 @@ def start_flask(find_relevant_chunks, query_ollama):
   def api_ask():
     data = request.get_json()
     user_query = data.get("question", "")
+    history_context = data.get("history", None)
     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)
+    rag_context = "\n\n".join([c[0] for c in top_chunks])
+    full_prompt = prompt(
+        rag_context, user_query, is_first_turn=(history_context is None))
+    answer, new_history_context = query_ollama(full_prompt, history_context)
     sources = ", ".join(sorted(set(meta['source'] for _, meta in top_chunks)))
 
-    return jsonify({"answer": answer, "sources": sources})
+    return jsonify({
+        "answer": answer,
+        "sources": sources,
+        "history": new_history_context
+    })
 
   app.run(host="0.0.0.0", port=5000, debug=False, use_reloader=False)
 
@@ -126,7 +155,7 @@ def main():
       return [(document, metadata)
               for document, metadata in zip(documents, metadatas)]
 
-    def query_ollama(prompt):
+    def query_ollama(prompt, context=None):
       url = f"{OLLAMA_URL}/api/generate"
       payload = {
           "model": OLLAMA_MODEL,
@@ -134,14 +163,25 @@ def main():
           "stream": False,
       }
 
+      if context:
+        payload["context"] = context
+
       try:
         response = requests.post(url, json=payload)
 
         response.raise_for_status()
 
-        return response.json().get('response', '').strip()
+        json_response = response.json()
+        answer = json_response.get('response', '').strip()
+        new_context = json_response.get('context')
+
+        return answer, new_context
       except Exception as error:
-        return f"Error communicating with Ollama: {error}"
+        error_message = f"Error communicating with Ollama: {error}"
+
+        logger.error(error_message)
+
+        return error_message, None
 
     flask_thread = Thread(
         target=start_flask,
@@ -153,6 +193,8 @@ def main():
         "Ready! Ask your Uma Musume build questions (type 'exit' to quit).")
     logger.info("Web chat available at http://localhost:5000/")
 
+    cli_history_context = None
+
     while True:
       user_query = input("\n> ")
 
@@ -160,9 +202,12 @@ def main():
         break
 
       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)
+      rag_context = "\n\n".join([c[0] for c in top_chunks])
+      full_prompt = prompt(
+          rag_context, user_query, is_first_turn=(cli_history_context is None))
+      answer, new_cli_history_context = query_ollama(full_prompt,
+                                                     cli_history_context)
+      cli_history_context = new_cli_history_context
 
       print("\n", answer)
       print(
diff --git a/src/umapyai/chat.html b/src/umapyai/chat.html
index db9857b..d5f50c2 100644
--- a/src/umapyai/chat.html
+++ b/src/umapyai/chat.html
@@ -184,6 +184,7 @@
         let prompt = document.getElementById("prompt");
         let sendButton = document.getElementById("send-button");
         let chat = [];
+        let historyContext = null;
 
         const render = () => {
           chatbox.innerHTML = "";
@@ -217,10 +218,12 @@
           let response = await fetch("/api/ask", {
             method: "POST",
             headers: { "Content-Type": "application/json" },
-            body: JSON.stringify({ question: query }),
+            body: JSON.stringify({ question: query, history: historyContext }),
           });
           let responseData = await response.json();
 
+          historyContext = responseData.history;
+
           chat.push({
             user: 0,
             text: responseData.answer,