asksql 2.1.0__py3-none-any.whl

vanna/legacy/pgvector/pgvector.py

@@ -0,0 +1,282 @@
+import ast
+import json
+import logging
+import uuid
+
+import pandas as pd
+from langchain_core.documents import Document
+from langchain_postgres.vectorstores import PGVector
+from sqlalchemy import create_engine, text
+
+from .. import ValidationError
+from ..base import VannaBase
+from ..types import TrainingPlan, TrainingPlanItem
+
+
+class PG_VectorStore(VannaBase):
+    def __init__(self, config=None):
+        if not config or "connection_string" not in config:
+            raise ValueError(
+                "A valid 'config' dictionary with a 'connection_string' is required."
+            )
+
+        VannaBase.__init__(self, config=config)
+
+        if config and "connection_string" in config:
+            self.connection_string = config.get("connection_string")
+            self.n_results = config.get("n_results", 10)
+
+        if config and "embedding_function" in config:
+            self.embedding_function = config.get("embedding_function")
+        else:
+            from langchain_huggingface import HuggingFaceEmbeddings
+
+            self.embedding_function = HuggingFaceEmbeddings(
+                model_name="all-MiniLM-L6-v2"
+            )
+
+        self.sql_collection = PGVector(
+            embeddings=self.embedding_function,
+            collection_name="sql",
+            connection=self.connection_string,
+        )
+        self.ddl_collection = PGVector(
+            embeddings=self.embedding_function,
+            collection_name="ddl",
+            connection=self.connection_string,
+        )
+        self.documentation_collection = PGVector(
+            embeddings=self.embedding_function,
+            collection_name="documentation",
+            connection=self.connection_string,
+        )
+
+    def add_question_sql(self, question: str, sql: str, **kwargs) -> str:
+        question_sql_json = json.dumps(
+            {
+                "question": question,
+                "sql": sql,
+            },
+            ensure_ascii=False,
+        )
+        id = str(uuid.uuid4()) + "-sql"
+        createdat = kwargs.get("createdat")
+        doc = Document(
+            page_content=question_sql_json,
+            metadata={"id": id, "createdat": createdat},
+        )
+        self.sql_collection.add_documents([doc], ids=[doc.metadata["id"]])
+
+        return id
+
+    def add_ddl(self, ddl: str, **kwargs) -> str:
+        _id = str(uuid.uuid4()) + "-ddl"
+        doc = Document(
+            page_content=ddl,
+            metadata={"id": _id},
+        )
+        self.ddl_collection.add_documents([doc], ids=[doc.metadata["id"]])
+        return _id
+
+    def add_documentation(self, documentation: str, **kwargs) -> str:
+        _id = str(uuid.uuid4()) + "-doc"
+        doc = Document(
+            page_content=documentation,
+            metadata={"id": _id},
+        )
+        self.documentation_collection.add_documents([doc], ids=[doc.metadata["id"]])
+        return _id
+
+    def get_collection(self, collection_name):
+        match collection_name:
+            case "sql":
+                return self.sql_collection
+            case "ddl":
+                return self.ddl_collection
+            case "documentation":
+                return self.documentation_collection
+            case _:
+                raise ValueError("Specified collection does not exist.")
+
+    def get_similar_question_sql(self, question: str) -> list:
+        documents = self.sql_collection.similarity_search(
+            query=question, k=self.n_results
+        )
+        return [ast.literal_eval(document.page_content) for document in documents]
+
+    def get_related_ddl(self, question: str, **kwargs) -> list:
+        documents = self.ddl_collection.similarity_search(
+            query=question, k=self.n_results
+        )
+        return [document.page_content for document in documents]
+
+    def get_related_documentation(self, question: str, **kwargs) -> list:
+        documents = self.documentation_collection.similarity_search(
+            query=question, k=self.n_results
+        )
+        return [document.page_content for document in documents]
+
+    def train(
+        self,
+        question: str | None = None,
+        sql: str | None = None,
+        ddl: str | None = None,
+        documentation: str | None = None,
+        plan: TrainingPlan | None = None,
+        createdat: str | None = None,
+    ):
+        if question and not sql:
+            raise ValidationError("Please provide a SQL query.")
+
+        if documentation:
+            logging.info(f"Adding documentation: {documentation}")
+            return self.add_documentation(documentation)
+
+        if sql and question:
+            return self.add_question_sql(
+                question=question, sql=sql, createdat=createdat
+            )
+
+        if ddl:
+            logging.info(f"Adding ddl: {ddl}")
+            return self.add_ddl(ddl)
+
+        if plan:
+            for item in plan._plan:
+                if item.item_type == TrainingPlanItem.ITEM_TYPE_DDL:
+                    self.add_ddl(item.item_value)
+                elif item.item_type == TrainingPlanItem.ITEM_TYPE_IS:
+                    self.add_documentation(item.item_value)
+                elif (
+                    item.item_type == TrainingPlanItem.ITEM_TYPE_SQL and item.item_name
+                ):
+                    self.add_question_sql(question=item.item_name, sql=item.item_value)
+
+    def get_training_data(self, **kwargs) -> pd.DataFrame:
+        # Establishing the connection
+        engine = create_engine(self.connection_string)
+
+        # Querying the 'langchain_pg_embedding' table
+        query_embedding = "SELECT cmetadata, document FROM langchain_pg_embedding"
+        df_embedding = pd.read_sql(query_embedding, engine)
+
+        # List to accumulate the processed rows
+        processed_rows = []
+
+        # Process each row in the DataFrame
+        for _, row in df_embedding.iterrows():
+            custom_id = row["cmetadata"]["id"]
+            document = row["document"]
+            training_data_type = (
+                "documentation" if custom_id[-3:] == "doc" else custom_id[-3:]
+            )
+
+            if training_data_type == "sql":
+                # Convert the document string to a dictionary
+                try:
+                    doc_dict = ast.literal_eval(document)
+                    question = doc_dict.get("question")
+                    content = doc_dict.get("sql")
+                except (ValueError, SyntaxError):
+                    logging.info(
+                        f"Skipping row with custom_id {custom_id} due to parsing error."
+                    )
+                    continue
+            elif training_data_type in ["documentation", "ddl"]:
+                question = None  # Default value for question
+                content = document
+            else:
+                # If the suffix is not recognized, skip this row
+                logging.info(
+                    f"Skipping row with custom_id {custom_id} due to unrecognized training data type."
+                )
+                continue
+
+            # Append the processed data to the list
+            processed_rows.append(
+                {
+                    "id": custom_id,
+                    "question": question,
+                    "content": content,
+                    "training_data_type": training_data_type,
+                }
+            )
+
+        # Create a DataFrame from the list of processed rows
+        df_processed = pd.DataFrame(processed_rows)
+
+        return df_processed
+
+    def remove_training_data(self, id: str, **kwargs) -> bool:
+        # Create the database engine
+        engine = create_engine(self.connection_string)
+
+        # SQL DELETE statement
+        delete_statement = text(
+            """
+            DELETE FROM langchain_pg_embedding
+            WHERE cmetadata ->> 'id' = :id
+        """
+        )
+
+        # Connect to the database and execute the delete statement
+        with engine.connect() as connection:
+            # Start a transaction
+            with connection.begin() as transaction:
+                try:
+                    result = connection.execute(delete_statement, {"id": id})
+                    # Commit the transaction if the delete was successful
+                    transaction.commit()
+                    # Check if any row was deleted and return True or False accordingly
+                    return result.rowcount > 0
+                except Exception as e:
+                    # Rollback the transaction in case of error
+                    logging.error(f"An error occurred: {e}")
+                    transaction.rollback()
+                    return False
+
+    def remove_collection(self, collection_name: str) -> bool:
+        engine = create_engine(self.connection_string)
+
+        # Determine the suffix to look for based on the collection name
+        suffix_map = {"ddl": "ddl", "sql": "sql", "documentation": "doc"}
+        suffix = suffix_map.get(collection_name)
+
+        if not suffix:
+            logging.info(
+                "Invalid collection name. Choose from 'ddl', 'sql', or 'documentation'."
+            )
+            return False
+
+        # SQL query to delete rows based on the condition
+        query = text(
+            f"""
+            DELETE FROM langchain_pg_embedding
+            WHERE cmetadata->>'id' LIKE '%{suffix}'
+        """
+        )
+
+        # Execute the deletion within a transaction block
+        with engine.connect() as connection:
+            with connection.begin() as transaction:
+                try:
+                    result = connection.execute(query)
+                    transaction.commit()  # Explicitly commit the transaction
+                    if result.rowcount > 0:
+                        logging.info(
+                            f"Deleted {result.rowcount} rows from "
+                            f"langchain_pg_embedding where collection is {collection_name}."
+                        )
+                        return True
+                    else:
+                        logging.info(
+                            f"No rows deleted for collection {collection_name}."
+                        )
+                        return False
+                except Exception as e:
+                    logging.error(f"An error occurred: {e}")
+                    transaction.rollback()  # Rollback in case of error
+                    return False
+
+    def generate_embedding(self, *args, **kwargs):
+        pass

vanna/legacy/pinecone/__init__.py

@@ -0,0 +1,3 @@
+from .pinecone_vector import PineconeDB_VectorStore
+
+__all__ = ["PineconeDB_VectorStore"]

vanna/legacy/pinecone/pinecone_vector.py

@@ -0,0 +1,275 @@
+import json
+from typing import List
+
+from pinecone import Pinecone, PodSpec, ServerlessSpec
+import pandas as pd
+from ..base import VannaBase
+from ..utils import deterministic_uuid
+
+from fastembed import TextEmbedding
+
+
+class PineconeDB_VectorStore(VannaBase):
+    """
+    Vectorstore using PineconeDB
+
+    Args:
+        config (dict): Configuration dictionary. Defaults to {}. You must provide either a Pinecone Client or an API key in the config.
+            - client (Pinecone, optional): Pinecone client. Defaults to None.
+            - api_key (str, optional): Pinecone API key. Defaults to None.
+            - n_results (int, optional): Number of results to return. Defaults to 10.
+            - dimensions (int, optional): Dimensions of the embeddings. Defaults to 384 which coresponds to the dimensions of BAAI/bge-small-en-v1.5.
+            - fastembed_model (str, optional): Fastembed model to use. Defaults to "BAAI/bge-small-en-v1.5".
+            - documentation_namespace (str, optional): Namespace for documentation. Defaults to "documentation".
+            - distance_metric (str, optional): Distance metric to use. Defaults to "cosine".
+            - ddl_namespace (str, optional): Namespace for DDL. Defaults to "ddl".
+            - sql_namespace (str, optional): Namespace for SQL. Defaults to "sql".
+            - index_name (str, optional): Name of the index. Defaults to "vanna-index".
+            - metadata_config (dict, optional): Metadata configuration if using a pinecone pod. Defaults to {}.
+            - server_type (str, optional): Type of Pinecone server to use. Defaults to "serverless". Options are "serverless" or "pod".
+            - podspec (PodSpec, optional): PodSpec configuration if using a pinecone pod. Defaults to PodSpec(environment="us-west-2", pod_type="p1.x1", metadata_config=self.metadata_config).
+            - serverless_spec (ServerlessSpec, optional): ServerlessSpec configuration if using a pinecone serverless index. Defaults to ServerlessSpec(cloud="aws", region="us-west-2").
+    Raises:
+        ValueError: If config is None, api_key is not provided OR client is not provided, client is not an instance of Pinecone, or server_type is not "serverless" or "pod".
+    """
+
+    def __init__(self, config=None):
+        VannaBase.__init__(self, config=config)
+        if config is None:
+            raise ValueError(
+                "config is required, pass either a Pinecone client or an API key in the config."
+            )
+        client = config.get("client")
+        api_key = config.get("api_key")
+        if not api_key and not client:
+            raise ValueError(
+                "api_key is required in config or pass a configured client"
+            )
+        if not client and api_key:
+            self._client = Pinecone(api_key=api_key)
+        elif not isinstance(client, Pinecone):
+            raise ValueError("client must be an instance of Pinecone")
+        else:
+            self._client = client
+
+        self.n_results = config.get("n_results", 10)
+        self.dimensions = config.get("dimensions", 384)
+        self.fastembed_model = config.get("fastembed_model", "BAAI/bge-small-en-v1.5")
+        self.documentation_namespace = config.get(
+            "documentation_namespace", "documentation"
+        )
+        self.distance_metric = config.get("distance_metric", "cosine")
+        self.ddl_namespace = config.get("ddl_namespace", "ddl")
+        self.sql_namespace = config.get("sql_namespace", "sql")
+        self.index_name = config.get("index_name", "vanna-index")
+        self.metadata_config = config.get("metadata_config", {})
+        self.server_type = config.get("server_type", "serverless")
+        if self.server_type not in ["serverless", "pod"]:
+            raise ValueError("server_type must be either 'serverless' or 'pod'")
+        self.podspec = config.get(
+            "podspec",
+            PodSpec(
+                environment="us-west-2",
+                pod_type="p1.x1",
+                metadata_config=self.metadata_config,
+            ),
+        )
+        self.serverless_spec = config.get(
+            "serverless_spec", ServerlessSpec(cloud="aws", region="us-west-2")
+        )
+        self._setup_index()
+
+    def _set_index_host(self, host: str) -> None:
+        self.Index = self._client.Index(host=host)
+
+    def _setup_index(self) -> None:
+        existing_indexes = self._get_indexes()
+        if self.index_name not in existing_indexes and self.server_type == "serverless":
+            self._client.create_index(
+                name=self.index_name,
+                dimension=self.dimensions,
+                metric=self.distance_metric,
+                spec=self.serverless_spec,
+            )
+            pinecone_index_host = self._client.describe_index(self.index_name)["host"]
+            self._set_index_host(pinecone_index_host)
+        elif self.index_name not in existing_indexes and self.server_type == "pod":
+            self._client.create_index(
+                name=self.index_name,
+                dimension=self.dimensions,
+                metric=self.distance_metric,
+                spec=self.podspec,
+            )
+            pinecone_index_host = self._client.describe_index(self.index_name)["host"]
+            self._set_index_host(pinecone_index_host)
+        else:
+            pinecone_index_host = self._client.describe_index(self.index_name)["host"]
+            self._set_index_host(pinecone_index_host)
+
+    def _get_indexes(self) -> list:
+        return [index["name"] for index in self._client.list_indexes()]
+
+    def _check_if_embedding_exists(self, id: str, namespace: str) -> bool:
+        fetch_response = self.Index.fetch(ids=[id], namespace=namespace)
+        if fetch_response.vectors == {}:
+            return False
+        return True
+
+    def add_ddl(self, ddl: str, **kwargs) -> str:
+        id = deterministic_uuid(ddl) + "-ddl"
+        if self._check_if_embedding_exists(id=id, namespace=self.ddl_namespace):
+            print(f"DDL with id: {id} already exists in the index. Skipping...")
+            return id
+        self.Index.upsert(
+            vectors=[(id, self.generate_embedding(ddl), {"ddl": ddl})],
+            namespace=self.ddl_namespace,
+        )
+        return id
+
+    def add_documentation(self, doc: str, **kwargs) -> str:
+        id = deterministic_uuid(doc) + "-doc"
+
+        if self._check_if_embedding_exists(
+            id=id, namespace=self.documentation_namespace
+        ):
+            print(
+                f"Documentation with id: {id} already exists in the index. Skipping..."
+            )
+            return id
+        self.Index.upsert(
+            vectors=[(id, self.generate_embedding(doc), {"documentation": doc})],
+            namespace=self.documentation_namespace,
+        )
+        return id
+
+    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"
+        if self._check_if_embedding_exists(id=id, namespace=self.sql_namespace):
+            print(
+                f"Question-SQL with id: {id} already exists in the index. Skipping..."
+            )
+            return id
+        self.Index.upsert(
+            vectors=[
+                (
+                    id,
+                    self.generate_embedding(question_sql_json),
+                    {"sql": question_sql_json},
+                )
+            ],
+            namespace=self.sql_namespace,
+        )
+        return id
+
+    def get_related_ddl(self, question: str, **kwargs) -> list:
+        res = self.Index.query(
+            namespace=self.ddl_namespace,
+            vector=self.generate_embedding(question),
+            top_k=self.n_results,
+            include_values=True,
+            include_metadata=True,
+        )
+        return [match["metadata"]["ddl"] for match in res["matches"]] if res else []
+
+    def get_related_documentation(self, question: str, **kwargs) -> list:
+        res = self.Index.query(
+            namespace=self.documentation_namespace,
+            vector=self.generate_embedding(question),
+            top_k=self.n_results,
+            include_values=True,
+            include_metadata=True,
+        )
+        return (
+            [match["metadata"]["documentation"] for match in res["matches"]]
+            if res
+            else []
+        )
+
+    def get_similar_question_sql(self, question: str, **kwargs) -> list:
+        res = self.Index.query(
+            namespace=self.sql_namespace,
+            vector=self.generate_embedding(question),
+            top_k=self.n_results,
+            include_values=True,
+            include_metadata=True,
+        )
+        return (
+            [
+                {
+                    key: value
+                    for key, value in json.loads(match["metadata"]["sql"]).items()
+                }
+                for match in res["matches"]
+            ]
+            if res
+            else []
+        )
+
+    def get_training_data(self, **kwargs) -> pd.DataFrame:
+        # Pinecone does not support getting all vectors in a namespace, so we have to query for the top_k vectors with a dummy vector
+        df = pd.DataFrame()
+        namespaces = {
+            "sql": self.sql_namespace,
+            "ddl": self.ddl_namespace,
+            "documentation": self.documentation_namespace,
+        }
+
+        for data_type, namespace in namespaces.items():
+            data = self.Index.query(
+                top_k=10000,  # max results that pinecone allows
+                namespace=namespace,
+                include_values=True,
+                include_metadata=True,
+                vector=[0.0] * self.dimensions,
+            )
+
+            if data is not None:
+                id_list = [match["id"] for match in data["matches"]]
+                content_list = [
+                    match["metadata"][data_type] for match in data["matches"]
+                ]
+                question_list = [
+                    (
+                        json.loads(match["metadata"][data_type])["question"]
+                        if data_type == "sql"
+                        else None
+                    )
+                    for match in data["matches"]
+                ]
+
+                df_data = pd.DataFrame(
+                    {
+                        "id": id_list,
+                        "question": question_list,
+                        "content": content_list,
+                    }
+                )
+                df_data["training_data_type"] = data_type
+                df = pd.concat([df, df_data])
+
+        return df
+
+    def remove_training_data(self, id: str, **kwargs) -> bool:
+        if id.endswith("-sql"):
+            self.Index.delete(ids=[id], namespace=self.sql_namespace)
+            return True
+        elif id.endswith("-ddl"):
+            self.Index.delete(ids=[id], namespace=self.ddl_namespace)
+            return True
+        elif id.endswith("-doc"):
+            self.Index.delete(ids=[id], namespace=self.documentation_namespace)
+            return True
+        else:
+            return False
+
+    def generate_embedding(self, data: str, **kwargs) -> List[float]:
+        embedding_model = TextEmbedding(model_name=self.fastembed_model)
+        embedding = next(embedding_model.embed(data))
+        return embedding.tolist()

vanna/legacy/qdrant/__init__.py

@@ -0,0 +1,3 @@
+from .qdrant import Qdrant_VectorStore
+
+__all__ = ["Qdrant_VectorStore"]