Remove corpus counts from chat chrome

apps/ai-service/rag/routing.py

@@ -53,6 +53,12 @@ class CatalogDrugResolver:
             for alias in aliases
             if (normalized := normalize_name(alias))
         ]
+        self._alias_to_drug_ids: dict[str, set[str]] = {}
+        for drug_id, alias in self._aliases:
+            self._alias_to_drug_ids.setdefault(alias, set()).add(drug_id)
+        self._max_alias_tokens = max(
+            (len(alias.split()) for alias in self._alias_to_drug_ids), default=0
+        )
         self._fuzzy_threshold = fuzzy_threshold
         self._ambiguity_margin = ambiguity_margin
 
@@ -69,18 +75,24 @@ class CatalogDrugResolver:
     # "Dịch vụ đang gặp sự cố" — not a provider outage at all. Caching by
     # exact input turns all but the newest turn's own text into a dict
     # lookup on every subsequent call.
-    @functools.lru_cache(maxsize=4096)
+    @functools.lru_cache(maxsize=4096)  # noqa: B019 - bounded process singleton
     def resolve(self, query: str) -> DrugResolution:
         normalized_query = normalize_name(query)
         query_tokens = normalized_query.split()
-        exact = [
-            (drug_id, alias, match.start(1), match.end(1))
-            for drug_id, alias in self._aliases
-            for match in [
-                re.search(rf"(?:^| )({re.escape(alias)})(?:$| )", normalized_query)
-            ]
-            if match
-        ]
+        # Exact matching used to compile and run one regex for every alias
+        # (~10k) on every new line. Enumerating the query's contiguous token
+        # spans and looking them up in an immutable alias index is equivalent
+        # at word boundaries and turns the common exact-name path into O(q²)
+        # in the short query rather than O(catalog).
+        exact: list[tuple[str, str, int, int]] = []
+        for start in range(len(query_tokens)):
+            last = min(len(query_tokens), start + self._max_alias_tokens)
+            for end in range(start + 1, last + 1):
+                alias = " ".join(query_tokens[start:end])
+                exact.extend(
+                    (drug_id, alias, start, end)
+                    for drug_id in self._alias_to_drug_ids.get(alias, ())
+                )
         if exact:
             maximal = [
                 row for row in exact
@@ -137,12 +149,22 @@ class CatalogDrugResolver:
         needle = normalize_name(prefix)
         if not needle:
             return []
-        matches: list[tuple[tuple[int, int], str]] = []
+        matches: list[tuple[tuple[int, int, int], str]] = []
         for drug_id, alias in self._aliases:
             position = alias.find(needle)
             if position < 0:
                 continue
-            matches.append(((0 if position == 0 else 1, len(alias)), drug_id))
+            canonical = normalize_name(drug_id.replace("_", " "))
+            canonical_words = canonical.split()
+            if canonical.startswith(needle):
+                source_rank = 0
+            elif any(word.startswith(needle) for word in canonical_words):
+                source_rank = 1
+            elif position == 0:
+                source_rank = 2
+            else:
+                source_rank = 3
+            matches.append(((source_rank, len(canonical), len(alias)), drug_id))
         matches.sort()
         ordered: list[str] = []
         seen: set[str] = set()
@@ -157,7 +179,7 @@ class CatalogDrugResolver:
 
     # See the comment on `resolve` above — same cost, same fix, same
     # single-caller read-only usage (safe to hand back a cached list).
-    @functools.lru_cache(maxsize=4096)
+    @functools.lru_cache(maxsize=4096)  # noqa: B019 - bounded process singleton
     def suggest(
         self, query: str, k: int = 3, min_score: float = 0.5
     ) -> list[tuple[str, float]]: