Coverage for agentos/observability/tracing.py: 46%

1""" 

2OpenTelemetry distributed tracing for AgentOS. 

3 

4Minimal setup — wraps standard OpenTelemetry SDK to trace: 

5- HTTP requests (inbound via FastAPI middleware, outbound via httpx instrumentor) 

6- Agent pipeline phases (PRE_LLM, POST_LLM, PRE_TOOL, POST_TOOL, PRE_EXEC, POST_EXEC) 

7- Database queries (SQLAlchemy instrumentor) 

8 

9Usage: 

10 from agentos.observability.tracing import setup_tracing, get_tracer 

11 

12 setup_tracing(service_name="agentos", otlp_endpoint="http://localhost:4317") 

13 tracer = get_tracer(__name__) 

14 

15 with tracer.start_as_current_span("agent.run") as span: 

16 span.set_attribute("agent.id", "agent-1") 

17 # ... do work ... 

18 

19With env vars: 

20 AGENTOS_OTLP_ENDPOINT=http://jaeger:4317 

21 AGENTOS_TRACE_ENABLED=true 

22""" 

23 

24from __future__ import annotations 

25 

26import logging 

27import os 

28from contextlib import contextmanager 

29from functools import wraps 

30from typing import Any, Callable, Optional 

31 

32logger = logging.getLogger(__name__) 

33 

34_tracer_provider: Any = None 

35_TRACE_ENABLED: bool = os.environ.get("AGENTOS_TRACE_ENABLED", "").lower() == "true" 

36 

37try: 

38 from opentelemetry import trace 

39 from opentelemetry.sdk.trace import TracerProvider 

40 from opentelemetry.sdk.resources import Resource, SERVICE_NAME 

41 from opentelemetry.sdk.trace.export import BatchSpanProcessor 

42 from opentelemetry.exporter.otlp.proto.grpc.trace_exporter import OTLPSpanExporter 

43 from opentelemetry.trace import SpanKind, Status, StatusCode 

44 

45 _OTEL_AVAILABLE = True 

46except ImportError: 

47 _OTEL_AVAILABLE = False 

48 logger.debug("opentelemetry not installed — tracing disabled") 

49 

50 

51def setup_tracing( 

52 service_name: str = "agentos", 

53 otlp_endpoint: Optional[str] = None, 

54 sample_rate: float = 1.0, 

55) -> None: 

56 """Initialize OpenTelemetry tracing. 

57 

58 Args: 

59 service_name: Logical service name for spans. 

60 otlp_endpoint: OTLP gRPC collector endpoint (default: AGENTOS_OTLP_ENDPOINT env). 

61 sample_rate: Trace sampling rate (1.0 = all). 

62 """ 

63 global _tracer_provider, _TRACE_ENABLED 

64 

65 if not _OTEL_AVAILABLE: 

66 logger.warning("OpenTelemetry SDK not available — tracing disabled") 

67 return 

68 

69 endpoint = otlp_endpoint or os.environ.get("AGENTOS_OTLP_ENDPOINT", "") 

70 if not endpoint: 

71 logger.debug("No OTLP endpoint configured — tracing disabled") 

72 return 

73 

74 resource = Resource(attributes={ 

75 SERVICE_NAME: service_name, 

76 "deployment.environment": os.environ.get("AGENTOS_ENV", "production"), 

77 }) 

78 

79 exporter = OTLPSpanExporter(endpoint=endpoint, insecure=True) 

80 processor = BatchSpanProcessor(exporter) 

81 

82 _tracer_provider = TracerProvider(resource=resource) 

83 _tracer_provider.add_span_processor(processor) 

84 trace.set_tracer_provider(_tracer_provider) 

85 

86 _TRACE_ENABLED = True 

87 logger.info(f"Tracing enabled → {endpoint}") 

88 

89 

90def shutdown_tracing() -> None: 

91 """Flush and shutdown the tracer provider.""" 

92 if _tracer_provider: 

93 _tracer_provider.shutdown() 

94 

95 

96def get_tracer(name: str = "agentos") -> Any: 

97 """Get a tracer instance (falls back to no-op if OTel not configured).""" 

98 if _OTEL_AVAILABLE and _TRACE_ENABLED: 

99 return trace.get_tracer(name) 

100 # No-op fallback 

101 return _NoOpTracer() 

102 

103 

104class _NoOpTracer: 

105 """Drop-in replacement when tracing is disabled.""" 

106 

107 @contextmanager 

108 def start_as_current_span(self, name: str, **kwargs): 

109 yield _NoOpSpan() 

110 

111 def start_span(self, name: str, **kwargs): 

112 return _NoOpSpan() 

113 

114 

115class _NoOpSpan: 

116 def set_attribute(self, key: str, value: Any) -> None: pass 

117 def set_status(self, status: Any) -> None: pass 

118 def add_event(self, name: str, attributes: dict = None) -> None: pass 

119 def end(self) -> None: pass 

120 def __enter__(self): return self 

121 def __exit__(self, *args): pass 

122 

123 

124def trace_function(name: Optional[str] = None, attributes: dict = None): 

125 """Decorator to trace a function as a span.""" 

126 def decorator(fn: Callable): 

127 span_name = name or f"{fn.__module__}.{fn.__qualname__}" 

128 

129 @wraps(fn) 

130 def wrapper(*args, **kwargs): 

131 tracer = get_tracer(fn.__module__) 

132 with tracer.start_as_current_span(span_name) as span: 

133 if attributes: 

134 for k, v in attributes.items(): 

135 span.set_attribute(k, v) 

136 return fn(*args, **kwargs) 

137 return wrapper 

138 return decorator 

139 

140 

141def get_current_span() -> Any: 

142 """Get the current active span (no-op safe).""" 

143 if _OTEL_AVAILABLE and _TRACE_ENABLED: 

144 return trace.get_current_span() 

145 return _NoOpSpan() 

146 

147 

148__all__ = [ 

149 "setup_tracing", 

150 "shutdown_tracing", 

151 "get_tracer", 

152 "trace_function", 

153 "get_current_span", 

154]