vert

langgenius · Dec 25, 2024 · 4fd5bb8 · 4fd5bb8
1 parent c508489
commit 4fd5bb8
Showing 1 changed file with 10 additions and 8 deletions.
diff --git a/api/core/rag/datasource/vdb/tidb_vector/tidb_vector.py b/api/core/rag/datasource/vdb/tidb_vector/tidb_vector.py
@@ -203,14 +203,16 @@ def search_by_vector(self, query_vector: list[float], **kwargs: Any) -> list[Doc
             tidb_dist_func = "VEC_COSINE_DISTANCE"
 
         with Session(self._engine) as session:
-            select_statement = sql_text(
-                f"""SELECT meta, text, distance FROM (
-                        SELECT meta, text, {tidb_func}(vector, "{query_vector_str}")  as distance
-                        FROM {self._collection_name}
-                        ORDER BY distance
-                        LIMIT {top_k}
-                    ) t WHERE distance < {distance};"""
-            )
+            select_statement = sql_text(f"""
+                SELECT
+                  meta,
+                  text,
+                  {tidb_dist_func}(vector, "{query_vector_str}") AS distance
+                FROM `{self._collection_name}`
+                WHERE distance <= {distance}
+                ORDER BY distance
+                LIMIT {top_k}
+                """)
             res = session.execute(select_statement)
             results = [(row[0], row[1], row[2]) for row in res]
             for meta, text, distance in results: