From 7eac4b9175edcccd633ed720e09663f04767e431 Mon Sep 17 00:00:00 2001 From: Jason Schleifer Date: Tue, 23 Jul 2024 03:29:54 +1200 Subject: [PATCH] fixed error loading VectorStoreDrivers if no OPENAI_API_KEY --- .../gtUIAmazonOpenSearchVectorStoreDriver.py | 10 +++++----- .../gtUIAzureMongoDbVectorStoreDriver.py | 10 +++++----- nodes/drivers/gtUIBaseVectorStoreDriver.py | 20 +++++++++++++++---- nodes/drivers/gtUILocalVectorStoreDriver.py | 9 +++++---- 4 files changed, 31 insertions(+), 18 deletions(-) diff --git a/nodes/drivers/gtUIAmazonOpenSearchVectorStoreDriver.py b/nodes/drivers/gtUIAmazonOpenSearchVectorStoreDriver.py index bd4d7fb..3949573 100644 --- a/nodes/drivers/gtUIAmazonOpenSearchVectorStoreDriver.py +++ b/nodes/drivers/gtUIAmazonOpenSearchVectorStoreDriver.py @@ -1,11 +1,11 @@ import os -from griptape.drivers import AmazonOpenSearchVectorStoreDriver, OpenAiEmbeddingDriver +from griptape.drivers import ( + AmazonOpenSearchVectorStoreDriver, +) from .gtUIBaseVectorStoreDriver import gtUIBaseVectorStoreDriver -default_embedding_driver = OpenAiEmbeddingDriver() - DEFAULT_HOST_ENV = "AMAZON_OPENSEARCH_HOST" DEFAULT_INDEX_ENV = "AMAZON_OPENSEARCH_INDEX_NAME" DEFAULT_AWS_ACCESS_KEY_ID = "AWS_ACCESS_KEY_ID" @@ -42,7 +42,7 @@ class gtUIAmazonOpenSearchVectorStoreDriver(gtUIBaseVectorStoreDriver): return inputs def create(self, **kwargs): - embedding_driver = kwargs.get("embedding_driver", default_embedding_driver) + embedding_driver = kwargs.get("embedding_driver", None) host_env = kwargs.get("host_env", DEFAULT_HOST_ENV) index_env_var = kwargs.get("index_name_env", DEFAULT_INDEX_ENV) @@ -58,6 +58,6 @@ class gtUIAmazonOpenSearchVectorStoreDriver(gtUIBaseVectorStoreDriver): if embedding_driver: params["embedding_driver"] = embedding_driver else: - params["embedding_driver"] = default_embedding_driver + params["embedding_driver"] = self.get_default_embedding_driver() driver = AmazonOpenSearchVectorStoreDriver(**params) return (driver,) diff --git a/nodes/drivers/gtUIAzureMongoDbVectorStoreDriver.py b/nodes/drivers/gtUIAzureMongoDbVectorStoreDriver.py index 6136af9..4689c95 100644 --- a/nodes/drivers/gtUIAzureMongoDbVectorStoreDriver.py +++ b/nodes/drivers/gtUIAzureMongoDbVectorStoreDriver.py @@ -1,9 +1,9 @@ -from griptape.drivers import AzureMongoDbVectorStoreDriver, OpenAiEmbeddingDriver +from griptape.drivers import ( + AzureMongoDbVectorStoreDriver, +) from .gtUIBaseVectorStoreDriver import gtUIBaseVectorStoreDriver -default_embedding_driver = OpenAiEmbeddingDriver() - DEFAULT_HOST_ENV = "AZURE_MONGODB_HOST" DEFAULT_USERNAME_ENV = "AZURE_MONGODB_USERNAME" DEFAULT_PASSWORD_ENV = "AZURE_MONGODB_PASSWORD" @@ -38,7 +38,7 @@ class gtUIAzureMongoDbVectorStoreDriver(gtUIBaseVectorStoreDriver): return inputs def create(self, **kwargs): - embedding_driver = kwargs.get("embedding_driver", default_embedding_driver) + embedding_driver = kwargs.get("embedding_driver", None) host_env = kwargs.get("host_env", DEFAULT_HOST_ENV) username_env = kwargs.get("username_env", DEFAULT_USERNAME_ENV) password_env = kwargs.get("password_env", DEFAULT_PASSWORD_ENV) @@ -78,6 +78,6 @@ class gtUIAzureMongoDbVectorStoreDriver(gtUIBaseVectorStoreDriver): if embedding_driver: params["embedding_driver"] = embedding_driver else: - params["embedding_driver"] = default_embedding_driver + params["embedding_driver"] = self.get_default_embedding_driver() driver = AzureMongoDbVectorStoreDriver(**params) return (driver,) diff --git a/nodes/drivers/gtUIBaseVectorStoreDriver.py b/nodes/drivers/gtUIBaseVectorStoreDriver.py index 1505ab4..d0c2af3 100644 --- a/nodes/drivers/gtUIBaseVectorStoreDriver.py +++ b/nodes/drivers/gtUIBaseVectorStoreDriver.py @@ -1,8 +1,14 @@ -from griptape.drivers import BaseVectorStoreDriver, OpenAiEmbeddingDriver +import os + +from griptape.drivers import ( + BaseVectorStoreDriver, + DummyEmbeddingDriver, + OpenAiEmbeddingDriver, +) from .gtUIBaseDriver import gtUIBaseDriver -default_embedding_driver = OpenAiEmbeddingDriver() +# default_embedding_driver = OpenAiEmbeddingDriver() class gtUIBaseVectorStoreDriver(gtUIBaseDriver): @@ -21,14 +27,20 @@ class gtUIBaseVectorStoreDriver(gtUIBaseDriver): CATEGORY = "Griptape/Drivers/Vector Store" RETURN_TYPES = ("VECTOR_STORE_DRIVER",) + def get_default_embedding_driver(self): + if os.getenv("OPENAI_API_KEY"): + return OpenAiEmbeddingDriver() + else: + return DummyEmbeddingDriver() + def create(self, **kwargs): - embedding_driver = kwargs.get("embedding_driver", default_embedding_driver) + embedding_driver = kwargs.get("embedding_driver", None) params = {} if embedding_driver: params["embedding_driver"] = embedding_driver else: - params["embedding_driver"] = default_embedding_driver + params["embedding_driver"] = self.get_default_embedding_driver() driver = BaseVectorStoreDriver(**params) return (driver,) diff --git a/nodes/drivers/gtUILocalVectorStoreDriver.py b/nodes/drivers/gtUILocalVectorStoreDriver.py index 5a2602a..1fd1c36 100644 --- a/nodes/drivers/gtUILocalVectorStoreDriver.py +++ b/nodes/drivers/gtUILocalVectorStoreDriver.py @@ -1,8 +1,9 @@ -from griptape.drivers import LocalVectorStoreDriver, OpenAiEmbeddingDriver +from griptape.drivers import ( + LocalVectorStoreDriver, +) from .gtUIBaseVectorStoreDriver import gtUIBaseVectorStoreDriver -default_embedding_driver = OpenAiEmbeddingDriver() default_filename = "griptape_local_vector_file.txt" @@ -26,7 +27,7 @@ class gtUILocalVectorStoreDriver(gtUIBaseVectorStoreDriver): return inputs def create(self, **kwargs): - embedding_driver = kwargs.get("embedding_driver", default_embedding_driver) + embedding_driver = kwargs.get("embedding_driver", None) persist_file = kwargs.get("persist_file", False) persist_filename = kwargs.get("persist_filename", default_filename) @@ -34,7 +35,7 @@ class gtUILocalVectorStoreDriver(gtUIBaseVectorStoreDriver): if embedding_driver: params["embedding_driver"] = embedding_driver else: - params["embedding_driver"] = default_embedding_driver + params["embedding_driver"] = self.get_default_embedding_driver() if persist_file and persist_filename: params["persist_file"] = persist_filename