8000 chore: allow custom tracer provider in Agent (#207) · StrandForge/sdk-python@68740c5 · GitHub
[go: up one dir, main page]
Skip to content

Commit 68740c5

Browse files
poshinchenposhinchen
authored
chore: allow custom tracer provider in Agent (strands-agents#207)
1 parent 7c5f7a7 commit 68740c5

File tree

3 files changed

+113
-44
lines changed

3 files changed

+113
-44
lines changed

‎src/strands/agent/conversation_manager/sliding_window_conversation_manager.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -140,7 +140,7 @@ def reduce_context(self, agent: "Agent", e: Optional[Exception] = None) -> None:
140140
if results_truncated:
141141
logger.debug("message_index=<%s> | tool results truncated", last_message_idx_with_tool_results)
142142
return
143-
143+
144144
# Try to trim index id when tool result cannot be truncated anymore
145145
# If the number of messages is less than the window_size, then we default to 2, otherwise, trim to window size
146146
trim_index = 2 if len(messages) <= self.window_size else len(messages) - self.window_size

‎src/strands/telemetry/tracer.py

Lines changed: 52 additions & 29 deletions
Original file line numberDiff line numberDiff line change
@@ -11,12 +11,16 @@
1111
from importlib.metadata import version
1212
from typing import Any, Dict, Mapping, Optional
1313

14-
from opentelemetry import trace
14+
import opentelemetry.trace as trace_api
15+
from opentelemetry import propagate
16+
from opentelemetry.baggage.propagation import W3CBaggagePropagator
1517
from opentelemetry.exporter.otlp.proto.http.trace_exporter import OTLPSpanExporter
18+
from opentelemetry.propagators.composite import CompositePropagator
1619
from opentelemetry.sdk.resources import Resource
17-
from opentelemetry.sdk.trace import TracerProvider
20+
from opentelemetry.sdk.trace import TracerProvider as SDKTracerProvider
1821
from opentelemetry.sdk.trace.export import BatchSpanProcessor, ConsoleSpanExporter, SimpleSpanProcessor
19-
from opentelemetry.trace import StatusCode
22+
from opentelemetry.trace import Span, StatusCode
23+
from opentelemetry.trace.propagation.tracecontext import TraceContextTextMapPropagator
2024

2125
from ..agent.agent_result import AgentResult
2226
from ..types.content import Message, Messages
@@ -133,16 +137,30 @@ def __init__(
133137

134138
self.service_name = service_name
135139
self.otlp_headers = otlp_headers or {}
136-
self.tracer_provider: Optional[TracerProvider] = None
137-
self.tracer: Optional[trace.Tracer] = None
138-
140+
self.tracer_provider: Optional[trace_api.TracerProvider] = None
141+
self.tracer: Optional[trace_api.Tracer] = None
142+
143+
propagate.set_global_textmap(
144+
CompositePropagator(
145+
[
146+
W3CBaggagePropagator(),
147+
TraceContextTextMapPropagator(),
148+
]
149+
)
150+
)
139151
if self.otlp_endpoint or self.enable_console_export:
152+
# Create our own tracer provider
140153
self._initialize_tracer()
141154

142155
def _initialize_tracer(self) -> None:
143156
"""Initialize the OpenTelemetry tracer."""
144157
logger.info("initializing tracer")
145158

159+
if self._is_initialized():
160+
self.tracer_provider = trace_api.get_tracer_provider()
161+
self.tracer = self.tracer_provider.get_tracer(self.service_name)
162+
return
163+
146164
# Create resource with service information
147165
resource = Resource.create(
148166
{
@@ -154,7 +172,7 @@ def _initialize_tracer(self) -> None:
154172
)
155173

156174
# Create tracer provider
157-
self.tracer_provider = TracerProvider(resource=resource)
175+
self.tracer_provider = SDKTracerProvider(resource=resource)
158176

159177
# Add console exporter if enabled
160178
if self.enable_console_export and self.tracer_provider:
@@ -190,15 +208,19 @@ def _initialize_tracer(self) -> None:
190208
logger.exception("error=<%s> | Failed to configure OTLP exporter", e)
191209

192210
# Set as global tracer provider
193-
trace.set_tracer_provider(self.tracer_provider)
194-
self.tracer = trace.get_tracer(self.service_name)
211+
trace_api.set_tracer_provider(self.tracer_provider)
212+
self.tracer = trace_api.get_tracer(self.service_name)
213+
214+
def _is_initialized(self) -> bool:
215+
tracer_provider = trace_api.get_tracer_provider()
216+
return isinstance(tracer_provider, SDKTracerProvider)
195217

196218
def _start_span(
197219
self,
198220
span_name: str,
199-
parent_span: Optional[trace.Span] = None,
221+
parent_span: Optional[Span] = None,
200222
attributes: Optional[Dict[str, AttributeValue]] = None,
201-
) -> Optional[trace.Span]:
223+
) -> Optional[Span]:
202224
"""Generic helper method to start a span with common attributes.
203225
204226
Args:
@@ -212,7 +234,7 @@ def _start_span(
212234
if self.tracer is None:
213235
return None
214236

215-
context = trace.set_span_in_context(parent_span) if parent_span else None
237+
context = trace_api.set_span_in_context(parent_span) if parent_span else None
216238
span = self.tracer.start_span(name=span_name, context=context)
217239

218240
# Set start time as a common attribute
@@ -224,7 +246,7 @@ def _start_span(
224246

225247
return span
226248

227-
def _set_attributes(self, span: trace.Span, attributes: Dict[str, AttributeValue]) -> None:
249+
def _set_attributes(self, span: Span, attributes: Dict[str, AttributeValue]) -> None:
228250
"""Set attributes on a span, handling different value types appropriately.
229251
230252
Args:
@@ -239,7 +261,7 @@ def _set_attributes(self, span: trace.Span, attributes: Dict[str, AttributeValue
239261

240262
def _end_span(
241263
self,
242-
span: trace.Span,
264+
span: Span,
243265
attributes: Optional[Dict[str, AttributeValue]] = None,
244266
error: Optional[Exception] = None,
245267
) -> None:
@@ -272,13 +294,13 @@ def _end_span(
272294
finally:
273295
span.end()
274296
# Force flush to ensure spans are exported
275-
if self.tracer_provider:
297+
if self.tracer_provider and hasattr(self.tracer_provider, 'force_flush'):
276298
try:
277299
self.tracer_provider.force_flush()
278300
except Exception as e:
279301
logger.warning("error=<%s> | failed to force flush tracer provider", e)
280302

281-
def end_span_with_error(self, span: trace.Span, error_message: str, exception: Optional[Exception] = None) -> None:
303+
def end_span_with_error(self, span: Span, error_message: str, exception: Optional[Exception] = None) -> None:
282304
"""End a span with error status.
283305
284306
Args:
@@ -294,12 +316,12 @@ def end_span_with_error(self, span: trace.Span, error_message: str, exception: O
294316

295317
def start_model_invoke_span(
296318
self,
297-
parent_span: Optional[trace.Span] = None,
319+
parent_span: Optional[Span] = None,
298320
agent_name: str = "Strands Agent",
299321
messages: Optional[Messages] = None,
300322
model_id: Optional[str] = None,
301323
**kwargs: Any,
302-
) -> Optional[trace.Span]:
324+
) -> Optional[Span]:
303325
"""Start a new span for a model invocation.
304326
305327
Args:
@@ -328,7 +350,7 @@ def start_model_invoke_span(
328350
return self._start_span("Model invoke", parent_span, attributes)
329351

330352
def end_model_invoke_span(
331-
self, span: trace.Span, message: Message, usage: Usage, error: Optional[Exception] = None
353+
self, span: Span, message: Message, usage: Usage, error: Optional[Exception] = None
332354
) -> None:
333355
"""End a model invocation span with results and metrics.
334356
@@ -347,9 +369,7 @@ def end_model_invoke_span(
347369

348370
self._end_span(span, attributes, error)
349371

350-
def start_tool_call_span(
351-
self, tool: ToolUse, parent_span: Optional[trace.Span] = None, **kwargs: Any
352-
) -> Optional[trace.Span]:
372+
def start_tool_call_span(self, tool: ToolUse, parent_span: Optional[Span] = None, **kwargs: Any) -> Optional[Span]:
353373
"""Start a new span for a tool call.
354374
355375
Args:
@@ -374,7 +394,7 @@ def start_tool_call_span(
374394
return self._start_span(span_name, parent_span, attributes)
375395

376396
def end_tool_call_span(
377-
self, span: trace.Span, tool_result: Optional[ToolResult], error: Optional[Exception] = None
397+
self, span: Span, tool_result: Optional[ToolResult], error: Optional[Exception] = None
378398
) -> None:
379399
"""End a tool call span with results.
380400
@@ -402,10 +422,10 @@ def end_tool_call_span(
402422
def start_event_loop_cycle_span(
403423
self,
404424
event_loop_kwargs: Any,
405-
parent_span: Optional[trace.Span] = None,
425+
parent_span: Optional[Span] = None,
406426
messages: Optional[Messages] = None,
407427
**kwargs: Any,
408-
) -> Optional[trace.Span]:
428+
) -> Optional[Span]:
409429
"""Start a new span for an event loop cycle.
410430
411431
Args:
@@ -436,7 +456,7 @@ def start_event_loop_cycle_span(
436456

437457
def end_event_loop_cycle_span(
438458
self,
439-
span: trace.Span,
459+
span: Span,
440460
message: Message,
441461
tool_result_message: Optional[Message] = None,
442462
error: Optional[Exception] = None,
@@ -466,7 +486,7 @@ def start_agent_span(
466486
tools: Optional[list] = None,
467487
custom_trace_attributes: Optional[Mapping[str, AttributeValue]] = None,
468488
**kwargs: Any,
469-
) -> Optional[trace.Span]:
489+
) -> Optional[Span]:
470490
"""Start a new span for an agent invocation.
471491
472492
Args:
@@ -506,7 +526,7 @@ def start_agent_span(
506526

507527
def end_agent_span(
508528
self,
509-
span: trace.Span,
529+
span: Span,
510530
response: Optional[AgentResult] = None,
511531
error: Optional[Exception] = None,
512532
) -> None:
@@ -557,13 +577,16 @@ def get_tracer(
557577
otlp_endpoint: OTLP endpoint URL for sending traces.
558578
otlp_headers: Headers to include with OTLP requests.
559579
enable_console_export: Whether to also export traces to console.
580+
tracer_provider: Optional existing TracerProvider to use instead of creating a new one.
560581
561582
Returns:
562583
The global tracer instance.
563584
"""
564585
global _tracer_instance
565586

566-
if _tracer_instance is None or (otlp_endpoint and _tracer_instance.otlp_endpoint != otlp_endpoint): # type: ignore[unreachable]
587+
if (
588+
_tracer_instance is None or (otlp_endpoint and _tracer_instance.otlp_endpoint != otlp_endpoint) # type: ignore[unreachable]
589+
):
567590
_tracer_instance = Tracer(
568591
service_name=service_name,
569592
otlp_endpoint=otlp_endpoint,

0 commit comments

Comments
 (0)
0