Inject user agent header at runtime

This commit is contained in:
Tao Chen
2026-04-22 21:33:02 -07:00
Unverified
parent adcd2d33f5
commit 38ff34cc17
8 changed files with 157 additions and 6 deletions
@@ -32,7 +32,7 @@ from agent_framework._clients import BaseChatClient
from agent_framework._compaction import CompactionStrategy, TokenizerProtocol
from agent_framework._middleware import ChatAndFunctionMiddlewareTypes, ChatMiddlewareLayer
from agent_framework._settings import SecretString
from agent_framework._telemetry import USER_AGENT_KEY
from agent_framework._telemetry import USER_AGENT_KEY, get_user_agent_extra_headers
from agent_framework._tools import (
SHELL_TOOL_KIND_VALUE,
FunctionInvocationConfiguration,
@@ -482,6 +482,13 @@ class RawOpenAIChatClient( # type: ignore[misc]
client = self.client
validated_options = await self._validate_options(options)
run_options = await self._prepare_options(messages, validated_options)
ua_headers = get_user_agent_extra_headers()
if ua_headers:
existing = run_options.get("extra_headers")
if existing is None:
run_options["extra_headers"] = ua_headers
elif USER_AGENT_KEY not in existing:
run_options["extra_headers"] = {**existing, **ua_headers}
return client, run_options, validated_options
def _handle_request_error(self, ex: Exception) -> NoReturn:
@@ -525,6 +532,7 @@ class RawOpenAIChatClient( # type: ignore[misc]
stream_response = await client.responses.retrieve(
continuation_token["response_id"],
stream=True,
extra_headers=get_user_agent_extra_headers(),
)
async for chunk in stream_response:
yield self._parse_chunk_from_openai(
@@ -572,7 +580,10 @@ class RawOpenAIChatClient( # type: ignore[misc]
client = self.client
validated_options = await self._validate_options(options)
try:
response = await client.responses.retrieve(continuation_token["response_id"])
response = await client.responses.retrieve(
continuation_token["response_id"],
extra_headers=get_user_agent_extra_headers(),
)
except Exception as ex:
self._handle_request_error(ex)
return self._parse_response_from_openai(response, options=validated_options)
@@ -22,7 +22,7 @@ from agent_framework._compaction import CompactionStrategy, TokenizerProtocol
from agent_framework._docstrings import apply_layered_docstring
from agent_framework._middleware import ChatAndFunctionMiddlewareTypes, ChatMiddlewareLayer
from agent_framework._settings import SecretString
from agent_framework._telemetry import USER_AGENT_KEY
from agent_framework._telemetry import USER_AGENT_KEY, get_user_agent_extra_headers
from agent_framework._tools import (
FunctionInvocationConfiguration,
FunctionInvocationLayer,
@@ -671,6 +671,16 @@ class RawOpenAIChatCompletionClient( # type: ignore[misc]
run_options["response_format"] = response_format
else:
run_options["response_format"] = type_to_response_format_param(response_format)
# runtime user-agent header
ua_headers = get_user_agent_extra_headers()
if ua_headers:
existing = run_options.get("extra_headers")
if existing is None:
run_options["extra_headers"] = ua_headers
elif USER_AGENT_KEY not in existing:
run_options["extra_headers"] = {**existing, **ua_headers}
return run_options
def _parse_response_from_openai(self, response: ChatCompletion, options: Mapping[str, Any]) -> ChatResponse:
@@ -10,7 +10,7 @@ from typing import TYPE_CHECKING, Any, ClassVar, Generic, Literal, TypedDict, ov
from agent_framework._clients import BaseEmbeddingClient
from agent_framework._settings import SecretString
from agent_framework._telemetry import USER_AGENT_KEY
from agent_framework._telemetry import USER_AGENT_KEY, get_user_agent_extra_headers
from agent_framework._types import Embedding, EmbeddingGenerationOptions, GeneratedEmbeddings, UsageDetails
from agent_framework.observability import EmbeddingTelemetryLayer
from openai import AsyncAzureOpenAI, AsyncOpenAI
@@ -282,6 +282,13 @@ class RawOpenAIEmbeddingClient(
kwargs["encoding_format"] = encoding_format
if user := opts.get("user"):
kwargs["user"] = user
ua_headers = get_user_agent_extra_headers()
if ua_headers:
existing = kwargs.get("extra_headers")
if existing is None:
kwargs["extra_headers"] = ua_headers
elif USER_AGENT_KEY not in existing:
kwargs["extra_headers"] = {**existing, **ua_headers}
response = await self.client.embeddings.create(**kwargs) # type: ignore[union-attr]
@@ -8,7 +8,7 @@ from copy import copy
from typing import TYPE_CHECKING, Any, Literal, Union
from agent_framework._settings import SecretString, load_settings
from agent_framework._telemetry import APP_INFO, prepend_agent_framework_to_user_agent
from agent_framework._telemetry import APP_INFO
from agent_framework.exceptions import SettingNotFoundError
from openai import AsyncAzureOpenAI, AsyncOpenAI, AsyncStream, _legacy_response # type: ignore
from openai.types import Completion
@@ -174,7 +174,6 @@ def load_openai_service_settings(
merged_headers = dict(copy(default_headers)) if default_headers else {}
if APP_INFO:
merged_headers.update(APP_INFO)
merged_headers = prepend_agent_framework_to_user_agent(merged_headers)
api_key_callable = api_key if callable(api_key) else None
api_key_str = api_key if not callable(api_key) else None