vanna 0.3.4__py3-none-any.whl → 0.4.0__py3-none-any.whl

vanna/flask/auth.py

@@ -0,0 +1,55 @@
+from abc import ABC, abstractmethod
+
+import flask
+
+
+class AuthInterface(ABC):
+    @abstractmethod
+    def get_user(self, flask_request) -> any:
+        pass
+
+    @abstractmethod
+    def is_logged_in(self, user: any) -> bool:
+        pass
+
+    @abstractmethod
+    def override_config_for_user(self, user: any, config: dict) -> dict:
+        pass
+
+    @abstractmethod
+    def login_form(self) -> str:
+        pass
+
+    @abstractmethod
+    def login_handler(self, flask_request) -> str:
+        pass
+
+    @abstractmethod
+    def callback_handler(self, flask_request) -> str:
+        pass
+
+    @abstractmethod
+    def logout_handler(self, flask_request) -> str:
+        pass
+
+class NoAuth(AuthInterface):
+    def get_user(self, flask_request) -> any:
+        return {}
+
+    def is_logged_in(self, user: any) -> bool:
+        return True
+
+    def override_config_for_user(self, user: any, config: dict) -> dict:
+        return config
+
+    def login_form(self) -> str:
+        return ''
+
+    def login_handler(self, flask_request) -> str:
+        return 'No login required'
+
+    def callback_handler(self, flask_request) -> str:
+        return 'No login required'
+
+    def logout_handler(self, flask_request) -> str:
+        return 'No login required'

vanna/google/__init__.py

@@ -0,0 +1 @@
+from .gemini_chat import GoogleGeminiChat

vanna/google/gemini_chat.py

@@ -0,0 +1,52 @@
+import os
+from ..base import VannaBase
+
+
+class GoogleGeminiChat(VannaBase):
+    def __init__(self, config=None):
+        VannaBase.__init__(self, config=config)
+
+        # default temperature - can be overrided using config
+        self.temperature = 0.7      
+
+        if "temperature" in config:
+            self.temperature = config["temperature"]
+
+        if "model_name" in config:
+            model_name = config["model_name"]
+        else:
+            model_name = "gemini-1.0-pro"
+
+        self.google_api_key = None
+
+        if "api_key" in config or os.getenv("GOOGLE_API_KEY"):
+            """
+            If Google api_key is provided through config
+            or set as an environment variable, assign it.
+            """
+            import google.generativeai as genai
+
+            genai.configure(api_key=config["api_key"])
+            self.chat_model = genai.GenerativeModel(model_name)
+        else:
+            # Authenticate using VertexAI
+            from vertexai.preview.generative_models import GenerativeModel
+            self.chat_model = GenerativeModel("gemini-pro")
+
+    def system_message(self, message: str) -> any:
+        return message
+
+    def user_message(self, message: str) -> any:
+        return message
+
+    def assistant_message(self, message: str) -> any:
+        return message
+
+    def submit_prompt(self, prompt, **kwargs) -> str:
+        response = self.chat_model.generate_content(
+            prompt,
+            generation_config={
+                "temperature": self.temperature,
+            },
+        )
+        return response.text

{vanna-0.3.4.dist-info → vanna-0.4.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: vanna
-Version: 0.3.4
+Version: 0.4.0
 Summary: Generate SQL queries from natural language
 Author-email: Zain Hoda <zain@vanna.ai>
 Requires-Python: >=3.9
@@ -28,11 +28,15 @@ Requires-Dist: chromadb ; extra == "all"
 Requires-Dist: anthropic ; extra == "all"
 Requires-Dist: zhipuai ; extra == "all"
 Requires-Dist: marqo ; extra == "all"
+Requires-Dist: google-generativeai ; extra == "all"
+Requires-Dist: google-cloud-aiplatform ; extra == "all"
 Requires-Dist: anthropic ; extra == "anthropic"
 Requires-Dist: google-cloud-bigquery ; extra == "bigquery"
 Requires-Dist: chromadb ; extra == "chromadb"
 Requires-Dist: duckdb ; extra == "duckdb"
 Requires-Dist: google-generativeai ; extra == "gemini"
+Requires-Dist: google-generativeai ; extra == "google"
+Requires-Dist: google-cloud-aiplatform ; extra == "google"
 Requires-Dist: marqo ; extra == "marqo"
 Requires-Dist: mistralai ; extra == "mistralai"
 Requires-Dist: PyMySQL ; extra == "mysql"
@@ -50,6 +54,7 @@ Provides-Extra: bigquery
 Provides-Extra: chromadb
 Provides-Extra: duckdb
 Provides-Extra: gemini
+Provides-Extra: google
 Provides-Extra: marqo
 Provides-Extra: mistralai
 Provides-Extra: mysql

{vanna-0.3.4.dist-info → vanna-0.4.0.dist-info}/RECORD

@@ -8,12 +8,15 @@ vanna/ZhipuAI/__init__.py,sha256=NlsijtcZp5Tj9jkOe9fNcOQND_QsGgu7otODsCLBPr0,116
 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=89XPWy97YVx6090mNmu1zvn4k8X1pusCuAIypHHexNc,58100
+vanna/base/base.py,sha256=_2vANGAcUe6IrsEhZyFnE6FdO8NMLylfyGlRI4XujWE,58143
 vanna/chromadb/__init__.py,sha256=-iL0nW_g4uM8nWKMuWnNePfN4nb9uk8P3WzGvezOqRg,50
-vanna/chromadb/chromadb_vector.py,sha256=1n4U4XpXThCFqyJf0zAYVA7mQu9rUkjOFtYn9e04JAo,8461
+vanna/chromadb/chromadb_vector.py,sha256=eKyPck99Y6Jt-BNWojvxLG-zvAERzLSm-3zY-bKXvaA,8792
 vanna/exceptions/__init__.py,sha256=N76unE7sjbGGBz6LmCrPQAugFWr9cUFv8ErJxBrCTts,717
-vanna/flask/__init__.py,sha256=tpwpA8596Uyn60FAy7I5oJ81c7kgCB2JG9X044P0_SA,21211
-vanna/flask/assets.py,sha256=pOOtPV8aWtFsTuxJneFHcfrXhXh6cOSvS-Y8JO2HYrY,180924
+vanna/flask/__init__.py,sha256=5Du2oK5s-VSLicRPvxMAL1-Gh_jdX847FjJOf5AVapo,23721
+vanna/flask/assets.py,sha256=ZESgn0-XrJl4_YV69Lu7Dr7-Y9Eql7xqb7PsoMzrofw,183889
+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/marqo/__init__.py,sha256=GaAWtJ0B-H5rTY607iLCCrLD7T0zMYM5qWIomEB9gLk,37
 vanna/marqo/marqo.py,sha256=W7WTtzWp4RJjZVy6OaXHqncUBIPdI4Q7qH7BRCxZ1_A,5242
 vanna/mistral/__init__.py,sha256=70rTY-69Z2ehkkMj84dNMCukPo6AWdflBGvIB_pztS0,29
@@ -26,6 +29,6 @@ vanna/openai/openai_embeddings.py,sha256=g4pNh9LVcYP9wOoO8ecaccDFWmCUYMInebfHucA
 vanna/types/__init__.py,sha256=Qhn_YscKtJh7mFPCyCDLa2K8a4ORLMGVnPpTbv9uB2U,4957
 vanna/vannadb/__init__.py,sha256=C6UkYocmO6dmzfPKZaWojN0mI5YlZZ9VIbdcquBE58A,48
 vanna/vannadb/vannadb_vector.py,sha256=9YwTO3Lh5owWQE7KPMBqLp2EkiGV0RC1sEYhslzJzgI,6168
-vanna-0.3.4.dist-info/WHEEL,sha256=EZbGkh7Ie4PoZfRQ8I0ZuP9VklN_TvcZ6DSE5Uar4z4,81
-vanna-0.3.4.dist-info/METADATA,sha256=FEg4vs5ZiSAvd5YkF5oEfFqod9n3UoNfi51Q_2WKotA,10107
-vanna-0.3.4.dist-info/RECORD,,
+vanna-0.4.0.dist-info/WHEEL,sha256=EZbGkh7Ie4PoZfRQ8I0ZuP9VklN_TvcZ6DSE5Uar4z4,81
+vanna-0.4.0.dist-info/METADATA,sha256=pz6wxMyiXUNXgQJlRxW-vyA3yC8j0t2OHiTBgdsOSeg,10352
+vanna-0.4.0.dist-info/RECORD,,