Update router.py

router.py

@@ -1,41 +1,128 @@
-from llm_clients import call_llama, call_gemini
+from llm_clients import (
+    call_llama,
+    call_gemini,
+    classify_prompt,
+    judge_answers
+)
+
+from memory import save_message, load_memory
 from search_tool import search_web
-from rag_engine import rag_response
+import requests
+
+
+# Optional: image microservice URL
+IMAGE_SPACE_URL = "https://your-image-space.hf.space/generate"
+
+
+def build_messages(system_prompt, memory, user_prompt):
+    messages = []
+
+    if system_prompt:
+        messages.append({"role": "system", "content": system_prompt})
+
+    messages.extend(memory)
+    messages.append({"role": "user", "content": user_prompt})
+
+    return messages
+
+
+def call_image_microservice(prompt):
+    try:
+        response = requests.post(
+            IMAGE_SPACE_URL,
+            json={"prompt": prompt},
+            timeout=60
+        )
+        return response.json()
+    except:
+        return {"error": "Image service unavailable"}
+
+
+def route_request(prompt, user_id):
+
+    # ==========================
+    # Image Command
+    # ==========================
+    if prompt.startswith("/image"):
+        clean_prompt = prompt.replace("/image", "").strip()
+        return call_image_microservice(clean_prompt)
+
+    # ==========================
+    # Load Memory
+    # ==========================
+    memory = load_memory(user_id)
+
+    # ==========================
+    # AI Classification
+    # ==========================
+    classification = classify_prompt(prompt)
+
+    intent = classification.get("intent", "chat")
+    needs_search = classification.get("needs_search", False)
+
+    system_prompt = "You are ZXAI, a helpful advanced AI assistant."
+
+    # ==========================
+    # Greeting (fast return)
+    # ==========================
+    if intent == "greeting":
+        response = "Hello 👋 I am ZXAI. How can I assist you today?"
+        save_message(user_id, "user", prompt)
+        save_message(user_id, "assistant", response)
+        return {"response": response}
+
+    # ==========================
+    # Reasoning → Gemini
+    # ==========================
+    if intent == "reasoning":
+        messages = build_messages(system_prompt, memory, prompt)
+        response = call_gemini(messages)
+
+        save_message(user_id, "user", prompt)
+        save_message(user_id, "assistant", response)
+
+        return {"response": response}
+
+    # ==========================
+    # Live Data
+    # ==========================
+    if intent == "live_data" or needs_search:
+
+        web_data = search_web(prompt)
+
+        enriched_prompt = f"""
+User Question:
+{prompt}
+
+Web Data:
+{web_data}
+
+Use the web data if relevant.
+"""
+
+        messages = build_messages(system_prompt, memory, enriched_prompt)
 
-def is_reasoning(prompt):
-    keywords = ["solve", "explain deeply", "why", "logic", "proof"]
-    return any(k in prompt.lower() for k in keywords)
+        llama_answer = call_llama(messages)
+        gemini_answer = call_gemini(messages)
 
-def is_live_data(prompt):
-    keywords = ["today", "latest", "current", "news", "price", "weather"]
-    return any(k in prompt.lower() for k in keywords)
+        winner = judge_answers(llama_answer, gemini_answer)
 
-def route_request(prompt, image=None):
+        final_answer = gemini_answer if winner == 2 else llama_answer
 
-    # RAG first
-    rag = rag_response(prompt)
-    if rag:
-        return rag
+        save_message(user_id, "user", prompt)
+        save_message(user_id, "assistant", final_answer)
 
-    # Image case
-    if image:
-        return call_gemini(prompt)
+        return {"response": final_answer}
 
-    # Reasoning
-    if is_reasoning(prompt):
-        return call_gemini(prompt)
+    # ==========================
+    # Default Chat → Llama
+    # ==========================
+    messages = build_messages(system_prompt, memory, prompt)
 
-    # Live data
-    if is_live_data(prompt):
-        search_data = search_web(prompt)
-        llama_answer = call_llama(prompt + "\n\nWeb Data:\n" + search_data)
-        gemini_answer = call_gemini(prompt + "\n\nWeb Data:\n" + search_data)
+    response = call_llama(messages)
 
-        # Simple judge (can improve later)
-        if len(gemini_answer) > len(llama_answer):
-            return gemini_answer
-        return llama_answer
+    save_message(user_id, "user", prompt)
+    save_message(user_id, "assistant", response)
 
-    # Default
-    return call_llama(prompt)
+    return {"response": response}