diff --git a/pyproject.toml b/pyproject.toml index e69e06c1..ccc10b3b 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -42,7 +42,7 @@ gemini = ["google-generativeai"] marqo = ["marqo"] zhipuai = ["zhipuai"] ollama = ["ollama", "httpx"] -qdrant = ["qdrant-client"] +qdrant = ["qdrant-client", "fastembed"] vllm = ["vllm"] opensearch = ["opensearch-py", "opensearch-dsl"] hf = ["transformers"] diff --git a/src/vanna/qdrant/qdrant.py b/src/vanna/qdrant/qdrant.py index fc5051f4..3730af0d 100644 --- a/src/vanna/qdrant/qdrant.py +++ b/src/vanna/qdrant/qdrant.py @@ -157,7 +157,7 @@ def get_training_data(self, **kwargs) -> pd.DataFrame: if ddl_data := self._get_all_points(DDL_COLLECTION_NAME): ddl_list = [data.payload["ddl"] for data in ddl_data] id_list = [ - self._format_point_id(data.id, DDL_COLLECTION_NAME) for data in sql_data + self._format_point_id(data.id, DDL_COLLECTION_NAME) for data in ddl_data ] df_ddl = pd.DataFrame( @@ -172,8 +172,6 @@ def get_training_data(self, **kwargs) -> pd.DataFrame: df = pd.concat([df, df_ddl]) - doc_data = self.documentation_collection.get() - if doc_data := self._get_all_points(DOCUMENTATION_COLLECTION_NAME): document_list = [data.payload["documentation"] for data in doc_data] id_list = [