vanna 0.5.0__tar.gz → 0.5.2__tar.gz

{vanna-0.5.0 → vanna-0.5.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: vanna
-Version: 0.5.0
+Version: 0.5.2
 Summary: Generate SQL queries from natural language
 Author-email: Zain Hoda <zain@vanna.ai>
 Requires-Python: >=3.9
@@ -57,6 +57,7 @@ Requires-Dist: opensearch-dsl ; extra == "opensearch"
 Requires-Dist: psycopg2-binary ; extra == "postgres"
 Requires-Dist: db-dtypes ; extra == "postgres"
 Requires-Dist: qdrant-client ; extra == "qdrant"
+Requires-Dist: fastembed ; extra == "qdrant"
 Requires-Dist: snowflake-connector-python ; extra == "snowflake"
 Requires-Dist: tox ; extra == "test"
 Requires-Dist: vllm ; extra == "vllm"
@@ -111,7 +112,7 @@ Vanna works in two easy steps - train a RAG "model" on your data, and then ask q
 
 If you don't know what RAG is, don't worry -- you don't need to know how this works under the hood to use it. You just need to know that you "train" a model, which stores some metadata and then use it to "ask" questions.
 
-See the [base class](src/vanna/base/base.py) for more details on how this works under the hood.
+See the [base class](https://github.com/vanna-ai/vanna/blob/main/src/vanna/base/base.py) for more details on how this works under the hood.
 
 ## User Interfaces
 These are some of the user interfaces that we've built using Vanna. You can use these as-is or as a starting point for your own custom interface.
@@ -304,7 +305,7 @@ Fine-Tuning
     - Expose to your end users via Slackbot, web app, Streamlit app, or a custom front end.
 
 ## Extending Vanna
-Vanna is designed to connect to any database, LLM, and vector database. There's a [VannaBase](src/vanna/base/base.py) abstract base class that defines some basic functionality. The package provides implementations for use with OpenAI and ChromaDB. You can easily extend Vanna to use your own LLM or vector database. See the [documentation](https://vanna.ai/docs/) for more details.
+Vanna is designed to connect to any database, LLM, and vector database. There's a [VannaBase](https://github.com/vanna-ai/vanna/blob/main/src/vanna/base/base.py) abstract base class that defines some basic functionality. The package provides implementations for use with OpenAI and ChromaDB. You can easily extend Vanna to use your own LLM or vector database. See the [documentation](https://vanna.ai/docs/) for more details.
 
 ## Vanna in 100 Seconds
 

{vanna-0.5.0 → vanna-0.5.2}/README.md

@@ -25,7 +25,7 @@ Vanna works in two easy steps - train a RAG "model" on your data, and then ask q
 
 If you don't know what RAG is, don't worry -- you don't need to know how this works under the hood to use it. You just need to know that you "train" a model, which stores some metadata and then use it to "ask" questions.
 
-See the [base class](src/vanna/base/base.py) for more details on how this works under the hood.
+See the [base class](https://github.com/vanna-ai/vanna/blob/main/src/vanna/base/base.py) for more details on how this works under the hood.
 
 ## User Interfaces
 These are some of the user interfaces that we've built using Vanna. You can use these as-is or as a starting point for your own custom interface.
@@ -218,7 +218,7 @@ Fine-Tuning
     - Expose to your end users via Slackbot, web app, Streamlit app, or a custom front end.
 
 ## Extending Vanna
-Vanna is designed to connect to any database, LLM, and vector database. There's a [VannaBase](src/vanna/base/base.py) abstract base class that defines some basic functionality. The package provides implementations for use with OpenAI and ChromaDB. You can easily extend Vanna to use your own LLM or vector database. See the [documentation](https://vanna.ai/docs/) for more details.
+Vanna is designed to connect to any database, LLM, and vector database. There's a [VannaBase](https://github.com/vanna-ai/vanna/blob/main/src/vanna/base/base.py) abstract base class that defines some basic functionality. The package provides implementations for use with OpenAI and ChromaDB. You can easily extend Vanna to use your own LLM or vector database. See the [documentation](https://vanna.ai/docs/) for more details.
 
 ## Vanna in 100 Seconds
 

{vanna-0.5.0 → vanna-0.5.2}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "flit_core.buildapi"
 
 [project]
 name = "vanna"
-version = "0.5.0"
+version = "0.5.2"
 authors = [
   { name="Zain Hoda", email="zain@vanna.ai" },
 ]
@@ -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"]

{vanna-0.5.0 → vanna-0.5.2}/src/vanna/base/base.py

@@ -1146,19 +1146,14 @@ class VannaBase(ABC):
 
         conn = None
 
-        try:
-            conn = bigquery.Client(project=project_id)
-        except:
-            print("Could not found any google cloud implicit credentials")
-
-        if cred_file_path:
+        if not cred_file_path:
+            try:
+                conn = bigquery.Client(project=project_id)
+            except:
+                print("Could not found any google cloud implicit credentials")
+        else:
             # Validate file path and pemissions
             validate_config_path(cred_file_path)
-        else:
-            if not conn:
-                raise ValidationError(
-                    "Pleae provide a service account credentials json file"
-                )
 
         if not conn:
             with open(cred_file_path, "r") as f:
@@ -1279,6 +1274,7 @@ class VannaBase(ABC):
             # Execute the SQL statement and return the result as a pandas DataFrame
             with engine.begin() as conn:
                 df = pd.read_sql_query(sa.text(sql), conn)
+                conn.close()
                 return df
 
             raise Exception("Couldn't run sql")
@@ -1286,6 +1282,116 @@ class VannaBase(ABC):
         self.dialect = "T-SQL / Microsoft SQL Server"
         self.run_sql = run_sql_mssql
         self.run_sql_is_set = True
+    def connect_to_presto(
+      self,
+      host: str,
+      catalog: str = 'hive',
+      schema: str = 'default',
+      user: str = None,
+      password: str = None,
+      port: int = None,
+      combined_pem_path: str = None,
+      protocol: str = 'https',
+      requests_kwargs: dict = None
+    ):
+      """
+        Connect to a Presto database using the specified parameters.
+
+        Args:
+            host (str): The host address of the Presto database.
+            catalog (str): The catalog to use in the Presto environment.
+            schema (str): The schema to use in the Presto environment.
+            user (str): The username for authentication.
+            password (str): The password for authentication.
+            port (int): The port number for the Presto connection.
+            combined_pem_path (str): The path to the combined pem file for SSL connection.
+            protocol (str): The protocol to use for the connection (default is 'https').
+            requests_kwargs (dict): Additional keyword arguments for requests.
+
+        Raises:
+            DependencyError: If required dependencies are not installed.
+            ImproperlyConfigured: If essential configuration settings are missing.
+
+        Returns:
+            None
+      """
+      try:
+        from pyhive import presto
+      except ImportError:
+        raise DependencyError(
+          "You need to install required dependencies to execute this method,"
+          " run command: \npip install pyhive"
+        )
+
+      if not host:
+        host = os.getenv("PRESTO_HOST")
+
+      if not host:
+        raise ImproperlyConfigured("Please set your presto host")
+
+      if not catalog:
+        catalog = os.getenv("PRESTO_CATALOG")
+
+      if not catalog:
+        raise ImproperlyConfigured("Please set your presto catalog")
+
+      if not user:
+        user = os.getenv("PRESTO_USER")
+
+      if not user:
+        raise ImproperlyConfigured("Please set your presto user")
+
+      if not password:
+        password = os.getenv("PRESTO_PASSWORD")
+
+      if not port:
+        port = os.getenv("PRESTO_PORT")
+
+      if not port:
+        raise ImproperlyConfigured("Please set your presto port")
+
+      conn = None
+
+      try:
+        if requests_kwargs is None and combined_pem_path is not None:
+          # use the combined pem file to verify the SSL connection
+          requests_kwargs = {
+            'verify': combined_pem_path,  # 使用转换后得到的 PEM 文件进行 SSL 验证
+          }
+        conn = presto.Connection(host=host,
+                                 username=user,
+                                 password=password,
+                                 catalog=catalog,
+                                 schema=schema,
+                                 port=port,
+                                 protocol=protocol,
+                                 requests_kwargs=requests_kwargs)
+      except presto.Error as e:
+        raise ValidationError(e)
+
+      def run_sql_presto(sql: str) -> Union[pd.DataFrame, None]:
+        if conn:
+          try:
+            cs = conn.cursor()
+            cs.execute(sql)
+            results = cs.fetchall()
+
+            # Create a pandas dataframe from the results
+            df = pd.DataFrame(
+              results, columns=[desc[0] for desc in cs.description]
+            )
+            return df
+
+          except presto.Error as e:
+            print(e)
+            raise ValidationError(e)
+
+          except Exception as e:
+            print(e)
+            raise e
+
+      self.run_sql_is_set = True
+      self.run_sql = run_sql_presto
 
     def run_sql(self, sql: str, **kwargs) -> pd.DataFrame:
         """

{vanna-0.5.0 → vanna-0.5.2}/src/vanna/opensearch/opensearch_vector.py

@@ -24,7 +24,77 @@ class OpenSearch_VectorStore(VannaBase):
     self.document_index = document_index
     self.ddl_index = ddl_index
     self.question_sql_index = question_sql_index
-    print("OpenSearch_VectorStore initialized with document_index: ", document_index, " ddl_index: ", ddl_index, " question_sql_index: ", question_sql_index)
+    print("OpenSearch_VectorStore initialized with document_index: ",
+          document_index, " ddl_index: ", ddl_index, " question_sql_index: ",
+          question_sql_index)
+
+    document_index_settings = {
+      "settings": {
+        "index": {
+          "number_of_shards": 6,
+          "number_of_replicas": 2
+        }
+      },
+      "mappings": {
+        "properties": {
+          "question": {
+            "type": "text",
+          },
+          "doc": {
+            "type": "text",
+          }
+        }
+      }
+    }
+
+    ddl_index_settings = {
+      "settings": {
+        "index": {
+          "number_of_shards": 6,
+          "number_of_replicas": 2
+        }
+      },
+      "mappings": {
+        "properties": {
+          "ddl": {
+            "type": "text",
+          },
+          "doc": {
+            "type": "text",
+          }
+        }
+      }
+    }
+
+    question_sql_index_settings = {
+      "settings": {
+        "index": {
+          "number_of_shards": 6,
+          "number_of_replicas": 2
+        }
+      },
+      "mappings": {
+        "properties": {
+          "question": {
+            "type": "text",
+          },
+          "sql": {
+            "type": "text",
+          }
+        }
+      }
+    }
+
+    if config is not None and "es_document_index_settings" in config:
+      document_index_settings = config["es_document_index_settings"]
+    if config is not None and "es_ddl_index_settings" in config:
+      ddl_index_settings = config["es_ddl_index_settings"]
+    if config is not None and "es_question_sql_index_settings" in config:
+      question_sql_index_settings = config["es_question_sql_index_settings"]
+
+    self.document_index_settings = document_index_settings
+    self.ddl_index_settings = ddl_index_settings
+    self.question_sql_index_settings = question_sql_index_settings
 
     es_urls = None
     if config is not None and "es_urls" in config:
@@ -85,6 +155,9 @@ class OpenSearch_VectorStore(VannaBase):
     else:
       max_retries = 10
 
+    print("OpenSearch_VectorStore initialized with es_urls: ", es_urls,
+          " host: ", host, " port: ", port, " ssl: ", ssl, " verify_certs: ",
+          verify_certs, " timeout: ", timeout, " max_retries: ", max_retries)
     if es_urls is not None:
       # Initialize the OpenSearch client by passing a list of URLs
       self.client = OpenSearch(
@@ -112,18 +185,26 @@ class OpenSearch_VectorStore(VannaBase):
         headers=headers
       )
 
+      print("OpenSearch_VectorStore initialized with client over ")
+
     # 执行一个简单的查询来检查连接
     try:
+      print('Connected to OpenSearch cluster:')
       info = self.client.info()
       print('OpenSearch cluster info:', info)
     except Exception as e:
       print('Error connecting to OpenSearch cluster:', e)
 
     # Create the indices if they don't exist
-    # self.create_index()
+    self.create_index_if_not_exists(self.document_index,
+                                    self.document_index_settings)
+    self.create_index_if_not_exists(self.ddl_index, self.ddl_index_settings)
+    self.create_index_if_not_exists(self.question_sql_index,
+                                    self.question_sql_index_settings)
 
   def create_index(self):
-    for index in [self.document_index, self.ddl_index, self.question_sql_index]:
+    for index in [self.document_index, self.ddl_index,
+                  self.question_sql_index]:
       try:
         self.client.indices.create(index)
       except Exception as e:
@@ -131,6 +212,20 @@ class OpenSearch_VectorStore(VannaBase):
         print(f"opensearch index {index} already exists")
         pass
 
+  def create_index_if_not_exists(self, index_name: str,
+                                 index_settings: dict) -> bool:
+    try:
+      if not self.client.indices.exists(index_name):
+        print(f"Index {index_name} does not exist. Creating...")
+        self.client.indices.create(index=index_name, body=index_settings)
+        return True
+      else:
+        print(f"Index {index_name} already exists.")
+        return False
+    except Exception as e:
+      print(f"Error creating index: {index_name} ", e)
+      return False
+
   def add_ddl(self, ddl: str, **kwargs) -> str:
     # Assuming that you have a DDL index in your OpenSearch
     id = str(uuid.uuid4()) + "-ddl"
@@ -278,7 +373,6 @@ class OpenSearch_VectorStore(VannaBase):
     # opensearch doesn't need to generate embeddings
     pass
 
-
 # OpenSearch_VectorStore.__init__(self, config={'es_urls':
 # "https://opensearch-node.test.com:9200", 'es_encoded_base64': True, 'es_user':
 # "admin", 'es_password': "admin", 'es_verify_certs': True})

{vanna-0.5.0 → vanna-0.5.2}/src/vanna/qdrant/qdrant.py

@@ -7,48 +7,51 @@ from qdrant_client import QdrantClient, grpc, models
 from ..base import VannaBase
 from ..utils import deterministic_uuid
 
-DOCUMENTATION_COLLECTION_NAME = "documentation"
-DDL_COLLECTION_NAME = "ddl"
-SQL_COLLECTION_NAME = "sql"
 SCROLL_SIZE = 1000
 
-ID_SUFFIXES = {
-    DDL_COLLECTION_NAME: "ddl",
-    DOCUMENTATION_COLLECTION_NAME: "doc",
-    SQL_COLLECTION_NAME: "sql",
-}
-
 
 class Qdrant_VectorStore(VannaBase):
-    """Vectorstore implementation using Qdrant - https://qdrant.tech/"""
+    """
+    Vectorstore implementation using Qdrant - https://qdrant.tech/
+
+    Args:
+        - config (dict, optional): Dictionary of `Qdrant_VectorStore config` options. Defaults to `{}`.
+            - client: A `qdrant_client.QdrantClient` instance. Overrides other config options.
+            - location: If `":memory:"` - use in-memory Qdrant instance. If `str` - use it as a `url` parameter.
+            - url: Either host or str of "Optional[scheme], host, Optional[port], Optional[prefix]". Eg. `"http://localhost:6333"`.
+            - prefer_grpc: If `true` - use gPRC interface whenever possible in custom methods.
+            - https: If `true` - use HTTPS(SSL) protocol. Default: `None`
+            - api_key: API key for authentication in Qdrant Cloud. Default: `None`
+            - timeout: Timeout for REST and gRPC API requests. Defaults to 5 seconds for REST and unlimited for gRPC.
+            - path: Persistence path for QdrantLocal. Default: `None`.
+            - prefix: Prefix to the REST URL paths. Example: `service/v1` will result in `http://localhost:6333/service/v1/{qdrant-endpoint}`.
+            - n_results: Number of results to return from similarity search. Defaults to 10.
+            - fastembed_model: [Model](https://qdrant.github.io/fastembed/examples/Supported_Models/#supported-text-embedding-models) to use for `fastembed.TextEmbedding`.
+              Defaults to `"BAAI/bge-small-en-v1.5"`.
+            - collection_params: Additional parameters to pass to `qdrant_client.QdrantClient#create_collection()` method.
+            - distance_metric: Distance metric to use when creating collections. Defaults to `qdrant_client.models.Distance.COSINE`.
+            - documentation_collection_name: Name of the collection to store documentation. Defaults to `"documentation"`.
+            - ddl_collection_name: Name of the collection to store DDL. Defaults to `"ddl"`.
+            - sql_collection_name: Name of the collection to store SQL. Defaults to `"sql"`.
+
+    Raises:
+        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={},
     ):
-        """
-        Vectorstore implementation using Qdrant - https://qdrant.tech/
-
-        Args:
-            - config (dict, optional): Dictionary of `Qdrant_VectorStore config` options. Defaults to `{}`.
-                - client: A `qdrant_client.QdrantClient` instance. Overrides other config options.
-                - location: If `":memory:"` - use in-memory Qdrant instance. If `str` - use it as a `url` parameter.
-                - url: Either host or str of "Optional[scheme], host, Optional[port], Optional[prefix]". Eg. `"http://localhost:6333"`.
-                - prefer_grpc: If `true` - use gPRC interface whenever possible in custom methods.
-                - https: If `true` - use HTTPS(SSL) protocol. Default: `None`
-                - api_key: API key for authentication in Qdrant Cloud. Default: `None`
-                - timeout: Timeout for REST and gRPC API requests. Defaults to 5 seconds for REST and unlimited for gRPC.
-                - path: Persistence path for QdrantLocal. Default: `None`.
-                - prefix: Prefix to the REST URL paths. Example: `service/v1` will result in `http://localhost:6333/service/v1/{qdrant-endpoint}`.
-                - n_results: Number of results to return from similarity search. Defaults to 10.
-                - fastembed_model: [Model](https://qdrant.github.io/fastembed/examples/Supported_Models/#supported-text-embedding-models) to use for `fastembed.TextEmbedding`.
-                  Defaults to `"BAAI/bge-small-en-v1.5"`.
-                - collection_params: Additional parameters to pass to `qdrant_client.QdrantClient#create_collection()` method.
-                - distance_metric: Distance metric to use when creating collections. Defaults to `qdrant_client.models.Distance.COSINE`.
-
-        Raises:
-            TypeError: If config["client"] is not a `qdrant_client.QdrantClient` instance
-        """
         VannaBase.__init__(self, config=config)
         client = config.get("client")
 
@@ -75,15 +78,24 @@ class Qdrant_VectorStore(VannaBase):
         self.fastembed_model = config.get("fastembed_model", "BAAI/bge-small-en-v1.5")
         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
+        )
+        self.ddl_collection_name = config.get(
+            "ddl_collection_name", self.ddl_collection_name
+        )
+        self.sql_collection_name = config.get(
+            "sql_collection_name", self.sql_collection_name
+        )
 
         self._setup_collections()
 
     def add_question_sql(self, question: str, sql: str, **kwargs) -> str:
-        question_answer = format("Question: {0}\n\nSQL: {1}", question, sql)
+        question_answer = "Question: {0}\n\nSQL: {1}".format(question, sql)
         id = deterministic_uuid(question_answer)
 
         self._client.upsert(
-            SQL_COLLECTION_NAME,
+            self.sql_collection_name,
             points=[
                 models.PointStruct(
                     id=id,
@@ -96,12 +108,12 @@ class Qdrant_VectorStore(VannaBase):
             ],
         )
 
-        return self._format_point_id(id, SQL_COLLECTION_NAME)
+        return self._format_point_id(id, self.sql_collection_name)
 
     def add_ddl(self, ddl: str, **kwargs) -> str:
         id = deterministic_uuid(ddl)
         self._client.upsert(
-            DDL_COLLECTION_NAME,
+            self.ddl_collection_name,
             points=[
                 models.PointStruct(
                     id=id,
@@ -112,13 +124,13 @@ class Qdrant_VectorStore(VannaBase):
                 )
             ],
         )
-        return self._format_point_id(id, DDL_COLLECTION_NAME)
+        return self._format_point_id(id, self.ddl_collection_name)
 
     def add_documentation(self, documentation: str, **kwargs) -> str:
         id = deterministic_uuid(documentation)
 
         self._client.upsert(
-            DOCUMENTATION_COLLECTION_NAME,
+            self.documentation_collection_name,
             points=[
                 models.PointStruct(
                     id=id,
@@ -130,16 +142,17 @@ class Qdrant_VectorStore(VannaBase):
             ],
         )
 
-        return self._format_point_id(id, DOCUMENTATION_COLLECTION_NAME)
+        return self._format_point_id(id, self.documentation_collection_name)
 
     def get_training_data(self, **kwargs) -> pd.DataFrame:
         df = pd.DataFrame()
 
-        if sql_data := self._get_all_points(SQL_COLLECTION_NAME):
+        if sql_data := self._get_all_points(self.sql_collection_name):
             question_list = [data.payload["question"] for data in sql_data]
             sql_list = [data.payload["sql"] for data in sql_data]
             id_list = [
-                self._format_point_id(data.id, SQL_COLLECTION_NAME) for data in sql_data
+                self._format_point_id(data.id, self.sql_collection_name)
+                for data in sql_data
             ]
 
             df_sql = pd.DataFrame(
@@ -154,10 +167,11 @@ class Qdrant_VectorStore(VannaBase):
 
             df = pd.concat([df, df_sql])
 
-        if ddl_data := self._get_all_points(DDL_COLLECTION_NAME):
+        if ddl_data := self._get_all_points(self.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, self.ddl_collection_name)
+                for data in ddl_data
             ]
 
             df_ddl = pd.DataFrame(
@@ -172,12 +186,10 @@ class Qdrant_VectorStore(VannaBase):
 
             df = pd.concat([df, df_ddl])
 
-        doc_data = self.documentation_collection.get()
-
-        if doc_data := self._get_all_points(DOCUMENTATION_COLLECTION_NAME):
+        if doc_data := self._get_all_points(self.documentation_collection_name):
             document_list = [data.payload["documentation"] for data in doc_data]
             id_list = [
-                self._format_point_id(data.id, DOCUMENTATION_COLLECTION_NAME)
+                self._format_point_id(data.id, self.documentation_collection_name)
                 for data in doc_data
             ]
 
@@ -212,7 +224,7 @@ class Qdrant_VectorStore(VannaBase):
         Returns:
             bool: True if collection is deleted, False otherwise
         """
-        if collection_name in ID_SUFFIXES.keys():
+        if collection_name in self.id_suffixes.keys():
             self._client.delete_collection(collection_name)
             self._setup_collections()
             return True
@@ -225,7 +237,7 @@ class Qdrant_VectorStore(VannaBase):
 
     def get_similar_question_sql(self, question: str, **kwargs) -> list:
         results = self._client.search(
-            SQL_COLLECTION_NAME,
+            self.sql_collection_name,
             query_vector=self.generate_embedding(question),
             limit=self.n_results,
             with_payload=True,
@@ -235,7 +247,7 @@ class Qdrant_VectorStore(VannaBase):
 
     def get_related_ddl(self, question: str, **kwargs) -> list:
         results = self._client.search(
-            DDL_COLLECTION_NAME,
+            self.ddl_collection_name,
             query_vector=self.generate_embedding(question),
             limit=self.n_results,
             with_payload=True,
@@ -245,7 +257,7 @@ class Qdrant_VectorStore(VannaBase):
 
     def get_related_documentation(self, question: str, **kwargs) -> list:
         results = self._client.search(
-            DOCUMENTATION_COLLECTION_NAME,
+            self.documentation_collection_name,
             query_vector=self.generate_embedding(question),
             limit=self.n_results,
             with_payload=True,
@@ -284,9 +296,9 @@ class Qdrant_VectorStore(VannaBase):
         return results
 
     def _setup_collections(self):
-        if not self._client.collection_exists(SQL_COLLECTION_NAME):
+        if not self._client.collection_exists(self.sql_collection_name):
             self._client.create_collection(
-                collection_name=SQL_COLLECTION_NAME,
+                collection_name=self.sql_collection_name,
                 vectors_config=models.VectorParams(
                     size=self.embeddings_dimension,
                     distance=self.distance_metric,
@@ -294,18 +306,18 @@ class Qdrant_VectorStore(VannaBase):
                 **self.collection_params,
             )
 
-        if not self._client.collection_exists(DDL_COLLECTION_NAME):
+        if not self._client.collection_exists(self.ddl_collection_name):
             self._client.create_collection(
-                collection_name=DDL_COLLECTION_NAME,
+                collection_name=self.ddl_collection_name,
                 vectors_config=models.VectorParams(
                     size=self.embeddings_dimension,
                     distance=self.distance_metric,
                 ),
                 **self.collection_params,
             )
-        if not self._client.collection_exists(DOCUMENTATION_COLLECTION_NAME):
+        if not self._client.collection_exists(self.documentation_collection_name):
             self._client.create_collection(
-                collection_name=DOCUMENTATION_COLLECTION_NAME,
+                collection_name=self.documentation_collection_name,
                 vectors_config=models.VectorParams(
                     size=self.embeddings_dimension,
                     distance=self.distance_metric,
@@ -314,11 +326,11 @@ class Qdrant_VectorStore(VannaBase):
             )
 
     def _format_point_id(self, id: str, collection_name: str) -> str:
-        return "{0}-{1}".format(id, ID_SUFFIXES[collection_name])
+        return "{0}-{1}".format(id, self.id_suffixes[collection_name])
 
     def _parse_point_id(self, id: str) -> Tuple[str, str]:
-        id, suffix = id.rsplit("-", 1)
-        for collection_name, suffix in ID_SUFFIXES.items():
-            if type == suffix:
+        id, curr_suffix = id.rsplit("-", 1)
+        for collection_name, suffix in self.id_suffixes.items():
+            if curr_suffix == suffix:
                 return id, collection_name
         raise ValueError(f"Invalid id {id}")