include openai, bedrock, anthropic test

Author: Klaijan

_test_unstructured_client/integration/test_integration_freemium.py

@@ -14,7 +14,7 @@
 
 
 # FREEMIUM_URL = "https://api.unstructured.io"
-FREEMIUM_URL = "http://127.0.0.1:5000"
+FREEMIUM_URL = "http://127.0.0.1:8009"
 
 
 @pytest.fixture(scope="module")
@@ -236,12 +236,13 @@ async def call_api():
     assert len(elements) > 0
 
 
-# TODO (klaijan) - still cannot test against vlm if the updated client vlm is generated
+# TODO (klaijan) - still cannot test against vlm if the updated client vlm is not generated
 # `vlm_model cannot be empty for vlm requests.` because partition_parameters do not have the param.
+# TODO (klaijan) - test with more models and providers
 @pytest.mark.parametrize("split_pdf", [True, False])
 @pytest.mark.parametrize("vlm_model", ["gpt-4o"])
 @pytest.mark.parametrize("vlm_model_provider", ["openai"])
-def test_partition_strategy_vlm(split_pdf, vlm_model, vlm_model_provider, client, doc_path):
+def test_partition_strategy_vlm_openai(split_pdf, vlm_model, vlm_model_provider, client, doc_path):
     filename = "layout-parser-paper-fast.pdf"
     with open(doc_path / filename, "rb") as f:
         files = shared.Files(
@@ -255,8 +256,49 @@ def test_partition_strategy_vlm(split_pdf, vlm_model, vlm_model_provider, client
             strategy="vlm",
             vlm_model=vlm_model,
             vlm_model_provider=vlm_model_provider,
-            # languages=["eng"],
-            # split_pdf_page=split_pdf,
+            languages=["eng"],
+            split_pdf_page=split_pdf,
+        )
+    )
+
+    response = client.general.partition(
+        server_url=FREEMIUM_URL,
+        request=req
+    )
+    assert response.status_code == 200
+    assert len(response.elements)
+
+
+@pytest.mark.parametrize("split_pdf", [True, False])
+@pytest.mark.parametrize("vlm_model", 
+    [
+        "us.amazon.nova-pro-v1:0", 
+        "us.amazon.nova-lite-v1:0", 
+        "us.anthropic.claude-3-5-sonnet-20241022-v2:0",
+        "us.anthropic.claude-3-opus-20240229-v1:0",
+        "us.anthropic.claude-3-haiku-20240307-v1:0",
+        "us.anthropic.claude-3-sonnet-20240229-v1:0",
+        "us.meta.llama3-2-90b-instruct-v1:0",
+        "us.meta.llama3-2-11b-instruct-v1:0",
+    ]
+)
+@pytest.mark.parametrize("vlm_model_provider", ["bedrock"])
+def test_partition_strategy_vlm_bedrock(split_pdf, vlm_model, vlm_model_provider, client, doc_path):
+    filename = "layout-parser-paper-fast.pdf"
+    with open(doc_path / filename, "rb") as f:
+        files = shared.Files(
+            content=f.read(),
+            file_name=filename,
+        )
+
+    req = operations.PartitionRequest(
+        partition_parameters=shared.PartitionParameters(
+            files=files,
+            strategy="vlm",
+            vlm_model=vlm_model,
+            vlm_model_provider=vlm_model_provider,
+            languages=["eng"],
+            split_pdf_page=split_pdf,
         )
     )
 
@@ -265,4 +307,38 @@ def test_partition_strategy_vlm(split_pdf, vlm_model, vlm_model_provider, client
         request=req
     )
     assert response.status_code == 200
-    assert len(response.elements)
+    assert len(response.elements)
+
+
+@pytest.mark.parametrize("split_pdf", [True, False])
+@pytest.mark.parametrize("vlm_model", 
+    [
+        "claude-3-5-sonnet-20241022",
+    ]
+)
+@pytest.mark.parametrize("vlm_model_provider", ["anthropic"])
+def test_partition_strategy_vlm_anthropic(split_pdf, vlm_model, vlm_model_provider, client, doc_path):
+    filename = "layout-parser-paper-fast.pdf"
+    with open(doc_path / filename, "rb") as f:
+        files = shared.Files(
+            content=f.read(),
+            file_name=filename,
+        )
+
+    req = operations.PartitionRequest(
+        partition_parameters=shared.PartitionParameters(
+            files=files,
+            strategy="vlm",
+            vlm_model=vlm_model,
+            vlm_model_provider=vlm_model_provider,
+            languages=["eng"],
+            split_pdf_page=split_pdf,
+        )
+    )
+
+    response = client.general.partition(
+        server_url=FREEMIUM_URL,
+        request=req
+    )
+    assert response.status_code == 200
+    assert len(response.elements)