Update main.py

main.py

@@ -8,7 +8,7 @@ from datetime import date, datetime
 from google import genai
 from google.genai import types
 
-warnings.filterwarnings("ignore")
+from pymysql.err import OperationalErrorwarnings.filterwarnings("ignore")
 
 # ───────────────────────────────────────────────────────────────────────────────
 # CONFIG
@@ -381,27 +381,81 @@ def keyword_search():
         rows = cur.fetchall()
     return jsonify({"ok": True, "query": q, "data": rows})
 
+
 @app.post("/api/similar")
 def similar_search():
     payload = request.get_json(force=True) or {}
     text = (payload.get("text") or "").strip()
     limit = max(1, min(int(payload.get("limit", 20)), 100))
+    candidates = int(payload.get("candidates", max(200, limit * 10)))  # ANN pre-topK
+    source_filter = (payload.get("source") or "").strip().upper()  # e.g., "AIC"
+
     if not text:
         return jsonify({"ok": False, "error": "text required"}), 400
-    vec = embed_text_to_vec1536(text)
-    vec_json = json.dumps(vec)
-    sql = (
-        "SELECT ps.object_id, ps.seq, ps.sentence, o.source, o.title, o.creator, "
-        f"VEC_COSINE_DISTANCE(ps.embedding, CAST(%s AS VECTOR({VEC_DIM}))) AS distance "
-        "FROM provenance_sentences ps "
-        "JOIN objects o ON o.object_id = ps.object_id "
-        "ORDER BY distance ASC "
-        "LIMIT %s"
+
+    # Embed without NumPy path
+    try:
+        import torch
+        vec_t = _load_model().encode([text], batch_size=1, show_progress_bar=False, convert_to_tensor=True)
+        if isinstance(vec_t, torch.Tensor):
+            vec = vec_t[0].detach().cpu().tolist()
+        else:
+            vec = list(vec_t[0])
+    except Exception as e:
+        return jsonify({"ok": False, "error": f"embedding_unavailable: {e}"}), 503
+
+    vec_json = json.dumps(_pad(vec, VEC_DIM))
+
+    # Build query with explicit HNSW usage and staged join
+    where_src = "WHERE o.source = %s" if source_filter else ""
+    sql = f"""
+    WITH nn AS (
+      SELECT /*+ USE_INDEX(ps, hnsw_vec) */
+             ps.sent_id, ps.object_id, ps.seq, ps.sentence,
+             VEC_COSINE_DISTANCE(ps.embedding, CAST(%s AS VECTOR({VEC_DIM}))) AS distance
+      FROM provenance_sentences ps
+      ORDER BY distance
+      LIMIT %s
     )
-    with cursor() as cur:
-        cur.execute(sql, (vec_json, limit))
-        rows = cur.fetchall()
-    return jsonify({"ok": True, "device": _DEVICE_INFO, "query": text, "data": rows})
+    SELECT nn.object_id, nn.seq, nn.sentence, o.source, o.title, o.creator, nn.distance
+    FROM nn
+    JOIN objects o ON o.object_id = nn.object_id
+    {where_src}
+    ORDER BY nn.distance
+    LIMIT %s
+    """
+
+    params = [vec_json, candidates]
+    if source_filter:
+        params.append(source_filter)
+    params.append(limit)
+
+    try:
+        with cursor() as cur:
+            cur.execute(sql, params)
+            rows = cur.fetchall()
+        return jsonify({"ok": True, "device": _DEVICE_INFO, "query": text, "data": rows,
+                        "meta": {"limit": limit, "candidates": candidates, "source": source_filter or None}})
+    except OperationalError as e:
+        # TiDB OOM (1105) → retry with smaller candidate set automatically
+        if e.args and e.args[0] == 1105 and candidates > max(100, limit * 4):
+            smaller = max(100, limit * 4)
+            params2 = [vec_json, smaller]
+            if source_filter:
+                params2.append(source_filter)
+            params2.append(limit)
+            try:
+                with cursor() as cur:
+                    cur.execute(sql, params2)
+                    rows = cur.fetchall()
+                return jsonify({"ok": True, "device": _DEVICE_INFO, "query": text, "data": rows,
+                                "meta": {"limit": limit, "candidates": smaller, "source": source_filter or None,
+                                         "note": "retried with smaller candidate set"}})
+            except Exception as e2:
+                return jsonify({"ok": False, "error": f"oom_retry_failed: {e2}"}), 500
+        # Not OOM or still failed
+        return jsonify({"ok": False, "error": f"query_failed: {e}"}), 500
+
 
 @app.get("/api/vocab")
 def vocab():