Update src/app.py

src/app.py

@@ -207,11 +207,22 @@ def clean_text(text):
     return text.strip()
 
 def ask_ai(user_prompt, system_persona, max_tokens):
+    # 1. Standardize Input: Convert the strings into the Message List format
+    # This ensures compatibility with our new memory-aware backend functions
+    messages_payload = [
+        {"role": "system", "content": system_persona},
+        {"role": "user", "content": user_prompt}
+    ]
+
+    # 2. Routing Logic
     if "GPT-4o" in model_choice:
-        return query_local_model(user_prompt, system_persona, max_tokens)
+        # CORRECTED: Now calls the OpenAI function
+        return query_openai_model(messages_payload, max_tokens)
     else:
+        # Lookup the technical name for Ollama
         technical_name = model_map[model_choice]
-        return query_local_model(user_prompt, system_persona, max_tokens, technical_name)
+        # Calls the Local function
+        return query_local_model(messages_payload, max_tokens, technical_name)
 
 # --- MAIN UI ---
 st.title("AI Toolkit")
@@ -272,7 +283,13 @@ with tab1:
                 st.session_state.email_draft = reply
                 
                 if usage:
-                    m_name = "Granite" if "Granite" in model_choice else "GPT-4o"
+                    # 1. Determine a clean name for the log
+                    if "GPT-4o" in model_choice:
+                         m_name = "GPT-4o"
+                    else:
+                         # Use the first word of the model choice (e.g., "Llama", "Gemma", "Granite")
+                         m_name = model_choice.split(" ")[0]
+                    # 2. Log it
                     tracker.log_usage(m_name, usage["input"], usage["output"])
                     update_sidebar_metrics() # Force update