Merge branch 'main' into bigquery-vector

vanna-ai · Oct 23, 2024 · d6d356f · d6d356f
2 parents bd7d55e + b7b030a
commit d6d356f
Show file tree

Hide file tree

Showing 3 changed files with 17 additions and 17 deletions.
diff --git 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"]

diff --git 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
 
@@ -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:

diff --git 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