Fix other providers

This commit is contained in:
Tao Chen
2026-04-23 14:31:26 -07:00
Unverified
parent 03f7dc86d3
commit 28a86d6d73
18 changed files with 48 additions and 45 deletions
@@ -6,13 +6,13 @@ from collections.abc import Sequence
from typing import Any, ClassVar, Generic, TypedDict
from agent_framework import (
AGENT_FRAMEWORK_USER_AGENT,
ChatAndFunctionMiddlewareTypes,
ChatMiddlewareLayer,
FunctionInvocationConfiguration,
FunctionInvocationLayer,
)
from agent_framework._settings import SecretString, load_settings
from agent_framework._telemetry import get_user_agent
from agent_framework.observability import ChatTelemetryLayer
from anthropic import AsyncAnthropicBedrock
@@ -94,7 +94,7 @@ class RawAnthropicBedrockClient(RawAnthropicClient[AnthropicOptionsT], Generic[A
aws_profile=settings.get("aws_profile"),
aws_session_token=session_token_secret.get_secret_value() if session_token_secret is not None else None,
base_url=settings.get("anthropic_bedrock_base_url"),
default_headers={"User-Agent": AGENT_FRAMEWORK_USER_AGENT},
default_headers={"User-Agent": get_user_agent()},
)
super().__init__(
@@ -8,7 +8,6 @@ from collections.abc import AsyncIterable, Awaitable, Callable, Mapping, Sequenc
from typing import Any, ClassVar, Final, Generic, Literal, TypedDict
from agent_framework import (
AGENT_FRAMEWORK_USER_AGENT,
Annotation,
BaseChatClient,
ChatAndFunctionMiddlewareTypes,
@@ -28,6 +27,7 @@ from agent_framework import (
tool,
)
from agent_framework._settings import SecretString, load_settings
from agent_framework._telemetry import get_user_agent
from agent_framework._tools import SHELL_TOOL_KIND_VALUE
from agent_framework._types import _get_data_bytes_as_str # type: ignore
from agent_framework.observability import ChatTelemetryLayer
@@ -332,7 +332,7 @@ class RawAnthropicClient(
anthropic_client = AsyncAnthropic(
api_key=api_key_secret.get_secret_value(),
default_headers={"User-Agent": AGENT_FRAMEWORK_USER_AGENT},
default_headers={"User-Agent": get_user_agent()},
)
# Initialize parent
@@ -604,7 +604,7 @@ class RawAnthropicClient(
run_options["betas"] = self._prepare_betas(options)
# extra headers
run_options["extra_headers"] = {"User-Agent": AGENT_FRAMEWORK_USER_AGENT}
run_options["extra_headers"] = {"User-Agent": get_user_agent()}
# Handle user option -> metadata.user_id (Anthropic uses metadata.user_id instead of user)
if user := run_options.pop("user", None):
@@ -6,13 +6,13 @@ from collections.abc import Awaitable, Callable, Sequence
from typing import Any, ClassVar, Generic, TypedDict
from agent_framework import (
AGENT_FRAMEWORK_USER_AGENT,
ChatAndFunctionMiddlewareTypes,
ChatMiddlewareLayer,
FunctionInvocationConfiguration,
FunctionInvocationLayer,
)
from agent_framework._settings import SecretString, load_settings
from agent_framework._telemetry import get_user_agent
from agent_framework.observability import ChatTelemetryLayer
from anthropic import AsyncAnthropicFoundry
@@ -91,14 +91,14 @@ class RawAnthropicFoundryClient(RawAnthropicClient[AnthropicOptionsT], Generic[A
base_url=base_url_setting,
api_key=api_key_value,
azure_ad_token_provider=azure_ad_token_provider,
default_headers={"User-Agent": AGENT_FRAMEWORK_USER_AGENT},
default_headers={"User-Agent": get_user_agent()},
)
else:
anthropic_client = AsyncAnthropicFoundry(
resource=resource_setting,
api_key=api_key_value,
azure_ad_token_provider=azure_ad_token_provider,
default_headers={"User-Agent": AGENT_FRAMEWORK_USER_AGENT},
default_headers={"User-Agent": get_user_agent()},
)
super().__init__(
@@ -6,13 +6,13 @@ from collections.abc import Sequence
from typing import TYPE_CHECKING, Any, ClassVar, Generic, TypedDict
from agent_framework import (
AGENT_FRAMEWORK_USER_AGENT,
ChatAndFunctionMiddlewareTypes,
ChatMiddlewareLayer,
FunctionInvocationConfiguration,
FunctionInvocationLayer,
)
from agent_framework._settings import load_settings
from agent_framework._telemetry import get_user_agent
from agent_framework.observability import ChatTelemetryLayer
from anthropic import NOT_GIVEN, AsyncAnthropicVertex
@@ -89,7 +89,7 @@ class RawAnthropicVertexClient(RawAnthropicClient[AnthropicOptionsT], Generic[An
access_token=access_token,
credentials=credentials,
base_url=settings.get("anthropic_vertex_base_url"),
default_headers={"User-Agent": AGENT_FRAMEWORK_USER_AGENT},
default_headers={"User-Agent": get_user_agent()},
)
super().__init__(
@@ -3,7 +3,8 @@
from unittest.mock import AsyncMock, MagicMock, patch
import pytest
from agent_framework import AGENT_FRAMEWORK_USER_AGENT, ChatMiddlewareLayer, FunctionInvocationLayer
from agent_framework import ChatMiddlewareLayer, FunctionInvocationLayer
from agent_framework._telemetry import get_user_agent
from agent_framework.observability import ChatTelemetryLayer
from agent_framework_anthropic import (
@@ -61,7 +62,7 @@ def test_raw_anthropic_foundry_client_creates_sdk_client_from_settings(tmp_path)
resource="test-resource",
api_key="test-key",
azure_ad_token_provider=None,
default_headers={"User-Agent": AGENT_FRAMEWORK_USER_AGENT},
default_headers={"User-Agent": get_user_agent()},
)
@@ -85,7 +86,7 @@ def test_raw_anthropic_foundry_client_creates_sdk_client_from_base_url_settings(
base_url="https://test-resource.services.ai.azure.com/anthropic/",
api_key="test-key",
azure_ad_token_provider=None,
default_headers={"User-Agent": AGENT_FRAMEWORK_USER_AGENT},
default_headers={"User-Agent": get_user_agent()},
)
@@ -130,7 +131,7 @@ def test_raw_anthropic_bedrock_client_creates_sdk_client_from_arguments() -> Non
aws_profile=None,
aws_session_token=None,
base_url=None,
default_headers={"User-Agent": AGENT_FRAMEWORK_USER_AGENT},
default_headers={"User-Agent": get_user_agent()},
)
@@ -152,5 +153,5 @@ def test_raw_anthropic_vertex_client_creates_sdk_client_from_arguments() -> None
access_token=None,
credentials=None,
base_url=None,
default_headers={"User-Agent": AGENT_FRAMEWORK_USER_AGENT},
default_headers={"User-Agent": get_user_agent()},
)
@@ -14,7 +14,6 @@ from collections.abc import Awaitable, Callable
from typing import TYPE_CHECKING, Any, ClassVar, Literal, TypedDict, overload
from agent_framework import (
AGENT_FRAMEWORK_USER_AGENT,
AgentSession,
Annotation,
Content,
@@ -25,6 +24,7 @@ from agent_framework import (
SupportsGetEmbeddings,
load_settings,
)
from agent_framework._telemetry import get_user_agent
from agent_framework.exceptions import SettingNotFoundError
from azure.core.credentials import AzureKeyCredential, TokenCredential
from azure.core.credentials_async import AsyncTokenCredential
@@ -535,7 +535,7 @@ class AzureAISearchContextProvider(ContextProvider):
endpoint=self.endpoint,
index_name=self.index_name,
credential=self.credential,
user_agent=AGENT_FRAMEWORK_USER_AGENT,
user_agent=get_user_agent(),
)
self._index_client: SearchIndexClient | None = None
@@ -544,7 +544,7 @@ class AzureAISearchContextProvider(ContextProvider):
self._index_client = SearchIndexClient(
endpoint=self.endpoint,
credential=self.credential,
user_agent=AGENT_FRAMEWORK_USER_AGENT,
user_agent=get_user_agent(),
)
self._knowledge_base_initialized = False
@@ -640,7 +640,7 @@ class AzureAISearchContextProvider(ContextProvider):
self._index_client = SearchIndexClient(
endpoint=self.endpoint,
credential=self.credential,
user_agent=AGENT_FRAMEWORK_USER_AGENT,
user_agent=get_user_agent(),
)
if not self.index_name:
logger.warning("Cannot auto-discover vector field: index_name is not set.")
@@ -740,7 +740,7 @@ class AzureAISearchContextProvider(ContextProvider):
endpoint=self.endpoint,
knowledge_base_name=knowledge_base_name,
credential=self.credential,
user_agent=AGENT_FRAMEWORK_USER_AGENT,
user_agent=get_user_agent(),
)
self._knowledge_base_initialized = True
return
@@ -802,7 +802,7 @@ class AzureAISearchContextProvider(ContextProvider):
endpoint=self.endpoint,
knowledge_base_name=knowledge_base_name,
credential=self.credential,
user_agent=AGENT_FRAMEWORK_USER_AGENT,
user_agent=get_user_agent(),
)
async def _agentic_search(self, messages: list[Message]) -> list[Message]:
@@ -7,8 +7,8 @@ from __future__ import annotations
import logging
from typing import Any, TypedDict
from agent_framework import AGENT_FRAMEWORK_USER_AGENT
from agent_framework._settings import SecretString, load_settings
from agent_framework._telemetry import get_user_agent
from agent_framework._workflows._checkpoint import CheckpointID, WorkflowCheckpoint
from agent_framework._workflows._checkpoint_encoding import decode_checkpoint_value, encode_checkpoint_value
from agent_framework.exceptions import WorkflowCheckpointException
@@ -194,7 +194,7 @@ class CosmosCheckpointStorage:
self._cosmos_client = CosmosClient(
url=settings["endpoint"], # type: ignore[arg-type]
credential=credential or settings["key"].get_secret_value(), # type: ignore[arg-type,union-attr]
user_agent_suffix=AGENT_FRAMEWORK_USER_AGENT,
user_agent_suffix=get_user_agent(),
)
self._owns_client = True
@@ -10,9 +10,10 @@ import uuid
from collections.abc import Sequence
from typing import Any, ClassVar, TypedDict
from agent_framework import AGENT_FRAMEWORK_USER_AGENT, Message
from agent_framework import Message
from agent_framework._sessions import HistoryProvider
from agent_framework._settings import SecretString, load_settings
from agent_framework._telemetry import get_user_agent
from azure.core.credentials import TokenCredential
from azure.core.credentials_async import AsyncTokenCredential
from azure.cosmos import PartitionKey
@@ -121,7 +122,7 @@ class CosmosHistoryProvider(HistoryProvider):
self._cosmos_client = CosmosClient(
url=settings["endpoint"], # type: ignore[arg-type]
credential=credential or settings["key"].get_secret_value(), # type: ignore[arg-type,union-attr]
user_agent_suffix=AGENT_FRAMEWORK_USER_AGENT,
user_agent_suffix=get_user_agent(),
)
self._owns_client = True
@@ -13,7 +13,6 @@ from typing import Any, ClassVar, Generic, Literal, TypedDict
from uuid import uuid4
from agent_framework import (
AGENT_FRAMEWORK_USER_AGENT,
BaseChatClient,
ChatAndFunctionMiddlewareTypes,
ChatMiddlewareLayer,
@@ -31,6 +30,7 @@ from agent_framework import (
validate_tool_mode,
)
from agent_framework._settings import SecretString, load_settings
from agent_framework._telemetry import get_user_agent
from agent_framework.exceptions import ChatClientInvalidResponseException
from agent_framework.observability import ChatTelemetryLayer
from boto3.session import Session as Boto3Session
@@ -299,7 +299,7 @@ class BedrockChatClient(
self._bedrock_client = session.client(
"bedrock-runtime",
region_name=region,
config=BotoConfig(user_agent_extra=AGENT_FRAMEWORK_USER_AGENT),
config=BotoConfig(user_agent_extra=get_user_agent()),
)
super().__init__(
@@ -11,7 +11,6 @@ from collections.abc import Sequence
from typing import Any, ClassVar, Generic, TypedDict
from agent_framework import (
AGENT_FRAMEWORK_USER_AGENT,
BaseEmbeddingClient,
Embedding,
EmbeddingGenerationOptions,
@@ -20,6 +19,7 @@ from agent_framework import (
UsageDetails,
load_settings,
)
from agent_framework._telemetry import get_user_agent
from agent_framework.observability import EmbeddingTelemetryLayer
from boto3.session import Session as Boto3Session
from botocore.client import BaseClient
@@ -140,7 +140,7 @@ class RawBedrockEmbeddingClient(
self._bedrock_client = boto3_session.client(
"bedrock-runtime",
region_name=region_name or resolved_region,
config=BotoConfig(user_agent_extra=AGENT_FRAMEWORK_USER_AGENT),
config=BotoConfig(user_agent_extra=get_user_agent()),
)
self.model: str = settings["embedding_model"] # type: ignore[assignment] # pyright: ignore[reportTypedDictNotRequiredAccess]
@@ -75,7 +75,7 @@ def _detect_hosted_environment() -> None:
_detect_hosted_environment()
def _get_user_agent() -> str:
def get_user_agent() -> str:
"""Return the full user agent string including any registered prefixes."""
if not _user_agent_prefixes:
return AGENT_FRAMEWORK_USER_AGENT
@@ -112,7 +112,7 @@ def prepend_agent_framework_to_user_agent(headers: dict[str, Any] | None = None)
"""
if not IS_TELEMETRY_ENABLED:
return headers or {}
user_agent = _get_user_agent()
user_agent = get_user_agent()
if not headers:
return {USER_AGENT_KEY: user_agent}
headers[USER_AGENT_KEY] = f"{user_agent} {headers[USER_AGENT_KEY]}" if USER_AGENT_KEY in headers else user_agent
@@ -15,7 +15,6 @@ from collections.abc import Awaitable, Callable, Mapping, MutableMapping, Sequen
from typing import TYPE_CHECKING, Any, ClassVar, Generic, cast
from agent_framework import (
AGENT_FRAMEWORK_USER_AGENT,
AgentMiddlewareLayer,
ChatAndFunctionMiddlewareTypes,
ChatMiddlewareLayer,
@@ -28,6 +27,7 @@ from agent_framework import (
load_settings,
)
from agent_framework._compaction import CompactionStrategy, TokenizerProtocol
from agent_framework._telemetry import get_user_agent
from agent_framework.observability import AgentTelemetryLayer, ChatTelemetryLayer
from agent_framework_openai._chat_client import OpenAIChatOptions, RawOpenAIChatClient
from azure.ai.projects.aio import AIProjectClient
@@ -190,7 +190,7 @@ class RawFoundryAgentChatClient( # type: ignore[misc]
project_client_kwargs: dict[str, Any] = {
"endpoint": resolved_endpoint,
"credential": credential,
"user_agent": AGENT_FRAMEWORK_USER_AGENT,
"user_agent": get_user_agent(),
}
if allow_preview is not None:
project_client_kwargs["allow_preview"] = allow_preview
@@ -8,7 +8,6 @@ from collections.abc import Awaitable, Callable, Mapping, Sequence
from typing import TYPE_CHECKING, Any, ClassVar, Generic, Literal
from agent_framework import (
AGENT_FRAMEWORK_USER_AGENT,
ChatMiddlewareLayer,
Content,
FunctionInvocationConfiguration,
@@ -17,6 +16,7 @@ from agent_framework import (
)
from agent_framework._compaction import CompactionStrategy, TokenizerProtocol
from agent_framework._feature_stage import ExperimentalFeature, experimental
from agent_framework._telemetry import get_user_agent
from agent_framework.observability import ChatTelemetryLayer
from agent_framework_openai._chat_client import OpenAIChatOptions, RawOpenAIChatClient
from azure.ai.projects.aio import AIProjectClient
@@ -198,7 +198,7 @@ class RawFoundryChatClient( # type: ignore[misc]
project_client_kwargs: dict[str, Any] = {
"endpoint": project_endpoint,
"credential": credential, # type: ignore[arg-type]
"user_agent": AGENT_FRAMEWORK_USER_AGENT,
"user_agent": get_user_agent(),
}
if allow_preview is not None:
project_client_kwargs["allow_preview"] = allow_preview
@@ -14,13 +14,13 @@ from contextlib import AbstractAsyncContextManager
from typing import TYPE_CHECKING, Any, ClassVar
from agent_framework import (
AGENT_FRAMEWORK_USER_AGENT,
AgentSession,
ContextProvider,
Message,
SessionContext,
load_settings,
)
from agent_framework._telemetry import get_user_agent
from azure.ai.projects.aio import AIProjectClient
from azure.core.credentials import TokenCredential
from azure.core.credentials_async import AsyncTokenCredential
@@ -119,7 +119,7 @@ class FoundryMemoryProvider(ContextProvider):
project_client_kwargs: dict[str, Any] = {
"endpoint": resolved_endpoint,
"credential": credential, # type: ignore[arg-type]
"user_agent": AGENT_FRAMEWORK_USER_AGENT,
"user_agent": get_user_agent(),
}
if allow_preview is not None:
project_client_kwargs["allow_preview"] = allow_preview
@@ -12,7 +12,7 @@ from unittest.mock import AsyncMock, MagicMock, patch
import pytest
from agent_framework import ChatResponse, Content, Message, SupportsChatGetResponse, tool
from agent_framework._telemetry import AGENT_FRAMEWORK_USER_AGENT
from agent_framework._telemetry import get_user_agent
from agent_framework.exceptions import ChatClientException, ChatClientInvalidRequestException
from agent_framework_openai import OpenAIContentFilterException
from azure.ai.projects.models import MCPTool as FoundryMCPTool
@@ -199,7 +199,7 @@ def test_init_with_project_endpoint_creates_project_client() -> None:
assert factory.call_args.kwargs["endpoint"] == _TEST_FOUNDRY_PROJECT_ENDPOINT
assert factory.call_args.kwargs["credential"] is credential
assert factory.call_args.kwargs["allow_preview"] is True
assert factory.call_args.kwargs["user_agent"] == AGENT_FRAMEWORK_USER_AGENT
assert factory.call_args.kwargs["user_agent"] == get_user_agent()
def test_init_with_empty_model_raises(monkeypatch: pytest.MonkeyPatch) -> None:
@@ -7,8 +7,9 @@ import os
from unittest.mock import AsyncMock, Mock, patch
import pytest
from agent_framework import AGENT_FRAMEWORK_USER_AGENT, AgentResponse, Message
from agent_framework import AgentResponse, Message
from agent_framework._sessions import AgentSession, SessionContext
from agent_framework._telemetry import get_user_agent
from agent_framework_foundry._memory_provider import FoundryMemoryProvider
@@ -94,7 +95,7 @@ def test_init_with_project_endpoint_and_credential(mock_project_client: AsyncMoc
endpoint="https://test.project.endpoint",
credential=mock_credential,
allow_preview=True,
user_agent=AGENT_FRAMEWORK_USER_AGENT,
user_agent=get_user_agent(),
)
@@ -10,7 +10,6 @@ from typing import Any, ClassVar, Generic, cast
from uuid import uuid4
from agent_framework import (
AGENT_FRAMEWORK_USER_AGENT,
BaseChatClient,
ChatAndFunctionMiddlewareTypes,
ChatMiddlewareLayer,
@@ -28,6 +27,7 @@ from agent_framework import (
validate_tool_mode,
)
from agent_framework._settings import SecretString, load_settings
from agent_framework._telemetry import get_user_agent
from agent_framework.observability import ChatTelemetryLayer
from google import genai
from google.auth.credentials import Credentials
@@ -355,7 +355,7 @@ class RawGeminiChatClient(
)
client_kwargs: dict[str, Any] = {
"http_options": {"headers": {"x-goog-api-client": AGENT_FRAMEWORK_USER_AGENT}},
"http_options": {"headers": {"x-goog-api-client": get_user_agent()}},
}
if configured_vertexai is not None:
client_kwargs["vertexai"] = configured_vertexai
@@ -11,7 +11,7 @@ from typing import Any, Literal, TypeVar, Union, overload
from uuid import uuid4
import httpx
from agent_framework import AGENT_FRAMEWORK_USER_AGENT
from agent_framework._telemetry import get_user_agent
from agent_framework.observability import get_tracer
from azure.core.credentials import TokenCredential
from azure.core.credentials_async import AsyncTokenCredential
@@ -189,7 +189,7 @@ class PurviewClient:
payload = model.model_dump(by_alias=True, exclude_none=True, mode="json")
request_headers = {
"Authorization": f"Bearer {token}",
"User-Agent": AGENT_FRAMEWORK_USER_AGENT,
"User-Agent": get_user_agent(),
"Content-Type": "application/json",
}
if correlation_id: