asksql 2.1.0__py3-none-any.whl

vanna/legacy/base/safe_exec.py

@@ -0,0 +1,190 @@
+"""Guarded execution of LLM-generated Plotly code (vanna-ai/vanna#1098).
+
+`get_plotly_figure` runs Python written by the LLM, and the LLM prompt can be
+influenced by users (e.g. the Flask `chart_instructions` query param). Without
+a guard, a prompt injection becomes remote code execution.
+
+This is defence in depth, NOT a true sandbox - CPython cannot be sandboxed
+in-process. It closes the easy paths: arbitrary imports, `_`/dunder tricks,
+file/network IO helpers, dangerous builtins and access to vanna's module globals.
+"""
+
+import ast
+import builtins
+
+import pandas as pd
+import plotly
+import plotly.express as px
+import plotly.graph_objects as go
+
+ALLOWED_IMPORT_ROOTS = {"plotly", "pandas", "numpy", "math"}
+
+# Attribute/import names that lead to the OS, the interpreter, or file/network IO.
+DENIED_NAMES = {
+    # modules reachable as attributes of other modules
+    "os",
+    "sys",
+    "subprocess",
+    "builtins",
+    "importlib",
+    "io",
+    "ctypes",
+    "ctypeslib",
+    "socket",
+    "shutil",
+    "pickle",
+    "marshal",
+    "pty",
+    "requests",
+    # process / code execution
+    "system",
+    "popen",
+    "Popen",
+    "spawn",
+    "fork",
+    "eval",
+    "exec",
+    "compile",
+    "query",
+    "globals",
+    "locals",
+    "vars",
+    "getattr",
+    "setattr",
+    "delattr",
+    "open",
+    "input",
+    "breakpoint",
+    # pandas / numpy / plotly file IO
+    "to_csv",
+    "to_excel",
+    "to_json",
+    "to_html",
+    "to_parquet",
+    "to_pickle",
+    "to_sql",
+    "to_hdf",
+    "to_feather",
+    "to_stata",
+    "to_xml",
+    "to_orc",
+    "to_latex",
+    "to_markdown",
+    "to_string",
+    "to_clipboard",
+    "tofile",
+    "write_html",
+    "write_image",
+    "write_json",
+    "save",
+    "savez",
+    "savez_compressed",
+    "savetxt",
+    "load",
+    "loadtxt",
+    "genfromtxt",
+    "fromfile",
+    "memmap",
+}
+
+SAFE_BUILTIN_NAMES = [
+    "abs",
+    "all",
+    "any",
+    "bool",
+    "dict",
+    "enumerate",
+    "filter",
+    "float",
+    "int",
+    "isinstance",
+    "len",
+    "list",
+    "map",
+    "max",
+    "min",
+    "print",
+    "range",
+    "reversed",
+    "round",
+    "set",
+    "sorted",
+    "str",
+    "sum",
+    "tuple",
+    "zip",
+    "ValueError",
+    "TypeError",
+    "KeyError",
+    "IndexError",
+    "Exception",
+]
+
+
+class UnsafePlotlyCodeError(ValueError):
+    pass
+
+
+def _is_denied(name: str) -> bool:
+    # A bare "_" is the usual throwaway loop variable; any other leading "_" is private/dunder.
+    if name != "_" and name.startswith("_"):
+        return True
+    return name in DENIED_NAMES or name.startswith("read_")
+
+
+def validate_plotly_code(code: str) -> None:
+    """Raise UnsafePlotlyCodeError if `code` uses anything outside the allowed subset."""
+    try:
+        tree = ast.parse(code)
+    except SyntaxError as e:
+        raise UnsafePlotlyCodeError(f"invalid Python: {e}") from e
+
+    for node in ast.walk(tree):
+        if isinstance(node, ast.Import):
+            for alias in node.names:
+                if alias.name.split(".")[0] not in ALLOWED_IMPORT_ROOTS:
+                    raise UnsafePlotlyCodeError(
+                        f"import of '{alias.name}' is not allowed"
+                    )
+        elif isinstance(node, ast.ImportFrom):
+            if (
+                node.level
+                or (node.module or "").split(".")[0] not in ALLOWED_IMPORT_ROOTS
+            ):
+                raise UnsafePlotlyCodeError(
+                    f"import from '{node.module}' is not allowed"
+                )
+            for alias in node.names:
+                if alias.name == "*" or _is_denied(alias.name):
+                    raise UnsafePlotlyCodeError(
+                        f"import of '{alias.name}' is not allowed"
+                    )
+        elif isinstance(node, ast.Attribute) and _is_denied(node.attr):
+            raise UnsafePlotlyCodeError(f"access to '.{node.attr}' is not allowed")
+        elif isinstance(node, ast.Name) and _is_denied(node.id):
+            raise UnsafePlotlyCodeError(f"use of '{node.id}' is not allowed")
+
+
+def _restricted_import(name, globals=None, locals=None, fromlist=(), level=0):
+    if level != 0 or name.split(".")[0] not in ALLOWED_IMPORT_ROOTS:
+        raise ImportError(f"import of '{name}' is not allowed in chart code")
+    return builtins.__import__(name, globals, locals, fromlist, level)
+
+
+def exec_plotly_code(code: str, df: pd.DataFrame):
+    """Validate and run LLM-generated Plotly code; return the `fig` it defines (or None)."""
+    validate_plotly_code(code)
+
+    safe_builtins = {name: getattr(builtins, name) for name in SAFE_BUILTIN_NAMES}
+    safe_builtins["__import__"] = _restricted_import
+
+    namespace = {
+        "__builtins__": safe_builtins,
+        "df": df,
+        "pd": pd,
+        "plotly": plotly,
+        "px": px,
+        "go": go,
+    }
+    exec(code, namespace)
+    return namespace.get("fig")

vanna/legacy/bedrock/__init__.py

@@ -0,0 +1 @@
+from .bedrock_converse import Bedrock_Converse

vanna/legacy/bedrock/bedrock_converse.py

@@ -0,0 +1,86 @@
+from ..base import VannaBase
+
+try:
+    import boto3
+    from botocore.exceptions import ClientError
+except ImportError:
+    raise ImportError("Please install boto3 and botocore to use Amazon Bedrock models")
+
+
+class Bedrock_Converse(VannaBase):
+    def __init__(self, client=None, config=None):
+        VannaBase.__init__(self, config=config)
+
+        # default parameters
+        self.temperature = 0.0
+        self.max_tokens = 500
+
+        if client is None:
+            raise ValueError(
+                "A valid Bedrock runtime client must be provided to invoke Bedrock models"
+            )
+        else:
+            self.client = client
+
+        if config is None:
+            raise ValueError(
+                "Config is required with model_id and inference parameters"
+            )
+
+        if "modelId" not in config:
+            raise ValueError("config must contain a modelId to invoke")
+        else:
+            self.model = config["modelId"]
+
+        if "temperature" in config:
+            self.temperature = config["temperature"]
+
+        if "max_tokens" in config:
+            self.max_tokens = config["max_tokens"]
+
+    def system_message(self, message: str) -> dict:
+        return {"role": "system", "content": message}
+
+    def user_message(self, message: str) -> dict:
+        return {"role": "user", "content": message}
+
+    def assistant_message(self, message: str) -> dict:
+        return {"role": "assistant", "content": message}
+
+    def submit_prompt(self, prompt, **kwargs) -> str:
+        inference_config = {
+            "temperature": self.temperature,
+            "maxTokens": self.max_tokens,
+        }
+        additional_model_fields = {
+            "top_p": 1,  # setting top_p value for nucleus sampling
+        }
+
+        system_message = None
+        no_system_prompt = []
+        for prompt_message in prompt:
+            role = prompt_message["role"]
+            if role == "system":
+                system_message = prompt_message["content"]
+            else:
+                no_system_prompt.append(
+                    {"role": role, "content": [{"text": prompt_message["content"]}]}
+                )
+
+        converse_api_params = {
+            "modelId": self.model,
+            "messages": no_system_prompt,
+            "inferenceConfig": inference_config,
+            "additionalModelRequestFields": additional_model_fields,
+        }
+
+        if system_message:
+            converse_api_params["system"] = [{"text": system_message}]
+
+        try:
+            response = self.client.converse(**converse_api_params)
+            text_content = response["output"]["message"]["content"][0]["text"]
+            return text_content
+        except ClientError as err:
+            message = err.response["Error"]["Message"]
+            raise Exception(f"A Bedrock client error occurred: {message}")

vanna/legacy/chromadb/__init__.py

@@ -0,0 +1 @@
+from .chromadb_vector import ChromaDB_VectorStore

vanna/legacy/chromadb/chromadb_vector.py

@@ -0,0 +1,259 @@
+import json
+from typing import List
+
+import chromadb
+import pandas as pd
+from chromadb.config import Settings
+from chromadb.utils import embedding_functions
+
+from ..base import VannaBase
+from ..utils import deterministic_uuid
+
+default_ef = embedding_functions.DefaultEmbeddingFunction()
+
+
+class ChromaDB_VectorStore(VannaBase):
+    def __init__(self, config=None):
+        VannaBase.__init__(self, config=config)
+        if config is None:
+            config = {}
+
+        path = config.get("path", ".")
+        self.embedding_function = config.get("embedding_function", default_ef)
+        curr_client = config.get("client", "persistent")
+        collection_metadata = config.get("collection_metadata", None)
+        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 curr_client == "persistent":
+            self.chroma_client = chromadb.PersistentClient(
+                path=path, settings=Settings(anonymized_telemetry=False)
+            )
+        elif curr_client == "in-memory":
+            self.chroma_client = chromadb.EphemeralClient(
+                settings=Settings(anonymized_telemetry=False)
+            )
+        elif isinstance(curr_client, chromadb.api.client.Client):
+            # allow providing client directly
+            self.chroma_client = curr_client
+        else:
+            raise ValueError(f"Unsupported client was set in config: {curr_client}")
+
+        self.documentation_collection = self.chroma_client.get_or_create_collection(
+            name="documentation",
+            embedding_function=self.embedding_function,
+            metadata=collection_metadata,
+        )
+        self.ddl_collection = self.chroma_client.get_or_create_collection(
+            name="ddl",
+            embedding_function=self.embedding_function,
+            metadata=collection_metadata,
+        )
+        self.sql_collection = self.chroma_client.get_or_create_collection(
+            name="sql",
+            embedding_function=self.embedding_function,
+            metadata=collection_metadata,
+        )
+
+    def generate_embedding(self, data: str, **kwargs) -> List[float]:
+        embedding = self.embedding_function([data])
+        if len(embedding) == 1:
+            return embedding[0]
+        return embedding
+
+    def add_question_sql(self, question: str, sql: str, **kwargs) -> str:
+        question_sql_json = json.dumps(
+            {
+                "question": question,
+                "sql": sql,
+            },
+            ensure_ascii=False,
+        )
+        id = deterministic_uuid(question_sql_json) + "-sql"
+        self.sql_collection.add(
+            documents=question_sql_json,
+            embeddings=self.generate_embedding(question_sql_json),
+            ids=id,
+        )
+
+        return id
+
+    def add_ddl(self, ddl: str, **kwargs) -> str:
+        id = deterministic_uuid(ddl) + "-ddl"
+        self.ddl_collection.add(
+            documents=ddl,
+            embeddings=self.generate_embedding(ddl),
+            ids=id,
+        )
+        return id
+
+    def add_documentation(self, documentation: str, **kwargs) -> str:
+        id = deterministic_uuid(documentation) + "-doc"
+        self.documentation_collection.add(
+            documents=documentation,
+            embeddings=self.generate_embedding(documentation),
+            ids=id,
+        )
+        return id
+
+    def get_training_data(self, **kwargs) -> pd.DataFrame:
+        sql_data = self.sql_collection.get()
+
+        df = pd.DataFrame()
+
+        if sql_data is not None:
+            # Extract the documents and ids
+            documents = [json.loads(doc) for doc in sql_data["documents"]]
+            ids = sql_data["ids"]
+
+            # Create a DataFrame
+            df_sql = pd.DataFrame(
+                {
+                    "id": ids,
+                    "question": [doc["question"] for doc in documents],
+                    "content": [doc["sql"] for doc in documents],
+                }
+            )
+
+            df_sql["training_data_type"] = "sql"
+
+            df = pd.concat([df, df_sql])
+
+        ddl_data = self.ddl_collection.get()
+
+        if ddl_data is not None:
+            # Extract the documents and ids
+            documents = [doc for doc in ddl_data["documents"]]
+            ids = ddl_data["ids"]
+
+            # Create a DataFrame
+            df_ddl = pd.DataFrame(
+                {
+                    "id": ids,
+                    "question": [None for doc in documents],
+                    "content": [doc for doc in documents],
+                }
+            )
+
+            df_ddl["training_data_type"] = "ddl"
+
+            df = pd.concat([df, df_ddl])
+
+        doc_data = self.documentation_collection.get()
+
+        if doc_data is not None:
+            # Extract the documents and ids
+            documents = [doc for doc in doc_data["documents"]]
+            ids = doc_data["ids"]
+
+            # Create a DataFrame
+            df_doc = pd.DataFrame(
+                {
+                    "id": ids,
+                    "question": [None for doc in documents],
+                    "content": [doc for doc in documents],
+                }
+            )
+
+            df_doc["training_data_type"] = "documentation"
+
+            df = pd.concat([df, df_doc])
+
+        return df
+
+    def remove_training_data(self, id: str, **kwargs) -> bool:
+        if id.endswith("-sql"):
+            self.sql_collection.delete(ids=id)
+            return True
+        elif id.endswith("-ddl"):
+            self.ddl_collection.delete(ids=id)
+            return True
+        elif id.endswith("-doc"):
+            self.documentation_collection.delete(ids=id)
+            return True
+        else:
+            return False
+
+    def remove_collection(self, collection_name: str) -> bool:
+        """
+        This function can reset the collection to empty state.
+
+        Args:
+            collection_name (str): sql or ddl or documentation
+
+        Returns:
+            bool: True if collection is deleted, False otherwise
+        """
+        if collection_name == "sql":
+            self.chroma_client.delete_collection(name="sql")
+            self.sql_collection = self.chroma_client.get_or_create_collection(
+                name="sql", embedding_function=self.embedding_function
+            )
+            return True
+        elif collection_name == "ddl":
+            self.chroma_client.delete_collection(name="ddl")
+            self.ddl_collection = self.chroma_client.get_or_create_collection(
+                name="ddl", embedding_function=self.embedding_function
+            )
+            return True
+        elif collection_name == "documentation":
+            self.chroma_client.delete_collection(name="documentation")
+            self.documentation_collection = self.chroma_client.get_or_create_collection(
+                name="documentation", embedding_function=self.embedding_function
+            )
+            return True
+        else:
+            return False
+
+    @staticmethod
+    def _extract_documents(query_results) -> list:
+        """
+        Static method to extract the documents from the results of a query.
+
+        Args:
+            query_results (pd.DataFrame): The dataframe to use.
+
+        Returns:
+            List[str] or None: The extracted documents, or an empty list or
+            single document if an error occurred.
+        """
+        if query_results is None:
+            return []
+
+        if "documents" in query_results:
+            documents = query_results["documents"]
+
+            if len(documents) == 1 and isinstance(documents[0], list):
+                try:
+                    documents = [json.loads(doc) for doc in documents[0]]
+                except Exception as e:
+                    return documents[0]
+
+            return documents
+
+    def get_similar_question_sql(self, question: str, **kwargs) -> list:
+        return ChromaDB_VectorStore._extract_documents(
+            self.sql_collection.query(
+                query_texts=[question],
+                n_results=self.n_results_sql,
+            )
+        )
+
+    def get_related_ddl(self, question: str, **kwargs) -> list:
+        return ChromaDB_VectorStore._extract_documents(
+            self.ddl_collection.query(
+                query_texts=[question],
+                n_results=self.n_results_ddl,
+            )
+        )
+
+    def get_related_documentation(self, question: str, **kwargs) -> list:
+        return ChromaDB_VectorStore._extract_documents(
+            self.documentation_collection.query(
+                query_texts=[question],
+                n_results=self.n_results_documentation,
+            )
+        )

vanna/legacy/cohere/__init__.py

@@ -0,0 +1,2 @@
+from .cohere_chat import Cohere_Chat
+from .cohere_embeddings import Cohere_Embeddings

vanna/legacy/cohere/cohere_chat.py

@@ -0,0 +1,99 @@
+import os
+
+from openai import OpenAI
+
+from ..base import VannaBase
+
+
+class Cohere_Chat(VannaBase):
+    def __init__(self, client=None, config=None):
+        VannaBase.__init__(self, config=config)
+
+        # default parameters - can be overridden using config
+        self.temperature = 0.2  # Lower temperature for more precise SQL generation
+        self.model = "command-a-03-2025"  # Cohere's default model
+
+        if config is not None:
+            if "temperature" in config:
+                self.temperature = config["temperature"]
+            if "model" in config:
+                self.model = config["model"]
+
+        if client is not None:
+            self.client = client
+            return
+
+        # Check for API key in environment variable
+        api_key = os.getenv("COHERE_API_KEY")
+
+        # Check for API key in config
+        if config is not None and "api_key" in config:
+            api_key = config["api_key"]
+
+        # Validate API key
+        if not api_key:
+            raise ValueError(
+                "Cohere API key is required. Please provide it via config or set the COHERE_API_KEY environment variable."
+            )
+
+        # Initialize client with validated API key
+        self.client = OpenAI(
+            base_url="https://api.cohere.ai/compatibility/v1",
+            api_key=api_key,
+        )
+
+    def system_message(self, message: str) -> any:
+        return {
+            "role": "developer",
+            "content": message,
+        }  # Cohere uses 'developer' for system role
+
+    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 submit_prompt(self, prompt, **kwargs) -> str:
+        if prompt is None:
+            raise Exception("Prompt is None")
+
+        if len(prompt) == 0:
+            raise Exception("Prompt is empty")
+
+        # Count the number of tokens in the message log
+        # Use 4 as an approximation for the number of characters per token
+        num_tokens = 0
+        for message in prompt:
+            num_tokens += len(message["content"]) / 4
+
+        # Use model from kwargs, config, or default
+        model = kwargs.get("model", self.model)
+        if self.config is not None and "model" in self.config and model == self.model:
+            model = self.config["model"]
+
+        print(f"Using model {model} for {num_tokens} tokens (approx)")
+        try:
+            response = self.client.chat.completions.create(
+                model=model,
+                messages=prompt,
+                temperature=self.temperature,
+            )
+
+            # Check if response has expected structure
+            if not response or not hasattr(response, "choices") or not response.choices:
+                raise ValueError("Received empty or malformed response from API")
+
+            if not response.choices[0] or not hasattr(response.choices[0], "message"):
+                raise ValueError("Response is missing expected 'message' field")
+
+            if not hasattr(response.choices[0].message, "content"):
+                raise ValueError("Response message is missing expected 'content' field")
+
+            return response.choices[0].message.content
+
+        except Exception as e:
+            # Log the error and raise a more informative exception
+            error_msg = f"Error processing Cohere chat response: {str(e)}"
+            print(error_msg)
+            raise Exception(error_msg)