feat/add precheck support for embedders (#473)

* add precheck support for embedders

* update tests

* fix bedrock

* fix syntax

Author: rbiseck3

CHANGELOG.md

@@ -1,3 +1,9 @@
+## 1.0.14
+
+### Enhancements
+
+* **Add precheck support for embedders that support listing models**
+
 ## 1.0.13
 
 ### Fixes

test/integration/connectors/test_google_drive.py

@@ -53,9 +53,7 @@ def google_drive_empty_folder(google_drive_connection_config):
     from googleapiclient.discovery import build
 
     access_config = google_drive_connection_config.access_config.get_secret_value()
-    creds = service_account.Credentials.from_service_account_info(
-        access_config.service_account_key
-    )
+    creds = service_account.Credentials.from_service_account_info(access_config.service_account_key)
     service = build("drive", "v3", credentials=creds)
 
     # Create an empty folder.
@@ -120,23 +118,16 @@ def test_google_drive_precheck_invalid_parameter(google_drive_connection_config)
         access_config=google_drive_connection_config.access_config,
     )
     index_config = GoogleDriveIndexerConfig(recursive=True)
-    indexer = GoogleDriveIndexer(
-        connection_config=connection_config, index_config=index_config
-    )
+    indexer = GoogleDriveIndexer(connection_config=connection_config, index_config=index_config)
     with pytest.raises(SourceConnectionError) as excinfo:
         indexer.precheck()
-    assert (
-        "invalid" in str(excinfo.value).lower()
-        or "not found" in str(excinfo.value).lower()
-    )
+    assert "invalid" in str(excinfo.value).lower() or "not found" in str(excinfo.value).lower()
 
 
 # Precheck fails due to lack of permission (simulate via monkeypatching).
 @pytest.mark.tags("google-drive", "precheck")
 @requires_env("GOOGLE_DRIVE_ID", "GOOGLE_DRIVE_SERVICE_KEY")
-def test_google_drive_precheck_no_permission(
-    google_drive_connection_config, monkeypatch
-):
+def test_google_drive_precheck_no_permission(google_drive_connection_config, monkeypatch):
     index_config = GoogleDriveIndexerConfig(recursive=True)
     indexer = GoogleDriveIndexer(
         connection_config=google_drive_connection_config,
@@ -153,10 +144,7 @@ def fake_get_root_info(files_client, object_id):
     monkeypatch.setattr(indexer, "get_root_info", fake_get_root_info)
     with pytest.raises(SourceConnectionError) as excinfo:
         indexer.precheck()
-    assert (
-        "forbidden" in str(excinfo.value).lower()
-        or "permission" in str(excinfo.value).lower()
-    )
+    assert "forbidden" in str(excinfo.value).lower() or "permission" in str(excinfo.value).lower()
 
 
 # Precheck fails when the folder is empty.
@@ -206,15 +194,10 @@ def test_google_drive_precheck_invalid_drive_id(google_drive_connection_config):
         access_config=google_drive_connection_config.access_config,
     )
     index_config = GoogleDriveIndexerConfig(recursive=True)
-    indexer = GoogleDriveIndexer(
-        connection_config=connection_config, index_config=index_config
-    )
+    indexer = GoogleDriveIndexer(connection_config=connection_config, index_config=index_config)
     with pytest.raises(SourceConnectionError) as excinfo:
         indexer.precheck()
-    assert (
-        "invalid" in str(excinfo.value).lower()
-        or "not found" in str(excinfo.value).lower()
-    )
+    assert "invalid" in str(excinfo.value).lower() or "not found" in str(excinfo.value).lower()
 
 
 MIME_TYPES_TO_TEST = [
@@ -244,19 +227,15 @@ async def test_google_drive_export_by_type(expected_mime, temp_dir):
     index_config = GoogleDriveIndexerConfig(recursive=True)
     download_config = GoogleDriveDownloaderConfig(download_dir=temp_dir)
 
-    indexer = GoogleDriveIndexer(
-        connection_config=connection_config, index_config=index_config
-    )
+    indexer = GoogleDriveIndexer(connection_config=connection_config, index_config=index_config)
     downloader = GoogleDriveDownloader(
         connection_config=connection_config, download_config=download_config
     )
 
     file_datas = list(indexer.run())
 
     # Filter only the target MIME type
-    target_files = [
-        f for f in file_datas if f.additional_metadata.get("mimeType") == expected_mime
-    ]
+    target_files = [f for f in file_datas if f.additional_metadata.get("mimeType") == expected_mime]
     assert target_files, f"No files found with MIME type: {expected_mime}"
 
     for file_data in target_files:

test/integration/embedders/test_azure_openai.py

@@ -38,6 +38,7 @@ def test_azure_openai_embedder(embedder_file: Path):
         embedding_azure_endpoint=azure_data.endpoint,
     )
     embedder = Embedder(config=embedder_config)
+    embedder.precheck()
     results = embedder.run(elements_filepath=embedder_file)
     assert results
     with embedder_file.open("r") as f:
@@ -54,4 +55,5 @@ def test_raw_azure_openai_embedder(embedder_file: Path):
             azure_endpoint=azure_data.endpoint,
         )
     )
+    embedder.precheck()
     validate_raw_embedder(embedder=embedder, embedder_file=embedder_file, expected_dimension=1536)

test/integration/embedders/test_bedrock.py

@@ -36,6 +36,7 @@ def test_bedrock_embedder(embedder_file: Path):
         embedding_aws_secret_access_key=aws_credentials["aws_secret_access_key"],
     )
     embedder = Embedder(config=embedder_config)
+    embedder.precheck()
     results = embedder.run(elements_filepath=embedder_file)
     assert results
     with embedder_file.open("r") as f:
@@ -52,6 +53,7 @@ def test_raw_bedrock_embedder(embedder_file: Path):
             aws_secret_access_key=aws_credentials["aws_secret_access_key"],
         )
     )
+    embedder.precheck()
     validate_raw_embedder(
         embedder=embedder,
         embedder_file=embedder_file,
@@ -82,7 +84,7 @@ def test_raw_bedrock_embedder_invalid_model(embedder_file: Path):
         )
     )
     with pytest.raises(UserError):
-        embedder.get_exemplary_embedding()
+        embedder.precheck()
 
 
 @requires_env("AWS_ACCESS_KEY_ID", "AWS_SECRET_ACCESS_KEY")
@@ -95,6 +97,7 @@ async def test_raw_async_bedrock_embedder(embedder_file: Path):
             aws_secret_access_key=aws_credentials["aws_secret_access_key"],
         )
     )
+    embedder.precheck()
     await validate_raw_embedder_async(
         embedder=embedder,
         embedder_file=embedder_file,

test/integration/embedders/test_huggingface.py

@@ -12,6 +12,7 @@
 def test_huggingface_embedder(embedder_file: Path):
     embedder_config = EmbedderConfig(embedding_provider="huggingface")
     embedder = Embedder(config=embedder_config)
+    embedder.precheck()
     results = embedder.run(elements_filepath=embedder_file)
     assert results
     with embedder_file.open("r") as f:
@@ -21,4 +22,5 @@ def test_huggingface_embedder(embedder_file: Path):
 
 def test_raw_hugginface_embedder(embedder_file: Path):
     embedder = HuggingFaceEmbeddingEncoder(config=HuggingFaceEmbeddingConfig())
+    embedder.precheck()
     validate_raw_embedder(embedder=embedder, embedder_file=embedder_file, expected_dimension=384)

test/integration/embedders/test_mixedbread.py

@@ -31,6 +31,7 @@ def test_mixedbread_embedder(embedder_file: Path):
     api_key = get_api_key()
     embedder_config = EmbedderConfig(embedding_provider="mixedbread-ai", embedding_api_key=api_key)
     embedder = Embedder(config=embedder_config)
+    embedder.precheck()
     results = embedder.run(elements_filepath=embedder_file)
     assert results
     with embedder_file.open("r") as f:
@@ -46,6 +47,7 @@ def test_raw_mixedbread_embedder(embedder_file: Path):
             api_key=api_key,
         )
     )
+    embedder.precheck()
     validate_raw_embedder(
         embedder=embedder,
         embedder_file=embedder_file,
@@ -63,6 +65,7 @@ async def test_raw_async_mixedbread_embedder(embedder_file: Path):
             api_key=api_key,
         )
     )
+    embedder.precheck()
     await validate_raw_embedder_async(
         embedder=embedder,
         embedder_file=embedder_file,

test/integration/embedders/test_octoai.py

@@ -15,7 +15,7 @@
     OctoAiEmbeddingConfig,
     OctoAIEmbeddingEncoder,
 )
-from unstructured_ingest.errors_v2 import UserAuthError
+from unstructured_ingest.errors_v2 import UserAuthError, UserError
 from unstructured_ingest.processes.embedder import Embedder, EmbedderConfig
 
 API_KEY = "OCTOAI_API_KEY"
@@ -32,6 +32,7 @@ def test_octoai_embedder(embedder_file: Path):
     api_key = get_api_key()
     embedder_config = EmbedderConfig(embedding_provider="octoai", embedding_api_key=api_key)
     embedder = Embedder(config=embedder_config)
+    embedder.precheck()
     results = embedder.run(elements_filepath=embedder_file)
     assert results
     with embedder_file.open("r") as f:
@@ -47,6 +48,7 @@ def test_raw_octoai_embedder(embedder_file: Path):
             api_key=api_key,
         )
     )
+    embedder.precheck()
     validate_raw_embedder(embedder=embedder, embedder_file=embedder_file, expected_dimension=1024)
 
 
@@ -58,7 +60,7 @@ def test_raw_octoai_embedder_invalid_credentials():
         )
     )
     with pytest.raises(UserAuthError):
-        embedder.get_exemplary_embedding()
+        embedder.precheck()
 
 
 @requires_env(API_KEY)
@@ -70,6 +72,17 @@ async def test_raw_async_octoai_embedder(embedder_file: Path):
             api_key=api_key,
         )
     )
+    embedder.precheck()
     await validate_raw_embedder_async(
         embedder=embedder, embedder_file=embedder_file, expected_dimension=1024
     )
+
+
+@requires_env(API_KEY)
+def test_octoai_wrong_model():
+    api_key = get_api_key()
+    embedder = OctoAIEmbeddingEncoder(
+        config=OctoAiEmbeddingConfig(api_key=api_key, model_name="fake_model_name")
+    )
+    with pytest.raises(UserError):
+        embedder.precheck()

test/integration/embedders/test_openai.py

@@ -15,7 +15,7 @@
     OpenAIEmbeddingConfig,
     OpenAIEmbeddingEncoder,
 )
-from unstructured_ingest.errors_v2 import UserAuthError
+from unstructured_ingest.errors_v2 import UserAuthError, UserError
 from unstructured_ingest.processes.embedder import Embedder, EmbedderConfig
 
 API_KEY = "OPENAI_API_KEY"
@@ -32,6 +32,7 @@ def test_openai_embedder(embedder_file: Path):
     api_key = get_api_key()
     embedder_config = EmbedderConfig(embedding_provider="openai", embedding_api_key=api_key)
     embedder = Embedder(config=embedder_config)
+    embedder.precheck()
     results = embedder.run(elements_filepath=embedder_file)
     assert results
     with embedder_file.open("r") as f:
@@ -47,6 +48,7 @@ def test_raw_openai_embedder(embedder_file: Path):
             api_key=api_key,
         )
     )
+    embedder.precheck()
     validate_raw_embedder(embedder=embedder, embedder_file=embedder_file, expected_dimension=1536)
 
 
@@ -69,6 +71,17 @@ async def test_raw_async_openai_embedder(embedder_file: Path):
             api_key=api_key,
         )
     )
+    embedder.precheck()
     await validate_raw_embedder_async(
         embedder=embedder, embedder_file=embedder_file, expected_dimension=1536
     )
+
+
+@requires_env(API_KEY)
+def test_openai_wrong_model():
+    api_key = get_api_key()
+    embedder = OpenAIEmbeddingEncoder(
+        config=OpenAIEmbeddingConfig(api_key=api_key, model_name="fake_model_name")
+    )
+    with pytest.raises(UserError):
+        embedder.precheck()

test/integration/embedders/test_togetherai.py

@@ -32,6 +32,7 @@ def test_togetherai_embedder(embedder_file: Path):
     api_key = get_api_key()
     embedder_config = EmbedderConfig(embedding_provider="togetherai", embedding_api_key=api_key)
     embedder = Embedder(config=embedder_config)
+    embedder.precheck()
     results = embedder.run(elements_filepath=embedder_file)
     assert results
     with embedder_file.open("r") as f:
@@ -43,6 +44,7 @@ def test_togetherai_embedder(embedder_file: Path):
 def test_raw_togetherai_embedder(embedder_file: Path):
     api_key = get_api_key()
     embedder = TogetherAIEmbeddingEncoder(config=TogetherAIEmbeddingConfig(api_key=api_key))
+    embedder.precheck()
     validate_raw_embedder(
         embedder=embedder,
         embedder_file=embedder_file,
@@ -63,6 +65,7 @@ def test_raw_togetherai_embedder_invalid_credentials():
 async def test_raw_async_togetherai_embedder(embedder_file: Path):
     api_key = get_api_key()
     embedder = AsyncTogetherAIEmbeddingEncoder(config=TogetherAIEmbeddingConfig(api_key=api_key))
+    embedder.precheck()
     await validate_raw_embedder_async(
         embedder=embedder,
         embedder_file=embedder_file,

test/integration/embedders/test_vertexai.py

@@ -31,6 +31,7 @@ def test_vertexai_embedder(embedder_file: Path):
     api_key = get_api_key()
     embedder_config = EmbedderConfig(embedding_provider="vertexai", embedding_api_key=api_key)
     embedder = Embedder(config=embedder_config)
+    embedder.precheck()
     results = embedder.run(elements_filepath=embedder_file)
     assert results
     with embedder_file.open("r") as f:
@@ -46,6 +47,7 @@ def test_raw_vertexai_embedder(embedder_file: Path):
             api_key=api_key,
         )
     )
+    embedder.precheck()
     validate_raw_embedder(embedder=embedder, embedder_file=embedder_file, expected_dimension=768)
 
 
@@ -58,6 +60,7 @@ async def test_raw_async_vertexai_embedder(embedder_file: Path):
             api_key=api_key,
         )
     )
+    embedder.precheck()
     await validate_raw_embedder_async(
         embedder=embedder, embedder_file=embedder_file, expected_dimension=768
     )