vanna 0.5.5__py3-none-any.whl → 0.6.0__py3-none-any.whl

vanna/qdrant/qdrant.py

@@ -3,6 +3,7 @@ from typing import List, Tuple
 
 import pandas as pd
 from qdrant_client import QdrantClient, grpc, models
+from qdrant_client.http.models.models import UpdateStatus
 
 from ..base import VannaBase
 from ..utils import deterministic_uuid
@@ -210,7 +211,8 @@ class Qdrant_VectorStore(VannaBase):
     def remove_training_data(self, id: str, **kwargs) -> bool:
         try:
             id, collection_name = self._parse_point_id(id)
-            self._client.delete(collection_name, points_selector=[id])
+            res = self._client.delete(collection_name, points_selector=[id])
+            res == UpdateStatus.COMPLETED
         except ValueError:
             return False
 

vanna/vannadb/vannadb_vector.py

@@ -5,6 +5,7 @@ from io import StringIO
 import pandas as pd
 import requests
 
+from ..advanced import VannaAdvanced
 from ..base import VannaBase
 from ..types import (
   DataFrameJSON,
@@ -20,7 +21,7 @@ from ..types import (
 from ..utils import sanitize_model_name
 
 
-class VannaDB_VectorStore(VannaBase):
+class VannaDB_VectorStore(VannaBase, VannaAdvanced):
     def __init__(self, vanna_model: str, vanna_api_key: str, config=None):
         VannaBase.__init__(self, config=config)
 
@@ -33,6 +34,12 @@ class VannaDB_VectorStore(VannaBase):
             else config["endpoint"]
         )
         self.related_training_data = {}
+        self._graphql_endpoint = "https://functionrag.com/query"
+        self._graphql_headers = {
+            "Content-Type": "application/json",
+            "API-KEY": self._api_key,
+            "NAMESPACE": self._model,
+        }
 
     def _rpc_call(self, method, params):
         if method != "list_orgs":
@@ -59,6 +66,177 @@ class VannaDB_VectorStore(VannaBase):
     def _dataclass_to_dict(self, obj):
         return dataclasses.asdict(obj)
 
+    def get_all_functions(self) -> list:
+        query = """
+            {
+                get_all_sql_functions {
+                    function_name
+                    description
+                    post_processing_code_template
+                    arguments {
+                        name
+                        description
+                        general_type
+                        is_user_editable
+                        available_values
+                    }
+                    sql_template
+                }
+            }
+        """
+
+        response = requests.post(self._graphql_endpoint, headers=self._graphql_headers, json={'query': query})
+        response_json = response.json()
+        if response.status_code == 200 and 'data' in response_json and 'get_all_sql_functions' in response_json['data']:
+            self.log(response_json['data']['get_all_sql_functions'])
+            resp = response_json['data']['get_all_sql_functions']
+
+            print(resp)
+
+            return resp
+        else:
+            raise Exception(f"Query failed to run by returning code of {response.status_code}. {response.text}")
+
+
+    def get_function(self, question: str, additional_data: dict = {}) -> dict:
+        query = """
+        query GetFunction($question: String!, $staticFunctionArguments: [StaticFunctionArgument]) {
+            get_and_instantiate_function(question: $question, static_function_arguments: $staticFunctionArguments) {
+                ... on SQLFunction {
+                function_name
+                description
+                post_processing_code_template
+                instantiated_post_processing_code
+                arguments {
+                    name
+                    description
+                    general_type
+                    is_user_editable
+                    instantiated_value
+                    available_values
+                }
+                sql_template
+                instantiated_sql
+            }
+            }
+        }
+        """
+        static_function_arguments = [{"name": key, "value": str(value)} for key, value in additional_data.items()]
+        variables = {"question": question, "staticFunctionArguments": static_function_arguments}
+        response = requests.post(self._graphql_endpoint, headers=self._graphql_headers, json={'query': query, 'variables': variables})
+        response_json = response.json()
+        if response.status_code == 200 and 'data' in response_json and 'get_and_instantiate_function' in response_json['data']:
+            self.log(response_json['data']['get_and_instantiate_function'])
+            resp = response_json['data']['get_and_instantiate_function']
+
+            print(resp)
+
+            return resp
+        else:
+            raise Exception(f"Query failed to run by returning code of {response.status_code}. {response.text}")
+
+    def create_function(self, question: str, sql: str, plotly_code: str, **kwargs) -> dict:
+        query = """
+        mutation CreateFunction($question: String!, $sql: String!, $plotly_code: String!) {
+            generate_and_create_sql_function(question: $question, sql: $sql, post_processing_code: $plotly_code) {
+                function_name
+                description
+                arguments {
+                    name
+                    description
+                    general_type
+                    is_user_editable
+                }
+                sql_template
+                post_processing_code_template
+            }
+        }
+        """
+        variables = {"question": question, "sql": sql, "plotly_code": plotly_code}
+        response = requests.post(self._graphql_endpoint, headers=self._graphql_headers, json={'query': query, 'variables': variables})
+        response_json = response.json()
+        if response.status_code == 200 and 'data' in response_json and response_json['data'] is not None and 'generate_and_create_sql_function' in response_json['data']:
+            resp = response_json['data']['generate_and_create_sql_function']
+
+            print(resp)
+
+            return resp
+        else:
+            raise Exception(f"Query failed to run by returning code of {response.status_code}. {response.text}")
+
+    def update_function(self, old_function_name: str, updated_function: dict) -> bool:
+        """
+        Update an existing SQL function based on the provided parameters.
+
+        Args:
+            old_function_name (str): The current name of the function to be updated.
+            updated_function (dict): A dictionary containing the updated function details. Expected keys:
+                - 'function_name': The new name of the function.
+                - 'description': The new description of the function.
+                - 'arguments': A list of dictionaries describing the function arguments.
+                - 'sql_template': The new SQL template for the function.
+                - 'post_processing_code_template': The new post-processing code template.
+
+        Returns:
+            bool: True if the function was successfully updated, False otherwise.
+        """
+        mutation = """
+        mutation UpdateSQLFunction($input: SQLFunctionUpdate!) {
+            update_sql_function(input: $input)
+        }
+        """
+
+        SQLFunctionUpdate = {
+            'function_name', 'description', 'arguments', 'sql_template', 'post_processing_code_template'
+        }
+
+        # Define the expected keys for each argument in the arguments list
+        ArgumentKeys = {'name', 'general_type', 'description', 'is_user_editable', 'available_values'}
+
+        # Function to validate and transform arguments
+        def validate_arguments(args):
+            return [
+                {key: arg[key] for key in arg if key in ArgumentKeys}
+                for arg in args
+            ]
+
+        # Keep only the keys that conform to the SQLFunctionUpdate GraphQL input type
+        updated_function = {key: value for key, value in updated_function.items() if key in SQLFunctionUpdate}
+
+        # Special handling for 'arguments' to ensure they conform to the spec
+        if 'arguments' in updated_function:
+            updated_function['arguments'] = validate_arguments(updated_function['arguments'])
+
+        variables = {
+            "input": {
+                "old_function_name": old_function_name,
+                **updated_function
+            }
+        }
+
+        print("variables", variables)
+
+        response = requests.post(self._graphql_endpoint, headers=self._graphql_headers, json={'query': mutation, 'variables': variables})
+        response_json = response.json()
+        if response.status_code == 200 and 'data' in response_json and response_json['data'] is not None and 'update_sql_function' in response_json['data']:
+            return response_json['data']['update_sql_function']
+        else:
+            raise Exception(f"Mutation failed to run by returning code of {response.status_code}. {response.text}")
+
+    def delete_function(self, function_name: str) -> bool:
+        mutation = """
+        mutation DeleteSQLFunction($function_name: String!) {
+            delete_sql_function(function_name: $function_name)
+        }
+        """
+        variables = {"function_name": function_name}
+        response = requests.post(self._graphql_endpoint, headers=self._graphql_headers, json={'query': mutation, 'variables': variables})
+        response_json = response.json()
+        if response.status_code == 200 and 'data' in response_json and response_json['data'] is not None and 'delete_sql_function' in response_json['data']:
+            return response_json['data']['delete_sql_function']
+        else:
+            raise Exception(f"Mutation failed to run by returning code of {response.status_code}. {response.text}")
+
     def create_model(self, model: str, **kwargs) -> bool:
         """
         **Example:**

vanna/vllm/vllm.py

@@ -17,6 +17,11 @@ class Vllm(VannaBase):
         else:
             self.model = config["model"]
 
+        if "auth-key" in config:
+            self.auth_key = config["auth-key"]
+        else:
+            self.auth_key = None
+
     def system_message(self, message: str) -> any:
         return {"role": "system", "content": message}
 
@@ -67,7 +72,17 @@ class Vllm(VannaBase):
             "messages": prompt,
         }
 
-        response = requests.post(url, json=data)
+        if self.auth_key is not None:
+            headers = { 
+            'Content-Type': 'application/json', 
+            'Authorization': f'Bearer {self.auth_key}' 
+            }
+
+            response = requests.post(url, headers=headers,json=data)
+
+
+        else:
+            response = requests.post(url, json=data)
 
         response_dict = response.json()
 

{vanna-0.5.5.dist-info → vanna-0.6.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: vanna
-Version: 0.5.5
+Version: 0.6.0
 Summary: Generate SQL queries from natural language
 Author-email: Zain Hoda <zain@vanna.ai>
 Requires-Python: >=3.9
@@ -42,7 +42,7 @@ Requires-Dist: pinecone-client ; extra == "all"
 Requires-Dist: anthropic ; extra == "anthropic"
 Requires-Dist: google-cloud-bigquery ; extra == "bigquery"
 Requires-Dist: chromadb ; extra == "chromadb"
-Requires-Dist: clickhouse_driver ; extra == "clickhouse"
+Requires-Dist: clickhouse_connect ; extra == "clickhouse"
 Requires-Dist: duckdb ; extra == "duckdb"
 Requires-Dist: google-generativeai ; extra == "gemini"
 Requires-Dist: google-generativeai ; extra == "google"

{vanna-0.5.5.dist-info → vanna-0.6.0.dist-info}/RECORD

@@ -5,15 +5,16 @@ vanna/utils.py,sha256=cs0B_0MwhmPI2nWjVHifDYCmCR0kkddylQ2vloaPDSw,2247
 vanna/ZhipuAI/ZhipuAI_Chat.py,sha256=WtZKUBIwlNH0BGbb4lZbVR7pTWIrn7b4RLIk-7u0SuQ,8725
 vanna/ZhipuAI/ZhipuAI_embeddings.py,sha256=lUqzJg9fOx7rVFhjdkFjXcDeVGV4aAB5Ss0oERsa8pE,2849
 vanna/ZhipuAI/__init__.py,sha256=NlsijtcZp5Tj9jkOe9fNcOQND_QsGgu7otODsCLBPr0,116
+vanna/advanced/__init__.py,sha256=oDj9g1JbrbCfp4WWdlr_bhgdMqNleyHgr6VXX6DcEbo,658
 vanna/anthropic/__init__.py,sha256=85s_2mAyyPxc0T_0JEvYeAkEKWJwkwqoyUwSC5dw9Gk,43
 vanna/anthropic/anthropic_chat.py,sha256=Wk0o-NMW1uvR2fhSWxrR_2FqNh-dLprNG4uuVqpqAkY,2615
 vanna/base/__init__.py,sha256=Sl-HM1RRYzAZoSqmL1CZQmF3ZF-byYTCFQP3JZ2A5MU,28
-vanna/base/base.py,sha256=YCL9MhhrGeoVv9da85NdWvEQtnqfkWSyj5AZo_wQ0TU,70853
+vanna/base/base.py,sha256=XiBJU4LOKoFnY1vnkr_jz3soAdVdoL0D1WIZdzZCVyU,70942
 vanna/chromadb/__init__.py,sha256=-iL0nW_g4uM8nWKMuWnNePfN4nb9uk8P3WzGvezOqRg,50
 vanna/chromadb/chromadb_vector.py,sha256=eKyPck99Y6Jt-BNWojvxLG-zvAERzLSm-3zY-bKXvaA,8792
 vanna/exceptions/__init__.py,sha256=dJ65xxxZh1lqBeg6nz6Tq_r34jLVmjvBvPO9Q6hFaQ8,685
-vanna/flask/__init__.py,sha256=2t9DbgL2-Anm05ESmZ1WAYMd1VSkOUsTf_-exVCebwQ,25192
-vanna/flask/assets.py,sha256=7wMnTS-UPZkNL8SSTrd4U0amuMgeOIIj94WGzBQW1fw,187711
+vanna/flask/__init__.py,sha256=wpAZg6emB7l8TmyMlxVLFMjxkUQReIIPTfqlkF-onlc,30267
+vanna/flask/assets.py,sha256=_UoUr57sS0QL2BuTxAOe9k4yy8T7-fp2NpbRSVtW3IM,451769
 vanna/flask/auth.py,sha256=UpKxh7W5cd43W0LGch0VqhncKwB78L6dtOQkl1JY5T0,1246
 vanna/google/__init__.py,sha256=M-dCxCZcKL4bTQyMLj6r6VRs65YNX9Tl2aoPCuqGm-8,41
 vanna/google/gemini_chat.py,sha256=ps3A-afFbCo3HeFTLL_nMoQO1PsGvRUUPRUppbMcDew,1584
@@ -37,12 +38,12 @@ vanna/opensearch/opensearch_vector.py,sha256=VhIcrSyNzWR9ZrqrJnyGFOyuQZs3swfbhr8
 vanna/pinecone/__init__.py,sha256=eO5l8aX8vKL6aIUMgAXGPt1jdqKxB_Hic6cmoVAUrD0,90
 vanna/pinecone/pinecone_vector.py,sha256=mpq1lzo3KRj2QfJEw8pwFclFQK1Oi_Nx-lDkx9Gp0mw,11448
 vanna/qdrant/__init__.py,sha256=PX_OsDOiPMvwCJ2iGER1drSdQ9AyM8iN5PEBhRb6qqY,73
-vanna/qdrant/qdrant.py,sha256=6M00nMiuOuftTDf3NsOrOcG9BA4DlIIDck2MNp9iEyg,12613
+vanna/qdrant/qdrant.py,sha256=RyICUvOO_jt8u9MB4oIYhqv3BicZ0d9pQkSFwkIfUjg,12719
 vanna/types/__init__.py,sha256=Qhn_YscKtJh7mFPCyCDLa2K8a4ORLMGVnPpTbv9uB2U,4957
 vanna/vannadb/__init__.py,sha256=C6UkYocmO6dmzfPKZaWojN0mI5YlZZ9VIbdcquBE58A,48
-vanna/vannadb/vannadb_vector.py,sha256=9YwTO3Lh5owWQE7KPMBqLp2EkiGV0RC1sEYhslzJzgI,6168
+vanna/vannadb/vannadb_vector.py,sha256=N8poMYvAojoaOF5gI4STD5pZWK9lBKPvyIjbh9dPBa0,14189
 vanna/vllm/__init__.py,sha256=aNlUkF9tbURdeXAJ8ytuaaF1gYwcG3ny1MfNl_cwQYg,23
-vanna/vllm/vllm.py,sha256=QerC3xF5eNzE_nGBDl6YrPYF4WYnjf0hHxxlDWdKX-0,2427
-vanna-0.5.5.dist-info/WHEEL,sha256=EZbGkh7Ie4PoZfRQ8I0ZuP9VklN_TvcZ6DSE5Uar4z4,81
-vanna-0.5.5.dist-info/METADATA,sha256=gmPrDKsawOtYazNHG5CQ_iURmdVLt1Jc8Rfwneue1w0,11505
-vanna-0.5.5.dist-info/RECORD,,
+vanna/vllm/vllm.py,sha256=oM_aA-1Chyl7T_Qc_yRKlL6oSX1etsijY9zQdjeMGMQ,2827
+vanna-0.6.0.dist-info/WHEEL,sha256=EZbGkh7Ie4PoZfRQ8I0ZuP9VklN_TvcZ6DSE5Uar4z4,81
+vanna-0.6.0.dist-info/METADATA,sha256=Nu63PMm3HT3HAkpxKU4fte9CKpIDgnL90Ww2AljWSOs,11506
+vanna-0.6.0.dist-info/RECORD,,