diff --git a/querybook/server/const/ai_assistant.py b/querybook/server/const/ai_assistant.py index 486a11442..1db8b888f 100644 --- a/querybook/server/const/ai_assistant.py +++ b/querybook/server/const/ai_assistant.py @@ -25,6 +25,6 @@ class AICommandType(Enum): # how many docs to fetch from vector store, it may include both table and query summary docs and they need additional processing. DEFAULT_VECTOR_STORE_FETCH_LIMIT = 30 # how many tables to return from vector table search eventually -DEFAUTL_TABLE_SEARCH_LIMIT = 10 +DEFAULT_TABLE_SEARCH_LIMIT = 10 # how many tables to select for text-to-sql -DEFAUTL_TABLE_SELECT_LIMIT = 3 +DEFAULT_TABLE_SELECT_LIMIT = 3 diff --git a/querybook/server/lib/ai_assistant/base_ai_assistant.py b/querybook/server/lib/ai_assistant/base_ai_assistant.py index 0096c107a..b5b1ae3f1 100644 --- a/querybook/server/lib/ai_assistant/base_ai_assistant.py +++ b/querybook/server/lib/ai_assistant/base_ai_assistant.py @@ -7,7 +7,7 @@ from app.db import with_session from const.ai_assistant import ( - DEFAUTL_TABLE_SELECT_LIMIT, + DEFAULT_TABLE_SELECT_LIMIT, MAX_SAMPLE_QUERY_COUNT_FOR_TABLE_SUMMARY, AICommandType, ) @@ -403,7 +403,7 @@ def summarize_query( def find_tables(self, metastore_id, question, session=None): """Search similar tables from vector store first, and then ask LLM to select most suitable tables for the question. - It will return at most `DEFAUTL_TABLE_SELECT_LIMIT` tables by default. + It will return at most `DEFAULT_TABLE_SELECT_LIMIT` tables by default. """ try: tables = get_vector_store().search_tables( @@ -443,7 +443,7 @@ def find_tables(self, metastore_id, question, session=None): table_docs[full_table_name] = summary prompt = self._get_table_select_prompt( - top_n=DEFAUTL_TABLE_SELECT_LIMIT, + top_n=DEFAULT_TABLE_SELECT_LIMIT, table_schemas=table_docs, question=question, ) diff --git a/querybook/server/lib/vector_store/base_vector_store.py b/querybook/server/lib/vector_store/base_vector_store.py index 333d5ead7..a2f82a506 100644 --- a/querybook/server/lib/vector_store/base_vector_store.py +++ b/querybook/server/lib/vector_store/base_vector_store.py @@ -3,7 +3,7 @@ from typing import Literal, Optional from const.ai_assistant import ( - DEFAUTL_TABLE_SEARCH_LIMIT, + DEFAULT_TABLE_SEARCH_LIMIT, DEFAULT_VECTOR_STORE_FETCH_LIMIT, DEFAULT_SIMILARITY_SCORE_THRESHOLD, DEFAULT_SIMILARITY_SCORE_THRESHOLD_GREAT_MATCH, @@ -60,7 +60,7 @@ def search_tables( text: str, search_type: Optional[Literal["table", "query"]] = None, threshold=DEFAULT_SIMILARITY_SCORE_THRESHOLD, - k=DEFAUTL_TABLE_SEARCH_LIMIT, + k=DEFAULT_TABLE_SEARCH_LIMIT, fetch_k=DEFAULT_VECTOR_STORE_FETCH_LIMIT, ) -> list[tuple[str, int]]: """Find tables using embedding based table search. diff --git a/querybook/server/logic/vector_store.py b/querybook/server/logic/vector_store.py index 88b94c48e..29669cd08 100644 --- a/querybook/server/logic/vector_store.py +++ b/querybook/server/logic/vector_store.py @@ -1,6 +1,6 @@ from app.db import with_session from const.ai_assistant import ( - DEFAUTL_TABLE_SEARCH_LIMIT, + DEFAULT_TABLE_SEARCH_LIMIT, MAX_SAMPLE_QUERY_COUNT_FOR_TABLE_SUMMARY, ) from langchain.docstore.document import Document @@ -154,7 +154,7 @@ def delete_table_doc(table_id: int): def search_tables( - metastore_id, keywords, filters=None, limit=DEFAUTL_TABLE_SEARCH_LIMIT + metastore_id, keywords, filters=None, limit=DEFAULT_TABLE_SEARCH_LIMIT ): """search tables from vector store and get the table details from elastic search."""