Create routing.py

rag/routing.py

@@ -0,0 +1,58 @@
+from typing import List, Tuple
+import re
+
+from rag.config import Settings
+from rag.logging_utils import get_logger
+
+logger = get_logger(__name__)
+
+_COMMON_FALSE_POSITIVES = {
+    "revenue", "income", "balance", "sheet", "notes", "q1", "q2", "q3", "q4",
+    "fy", "cash", "assets", "liabilities", "growth", "profit", "loss",
+}
+
+def _looks_like_company_token(s: str) -> bool:
+    s = s.strip().lower()
+    if not s or len(s) < 3:
+        return False
+    if s in _COMMON_FALSE_POSITIVES:
+        return False
+    # avoid pure numbers / quarter tokens etc
+    if re.fullmatch(r"[\d\W_]+", s):
+        return False
+    return True
+
+def route_companies(query: str, router_model, settings: Settings) -> Tuple[List[str], List[str]]:
+    """
+    Returns (detected_companies_in_scope, out_of_scope_mentions)
+    Soft behavior:
+      - If out-of-scope is detected, we DON'T hard-fail automatically.
+        We'll mention it + still allow broad search, unless user insists.
+    """
+    labels = list(settings.router_labels)
+    entities = router_model.predict_entities(query, labels, threshold=settings.router_threshold)
+
+    detected = set()
+    out_of_scope = set()
+
+    for e in entities:
+        text = (e.get("text") or "").strip()
+        if not _looks_like_company_token(text):
+            continue
+        name = text.lower()
+
+        if "microsoft" in name or "msft" in name:
+            detected.add("Microsoft")
+        elif "apple" in name or "aapl" in name:
+            detected.add("Apple")
+        else:
+            out_of_scope.add(text)
+
+    # Also fallback on simple keyword scan (super robust, low cost)
+    ql = query.lower()
+    if "microsoft" in ql or "msft" in ql:
+        detected.add("Microsoft")
+    if "apple" in ql or "aapl" in ql:
+        detected.add("Apple")
+
+    return sorted(detected), sorted(out_of_scope)