feat(chat): retrieve top-k passages through the Librarian seams (#1075)

application/database/db.py

@@ -3158,6 +3158,19 @@ def can_use_pgvector_similarity(self) -> bool:
         self._pgvector_similarity_ready = ready
         return ready
 
+    def _rollback_quietly(self) -> None:
+        """Best-effort ``session.rollback()`` after a failed pgvector statement.
+
+        On Postgres a failed statement aborts the open transaction; without a
+        rollback the next query on the shared session raises
+        ``InFailedSqlTransaction``. Never raises: the caller is already on the
+        degrade-to-no-match path.
+        """
+        try:
+            self.session.rollback()
+        except Exception:  # noqa: BLE001 - best effort only
+            pass
+
     def find_most_similar_embedding_id(
         self,
         query_embedding: List[float],
@@ -3208,12 +3221,76 @@ def find_most_similar_embedding_id(
         except Exception as exc:
             # Match prior sklearn-path resilience: chat/import should degrade
             # to "no match" rather than 500 on a transient driver/query error.
+            # A failed statement leaves the Postgres transaction aborted, so
+            # roll back or every later query on this session fails too.
+            self._rollback_quietly()
             logger.warning(
                 "pgvector similarity query failed (%s); returning no match",
                 exc,
             )
             return None, None
 
+    def find_top_k_similar_embedding_ids(
+        self,
+        query_embedding: List[float],
+        *,
+        doc_type: str,
+        id_column: str,
+        k: int,
+        similarity_threshold: float,
+    ) -> List[Tuple[str, float]]:
+        """Top-K cosine matches via pgvector ``<=>`` (Postgres only). #1075.
+
+        Returns ``[(object_id, score), ...]`` rank-ordered (best first), only
+        rows with ``score >= similarity_threshold``; ``[]`` when nothing
+        clears it or on a transient query error. Refuses SQLite / missing
+        ``embedding_vec`` with ``SystemExit`` exactly like
+        :meth:`find_most_similar_embedding_id`; callers gate on
+        ``can_use_pgvector_similarity()`` and fall back to sklearn otherwise.
+        """
+        from application.database.pgvector_utils import (
+            fail_pgvector_unavailable,
+            to_pgvector_literal,
+            top_k_similar_ids_sql,
+        )
+        from sqlalchemy import text as sql_text
+
+        if k <= 0:
+            raise ValueError(f"k must be > 0, got {k}")
+        if not self.can_use_pgvector_similarity():
+            fail_pgvector_unavailable(context="find_top_k_similar_embedding_ids")
+
+        sql = top_k_similar_ids_sql(id_column)
+        bind = self.session.get_bind()
+        params = {
+            "q": to_pgvector_literal(query_embedding),
+            "doc_type": doc_type,
+            "k": int(k),
+        }
+        try:
+            if hasattr(self.session, "execute"):
+                rows = self.session.execute(sql_text(sql), params).fetchall()
+            else:
+                with bind.connect() as conn:
+                    rows = conn.execute(sql_text(sql), params).fetchall()
+            out: List[Tuple[str, float]] = []
+            for row in rows or []:
+                score = float(row.score)
+                if score < similarity_threshold:
+                    # rows are rank-ordered, so everything after is lower too
+                    break
+                out.append((str(row.object_id), score))
+            return out
+        except SystemExit:
+            raise
+        except Exception as exc:
+            self._rollback_quietly()
+            logger.warning(
+                "pgvector top-k similarity query failed (%s); returning no matches",
+                exc,
+            )
+            return []
+
     def assert_embedding_contract(
         self,
         *,

application/database/pgvector_utils.py

@@ -300,3 +300,23 @@ def most_similar_id_sql(id_column: str) -> str:
         f"ORDER BY embedding_vec <=> CAST(:q AS vector) "
         f"LIMIT 1"
     )
+
+
+def top_k_similar_ids_sql(id_column: str) -> str:
+    """SQL for top-K cosine similarity over ``embedding_vec`` (#1075).
+
+    Same shape as :func:`most_similar_id_sql` but ``LIMIT :k`` so chat can
+    pull several passages in one round trip. Rows come back rank-ordered
+    (closest first); the caller applies the threshold / break-off.
+    """
+    if id_column not in ("cre_id", "node_id"):
+        raise ValueError(f"id_column must be cre_id or node_id, got {id_column!r}")
+    return (
+        f"SELECT {id_column} AS object_id, "
+        f"1 - (embedding_vec <=> CAST(:q AS vector)) AS score "
+        f"FROM embeddings "
+        f"WHERE doc_type = :doc_type AND {id_column} IS NOT NULL "
+        f"AND embedding_vec IS NOT NULL "
+        f"ORDER BY embedding_vec <=> CAST(:q AS vector) "
+        f"LIMIT :k"
+    )