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>
This commit is contained in:
Eduard van Valkenburg
2026-02-24 08:40:20 +01:00
committed by GitHub
Unverified
parent 7b24d9160d
commit 6305e3e092
18 changed files with 1890 additions and 3 deletions
@@ -20,9 +20,11 @@ __version__: Final[str] = _version
from ._agents import Agent, BaseAgent, RawAgent, SupportsAgentRun
from ._clients import (
BaseChatClient,
BaseEmbeddingClient,
SupportsChatGetResponse,
SupportsCodeInterpreterTool,
SupportsFileSearchTool,
SupportsGetEmbeddings,
SupportsImageGenerationTool,
SupportsMCPTool,
SupportsWebSearchTool,
@@ -82,9 +84,14 @@ from ._types import (
ChatResponseUpdate,
Content,
ContinuationToken,
Embedding,
EmbeddingGenerationOptions,
EmbeddingInputT,
EmbeddingT,
FinalT,
FinishReason,
FinishReasonLiteral,
GeneratedEmbeddings,
Message,
OuterFinalT,
OuterUpdateT,
@@ -201,6 +208,7 @@ __all__ = [
"BaseAgent",
"BaseChatClient",
"BaseContextProvider",
"BaseEmbeddingClient",
"BaseHistoryProvider",
"Case",
"ChatAndFunctionMiddlewareTypes",
@@ -218,6 +226,10 @@ __all__ = [
"Edge",
"EdgeCondition",
"EdgeDuplicationError",
"Embedding",
"EmbeddingGenerationOptions",
"EmbeddingInputT",
"EmbeddingT",
"Executor",
"FanInEdgeGroup",
"FanOutEdgeGroup",
@@ -232,6 +244,7 @@ __all__ = [
"FunctionMiddleware",
"FunctionMiddlewareTypes",
"FunctionTool",
"GeneratedEmbeddings",
"GraphConnectivityError",
"InMemoryCheckpointStorage",
"InMemoryHistoryProvider",
@@ -261,6 +274,7 @@ __all__ = [
"SupportsChatGetResponse",
"SupportsCodeInterpreterTool",
"SupportsFileSearchTool",
"SupportsGetEmbeddings",
"SupportsImageGenerationTool",
"SupportsMCPTool",
"SupportsWebSearchTool",
@@ -35,6 +35,10 @@ from ._tools import (
from ._types import (
ChatResponse,
ChatResponseUpdate,
EmbeddingGenerationOptions,
EmbeddingInputT,
EmbeddingT,
GeneratedEmbeddings,
Message,
ResponseStream,
validate_chat_options,
@@ -56,7 +60,6 @@ if TYPE_CHECKING:
InputT = TypeVar("InputT", contravariant=True)
EmbeddingT = TypeVar("EmbeddingT")
BaseChatClientT = TypeVar("BaseChatClientT", bound="BaseChatClient")
logger = logging.getLogger("agent_framework")
@@ -660,3 +663,140 @@ class SupportsFileSearchTool(Protocol):
# endregion
# region SupportsGetEmbeddings Protocol
# Contravariant/covariant TypeVars for the Protocol
EmbeddingInputContraT = TypeVar(
"EmbeddingInputContraT",
default="str",
contravariant=True,
)
EmbeddingCoT = TypeVar(
"EmbeddingCoT",
default="list[float]",
)
EmbeddingOptionsContraT = TypeVar(
"EmbeddingOptionsContraT",
bound=TypedDict, # type: ignore[valid-type]
default="EmbeddingGenerationOptions",
contravariant=True,
)
@runtime_checkable
class SupportsGetEmbeddings(Protocol[EmbeddingInputContraT, EmbeddingCoT, EmbeddingOptionsContraT]):
"""Protocol for an embedding client that can generate embeddings.
This protocol enables duck-typing for embedding generation. Any class that
implements ``get_embeddings`` with a compatible signature satisfies this protocol.
Generic over the input type (defaults to ``str``), output embedding type
(defaults to ``list[float]``), and options type.
Examples:
.. code-block:: python
from agent_framework import SupportsGetEmbeddings
async def use_embeddings(client: SupportsGetEmbeddings) -> None:
result = await client.get_embeddings(["Hello, world!"])
for embedding in result:
print(embedding.vector)
"""
additional_properties: dict[str, Any]
def get_embeddings(
self,
values: Sequence[EmbeddingInputContraT],
*,
options: EmbeddingOptionsContraT | None = None,
) -> Awaitable[GeneratedEmbeddings[EmbeddingCoT]]:
"""Generate embeddings for the given values.
Args:
values: The values to generate embeddings for.
options: Optional embedding generation options.
Returns:
Generated embeddings with metadata.
"""
...
# endregion
# region BaseEmbeddingClient
# Covariant for the BaseEmbeddingClient
EmbeddingOptionsCoT = TypeVar(
"EmbeddingOptionsCoT",
bound=TypedDict, # type: ignore[valid-type]
default="EmbeddingGenerationOptions",
covariant=True,
)
class BaseEmbeddingClient(SerializationMixin, ABC, Generic[EmbeddingInputT, EmbeddingT, EmbeddingOptionsCoT]):
"""Abstract base class for embedding clients.
Subclasses implement ``get_embeddings`` to provide the actual
embedding generation logic.
Generic over the input type (defaults to ``str``), output embedding type
(defaults to ``list[float]``), and options type.
Examples:
.. code-block:: python
from agent_framework import BaseEmbeddingClient, Embedding, GeneratedEmbeddings
from collections.abc import Sequence
class CustomEmbeddingClient(BaseEmbeddingClient):
async def get_embeddings(self, values, *, options=None):
return GeneratedEmbeddings([Embedding(vector=[0.1, 0.2, 0.3]) for _ in values])
"""
OTEL_PROVIDER_NAME: ClassVar[str] = "unknown"
DEFAULT_EXCLUDE: ClassVar[set[str]] = {"additional_properties"}
def __init__(
self,
*,
additional_properties: dict[str, Any] | None = None,
**kwargs: Any,
) -> None:
"""Initialize a BaseEmbeddingClient instance.
Args:
additional_properties: Additional properties to pass to the client.
**kwargs: Additional keyword arguments passed to parent classes (for MRO).
"""
self.additional_properties = additional_properties or {}
super().__init__(**kwargs)
@abstractmethod
async def get_embeddings(
self,
values: Sequence[EmbeddingInputT],
*,
options: EmbeddingOptionsCoT | None = None,
) -> GeneratedEmbeddings[EmbeddingT]:
"""Generate embeddings for the given values.
Args:
values: The values to generate embeddings for.
options: Optional embedding generation options.
Returns:
Generated embeddings with metadata.
"""
...
# endregion
+139 -2
View File
@@ -8,8 +8,18 @@ import logging
import re
import sys
from asyncio import iscoroutine
from collections.abc import AsyncIterable, AsyncIterator, Awaitable, Callable, Mapping, MutableMapping, Sequence
from collections.abc import (
AsyncIterable,
AsyncIterator,
Awaitable,
Callable,
Iterable,
Mapping,
MutableMapping,
Sequence,
)
from copy import deepcopy
from datetime import datetime
from typing import TYPE_CHECKING, Any, ClassVar, Final, Generic, Literal, NewType, cast, overload
from pydantic import BaseModel
@@ -272,7 +282,8 @@ def _serialize_value(value: Any, exclude_none: bool) -> Any:
# region Constants and types
_T = TypeVar("_T")
EmbeddingT = TypeVar("EmbeddingT")
EmbeddingT = TypeVar("EmbeddingT", default="list[float]")
EmbeddingInputT = TypeVar("EmbeddingInputT", default="str")
ChatResponseT = TypeVar("ChatResponseT", bound="ChatResponse")
ToolModeT = TypeVar("ToolModeT", bound="ToolMode")
AgentResponseT = TypeVar("AgentResponseT", bound="AgentResponse")
@@ -3162,3 +3173,129 @@ def merge_chat_options(
result[key] = value
return result
# region Embedding Types
class EmbeddingGenerationOptions(TypedDict, total=False):
"""Common request settings for embedding generation.
All fields are optional (total=False) to allow partial specification.
Provider-specific TypedDicts extend this with additional options.
Examples:
.. code-block:: python
from agent_framework import EmbeddingGenerationOptions
options: EmbeddingGenerationOptions = {
"model_id": "text-embedding-3-small",
"dimensions": 1536,
}
"""
model_id: str
dimensions: int
class Embedding(Generic[EmbeddingT]):
"""A single embedding vector with metadata.
Generic over the embedding vector type, e.g. ``Embedding[list[float]]``,
``Embedding[list[int]]``, or ``Embedding[bytes]``.
Args:
vector: The embedding vector data.
model_id: The model used to generate this embedding.
dimensions: Explicit dimension count (computed from vector length if omitted).
created_at: Timestamp of when the embedding was generated.
additional_properties: Additional metadata.
Examples:
.. code-block:: python
from agent_framework import Embedding
embedding = Embedding(
vector=[0.1, 0.2, 0.3],
model_id="text-embedding-3-small",
)
assert embedding.dimensions == 3
"""
def __init__(
self,
vector: EmbeddingT,
*,
model_id: str | None = None,
dimensions: int | None = None,
created_at: datetime | None = None,
additional_properties: dict[str, Any] | None = None,
) -> None:
self.vector = vector
self._dimensions = dimensions
self.model_id = model_id
self.created_at = created_at
self.additional_properties = additional_properties or {}
@property
def dimensions(self) -> int | None:
"""Return the number of dimensions in the embedding vector.
Uses the explicitly provided value if set, otherwise computes from vector length.
"""
if self._dimensions is not None:
return self._dimensions
if isinstance(self.vector, (list, tuple, bytes)):
return len(self.vector)
return None
EmbeddingOptionsT = TypeVar(
"EmbeddingOptionsT",
bound=TypedDict, # type: ignore[valid-type]
default="EmbeddingGenerationOptions",
)
class GeneratedEmbeddings(list[Embedding[EmbeddingT]], Generic[EmbeddingT, EmbeddingOptionsT]):
"""A list of generated embeddings with usage metadata.
Extends list for direct iteration and indexing.
Generic over both the embedding vector type and the options type used for generation.
Args:
embeddings: Sequence of Embedding objects.
options: The options used to generate these embeddings.
usage: Token usage information (e.g. prompt_tokens, total_tokens).
additional_properties: Additional metadata.
Examples:
.. code-block:: python
from agent_framework import Embedding, GeneratedEmbeddings
embeddings = GeneratedEmbeddings(
[Embedding(vector=[0.1, 0.2]), Embedding(vector=[0.3, 0.4])],
usage={"prompt_tokens": 10, "total_tokens": 10},
)
assert len(embeddings) == 2
assert embeddings.usage["prompt_tokens"] == 10
"""
def __init__(
self,
embeddings: Iterable[Embedding[EmbeddingT]] | None = None,
*,
options: EmbeddingOptionsT | None = None,
usage: dict[str, Any] | None = None,
additional_properties: dict[str, Any] | None = None,
) -> None:
super().__init__(embeddings or [])
self.options = options
self.usage = usage
self.additional_properties = additional_properties or {}
# endregion
@@ -37,6 +37,7 @@ _IMPORTS: dict[str, tuple[str, str]] = {
"AzureOpenAIAssistantsOptions": ("agent_framework.azure._assistants_client", "agent-framework-core"),
"AzureOpenAIChatClient": ("agent_framework.azure._chat_client", "agent-framework-core"),
"AzureOpenAIChatOptions": ("agent_framework.azure._chat_client", "agent-framework-core"),
"AzureOpenAIEmbeddingClient": ("agent_framework.azure._embedding_client", "agent-framework-core"),
"AzureOpenAIResponsesClient": ("agent_framework.azure._responses_client", "agent-framework-core"),
"AzureOpenAIResponsesOptions": ("agent_framework.azure._responses_client", "agent-framework-core"),
"AzureOpenAISettings": ("agent_framework.azure._shared", "agent-framework-core"),
@@ -22,6 +22,7 @@ from agent_framework_durabletask import (
from agent_framework.azure._assistants_client import AzureOpenAIAssistantsClient
from agent_framework.azure._chat_client import AzureOpenAIChatClient
from agent_framework.azure._embedding_client import AzureOpenAIEmbeddingClient
from agent_framework.azure._entra_id_authentication import AzureCredentialTypes, AzureTokenProvider
from agent_framework.azure._responses_client import AzureOpenAIResponsesClient
from agent_framework.azure._shared import AzureOpenAISettings
@@ -41,6 +42,7 @@ __all__ = [
"AzureCredentialTypes",
"AzureOpenAIAssistantsClient",
"AzureOpenAIChatClient",
"AzureOpenAIEmbeddingClient",
"AzureOpenAIResponsesClient",
"AzureOpenAISettings",
"AzureTokenProvider",
@@ -0,0 +1,136 @@
# Copyright (c) Microsoft. All rights reserved.
from __future__ import annotations
import sys
from collections.abc import Mapping
from typing import Generic
from openai.lib.azure import AsyncAzureOpenAI
from agent_framework.observability import EmbeddingTelemetryLayer
from agent_framework.openai import OpenAIEmbeddingOptions
from agent_framework.openai._embedding_client import RawOpenAIEmbeddingClient
from .._settings import load_settings
from ._entra_id_authentication import AzureCredentialTypes, AzureTokenProvider
from ._shared import (
AzureOpenAIConfigMixin,
AzureOpenAISettings,
_apply_azure_defaults,
)
if sys.version_info >= (3, 13):
from typing import TypeVar # type: ignore # pragma: no cover
else:
from typing_extensions import TypeVar # type: ignore # pragma: no cover
if sys.version_info >= (3, 11):
from typing import TypedDict # type: ignore # pragma: no cover
else:
from typing_extensions import TypedDict # type: ignore # pragma: no cover
AzureOpenAIEmbeddingOptionsT = TypeVar(
"AzureOpenAIEmbeddingOptionsT",
bound=TypedDict, # type: ignore[valid-type]
default="OpenAIEmbeddingOptions",
covariant=True,
)
class AzureOpenAIEmbeddingClient(
AzureOpenAIConfigMixin,
EmbeddingTelemetryLayer[str, list[float], AzureOpenAIEmbeddingOptionsT],
RawOpenAIEmbeddingClient[AzureOpenAIEmbeddingOptionsT],
Generic[AzureOpenAIEmbeddingOptionsT],
):
"""Azure OpenAI embedding client with telemetry support.
Keyword Args:
api_key: The API key. If provided, will override the value in the env vars or .env file.
Can also be set via environment variable AZURE_OPENAI_API_KEY.
deployment_name: The deployment name. If provided, will override the value
(embedding_deployment_name) in the env vars or .env file.
Can also be set via environment variable AZURE_OPENAI_EMBEDDING_DEPLOYMENT_NAME.
endpoint: The deployment endpoint.
Can also be set via environment variable AZURE_OPENAI_ENDPOINT.
base_url: The deployment base URL.
Can also be set via environment variable AZURE_OPENAI_BASE_URL.
api_version: The deployment API version.
Can also be set via environment variable AZURE_OPENAI_API_VERSION.
token_endpoint: The token endpoint to request an Azure token.
Can also be set via environment variable AZURE_OPENAI_TOKEN_ENDPOINT.
credential: Azure credential or token provider for authentication.
default_headers: Default headers for HTTP requests.
async_client: An existing client to use.
env_file_path: Path to .env file for settings.
env_file_encoding: Encoding for .env file.
Examples:
.. code-block:: python
from agent_framework.azure import AzureOpenAIEmbeddingClient
# Using environment variables
# Set AZURE_OPENAI_ENDPOINT=https://your-endpoint.openai.azure.com
# Set AZURE_OPENAI_EMBEDDING_DEPLOYMENT_NAME=text-embedding-3-small
# Set AZURE_OPENAI_API_KEY=your-key
client = AzureOpenAIEmbeddingClient()
# Or passing parameters directly
client = AzureOpenAIEmbeddingClient(
endpoint="https://your-endpoint.openai.azure.com",
deployment_name="text-embedding-3-small",
api_key="your-key",
)
result = await client.get_embeddings(["Hello, world!"])
"""
def __init__(
self,
*,
api_key: str | None = None,
deployment_name: str | None = None,
endpoint: str | None = None,
base_url: str | None = None,
api_version: str | None = None,
token_endpoint: str | None = None,
credential: AzureCredentialTypes | AzureTokenProvider | None = None,
default_headers: Mapping[str, str] | None = None,
async_client: AsyncAzureOpenAI | None = None,
env_file_path: str | None = None,
env_file_encoding: str | None = None,
) -> None:
"""Initialize an Azure OpenAI embedding client."""
azure_openai_settings = load_settings(
AzureOpenAISettings,
env_prefix="AZURE_OPENAI_",
api_key=api_key,
base_url=base_url,
endpoint=endpoint,
embedding_deployment_name=deployment_name,
api_version=api_version,
env_file_path=env_file_path,
env_file_encoding=env_file_encoding,
token_endpoint=token_endpoint,
)
_apply_azure_defaults(azure_openai_settings)
if not azure_openai_settings.get("embedding_deployment_name"):
raise ValueError(
"Azure OpenAI embedding deployment name is required. Set via 'deployment_name' parameter "
"or 'AZURE_OPENAI_EMBEDDING_DEPLOYMENT_NAME' environment variable."
)
super().__init__(
deployment_name=azure_openai_settings["embedding_deployment_name"], # type: ignore[arg-type]
endpoint=azure_openai_settings["endpoint"],
base_url=azure_openai_settings["base_url"],
api_version=azure_openai_settings["api_version"], # type: ignore
api_key=azure_openai_settings["api_key"].get_secret_value() if azure_openai_settings["api_key"] else None,
token_endpoint=azure_openai_settings["token_endpoint"],
credential=credential,
default_headers=default_headers,
client=async_client,
)
@@ -53,6 +53,8 @@ class AzureOpenAISettings(TypedDict, total=False):
Resource Management > Deployments in the Azure portal or, alternatively,
under Management > Deployments in Azure AI Foundry.
Can be set via environment variable AZURE_OPENAI_RESPONSES_DEPLOYMENT_NAME.
embedding_deployment_name: The name of the Azure Embedding deployment.
Can be set via environment variable AZURE_OPENAI_EMBEDDING_DEPLOYMENT_NAME.
api_key: The API key for the Azure deployment. This value can be
found in the Keys & Endpoint section when examining your resource in
the Azure portal. You can use either KEY1 or KEY2.
@@ -95,6 +97,7 @@ class AzureOpenAISettings(TypedDict, total=False):
chat_deployment_name: str | None
responses_deployment_name: str | None
embedding_deployment_name: str | None
endpoint: str | None
base_url: str | None
api_key: SecretString | None
@@ -59,7 +59,9 @@ if TYPE_CHECKING: # pragma: no cover
ChatResponse,
ChatResponseUpdate,
Content,
EmbeddingGenerationOptions,
FinishReason,
GeneratedEmbeddings,
Message,
ResponseStream,
)
@@ -70,6 +72,7 @@ __all__ = [
"OBSERVABILITY_SETTINGS",
"AgentTelemetryLayer",
"ChatTelemetryLayer",
"EmbeddingTelemetryLayer",
"OtelAttr",
"configure_otel_providers",
"create_metric_views",
@@ -80,6 +83,8 @@ __all__ = [
]
EmbeddingInputT = TypeVar("EmbeddingInputT", default="str")
EmbeddingT = TypeVar("EmbeddingT", default="list[float]")
AgentT = TypeVar("AgentT", bound="SupportsAgentRun")
ChatClientT = TypeVar("ChatClientT", bound="SupportsChatGetResponse[Any]")
@@ -259,6 +264,7 @@ class OtelAttr(str, Enum):
# Operation names
CHAT_COMPLETION_OPERATION = "chat"
EMBEDDING_OPERATION = "embeddings"
TOOL_EXECUTION_OPERATION = "execute_tool"
# Describes GenAI agent creation and is usually applicable when working with remote agent services.
AGENT_CREATE_OPERATION = "create_agent"
@@ -1273,6 +1279,70 @@ class ChatTelemetryLayer(Generic[OptionsCoT]):
return _get_response()
EmbeddingOptionsCoT = TypeVar(
"EmbeddingOptionsCoT",
bound=TypedDict, # type: ignore[valid-type]
default="EmbeddingGenerationOptions",
covariant=True,
)
class EmbeddingTelemetryLayer(Generic[EmbeddingInputT, EmbeddingT, EmbeddingOptionsCoT]):
"""Layer that wraps embedding client get_embeddings with OpenTelemetry tracing."""
def __init__(self, *args: Any, otel_provider_name: str | None = None, **kwargs: Any) -> None:
"""Initialize telemetry attributes and histograms."""
super().__init__(*args, **kwargs)
self.token_usage_histogram = _get_token_usage_histogram()
self.duration_histogram = _get_duration_histogram()
self.otel_provider_name = otel_provider_name or getattr(self, "OTEL_PROVIDER_NAME", "unknown")
async def get_embeddings(
self,
values: Sequence[EmbeddingInputT],
*,
options: EmbeddingOptionsCoT | None = None,
) -> GeneratedEmbeddings[EmbeddingT]:
"""Trace embedding generation with OpenTelemetry spans and metrics."""
global OBSERVABILITY_SETTINGS
super_get_embeddings = super().get_embeddings # type: ignore[misc]
if not OBSERVABILITY_SETTINGS.ENABLED:
return await super_get_embeddings(values, options=options) # type: ignore[no-any-return]
opts: dict[str, Any] = options or {} # type: ignore[assignment]
provider_name = str(self.otel_provider_name)
model_id = opts.get("model_id") or getattr(self, "model_id", None) or "unknown"
service_url_func = getattr(self, "service_url", None)
service_url = str(service_url_func() if callable(service_url_func) else "unknown")
attributes = _get_span_attributes(
operation_name=OtelAttr.EMBEDDING_OPERATION,
provider_name=provider_name,
model=model_id,
service_url=service_url,
)
with _get_span(attributes=attributes, span_name_attribute=OtelAttr.REQUEST_MODEL) as span:
start_time_stamp = perf_counter()
try:
result = await super_get_embeddings(values, options=options)
except Exception as exception:
capture_exception(span=span, exception=exception, timestamp=time_ns())
raise
duration = perf_counter() - start_time_stamp
response_attributes: dict[str, Any] = {**attributes}
if result.usage and "prompt_tokens" in result.usage:
response_attributes[OtelAttr.INPUT_TOKENS] = result.usage["prompt_tokens"]
_capture_response(
span=span,
attributes=response_attributes,
token_usage_histogram=self.token_usage_histogram,
operation_duration_histogram=self.duration_histogram,
duration=duration,
)
return result # type: ignore[no-any-return]
class AgentTelemetryLayer:
"""Layer that wraps agent run with OpenTelemetry tracing."""
@@ -19,6 +19,7 @@ from ._assistants_client import (
OpenAIAssistantsOptions,
)
from ._chat_client import OpenAIChatClient, OpenAIChatOptions
from ._embedding_client import OpenAIEmbeddingClient, OpenAIEmbeddingOptions
from ._exceptions import ContentFilterResultSeverity, OpenAIContentFilterException
from ._responses_client import (
OpenAIContinuationToken,
@@ -38,6 +39,8 @@ __all__ = [
"OpenAIChatOptions",
"OpenAIContentFilterException",
"OpenAIContinuationToken",
"OpenAIEmbeddingClient",
"OpenAIEmbeddingOptions",
"OpenAIResponsesClient",
"OpenAIResponsesOptions",
"OpenAISettings",
@@ -0,0 +1,211 @@
# Copyright (c) Microsoft. All rights reserved.
from __future__ import annotations
import base64
import struct
import sys
from collections.abc import Awaitable, Callable, Mapping, Sequence
from typing import Any, Generic, Literal, TypedDict
from openai import AsyncOpenAI
from .._clients import BaseEmbeddingClient
from .._settings import load_settings
from .._types import Embedding, EmbeddingGenerationOptions, GeneratedEmbeddings
from ..observability import EmbeddingTelemetryLayer
from ._shared import OpenAIBase, OpenAIConfigMixin, OpenAISettings
if sys.version_info >= (3, 13):
from typing import TypeVar # type: ignore # pragma: no cover
else:
from typing_extensions import TypeVar # type: ignore # pragma: no cover
class OpenAIEmbeddingOptions(EmbeddingGenerationOptions, total=False):
"""OpenAI-specific embedding options.
Extends EmbeddingGenerationOptions with OpenAI-specific fields.
Examples:
.. code-block:: python
from agent_framework.openai import OpenAIEmbeddingOptions
options: OpenAIEmbeddingOptions = {
"model_id": "text-embedding-3-small",
"dimensions": 1536,
"encoding_format": "float",
}
"""
encoding_format: Literal["float", "base64"]
user: str
OpenAIEmbeddingOptionsT = TypeVar(
"OpenAIEmbeddingOptionsT",
bound=TypedDict, # type: ignore[valid-type]
default="OpenAIEmbeddingOptions",
covariant=True,
)
class RawOpenAIEmbeddingClient(
OpenAIBase,
BaseEmbeddingClient[str, list[float], OpenAIEmbeddingOptionsT],
Generic[OpenAIEmbeddingOptionsT],
):
"""Raw OpenAI embedding client without telemetry."""
def service_url(self) -> str:
"""Get the URL of the service."""
return str(self.client.base_url) if self.client else "Unknown"
async def get_embeddings(
self,
values: Sequence[str],
*,
options: OpenAIEmbeddingOptionsT | None = None,
) -> GeneratedEmbeddings[list[float]]:
"""Call the OpenAI embeddings API.
Args:
values: The text values to generate embeddings for.
options: Optional embedding generation options.
Returns:
Generated embeddings with usage metadata.
Raises:
ValueError: If model_id is not provided or values is empty.
"""
if not values:
return GeneratedEmbeddings([], options=options)
opts: dict[str, Any] = dict(options) if options else {}
model = opts.get("model_id") or self.model_id
if not model:
raise ValueError("model_id is required")
kwargs: dict[str, Any] = {"input": list(values), "model": model}
if dimensions := opts.get("dimensions"):
kwargs["dimensions"] = dimensions
if encoding_format := opts.get("encoding_format"):
kwargs["encoding_format"] = encoding_format
if user := opts.get("user"):
kwargs["user"] = user
response = await (await self._ensure_client()).embeddings.create(**kwargs)
encoding = kwargs.get("encoding_format", "float")
embeddings: list[Embedding[list[float]]] = []
for item in response.data:
vector: list[float]
if encoding == "base64" and isinstance(item.embedding, str):
# Decode base64-encoded floats (little-endian IEEE 754)
raw = base64.b64decode(item.embedding)
vector = list(struct.unpack(f"<{len(raw) // 4}f", raw))
else:
vector = item.embedding # type: ignore[assignment]
embeddings.append(
Embedding(
vector=vector,
dimensions=len(vector),
model_id=response.model,
)
)
usage_dict: dict[str, Any] | None = None
if response.usage:
usage_dict = {
"prompt_tokens": response.usage.prompt_tokens,
"total_tokens": response.usage.total_tokens,
}
return GeneratedEmbeddings(embeddings, options=options, usage=usage_dict)
class OpenAIEmbeddingClient(
OpenAIConfigMixin,
EmbeddingTelemetryLayer[str, list[float], OpenAIEmbeddingOptionsT],
RawOpenAIEmbeddingClient[OpenAIEmbeddingOptionsT],
Generic[OpenAIEmbeddingOptionsT],
):
"""OpenAI embedding client with telemetry support.
Keyword Args:
model_id: The embedding model ID (e.g. "text-embedding-3-small").
Can also be set via environment variable OPENAI_EMBEDDING_MODEL_ID.
api_key: OpenAI API key.
Can also be set via environment variable OPENAI_API_KEY.
org_id: OpenAI organization ID.
default_headers: Additional HTTP headers.
async_client: Pre-configured AsyncOpenAI client.
base_url: Custom API base URL.
env_file_path: Path to .env file for settings.
env_file_encoding: Encoding for .env file.
Examples:
.. code-block:: python
from agent_framework.openai import OpenAIEmbeddingClient
# Using environment variables
# Set OPENAI_API_KEY=sk-...
# Set OPENAI_EMBEDDING_MODEL_ID=text-embedding-3-small
client = OpenAIEmbeddingClient()
# Or passing parameters directly
client = OpenAIEmbeddingClient(
model_id="text-embedding-3-small",
api_key="sk-...",
)
# Generate embeddings
result = await client.get_embeddings(["Hello, world!"])
print(result[0].vector)
"""
def __init__(
self,
*,
model_id: str | None = None,
api_key: str | Callable[[], str | Awaitable[str]] | None = None,
org_id: str | None = None,
default_headers: Mapping[str, str] | None = None,
async_client: AsyncOpenAI | None = None,
base_url: str | None = None,
env_file_path: str | None = None,
env_file_encoding: str | None = None,
) -> None:
"""Initialize an OpenAI embedding client."""
openai_settings = load_settings(
OpenAISettings,
env_prefix="OPENAI_",
api_key=api_key,
base_url=base_url,
org_id=org_id,
embedding_model_id=model_id,
env_file_path=env_file_path,
env_file_encoding=env_file_encoding,
)
if not async_client and not openai_settings["api_key"]:
raise ValueError(
"OpenAI API key is required. Set via 'api_key' parameter or 'OPENAI_API_KEY' environment variable."
)
if not openai_settings["embedding_model_id"]:
raise ValueError(
"OpenAI embedding model ID is required. "
"Set via 'model_id' parameter or 'OPENAI_EMBEDDING_MODEL_ID' environment variable."
)
super().__init__(
model_id=openai_settings["embedding_model_id"],
api_key=self._get_api_key(openai_settings["api_key"]),
base_url=openai_settings["base_url"] if openai_settings["base_url"] else None,
org_id=openai_settings["org_id"],
default_headers=default_headers,
client=async_client,
)
@@ -92,6 +92,8 @@ class OpenAISettings(TypedDict, total=False):
Can be set via environment variable OPENAI_CHAT_MODEL_ID.
responses_model_id: The OpenAI responses model ID to use, for example, gpt-4o or o1.
Can be set via environment variable OPENAI_RESPONSES_MODEL_ID.
embedding_model_id: The OpenAI embedding model ID to use, for example, text-embedding-3-small.
Can be set via environment variable OPENAI_EMBEDDING_MODEL_ID.
Examples:
.. code-block:: python
@@ -115,6 +117,7 @@ class OpenAISettings(TypedDict, total=False):
org_id: str | None
chat_model_id: str | None
responses_model_id: str | None
embedding_model_id: str | None
class OpenAIBase(SerializationMixin):
@@ -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
@@ -0,0 +1,70 @@
# Copyright (c) Microsoft. All rights reserved.
# Run with: uv run samples/02-agents/embeddings/azure_openai_embeddings.py
import asyncio
from agent_framework.azure import AzureOpenAIEmbeddingClient
from dotenv import load_dotenv
load_dotenv()
"""Azure OpenAI Embedding Client Example
This sample demonstrates how to generate embeddings using the Azure OpenAI embedding client.
It supports both API key and Azure credential authentication.
Prerequisites:
Set the following environment variables or add them to a .env file:
- AZURE_OPENAI_ENDPOINT: Your Azure OpenAI endpoint URL
- AZURE_OPENAI_EMBEDDING_DEPLOYMENT_NAME: The embedding model deployment name
- AZURE_OPENAI_API_KEY: Your API key (or use Azure credential instead)
"""
async def main() -> None:
"""Generate embeddings with Azure OpenAI."""
# 1. Create a client using environment variables.
# Reads AZURE_OPENAI_ENDPOINT, AZURE_OPENAI_EMBEDDING_DEPLOYMENT_NAME,
# and AZURE_OPENAI_API_KEY from environment.
client = AzureOpenAIEmbeddingClient()
# 2. Generate a single embedding.
result = await client.get_embeddings(["Hello, world!"])
print(f"Single embedding dimensions: {result[0].dimensions}")
print(f"First 5 values: {result[0].vector[:5]}")
print(f"Model: {result[0].model_id}")
print(f"Usage: {result.usage}")
print()
# 3. Generate embeddings for multiple inputs.
texts = [
"The weather is sunny today.",
"It is raining outside.",
"Machine learning is fascinating.",
]
result = await client.get_embeddings(texts)
print(f"Batch of {len(result)} embeddings, each with {result[0].dimensions} dimensions")
print()
# 4. Generate embeddings with custom dimensions.
result = await client.get_embeddings(["Custom dimensions example"], options={"dimensions": 256})
print(f"Custom dimensions: {result[0].dimensions}")
if __name__ == "__main__":
asyncio.run(main())
"""
Sample output:
Single embedding dimensions: 1536
First 5 values: [0.012, -0.034, 0.056, -0.078, 0.090]
Model: text-embedding-3-small
Usage: {'prompt_tokens': 4, 'total_tokens': 4}
Batch of 3 embeddings, each with 1536 dimensions
Custom dimensions: 256
"""
@@ -0,0 +1,65 @@
# Copyright (c) Microsoft. All rights reserved.
# Run with: uv run samples/02-agents/embeddings/openai_embeddings.py
import asyncio
from agent_framework.openai import OpenAIEmbeddingClient
from dotenv import load_dotenv
load_dotenv()
"""OpenAI Embedding Client Example
This sample demonstrates how to generate embeddings using the OpenAI embedding client.
It shows single and batch embedding generation, as well as custom dimensions.
Prerequisites:
Set the OPENAI_API_KEY environment variable or add it to a .env file.
"""
async def main() -> None:
"""Generate embeddings with OpenAI."""
client = OpenAIEmbeddingClient(model_id="text-embedding-3-small")
# 1. Generate a single embedding.
result = await client.get_embeddings(["Hello, world!"])
print(f"Single embedding dimensions: {result[0].dimensions}")
print(f"First 5 values: {result[0].vector[:5]}")
print(f"Model: {result[0].model_id}")
print(f"Usage: {result.usage}")
print()
# 2. Generate embeddings for multiple inputs.
texts = [
"The weather is sunny today.",
"It is raining outside.",
"Machine learning is fascinating.",
]
result = await client.get_embeddings(texts)
print(f"Batch of {len(result)} embeddings, each with {result[0].dimensions} dimensions")
print(f"First embedding vector: {result[0].vector[:5]}") # Print first 5 values of the first embedding
print()
# 3. Generate embeddings with custom dimensions.
result = await client.get_embeddings(["Custom dimensions example"], options={"dimensions": 256})
print(f"Custom dimensions: {result[0].dimensions}")
print(f"First 5 values: {result[0].vector[:5]}")
if __name__ == "__main__":
asyncio.run(main())
"""
Sample output:
Single embedding dimensions: 1536
First 5 values: [0.012, -0.034, 0.056, -0.078, 0.090]
Model: text-embedding-3-small
Usage: {'prompt_tokens': 4, 'total_tokens': 4}
Batch of 3 embeddings, each with 1536 dimensions
Custom dimensions: 256
"""