Python: Capture context provider instructions in agent telemetry (#6515)

* Fix agent instructions telemetry

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

* Simplify agent instructions telemetry guard

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

* Fix observability mypy cast

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-06-17 10:28:16 +02:00
committed by GitHub
parent 7bf2d2a6d0
commit 4ff952e100
2 changed files with 246 additions and 35 deletions
@@ -1488,14 +1488,21 @@ class ChatTelemetryLayer(Generic[OptionsCoT]):
)
if stream:
agent_span = trace.get_current_span()
span = _start_streaming_span(attributes, OtelAttr.REQUEST_MODEL)
if OBSERVABILITY_SETTINGS.SENSITIVE_DATA_ENABLED and messages and span.is_recording():
system_instructions = _get_instructions_from_options(opts)
_capture_current_agent_system_instructions(
agent_span,
span,
system_instructions,
)
_capture_messages(
span=span,
provider_name=provider_name,
messages=messages,
system_instructions=opts.get("instructions"),
system_instructions=system_instructions,
)
span_state = {"closed": False}
@@ -1585,13 +1592,20 @@ class ChatTelemetryLayer(Generic[OptionsCoT]):
return wrapped_stream
async def _get_response() -> ChatResponse:
agent_span = trace.get_current_span()
with _get_span(attributes=attributes, span_name_attribute=OtelAttr.REQUEST_MODEL) as span:
if OBSERVABILITY_SETTINGS.SENSITIVE_DATA_ENABLED and messages and span.is_recording():
system_instructions = _get_instructions_from_options(opts)
_capture_current_agent_system_instructions(
agent_span,
span,
system_instructions,
)
_capture_messages(
span=span,
provider_name=provider_name,
messages=messages,
system_instructions=opts.get("instructions"),
system_instructions=system_instructions,
)
start_time_stamp = perf_counter()
try:
@@ -1761,7 +1775,6 @@ class AgentTelemetryLayer:
inner_response_telemetry_captured_fields
)
inner_accumulated_usage_token = INNER_ACCUMULATED_USAGE.set({})
if stream:
span = _start_streaming_span(attributes, OtelAttr.AGENT_NAME)
@@ -1859,38 +1872,44 @@ class AgentTelemetryLayer:
async def _run() -> AgentResponse[Any]:
try:
with _get_span(attributes=attributes, span_name_attribute=OtelAttr.AGENT_NAME) as span:
if OBSERVABILITY_SETTINGS.SENSITIVE_DATA_ENABLED and messages and span.is_recording():
_capture_messages(
span=span,
provider_name=provider_name,
messages=messages,
system_instructions=_get_instructions_from_options(dict(merged_options)),
)
start_time_stamp = perf_counter()
try:
response: AgentResponse[Any] = await execute()
except Exception as exception:
capture_exception(span=span, exception=exception, timestamp=time_ns())
raise
duration = perf_counter() - start_time_stamp
if response:
response_attributes = _get_response_attributes(
attributes,
response,
capture_response_id=INNER_RESPONSE_ID_CAPTURED_FIELD
not in inner_response_telemetry_captured_fields,
capture_usage=INNER_USAGE_CAPTURED_FIELD not in inner_response_telemetry_captured_fields,
)
_apply_accumulated_usage(response_attributes, inner_response_telemetry_captured_fields)
_capture_response(span=span, attributes=response_attributes, duration=duration)
if OBSERVABILITY_SETTINGS.SENSITIVE_DATA_ENABLED and response.messages and span.is_recording():
if OBSERVABILITY_SETTINGS.SENSITIVE_DATA_ENABLED and messages and span.is_recording():
_capture_messages(
span=span,
provider_name=provider_name,
messages=response.messages,
output=True,
messages=messages,
system_instructions=_get_instructions_from_options(dict(merged_options)),
)
return response # type: ignore[return-value,no-any-return]
start_time_stamp = perf_counter()
response: AgentResponse[Any] = await execute()
duration = perf_counter() - start_time_stamp
if response:
response_attributes = _get_response_attributes(
attributes,
response,
capture_response_id=INNER_RESPONSE_ID_CAPTURED_FIELD
not in inner_response_telemetry_captured_fields,
capture_usage=(
INNER_USAGE_CAPTURED_FIELD not in inner_response_telemetry_captured_fields
),
)
_apply_accumulated_usage(response_attributes, inner_response_telemetry_captured_fields)
_capture_response(span=span, attributes=response_attributes, duration=duration)
if (
OBSERVABILITY_SETTINGS.SENSITIVE_DATA_ENABLED
and response.messages
and span.is_recording()
):
_capture_messages(
span=span,
provider_name=provider_name,
messages=response.messages,
output=True,
)
return response # type: ignore[return-value,no-any-return]
except Exception as exception:
capture_exception(span=span, exception=exception, timestamp=time_ns())
raise
finally:
INNER_RESPONSE_TELEMETRY_CAPTURED_FIELDS.reset(inner_response_telemetry_captured_fields_token)
INNER_ACCUMULATED_USAGE.reset(inner_accumulated_usage_token)
@@ -2263,6 +2282,83 @@ def capture_exception(span: trace.Span, exception: Exception, timestamp: int | N
span.set_status(status=trace.StatusCode.ERROR, description=repr(exception))
def _capture_system_instructions(span: trace.Span, system_instructions: str | list[str] | None) -> None:
"""Capture system instructions on a span."""
if not system_instructions:
return
otel_sys_instructions = [
{"type": "text", "content": instruction} for instruction in _normalize_instructions(system_instructions)
]
span.set_attribute(OtelAttr.SYSTEM_INSTRUCTIONS, json.dumps(otel_sys_instructions, ensure_ascii=False))
def _capture_current_agent_system_instructions(
agent_span: trace.Span,
chat_span: trace.Span,
system_instructions: str | list[str] | None,
) -> None:
"""Capture final chat instructions on the current agent span when the chat span belongs to it."""
if not system_instructions or not agent_span.is_recording():
return
agent_attributes_obj = getattr(agent_span, "attributes", None)
if not isinstance(agent_attributes_obj, Mapping):
return
agent_attributes = cast(Mapping[str, Any], agent_attributes_obj)
if agent_attributes.get(OtelAttr.OPERATION.value) != OtelAttr.AGENT_INVOKE_OPERATION:
return
if not _instructions_preserve_existing_agent_instructions(agent_attributes, system_instructions):
return
chat_parent = getattr(chat_span, "parent", None)
agent_context = agent_span.get_span_context()
if (
chat_parent is None
or chat_parent.span_id != agent_context.span_id
or chat_parent.trace_id != agent_context.trace_id
):
return
_capture_system_instructions(agent_span, system_instructions)
def _normalize_instructions(system_instructions: str | list[str]) -> list[str]:
"""Normalize system instructions to telemetry text items."""
return system_instructions if isinstance(system_instructions, list) else [system_instructions]
def _instructions_preserve_existing_agent_instructions(
agent_attributes: Mapping[str, Any],
system_instructions: str | list[str],
) -> bool:
"""Return True when chat instructions preserve the agent span's existing instructions."""
existing = agent_attributes.get(OtelAttr.SYSTEM_INSTRUCTIONS)
if not isinstance(existing, str):
return True
try:
existing_items_obj = json.loads(existing)
except json.JSONDecodeError:
return False
if not isinstance(existing_items_obj, list):
return False
existing_items = cast(list[object], existing_items_obj)
existing_contents: list[str] = []
for item in existing_items:
if not isinstance(item, Mapping):
continue
content = cast(Mapping[str, Any], item).get("content")
if isinstance(content, str):
existing_contents.append(content)
existing_text = "\n".join(existing_contents)
new_text = "\n".join(_normalize_instructions(system_instructions))
return new_text == existing_text or new_text.startswith(f"{existing_text}\n")
def _capture_messages(
span: trace.Span,
provider_name: str,
@@ -2294,11 +2390,7 @@ def _capture_messages(
span.set_attribute(
OtelAttr.OUTPUT_MESSAGES if output else OtelAttr.INPUT_MESSAGES, json.dumps(otel_messages, ensure_ascii=False)
)
if system_instructions:
if not isinstance(system_instructions, list):
system_instructions = [system_instructions]
otel_sys_instructions = [{"type": "text", "content": instruction} for instruction in system_instructions]
span.set_attribute(OtelAttr.SYSTEM_INSTRUCTIONS, json.dumps(otel_sys_instructions, ensure_ascii=False))
_capture_system_instructions(span, system_instructions)
def _to_otel_message(message: Message) -> dict[str, Any]:
@@ -17,6 +17,7 @@ from agent_framework import (
ChatResponse,
ChatResponseUpdate,
Content,
ContextProvider,
Message,
RawAgent,
ResponseStream,
@@ -3647,6 +3648,124 @@ async def test_agent_streaming_instructions_merged_from_default_and_options(
assert "Stream override." in system_instructions[0]["content"]
@pytest.mark.parametrize("enable_sensitive_data", [True], indirect=True)
@pytest.mark.parametrize("stream", [False, True])
async def test_agent_instructions_include_context_provider_extensions(
mock_chat_client,
span_exporter: InMemorySpanExporter,
enable_sensitive_data,
stream: bool,
) -> None:
"""Agent span instructions include instructions added by context providers."""
import json
class UserMemoryProvider(ContextProvider):
def __init__(self) -> None:
super().__init__(source_id="user-memory")
async def before_run(
self,
*,
agent: Any,
session: Any,
context: Any,
state: dict[str, Any],
) -> None:
context.extend_instructions(self.source_id, "The user's name is Alice.")
agent = Agent(
client=mock_chat_client(),
name="memory_agent",
instructions="You are a friendly assistant.",
context_providers=[UserMemoryProvider()],
)
span_exporter.clear()
if stream:
result_stream = agent.run("Hello", stream=True)
async for _ in result_stream:
pass
await result_stream.get_final_response()
else:
await agent.run("Hello")
spans = span_exporter.get_finished_spans()
agent_spans = [
span for span in spans if span.attributes.get(OtelAttr.OPERATION.value) == OtelAttr.AGENT_INVOKE_OPERATION
]
assert len(agent_spans) == 1
system_instructions = json.loads(agent_spans[0].attributes[OtelAttr.SYSTEM_INSTRUCTIONS])
contents = [item["content"] for item in system_instructions]
assert any("You are a friendly assistant." in content for content in contents)
assert any("The user's name is Alice." in content for content in contents)
@pytest.mark.parametrize("enable_sensitive_data", [True], indirect=True)
async def test_agent_instructions_not_overwritten_by_unrelated_nested_chat(
mock_chat_client,
span_exporter: InMemorySpanExporter,
enable_sensitive_data,
) -> None:
"""Unrelated nested chat calls must not overwrite agent span instructions."""
import json
class NestedChatProvider(ContextProvider):
def __init__(self, nested_client: BaseChatClient[Any]) -> None:
super().__init__(source_id="nested-chat")
self.nested_client = nested_client
async def before_run(
self,
*,
agent: Any,
session: Any,
context: Any,
state: dict[str, Any],
) -> None:
context.extend_instructions(self.source_id, "Context-provided instructions.")
async def after_run(
self,
*,
agent: Any,
session: Any,
context: Any,
state: dict[str, Any],
) -> None:
await self.nested_client.get_response(
messages=[Message(role="user", contents=["Nested request"])],
options={"model": "NestedModel", "instructions": "Unrelated nested instructions."},
client_kwargs={"session": session},
)
agent = Agent(
client=mock_chat_client(),
name="guarded_agent",
instructions="Base agent instructions.",
context_providers=[NestedChatProvider(mock_chat_client())],
)
span_exporter.clear()
await agent.run("Hello")
spans = span_exporter.get_finished_spans()
agent_spans = [
span for span in spans if span.attributes.get(OtelAttr.OPERATION.value) == OtelAttr.AGENT_INVOKE_OPERATION
]
assert len(agent_spans) == 1
chat_spans = [
span for span in spans if span.attributes.get(OtelAttr.OPERATION.value) == OtelAttr.CHAT_COMPLETION_OPERATION
]
assert len(chat_spans) == 2
system_instructions = json.loads(agent_spans[0].attributes[OtelAttr.SYSTEM_INSTRUCTIONS])
contents = [item["content"] for item in system_instructions]
assert any("Base agent instructions." in content for content in contents)
assert any("Context-provided instructions." in content for content in contents)
assert all("Unrelated nested instructions." not in content for content in contents)
@pytest.mark.parametrize("enable_sensitive_data", [True], indirect=True)
async def test_agent_no_instructions_in_default_or_options(
mock_chat_agent, span_exporter: InMemorySpanExporter, enable_sensitive_data