Coverage for src/infra/telemetry.py: 73%

1"""

2 Telemetry abstraction for agent tracing.

3

4 Provides a pluggable telemetry system with:

5- TelemetryProvider protocol for abstraction

6- TelemetrySpan protocol for span context managers

7- NullTelemetryProvider for testing and opt-out

8

9 For production with Braintrust, use BraintrustProvider from braintrust_integration.py:

10

11 from src.infra.clients.braintrust_integration import BraintrustProvider

12

13 Usage:

14 # For tests or opt-out:

15 provider = NullTelemetryProvider()

16

17 # Use via protocol:

18 if provider.is_enabled():

19 with provider.create_span("task-123", {"agent_id": "agent-1"}):

20 # Work happens here

21 pass

22 provider.flush()

23"""

24

25 from __future__ import annotations

26

27 from typing import TYPE_CHECKING, Any, Protocol, Self

28

29 if TYPE_CHECKING:

30 from types import TracebackType

31

32

33 class TelemetrySpan(Protocol):

34 """Protocol for a telemetry span context manager."""

35

36 def __enter__(self) -> Self:

37 """Enter the span context."""

38 ...

39

40 def __exit__(

41 self,

42 exc_type: type[BaseException] | None,

43 exc_val: BaseException | None,

44 exc_tb: TracebackType | None,

45 ) -> None:

46 """Exit the span context."""

47 ...

48

49 def log_input(self, prompt: str) -> None:

50 """Log the initial user prompt."""

51 ...

52

53 def log_message(self, message: object) -> None:

54 """Log a message from the SDK."""

55 ...

56

57 def set_success(self, success: bool) -> None:

58 """Mark the execution as successful or failed."""

59 ...

60

61 def set_error(self, error: str) -> None:

62 """Record an error message."""

63 ...

64

65

66 class TelemetryProvider(Protocol):

67 """Protocol for telemetry providers.

68

69 Telemetry providers abstract the underlying tracing system,

70 allowing tests to use a null implementation and production

71 code to use Braintrust or other backends.

72 """

73

74 def is_enabled(self) -> bool:

75 """Check if telemetry is active and configured."""

76 ...

77

78 def create_span(

79 self, name: str, metadata: dict[str, Any] | None = None

80 ) -> TelemetrySpan:

81 """Create a span context manager for tracing an operation.

82

83 Args:

84 name: Span name (typically issue_id for agent executions)

85 metadata: Optional metadata dict (e.g., agent_id, custom fields)

86

87 Returns:

88 A TelemetrySpan context manager for the operation

89 """

90 ...

91

92 def flush(self) -> None:

93 """Flush pending telemetry data."""

94 ...

95

96

97 class NullSpan:

98 """No-op span implementation for testing and opt-out.

99

100 All methods are no-ops that return immediately.

101 """

102

103 def __enter__(self) -> Self:

104 return self

105

106 def __exit__(

107 self,

108 exc_type: type[BaseException] | None,

109 exc_val: BaseException | None,

110 exc_tb: TracebackType | None,

111 ) -> None:

112 pass

113

114 def log_input(self, prompt: str) -> None:

115 pass

116

117 def log_message(self, message: object) -> None:

118 pass

119

120 def set_success(self, success: bool) -> None:

121 pass

122

123 def set_error(self, error: str) -> None:

124 pass

125

126

127 class NullTelemetryProvider:

128 """No-op telemetry provider for testing and opt-out.

129

130 This provider is completely stateless and has no side effects.

131 All operations return immediately without doing anything.

132

133 Usage:

134 provider = NullTelemetryProvider()

135 with provider.create_span("task-123"):

136 # Work happens here - nothing is traced

137 pass

138 """

139

140 def is_enabled(self) -> bool:

141 """Always returns False - null provider is never 'enabled'."""

142 return False

143

144 def create_span(

145 self, name: str, metadata: dict[str, Any] | None = None

146 ) -> NullSpan:

147 """Return a no-op span that ignores all operations."""

148 return NullSpan()

149

150 def flush(self) -> None:

151 """No-op flush."""

152 pass