From df24227bc87c103f7f0cfa3548408ac3d178f84c Mon Sep 17 00:00:00 2001 From: jannesDoege Date: Fri, 4 Oct 2024 12:34:56 +0200 Subject: [PATCH 1/6] Fixed typo in ai-sql-accuracy-2023-08-17.md The relevance search is depicted in a green box, not a red box. --- papers/ai-sql-accuracy-2023-08-17.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/papers/ai-sql-accuracy-2023-08-17.md b/papers/ai-sql-accuracy-2023-08-17.md index faf6a3b9..082b5edb 100644 --- a/papers/ai-sql-accuracy-2023-08-17.md +++ b/papers/ai-sql-accuracy-2023-08-17.md @@ -229,7 +229,7 @@ By providing just those 3 example queries, we see substantial improvements to th Enterprise data warehouses often contain 100s (or even 1000s) of tables, and an order of magnitude more queries that cover all the use cases within their organizations. Given the limited size of the context windows of modern LLMs, we can’t just shove all the prior queries and schema definitions into the prompt. -Our final approach to context is a more sophisticated ML approach - load embeddings of prior queries and the table schemas into a vector database, and only choose the most relevant queries / tables to the question asked. Here's a diagram of what we are doing - note the contextual relevance search in the red box - +Our final approach to context is a more sophisticated ML approach - load embeddings of prior queries and the table schemas into a vector database, and only choose the most relevant queries / tables to the question asked. Here's a diagram of what we are doing - note the contextual relevance search in the green box - ![](https://raw.githubusercontent.com/vanna-ai/vanna/main/papers/img/using-contextually-relevant-examples.png) From 55563c83a68fbc55dd18b23d0aa09b5a006bc385 Mon Sep 17 00:00:00 2001 From: Gui Quental Date: Tue, 15 Oct 2024 10:21:11 -0300 Subject: [PATCH 2/6] Update gemini version to 1.5 --- src/vanna/google/gemini_chat.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/vanna/google/gemini_chat.py b/src/vanna/google/gemini_chat.py index b33ce5bd..95bffb70 100644 --- a/src/vanna/google/gemini_chat.py +++ b/src/vanna/google/gemini_chat.py @@ -15,7 +15,7 @@ def __init__(self, config=None): if "model_name" in config: model_name = config["model_name"] else: - model_name = "gemini-1.0-pro" + model_name = "gemini-1.5-pro" self.google_api_key = None From baf3dbeec1b10c3bb2111f39f5a9d53ebaaa1f39 Mon Sep 17 00:00:00 2001 From: Gui Quental Date: Tue, 15 Oct 2024 10:21:24 -0300 Subject: [PATCH 3/6] Change SDK to GA instead of preview --- src/vanna/google/gemini_chat.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/vanna/google/gemini_chat.py b/src/vanna/google/gemini_chat.py index 95bffb70..60732472 100644 --- a/src/vanna/google/gemini_chat.py +++ b/src/vanna/google/gemini_chat.py @@ -30,7 +30,7 @@ def __init__(self, config=None): self.chat_model = genai.GenerativeModel(model_name) else: # Authenticate using VertexAI - from vertexai.preview.generative_models import GenerativeModel + from vertexai.generative_models import GenerativeModel self.chat_model = GenerativeModel(model_name) def system_message(self, message: str) -> any: From bd7d55e5b25d41e77c593588ef6bbb4645af1f5f Mon Sep 17 00:00:00 2001 From: Gui Quental Date: Tue, 15 Oct 2024 13:46:34 -0300 Subject: [PATCH 4/6] Add support to embeddings from VertexAI --- src/vanna/google/bigquery_vector.py | 49 ++++++++++++++++++++--------- 1 file changed, 34 insertions(+), 15 deletions(-) diff --git a/src/vanna/google/bigquery_vector.py b/src/vanna/google/bigquery_vector.py index 09cbf391..df68835d 100644 --- a/src/vanna/google/bigquery_vector.py +++ b/src/vanna/google/bigquery_vector.py @@ -2,6 +2,10 @@ import os import uuid from typing import List, Optional +from vertexai.language_models import ( + TextEmbeddingInput, + TextEmbeddingModel +) import pandas as pd from google.cloud import bigquery @@ -23,17 +27,15 @@ def __init__(self, config: dict, **kwargs): or set as an environment variable, assign it. """ print("Configuring genai") + self.type = "GEMINI" import google.generativeai as genai genai.configure(api_key=config["api_key"]) self.genai = genai else: + self.type = "VERTEX_AI" # Authenticate using VertexAI - from vertexai.language_models import ( - TextEmbeddingInput, - TextEmbeddingModel, - ) if self.config.get("project_id"): self.project_id = self.config.get("project_id") @@ -139,25 +141,42 @@ def fetch_similar_training_data(self, training_data_type: str, question: str, n_ results = self.conn.query(query).result().to_dataframe() return results - def generate_question_embedding(self, data: str, **kwargs) -> List[float]: - result = self.genai.embed_content( + def get_embeddings(self, data: str, task: str) -> List[float]: + embeddings = None + + if self.type == "VERTEX_AI": + input = [TextEmbeddingInput(data, task)] + model = TextEmbeddingModel.from_pretrained("text-embedding-004") + + result = model.get_embeddings(input) + + if len(result) > 0: + embeddings = result[0].values + else: + # Use Gemini Consumer API + result = self.genai.embed_content( model="models/text-embedding-004", content=data, - task_type="retrieval_query") + task_type=task) - if 'embedding' in result: - return result['embedding'] + if 'embedding' in result: + embeddings = result['embedding'] + + return embeddings + + def generate_question_embedding(self, data: str, **kwargs) -> List[float]: + result = self.get_embeddings(data, "RETRIEVAL_QUERY") + + if result != None: + return result else: raise ValueError("No embeddings returned") def generate_storage_embedding(self, data: str, **kwargs) -> List[float]: - result = self.genai.embed_content( - model="models/text-embedding-004", - content=data, - task_type="retrieval_document") + result = self.get_embeddings(data, "RETRIEVAL_DOCUMENT") - if 'embedding' in result: - return result['embedding'] + if result != None: + return result else: raise ValueError("No embeddings returned") From e93e5d39ee6bc285dedfaae2bd551d396ef9d997 Mon Sep 17 00:00:00 2001 From: eux Date: Wed, 23 Oct 2024 12:01:29 +0800 Subject: [PATCH 5/6] fix: some tests due to missing modules --- pyproject.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pyproject.toml b/pyproject.toml index 4a4fb927..af25d8ca 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -33,7 +33,7 @@ bigquery = ["google-cloud-bigquery"] snowflake = ["snowflake-connector-python"] duckdb = ["duckdb"] google = ["google-generativeai", "google-cloud-aiplatform"] -all = ["psycopg2-binary", "db-dtypes", "PyMySQL", "google-cloud-bigquery", "snowflake-connector-python", "duckdb", "openai", "qianfan", "mistralai>=1.0.0", "chromadb", "anthropic", "zhipuai", "marqo", "google-generativeai", "google-cloud-aiplatform", "qdrant-client", "fastembed", "ollama", "httpx", "opensearch-py", "opensearch-dsl", "transformers", "pinecone-client", "pymilvus[model]","weaviate-client", "azure-search-documents", "azure-identity", "azure-common", "faiss-cpu", "boto", "botocore"] +all = ["psycopg2-binary", "db-dtypes", "PyMySQL", "google-cloud-bigquery", "snowflake-connector-python", "duckdb", "openai", "qianfan", "mistralai>=1.0.0", "chromadb", "anthropic", "zhipuai", "marqo", "google-generativeai", "google-cloud-aiplatform", "qdrant-client", "fastembed", "ollama", "httpx", "opensearch-py", "opensearch-dsl", "transformers", "pinecone-client", "pymilvus[model]","weaviate-client", "azure-search-documents", "azure-identity", "azure-common", "faiss-cpu", "boto", "boto3", "botocore", "langchain_core", "langchain_postgres"] test = ["tox"] chromadb = ["chromadb"] openai = ["openai"] From 932131f7e571f6a28ddf1ec1dcfe8ac206ed7c0a Mon Sep 17 00:00:00 2001 From: Zain Hoda <7146154+zainhoda@users.noreply.github.com> Date: Wed, 23 Oct 2024 09:02:07 -0400 Subject: [PATCH 6/6] remove pgvector tests --- tests/test_pgvector.py | 28 ++++++++++++++-------------- 1 file changed, 14 insertions(+), 14 deletions(-) diff --git a/tests/test_pgvector.py b/tests/test_pgvector.py index 4bc1dea9..8c9344a5 100644 --- a/tests/test_pgvector.py +++ b/tests/test_pgvector.py @@ -2,23 +2,23 @@ from dotenv import load_dotenv -from vanna.pgvector import PG_VectorStore +# from vanna.pgvector import PG_VectorStore load_dotenv() +# Removing thiese tests for now until the dependencies are sorted out +# def get_vanna_connection_string(): +# server = os.environ.get("PG_SERVER") +# driver = "psycopg" +# port = 5434 +# database = os.environ.get("PG_DATABASE") +# username = os.environ.get("PG_USERNAME") +# password = os.environ.get("PG_PASSWORD") -def get_vanna_connection_string(): - server = os.environ.get("PG_SERVER") - driver = "psycopg" - port = 5434 - database = os.environ.get("PG_DATABASE") - username = os.environ.get("PG_USERNAME") - password = os.environ.get("PG_PASSWORD") +# return f"postgresql+psycopg://{username}:{password}@{server}:{port}/{database}" - return f"postgresql+psycopg://{username}:{password}@{server}:{port}/{database}" - -def test_pgvector(): - connection_string = get_vanna_connection_string() - pgclient = PG_VectorStore(config={"connection_string": connection_string}) - assert pgclient is not None +# def test_pgvector(): +# connection_string = get_vanna_connection_string() +# pgclient = PG_VectorStore(config={"connection_string": connection_string}) +# assert pgclient is not None