vanna 0.7.8__py3-none-any.whl → 2.0.0__py3-none-any.whl

vanna/opensearch/opensearch_vector.py

@@ -1,392 +0,0 @@
-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(query)
-    # 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})

vanna/opensearch/opensearch_vector_semantic.py

@@ -1,175 +0,0 @@
-import json
-
-import pandas as pd
-from langchain_community.vectorstores import OpenSearchVectorSearch
-
-from ..base import VannaBase
-from ..utils import deterministic_uuid
-
-
-class OpenSearch_Semantic_VectorStore(VannaBase):
-  def __init__(self, config=None):
-    VannaBase.__init__(self, config=config)
-    if config is None:
-      config = {}
-
-    if "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.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))
-
-    self.document_index = config.get("es_document_index", "vanna_document_index")
-    self.ddl_index = config.get("es_ddl_index", "vanna_ddl_index")
-    self.question_sql_index = config.get("es_question_sql_index", "vanna_questions_sql_index")
-
-    self.log(f"OpenSearch_Semantic_VectorStore initialized with document_index: {self.document_index}, ddl_index: {self.ddl_index}, question_sql_index: {self.question_sql_index}")
-
-    es_urls = config.get("es_urls", "https://localhost:9200")
-    ssl = config.get("es_ssl", True)
-    verify_certs = config.get("es_verify_certs", True)
-
-    if "es_user" in config:
-      auth = (config["es_user"], config["es_password"])
-    else:
-      auth = None
-
-    headers = config.get("es_headers", None)
-    timeout = config.get("es_timeout", 60)
-    max_retries = config.get("es_max_retries", 10)
-
-    common_args = {
-        "opensearch_url": es_urls,
-        "embedding_function": self.embedding_function,
-        "engine": "faiss",
-        "http_auth": auth,
-        "use_ssl": ssl,
-        "verify_certs": verify_certs,
-        "timeout": timeout,
-        "max_retries": max_retries,
-        "retry_on_timeout": True,
-        "headers": headers,
-    }
-
-    self.documentation_store = OpenSearchVectorSearch(index_name=self.document_index, **common_args)
-    self.ddl_store = OpenSearchVectorSearch(index_name=self.ddl_index, **common_args)
-    self.sql_store = OpenSearchVectorSearch(index_name=self.question_sql_index, **common_args)
-
-  def add_ddl(self, ddl: str, **kwargs) -> str:
-    _id = deterministic_uuid(ddl) + "-ddl"
-    self.ddl_store.add_texts(texts=[ddl], ids=[_id], **kwargs)
-    return _id
-
-  def add_documentation(self, documentation: str, **kwargs) -> str:
-    _id = deterministic_uuid(documentation) + "-doc"
-    self.documentation_store.add_texts(texts=[documentation], ids=[_id], **kwargs)
-    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"
-    self.sql_store.add_texts(texts=[question_sql_json], ids=[_id], **kwargs)
-    return _id
-
-  def get_related_ddl(self, question: str, **kwargs) -> list:
-    documents = self.ddl_store.similarity_search(query=question, k=self.n_results_ddl)
-    return [document.page_content for document in documents]
-
-  def get_related_documentation(self, question: str, **kwargs) -> list:
-    documents = self.documentation_store.similarity_search(query=question, k=self.n_results_documentation)
-    return [document.page_content for document in documents]
-
-  def get_similar_question_sql(self, question: str, **kwargs) -> list:
-    documents = self.sql_store.similarity_search(query=question, k=self.n_results_sql)
-    return [json.loads(document.page_content) for document in documents]
-
-  def get_training_data(self, **kwargs) -> pd.DataFrame:
-    data = []
-    query = {
-      "query": {
-        "match_all": {}
-      }
-    }
-
-    indices = [
-      {"index": self.document_index, "type": "documentation"},
-      {"index": self.question_sql_index, "type": "sql"},
-      {"index": self.ddl_index, "type": "ddl"},
-    ]
-
-    # Use documentation_store.client consistently for search on all indices
-    opensearch_client = self.documentation_store.client
-
-    for index_info in indices:
-      index_name = index_info["index"]
-      training_data_type = index_info["type"]
-      scroll = '1m'  # keep scroll context for 1 minute
-      response = opensearch_client.search(
-        index=index_name,
-        ignore_unavailable=True,
-        body=query,
-        scroll=scroll,
-        size=1000
-      )
-
-      scroll_id = response.get('_scroll_id')
-
-      while scroll_id:
-        hits = response['hits']['hits']
-        if not hits:
-          break  # No more hits, exit loop
-
-        for hit in hits:
-          source = hit['_source']
-          if training_data_type == "sql":
-            try:
-              doc_dict = json.loads(source['text'])
-              content = doc_dict.get("sql")
-              question = doc_dict.get("question")
-            except json.JSONDecodeError as e:
-              self.log(f"Skipping row with custom_id {hit['_id']} due to JSON parsing error: {e}","Error")
-              continue
-          else:  # documentation or ddl
-            content = source['text']
-            question = None
-
-          data.append({
-            "id": hit["_id"],
-            "training_data_type": training_data_type,
-            "question": question,
-            "content": content,
-          })
-
-        # Get next batch of results, using documentation_store.client.scroll
-        response = opensearch_client.scroll(scroll_id=scroll_id, scroll=scroll)
-        scroll_id = response.get('_scroll_id')
-
-    return pd.DataFrame(data)
-
-  def remove_training_data(self, id: str, **kwargs) -> bool:
-    try:
-      if id.endswith("-sql"):
-        return self.sql_store.delete(ids=[id], **kwargs)
-      elif id.endswith("-ddl"):
-        return self.ddl_store.delete(ids=[id], **kwargs)
-      elif id.endswith("-doc"):
-        return self.documentation_store.delete(ids=[id], **kwargs)
-      else:
-        return False
-    except Exception as e:
-      self.log(f"Error deleting training dataError deleting training data: {e}", "Error")
-      return False
-
-  def generate_embedding(self, data: str, **kwargs) -> list[float]:
-    pass