Python: feat(python): Add embedding abstractions and OpenAI implementation (Phase 1) (#4153)

* feat(python): Add embedding abstractions and OpenAI implementation (Phase 1) This PR contains two parts: 1. **Overall migration plan** for porting vector stores and embeddings from Semantic Kernel to Agent Framework (docs/features/vector-stores-and-embeddings/README.md) covering all 10 phases from core abstractions through connectors and TextSearch. 2. **Phase 1 implementation** — core embedding abstractions and OpenAI/Azure OpenAI embedding clients: Core types (_types.py): - EmbeddingGenerationOptions TypedDict (total=False) - Embedding[EmbeddingT] generic class with model_id, dimensions, created_at - GeneratedEmbeddings[EmbeddingT, EmbeddingOptionsT] list container with options, usage - EmbeddingInputT (default str) and EmbeddingT (default list[float]) TypeVars Protocol + base class (_clients.py): - SupportsGetEmbeddings protocol — Generic[EmbeddingInputT, EmbeddingT, OptionsContraT] - BaseEmbeddingClient ABC — Generic[EmbeddingInputT, EmbeddingT, OptionsCoT] Telemetry (observability.py): - EmbeddingTelemetryLayer with gen_ai.operation.name = "embeddings" OpenAI implementation (openai/_embedding_client.py): - RawOpenAIEmbeddingClient, OpenAIEmbeddingClient, OpenAIEmbeddingOptions - Uses _ensure_client() factory pattern Azure OpenAI implementation (azure/_embedding_client.py): - AzureOpenAIEmbeddingClient following AzureOpenAIChatClient pattern - Supports API key, Entra ID credentials, env var configuration Tests: - 47 unit tests for types, protocol, base class, OpenAI, and Azure clients - 6 integration tests (gated behind RUN_INTEGRATION_TESTS + credentials) Samples: - samples/02-agents/embeddings/openai_embeddings.py - samples/02-agents/embeddings/azure_openai_embeddings.py Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> * fix: Add AzureOpenAIEmbeddingClient to azure __init__.pyi stub Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> * ci: Add embedding env vars to Python integration tests Map OPENAI_EMBEDDING_MODEL_ID and AZURE_OPENAI_EMBEDDING_DEPLOYMENT_NAME from GitHub vars to the integration test environment. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> * fix: Handle base64 encoding_format in OpenAI embedding client When encoding_format='base64' is used, the OpenAI API returns base64-encoded floats instead of a JSON array. Decode these automatically to list[float] so the return type stays consistent regardless of encoding format. Also adds a unit test for base64 decoding and fixes minor docstring/import issues. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> * fix: Only record INPUT_TOKENS for embedding telemetry Embeddings have no output/completion tokens. Remove OUTPUT_TOKENS recording which was double-counting prompt_tokens via the total_tokens fallback. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> * fix: Resolve mypy variance error and lint warning Use contravariant/covariant TypeVars for SupportsGetEmbeddings Protocol. Combine nested if into single statement in telemetry layer. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> * fix: Make EmbeddingCoT invariant for mypy compatibility GeneratedEmbeddings is invariant in its type param, so the Protocol TypeVar cannot be covariant. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> * fix: Address PR review - empty values guard, service_url for telemetry - Add early return for empty values in get_embeddings to avoid unnecessary API calls - Add service_url() method to RawOpenAIEmbeddingClient for proper telemetry endpoint reporting - Add test for empty values behavior Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> * Python: Fix OpenAI chat client compatibility with third-party endpoints and OTel 0.4.14 (#4161) * Fix system message content sent as list instead of string Some OpenAI-compatible endpoints (e.g. NVIDIA NIM) reject system messages when content is a list of content parts. This change flattens system and developer message content to a plain string in the Chat Completions client. Fixes https://github.com/microsoft/agent-framework/issues/1407 Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> * Fix compatibility with opentelemetry-semantic-conventions-ai 0.4.14 Version 0.4.14 removed several LLM_* attributes from SpanAttributes (LLM_SYSTEM, LLM_REQUEST_MODEL, LLM_RESPONSE_MODEL, LLM_REQUEST_MAX_TOKENS, LLM_REQUEST_TEMPERATURE, LLM_REQUEST_TOP_P, LLM_TOKEN_TYPE). Move these to the OtelAttr enum with their well-known gen_ai.* string values and update all references in observability.py and tests. Fixes https://github.com/microsoft/agent-framework/issues/4160 Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> * Flatten text-only message content to string for all roles Extend the system/developer fix to all message roles. Text-only content lists are now post-processed into plain strings, while multimodal content (text + images/audio) remains as a list. This fixes compatibility with OpenAI-like endpoints that cannot deserialize list content (e.g. Foundry Local's Neutron backend). Partially fixes https://github.com/microsoft/agent-framework/issues/4084 Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> * Fix streaming text lost when usage data in same chunk Some providers (e.g. Gemini) include both usage data and text content in the same streaming chunk. The early return on chunk.usage caused text and tool call parsing to be skipped entirely. Remove the early return and process usage alongside text/tool calls. Fixes https://github.com/microsoft/agent-framework/issues/3434 Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> * Fix mypy errors in _chat_client.py Rename shadowed variable 'args' in system/developer branch to 'sys_args' and rename loop variable 'content' to 'msg_content' to avoid type conflict. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --------- Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> * reorder imports * fix: Use OtelAttr.REQUEST_MODEL instead of removed SpanAttributes.LLM_REQUEST_MODEL Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> * docs: Add score_threshold to vector store plan Reference SK .NET PR #13501 for score threshold filtering semantics. Include score_threshold in SearchOptions from Phase 3. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> * docs: Add reference to roji's SK .NET MEVD work for SQL connectors Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> * fix: Clear env vars in construction tests to avoid CI leakage Tests for missing API key / model ID now use monkeypatch.delenv to ensure env vars from the integration test environment don't prevent the expected ValueError from being raised. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --------- Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
2026-06-16 21:04:09 +08:00 · 2026-02-24 08:40:20 +01:00
parent 7b24d9160d
commit 6305e3e092
18 changed files with 1890 additions and 3 deletions
@@ -0,0 +1,97 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+from __future__ import annotations
+
+from collections.abc import Sequence
+
+from agent_framework import (
+    BaseEmbeddingClient,
+    Embedding,
+    EmbeddingGenerationOptions,
+    GeneratedEmbeddings,
+    SupportsGetEmbeddings,
+)
+
+
+class MockEmbeddingClient(BaseEmbeddingClient):
+    """A simple mock embedding client for testing."""
+
+    async def get_embeddings(
+        self,
+        values: Sequence[str],
+        *,
+        options: EmbeddingGenerationOptions | None = None,
+    ) -> GeneratedEmbeddings[list[float]]:
+        return GeneratedEmbeddings(
+            [Embedding(vector=[0.1, 0.2, 0.3], model_id="mock-model") for _ in values],
+            usage={"prompt_tokens": len(values), "total_tokens": len(values)},
+        )
+
+
+# --- BaseEmbeddingClient tests ---
+
+
+async def test_base_get_embeddings() -> None:
+    client = MockEmbeddingClient()
+    result = await client.get_embeddings(["hello", "world"])
+    assert len(result) == 2
+    assert result[0].vector == [0.1, 0.2, 0.3]
+    assert result[0].model_id == "mock-model"
+
+
+async def test_base_get_embeddings_with_options() -> None:
+    client = MockEmbeddingClient()
+    options: EmbeddingGenerationOptions = {"model_id": "test", "dimensions": 3}
+    result = await client.get_embeddings(["hello"], options=options)
+    assert len(result) == 1
+
+
+async def test_base_get_embeddings_usage() -> None:
+    client = MockEmbeddingClient()
+    result = await client.get_embeddings(["a", "b", "c"])
+    assert result.usage is not None
+    assert result.usage["prompt_tokens"] == 3
+
+
+def test_base_additional_properties_default() -> None:
+    client = MockEmbeddingClient()
+    assert client.additional_properties == {}
+
+
+def test_base_additional_properties_custom() -> None:
+    client = MockEmbeddingClient(additional_properties={"key": "value"})
+    assert client.additional_properties == {"key": "value"}
+
+
+# --- SupportsGetEmbeddings protocol tests ---
+
+
+def test_mock_client_satisfies_protocol() -> None:
+    client = MockEmbeddingClient()
+    assert isinstance(client, SupportsGetEmbeddings)
+
+
+def test_plain_class_satisfies_protocol() -> None:
+    """A plain class with the right signature should satisfy the protocol."""
+
+    class PlainEmbeddingClient:
+        additional_properties: dict = {}
+
+        async def get_embeddings(self, values, *, options=None):
+            return GeneratedEmbeddings()
+
+    client = PlainEmbeddingClient()
+    assert isinstance(client, SupportsGetEmbeddings)
+
+
+def test_wrong_class_does_not_satisfy_protocol() -> None:
+    """A class without get_embeddings should not satisfy the protocol."""
+
+    class NotAnEmbeddingClient:
+        additional_properties: dict = {}
+
+        async def generate(self, values):
+            pass
+
+    client = NotAnEmbeddingClient()
+    assert not isinstance(client, SupportsGetEmbeddings)
@@ -0,0 +1,182 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+from __future__ import annotations
+
+from datetime import datetime
+
+from agent_framework import Embedding, EmbeddingGenerationOptions, GeneratedEmbeddings
+
+# --- Embedding tests ---
+
+
+def test_embedding_basic_construction() -> None:
+    embedding = Embedding(vector=[0.1, 0.2, 0.3])
+    assert embedding.vector == [0.1, 0.2, 0.3]
+    assert embedding.model_id is None
+    assert embedding.created_at is None
+    assert embedding.additional_properties == {}
+
+
+def test_embedding_construction_with_metadata() -> None:
+    now = datetime.now()
+    embedding = Embedding(
+        vector=[0.1, 0.2],
+        model_id="text-embedding-3-small",
+        created_at=now,
+        additional_properties={"key": "value"},
+    )
+    assert embedding.model_id == "text-embedding-3-small"
+    assert embedding.created_at == now
+    assert embedding.additional_properties == {"key": "value"}
+
+
+def test_embedding_dimensions_computed_from_list() -> None:
+    embedding = Embedding(vector=[0.1, 0.2, 0.3])
+    assert embedding.dimensions == 3
+
+
+def test_embedding_dimensions_computed_from_tuple() -> None:
+    embedding = Embedding(vector=(0.1, 0.2, 0.3, 0.4))
+    assert embedding.dimensions == 4
+
+
+def test_embedding_dimensions_computed_from_bytes() -> None:
+    embedding = Embedding(vector=b"\x00\x01\x02")
+    assert embedding.dimensions == 3
+
+
+def test_embedding_dimensions_explicit_overrides_computed() -> None:
+    embedding = Embedding(vector=[0.1, 0.2, 0.3], dimensions=1536)
+    assert embedding.dimensions == 1536
+
+
+def test_embedding_dimensions_none_for_unknown_type() -> None:
+    embedding = Embedding(vector="not a list")  # type: ignore[arg-type]
+    assert embedding.dimensions is None
+
+
+def test_embedding_dimensions_explicit_with_unknown_type() -> None:
+    embedding = Embedding(vector="not a list", dimensions=100)  # type: ignore[arg-type]
+    assert embedding.dimensions == 100
+
+
+def test_embedding_empty_vector() -> None:
+    embedding = Embedding(vector=[])
+    assert embedding.dimensions == 0
+
+
+def test_embedding_int_vector() -> None:
+    embedding = Embedding(vector=[1, 2, 3])
+    assert embedding.vector == [1, 2, 3]
+    assert embedding.dimensions == 3
+
+
+# --- GeneratedEmbeddings tests ---
+
+
+def test_generated_basic_construction() -> None:
+    embeddings = GeneratedEmbeddings()
+    assert len(embeddings) == 0
+    assert embeddings.options is None
+    assert embeddings.usage is None
+    assert embeddings.additional_properties == {}
+
+
+def test_generated_construction_with_embeddings() -> None:
+    items = [Embedding(vector=[0.1, 0.2]), Embedding(vector=[0.3, 0.4])]
+    embeddings = GeneratedEmbeddings(items)
+    assert len(embeddings) == 2
+    assert embeddings[0].vector == [0.1, 0.2]
+    assert embeddings[1].vector == [0.3, 0.4]
+
+
+def test_generated_construction_with_usage() -> None:
+    usage = {"prompt_tokens": 10, "total_tokens": 10}
+    embeddings = GeneratedEmbeddings(
+        [
+            Embedding(
+                vector=[0.1],
+                model_id="test-model",
+            )
+        ],
+        usage=usage,
+    )
+    assert embeddings.usage == usage
+    assert embeddings.usage["prompt_tokens"] == 10
+
+
+def test_generated_construction_with_additional_properties() -> None:
+    embeddings = GeneratedEmbeddings(
+        additional_properties={"model": "test"},
+    )
+    assert embeddings.additional_properties == {"model": "test"}
+
+
+def test_generated_construction_with_options() -> None:
+    opts: EmbeddingGenerationOptions = {"model_id": "text-embedding-3-small", "dimensions": 256}
+    embeddings = GeneratedEmbeddings(
+        [Embedding(vector=[0.1])],
+        options=opts,
+    )
+    assert embeddings.options is not None
+    assert embeddings.options["model_id"] == "text-embedding-3-small"
+    assert embeddings.options["dimensions"] == 256
+
+
+def test_generated_list_behavior_iteration() -> None:
+    items = [Embedding(vector=[float(i)]) for i in range(5)]
+    embeddings = GeneratedEmbeddings(items)
+    vectors = [e.vector for e in embeddings]
+    assert vectors == [[0.0], [1.0], [2.0], [3.0], [4.0]]
+
+
+def test_generated_list_behavior_indexing() -> None:
+    items = [Embedding(vector=[0.1]), Embedding(vector=[0.2])]
+    embeddings = GeneratedEmbeddings(items)
+    assert embeddings[0].vector == [0.1]
+    assert embeddings[-1].vector == [0.2]
+
+
+def test_generated_list_behavior_slicing() -> None:
+    items = [Embedding(vector=[float(i)]) for i in range(5)]
+    embeddings = GeneratedEmbeddings(items)
+    sliced = embeddings[1:3]
+    assert len(sliced) == 2
+
+
+def test_generated_list_behavior_append() -> None:
+    embeddings = GeneratedEmbeddings()
+    embeddings.append(Embedding(vector=[0.1]))
+    assert len(embeddings) == 1
+
+
+def test_generated_none_embeddings_creates_empty_list() -> None:
+    embeddings = GeneratedEmbeddings(None)
+    assert len(embeddings) == 0
+
+
+# --- EmbeddingGenerationOptions tests ---
+
+
+def test_options_empty() -> None:
+    options: EmbeddingGenerationOptions = {}
+    assert "model_id" not in options
+
+
+def test_options_with_model_id() -> None:
+    options: EmbeddingGenerationOptions = {"model_id": "text-embedding-3-small"}
+    assert options["model_id"] == "text-embedding-3-small"
+
+
+def test_options_with_dimensions() -> None:
+    options: EmbeddingGenerationOptions = {"dimensions": 1536}
+    assert options["dimensions"] == 1536
+
+
+def test_options_with_all_fields() -> None:
+    options: EmbeddingGenerationOptions = {
+        "model_id": "text-embedding-3-small",
+        "dimensions": 1536,
+    }
+    assert options["model_id"] == "text-embedding-3-small"
+    assert options["dimensions"] == 1536
@@ -0,0 +1,362 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+from __future__ import annotations
+
+import os
+from unittest.mock import AsyncMock, MagicMock
+
+import pytest
+from openai.types import CreateEmbeddingResponse
+from openai.types import Embedding as OpenAIEmbedding
+from openai.types.create_embedding_response import Usage
+
+from agent_framework.azure import AzureOpenAIEmbeddingClient
+from agent_framework.openai import (
+    OpenAIEmbeddingClient,
+    OpenAIEmbeddingOptions,
+)
+
+
+def _make_openai_response(
+    embeddings: list[list[float]],
+    model: str = "text-embedding-3-small",
+    prompt_tokens: int = 5,
+    total_tokens: int = 5,
+) -> CreateEmbeddingResponse:
+    """Helper to create a mock OpenAI embeddings response."""
+    data = [OpenAIEmbedding(embedding=emb, index=i, object="embedding") for i, emb in enumerate(embeddings)]
+    return CreateEmbeddingResponse(
+        data=data,
+        model=model,
+        object="list",
+        usage=Usage(prompt_tokens=prompt_tokens, total_tokens=total_tokens),
+    )
+
+
+@pytest.fixture
+def openai_unit_test_env(monkeypatch: pytest.MonkeyPatch) -> None:
+    """Set up environment variables for OpenAI embedding client."""
+    monkeypatch.setenv("OPENAI_API_KEY", "test-api-key")
+    monkeypatch.setenv("OPENAI_EMBEDDING_MODEL_ID", "text-embedding-3-small")
+
+
+# --- OpenAI unit tests ---
+
+
+def test_openai_construction_with_explicit_params() -> None:
+    client = OpenAIEmbeddingClient(
+        model_id="text-embedding-3-small",
+        api_key="test-key",
+    )
+    assert client.model_id == "text-embedding-3-small"
+
+
+def test_openai_construction_from_env(openai_unit_test_env: None) -> None:
+    client = OpenAIEmbeddingClient()
+    assert client.model_id == "text-embedding-3-small"
+
+
+def test_openai_construction_missing_api_key_raises(monkeypatch: pytest.MonkeyPatch) -> None:
+    monkeypatch.delenv("OPENAI_API_KEY", raising=False)
+    with pytest.raises(ValueError, match="API key is required"):
+        OpenAIEmbeddingClient(model_id="text-embedding-3-small")
+
+
+def test_openai_construction_missing_model_raises(monkeypatch: pytest.MonkeyPatch) -> None:
+    monkeypatch.delenv("OPENAI_EMBEDDING_MODEL_ID", raising=False)
+    with pytest.raises(ValueError, match="model ID is required"):
+        OpenAIEmbeddingClient(api_key="test-key")
+
+
+async def test_openai_get_embeddings(openai_unit_test_env: None) -> None:
+    mock_response = _make_openai_response(
+        embeddings=[[0.1, 0.2, 0.3], [0.4, 0.5, 0.6]],
+    )
+    client = OpenAIEmbeddingClient()
+    client.client = MagicMock()
+    client.client.embeddings = MagicMock()
+    client.client.embeddings.create = AsyncMock(return_value=mock_response)
+
+    result = await client.get_embeddings(["hello", "world"])
+
+    assert len(result) == 2
+    assert result[0].vector == [0.1, 0.2, 0.3]
+    assert result[1].vector == [0.4, 0.5, 0.6]
+    assert result[0].model_id == "text-embedding-3-small"
+    assert result[0].dimensions == 3
+
+
+async def test_openai_get_embeddings_usage(openai_unit_test_env: None) -> None:
+    mock_response = _make_openai_response(
+        embeddings=[[0.1]],
+        prompt_tokens=10,
+        total_tokens=10,
+    )
+    client = OpenAIEmbeddingClient()
+    client.client = MagicMock()
+    client.client.embeddings = MagicMock()
+    client.client.embeddings.create = AsyncMock(return_value=mock_response)
+
+    result = await client.get_embeddings(["test"])
+
+    assert result.usage is not None
+    assert result.usage["prompt_tokens"] == 10
+    assert result.usage["total_tokens"] == 10
+
+
+async def test_openai_options_passthrough_dimensions(openai_unit_test_env: None) -> None:
+    mock_response = _make_openai_response(embeddings=[[0.1]])
+    client = OpenAIEmbeddingClient()
+    client.client = MagicMock()
+    client.client.embeddings = MagicMock()
+    client.client.embeddings.create = AsyncMock(return_value=mock_response)
+
+    options: OpenAIEmbeddingOptions = {"dimensions": 256}
+    result = await client.get_embeddings(["test"], options=options)
+
+    call_kwargs = client.client.embeddings.create.call_args[1]
+    assert call_kwargs["dimensions"] == 256
+    assert result.options is options
+
+
+async def test_openai_options_passthrough_encoding_format(openai_unit_test_env: None) -> None:
+    mock_response = _make_openai_response(embeddings=[[0.1]])
+    client = OpenAIEmbeddingClient()
+    client.client = MagicMock()
+    client.client.embeddings = MagicMock()
+    client.client.embeddings.create = AsyncMock(return_value=mock_response)
+
+    options: OpenAIEmbeddingOptions = {"encoding_format": "base64"}
+    await client.get_embeddings(["test"], options=options)
+
+    call_kwargs = client.client.embeddings.create.call_args[1]
+    assert call_kwargs["encoding_format"] == "base64"
+
+
+async def test_openai_base64_decoding(openai_unit_test_env: None) -> None:
+    import base64
+    import struct
+
+    # Encode [0.1, 0.2, 0.3] as base64 little-endian floats
+    raw_floats = [0.1, 0.2, 0.3]
+    b64_str = base64.b64encode(struct.pack(f"<{len(raw_floats)}f", *raw_floats)).decode()
+
+    # Mock the embedding item to return a base64 string (as the API does with encoding_format=base64)
+    mock_item = MagicMock()
+    mock_item.embedding = b64_str
+    mock_item.index = 0
+
+    mock_response = MagicMock()
+    mock_response.data = [mock_item]
+    mock_response.model = "text-embedding-3-small"
+    mock_response.usage = MagicMock(prompt_tokens=3, total_tokens=3)
+
+    client = OpenAIEmbeddingClient()
+    client.client = MagicMock()
+    client.client.embeddings = MagicMock()
+    client.client.embeddings.create = AsyncMock(return_value=mock_response)
+
+    options: OpenAIEmbeddingOptions = {"encoding_format": "base64"}
+    result = await client.get_embeddings(["test"], options=options)
+
+    assert len(result) == 1
+    assert len(result[0].vector) == 3
+    assert result[0].dimensions == 3
+    for expected, actual in zip(raw_floats, result[0].vector):
+        assert abs(expected - actual) < 1e-6
+
+
+async def test_openai_error_when_no_model_id() -> None:
+    client = OpenAIEmbeddingClient.__new__(OpenAIEmbeddingClient)
+    client.model_id = None
+    client.client = MagicMock()
+    client.additional_properties = {}
+    client.otel_provider_name = "openai"
+
+    with pytest.raises(ValueError, match="model_id is required"):
+        await client.get_embeddings(["test"])
+
+
+async def test_openai_empty_values_returns_empty(openai_unit_test_env: None) -> None:
+    client = OpenAIEmbeddingClient()
+    client.client = MagicMock()
+    client.client.embeddings = MagicMock()
+    client.client.embeddings.create = AsyncMock()
+
+    result = await client.get_embeddings([])
+
+    assert len(result) == 0
+    assert result.usage is None
+    client.client.embeddings.create.assert_not_called()
+
+
+# --- Azure OpenAI unit tests ---
+
+
+def test_azure_construction_with_deployment_name() -> None:
+    client = AzureOpenAIEmbeddingClient(
+        deployment_name="text-embedding-3-small",
+        api_key="test-key",
+        endpoint="https://test.openai.azure.com/",
+    )
+    assert client.model_id == "text-embedding-3-small"
+
+
+def test_azure_construction_with_existing_client() -> None:
+    mock_client = MagicMock()
+    client = AzureOpenAIEmbeddingClient(
+        deployment_name="my-deployment",
+        async_client=mock_client,
+    )
+    assert client.model_id == "my-deployment"
+    assert client.client is mock_client
+
+
+def test_azure_construction_missing_deployment_name_raises() -> None:
+    with pytest.raises(ValueError, match="deployment name is required"):
+        AzureOpenAIEmbeddingClient(
+            api_key="test-key",
+            endpoint="https://test.openai.azure.com/",
+        )
+
+
+def test_azure_construction_missing_credentials_raises() -> None:
+    with pytest.raises(ValueError, match="api_key, credential, or a client"):
+        AzureOpenAIEmbeddingClient(
+            deployment_name="test",
+            endpoint="https://test.openai.azure.com/",
+        )
+
+
+async def test_azure_get_embeddings() -> None:
+    mock_response = _make_openai_response(
+        embeddings=[[0.1, 0.2]],
+    )
+    mock_async_client = MagicMock()
+    mock_async_client.embeddings = MagicMock()
+    mock_async_client.embeddings.create = AsyncMock(return_value=mock_response)
+
+    client = AzureOpenAIEmbeddingClient(
+        deployment_name="text-embedding-3-small",
+        async_client=mock_async_client,
+    )
+
+    result = await client.get_embeddings(["hello"])
+
+    assert len(result) == 1
+    assert result[0].vector == [0.1, 0.2]
+
+
+def test_azure_otel_provider_name() -> None:
+    mock_client = MagicMock()
+    client = AzureOpenAIEmbeddingClient(
+        deployment_name="test",
+        async_client=mock_client,
+    )
+    assert client.OTEL_PROVIDER_NAME == "azure.ai.openai"
+
+
+# --- Integration tests ---
+
+skip_if_openai_integration_tests_disabled = pytest.mark.skipif(
+    os.getenv("RUN_INTEGRATION_TESTS", "false").lower() != "true"
+    or os.getenv("OPENAI_API_KEY", "") in ("", "test-dummy-key"),
+    reason="No real OPENAI_API_KEY provided; skipping integration tests."
+    if os.getenv("RUN_INTEGRATION_TESTS", "false").lower() == "true"
+    else "Integration tests are disabled.",
+)
+
+skip_if_azure_openai_integration_tests_disabled = pytest.mark.skipif(
+    os.getenv("RUN_INTEGRATION_TESTS", "false").lower() != "true"
+    or not os.getenv("AZURE_OPENAI_ENDPOINT")
+    or (not os.getenv("AZURE_OPENAI_API_KEY") and not os.getenv("AZURE_OPENAI_EMBEDDING_DEPLOYMENT_NAME")),
+    reason="No Azure OpenAI credentials provided; skipping integration tests."
+    if os.getenv("RUN_INTEGRATION_TESTS", "false").lower() == "true"
+    else "Integration tests are disabled.",
+)
+
+
+@skip_if_openai_integration_tests_disabled
+@pytest.mark.flaky
+async def test_integration_openai_get_embeddings() -> None:
+    """End-to-end test of OpenAI embedding generation."""
+    client = OpenAIEmbeddingClient(model_id="text-embedding-3-small")
+
+    result = await client.get_embeddings(["hello world"])
+
+    assert len(result) == 1
+    assert isinstance(result[0].vector, list)
+    assert len(result[0].vector) > 0
+    assert all(isinstance(v, float) for v in result[0].vector)
+    assert result[0].model_id is not None
+    assert result.usage is not None
+    assert result.usage["prompt_tokens"] > 0
+
+
+@skip_if_openai_integration_tests_disabled
+@pytest.mark.flaky
+async def test_integration_openai_get_embeddings_multiple() -> None:
+    """Test embedding generation for multiple inputs."""
+    client = OpenAIEmbeddingClient(model_id="text-embedding-3-small")
+
+    result = await client.get_embeddings(["hello", "world", "test"])
+
+    assert len(result) == 3
+    dims = [len(e.vector) for e in result]
+    assert all(d == dims[0] for d in dims)
+
+
+@skip_if_openai_integration_tests_disabled
+@pytest.mark.flaky
+async def test_integration_openai_get_embeddings_with_dimensions() -> None:
+    """Test embedding generation with custom dimensions."""
+    client = OpenAIEmbeddingClient(model_id="text-embedding-3-small")
+
+    options: OpenAIEmbeddingOptions = {"dimensions": 256}
+    result = await client.get_embeddings(["hello world"], options=options)
+
+    assert len(result) == 1
+    assert len(result[0].vector) == 256
+
+
+@skip_if_azure_openai_integration_tests_disabled
+@pytest.mark.flaky
+async def test_integration_azure_openai_get_embeddings() -> None:
+    """End-to-end test of Azure OpenAI embedding generation."""
+    client = AzureOpenAIEmbeddingClient()
+
+    result = await client.get_embeddings(["hello world"])
+
+    assert len(result) == 1
+    assert isinstance(result[0].vector, list)
+    assert len(result[0].vector) > 0
+    assert all(isinstance(v, float) for v in result[0].vector)
+    assert result[0].model_id is not None
+    assert result.usage is not None
+    assert result.usage["prompt_tokens"] > 0
+
+
+@skip_if_azure_openai_integration_tests_disabled
+@pytest.mark.flaky
+async def test_integration_azure_openai_get_embeddings_multiple() -> None:
+    """Test Azure OpenAI embedding generation for multiple inputs."""
+    client = AzureOpenAIEmbeddingClient()
+
+    result = await client.get_embeddings(["hello", "world", "test"])
+
+    assert len(result) == 3
+    dims = [len(e.vector) for e in result]
+    assert all(d == dims[0] for d in dims)
+
+
+@skip_if_azure_openai_integration_tests_disabled
+@pytest.mark.flaky
+async def test_integration_azure_openai_get_embeddings_with_dimensions() -> None:
+    """Test Azure OpenAI embedding generation with custom dimensions."""
+    client = AzureOpenAIEmbeddingClient()
+
+    options: OpenAIEmbeddingOptions = {"dimensions": 256}
+    result = await client.get_embeddings(["hello world"], options=options)
+
+    assert len(result) == 1
+    assert len(result[0].vector) == 256