From e1754dc4a1c2b343899d07bc4eada57da43353b2 Mon Sep 17 00:00:00 2001
From: Victor <rongzhang>
Date: Tue, 7 Jan 2025 14:26:58 -0500
Subject: [PATCH] address comments

---
 docs_website/docs/user_guide/ai_assistant.mdx          | 4 ++++
 querybook/server/const/ai_assistant.py                 | 2 +-
 querybook/server/lib/vector_store/base_vector_store.py | 3 ++-
 3 files changed, 7 insertions(+), 2 deletions(-)

diff --git a/docs_website/docs/user_guide/ai_assistant.mdx b/docs_website/docs/user_guide/ai_assistant.mdx
index a867fdddb..97489dce2 100644
--- a/docs_website/docs/user_guide/ai_assistant.mdx
+++ b/docs_website/docs/user_guide/ai_assistant.mdx
@@ -48,4 +48,8 @@ This feature is disabled by default. You can enable it in the Editor tab of user
 
 If [vector store](../integrations/add_ai_assistant.mdx#vector-store) of the AI assistant plugin is also enabled, you'll be able to search the tables by natual language as well as keyword based search.
 
+## Search Tables and Queries by Natural Language
+
+If the [vector store](../integrations/add_ai_assistant.mdx#vector-store) of the AI assistant plugin is enabled, you'll be able to search both tables and queries using natural language in addition to traditional keyword-based search.
+
 ![](/img/user_guide/table_vector_search.png)
diff --git a/querybook/server/const/ai_assistant.py b/querybook/server/const/ai_assistant.py
index d1de851fc..eede34f92 100644
--- a/querybook/server/const/ai_assistant.py
+++ b/querybook/server/const/ai_assistant.py
@@ -27,7 +27,7 @@ class AICommandType(Enum):
 DEFAULT_VECTOR_STORE_FETCH_LIMIT = 30
 # how many tables to return from vector table search eventually
 DEFAULT_TABLE_SEARCH_LIMIT = 10
-# how many tables to return from vector query search eventually
+# how many queries to return from vector query search eventually
 DEFAULT_QUERY_SEARCH_LIMIT = 10
 # how many tables to select for text-to-sql
 DEFAULT_TABLE_SELECT_LIMIT = 3
diff --git a/querybook/server/lib/vector_store/base_vector_store.py b/querybook/server/lib/vector_store/base_vector_store.py
index 301a371ad..4eb8ee77d 100644
--- a/querybook/server/lib/vector_store/base_vector_store.py
+++ b/querybook/server/lib/vector_store/base_vector_store.py
@@ -3,6 +3,7 @@
 from typing import Literal, Optional
 
 from const.ai_assistant import (
+    DEFAULT_QUERY_SEARCH_LIMIT,
     DEFAULT_TABLE_SEARCH_LIMIT,
     DEFAULT_VECTOR_STORE_FETCH_LIMIT,
     DEFAULT_SIMILARITY_SCORE_THRESHOLD,
@@ -113,7 +114,7 @@ def search_query(
         self,
         text: str,
         threshold: float = DEFAULT_SIMILARITY_SCORE_THRESHOLD,
-        k=DEFAULT_TABLE_SEARCH_LIMIT,
+        k=DEFAULT_QUERY_SEARCH_LIMIT,
         fetch_k=DEFAULT_VECTOR_STORE_FETCH_LIMIT,
     ) -> list[tuple[int, float]]:
         """