Merge pull request #401 from vanna-ai/qdrant-get-training-data

Fix get_training_data for qdrant
vanna-ai · May 2, 2024 · 2a3bedd · 2a3bedd
2 parents c2e7ea2 + 089b546
commit 2a3bedd
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 4 deletions.
diff --git 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
@@ -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 = [