Upd translation modules. Increase kb retrieval threshold

app.py

@@ -10,6 +10,7 @@ from pymongo import MongoClient
 from google import genai
 from sentence_transformers import SentenceTransformer
 from memory import MemoryManager
+from translation import translate_query
 
 # ✅ Enable Logging for Debugging
 import logging
@@ -129,8 +130,8 @@ def load_faiss_index():
             logger.error("[KB] ❌ FAISS index not found in GridFS.")
     return index
 
-# ✅ Retrieve Medical Info
-def retrieve_medical_info(query, k=5, min_sim=0.6): # Min similarity between query and kb is to be 60%
+# ✅ Retrieve Medical Info (256,916 scenario)
+def retrieve_medical_info(query, k=5, min_sim=0.8): # Min similarity between query and kb is to be 80%
     global index
     index = load_faiss_index()
     if index is None:
@@ -140,15 +141,19 @@ def retrieve_medical_info(query, k=5, min_sim=0.6): # Min similarity between que
     D, I = index.search(query_vec, k=k)
     # Filter by cosine threshold
     results = []
+    seen = set()  # avoid near-duplicate KB responses
     for score, idx in zip(D[0], I[0]):
         if score < min_sim:
             continue
         doc = qa_collection.find_one({"i": int(idx)})
         if doc:
-            results.append(doc.get("Doctor", "No answer available."))
-    return results if results else ["No relevant medical entries found."]
+            answer = doc.get("Doctor", "No answer available.")
+            if answer not in seen:
+                seen.add(answer)
+                results.append(answer)
 
-# ✅ Retrieve Sym-Dia Info (4962 scenario)
+
+# ✅ Retrieve Sym-Dia Info (4,962 scenario)
 def retrieve_diagnosis_from_symptoms(symptom_text, top_k=5, min_sim=0.4):
     global SYMPTOM_VECTORS, SYMPTOM_DOCS
     # Lazy load
@@ -162,12 +167,17 @@ def retrieve_diagnosis_from_symptoms(symptom_text, top_k=5, min_sim=0.4):
     # Similarity compute
     sims = SYMPTOM_VECTORS @ qvec  # cosine
     sorted_idx = np.argsort(sims)[-top_k:][::-1]
-    # Final
-    return [
-        SYMPTOM_DOCS[i]["answer"]
-        for i in sorted_idx
-        if sims[i] >= min_sim
-    ]
+    seen_diag = set()
+    final = [] # Dedup
+    for i in sorted_idx:
+        sim = sims[i]
+        if sim < min_sim:
+            continue
+        label = SYMPTOM_DOCS[i]["prognosis"]
+        if label not in seen_diag:
+            final.append(SYMPTOM_DOCS[i]["answer"])
+            seen_diag.add(label)
+    return final
 
 # ✅ Gemini Flash API Call
 def gemini_flash_completion(prompt, model, temperature=0.7):
@@ -186,6 +196,10 @@ class RAGMedicalChatbot:
         self.retrieve = retrieve_function
 
     def chat(self, user_id: str, user_query: str, lang: str = "EN") -> str:
+        # 0. Translate query if not EN, this help our RAG system
+        if lang.upper() in {"VI", "ZH"}:
+            user_query = translate_query(user_query, lang.lower())
+
         # 1. Fetch knowledge
         ## a. KB for generic QA retrieval
         retrieved_info = self.retrieve(user_query)
@@ -205,7 +219,7 @@ class RAGMedicalChatbot:
             parts.append("Relevant context from prior conversation:\n" + "\n".join(context))
         # Load up guideline
         if knowledge_base:
-            parts.append(f"Medical knowledge (256,916 medical scenario): {knowledge_base}")
+            parts.append(f"Medical scenario knowledge: {knowledge_base}")
         # Symptom-Diagnosis prediction RAG
         if diagnosis_guides:
             parts.append("Symptom-based diagnosis guidance:\n" + "\n".join(diagnosis_guides))

translation.py

@@ -0,0 +1,26 @@
+# translation.py
+from transformers import pipeline
+import logging
+
+logger = logging.getLogger("translation-agent")
+logging.basicConfig(level=logging.INFO, format="%(asctime)s — %(name)s — %(levelname)s — %(message)s", force=True) # Change INFO to DEBUG for full-ctx JSON loader
+
+# To use lazy model loader
+vi_en = None
+zh_en = None
+
+def translate_query(text: str, lang_code: str) -> str:
+    global vi_en, zh_en
+    if lang_code == "vi":
+        if vi_en is None:
+            vi_en = pipeline("translation", model="VietAI/envit5-translation", src_lang="vi", tgt_lang="en", device=-1)
+        result = vi_en(text, max_length=512)[0]["translation_text"]
+        logger.info(f"[En-Vi] Query in `{lang_code}` translated to: {result}")
+        return result
+    elif lang_code == "zh":
+        if zh_en is None:
+            zh_en = pipeline("translation", model="Helsinki-NLP/opus-mt-zh-en", device=-1)
+        result = zh_en(text, max_length=512)[0]["translation_text"]
+        logger.info(f"[En-Zh] Query in `{lang_code}` translated to: {result}")
+        return result
+    return text