Use ImageEmbeddings client directly

Author: pamelafox

app/backend/app.py

@@ -103,6 +103,7 @@
     setup_search_info,
 )
 from prepdocslib.blobmanager import AdlsBlobManager
+from prepdocslib.embeddings import ImageEmbeddings
 from prepdocslib.filestrategy import UploadUserFileStrategy
 from prepdocslib.listfilestrategy import File
 
@@ -624,6 +625,10 @@ async def setup_clients():
         )
         current_app.config[CONFIG_INGESTER] = ingester
 
+    image_embeddings_client = None
+    if USE_MULTIMODAL:
+        image_embeddings_client = ImageEmbeddings(AZURE_VISION_ENDPOINT, azure_ai_token_provider)
+
     current_app.config[CONFIG_OPENAI_CLIENT] = openai_client
     current_app.config[CONFIG_SEARCH_CLIENT] = search_client
     current_app.config[CONFIG_AGENT_CLIENT] = agent_client
@@ -659,6 +664,7 @@ async def setup_clients():
 
     # Set up the two default RAG approaches for /ask and /chat
     # RetrieveThenReadApproach is used by /ask for single-turn Q&A
+
     current_app.config[CONFIG_ASK_APPROACH] = RetrieveThenReadApproach(
         search_client=search_client,
         search_index_name=AZURE_SEARCH_INDEX,
@@ -667,8 +673,6 @@ async def setup_clients():
         agent_client=agent_client,
         openai_client=openai_client,
         auth_helper=auth_helper,
-        image_blob_container_client=image_blob_container_client,
-        image_datalake_client=user_blob_container_client,
         chatgpt_model=OPENAI_CHATGPT_MODEL,
         chatgpt_deployment=AZURE_OPENAI_CHATGPT_DEPLOYMENT,
         embedding_model=OPENAI_EMB_MODEL,
@@ -681,9 +685,10 @@ async def setup_clients():
         query_speller=AZURE_SEARCH_QUERY_SPELLER,
         prompt_manager=prompt_manager,
         reasoning_effort=OPENAI_REASONING_EFFORT,
-        vision_endpoint=AZURE_VISION_ENDPOINT,
-        vision_token_provider=azure_ai_token_provider,
         multimodal_enabled=USE_MULTIMODAL,
+        image_embeddings_client=image_embeddings_client,
+        image_blob_container_client=image_blob_container_client,
+        image_datalake_client=user_blob_container_client,
     )
 
     # ChatReadRetrieveReadApproach is used by /chat for multi-turn conversation
@@ -695,8 +700,6 @@ async def setup_clients():
         agent_client=agent_client,
         openai_client=openai_client,
         auth_helper=auth_helper,
-        image_blob_container_client=image_blob_container_client,
-        image_datalake_client=user_blob_container_client,
         chatgpt_model=OPENAI_CHATGPT_MODEL,
         chatgpt_deployment=AZURE_OPENAI_CHATGPT_DEPLOYMENT,
         embedding_model=OPENAI_EMB_MODEL,
@@ -709,9 +712,10 @@ async def setup_clients():
         query_speller=AZURE_SEARCH_QUERY_SPELLER,
         prompt_manager=prompt_manager,
         reasoning_effort=OPENAI_REASONING_EFFORT,
-        vision_endpoint=AZURE_VISION_ENDPOINT,
-        vision_token_provider=azure_ai_token_provider,
         multimodal_enabled=USE_MULTIMODAL,
+        image_embeddings_client=image_embeddings_client,
+        image_blob_container_client=image_blob_container_client,
+        image_datalake_client=user_blob_container_client,
     )
 
 

app/backend/approaches/approach.py

@@ -2,10 +2,8 @@
 from collections.abc import AsyncGenerator, Awaitable
 from dataclasses import dataclass, field
 from enum import Enum
-from typing import Any, Callable, Optional, TypedDict, Union, cast
-from urllib.parse import urljoin
+from typing import Any, Optional, TypedDict, Union, cast
 
-import aiohttp
 from azure.search.documents.agent.aio import KnowledgeAgentRetrievalClient
 from azure.search.documents.agent.models import (
     KnowledgeAgentAzureSearchDocReference,
@@ -38,6 +36,7 @@
 from approaches.promptmanager import PromptManager
 from core.authentication import AuthenticationHelper
 from core.imageshelper import download_blob_as_base64
+from prepdocslib.embeddings import ImageEmbeddings
 
 
 class LLMInputType(str, Enum):
@@ -174,8 +173,7 @@ def __init__(
         prompt_manager: PromptManager,
         reasoning_effort: Optional[str] = None,
         multimodal_enabled: bool = False,
-        vision_endpoint: Optional[str] = None,
-        vision_token_provider: Optional[Callable[[], Awaitable[str]]] = None,
+        image_embeddings_client: Optional[ImageEmbeddings] = None,
         image_blob_container_client: Optional[ContainerClient] = None,
         image_datalake_client: Optional[FileSystemClient] = None,
     ):
@@ -193,8 +191,7 @@ def __init__(
         self.reasoning_effort = reasoning_effort
         self.include_token_usage = True
         self.multimodal_enabled = multimodal_enabled
-        self.vision_endpoint = vision_endpoint
-        self.vision_token_provider = vision_token_provider
+        self.image_embeddings_client = image_embeddings_client
         self.image_blob_container_client = image_blob_container_client
         self.image_datalake_client = image_datalake_client
 
@@ -462,25 +459,9 @@ class ExtraArgs(TypedDict, total=False):
         # so we do not need to explicitly pass in an oversampling parameter here
         return VectorizedQuery(vector=query_vector, k_nearest_neighbors=50, fields=self.embedding_field)
 
-    async def compute_image_embedding(self, q: str):
-        if not self.vision_endpoint:
-            raise ValueError("Azure AI Vision endpoint must be set to compute image embedding.")
-        endpoint = urljoin(self.vision_endpoint, "computervision/retrieval:vectorizeText")
-        headers = {"Content-Type": "application/json"}
-        params = {"api-version": "2024-02-01", "model-version": "2023-04-15"}
-        data = {"text": q}
-
-        if not self.vision_token_provider:
-            raise ValueError("Azure AI Vision token provider must be set to compute image embedding.")
-        headers["Authorization"] = "Bearer " + await self.vision_token_provider()
-
-        async with aiohttp.ClientSession() as session:
-            async with session.post(
-                url=endpoint, params=params, headers=headers, json=data, raise_for_status=True
-            ) as response:
-                json = await response.json()
-                image_query_vector = json["vector"]
-        return VectorizedQuery(vector=image_query_vector, k_nearest_neighbors=50, fields="images/embedding")
+    async def compute_multimodal_embedding(self, q: str):
+        multimodal_query_vector = await self.image_embeddings_client.create_embedding_for_text(q)
+        return VectorizedQuery(vector=multimodal_query_vector, k_nearest_neighbors=50, fields="images/embedding")
 
     def get_system_prompt_variables(self, override_prompt: Optional[str]) -> dict[str, str]:
         # Allows client to replace the entire prompt, or to inject into the existing prompt using >>>

app/backend/approaches/chatreadretrieveread.py

@@ -1,7 +1,7 @@
 import json
 import re
 from collections.abc import AsyncGenerator, Awaitable
-from typing import Any, Callable, Optional, Union, cast
+from typing import Any, Optional, Union, cast
 
 from azure.search.documents.agent.aio import KnowledgeAgentRetrievalClient
 from azure.search.documents.aio import SearchClient
@@ -26,6 +26,7 @@
 )
 from approaches.promptmanager import PromptManager
 from core.authentication import AuthenticationHelper
+from prepdocslib.embeddings import ImageEmbeddings
 
 
 class ChatReadRetrieveReadApproach(Approach):
@@ -60,8 +61,7 @@ def __init__(
         prompt_manager: PromptManager,
         reasoning_effort: Optional[str] = None,
         multimodal_enabled: bool = False,
-        vision_endpoint: Optional[str] = None,
-        vision_token_provider: Optional[Callable[[], Awaitable[str]]] = None,
+        image_embeddings_client: Optional[ImageEmbeddings] = None,
         image_blob_container_client: Optional[ContainerClient] = None,
         image_datalake_client: Optional[FileSystemClient] = None,
     ):
@@ -72,8 +72,7 @@ def __init__(
         self.agent_client = agent_client
         self.openai_client = openai_client
         self.auth_helper = auth_helper
-        self.image_blob_container_client = image_blob_container_client
-        self.image_datalake_client = image_datalake_client
+
         self.chatgpt_model = chatgpt_model
         self.chatgpt_deployment = chatgpt_deployment
         self.embedding_deployment = embedding_deployment
@@ -90,9 +89,10 @@ def __init__(
         self.answer_prompt = self.prompt_manager.load_prompt("chat_answer_question.prompty")
         self.reasoning_effort = reasoning_effort
         self.include_token_usage = True
-        self.vision_endpoint = vision_endpoint
-        self.vision_token_provider = vision_token_provider
         self.multimodal_enabled = multimodal_enabled
+        self.image_embeddings_client = image_embeddings_client
+        self.image_blob_container_client = image_blob_container_client
+        self.image_datalake_client = image_datalake_client
 
     def get_search_query(self, chat_completion: ChatCompletion, user_query: str):
         response_message = chat_completion.choices[0].message
@@ -340,7 +340,7 @@ async def run_search_approach(
         if use_vector_search:
             vectors.append(await self.compute_text_embedding(query_text))
             if use_image_embeddings:
-                vectors.append(await self.compute_image_embedding(query_text))
+                vectors.append(await self.compute_multimodal_embedding(query_text))
 
         results = await self.search(
             top,

app/backend/approaches/retrievethenread.py

@@ -1,5 +1,4 @@
-from collections.abc import Awaitable
-from typing import Any, Callable, Optional, cast
+from typing import Any, Optional, cast
 
 from azure.search.documents.agent.aio import KnowledgeAgentRetrievalClient
 from azure.search.documents.aio import SearchClient
@@ -19,6 +18,7 @@
 )
 from approaches.promptmanager import PromptManager
 from core.authentication import AuthenticationHelper
+from prepdocslib.embeddings import ImageEmbeddings
 
 
 class RetrieveThenReadApproach(Approach):
@@ -51,8 +51,7 @@ def __init__(
         prompt_manager: PromptManager,
         reasoning_effort: Optional[str] = None,
         multimodal_enabled: bool = False,
-        vision_endpoint: Optional[str] = None,
-        vision_token_provider: Optional[Callable[[], Awaitable[str]]] = None,
+        image_embeddings_client: Optional[ImageEmbeddings] = None,
         image_blob_container_client: Optional[ContainerClient] = None,
         image_datalake_client: Optional[FileSystemClient] = None,
     ):
@@ -64,8 +63,6 @@ def __init__(
         self.chatgpt_deployment = chatgpt_deployment
         self.openai_client = openai_client
         self.auth_helper = auth_helper
-        self.image_blob_container_client = image_blob_container_client
-        self.image_datalake_client = image_datalake_client
         self.chatgpt_model = chatgpt_model
         self.embedding_model = embedding_model
         self.embedding_dimensions = embedding_dimensions
@@ -80,9 +77,10 @@ def __init__(
         self.answer_prompt = self.prompt_manager.load_prompt("ask_answer_question.prompty")
         self.reasoning_effort = reasoning_effort
         self.include_token_usage = True
-        self.vision_endpoint = vision_endpoint
-        self.vision_token_provider = vision_token_provider
         self.multimodal_enabled = multimodal_enabled
+        self.image_embeddings_client = image_embeddings_client
+        self.image_blob_container_client = image_blob_container_client
+        self.image_datalake_client = image_datalake_client
 
     async def run(
         self,
@@ -186,7 +184,7 @@ async def run_search_approach(
             if vector_fields_enum != VectorFieldType.IMAGE_EMBEDDING:
                 vectors.append(await self.compute_text_embedding(q))
             if use_image_embeddings:
-                vectors.append(await self.compute_image_embedding(q))
+                vectors.append(await self.compute_multimodal_embedding(q))
 
         results = await self.search(
             top,

app/backend/prepdocslib/embeddings.py

@@ -236,7 +236,7 @@ def __init__(self, endpoint: str, token_provider: Callable[[], Awaitable[str]]):
         self.token_provider = token_provider
         self.endpoint = endpoint
 
-    async def create_embedding(self, image_bytes: bytes) -> list[float]:
+    async def create_embedding_for_image(self, image_bytes: bytes) -> list[float]:
         endpoint = urljoin(self.endpoint, "computervision/retrieval:vectorizeImage")
         params = {"api-version": "2024-02-01", "model-version": "2023-04-15"}
         headers = {"Authorization": "Bearer " + await self.token_provider()}
@@ -254,5 +254,25 @@ async def create_embedding(self, image_bytes: bytes) -> list[float]:
                         return resp_json["vector"]
         raise ValueError("Failed to get image embedding after multiple retries.")
 
+    async def create_embedding_for_text(self, q: str):
+        if not self.endpoint:
+            raise ValueError("Azure AI Vision endpoint must be set to compute image embedding.")
+        endpoint = urljoin(self.endpoint, "computervision/retrieval:vectorizeText")
+        headers = {"Content-Type": "application/json"}
+        params = {"api-version": "2024-02-01", "model-version": "2023-04-15"}
+        data = {"text": q}
+
+        if not self.token_provider:
+            raise ValueError("Azure AI Vision token provider must be set to compute image embedding.")
+        headers["Authorization"] = "Bearer " + await self.token_provider()
+
+        async with aiohttp.ClientSession() as session:
+            async with session.post(
+                url=endpoint, params=params, headers=headers, json=data, raise_for_status=True
+            ) as response:
+                json = await response.json()
+                return json["vector"]
+        raise ValueError("Failed to get image embedding after multiple retries.")
+
     def before_retry_sleep(self, retry_state):
         logger.info("Rate limited on the Vision embeddings API, sleeping before retrying...")

app/backend/prepdocslib/filestrategy.py

@@ -38,7 +38,7 @@ async def parse_file(
                     file.filename(), image.bytes, image.filename, image.page_num, user_oid=user_oid
                 )
             if image_embeddings_client:
-                image.embedding = await image_embeddings_client.create_embedding(image.bytes)
+                image.embedding = await image_embeddings_client.create_embedding_for_image(image.bytes)
     logger.info("Splitting '%s' into sections", file.filename())
     sections = [
         Section(split_page, content=file, category=category) for split_page in processor.splitter.split_pages(pages)

tests/test_prepdocs.py

@@ -232,7 +232,7 @@ async def test_image_embeddings_success(mock_azurehttp_calls):
 
     # Call the create_embedding method with fake image bytes
     image_bytes = b"fake_image_data"
-    embedding = await image_embeddings.create_embedding(image_bytes)
+    embedding = await image_embeddings.create_embedding_for_image(image_bytes)
 
     # Verify the result
     assert embedding == [

todo.txt

@@ -3,10 +3,9 @@ TODO:
 * Test with integrated vectorization
    * Multivector is working
    * Can we get images mapped??
+   * We need DocIntelligence skill
 * Update all TODOs in the code/docs
 * Fix/add unit tests - check coverage
-* In conftest, should I make a new env for vision? Currently I mashed it into the existing env, but it might be cleaner to have a separate one, as now I have to pass llm_inputs explicitly in the tests to turn off image responses.
-   * vote: make a new env
 * LLMInputType and VectorFields have inconsistently named values
    * # Vector fields:
       # [X] text embedding field (embedding3) use_text_vector=True
@@ -16,9 +15,6 @@ TODO:
       # [X] text sources , use_text_sources = True
       # [X] image sources , use_image_sources = True
 
-* Should we make an Azure AI Vision client class? So we dont have to pass two things around, just one?
-   * vision_endpoint and vision_token_provider
-   * we have one! its in embeddings.py, add compute_image_embeddings method to it and use it instead
 
 To decide:
 * For user data lake client, how often should we double check the ACL matches the oid, versus assuming the URLs convey that? (Like when fetching the image?)