asksql 2.1.0__py3-none-any.whl

vanna/legacy/flask/auth.py

@@ -0,0 +1,56 @@
+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/legacy/google/__init__.py

@@ -0,0 +1,2 @@
+from .bigquery_vector import BigQuery_VectorStore
+from .gemini_chat import GoogleGeminiChat

vanna/legacy/google/bigquery_vector.py

@@ -0,0 +1,293 @@
+import datetime
+import os
+import uuid
+from typing import List, Optional
+from vertexai.language_models import TextEmbeddingInput, TextEmbeddingModel
+
+import pandas as pd
+from google.cloud import bigquery
+
+from ..base import VannaBase
+
+
+class BigQuery_VectorStore(VannaBase):
+    def __init__(self, config: dict, **kwargs):
+        self.config = config
+
+        self.n_results_sql = config.get("n_results_sql", config.get("n_results", 10))
+        self.n_results_documentation = config.get(
+            "n_results_documentation", config.get("n_results", 10)
+        )
+        self.n_results_ddl = config.get("n_results_ddl", config.get("n_results", 10))
+
+        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.
+            """
+            print("Configuring genai")
+            self.type = "GEMINI"
+            import google.generativeai as genai
+
+            genai.configure(api_key=config["api_key"])
+
+            self.genai = genai
+        else:
+            self.type = "VERTEX_AI"
+            # Authenticate using VertexAI
+
+        if self.config.get("project_id"):
+            self.project_id = self.config.get("project_id")
+        else:
+            self.project_id = os.getenv("GOOGLE_CLOUD_PROJECT")
+
+        if self.project_id is None:
+            raise ValueError("Project ID is not set")
+
+        self.conn = bigquery.Client(project=self.project_id)
+
+        dataset_name = self.config.get("bigquery_dataset_name", "vanna_managed")
+        self.dataset_id = f"{self.project_id}.{dataset_name}"
+        dataset = bigquery.Dataset(self.dataset_id)
+
+        try:
+            self.conn.get_dataset(self.dataset_id)  # Make an API request.
+            print(f"Dataset {self.dataset_id} already exists")
+        except Exception:
+            # Dataset does not exist, create it
+            dataset.location = "US"
+            self.conn.create_dataset(dataset, timeout=30)  # Make an API request.
+            print(f"Created dataset {self.dataset_id}")
+
+        # Create a table called training_data in the dataset that contains the columns:
+        # id, training_data_type, question, content, embedding, created_at
+
+        self.table_id = f"{self.dataset_id}.training_data"
+        schema = [
+            bigquery.SchemaField("id", "STRING", mode="REQUIRED"),
+            bigquery.SchemaField("training_data_type", "STRING", mode="REQUIRED"),
+            bigquery.SchemaField("question", "STRING", mode="REQUIRED"),
+            bigquery.SchemaField("content", "STRING", mode="REQUIRED"),
+            bigquery.SchemaField("embedding", "FLOAT64", mode="REPEATED"),
+            bigquery.SchemaField("created_at", "TIMESTAMP", mode="REQUIRED"),
+        ]
+
+        table = bigquery.Table(self.table_id, schema=schema)
+
+        try:
+            self.conn.get_table(self.table_id)  # Make an API request.
+            print(f"Table {self.table_id} already exists")
+        except Exception:
+            # Table does not exist, create it
+            self.conn.create_table(table, timeout=30)  # Make an API request.
+            print(f"Created table {self.table_id}")
+
+        # Create VECTOR INDEX IF NOT EXISTS
+        # TODO: This requires 5000 rows before it can be created
+        # vector_index_query = f"""
+        # CREATE VECTOR INDEX IF NOT EXISTS my_index
+        # ON `{self.table_id}`(embedding)
+        # OPTIONS(
+        #     distance_type='COSINE',
+        #     index_type='IVF',
+        #     ivf_options='{{"num_lists": 1000}}'
+        # )
+        # """
+
+        # try:
+        #     self.conn.query(vector_index_query).result()  # Make an API request.
+        #     print(f"Vector index on {self.table_id} created or already exists")
+        # except Exception as e:
+        #     print(f"Failed to create vector index: {e}")
+
+    def store_training_data(
+        self,
+        training_data_type: str,
+        question: str,
+        content: str,
+        embedding: List[float],
+        **kwargs,
+    ) -> str:
+        id = str(uuid.uuid4())
+        created_at = datetime.datetime.now()
+        self.conn.insert_rows_json(
+            self.table_id,
+            [
+                {
+                    "id": id,
+                    "training_data_type": training_data_type,
+                    "question": question,
+                    "content": content,
+                    "embedding": embedding,
+                    "created_at": created_at.isoformat(),
+                }
+            ],
+        )
+
+        return id
+
+    def fetch_similar_training_data(
+        self, training_data_type: str, question: str, n_results, **kwargs
+    ) -> pd.DataFrame:
+        question_embedding = self.generate_question_embedding(question)
+
+        query = f"""
+        SELECT
+            base.id as id,
+            base.question as question,
+            base.training_data_type as training_data_type,
+            base.content as content,
+            distance
+        FROM
+            VECTOR_SEARCH(
+                TABLE `{self.table_id}`,
+                'embedding',
+                (SELECT * FROM UNNEST([STRUCT({question_embedding})])),
+                top_k => 5,
+                distance_type => 'COSINE',
+                options => '{{"use_brute_force":true}}'
+            )
+        WHERE
+            base.training_data_type = @training_data_type
+        """
+        job_config = bigquery.QueryJobConfig(
+            query_parameters=[
+                bigquery.ScalarQueryParameter(
+                    "training_data_type", "STRING", training_data_type
+                )
+            ]
+        )
+
+        results = self.conn.query(query, job_config=job_config).result().to_dataframe()
+        return results
+
+    def get_embeddings(self, data: str, task: str) -> List[float]:
+        embeddings = None
+
+        if self.type == "VERTEX_AI":
+            input = [TextEmbeddingInput(data, task)]
+            model = TextEmbeddingModel.from_pretrained("text-embedding-004")
+
+            result = model.get_embeddings(input)
+
+            if len(result) > 0:
+                embeddings = result[0].values
+        else:
+            # Use Gemini Consumer API
+            result = self.genai.embed_content(
+                model="models/text-embedding-004", content=data, task_type=task
+            )
+
+            if "embedding" in result:
+                embeddings = result["embedding"]
+
+        return embeddings
+
+    def generate_question_embedding(self, data: str, **kwargs) -> List[float]:
+        result = self.get_embeddings(data, "RETRIEVAL_QUERY")
+
+        if result is not None:
+            return result
+        else:
+            raise ValueError("No embeddings returned")
+
+    def generate_storage_embedding(self, data: str, **kwargs) -> List[float]:
+        result = self.get_embeddings(data, "RETRIEVAL_DOCUMENT")
+
+        if result is not None:
+            return result
+        else:
+            raise ValueError("No embeddings returned")
+
+        # task = "RETRIEVAL_DOCUMENT"
+        # inputs = [TextEmbeddingInput(data, task)]
+        # embeddings = self.vertex_embedding_model.get_embeddings(inputs)
+
+        # if len(embeddings) == 0:
+        #     raise ValueError("No embeddings returned")
+
+        # return embeddings[0].values
+
+        return result
+
+    def generate_embedding(self, data: str, **kwargs) -> List[float]:
+        return self.generate_storage_embedding(data, **kwargs)
+
+    def get_similar_question_sql(self, question: str, **kwargs) -> list:
+        df = self.fetch_similar_training_data(
+            training_data_type="sql", question=question, n_results=self.n_results_sql
+        )
+
+        # Return a list of dictionaries with only question, sql fields. The content field needs to be renamed to sql
+        return df.rename(columns={"content": "sql"})[["question", "sql"]].to_dict(
+            orient="records"
+        )
+
+    def get_related_ddl(self, question: str, **kwargs) -> list:
+        df = self.fetch_similar_training_data(
+            training_data_type="ddl", question=question, n_results=self.n_results_ddl
+        )
+
+        # Return a list of strings of the content
+        return df["content"].tolist()
+
+    def get_related_documentation(self, question: str, **kwargs) -> list:
+        df = self.fetch_similar_training_data(
+            training_data_type="documentation",
+            question=question,
+            n_results=self.n_results_documentation,
+        )
+
+        # Return a list of strings of the content
+        return df["content"].tolist()
+
+    def add_question_sql(self, question: str, sql: str, **kwargs) -> str:
+        doc = {"question": question, "sql": sql}
+
+        embedding = self.generate_embedding(str(doc))
+
+        return self.store_training_data(
+            training_data_type="sql",
+            question=question,
+            content=sql,
+            embedding=embedding,
+        )
+
+    def add_ddl(self, ddl: str, **kwargs) -> str:
+        embedding = self.generate_embedding(ddl)
+
+        return self.store_training_data(
+            training_data_type="ddl", question="", content=ddl, embedding=embedding
+        )
+
+    def add_documentation(self, documentation: str, **kwargs) -> str:
+        embedding = self.generate_embedding(documentation)
+
+        return self.store_training_data(
+            training_data_type="documentation",
+            question="",
+            content=documentation,
+            embedding=embedding,
+        )
+
+    def get_training_data(self, **kwargs) -> pd.DataFrame:
+        query = (
+            f"SELECT id, training_data_type, question, content FROM `{self.table_id}`"
+        )
+
+        return self.conn.query(query).result().to_dataframe()
+
+    def remove_training_data(self, id: str, **kwargs) -> bool:
+        # Pass id as a query parameter, never inside the SQL string (CVE-2026-4229)
+        query = f"DELETE FROM `{self.table_id}` WHERE id = @id"
+        job_config = bigquery.QueryJobConfig(
+            query_parameters=[bigquery.ScalarQueryParameter("id", "STRING", id)]
+        )
+
+        try:
+            self.conn.query(query, job_config=job_config).result()
+            return True
+
+        except Exception as e:
+            print(f"Failed to remove training data: {e}")
+            return False

vanna/legacy/google/gemini_chat.py

@@ -0,0 +1,78 @@
+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.5-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
+            import google.auth
+            import vertexai
+            from vertexai.generative_models import GenerativeModel
+
+            json_file_path = config.get(
+                "google_credentials"
+            )  # Assuming the JSON file path is provided in the config
+
+            if not json_file_path or not os.path.exists(json_file_path):
+                raise FileNotFoundError(
+                    f"JSON credentials file not found at: {json_file_path}"
+                )
+
+            try:
+                # Validate and set the JSON file path for GOOGLE_APPLICATION_CREDENTIALS
+                os.environ["GOOGLE_APPLICATION_CREDENTIALS"] = json_file_path
+
+                # Initialize VertexAI with the credentials
+                credentials, _ = google.auth.default()
+                vertexai.init(credentials=credentials)
+                self.chat_model = GenerativeModel(model_name)
+            except google.auth.exceptions.DefaultCredentialsError as e:
+                raise RuntimeError(f"Default credentials error: {e}")
+            except google.auth.exceptions.TransportError as e:
+                raise RuntimeError(f"Transport error during authentication: {e}")
+            except Exception as e:
+                raise RuntimeError(f"Failed to authenticate using JSON file: {e}")
+
+    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/legacy/hf/__init__.py

@@ -0,0 +1 @@
+from .hf import Hf

vanna/legacy/hf/hf.py

@@ -0,0 +1,80 @@
+import re
+from transformers import AutoTokenizer, AutoModelForCausalLM
+
+from ..base import VannaBase
+
+
+class Hf(VannaBase):
+    def __init__(self, config=None):
+        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_or_path,
+            quantization_config=quantization_config,
+            device_map="auto",
+        )
+
+    def system_message(self, message: str) -> any:
+        return {"role": "system", "content": message}
+
+    def user_message(self, message: str) -> any:
+        return {"role": "user", "content": message}
+
+    def assistant_message(self, message: str) -> any:
+        return {"role": "assistant", "content": message}
+
+    def extract_sql_query(self, text):
+        """
+        Extracts the first SQL statement after the word 'select', ignoring case,
+        matches until the first semicolon, three backticks, or the end of the string,
+        and removes three backticks if they exist in the extracted string.
+
+        Args:
+        - text (str): The string to search within for an SQL statement.
+
+        Returns:
+        - str: The first SQL statement found, with three backticks removed, or an empty string if no match is found.
+        """
+        # Regular expression to find 'select' (ignoring case) and capture until ';', '```', or end of string
+        pattern = re.compile(r"select.*?(?:;|```|$)", re.IGNORECASE | re.DOTALL)
+
+        match = pattern.search(text)
+        if match:
+            # Remove three backticks from the matched string if they exist
+            return match.group(0).replace("```", "")
+        else:
+            return text
+
+    def generate_sql(self, question: str, **kwargs) -> str:
+        # Use the super generate_sql
+        sql = super().generate_sql(question, **kwargs)
+
+        # Replace "\_" with "_"
+        sql = sql.replace("\\_", "_")
+
+        sql = sql.replace("\\", "")
+
+        return self.extract_sql_query(sql)
+
+    def submit_prompt(self, prompt, **kwargs) -> str:
+        input_ids = self.tokenizer.apply_chat_template(
+            prompt, add_generation_prompt=True, return_tensors="pt"
+        ).to(self.model.device)
+
+        outputs = self.model.generate(
+            input_ids,
+            max_new_tokens=512,
+            eos_token_id=self.tokenizer.eos_token_id,
+            do_sample=True,
+            temperature=1,
+            top_p=0.9,
+        )
+        response = outputs[0][input_ids.shape[-1] :]
+        response = self.tokenizer.decode(response, skip_special_tokens=True)
+        self.log(response)
+
+        return response

vanna/legacy/local.py

@@ -0,0 +1,8 @@
+from .chromadb.chromadb_vector import ChromaDB_VectorStore
+from .openai.openai_chat import OpenAI_Chat
+
+
+class LocalContext_OpenAI(ChromaDB_VectorStore, OpenAI_Chat):
+    def __init__(self, config=None):
+        ChromaDB_VectorStore.__init__(self, config=config)
+        OpenAI_Chat.__init__(self, config=config)

vanna/legacy/marqo/__init__.py

@@ -0,0 +1 @@
+from .marqo import Marqo_VectorStore

vanna/legacy/marqo/marqo.py

@@ -0,0 +1,168 @@
+import uuid
+
+import marqo
+import pandas as pd
+
+from ..base import VannaBase
+
+
+class Marqo_VectorStore(VannaBase):
+    def __init__(self, config=None):
+        VannaBase.__init__(self, config=config)
+
+        if config is not None and "marqo_url" in config:
+            marqo_url = config["marqo_url"]
+        else:
+            marqo_url = "http://localhost:8882"
+
+        if config is not None and "marqo_model" in config:
+            marqo_model = config["marqo_model"]
+        else:
+            marqo_model = "hf/all_datasets_v4_MiniLM-L6"
+
+        self.mq = marqo.Client(url=marqo_url)
+
+        for index in ["vanna-sql", "vanna-ddl", "vanna-doc"]:
+            try:
+                self.mq.create_index(index, model=marqo_model)
+            except Exception as e:
+                print(e)
+                print(f"Marqo index {index} already exists")
+                pass
+
+    def generate_embedding(self, data: str, **kwargs) -> list[float]:
+        # Marqo doesn't need to generate embeddings
+        pass
+
+    def add_question_sql(self, question: str, sql: str, **kwargs) -> str:
+        id = str(uuid.uuid4()) + "-sql"
+        question_sql_dict = {
+            "question": question,
+            "sql": sql,
+            "_id": id,
+        }
+
+        self.mq.index("vanna-sql").add_documents(
+            [question_sql_dict],
+            tensor_fields=["question", "sql"],
+        )
+
+        return id
+
+    def add_ddl(self, ddl: str, **kwargs) -> str:
+        id = str(uuid.uuid4()) + "-ddl"
+        ddl_dict = {
+            "ddl": ddl,
+            "_id": id,
+        }
+
+        self.mq.index("vanna-ddl").add_documents(
+            [ddl_dict],
+            tensor_fields=["ddl"],
+        )
+        return id
+
+    def add_documentation(self, documentation: str, **kwargs) -> str:
+        id = str(uuid.uuid4()) + "-doc"
+        doc_dict = {
+            "doc": documentation,
+            "_id": id,
+        }
+
+        self.mq.index("vanna-doc").add_documents(
+            [doc_dict],
+            tensor_fields=["doc"],
+        )
+        return id
+
+    def get_training_data(self, **kwargs) -> pd.DataFrame:
+        data = []
+
+        for hit in self.mq.index("vanna-doc").search("", limit=1000)["hits"]:
+            data.append(
+                {
+                    "id": hit["_id"],
+                    "training_data_type": "documentation",
+                    "question": "",
+                    "content": hit["doc"],
+                }
+            )
+
+        for hit in self.mq.index("vanna-ddl").search("", limit=1000)["hits"]:
+            data.append(
+                {
+                    "id": hit["_id"],
+                    "training_data_type": "ddl",
+                    "question": "",
+                    "content": hit["ddl"],
+                }
+            )
+
+        for hit in self.mq.index("vanna-sql").search("", limit=1000)["hits"]:
+            data.append(
+                {
+                    "id": hit["_id"],
+                    "training_data_type": "sql",
+                    "question": hit["question"],
+                    "content": hit["sql"],
+                }
+            )
+
+        df = pd.DataFrame(data)
+
+        return df
+
+    def remove_training_data(self, id: str, **kwargs) -> bool:
+        if id.endswith("-sql"):
+            self.mq.index("vanna-sql").delete_documents(ids=[id])
+            return True
+        elif id.endswith("-ddl"):
+            self.mq.index("vanna-ddl").delete_documents(ids=[id])
+            return True
+        elif id.endswith("-doc"):
+            self.mq.index("vanna-doc").delete_documents(ids=[id])
+            return True
+        else:
+            return False
+
+    # Static method to extract the documents from the results of a query
+    @staticmethod
+    def _extract_documents(data) -> list:
+        # Check if 'hits' key is in the dictionary and if it's a list
+        if "hits" in data and isinstance(data["hits"], list):
+            # Iterate over each item in 'hits'
+
+            if len(data["hits"]) == 0:
+                return []
+
+            # If there is a "doc" key, return the value of that key
+            if "doc" in data["hits"][0]:
+                return [hit["doc"] for hit in data["hits"]]
+
+            # If there is a "ddl" key, return the value of that key
+            if "ddl" in data["hits"][0]:
+                return [hit["ddl"] for hit in data["hits"]]
+
+            # Otherwise return the entire hit
+            return [
+                {key: value for key, value in hit.items() if not key.startswith("_")}
+                for hit in data["hits"]
+            ]
+        else:
+            # Return an empty list if 'hits' is not found or not a list
+            return []
+
+    def get_similar_question_sql(self, question: str, **kwargs) -> list:
+        return Marqo_VectorStore._extract_documents(
+            self.mq.index("vanna-sql").search(question)
+        )
+
+    def get_related_ddl(self, question: str, **kwargs) -> list:
+        return Marqo_VectorStore._extract_documents(
+            self.mq.index("vanna-ddl").search(question)
+        )
+
+    def get_related_documentation(self, question: str, **kwargs) -> list:
+        return Marqo_VectorStore._extract_documents(
+            self.mq.index("vanna-doc").search(question)
+        )