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

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 collections.abc import Callable 

29from contextlib import contextmanager 

30from functools import wraps 

31from typing import Any 

32 

33logger = logging.getLogger(__name__) 

34 

35_tracer_provider: Any = None 

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

37 

38try: 

39 from opentelemetry import trace 

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

41 from opentelemetry.sdk.resources import SERVICE_NAME, Resource 

42 from opentelemetry.sdk.trace import TracerProvider 

43 from opentelemetry.sdk.trace.export import BatchSpanProcessor 

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: str | None = 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( 

75 attributes={ 

76 SERVICE_NAME: service_name, 

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

78 } 

79 ) 

80 

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

82 processor = BatchSpanProcessor(exporter) 

83 

84 _tracer_provider = TracerProvider(resource=resource) 

85 _tracer_provider.add_span_processor(processor) 

86 trace.set_tracer_provider(_tracer_provider) 

87 

88 _TRACE_ENABLED = True 

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

90 

91 

92def shutdown_tracing() -> None: 

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

94 if _tracer_provider: 

95 _tracer_provider.shutdown() 

96 

97 

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

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

100 if _OTEL_AVAILABLE and _TRACE_ENABLED: 

101 return trace.get_tracer(name) 

102 # No-op fallback 

103 return _NoOpTracer() 

104 

105 

106class _NoOpTracer: 

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

108 

109 @contextmanager 

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

111 yield _NoOpSpan() 

112 

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

114 return _NoOpSpan() 

115 

116 

117class _NoOpSpan: 

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

119 pass 

120 

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

122 pass 

123 

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

125 pass 

126 

127 def end(self) -> None: 

128 pass 

129 

130 def __enter__(self): 

131 return self 

132 

133 def __exit__(self, *args): 

134 pass 

135 

136 

137def trace_function(name: str | None = None, attributes: dict = None): 

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

139 

140 def decorator(fn: Callable): 

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

142 

143 @wraps(fn) 

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

145 tracer = get_tracer(fn.__module__) 

146 with tracer.start_as_current_span(span_name) as span: 

147 if attributes: 

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

149 span.set_attribute(k, v) 

150 return fn(*args, **kwargs) 

151 

152 return wrapper 

153 

154 return decorator 

155 

156 

157def get_current_span() -> Any: 

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

159 if _OTEL_AVAILABLE and _TRACE_ENABLED: 

160 return trace.get_current_span() 

161 return _NoOpSpan() 

162 

163 

164__all__ = [ 

165 "setup_tracing", 

166 "shutdown_tracing", 

167 "get_tracer", 

168 "trace_function", 

169 "get_current_span", 

170]