Add media describer and embeddings tests

Author: pamelafox

app/backend/prepdocslib/embeddings.py

@@ -244,16 +244,14 @@ async def create_embedding(self, image_bytes: bytes) -> list[float]:
         async with aiohttp.ClientSession(headers=headers) as session:
             async for attempt in AsyncRetrying(
                 retry=retry_if_exception_type(Exception),
-                    wait=wait_random_exponential(min=15, max=60),
-                    stop=stop_after_attempt(15),
-                    before_sleep=self.before_retry_sleep,
-                ):
+                wait=wait_random_exponential(min=15, max=60),
+                stop=stop_after_attempt(15),
+                before_sleep=self.before_retry_sleep,
+            ):
                 with attempt:
                     async with session.post(url=endpoint, params=params, data=image_bytes) as resp:
                         resp_json = await resp.json()
                         return resp_json["vector"]
-                    
-        return []
 
     def before_retry_sleep(self, retry_state):
         logger.info("Rate limited on the Vision embeddings API, sleeping before retrying...")

tests/conftest.py

@@ -74,6 +74,8 @@ def mock_azurehttp_calls(monkeypatch):
     def mock_post(*args, **kwargs):
         if kwargs.get("url").endswith("computervision/retrieval:vectorizeText"):
             return mock_computervision_response()
+        elif kwargs.get("url").endswith("computervision/retrieval:vectorizeImage"):
+            return mock_computervision_response()
         else:
             raise Exception("Unexpected URL for mock call to ClientSession.post()")
 
@@ -424,10 +426,6 @@ def mock_env(monkeypatch, request):
 
         with mock.patch("app.AzureDeveloperCliCredential") as mock_default_azure_credential:
             mock_default_azure_credential.return_value = MockAzureCredential()
-            # Patch the token_provider in the app to avoid the error
-            monkeypatch.setattr(
-                "azure.identity.aio.get_bearer_token_provider", lambda *args, **kwargs: mock_token_provider
-            )
             yield
 
 
@@ -452,10 +450,6 @@ def mock_reasoning_env(monkeypatch, request):
 
         with mock.patch("app.AzureDeveloperCliCredential") as mock_default_azure_credential:
             mock_default_azure_credential.return_value = MockAzureCredential()
-            # Patch the token_provider in the app to avoid the error
-            monkeypatch.setattr(
-                "azure.identity.aio.get_bearer_token_provider", lambda *args, **kwargs: mock_token_provider
-            )
             yield
 
 
@@ -480,10 +474,6 @@ def mock_agent_env(monkeypatch, request):
 
         with mock.patch("app.AzureDeveloperCliCredential") as mock_default_azure_credential:
             mock_default_azure_credential.return_value = MockAzureCredential()
-            # Patch the token_provider in the app to avoid the error
-            monkeypatch.setattr(
-                "azure.identity.aio.get_bearer_token_provider", lambda *args, **kwargs: mock_token_provider
-            )
             yield
 
 
@@ -508,10 +498,6 @@ def mock_agent_auth_env(monkeypatch, request):
 
         with mock.patch("app.AzureDeveloperCliCredential") as mock_default_azure_credential:
             mock_default_azure_credential.return_value = MockAzureCredential()
-            # Patch the token_provider in the app to avoid the error
-            monkeypatch.setattr(
-                "azure.identity.aio.get_bearer_token_provider", lambda *args, **kwargs: mock_token_provider
-            )
             yield
 
 

tests/test_mediadescriber.py

@@ -3,8 +3,14 @@
 
 import aiohttp
 import pytest
+from openai.types import CompletionUsage
+from openai.types.chat import ChatCompletion, ChatCompletionMessage
+from openai.types.chat.chat_completion import Choice
 
-from prepdocslib.mediadescriber import ContentUnderstandingDescriber
+from prepdocslib.mediadescriber import (
+    ContentUnderstandingDescriber,
+    MultimodalModelDescriber,
+)
 
 from .mocks import MockAzureCredential, MockResponse
 
@@ -133,3 +139,115 @@ def mock_put(self, *args, **kwargs):
     )
     with pytest.raises(Exception):
         await describer_bad_analyze.describe_image(b"imagebytes")
+
+
+class MockAsyncOpenAI:
+    def __init__(self, test_response):
+        self.chat = type("MockChat", (), {})()
+        self.chat.completions = MockChatCompletions(test_response)
+
+
+class MockChatCompletions:
+    def __init__(self, test_response):
+        self.test_response = test_response
+        self.create_calls = []
+
+    async def create(self, *args, **kwargs):
+        self.create_calls.append(kwargs)
+        return self.test_response
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+    "model, deployment, expected_model_param",
+    [
+        ("gpt-4o-mini", None, "gpt-4o-mini"),  # Test with model name only
+        ("gpt-4-vision-preview", "my-vision-deployment", "my-vision-deployment"),  # Test with deployment name
+    ],
+)
+async def test_multimodal_model_describer(monkeypatch, model, deployment, expected_model_param):
+    # Sample image bytes - a minimal valid PNG
+    image_bytes = b"\x89PNG\r\n\x1a\n\x00\x00\x00\rIHDR\x00\x00\x00\x01\x00\x00\x00\x01\x08\x04\x00\x00\x00\xb5\x1c\x0c\x02\x00\x00\x00\x0bIDATx\xdac\xfc\xff\xff?\x00\x05\xfe\x02\xfe\xa3\xb8\xfb\x26\x00\x00\x00\x00IEND\xaeB`\x82"
+
+    # Expected description from the model
+    expected_description = "This is a chart showing financial data trends over time."
+
+    # Create a mock OpenAI chat completion response
+    mock_response = ChatCompletion(
+        id="chatcmpl-123",
+        choices=[
+            Choice(
+                index=0,
+                message=ChatCompletionMessage(content=expected_description, role="assistant"),
+                finish_reason="stop",
+            )
+        ],
+        created=1677652288,
+        model=expected_model_param,
+        object="chat.completion",
+        usage=CompletionUsage(completion_tokens=25, prompt_tokens=50, total_tokens=75),
+    )
+
+    # Create mock OpenAI client
+    mock_openai_client = MockAsyncOpenAI(mock_response)
+
+    # Create the describer with the mock client
+    describer = MultimodalModelDescriber(openai_client=mock_openai_client, model=model, deployment=deployment)
+
+    # Call the method under test
+    result = await describer.describe_image(image_bytes)
+
+    # Verify the result matches our expected description
+    assert result == expected_description
+
+    # Verify the API was called with the correct parameters
+    assert len(mock_openai_client.chat.completions.create_calls) == 1
+    call_args = mock_openai_client.chat.completions.create_calls[0]
+
+    # Check model parameter - should be either the model or deployment based on our test case
+    assert call_args["model"] == expected_model_param
+
+    # Check that max_tokens was set
+    assert call_args["max_tokens"] == 500
+
+    # Check system message
+    messages = call_args["messages"]
+    assert len(messages) == 2
+    assert messages[0]["role"] == "system"
+    assert "helpful assistant" in messages[0]["content"]
+
+    # Check user message with image
+    assert messages[1]["role"] == "user"
+    assert len(messages[1]["content"]) == 2
+    assert messages[1]["content"][0]["type"] == "text"
+    assert "Describe image" in messages[1]["content"][0]["text"]
+    assert messages[1]["content"][1]["type"] == "image_url"
+    assert "data:image/png;base64," in messages[1]["content"][1]["image_url"]["url"]
+
+
+@pytest.mark.asyncio
+async def test_multimodal_model_describer_empty_response(monkeypatch):
+    # Sample image bytes
+    image_bytes = b"\x89PNG\r\n\x1a\n\x00\x00\x00\rIHDR\x00\x00\x00\x01\x00\x00\x00\x01\x08\x04\x00\x00\x00\xb5\x1c\x0c\x02\x00\x00\x00\x0bIDATx\xdac\xfc\xff\xff?\x00\x05\xfe\x02\xfe\xa3\xb8\xfb\x26\x00\x00\x00\x00IEND\xaeB`\x82"
+
+    # Create mock response with empty content
+    mock_response = ChatCompletion(
+        id="chatcmpl-789",
+        choices=[],  # Empty choices array
+        created=1677652288,
+        model="gpt-4o-mini",
+        object="chat.completion",
+        usage=CompletionUsage(completion_tokens=0, prompt_tokens=50, total_tokens=50),
+    )
+
+    # Create mock OpenAI client
+    mock_openai_client = MockAsyncOpenAI(mock_response)
+
+    # Create the describer
+    describer = MultimodalModelDescriber(openai_client=mock_openai_client, model="gpt-4o-mini", deployment=None)
+
+    # Call the method under test
+    result = await describer.describe_image(image_bytes)
+
+    # Verify that an empty string is returned when no choices in response
+    assert result == ""

tests/test_prepdocs.py

@@ -1,4 +1,5 @@
 import logging
+from unittest.mock import AsyncMock
 
 import openai
 import openai.types
@@ -9,6 +10,7 @@
 
 from prepdocslib.embeddings import (
     AzureOpenAIEmbeddingService,
+    ImageEmbeddings,
     OpenAIEmbeddingService,
 )
 
@@ -216,3 +218,33 @@ async def test_compute_embedding_autherror(monkeypatch, capsys):
         )
         monkeypatch.setattr(embeddings, "create_client", create_auth_error_limit_client)
         await embeddings.create_embeddings(texts=["foo"])
+
+
+@pytest.mark.asyncio
+async def test_image_embeddings_success(mock_azurehttp_calls):
+    mock_token_provider = AsyncMock(return_value="fake_token")
+
+    # Create the ImageEmbeddings instance
+    image_embeddings = ImageEmbeddings(
+        endpoint="https://fake-endpoint.azure.com/",
+        token_provider=mock_token_provider,
+    )
+
+    # Call the create_embedding method with fake image bytes
+    image_bytes = b"fake_image_data"
+    embedding = await image_embeddings.create_embedding(image_bytes)
+
+    # Verify the result
+    assert embedding == [
+        0.011925711,
+        0.023533698,
+        0.010133852,
+        0.0063544377,
+        -0.00038590943,
+        0.0013952175,
+        0.009054946,
+        -0.033573493,
+        -0.002028305,
+    ]
+
+    mock_token_provider.assert_called_once()

todo.txt

@@ -1,6 +1,6 @@
 TODO:
 
-* Fix/add unit tests
+* Fix/add unit tests - check coverage
 * Add documentation
 * Test with agentic
   * Add vectorizer for images field - special from https://learn.microsoft.com/en-us/azure/search/vector-search-vectorizer-ai-services-vision