asksql 2.1.0__py3-none-any.whl

vanna/legacy/openai/openai_chat.py

@@ -0,0 +1,124 @@
+import os
+
+from openai import OpenAI
+
+from ..base import VannaBase
+
+
+class OpenAI_Chat(VannaBase):
+    def __init__(self, client=None, config=None):
+        VannaBase.__init__(self, config=config)
+
+        # default parameters - can be overrided using config
+        self.temperature = 0.7
+
+        if "temperature" in config:
+            self.temperature = config["temperature"]
+
+        if "api_type" in config:
+            raise Exception(
+                "Passing api_type is now deprecated. Please pass an OpenAI client instead."
+            )
+
+        if "api_base" in config:
+            raise Exception(
+                "Passing api_base is now deprecated. Please pass an OpenAI client instead."
+            )
+
+        if "api_version" in config:
+            raise Exception(
+                "Passing api_version is now deprecated. Please pass an OpenAI client instead."
+            )
+
+        if client is not None:
+            self.client = client
+            return
+
+        if config is None and client is None:
+            self.client = OpenAI(api_key=os.getenv("OPENAI_API_KEY"))
+            return
+
+        if "api_key" in config:
+            self.client = OpenAI(api_key=config["api_key"])
+
+    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:
+        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
+
+        if kwargs.get("model", None) is not None:
+            model = kwargs.get("model", None)
+            print(f"Using model {model} for {num_tokens} tokens (approx)")
+            response = self.client.chat.completions.create(
+                model=model,
+                messages=prompt,
+                stop=None,
+                temperature=self.temperature,
+            )
+        elif kwargs.get("engine", None) is not None:
+            engine = kwargs.get("engine", None)
+            print(f"Using model {engine} for {num_tokens} tokens (approx)")
+            response = self.client.chat.completions.create(
+                engine=engine,
+                messages=prompt,
+                stop=None,
+                temperature=self.temperature,
+            )
+        elif self.config is not None and "engine" in self.config:
+            print(
+                f"Using engine {self.config['engine']} for {num_tokens} tokens (approx)"
+            )
+            response = self.client.chat.completions.create(
+                engine=self.config["engine"],
+                messages=prompt,
+                stop=None,
+                temperature=self.temperature,
+            )
+        elif self.config is not None and "model" in self.config:
+            print(
+                f"Using model {self.config['model']} for {num_tokens} tokens (approx)"
+            )
+            response = self.client.chat.completions.create(
+                model=self.config["model"],
+                messages=prompt,
+                stop=None,
+                temperature=self.temperature,
+            )
+        else:
+            if num_tokens > 3500:
+                model = "gpt-3.5-turbo-16k"
+            else:
+                model = "gpt-3.5-turbo"
+
+            print(f"Using model {model} for {num_tokens} tokens (approx)")
+            response = self.client.chat.completions.create(
+                model=model,
+                messages=prompt,
+                stop=None,
+                temperature=self.temperature,
+            )
+
+        # Find the first response from the chatbot that has text in it (some responses may not have text)
+        for choice in response.choices:
+            if "text" in choice:
+                return choice.text
+
+        # If no response with text is found, return the first response's content (which may be empty)
+        return response.choices[0].message.content

vanna/legacy/openai/openai_embeddings.py

@@ -0,0 +1,46 @@
+from openai import OpenAI
+
+from ..base import VannaBase
+
+
+class OpenAI_Embeddings(VannaBase):
+    def __init__(self, client=None, config=None):
+        VannaBase.__init__(self, config=config)
+
+        if client is not None:
+            self.client = client
+            return
+
+        if self.client is not None:
+            return
+
+        self.client = OpenAI()
+
+        if config is None:
+            return
+
+        if "api_type" in config:
+            self.client.api_type = config["api_type"]
+
+        if "api_base" in config:
+            self.client.api_base = config["api_base"]
+
+        if "api_version" in config:
+            self.client.api_version = config["api_version"]
+
+        if "api_key" in config:
+            self.client.api_key = config["api_key"]
+
+    def generate_embedding(self, data: str, **kwargs) -> list[float]:
+        if self.config is not None and "engine" in self.config:
+            embedding = self.client.embeddings.create(
+                engine=self.config["engine"],
+                input=data,
+            )
+        else:
+            embedding = self.client.embeddings.create(
+                model="text-embedding-ada-002",
+                input=data,
+            )
+
+        return embedding.get("data")[0]["embedding"]

vanna/legacy/opensearch/__init__.py

@@ -0,0 +1,2 @@
+from .opensearch_vector import OpenSearch_VectorStore
+from .opensearch_vector_semantic import OpenSearch_Semantic_VectorStore

vanna/legacy/opensearch/opensearch_vector.py

@@ -0,0 +1,369 @@
+import base64
+import uuid
+from typing import List
+
+import pandas as pd
+from opensearchpy import OpenSearch
+
+from ..base import VannaBase
+
+
+class OpenSearch_VectorStore(VannaBase):
+    def __init__(self, config=None):
+        VannaBase.__init__(self, config=config)
+        document_index = "vanna_document_index"
+        ddl_index = "vanna_ddl_index"
+        question_sql_index = "vanna_questions_sql_index"
+        if config is not None and "es_document_index" in config:
+            document_index = config["es_document_index"]
+        if config is not None and "es_ddl_index" in config:
+            ddl_index = config["es_ddl_index"]
+        if config is not None and "es_question_sql_index" in config:
+            question_sql_index = config["es_question_sql_index"]
+
+        self.document_index = document_index
+        self.ddl_index = ddl_index
+        self.question_sql_index = question_sql_index
+        print(
+            "OpenSearch_VectorStore initialized with document_index: ",
+            document_index,
+            " ddl_index: ",
+            ddl_index,
+            " question_sql_index: ",
+            question_sql_index,
+        )
+
+        document_index_settings = {
+            "settings": {"index": {"number_of_shards": 6, "number_of_replicas": 2}},
+            "mappings": {
+                "properties": {
+                    "question": {
+                        "type": "text",
+                    },
+                    "doc": {
+                        "type": "text",
+                    },
+                }
+            },
+        }
+
+        ddl_index_settings = {
+            "settings": {"index": {"number_of_shards": 6, "number_of_replicas": 2}},
+            "mappings": {
+                "properties": {
+                    "ddl": {
+                        "type": "text",
+                    },
+                    "doc": {
+                        "type": "text",
+                    },
+                }
+            },
+        }
+
+        question_sql_index_settings = {
+            "settings": {"index": {"number_of_shards": 6, "number_of_replicas": 2}},
+            "mappings": {
+                "properties": {
+                    "question": {
+                        "type": "text",
+                    },
+                    "sql": {
+                        "type": "text",
+                    },
+                }
+            },
+        }
+
+        if config is not None and "es_document_index_settings" in config:
+            document_index_settings = config["es_document_index_settings"]
+        if config is not None and "es_ddl_index_settings" in config:
+            ddl_index_settings = config["es_ddl_index_settings"]
+        if config is not None and "es_question_sql_index_settings" in config:
+            question_sql_index_settings = config["es_question_sql_index_settings"]
+
+        self.document_index_settings = document_index_settings
+        self.ddl_index_settings = ddl_index_settings
+        self.question_sql_index_settings = question_sql_index_settings
+
+        es_urls = None
+        if config is not None and "es_urls" in config:
+            es_urls = config["es_urls"]
+
+        # Host and port
+        if config is not None and "es_host" in config:
+            host = config["es_host"]
+        else:
+            host = "localhost"
+
+        if config is not None and "es_port" in config:
+            port = config["es_port"]
+        else:
+            port = 9200
+
+        if config is not None and "es_ssl" in config:
+            ssl = config["es_ssl"]
+        else:
+            ssl = False
+
+        if config is not None and "es_verify_certs" in config:
+            verify_certs = config["es_verify_certs"]
+        else:
+            verify_certs = False
+
+        # Authentication
+        if config is not None and "es_user" in config:
+            auth = (config["es_user"], config["es_password"])
+        else:
+            # Default to admin:admin
+            auth = None
+
+        headers = None
+        # base64 authentication
+        if (
+            config is not None
+            and "es_encoded_base64" in config
+            and "es_user" in config
+            and "es_password" in config
+        ):
+            if config["es_encoded_base64"]:
+                encoded_credentials = base64.b64encode(
+                    (config["es_user"] + ":" + config["es_password"]).encode("utf-8")
+                ).decode("utf-8")
+                headers = {"Authorization": "Basic " + encoded_credentials}
+                # remove auth from config
+                auth = None
+
+        # custom headers
+        if config is not None and "es_headers" in config:
+            headers = config["es_headers"]
+
+        if config is not None and "es_timeout" in config:
+            timeout = config["es_timeout"]
+        else:
+            timeout = 60
+
+        if config is not None and "es_max_retries" in config:
+            max_retries = config["es_max_retries"]
+        else:
+            max_retries = 10
+
+        if config is not None and "es_http_compress" in config:
+            es_http_compress = config["es_http_compress"]
+        else:
+            es_http_compress = False
+
+        print(
+            "OpenSearch_VectorStore initialized with es_urls: ",
+            es_urls,
+            " host: ",
+            host,
+            " port: ",
+            port,
+            " ssl: ",
+            ssl,
+            " verify_certs: ",
+            verify_certs,
+            " timeout: ",
+            timeout,
+            " max_retries: ",
+            max_retries,
+        )
+        if es_urls is not None:
+            # Initialize the OpenSearch client by passing a list of URLs
+            self.client = OpenSearch(
+                hosts=[es_urls],
+                http_compress=es_http_compress,
+                use_ssl=ssl,
+                verify_certs=verify_certs,
+                timeout=timeout,
+                max_retries=max_retries,
+                retry_on_timeout=True,
+                http_auth=auth,
+                headers=headers,
+            )
+        else:
+            # Initialize the OpenSearch client by passing a host and port
+            self.client = OpenSearch(
+                hosts=[{"host": host, "port": port}],
+                http_compress=es_http_compress,
+                use_ssl=ssl,
+                verify_certs=verify_certs,
+                timeout=timeout,
+                max_retries=max_retries,
+                retry_on_timeout=True,
+                http_auth=auth,
+                headers=headers,
+            )
+
+            print("OpenSearch_VectorStore initialized with client over ")
+
+        # 执行一个简单的查询来检查连接
+        try:
+            print("Connected to OpenSearch cluster:")
+            info = self.client.info()
+            print("OpenSearch cluster info:", info)
+        except Exception as e:
+            print("Error connecting to OpenSearch cluster:", e)
+
+        # Create the indices if they don't exist
+        self.create_index_if_not_exists(
+            self.document_index, self.document_index_settings
+        )
+        self.create_index_if_not_exists(self.ddl_index, self.ddl_index_settings)
+        self.create_index_if_not_exists(
+            self.question_sql_index, self.question_sql_index_settings
+        )
+
+    def create_index(self):
+        for index in [self.document_index, self.ddl_index, self.question_sql_index]:
+            try:
+                self.client.indices.create(index)
+            except Exception as e:
+                print("Error creating index: ", e)
+                print(f"opensearch index {index} already exists")
+                pass
+
+    def create_index_if_not_exists(self, index_name: str, index_settings: dict) -> bool:
+        try:
+            if not self.client.indices.exists(index_name):
+                print(f"Index {index_name} does not exist. Creating...")
+                self.client.indices.create(index=index_name, body=index_settings)
+                return True
+            else:
+                print(f"Index {index_name} already exists.")
+                return False
+        except Exception as e:
+            print(f"Error creating index: {index_name} ", e)
+            return False
+
+    def add_ddl(self, ddl: str, **kwargs) -> str:
+        # Assuming that you have a DDL index in your OpenSearch
+        id = str(uuid.uuid4()) + "-ddl"
+        ddl_dict = {"ddl": ddl}
+        response = self.client.index(
+            index=self.ddl_index, body=ddl_dict, id=id, **kwargs
+        )
+        return response["_id"]
+
+    def add_documentation(self, doc: str, **kwargs) -> str:
+        # Assuming you have a documentation index in your OpenSearch
+        id = str(uuid.uuid4()) + "-doc"
+        doc_dict = {"doc": doc}
+        response = self.client.index(
+            index=self.document_index, id=id, body=doc_dict, **kwargs
+        )
+        return response["_id"]
+
+    def add_question_sql(self, question: str, sql: str, **kwargs) -> str:
+        # Assuming you have a Questions and SQL index in your OpenSearch
+        id = str(uuid.uuid4()) + "-sql"
+        question_sql_dict = {"question": question, "sql": sql}
+        response = self.client.index(
+            index=self.question_sql_index, body=question_sql_dict, id=id, **kwargs
+        )
+        return response["_id"]
+
+    def get_related_ddl(self, question: str, **kwargs) -> List[str]:
+        # Assume you have some vector search mechanism associated with your data
+        query = {"query": {"match": {"ddl": question}}}
+        print(query)
+        response = self.client.search(index=self.ddl_index, body=query, **kwargs)
+        return [hit["_source"]["ddl"] for hit in response["hits"]["hits"]]
+
+    def get_related_documentation(self, question: str, **kwargs) -> List[str]:
+        query = {"query": {"match": {"doc": question}}}
+        print(query)
+        response = self.client.search(index=self.document_index, body=query, **kwargs)
+        return [hit["_source"]["doc"] for hit in response["hits"]["hits"]]
+
+    def get_similar_question_sql(self, question: str, **kwargs) -> List[str]:
+        query = {"query": {"match": {"question": question}}}
+        print(query)
+        response = self.client.search(
+            index=self.question_sql_index, body=query, **kwargs
+        )
+        return [
+            (hit["_source"]["question"], hit["_source"]["sql"])
+            for hit in response["hits"]["hits"]
+        ]
+
+    def get_training_data(self, **kwargs) -> pd.DataFrame:
+        # This will be a simple example pulling all data from an index
+        # WARNING: Do not use this approach in production for large indices!
+        data = []
+        response = self.client.search(
+            index=self.document_index, body={"query": {"match_all": {}}}, size=1000
+        )
+        print(response)
+        # records = [hit['_source'] for hit in response['hits']['hits']]
+        for hit in response["hits"]["hits"]:
+            data.append(
+                {
+                    "id": hit["_id"],
+                    "training_data_type": "documentation",
+                    "question": "",
+                    "content": hit["_source"]["doc"],
+                }
+            )
+
+        response = self.client.search(
+            index=self.question_sql_index, body={"query": {"match_all": {}}}, size=1000
+        )
+        # records = [hit['_source'] for hit in response['hits']['hits']]
+        for hit in response["hits"]["hits"]:
+            data.append(
+                {
+                    "id": hit["_id"],
+                    "training_data_type": "sql",
+                    "question": hit.get("_source", {}).get("question", ""),
+                    "content": hit.get("_source", {}).get("sql", ""),
+                }
+            )
+
+        response = self.client.search(
+            index=self.ddl_index, body={"query": {"match_all": {}}}, size=1000
+        )
+        # records = [hit['_source'] for hit in response['hits']['hits']]
+        for hit in response["hits"]["hits"]:
+            data.append(
+                {
+                    "id": hit["_id"],
+                    "training_data_type": "ddl",
+                    "question": "",
+                    "content": hit["_source"]["ddl"],
+                }
+            )
+
+        return pd.DataFrame(data)
+
+    def remove_training_data(self, id: str, **kwargs) -> bool:
+        try:
+            if id.endswith("-sql"):
+                self.client.delete(index=self.question_sql_index, id=id)
+                return True
+            elif id.endswith("-ddl"):
+                self.client.delete(index=self.ddl_index, id=id, **kwargs)
+                return True
+            elif id.endswith("-doc"):
+                self.client.delete(index=self.document_index, id=id, **kwargs)
+                return True
+            else:
+                return False
+        except Exception as e:
+            print("Error deleting training dataError deleting training data: ", e)
+            return False
+
+    def generate_embedding(self, data: str, **kwargs) -> list[float]:
+        # opensearch doesn't need to generate embeddings
+        pass
+
+
+# OpenSearch_VectorStore.__init__(self, config={'es_urls':
+# "https://opensearch-node.test.com:9200", 'es_encoded_base64': True, 'es_user':
+# "admin", 'es_password': "admin", 'es_verify_certs': True})
+
+
+# OpenSearch_VectorStore.__init__(self, config={'es_host':
+# "https://opensearch-node.test.com", 'es_port': 9200, 'es_user': "admin",
+# 'es_password': "admin", 'es_verify_certs': True})