vanna 0.5.1__tar.gz → 0.5.3__tar.gz

{vanna-0.5.1 → vanna-0.5.3}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: vanna
-Version: 0.5.1
+Version: 0.5.3
 Summary: Generate SQL queries from natural language
 Author-email: Zain Hoda <zain@vanna.ai>
 Requires-Python: >=3.9

{vanna-0.5.1 → vanna-0.5.3}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "flit_core.buildapi"
 
 [project]
 name = "vanna"
-version = "0.5.1"
+version = "0.5.3"
 authors = [
   { name="Zain Hoda", email="zain@vanna.ai" },
 ]

{vanna-0.5.1 → vanna-0.5.3}/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.1 → vanna-0.5.3}/src/vanna/ollama/ollama.py

@@ -24,7 +24,7 @@ class Ollama(VannaBase):
       raise ValueError("config must contain at least Ollama model")
     self.host = config.get("ollama_host", "http://localhost:11434")
     self.model = config["model"]
-    if ":" in self.model:
+    if ":" not in self.model:
       self.model += ":latest"
 
     self.ollama_client = ollama.Client(self.host, timeout=Timeout(240.0))

{vanna-0.5.1 → vanna-0.5.3}/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.1 → vanna-0.5.3}/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,6 +78,15 @@ 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()
 
@@ -83,7 +95,7 @@ class Qdrant_VectorStore(VannaBase):
         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 ddl_data
+                self._format_point_id(data.id, self.ddl_collection_name)
+                for data in ddl_data
             ]
 
             df_ddl = pd.DataFrame(
@@ -172,10 +186,10 @@ class Qdrant_VectorStore(VannaBase):
 
             df = pd.concat([df, df_ddl])
 
-        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
             ]
 
@@ -210,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
@@ -223,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,
@@ -233,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,
@@ -243,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,
@@ -282,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,
@@ -292,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,
@@ -312,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}")