Update router.py

router.py

@@ -1,3 +1,7 @@
+import asyncio
+import time
+import requests
+
 from llm_clients import (
     call_llama,
     call_gemini,
@@ -9,16 +13,46 @@ from memory import save_message, load_memory
 from search_tool import search_web
 from rag_engine import rag_response
 
-import requests
 
+# =====================================
+# CONFIG
+# =====================================
 
-# 🔥 Image microservice endpoint
 IMAGE_SPACE_URL = "https://your-image-space.hf.space/generate"
 
-# 🔥 Simple in-memory cache
+CACHE_TTL_SECONDS = 300  # 5 minutes
+
 response_cache = {}
 
 
+# =====================================
+# CACHE HELPERS
+# =====================================
+
+def get_cached_response(cache_key):
+    entry = response_cache.get(cache_key)
+
+    if not entry:
+        return None
+
+    if time.time() > entry["expires_at"]:
+        del response_cache[cache_key]
+        return None
+
+    return entry["response"]
+
+
+def set_cache(cache_key, response):
+    response_cache[cache_key] = {
+        "response": response,
+        "expires_at": time.time() + CACHE_TTL_SECONDS
+    }
+
+
+# =====================================
+# MESSAGE BUILDER
+# =====================================
+
 def build_messages(system_prompt, memory, user_prompt):
 
     messages = []
@@ -32,6 +66,10 @@ def build_messages(system_prompt, memory, user_prompt):
     return messages
 
 
+# =====================================
+# IMAGE SERVICE
+# =====================================
+
 def call_image_microservice(prompt):
 
     try:
@@ -45,6 +83,22 @@ def call_image_microservice(prompt):
         return {"error": "Image service unavailable"}
 
 
+# =====================================
+# ASYNC LLM WRAPPERS
+# =====================================
+
+async def async_llama(messages):
+    return await asyncio.to_thread(call_llama, messages)
+
+
+async def async_gemini(messages):
+    return await asyncio.to_thread(call_gemini, messages)
+
+
+# =====================================
+# MAIN ROUTER
+# =====================================
+
 def route_request(prompt, user_id):
 
     cache_key = f"{user_id}:{prompt}"
@@ -52,8 +106,9 @@ def route_request(prompt, user_id):
     # ==========================
     # CACHE CHECK
     # ==========================
-    if cache_key in response_cache:
-        return {"response": response_cache[cache_key]}
+    cached = get_cached_response(cache_key)
+    if cached:
+        return {"response": cached}
 
     # ==========================
     # IMAGE COMMAND
@@ -66,9 +121,8 @@ def route_request(prompt, user_id):
     # RAG QUICK RESPONSE
     # ==========================
     rag_answer = rag_response(prompt)
-
     if rag_answer:
-        response_cache[cache_key] = rag_answer
+        set_cache(cache_key, rag_answer)
         return {"response": rag_answer}
 
     # ==========================
@@ -77,25 +131,25 @@ def route_request(prompt, user_id):
     memory = load_memory(user_id)
 
     # ==========================
-    # AI CLASSIFICATION
+    # CLASSIFY
     # ==========================
     classification = classify_prompt(prompt)
 
     intent = classification.get("intent", "chat")
     needs_search = classification.get("needs_search", False)
 
-    system_prompt = "You are ZXAI, a powerful advanced AI assistant."
+    system_prompt = "You are ZXAI, an advanced AI assistant."
 
     # ==========================
     # GREETING FAST PATH
     # ==========================
     if intent == "greeting":
-        response = "Hello 👋 I am ZXAI. How can I assist you today?"
+        response = "Hello 👋 I am ZXAI. How can I help you today?"
 
         save_message(user_id, "user", prompt)
         save_message(user_id, "assistant", response)
 
-        response_cache[cache_key] = response
+        set_cache(cache_key, response)
         return {"response": response}
 
     # ==========================
@@ -109,11 +163,11 @@ def route_request(prompt, user_id):
         save_message(user_id, "user", prompt)
         save_message(user_id, "assistant", response)
 
-        response_cache[cache_key] = response
+        set_cache(cache_key, response)
         return {"response": response}
 
     # ==========================
-    # LIVE DATA / SEARCH
+    # LIVE DATA (Parallel LLM)
     # ==========================
     if intent == "live_data" or needs_search:
 
@@ -131,17 +185,25 @@ Use web data if helpful.
 
         messages = build_messages(system_prompt, memory, enriched_prompt)
 
-        llama_answer = call_llama(messages)
-        gemini_answer = call_gemini(messages)
+        async def run_parallel():
+            llama_task = asyncio.create_task(async_llama(messages))
+            gemini_task = asyncio.create_task(async_gemini(messages))
 
-        winner = judge_answers(llama_answer, gemini_answer)
+            llama_answer = await llama_task
+            gemini_answer = await gemini_task
+
+            return llama_answer, gemini_answer
 
+        llama_answer, gemini_answer = asyncio.run(run_parallel())
+
+        winner = judge_answers(llama_answer, gemini_answer)
         final_answer = gemini_answer if winner == 2 else llama_answer
 
         save_message(user_id, "user", prompt)
         save_message(user_id, "assistant", final_answer)
 
-        response_cache[cache_key] = final_answer
+        set_cache(cache_key, final_answer)
+
         return {"response": final_answer}
 
     # ==========================
@@ -154,6 +216,6 @@ Use web data if helpful.
     save_message(user_id, "user", prompt)
     save_message(user_id, "assistant", response)
 
-    response_cache[cache_key] = response
+    set_cache(cache_key, response)
 
     return {"response": response}