Update rag/pipeline.py

rag/pipeline.py

@@ -23,41 +23,59 @@ def _route_companies(
         "msft": "Microsoft"
     }
 
-    labels = list(settings.router_labels)
-    entities = router_model.predict_entities(query, labels, threshold=settings.router_threshold)
+    try:
+        labels = list(settings.router_labels)
+        entities = router_model.predict_entities(query, labels, threshold=settings.router_threshold)
+    except Exception as e:
+        print(f"Warning: GLiNER failed with error {e}, falling back to keywords.")
+        entities = []
     
     detected_targets = []
     unsupported_targets = []
     
     for e in entities:
-        name_clean = (e.get("text") or "").lower().strip()
-        found_match = False
+        text = e.get("text", "")
+        name_clean = text.lower().strip().strip(".,?!'s")
         
-        for key, canonical_name in allowed_companies.items():
+        found_match = False
+        for key, canonical_name in allowed_companies.items(): 
+            
             if key in name_clean:
                 detected_targets.append(canonical_name)
                 found_match = True
                 break
         
         if not found_match:
-            unsupported_targets.append(e.get("text"))
+            unsupported_targets.append(text)
 
+    
+    if not detected_targets:
+        query_lower = query.lower()
+        for key, canonical_name in allowed_companies.items():
+            if key in query_lower:
+                detected_targets.append(canonical_name)  
+   
     detected_targets = list(set(detected_targets))
-    unsupported_targets = list(set(unsupported_targets))
 
+   
+    
+    
+    if detected_targets:
+        return detected_targets, None
+
+    
     if unsupported_targets:
+        unique_unsupported = list(set(unsupported_targets))
         return [], (
-            f"⛔ **Out of Scope:** I detected a request for **{', '.join(unsupported_targets)}**. "
+            f"⛔ **Out of Scope:** I detected a request for **{', '.join(unique_unsupported)}**. "
             "This system only has access to **Microsoft** and **Apple** data."
         )
 
-    if not detected_targets:
-        return [], (
-            "❓ **Ambiguous Query:** I could not identify a specific company (Apple or Microsoft). "
-            "Please name the company you want to analyze."
-        )
-
-    return detected_targets, None
+    # Cas C : Le désert total (vraiment rien trouvé)
+    return [], (
+        "❓ **Ambiguous Query:** I could not identify a specific company (Apple or Microsoft). "
+        "Please name the company explicitly in your question."
+    )
 
 def _filter_docs(
     dataset: List[Dict[str, Any]],