Merge pull request vanna-ai#490 from Anush008/qdrant-suffix-fix

fix: Qdrant id_suffixes when using custom collection names
boldsbrain · Jun 10, 2024 · 56b782b · 56b782b
2 parents 7ad88c1 + e0ff865
commit 56b782b
Showing 1 changed file with 9 additions and 13 deletions.
diff --git a/src/vanna/qdrant/qdrant.py b/src/vanna/qdrant/qdrant.py
@@ -39,16 +39,6 @@ class Qdrant_VectorStore(VannaBase):
         TypeError: If config["client"] is not a `qdrant_client.QdrantClient` instance
     """
 
-    documentation_collection_name = "documentation"
-    ddl_collection_name = "ddl"
-    sql_collection_name = "sql"
-
-    id_suffixes = {
-        ddl_collection_name: "ddl",
-        documentation_collection_name: "doc",
-        sql_collection_name: "sql",
-    }
-
     def __init__(
         self,
         config={},
@@ -80,15 +70,21 @@ def __init__(
         self.collection_params = config.get("collection_params", {})
         self.distance_metric = config.get("distance_metric", models.Distance.COSINE)
         self.documentation_collection_name = config.get(
-            "documentation_collection_name", self.documentation_collection_name
+            "documentation_collection_name", "documentation"
         )
         self.ddl_collection_name = config.get(
-            "ddl_collection_name", self.ddl_collection_name
+            "ddl_collection_name", "ddl"
         )
         self.sql_collection_name = config.get(
-            "sql_collection_name", self.sql_collection_name
+            "sql_collection_name", "sql"
         )
 
+        self.id_suffixes = {
+            self.ddl_collection_name: "ddl",
+            self.documentation_collection_name: "doc",
+            self.sql_collection_name: "sql",
+        }
+
         self._setup_collections()
 
     def add_question_sql(self, question: str, sql: str, **kwargs) -> str: