From 089b5467e301486a76fc49bf91183666f94e758b Mon Sep 17 00:00:00 2001
From: Zain Hoda <7146154+zainhoda@users.noreply.github.com>
Date: Thu, 2 May 2024 09:09:14 -0400
Subject: [PATCH] Fix get_training_data for qdrant

---
 pyproject.toml             | 2 +-
 src/vanna/qdrant/qdrant.py | 4 +---
 2 files changed, 2 insertions(+), 4 deletions(-)

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 = [