asksql 2.1.0__py3-none-any.whl

vanna/legacy/milvus/__init__.py

@@ -0,0 +1 @@
+from .milvus_vector import Milvus_VectorStore

vanna/legacy/milvus/milvus_vector.py

@@ -0,0 +1,328 @@
+import uuid
+from typing import List
+
+import pandas as pd
+from pymilvus import DataType, MilvusClient, model
+
+from ..base import VannaBase
+
+# Setting the URI as a local file, e.g.`./milvus.db`,
+# is the most convenient method, as it automatically utilizes Milvus Lite
+# to store all data in this file.
+#
+# If you have large scale of data such as more than a million docs, we
+# recommend setting up a more performant Milvus server on docker or kubernetes.
+# When using this setup, please use the server URI,
+# e.g.`http://localhost:19530`, as your URI.
+
+DEFAULT_MILVUS_URI = "./milvus.db"
+# DEFAULT_MILVUS_URI = "http://localhost:19530"
+
+MAX_LIMIT_SIZE = 10_000
+
+
+class Milvus_VectorStore(VannaBase):
+    """
+    Vectorstore implementation using Milvus - https://milvus.io/docs/quickstart.md
+
+    Args:
+        - config (dict, optional): Dictionary of `Milvus_VectorStore config` options. Defaults to `None`.
+            - milvus_client: A `pymilvus.MilvusClient` instance.
+            - embedding_function:
+                A `milvus_model.base.BaseEmbeddingFunction` instance. Defaults to `DefaultEmbeddingFunction()`.
+                For more models, please refer to:
+                https://milvus.io/docs/embeddings.md
+    """
+
+    def __init__(self, config=None):
+        VannaBase.__init__(self, config=config)
+
+        if "milvus_client" in config:
+            self.milvus_client = config["milvus_client"]
+        else:
+            self.milvus_client = MilvusClient(uri=DEFAULT_MILVUS_URI)
+
+        if "embedding_function" in config:
+            self.embedding_function = config.get("embedding_function")
+        else:
+            self.embedding_function = model.DefaultEmbeddingFunction()
+        self._embedding_dim = self.embedding_function.encode_documents(["foo"])[
+            0
+        ].shape[0]
+        self._create_collections()
+        self.n_results = config.get("n_results", 10)
+
+    def _create_collections(self):
+        self._create_sql_collection("vannasql")
+        self._create_ddl_collection("vannaddl")
+        self._create_doc_collection("vannadoc")
+
+    def generate_embedding(self, data: str, **kwargs) -> List[float]:
+        return self.embedding_function.encode_documents(data).tolist()
+
+    def _create_sql_collection(self, name: str):
+        if not self.milvus_client.has_collection(collection_name=name):
+            vannasql_schema = MilvusClient.create_schema(
+                auto_id=False,
+                enable_dynamic_field=False,
+            )
+            vannasql_schema.add_field(
+                field_name="id",
+                datatype=DataType.VARCHAR,
+                max_length=65535,
+                is_primary=True,
+            )
+            vannasql_schema.add_field(
+                field_name="text", datatype=DataType.VARCHAR, max_length=65535
+            )
+            vannasql_schema.add_field(
+                field_name="sql", datatype=DataType.VARCHAR, max_length=65535
+            )
+            vannasql_schema.add_field(
+                field_name="vector",
+                datatype=DataType.FLOAT_VECTOR,
+                dim=self._embedding_dim,
+            )
+
+            vannasql_index_params = self.milvus_client.prepare_index_params()
+            vannasql_index_params.add_index(
+                field_name="vector",
+                index_name="vector",
+                index_type="AUTOINDEX",
+                metric_type="L2",
+            )
+            self.milvus_client.create_collection(
+                collection_name=name,
+                schema=vannasql_schema,
+                index_params=vannasql_index_params,
+                consistency_level="Strong",
+            )
+
+    def _create_ddl_collection(self, name: str):
+        if not self.milvus_client.has_collection(collection_name=name):
+            vannaddl_schema = MilvusClient.create_schema(
+                auto_id=False,
+                enable_dynamic_field=False,
+            )
+            vannaddl_schema.add_field(
+                field_name="id",
+                datatype=DataType.VARCHAR,
+                max_length=65535,
+                is_primary=True,
+            )
+            vannaddl_schema.add_field(
+                field_name="ddl", datatype=DataType.VARCHAR, max_length=65535
+            )
+            vannaddl_schema.add_field(
+                field_name="vector",
+                datatype=DataType.FLOAT_VECTOR,
+                dim=self._embedding_dim,
+            )
+
+            vannaddl_index_params = self.milvus_client.prepare_index_params()
+            vannaddl_index_params.add_index(
+                field_name="vector",
+                index_name="vector",
+                index_type="AUTOINDEX",
+                metric_type="L2",
+            )
+            self.milvus_client.create_collection(
+                collection_name=name,
+                schema=vannaddl_schema,
+                index_params=vannaddl_index_params,
+                consistency_level="Strong",
+            )
+
+    def _create_doc_collection(self, name: str):
+        if not self.milvus_client.has_collection(collection_name=name):
+            vannadoc_schema = MilvusClient.create_schema(
+                auto_id=False,
+                enable_dynamic_field=False,
+            )
+            vannadoc_schema.add_field(
+                field_name="id",
+                datatype=DataType.VARCHAR,
+                max_length=65535,
+                is_primary=True,
+            )
+            vannadoc_schema.add_field(
+                field_name="doc", datatype=DataType.VARCHAR, max_length=65535
+            )
+            vannadoc_schema.add_field(
+                field_name="vector",
+                datatype=DataType.FLOAT_VECTOR,
+                dim=self._embedding_dim,
+            )
+
+            vannadoc_index_params = self.milvus_client.prepare_index_params()
+            vannadoc_index_params.add_index(
+                field_name="vector",
+                index_name="vector",
+                index_type="AUTOINDEX",
+                metric_type="L2",
+            )
+            self.milvus_client.create_collection(
+                collection_name=name,
+                schema=vannadoc_schema,
+                index_params=vannadoc_index_params,
+                consistency_level="Strong",
+            )
+
+    def add_question_sql(self, question: str, sql: str, **kwargs) -> str:
+        if len(question) == 0 or len(sql) == 0:
+            raise Exception("pair of question and sql can not be null")
+        _id = str(uuid.uuid4()) + "-sql"
+        embedding = self.embedding_function.encode_documents([question])[0]
+        self.milvus_client.insert(
+            collection_name="vannasql",
+            data={"id": _id, "text": question, "sql": sql, "vector": embedding},
+        )
+        return _id
+
+    def add_ddl(self, ddl: str, **kwargs) -> str:
+        if len(ddl) == 0:
+            raise Exception("ddl can not be null")
+        _id = str(uuid.uuid4()) + "-ddl"
+        embedding = self.embedding_function.encode_documents([ddl])[0]
+        self.milvus_client.insert(
+            collection_name="vannaddl",
+            data={"id": _id, "ddl": ddl, "vector": embedding},
+        )
+        return _id
+
+    def add_documentation(self, documentation: str, **kwargs) -> str:
+        if len(documentation) == 0:
+            raise Exception("documentation can not be null")
+        _id = str(uuid.uuid4()) + "-doc"
+        embedding = self.embedding_function.encode_documents([documentation])[0]
+        self.milvus_client.insert(
+            collection_name="vannadoc",
+            data={"id": _id, "doc": documentation, "vector": embedding},
+        )
+        return _id
+
+    def get_training_data(self, **kwargs) -> pd.DataFrame:
+        sql_data = self.milvus_client.query(
+            collection_name="vannasql",
+            output_fields=["*"],
+            limit=MAX_LIMIT_SIZE,
+        )
+        df = pd.DataFrame()
+        df_sql = pd.DataFrame(
+            {
+                "id": [doc["id"] for doc in sql_data],
+                "question": [doc["text"] for doc in sql_data],
+                "content": [doc["sql"] for doc in sql_data],
+            }
+        )
+        df = pd.concat([df, df_sql])
+
+        ddl_data = self.milvus_client.query(
+            collection_name="vannaddl",
+            output_fields=["*"],
+            limit=MAX_LIMIT_SIZE,
+        )
+
+        df_ddl = pd.DataFrame(
+            {
+                "id": [doc["id"] for doc in ddl_data],
+                "question": [None for doc in ddl_data],
+                "content": [doc["ddl"] for doc in ddl_data],
+            }
+        )
+        df = pd.concat([df, df_ddl])
+
+        doc_data = self.milvus_client.query(
+            collection_name="vannadoc",
+            output_fields=["*"],
+            limit=MAX_LIMIT_SIZE,
+        )
+
+        df_doc = pd.DataFrame(
+            {
+                "id": [doc["id"] for doc in doc_data],
+                "question": [None for doc in doc_data],
+                "content": [doc["doc"] for doc in doc_data],
+            }
+        )
+        df = pd.concat([df, df_doc])
+        return df
+
+    def get_similar_question_sql(self, question: str, **kwargs) -> list:
+        search_params = {
+            "metric_type": "L2",
+            "params": {"nprobe": 128},
+        }
+        embeddings = self.embedding_function.encode_queries([question])
+        res = self.milvus_client.search(
+            collection_name="vannasql",
+            anns_field="vector",
+            data=embeddings,
+            limit=self.n_results,
+            output_fields=["text", "sql"],
+            search_params=search_params,
+        )
+        res = res[0]
+
+        list_sql = []
+        for doc in res:
+            dict = {}
+            dict["question"] = doc["entity"]["text"]
+            dict["sql"] = doc["entity"]["sql"]
+            list_sql.append(dict)
+        return list_sql
+
+    def get_related_ddl(self, question: str, **kwargs) -> list:
+        search_params = {
+            "metric_type": "L2",
+            "params": {"nprobe": 128},
+        }
+        embeddings = self.embedding_function.encode_queries([question])
+        res = self.milvus_client.search(
+            collection_name="vannaddl",
+            anns_field="vector",
+            data=embeddings,
+            limit=self.n_results,
+            output_fields=["ddl"],
+            search_params=search_params,
+        )
+        res = res[0]
+
+        list_ddl = []
+        for doc in res:
+            list_ddl.append(doc["entity"]["ddl"])
+        return list_ddl
+
+    def get_related_documentation(self, question: str, **kwargs) -> list:
+        search_params = {
+            "metric_type": "L2",
+            "params": {"nprobe": 128},
+        }
+        embeddings = self.embedding_function.encode_queries([question])
+        res = self.milvus_client.search(
+            collection_name="vannadoc",
+            anns_field="vector",
+            data=embeddings,
+            limit=self.n_results,
+            output_fields=["doc"],
+            search_params=search_params,
+        )
+        res = res[0]
+
+        list_doc = []
+        for doc in res:
+            list_doc.append(doc["entity"]["doc"])
+        return list_doc
+
+    def remove_training_data(self, id: str, **kwargs) -> bool:
+        if id.endswith("-sql"):
+            self.milvus_client.delete(collection_name="vannasql", ids=[id])
+            return True
+        elif id.endswith("-ddl"):
+            self.milvus_client.delete(collection_name="vannaddl", ids=[id])
+            return True
+        elif id.endswith("-doc"):
+            self.milvus_client.delete(collection_name="vannadoc", ids=[id])
+            return True
+        else:
+            return False

vanna/legacy/mistral/__init__.py

@@ -0,0 +1 @@
+from .mistral import Mistral

vanna/legacy/mistral/mistral.py

@@ -0,0 +1,51 @@
+import os
+
+from mistralai import Mistral as MistralClient
+from mistralai import UserMessage
+
+from ..base import VannaBase
+
+
+class Mistral(VannaBase):
+    def __init__(self, config=None):
+        if config is None:
+            raise ValueError(
+                "For Mistral, config must be provided with an api_key and model"
+            )
+
+        if "api_key" not in config:
+            raise ValueError("config must contain a Mistral api_key")
+
+        if "model" not in config:
+            raise ValueError("config must contain a Mistral model")
+
+        api_key = config["api_key"]
+        model = config["model"]
+        self.client = MistralClient(api_key=api_key)
+        self.model = model
+
+    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 generate_sql(self, question: str, **kwargs) -> str:
+        # Use the super generate_sql
+        sql = super().generate_sql(question, **kwargs)
+
+        # Replace "\_" with "_"
+        sql = sql.replace("\\_", "_")
+
+        return sql
+
+    def submit_prompt(self, prompt, **kwargs) -> str:
+        chat_response = self.client.chat.complete(
+            model=self.model,
+            messages=prompt,
+        )
+
+        return chat_response.choices[0].message.content

vanna/legacy/mock/__init__.py

@@ -0,0 +1,3 @@
+from .embedding import MockEmbedding
+from .llm import MockLLM
+from .vectordb import MockVectorDB

vanna/legacy/mock/embedding.py

@@ -0,0 +1,11 @@
+from typing import List
+
+from ..base import VannaBase
+
+
+class MockEmbedding(VannaBase):
+    def __init__(self, config=None):
+        pass
+
+    def generate_embedding(self, data: str, **kwargs) -> List[float]:
+        return [1.0, 2.0, 3.0, 4.0, 5.0]

vanna/legacy/mock/llm.py

@@ -0,0 +1,18 @@
+from ..base import VannaBase
+
+
+class MockLLM(VannaBase):
+    def __init__(self, config=None):
+        pass
+
+    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 submit_prompt(self, prompt, **kwargs) -> str:
+        return "Mock LLM response"

vanna/legacy/mock/vectordb.py

@@ -0,0 +1,67 @@
+import pandas as pd
+
+from ..base import VannaBase
+
+
+class MockVectorDB(VannaBase):
+    def __init__(self, config=None):
+        pass
+
+    def _get_id(self, value: str, **kwargs) -> str:
+        # Hash the value and return the ID
+        return str(hash(value))
+
+    def add_ddl(self, ddl: str, **kwargs) -> str:
+        return self._get_id(ddl)
+
+    def add_documentation(self, doc: str, **kwargs) -> str:
+        return self._get_id(doc)
+
+    def add_question_sql(self, question: str, sql: str, **kwargs) -> str:
+        return self._get_id(question)
+
+    def get_related_ddl(self, question: str, **kwargs) -> list:
+        return []
+
+    def get_related_documentation(self, question: str, **kwargs) -> list:
+        return []
+
+    def get_similar_question_sql(self, question: str, **kwargs) -> list:
+        return []
+
+    def get_training_data(self, **kwargs) -> pd.DataFrame:
+        return pd.DataFrame(
+            {
+                "id": {
+                    0: "19546-ddl",
+                    1: "91597-sql",
+                    2: "133976-sql",
+                    3: "59851-doc",
+                    4: "73046-sql",
+                },
+                "training_data_type": {
+                    0: "ddl",
+                    1: "sql",
+                    2: "sql",
+                    3: "documentation",
+                    4: "sql",
+                },
+                "question": {
+                    0: None,
+                    1: "What are the top selling genres?",
+                    2: "What are the low 7 artists by sales?",
+                    3: None,
+                    4: "What is the total sales for each customer?",
+                },
+                "content": {
+                    0: "CREATE TABLE [Invoice]\n(\n    [InvoiceId] INTEGER  NOT NULL,\n    [CustomerId] INTEGER  NOT NULL,\n    [InvoiceDate] DATETIME  NOT NULL,\n    [BillingAddress] NVARCHAR(70),\n    [BillingCity] NVARCHAR(40),\n    [BillingState] NVARCHAR(40),\n    [BillingCountry] NVARCHAR(40),\n    [BillingPostalCode] NVARCHAR(10),\n    [Total] NUMERIC(10,2)  NOT NULL,\n    CONSTRAINT [PK_Invoice] PRIMARY KEY  ([InvoiceId]),\n    FOREIGN KEY ([CustomerId]) REFERENCES [Customer] ([CustomerId]) \n\t\tON DELETE NO ACTION ON UPDATE NO ACTION\n)",
+                    1: "SELECT g.Name AS Genre, SUM(il.Quantity) AS TotalSales\nFROM Genre g\nJOIN Track t ON g.GenreId = t.GenreId\nJOIN InvoiceLine il ON t.TrackId = il.TrackId\nGROUP BY g.GenreId, g.Name\nORDER BY TotalSales DESC;",
+                    2: "SELECT a.ArtistId, a.Name, SUM(il.Quantity) AS TotalSales\nFROM Artist a\nINNER JOIN Album al ON a.ArtistId = al.ArtistId\nINNER JOIN Track t ON al.AlbumId = t.AlbumId\nINNER JOIN InvoiceLine il ON t.TrackId = il.TrackId\nGROUP BY a.ArtistId, a.Name\nORDER BY TotalSales ASC\nLIMIT 7;",
+                    3: "This is a SQLite database. For dates rememeber to use SQLite syntax.",
+                    4: "SELECT c.CustomerId, c.FirstName, c.LastName, SUM(i.Total) AS TotalSales\nFROM Customer c\nJOIN Invoice i ON c.CustomerId = i.CustomerId\nGROUP BY c.CustomerId, c.FirstName, c.LastName;",
+                },
+            }
+        )
+
+    def remove_training_data(id: str, **kwargs) -> bool:
+        return True

vanna/legacy/ollama/__init__.py

@@ -0,0 +1 @@
+from .ollama import Ollama

vanna/legacy/ollama/ollama.py

@@ -0,0 +1,110 @@
+import json
+import re
+
+from httpx import Timeout
+
+from ..base import VannaBase
+from ..exceptions import DependencyError
+
+
+class Ollama(VannaBase):
+    def __init__(self, config=None):
+        try:
+            ollama = __import__("ollama")
+        except ImportError:
+            raise DependencyError(
+                "You need to install required dependencies to execute this method, run command:"
+                " \npip install ollama"
+            )
+
+        if not config:
+            raise ValueError("config must contain at least Ollama model")
+        if "model" not in config.keys():
+            raise ValueError("config must contain at least Ollama model")
+        self.host = config.get("ollama_host", "http://localhost:11434")
+        self.model = config["model"]
+        if ":" not in self.model:
+            self.model += ":latest"
+
+        self.ollama_timeout = config.get("ollama_timeout", 240.0)
+
+        self.ollama_client = ollama.Client(
+            self.host, timeout=Timeout(self.ollama_timeout)
+        )
+        self.keep_alive = config.get("keep_alive", None)
+        self.ollama_options = config.get("options", {})
+        self.num_ctx = self.ollama_options.get("num_ctx", 2048)
+        self.__pull_model_if_ne(self.ollama_client, self.model)
+
+    @staticmethod
+    def __pull_model_if_ne(ollama_client, model):
+        model_response = ollama_client.list()
+        model_lists = [
+            model_element["model"] for model_element in model_response.get("models", [])
+        ]
+        if model not in model_lists:
+            ollama_client.pull(model)
+
+    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(self, llm_response):
+        """
+        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:
+        - llm_response (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.
+        """
+        # Remove ollama-generated extra characters
+        llm_response = llm_response.replace("\\_", "_")
+        llm_response = llm_response.replace("\\", "")
+
+        # Regular expression to find ```sql' and capture until '```'
+        sql = re.search(r"```sql\n((.|\n)*?)(?=;|\[|```)", llm_response, re.DOTALL)
+        # Regular expression to find 'select, with (ignoring case) and capture until ';', [ (this happens in case of mistral) or end of string
+        select_with = re.search(
+            r"(select|(with.*?as \())(.*?)(?=;|\[|```)",
+            llm_response,
+            re.IGNORECASE | re.DOTALL,
+        )
+        if sql:
+            self.log(f"Output from LLM: {llm_response} \nExtracted SQL: {sql.group(1)}")
+            return sql.group(1).replace("```", "")
+        elif select_with:
+            self.log(
+                f"Output from LLM: {llm_response} \nExtracted SQL: {select_with.group(0)}"
+            )
+            return select_with.group(0)
+        else:
+            return llm_response
+
+    def submit_prompt(self, prompt, **kwargs) -> str:
+        self.log(
+            f"Ollama parameters:\n"
+            f"model={self.model},\n"
+            f"options={self.ollama_options},\n"
+            f"keep_alive={self.keep_alive}"
+        )
+        self.log(f"Prompt Content:\n{json.dumps(prompt, ensure_ascii=False)}")
+        response_dict = self.ollama_client.chat(
+            model=self.model,
+            messages=prompt,
+            stream=False,
+            options=self.ollama_options,
+            keep_alive=self.keep_alive,
+        )
+
+        self.log(f"Ollama Response:\n{str(response_dict)}")
+
+        return response_dict["message"]["content"]

vanna/legacy/openai/__init__.py

@@ -0,0 +1,2 @@
+from .openai_chat import OpenAI_Chat
+from .openai_embeddings import OpenAI_Embeddings