vanna 0.6.2__py3-none-any.whl → 0.6.4__py3-none-any.whl

vanna/google/gemini_chat.py

@@ -7,7 +7,7 @@ class GoogleGeminiChat(VannaBase):
         VannaBase.__init__(self, config=config)
 
         # default temperature - can be overrided using config
-        self.temperature = 0.7      
+        self.temperature = 0.7
 
         if "temperature" in config:
             self.temperature = config["temperature"]
@@ -31,7 +31,7 @@ class GoogleGeminiChat(VannaBase):
         else:
             # Authenticate using VertexAI
             from vertexai.preview.generative_models import GenerativeModel
-            self.chat_model = GenerativeModel("gemini-pro")
+            self.chat_model = GenerativeModel(model_name)
 
     def system_message(self, message: str) -> any:
         return message

vanna/hf/hf.py

@@ -6,13 +6,15 @@ from ..base import VannaBase
 
 class Hf(VannaBase):
     def __init__(self, config=None):
-        model_name = self.config.get(
-            "model_name", None
-        )  # e.g. meta-llama/Meta-Llama-3-8B-Instruct
-        self.tokenizer = AutoTokenizer.from_pretrained(model_name)
+        model_name_or_path = self.config.get(
+            "model_name_or_path", None
+        )  # e.g. meta-llama/Meta-Llama-3-8B-Instruct or local path to the model checkpoint files
+        # list of quantization methods supported by transformers package: https://huggingface.co/docs/transformers/main/en/quantization/overview
+        quantization_config = self.config.get("quantization_config", None)
+        self.tokenizer = AutoTokenizer.from_pretrained(model_name_or_path)
         self.model = AutoModelForCausalLM.from_pretrained(
-            model_name,
-            torch_dtype="auto",
+            model_name_or_path,
+            quantization_config=quantization_config,
             device_map="auto",
         )
 

vanna/openai/openai_chat.py

@@ -11,14 +11,10 @@ class OpenAI_Chat(VannaBase):
 
         # default parameters - can be overrided using config
         self.temperature = 0.7
-        self.max_tokens = 500
 
         if "temperature" in config:
             self.temperature = config["temperature"]
 
-        if "max_tokens" in config:
-            self.max_tokens = config["max_tokens"]
-
         if "api_type" in config:
             raise Exception(
                 "Passing api_type is now deprecated. Please pass an OpenAI client instead."
@@ -75,7 +71,6 @@ class OpenAI_Chat(VannaBase):
             response = self.client.chat.completions.create(
                 model=model,
                 messages=prompt,
-                max_tokens=self.max_tokens,
                 stop=None,
                 temperature=self.temperature,
             )
@@ -87,7 +82,6 @@ class OpenAI_Chat(VannaBase):
             response = self.client.chat.completions.create(
                 engine=engine,
                 messages=prompt,
-                max_tokens=self.max_tokens,
                 stop=None,
                 temperature=self.temperature,
             )
@@ -98,7 +92,6 @@ class OpenAI_Chat(VannaBase):
             response = self.client.chat.completions.create(
                 engine=self.config["engine"],
                 messages=prompt,
-                max_tokens=self.max_tokens,
                 stop=None,
                 temperature=self.temperature,
             )
@@ -109,7 +102,6 @@ class OpenAI_Chat(VannaBase):
             response = self.client.chat.completions.create(
                 model=self.config["model"],
                 messages=prompt,
-                max_tokens=self.max_tokens,
                 stop=None,
                 temperature=self.temperature,
             )
@@ -123,7 +115,6 @@ class OpenAI_Chat(VannaBase):
             response = self.client.chat.completions.create(
                 model=model,
                 messages=prompt,
-                max_tokens=self.max_tokens,
                 stop=None,
                 temperature=self.temperature,
             )

vanna/weaviate/__init__.py

@@ -0,0 +1 @@
+from .weaviate_vector import WeaviateDatabase

vanna/weaviate/weaviate_vector.py

@@ -0,0 +1,174 @@
+import weaviate
+import weaviate.classes as wvc
+from fastembed import TextEmbedding
+
+from vanna.base import VannaBase
+
+
+class WeaviateDatabase(VannaBase):
+
+    def __init__(self, config=None):
+        """
+        Initialize the VannaEnhanced class with the provided configuration.
+
+        :param config: Dictionary containing configuration parameters.
+
+        params:
+        weaviate_url (str): Weaviate cluster URL while using weaviate cloud,
+        weaviate_api_key (str): Weaviate API key while using weaviate cloud,
+        weaviate_port (num): Weaviate port while using local weaviate,
+        weaviate_grpc (num): Weaviate gRPC port while using local weaviate,
+        fastembed_model (str): Fastembed model name for text embeddings. BAAI/bge-small-en-v1.5 by default.
+
+        """
+        super().__init__(config=config)
+
+        if config is None:
+            raise ValueError("config is required")
+
+        self.fastembed_model = config.get("fastembed_model", "BAAI/bge-small-en-v1.5")
+        self.weaviate_api_key = config.get("weaviate_api_key")
+        self.weaviate_url = config.get("weaviate_url")
+        self.weaviate_port = config.get("weaviate_port")
+        self.weaviate_grpc_port = config.get("weaviate_grpc", 50051)
+
+        if not self.weaviate_api_key and not self.weaviate_port:
+            raise ValueError("Add proper credentials to connect to weaviate")
+
+        self.weaviate_client = self._initialize_weaviate_client()
+        self.embeddings = TextEmbedding(model_name=self.fastembed_model)
+
+        self.training_data_cluster = {
+            "sql": "SQLTrainingDataEntry",
+            "ddl": "DDLEntry",
+            "doc": "DocumentationEntry"
+        }
+
+        self._create_collections_if_not_exist()
+
+    def _create_collections_if_not_exist(self):
+        properties_dict = {
+            self.training_data_cluster['ddl']: [
+                wvc.config.Property(name="description", data_type=wvc.config.DataType.TEXT),
+            ],
+            self.training_data_cluster['doc']: [
+                wvc.config.Property(name="description", data_type=wvc.config.DataType.TEXT),
+            ],
+            self.training_data_cluster['sql']: [
+                wvc.config.Property(name="sql", data_type=wvc.config.DataType.TEXT),
+                wvc.config.Property(name="natural_language_question", data_type=wvc.config.DataType.TEXT),
+            ]
+        }
+
+        for cluster, properties in properties_dict.items():
+            if not self.weaviate_client.collections.exists(cluster):
+                self.weaviate_client.collections.create(
+                    name=cluster,
+                    properties=properties
+                )
+
+    def _initialize_weaviate_client(self):
+        if self.weaviate_api_key:
+            return weaviate.connect_to_wcs(
+                cluster_url=self.weaviate_url,
+                auth_credentials=weaviate.auth.AuthApiKey(self.weaviate_api_key),
+                additional_config=weaviate.config.AdditionalConfig(timeout=(10, 300)),
+                skip_init_checks=True
+            )
+        else:
+            return weaviate.connect_to_local(
+                port=self.weaviate_port,
+                grpc_port=self.weaviate_grpc_port,
+                additional_config=weaviate.config.AdditionalConfig(timeout=(10, 300)),
+                skip_init_checks=True
+            )
+
+    def generate_embedding(self, data: str, **kwargs):
+            embedding_model = TextEmbedding(model_name=self.fastembed_model)
+            embedding = next(embedding_model.embed(data))
+            return embedding.tolist()
+
+
+    def _insert_data(self, cluster_key: str, data_object: dict, vector: list) -> str:
+        self.weaviate_client.connect()
+        response = self.weaviate_client.collections.get(self.training_data_cluster[cluster_key]).data.insert(
+            properties=data_object,
+            vector=vector
+        )
+        self.weaviate_client.close()
+        return response
+
+    def add_ddl(self, ddl: str, **kwargs) -> str:
+        data_object = {
+            "description": ddl,
+        }
+        response = self._insert_data('ddl', data_object, self.generate_embedding(ddl))
+        return f'{response}-ddl'
+
+    def add_documentation(self, doc: str, **kwargs) -> str:
+        data_object = {
+            "description": doc,
+        }
+        response = self._insert_data('doc', data_object, self.generate_embedding(doc))
+        return f'{response}-doc'
+
+    def add_question_sql(self, question: str, sql: str, **kwargs) -> str:
+        data_object = {
+            "sql": sql,
+            "natural_language_question": question,
+        }
+        response = self._insert_data('sql', data_object, self.generate_embedding(question))
+        return f'{response}-sql'
+
+    def _query_collection(self, cluster_key: str, vector_input: list, return_properties: list, limit: int = 3) -> list:
+        self.weaviate_client.connect()
+        collection = self.weaviate_client.collections.get(self.training_data_cluster[cluster_key])
+        response = collection.query.near_vector(
+            near_vector=vector_input,
+            limit=limit,
+            return_properties=return_properties
+        )
+        response_list = [item.properties for item in response.objects]
+        self.weaviate_client.close()
+        return response_list
+
+    def get_related_ddl(self, question: str, **kwargs) -> list:
+        vector_input = self.generate_embedding(question)
+        response_list = self._query_collection('ddl', vector_input, ["description"])
+        return [item["description"] for item in response_list]
+
+    def get_related_documentation(self, question: str, **kwargs) -> list:
+        vector_input = self.generate_embedding(question)
+        response_list = self._query_collection('doc', vector_input, ["description"])
+        return [item["description"] for item in response_list]
+
+    def get_similar_question_sql(self, question: str, **kwargs) -> list:
+        vector_input = self.generate_embedding(question)
+        response_list = self._query_collection('sql', vector_input, ["sql", "natural_language_question"])
+        return [{"question": item["natural_language_question"], "sql": item["sql"]} for item in response_list]
+
+    def get_training_data(self, **kwargs) -> list:
+        self.weaviate_client.connect()
+        combined_response_list = []
+        for collection_name in self.training_data_cluster.values():
+            if self.weaviate_client.collections.exists(collection_name):
+                collection = self.weaviate_client.collections.get(collection_name)
+                response_list = [item.properties for item in collection.iterator()]
+                combined_response_list.extend(response_list)
+        self.weaviate_client.close()
+        return combined_response_list
+
+    def remove_training_data(self, id: str, **kwargs) -> bool:
+        self.weaviate_client.connect()
+        success = False
+        if id.endswith("-sql"):
+            id = id.replace('-sql', '')
+            success = self.weaviate_client.collections.get(self.training_data_cluster['sql']).data.delete_by_id(id)
+        elif id.endswith("-ddl"):
+            id = id.replace('-ddl', '')
+            success = self.weaviate_client.collections.get(self.training_data_cluster['ddl']).data.delete_by_id(id)
+        elif id.endswith("-doc"):
+            id = id.replace('-doc', '')
+            success = self.weaviate_client.collections.get(self.training_data_cluster['doc']).data.delete_by_id(id)
+        self.weaviate_client.close()
+        return success

{vanna-0.6.2.dist-info → vanna-0.6.4.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: vanna
-Version: 0.6.2
+Version: 0.6.4
 Summary: Generate SQL queries from natural language
 Author-email: Zain Hoda <zain@vanna.ai>
 Requires-Python: >=3.9
@@ -16,6 +16,7 @@ Requires-Dist: sqlparse
 Requires-Dist: kaleido
 Requires-Dist: flask
 Requires-Dist: flask-sock
+Requires-Dist: flasgger
 Requires-Dist: sqlalchemy
 Requires-Dist: psycopg2-binary ; extra == "all"
 Requires-Dist: db-dtypes ; extra == "all"
@@ -40,7 +41,10 @@ Requires-Dist: opensearch-dsl ; extra == "all"
 Requires-Dist: transformers ; extra == "all"
 Requires-Dist: pinecone-client ; extra == "all"
 Requires-Dist: pymilvus[model] ; extra == "all"
+Requires-Dist: weaviate-client ; extra == "all"
 Requires-Dist: anthropic ; extra == "anthropic"
+Requires-Dist: boto3 ; extra == "bedrock"
+Requires-Dist: botocore ; extra == "bedrock"
 Requires-Dist: google-cloud-bigquery ; extra == "bigquery"
 Requires-Dist: chromadb ; extra == "chromadb"
 Requires-Dist: clickhouse_connect ; extra == "clickhouse"
@@ -67,11 +71,13 @@ Requires-Dist: fastembed ; extra == "qdrant"
 Requires-Dist: snowflake-connector-python ; extra == "snowflake"
 Requires-Dist: tox ; extra == "test"
 Requires-Dist: vllm ; extra == "vllm"
+Requires-Dist: weaviate-client ; extra == "weaviate"
 Requires-Dist: zhipuai ; extra == "zhipuai"
 Project-URL: Bug Tracker, https://github.com/vanna-ai/vanna/issues
 Project-URL: Homepage, https://github.com/vanna-ai/vanna
 Provides-Extra: all
 Provides-Extra: anthropic
+Provides-Extra: bedrock
 Provides-Extra: bigquery
 Provides-Extra: chromadb
 Provides-Extra: clickhouse
@@ -92,6 +98,7 @@ Provides-Extra: qdrant
 Provides-Extra: snowflake
 Provides-Extra: test
 Provides-Extra: vllm
+Provides-Extra: weaviate
 Provides-Extra: zhipuai
 
 

{vanna-0.6.2.dist-info → vanna-0.6.4.dist-info}/RECORD

@@ -7,19 +7,21 @@ vanna/ZhipuAI/ZhipuAI_embeddings.py,sha256=lUqzJg9fOx7rVFhjdkFjXcDeVGV4aAB5Ss0oE
 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/anthropic/anthropic_chat.py,sha256=7X3x8SYwDY28aGyBnt0YNRMG8YY1p_t-foMfKGj8_Oo,2627
 vanna/base/__init__.py,sha256=Sl-HM1RRYzAZoSqmL1CZQmF3ZF-byYTCFQP3JZ2A5MU,28
-vanna/base/base.py,sha256=l1H0TKsK9DN3n5XgDkUckdLois4dTCAUwrVsRa_6SlQ,70988
+vanna/base/base.py,sha256=3Du70NrXQMn_LOif2YFPRRVKo4wH5-f6eZcLlXEX0X8,71705
+vanna/bedrock/__init__.py,sha256=hRT2bgJbHEqViLdL-t9hfjSfFdIOkPU2ADBt-B1En-8,46
+vanna/bedrock/bedrock_converse.py,sha256=Nx5kYm-diAfYmsWAnTP5xnv7V84Og69-AP9b3seIe0E,2869
 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=urPrHUqM1mpx96VHiQWVXCy3NQwDh6OsSkm4V4wqccY,30211
+vanna/flask/__init__.py,sha256=r1ucQupb6wuTcjVVKpkdrg6R38FZe6KQoKw9AtcghDQ,42889
 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
+vanna/google/gemini_chat.py,sha256=j1szC2PamMLFrs0Z4lYPS69i017FYICe-mNObNYFBPQ,1576
 vanna/hf/__init__.py,sha256=vD0bIhfLkA1UsvVSF4MAz3Da8aQunkQo3wlDztmMuj0,19
-vanna/hf/hf.py,sha256=v1v6sZnbj5xcrjgmvLP_ytS9NM7E5d0GyMfXXtr6BMU,2703
+vanna/hf/hf.py,sha256=N8N5g3xvKDBt3dez2r_U0qATxbl2pN8SVLTZK9CSRA0,3020
 vanna/marqo/__init__.py,sha256=GaAWtJ0B-H5rTY607iLCCrLD7T0zMYM5qWIomEB9gLk,37
 vanna/marqo/marqo.py,sha256=W7WTtzWp4RJjZVy6OaXHqncUBIPdI4Q7qH7BRCxZ1_A,5242
 vanna/milvus/__init__.py,sha256=VBasJG2eTKbJI6CEand7kPLNBrqYrn0QCAhSYVz814s,46
@@ -33,7 +35,7 @@ vanna/mock/vectordb.py,sha256=h45znfYMUnttE2BBC8v6TKeMaA58pFJL-5B3OGeRNFI,2681
 vanna/ollama/__init__.py,sha256=4xyu8aHPdnEHg5a-QAMwr5o0ns5wevsp_zkI-ndMO2k,27
 vanna/ollama/ollama.py,sha256=rXa7cfvdlO1E5SLysXIl3IZpIaA2r0RBvV5jX2-upiE,3794
 vanna/openai/__init__.py,sha256=tGkeQ7wTIPsando7QhoSHehtoQVdYLwFbKNlSmCmNeQ,86
-vanna/openai/openai_chat.py,sha256=lm-hUsQxu6Q1t06A2csC037zI4VkMk0wFbQ-_Lj74Wg,4764
+vanna/openai/openai_chat.py,sha256=KU6ynOQ5v7vwrQQ13phXoUXeQUrH6_vmhfiPvWddTrQ,4427
 vanna/openai/openai_embeddings.py,sha256=g4pNh9LVcYP9wOoO8ecaccDFWmCUYMInebfHucAa2Gc,1260
 vanna/opensearch/__init__.py,sha256=0unDevWOTs7o8S79TOHUKF1mSiuQbBUVm-7k9jV5WW4,54
 vanna/opensearch/opensearch_vector.py,sha256=VhIcrSyNzWR9ZrqrJnyGFOyuQZs3swfbhr8QyVGI0eI,12226
@@ -46,6 +48,8 @@ vanna/vannadb/__init__.py,sha256=C6UkYocmO6dmzfPKZaWojN0mI5YlZZ9VIbdcquBE58A,48
 vanna/vannadb/vannadb_vector.py,sha256=N8poMYvAojoaOF5gI4STD5pZWK9lBKPvyIjbh9dPBa0,14189
 vanna/vllm/__init__.py,sha256=aNlUkF9tbURdeXAJ8ytuaaF1gYwcG3ny1MfNl_cwQYg,23
 vanna/vllm/vllm.py,sha256=oM_aA-1Chyl7T_Qc_yRKlL6oSX1etsijY9zQdjeMGMQ,2827
-vanna-0.6.2.dist-info/WHEEL,sha256=EZbGkh7Ie4PoZfRQ8I0ZuP9VklN_TvcZ6DSE5Uar4z4,81
-vanna-0.6.2.dist-info/METADATA,sha256=RVle66HeuhBS8iaO0vD8_iDQqk9NbeO1pZOgCgKwh54,11628
-vanna-0.6.2.dist-info/RECORD,,
+vanna/weaviate/__init__.py,sha256=HL6PAl7ePBAkeG8uln-BmM7IUtWohyTPvDfcPzSGSCg,46
+vanna/weaviate/weaviate_vector.py,sha256=GEiu4Vd9w-7j10aB-zTxJ8gefqe_F-LUUGvttFs1vlg,7539
+vanna-0.6.4.dist-info/WHEEL,sha256=EZbGkh7Ie4PoZfRQ8I0ZuP9VklN_TvcZ6DSE5Uar4z4,81
+vanna-0.6.4.dist-info/METADATA,sha256=LqIi4Hg1y_aTEH79PX48nnY1TM-u6ese9K8Os9Cqkg0,11889
+vanna-0.6.4.dist-info/RECORD,,