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:
committed by
GitHub
parent
7bf2d2a6d0
commit
4ff952e100
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user