alma-memory 0.5.1__py3-none-any.whl → 0.7.0__py3-none-any.whl
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- alma/__init__.py +296 -226
- alma/compression/__init__.py +33 -0
- alma/compression/pipeline.py +980 -0
- alma/confidence/__init__.py +47 -47
- alma/confidence/engine.py +540 -540
- alma/confidence/types.py +351 -351
- alma/config/loader.py +157 -157
- alma/consolidation/__init__.py +23 -23
- alma/consolidation/engine.py +678 -678
- alma/consolidation/prompts.py +84 -84
- alma/core.py +1189 -430
- alma/domains/__init__.py +30 -30
- alma/domains/factory.py +359 -359
- alma/domains/schemas.py +448 -448
- alma/domains/types.py +272 -272
- alma/events/__init__.py +75 -75
- alma/events/emitter.py +285 -284
- alma/events/storage_mixin.py +246 -246
- alma/events/types.py +126 -126
- alma/events/webhook.py +425 -425
- alma/exceptions.py +49 -49
- alma/extraction/__init__.py +31 -31
- alma/extraction/auto_learner.py +265 -265
- alma/extraction/extractor.py +420 -420
- alma/graph/__init__.py +106 -106
- alma/graph/backends/__init__.py +32 -32
- alma/graph/backends/kuzu.py +624 -624
- alma/graph/backends/memgraph.py +432 -432
- alma/graph/backends/memory.py +236 -236
- alma/graph/backends/neo4j.py +417 -417
- alma/graph/base.py +159 -159
- alma/graph/extraction.py +198 -198
- alma/graph/store.py +860 -860
- alma/harness/__init__.py +35 -35
- alma/harness/base.py +386 -386
- alma/harness/domains.py +705 -705
- alma/initializer/__init__.py +37 -37
- alma/initializer/initializer.py +418 -418
- alma/initializer/types.py +250 -250
- alma/integration/__init__.py +62 -62
- alma/integration/claude_agents.py +444 -444
- alma/integration/helena.py +423 -423
- alma/integration/victor.py +471 -471
- alma/learning/__init__.py +101 -86
- alma/learning/decay.py +878 -0
- alma/learning/forgetting.py +1446 -1446
- alma/learning/heuristic_extractor.py +390 -390
- alma/learning/protocols.py +374 -374
- alma/learning/validation.py +346 -346
- alma/mcp/__init__.py +123 -45
- alma/mcp/__main__.py +156 -156
- alma/mcp/resources.py +122 -122
- alma/mcp/server.py +955 -591
- alma/mcp/tools.py +3254 -509
- alma/observability/__init__.py +91 -84
- alma/observability/config.py +302 -302
- alma/observability/guidelines.py +170 -0
- alma/observability/logging.py +424 -424
- alma/observability/metrics.py +583 -583
- alma/observability/tracing.py +440 -440
- alma/progress/__init__.py +21 -21
- alma/progress/tracker.py +607 -607
- alma/progress/types.py +250 -250
- alma/retrieval/__init__.py +134 -53
- alma/retrieval/budget.py +525 -0
- alma/retrieval/cache.py +1304 -1061
- alma/retrieval/embeddings.py +202 -202
- alma/retrieval/engine.py +850 -427
- alma/retrieval/modes.py +365 -0
- alma/retrieval/progressive.py +560 -0
- alma/retrieval/scoring.py +344 -344
- alma/retrieval/trust_scoring.py +637 -0
- alma/retrieval/verification.py +797 -0
- alma/session/__init__.py +19 -19
- alma/session/manager.py +442 -399
- alma/session/types.py +288 -288
- alma/storage/__init__.py +101 -90
- alma/storage/archive.py +233 -0
- alma/storage/azure_cosmos.py +1259 -1259
- alma/storage/base.py +1083 -583
- alma/storage/chroma.py +1443 -1443
- alma/storage/constants.py +103 -103
- alma/storage/file_based.py +614 -614
- alma/storage/migrations/__init__.py +21 -21
- alma/storage/migrations/base.py +321 -321
- alma/storage/migrations/runner.py +323 -323
- alma/storage/migrations/version_stores.py +337 -337
- alma/storage/migrations/versions/__init__.py +11 -11
- alma/storage/migrations/versions/v1_0_0.py +373 -373
- alma/storage/migrations/versions/v1_1_0_workflow_context.py +551 -0
- alma/storage/pinecone.py +1080 -1080
- alma/storage/postgresql.py +1948 -1559
- alma/storage/qdrant.py +1306 -1306
- alma/storage/sqlite_local.py +3041 -1457
- alma/testing/__init__.py +46 -46
- alma/testing/factories.py +301 -301
- alma/testing/mocks.py +389 -389
- alma/types.py +292 -264
- alma/utils/__init__.py +19 -0
- alma/utils/tokenizer.py +521 -0
- alma/workflow/__init__.py +83 -0
- alma/workflow/artifacts.py +170 -0
- alma/workflow/checkpoint.py +311 -0
- alma/workflow/context.py +228 -0
- alma/workflow/outcomes.py +189 -0
- alma/workflow/reducers.py +393 -0
- {alma_memory-0.5.1.dist-info → alma_memory-0.7.0.dist-info}/METADATA +210 -72
- alma_memory-0.7.0.dist-info/RECORD +112 -0
- alma_memory-0.5.1.dist-info/RECORD +0 -93
- {alma_memory-0.5.1.dist-info → alma_memory-0.7.0.dist-info}/WHEEL +0 -0
- {alma_memory-0.5.1.dist-info → alma_memory-0.7.0.dist-info}/top_level.txt +0 -0
alma/observability/tracing.py
CHANGED
|
@@ -1,440 +1,440 @@
|
|
|
1
|
-
"""
|
|
2
|
-
ALMA Distributed Tracing.
|
|
3
|
-
|
|
4
|
-
Provides distributed tracing using OpenTelemetry with fallback
|
|
5
|
-
to logging when OTel is not available.
|
|
6
|
-
"""
|
|
7
|
-
|
|
8
|
-
import functools
|
|
9
|
-
import logging
|
|
10
|
-
from contextlib import contextmanager
|
|
11
|
-
from enum import Enum
|
|
12
|
-
from typing import Any, Callable, Dict, Optional, TypeVar, Union
|
|
13
|
-
|
|
14
|
-
# Try to import OpenTelemetry
|
|
15
|
-
_otel_available = False
|
|
16
|
-
_NoOpSpan = None
|
|
17
|
-
_NoOpTracer = None
|
|
18
|
-
|
|
19
|
-
try:
|
|
20
|
-
from opentelemetry import trace
|
|
21
|
-
from opentelemetry.trace import SpanKind as OTelSpanKind
|
|
22
|
-
from opentelemetry.trace import Status, StatusCode
|
|
23
|
-
|
|
24
|
-
_otel_available = True
|
|
25
|
-
except ImportError:
|
|
26
|
-
pass
|
|
27
|
-
|
|
28
|
-
|
|
29
|
-
logger = logging.getLogger(__name__)
|
|
30
|
-
|
|
31
|
-
# Type variable for decorated functions
|
|
32
|
-
F = TypeVar("F", bound=Callable[..., Any])
|
|
33
|
-
|
|
34
|
-
|
|
35
|
-
class SpanKind(Enum):
|
|
36
|
-
"""Span kind enum (mirrors OpenTelemetry SpanKind)."""
|
|
37
|
-
|
|
38
|
-
INTERNAL = "internal"
|
|
39
|
-
SERVER = "server"
|
|
40
|
-
CLIENT = "client"
|
|
41
|
-
PRODUCER = "producer"
|
|
42
|
-
CONSUMER = "consumer"
|
|
43
|
-
|
|
44
|
-
|
|
45
|
-
class NoOpSpan:
|
|
46
|
-
"""No-op span implementation when OpenTelemetry is not available."""
|
|
47
|
-
|
|
48
|
-
def __init__(self, name: str, attributes: Optional[Dict[str, Any]] = None):
|
|
49
|
-
self.name = name
|
|
50
|
-
self.attributes = attributes or {}
|
|
51
|
-
self._logger = logging.getLogger(f"alma.trace.{name}")
|
|
52
|
-
|
|
53
|
-
def set_attribute(self, key: str, value: Any):
|
|
54
|
-
"""Set a span attribute."""
|
|
55
|
-
self.attributes[key] = value
|
|
56
|
-
|
|
57
|
-
def set_attributes(self, attributes: Dict[str, Any]):
|
|
58
|
-
"""Set multiple span attributes."""
|
|
59
|
-
self.attributes.update(attributes)
|
|
60
|
-
|
|
61
|
-
def add_event(self, name: str, attributes: Optional[Dict[str, Any]] = None):
|
|
62
|
-
"""Add an event to the span."""
|
|
63
|
-
self._logger.debug(f"Event: {name}", extra={"event_attributes": attributes})
|
|
64
|
-
|
|
65
|
-
def set_status(self, status: Any, description: Optional[str] = None):
|
|
66
|
-
"""Set span status."""
|
|
67
|
-
pass
|
|
68
|
-
|
|
69
|
-
def record_exception(
|
|
70
|
-
self, exception: BaseException, attributes: Optional[Dict[str, Any]] = None
|
|
71
|
-
):
|
|
72
|
-
"""Record an exception."""
|
|
73
|
-
self._logger.error(f"Exception in span {self.name}: {exception}", exc_info=True)
|
|
74
|
-
|
|
75
|
-
def end(self, end_time: Optional[int] = None):
|
|
76
|
-
"""End the span."""
|
|
77
|
-
pass
|
|
78
|
-
|
|
79
|
-
def __enter__(self):
|
|
80
|
-
return self
|
|
81
|
-
|
|
82
|
-
def __exit__(self, exc_type, exc_val, exc_tb):
|
|
83
|
-
if exc_type is not None:
|
|
84
|
-
self.record_exception(exc_val)
|
|
85
|
-
return False
|
|
86
|
-
|
|
87
|
-
|
|
88
|
-
class NoOpTracer:
|
|
89
|
-
"""No-op tracer implementation when OpenTelemetry is not available."""
|
|
90
|
-
|
|
91
|
-
def __init__(self, name: str):
|
|
92
|
-
self.name = name
|
|
93
|
-
|
|
94
|
-
def start_span(
|
|
95
|
-
self,
|
|
96
|
-
name: str,
|
|
97
|
-
context: Optional[Any] = None,
|
|
98
|
-
kind: SpanKind = SpanKind.INTERNAL,
|
|
99
|
-
attributes: Optional[Dict[str, Any]] = None,
|
|
100
|
-
start_time: Optional[int] = None,
|
|
101
|
-
) -> NoOpSpan:
|
|
102
|
-
"""Start a new span."""
|
|
103
|
-
return NoOpSpan(name, attributes)
|
|
104
|
-
|
|
105
|
-
@contextmanager
|
|
106
|
-
def start_as_current_span(
|
|
107
|
-
self,
|
|
108
|
-
name: str,
|
|
109
|
-
context: Optional[Any] = None,
|
|
110
|
-
kind: SpanKind = SpanKind.INTERNAL,
|
|
111
|
-
attributes: Optional[Dict[str, Any]] = None,
|
|
112
|
-
start_time: Optional[int] = None,
|
|
113
|
-
):
|
|
114
|
-
"""Start a span as the current span."""
|
|
115
|
-
span = NoOpSpan(name, attributes)
|
|
116
|
-
try:
|
|
117
|
-
yield span
|
|
118
|
-
except Exception as e:
|
|
119
|
-
span.record_exception(e)
|
|
120
|
-
raise
|
|
121
|
-
finally:
|
|
122
|
-
span.end()
|
|
123
|
-
|
|
124
|
-
|
|
125
|
-
class TracingContext:
|
|
126
|
-
"""
|
|
127
|
-
Context for managing trace propagation and span creation.
|
|
128
|
-
|
|
129
|
-
Provides a unified interface for tracing regardless of
|
|
130
|
-
whether OpenTelemetry is available.
|
|
131
|
-
"""
|
|
132
|
-
|
|
133
|
-
def __init__(self, tracer_name: str = "alma"):
|
|
134
|
-
"""
|
|
135
|
-
Initialize tracing context.
|
|
136
|
-
|
|
137
|
-
Args:
|
|
138
|
-
tracer_name: Name for the tracer
|
|
139
|
-
"""
|
|
140
|
-
self.tracer_name = tracer_name
|
|
141
|
-
self._tracer = None
|
|
142
|
-
|
|
143
|
-
@property
|
|
144
|
-
def tracer(self):
|
|
145
|
-
"""Get the tracer (lazy initialization)."""
|
|
146
|
-
if self._tracer is None:
|
|
147
|
-
self._tracer = get_tracer(self.tracer_name)
|
|
148
|
-
return self._tracer
|
|
149
|
-
|
|
150
|
-
@contextmanager
|
|
151
|
-
def span(
|
|
152
|
-
self,
|
|
153
|
-
name: str,
|
|
154
|
-
kind: SpanKind = SpanKind.INTERNAL,
|
|
155
|
-
attributes: Optional[Dict[str, Any]] = None,
|
|
156
|
-
):
|
|
157
|
-
"""
|
|
158
|
-
Create a span context manager.
|
|
159
|
-
|
|
160
|
-
Args:
|
|
161
|
-
name: Span name
|
|
162
|
-
kind: Span kind
|
|
163
|
-
attributes: Initial span attributes
|
|
164
|
-
|
|
165
|
-
Yields:
|
|
166
|
-
The created span
|
|
167
|
-
"""
|
|
168
|
-
if _otel_available:
|
|
169
|
-
otel_kind = _map_span_kind(kind)
|
|
170
|
-
with self.tracer.start_as_current_span(
|
|
171
|
-
name,
|
|
172
|
-
kind=otel_kind,
|
|
173
|
-
attributes=attributes,
|
|
174
|
-
) as span:
|
|
175
|
-
yield span
|
|
176
|
-
else:
|
|
177
|
-
with self.tracer.start_as_current_span(
|
|
178
|
-
name,
|
|
179
|
-
kind=kind,
|
|
180
|
-
attributes=attributes,
|
|
181
|
-
) as span:
|
|
182
|
-
yield span
|
|
183
|
-
|
|
184
|
-
def create_span(
|
|
185
|
-
self,
|
|
186
|
-
name: str,
|
|
187
|
-
kind: SpanKind = SpanKind.INTERNAL,
|
|
188
|
-
attributes: Optional[Dict[str, Any]] = None,
|
|
189
|
-
):
|
|
190
|
-
"""
|
|
191
|
-
Create a span (not automatically set as current).
|
|
192
|
-
|
|
193
|
-
Args:
|
|
194
|
-
name: Span name
|
|
195
|
-
kind: Span kind
|
|
196
|
-
attributes: Initial span attributes
|
|
197
|
-
|
|
198
|
-
Returns:
|
|
199
|
-
The created span
|
|
200
|
-
"""
|
|
201
|
-
if _otel_available:
|
|
202
|
-
otel_kind = _map_span_kind(kind)
|
|
203
|
-
return self.tracer.start_span(
|
|
204
|
-
name,
|
|
205
|
-
kind=otel_kind,
|
|
206
|
-
attributes=attributes,
|
|
207
|
-
)
|
|
208
|
-
else:
|
|
209
|
-
return self.tracer.start_span(
|
|
210
|
-
name,
|
|
211
|
-
kind=kind,
|
|
212
|
-
attributes=attributes,
|
|
213
|
-
)
|
|
214
|
-
|
|
215
|
-
|
|
216
|
-
def _map_span_kind(kind: SpanKind):
|
|
217
|
-
"""Map our SpanKind to OpenTelemetry SpanKind."""
|
|
218
|
-
if not _otel_available:
|
|
219
|
-
return kind
|
|
220
|
-
|
|
221
|
-
mapping = {
|
|
222
|
-
SpanKind.INTERNAL: OTelSpanKind.INTERNAL,
|
|
223
|
-
SpanKind.SERVER: OTelSpanKind.SERVER,
|
|
224
|
-
SpanKind.CLIENT: OTelSpanKind.CLIENT,
|
|
225
|
-
SpanKind.PRODUCER: OTelSpanKind.PRODUCER,
|
|
226
|
-
SpanKind.CONSUMER: OTelSpanKind.CONSUMER,
|
|
227
|
-
}
|
|
228
|
-
return mapping.get(kind, OTelSpanKind.INTERNAL)
|
|
229
|
-
|
|
230
|
-
|
|
231
|
-
def get_tracer(name: str = "alma") -> Union["NoOpTracer", Any]:
|
|
232
|
-
"""
|
|
233
|
-
Get a tracer for the given name.
|
|
234
|
-
|
|
235
|
-
Uses OpenTelemetry tracer if available, otherwise returns
|
|
236
|
-
a no-op tracer that logs span information.
|
|
237
|
-
|
|
238
|
-
Args:
|
|
239
|
-
name: Tracer name (typically module name)
|
|
240
|
-
|
|
241
|
-
Returns:
|
|
242
|
-
Tracer instance
|
|
243
|
-
"""
|
|
244
|
-
if _otel_available:
|
|
245
|
-
return trace.get_tracer(name)
|
|
246
|
-
return NoOpTracer(name)
|
|
247
|
-
|
|
248
|
-
|
|
249
|
-
def get_current_span():
|
|
250
|
-
"""
|
|
251
|
-
Get the current span.
|
|
252
|
-
|
|
253
|
-
Returns:
|
|
254
|
-
Current span or NoOpSpan if no span is active
|
|
255
|
-
"""
|
|
256
|
-
if _otel_available:
|
|
257
|
-
return trace.get_current_span()
|
|
258
|
-
return NoOpSpan("current")
|
|
259
|
-
|
|
260
|
-
|
|
261
|
-
def trace_method(
|
|
262
|
-
name: Optional[str] = None,
|
|
263
|
-
kind: SpanKind = SpanKind.INTERNAL,
|
|
264
|
-
record_args: bool = True,
|
|
265
|
-
record_result: bool = False,
|
|
266
|
-
) -> Callable[[F], F]:
|
|
267
|
-
"""
|
|
268
|
-
Decorator to trace a synchronous method.
|
|
269
|
-
|
|
270
|
-
Args:
|
|
271
|
-
name: Span name (defaults to function name)
|
|
272
|
-
kind: Span kind
|
|
273
|
-
record_args: Whether to record function arguments as attributes
|
|
274
|
-
record_result: Whether to record the return value
|
|
275
|
-
|
|
276
|
-
Usage:
|
|
277
|
-
@trace_method(name="my_operation")
|
|
278
|
-
def my_function(arg1, arg2):
|
|
279
|
-
return result
|
|
280
|
-
"""
|
|
281
|
-
|
|
282
|
-
def decorator(func: F) -> F:
|
|
283
|
-
span_name = name or func.__qualname__
|
|
284
|
-
|
|
285
|
-
@functools.wraps(func)
|
|
286
|
-
def wrapper(*args, **kwargs):
|
|
287
|
-
tracer = get_tracer(func.__module__)
|
|
288
|
-
|
|
289
|
-
attributes: Dict[str, Any] = {
|
|
290
|
-
"code.function": func.__name__,
|
|
291
|
-
"code.namespace": func.__module__,
|
|
292
|
-
}
|
|
293
|
-
|
|
294
|
-
if record_args:
|
|
295
|
-
# Record positional args (skip 'self' for methods)
|
|
296
|
-
arg_names = func.__code__.co_varnames[: func.__code__.co_argcount]
|
|
297
|
-
start_idx = 1 if arg_names and arg_names[0] in ("self", "cls") else 0
|
|
298
|
-
for i, arg in enumerate(args[start_idx:], start=start_idx):
|
|
299
|
-
if i < len(arg_names):
|
|
300
|
-
arg_val = _safe_attribute_value(arg)
|
|
301
|
-
if arg_val is not None:
|
|
302
|
-
attributes[f"arg.{arg_names[i]}"] = arg_val
|
|
303
|
-
|
|
304
|
-
# Record keyword args
|
|
305
|
-
for key, value in kwargs.items():
|
|
306
|
-
arg_val = _safe_attribute_value(value)
|
|
307
|
-
if arg_val is not None:
|
|
308
|
-
attributes[f"arg.{key}"] = arg_val
|
|
309
|
-
|
|
310
|
-
if _otel_available:
|
|
311
|
-
otel_kind = _map_span_kind(kind)
|
|
312
|
-
with tracer.start_as_current_span(
|
|
313
|
-
span_name,
|
|
314
|
-
kind=otel_kind,
|
|
315
|
-
attributes=attributes,
|
|
316
|
-
) as span:
|
|
317
|
-
try:
|
|
318
|
-
result = func(*args, **kwargs)
|
|
319
|
-
if record_result:
|
|
320
|
-
result_val = _safe_attribute_value(result)
|
|
321
|
-
if result_val is not None:
|
|
322
|
-
span.set_attribute("result", result_val)
|
|
323
|
-
return result
|
|
324
|
-
except Exception as e:
|
|
325
|
-
span.record_exception(e)
|
|
326
|
-
span.set_status(Status(StatusCode.ERROR, str(e)))
|
|
327
|
-
raise
|
|
328
|
-
else:
|
|
329
|
-
with tracer.start_as_current_span(
|
|
330
|
-
span_name,
|
|
331
|
-
kind=kind,
|
|
332
|
-
attributes=attributes,
|
|
333
|
-
) as span:
|
|
334
|
-
return func(*args, **kwargs)
|
|
335
|
-
|
|
336
|
-
return wrapper # type: ignore
|
|
337
|
-
|
|
338
|
-
return decorator
|
|
339
|
-
|
|
340
|
-
|
|
341
|
-
def trace_async(
|
|
342
|
-
name: Optional[str] = None,
|
|
343
|
-
kind: SpanKind = SpanKind.INTERNAL,
|
|
344
|
-
record_args: bool = True,
|
|
345
|
-
record_result: bool = False,
|
|
346
|
-
) -> Callable[[F], F]:
|
|
347
|
-
"""
|
|
348
|
-
Decorator to trace an async method.
|
|
349
|
-
|
|
350
|
-
Args:
|
|
351
|
-
name: Span name (defaults to function name)
|
|
352
|
-
kind: Span kind
|
|
353
|
-
record_args: Whether to record function arguments as attributes
|
|
354
|
-
record_result: Whether to record the return value
|
|
355
|
-
|
|
356
|
-
Usage:
|
|
357
|
-
@trace_async(name="my_async_operation")
|
|
358
|
-
async def my_async_function(arg1, arg2):
|
|
359
|
-
return result
|
|
360
|
-
"""
|
|
361
|
-
|
|
362
|
-
def decorator(func: F) -> F:
|
|
363
|
-
span_name = name or func.__qualname__
|
|
364
|
-
|
|
365
|
-
@functools.wraps(func)
|
|
366
|
-
async def wrapper(*args, **kwargs):
|
|
367
|
-
tracer = get_tracer(func.__module__)
|
|
368
|
-
|
|
369
|
-
attributes: Dict[str, Any] = {
|
|
370
|
-
"code.function": func.__name__,
|
|
371
|
-
"code.namespace": func.__module__,
|
|
372
|
-
}
|
|
373
|
-
|
|
374
|
-
if record_args:
|
|
375
|
-
# Record positional args (skip 'self' for methods)
|
|
376
|
-
arg_names = func.__code__.co_varnames[: func.__code__.co_argcount]
|
|
377
|
-
start_idx = 1 if arg_names and arg_names[0] in ("self", "cls") else 0
|
|
378
|
-
for i, arg in enumerate(args[start_idx:], start=start_idx):
|
|
379
|
-
if i < len(arg_names):
|
|
380
|
-
arg_val = _safe_attribute_value(arg)
|
|
381
|
-
if arg_val is not None:
|
|
382
|
-
attributes[f"arg.{arg_names[i]}"] = arg_val
|
|
383
|
-
|
|
384
|
-
# Record keyword args
|
|
385
|
-
for key, value in kwargs.items():
|
|
386
|
-
arg_val = _safe_attribute_value(value)
|
|
387
|
-
if arg_val is not None:
|
|
388
|
-
attributes[f"arg.{key}"] = arg_val
|
|
389
|
-
|
|
390
|
-
if _otel_available:
|
|
391
|
-
otel_kind = _map_span_kind(kind)
|
|
392
|
-
with tracer.start_as_current_span(
|
|
393
|
-
span_name,
|
|
394
|
-
kind=otel_kind,
|
|
395
|
-
attributes=attributes,
|
|
396
|
-
) as span:
|
|
397
|
-
try:
|
|
398
|
-
result = await func(*args, **kwargs)
|
|
399
|
-
if record_result:
|
|
400
|
-
result_val = _safe_attribute_value(result)
|
|
401
|
-
if result_val is not None:
|
|
402
|
-
span.set_attribute("result", result_val)
|
|
403
|
-
return result
|
|
404
|
-
except Exception as e:
|
|
405
|
-
span.record_exception(e)
|
|
406
|
-
span.set_status(Status(StatusCode.ERROR, str(e)))
|
|
407
|
-
raise
|
|
408
|
-
else:
|
|
409
|
-
with tracer.start_as_current_span(
|
|
410
|
-
span_name,
|
|
411
|
-
kind=kind,
|
|
412
|
-
attributes=attributes,
|
|
413
|
-
) as span:
|
|
414
|
-
return await func(*args, **kwargs)
|
|
415
|
-
|
|
416
|
-
return wrapper # type: ignore
|
|
417
|
-
|
|
418
|
-
return decorator
|
|
419
|
-
|
|
420
|
-
|
|
421
|
-
def _safe_attribute_value(value: Any) -> Optional[Union[str, int, float, bool]]:
|
|
422
|
-
"""
|
|
423
|
-
Convert a value to a safe attribute value for tracing.
|
|
424
|
-
|
|
425
|
-
OpenTelemetry only supports certain types for attributes.
|
|
426
|
-
"""
|
|
427
|
-
if value is None:
|
|
428
|
-
return None
|
|
429
|
-
if isinstance(value, (str, int, float, bool)):
|
|
430
|
-
return value
|
|
431
|
-
if isinstance(value, (list, tuple)):
|
|
432
|
-
if len(value) <= 10: # Limit list size
|
|
433
|
-
return str(value)
|
|
434
|
-
return f"[{len(value)} items]"
|
|
435
|
-
if isinstance(value, dict):
|
|
436
|
-
if len(value) <= 5: # Limit dict size
|
|
437
|
-
return str(value)
|
|
438
|
-
return f"{{{len(value)} items}}"
|
|
439
|
-
# For complex objects, return type and id
|
|
440
|
-
return f"<{type(value).__name__}>"
|
|
1
|
+
"""
|
|
2
|
+
ALMA Distributed Tracing.
|
|
3
|
+
|
|
4
|
+
Provides distributed tracing using OpenTelemetry with fallback
|
|
5
|
+
to logging when OTel is not available.
|
|
6
|
+
"""
|
|
7
|
+
|
|
8
|
+
import functools
|
|
9
|
+
import logging
|
|
10
|
+
from contextlib import contextmanager
|
|
11
|
+
from enum import Enum
|
|
12
|
+
from typing import Any, Callable, Dict, Optional, TypeVar, Union
|
|
13
|
+
|
|
14
|
+
# Try to import OpenTelemetry
|
|
15
|
+
_otel_available = False
|
|
16
|
+
_NoOpSpan = None
|
|
17
|
+
_NoOpTracer = None
|
|
18
|
+
|
|
19
|
+
try:
|
|
20
|
+
from opentelemetry import trace
|
|
21
|
+
from opentelemetry.trace import SpanKind as OTelSpanKind
|
|
22
|
+
from opentelemetry.trace import Status, StatusCode
|
|
23
|
+
|
|
24
|
+
_otel_available = True
|
|
25
|
+
except ImportError:
|
|
26
|
+
pass
|
|
27
|
+
|
|
28
|
+
|
|
29
|
+
logger = logging.getLogger(__name__)
|
|
30
|
+
|
|
31
|
+
# Type variable for decorated functions
|
|
32
|
+
F = TypeVar("F", bound=Callable[..., Any])
|
|
33
|
+
|
|
34
|
+
|
|
35
|
+
class SpanKind(Enum):
|
|
36
|
+
"""Span kind enum (mirrors OpenTelemetry SpanKind)."""
|
|
37
|
+
|
|
38
|
+
INTERNAL = "internal"
|
|
39
|
+
SERVER = "server"
|
|
40
|
+
CLIENT = "client"
|
|
41
|
+
PRODUCER = "producer"
|
|
42
|
+
CONSUMER = "consumer"
|
|
43
|
+
|
|
44
|
+
|
|
45
|
+
class NoOpSpan:
|
|
46
|
+
"""No-op span implementation when OpenTelemetry is not available."""
|
|
47
|
+
|
|
48
|
+
def __init__(self, name: str, attributes: Optional[Dict[str, Any]] = None):
|
|
49
|
+
self.name = name
|
|
50
|
+
self.attributes = attributes or {}
|
|
51
|
+
self._logger = logging.getLogger(f"alma.trace.{name}")
|
|
52
|
+
|
|
53
|
+
def set_attribute(self, key: str, value: Any):
|
|
54
|
+
"""Set a span attribute."""
|
|
55
|
+
self.attributes[key] = value
|
|
56
|
+
|
|
57
|
+
def set_attributes(self, attributes: Dict[str, Any]):
|
|
58
|
+
"""Set multiple span attributes."""
|
|
59
|
+
self.attributes.update(attributes)
|
|
60
|
+
|
|
61
|
+
def add_event(self, name: str, attributes: Optional[Dict[str, Any]] = None):
|
|
62
|
+
"""Add an event to the span."""
|
|
63
|
+
self._logger.debug(f"Event: {name}", extra={"event_attributes": attributes})
|
|
64
|
+
|
|
65
|
+
def set_status(self, status: Any, description: Optional[str] = None):
|
|
66
|
+
"""Set span status."""
|
|
67
|
+
pass
|
|
68
|
+
|
|
69
|
+
def record_exception(
|
|
70
|
+
self, exception: BaseException, attributes: Optional[Dict[str, Any]] = None
|
|
71
|
+
):
|
|
72
|
+
"""Record an exception."""
|
|
73
|
+
self._logger.error(f"Exception in span {self.name}: {exception}", exc_info=True)
|
|
74
|
+
|
|
75
|
+
def end(self, end_time: Optional[int] = None):
|
|
76
|
+
"""End the span."""
|
|
77
|
+
pass
|
|
78
|
+
|
|
79
|
+
def __enter__(self):
|
|
80
|
+
return self
|
|
81
|
+
|
|
82
|
+
def __exit__(self, exc_type, exc_val, exc_tb):
|
|
83
|
+
if exc_type is not None:
|
|
84
|
+
self.record_exception(exc_val)
|
|
85
|
+
return False
|
|
86
|
+
|
|
87
|
+
|
|
88
|
+
class NoOpTracer:
|
|
89
|
+
"""No-op tracer implementation when OpenTelemetry is not available."""
|
|
90
|
+
|
|
91
|
+
def __init__(self, name: str):
|
|
92
|
+
self.name = name
|
|
93
|
+
|
|
94
|
+
def start_span(
|
|
95
|
+
self,
|
|
96
|
+
name: str,
|
|
97
|
+
context: Optional[Any] = None,
|
|
98
|
+
kind: SpanKind = SpanKind.INTERNAL,
|
|
99
|
+
attributes: Optional[Dict[str, Any]] = None,
|
|
100
|
+
start_time: Optional[int] = None,
|
|
101
|
+
) -> NoOpSpan:
|
|
102
|
+
"""Start a new span."""
|
|
103
|
+
return NoOpSpan(name, attributes)
|
|
104
|
+
|
|
105
|
+
@contextmanager
|
|
106
|
+
def start_as_current_span(
|
|
107
|
+
self,
|
|
108
|
+
name: str,
|
|
109
|
+
context: Optional[Any] = None,
|
|
110
|
+
kind: SpanKind = SpanKind.INTERNAL,
|
|
111
|
+
attributes: Optional[Dict[str, Any]] = None,
|
|
112
|
+
start_time: Optional[int] = None,
|
|
113
|
+
):
|
|
114
|
+
"""Start a span as the current span."""
|
|
115
|
+
span = NoOpSpan(name, attributes)
|
|
116
|
+
try:
|
|
117
|
+
yield span
|
|
118
|
+
except Exception as e:
|
|
119
|
+
span.record_exception(e)
|
|
120
|
+
raise
|
|
121
|
+
finally:
|
|
122
|
+
span.end()
|
|
123
|
+
|
|
124
|
+
|
|
125
|
+
class TracingContext:
|
|
126
|
+
"""
|
|
127
|
+
Context for managing trace propagation and span creation.
|
|
128
|
+
|
|
129
|
+
Provides a unified interface for tracing regardless of
|
|
130
|
+
whether OpenTelemetry is available.
|
|
131
|
+
"""
|
|
132
|
+
|
|
133
|
+
def __init__(self, tracer_name: str = "alma"):
|
|
134
|
+
"""
|
|
135
|
+
Initialize tracing context.
|
|
136
|
+
|
|
137
|
+
Args:
|
|
138
|
+
tracer_name: Name for the tracer
|
|
139
|
+
"""
|
|
140
|
+
self.tracer_name = tracer_name
|
|
141
|
+
self._tracer = None
|
|
142
|
+
|
|
143
|
+
@property
|
|
144
|
+
def tracer(self):
|
|
145
|
+
"""Get the tracer (lazy initialization)."""
|
|
146
|
+
if self._tracer is None:
|
|
147
|
+
self._tracer = get_tracer(self.tracer_name)
|
|
148
|
+
return self._tracer
|
|
149
|
+
|
|
150
|
+
@contextmanager
|
|
151
|
+
def span(
|
|
152
|
+
self,
|
|
153
|
+
name: str,
|
|
154
|
+
kind: SpanKind = SpanKind.INTERNAL,
|
|
155
|
+
attributes: Optional[Dict[str, Any]] = None,
|
|
156
|
+
):
|
|
157
|
+
"""
|
|
158
|
+
Create a span context manager.
|
|
159
|
+
|
|
160
|
+
Args:
|
|
161
|
+
name: Span name
|
|
162
|
+
kind: Span kind
|
|
163
|
+
attributes: Initial span attributes
|
|
164
|
+
|
|
165
|
+
Yields:
|
|
166
|
+
The created span
|
|
167
|
+
"""
|
|
168
|
+
if _otel_available:
|
|
169
|
+
otel_kind = _map_span_kind(kind)
|
|
170
|
+
with self.tracer.start_as_current_span(
|
|
171
|
+
name,
|
|
172
|
+
kind=otel_kind,
|
|
173
|
+
attributes=attributes,
|
|
174
|
+
) as span:
|
|
175
|
+
yield span
|
|
176
|
+
else:
|
|
177
|
+
with self.tracer.start_as_current_span(
|
|
178
|
+
name,
|
|
179
|
+
kind=kind,
|
|
180
|
+
attributes=attributes,
|
|
181
|
+
) as span:
|
|
182
|
+
yield span
|
|
183
|
+
|
|
184
|
+
def create_span(
|
|
185
|
+
self,
|
|
186
|
+
name: str,
|
|
187
|
+
kind: SpanKind = SpanKind.INTERNAL,
|
|
188
|
+
attributes: Optional[Dict[str, Any]] = None,
|
|
189
|
+
):
|
|
190
|
+
"""
|
|
191
|
+
Create a span (not automatically set as current).
|
|
192
|
+
|
|
193
|
+
Args:
|
|
194
|
+
name: Span name
|
|
195
|
+
kind: Span kind
|
|
196
|
+
attributes: Initial span attributes
|
|
197
|
+
|
|
198
|
+
Returns:
|
|
199
|
+
The created span
|
|
200
|
+
"""
|
|
201
|
+
if _otel_available:
|
|
202
|
+
otel_kind = _map_span_kind(kind)
|
|
203
|
+
return self.tracer.start_span(
|
|
204
|
+
name,
|
|
205
|
+
kind=otel_kind,
|
|
206
|
+
attributes=attributes,
|
|
207
|
+
)
|
|
208
|
+
else:
|
|
209
|
+
return self.tracer.start_span(
|
|
210
|
+
name,
|
|
211
|
+
kind=kind,
|
|
212
|
+
attributes=attributes,
|
|
213
|
+
)
|
|
214
|
+
|
|
215
|
+
|
|
216
|
+
def _map_span_kind(kind: SpanKind):
|
|
217
|
+
"""Map our SpanKind to OpenTelemetry SpanKind."""
|
|
218
|
+
if not _otel_available:
|
|
219
|
+
return kind
|
|
220
|
+
|
|
221
|
+
mapping = {
|
|
222
|
+
SpanKind.INTERNAL: OTelSpanKind.INTERNAL,
|
|
223
|
+
SpanKind.SERVER: OTelSpanKind.SERVER,
|
|
224
|
+
SpanKind.CLIENT: OTelSpanKind.CLIENT,
|
|
225
|
+
SpanKind.PRODUCER: OTelSpanKind.PRODUCER,
|
|
226
|
+
SpanKind.CONSUMER: OTelSpanKind.CONSUMER,
|
|
227
|
+
}
|
|
228
|
+
return mapping.get(kind, OTelSpanKind.INTERNAL)
|
|
229
|
+
|
|
230
|
+
|
|
231
|
+
def get_tracer(name: str = "alma") -> Union["NoOpTracer", Any]:
|
|
232
|
+
"""
|
|
233
|
+
Get a tracer for the given name.
|
|
234
|
+
|
|
235
|
+
Uses OpenTelemetry tracer if available, otherwise returns
|
|
236
|
+
a no-op tracer that logs span information.
|
|
237
|
+
|
|
238
|
+
Args:
|
|
239
|
+
name: Tracer name (typically module name)
|
|
240
|
+
|
|
241
|
+
Returns:
|
|
242
|
+
Tracer instance
|
|
243
|
+
"""
|
|
244
|
+
if _otel_available:
|
|
245
|
+
return trace.get_tracer(name)
|
|
246
|
+
return NoOpTracer(name)
|
|
247
|
+
|
|
248
|
+
|
|
249
|
+
def get_current_span():
|
|
250
|
+
"""
|
|
251
|
+
Get the current span.
|
|
252
|
+
|
|
253
|
+
Returns:
|
|
254
|
+
Current span or NoOpSpan if no span is active
|
|
255
|
+
"""
|
|
256
|
+
if _otel_available:
|
|
257
|
+
return trace.get_current_span()
|
|
258
|
+
return NoOpSpan("current")
|
|
259
|
+
|
|
260
|
+
|
|
261
|
+
def trace_method(
|
|
262
|
+
name: Optional[str] = None,
|
|
263
|
+
kind: SpanKind = SpanKind.INTERNAL,
|
|
264
|
+
record_args: bool = True,
|
|
265
|
+
record_result: bool = False,
|
|
266
|
+
) -> Callable[[F], F]:
|
|
267
|
+
"""
|
|
268
|
+
Decorator to trace a synchronous method.
|
|
269
|
+
|
|
270
|
+
Args:
|
|
271
|
+
name: Span name (defaults to function name)
|
|
272
|
+
kind: Span kind
|
|
273
|
+
record_args: Whether to record function arguments as attributes
|
|
274
|
+
record_result: Whether to record the return value
|
|
275
|
+
|
|
276
|
+
Usage:
|
|
277
|
+
@trace_method(name="my_operation")
|
|
278
|
+
def my_function(arg1, arg2):
|
|
279
|
+
return result
|
|
280
|
+
"""
|
|
281
|
+
|
|
282
|
+
def decorator(func: F) -> F:
|
|
283
|
+
span_name = name or func.__qualname__
|
|
284
|
+
|
|
285
|
+
@functools.wraps(func)
|
|
286
|
+
def wrapper(*args, **kwargs):
|
|
287
|
+
tracer = get_tracer(func.__module__)
|
|
288
|
+
|
|
289
|
+
attributes: Dict[str, Any] = {
|
|
290
|
+
"code.function": func.__name__,
|
|
291
|
+
"code.namespace": func.__module__,
|
|
292
|
+
}
|
|
293
|
+
|
|
294
|
+
if record_args:
|
|
295
|
+
# Record positional args (skip 'self' for methods)
|
|
296
|
+
arg_names = func.__code__.co_varnames[: func.__code__.co_argcount]
|
|
297
|
+
start_idx = 1 if arg_names and arg_names[0] in ("self", "cls") else 0
|
|
298
|
+
for i, arg in enumerate(args[start_idx:], start=start_idx):
|
|
299
|
+
if i < len(arg_names):
|
|
300
|
+
arg_val = _safe_attribute_value(arg)
|
|
301
|
+
if arg_val is not None:
|
|
302
|
+
attributes[f"arg.{arg_names[i]}"] = arg_val
|
|
303
|
+
|
|
304
|
+
# Record keyword args
|
|
305
|
+
for key, value in kwargs.items():
|
|
306
|
+
arg_val = _safe_attribute_value(value)
|
|
307
|
+
if arg_val is not None:
|
|
308
|
+
attributes[f"arg.{key}"] = arg_val
|
|
309
|
+
|
|
310
|
+
if _otel_available:
|
|
311
|
+
otel_kind = _map_span_kind(kind)
|
|
312
|
+
with tracer.start_as_current_span(
|
|
313
|
+
span_name,
|
|
314
|
+
kind=otel_kind,
|
|
315
|
+
attributes=attributes,
|
|
316
|
+
) as span:
|
|
317
|
+
try:
|
|
318
|
+
result = func(*args, **kwargs)
|
|
319
|
+
if record_result:
|
|
320
|
+
result_val = _safe_attribute_value(result)
|
|
321
|
+
if result_val is not None:
|
|
322
|
+
span.set_attribute("result", result_val)
|
|
323
|
+
return result
|
|
324
|
+
except Exception as e:
|
|
325
|
+
span.record_exception(e)
|
|
326
|
+
span.set_status(Status(StatusCode.ERROR, str(e)))
|
|
327
|
+
raise
|
|
328
|
+
else:
|
|
329
|
+
with tracer.start_as_current_span(
|
|
330
|
+
span_name,
|
|
331
|
+
kind=kind,
|
|
332
|
+
attributes=attributes,
|
|
333
|
+
) as span:
|
|
334
|
+
return func(*args, **kwargs)
|
|
335
|
+
|
|
336
|
+
return wrapper # type: ignore
|
|
337
|
+
|
|
338
|
+
return decorator
|
|
339
|
+
|
|
340
|
+
|
|
341
|
+
def trace_async(
|
|
342
|
+
name: Optional[str] = None,
|
|
343
|
+
kind: SpanKind = SpanKind.INTERNAL,
|
|
344
|
+
record_args: bool = True,
|
|
345
|
+
record_result: bool = False,
|
|
346
|
+
) -> Callable[[F], F]:
|
|
347
|
+
"""
|
|
348
|
+
Decorator to trace an async method.
|
|
349
|
+
|
|
350
|
+
Args:
|
|
351
|
+
name: Span name (defaults to function name)
|
|
352
|
+
kind: Span kind
|
|
353
|
+
record_args: Whether to record function arguments as attributes
|
|
354
|
+
record_result: Whether to record the return value
|
|
355
|
+
|
|
356
|
+
Usage:
|
|
357
|
+
@trace_async(name="my_async_operation")
|
|
358
|
+
async def my_async_function(arg1, arg2):
|
|
359
|
+
return result
|
|
360
|
+
"""
|
|
361
|
+
|
|
362
|
+
def decorator(func: F) -> F:
|
|
363
|
+
span_name = name or func.__qualname__
|
|
364
|
+
|
|
365
|
+
@functools.wraps(func)
|
|
366
|
+
async def wrapper(*args, **kwargs):
|
|
367
|
+
tracer = get_tracer(func.__module__)
|
|
368
|
+
|
|
369
|
+
attributes: Dict[str, Any] = {
|
|
370
|
+
"code.function": func.__name__,
|
|
371
|
+
"code.namespace": func.__module__,
|
|
372
|
+
}
|
|
373
|
+
|
|
374
|
+
if record_args:
|
|
375
|
+
# Record positional args (skip 'self' for methods)
|
|
376
|
+
arg_names = func.__code__.co_varnames[: func.__code__.co_argcount]
|
|
377
|
+
start_idx = 1 if arg_names and arg_names[0] in ("self", "cls") else 0
|
|
378
|
+
for i, arg in enumerate(args[start_idx:], start=start_idx):
|
|
379
|
+
if i < len(arg_names):
|
|
380
|
+
arg_val = _safe_attribute_value(arg)
|
|
381
|
+
if arg_val is not None:
|
|
382
|
+
attributes[f"arg.{arg_names[i]}"] = arg_val
|
|
383
|
+
|
|
384
|
+
# Record keyword args
|
|
385
|
+
for key, value in kwargs.items():
|
|
386
|
+
arg_val = _safe_attribute_value(value)
|
|
387
|
+
if arg_val is not None:
|
|
388
|
+
attributes[f"arg.{key}"] = arg_val
|
|
389
|
+
|
|
390
|
+
if _otel_available:
|
|
391
|
+
otel_kind = _map_span_kind(kind)
|
|
392
|
+
with tracer.start_as_current_span(
|
|
393
|
+
span_name,
|
|
394
|
+
kind=otel_kind,
|
|
395
|
+
attributes=attributes,
|
|
396
|
+
) as span:
|
|
397
|
+
try:
|
|
398
|
+
result = await func(*args, **kwargs)
|
|
399
|
+
if record_result:
|
|
400
|
+
result_val = _safe_attribute_value(result)
|
|
401
|
+
if result_val is not None:
|
|
402
|
+
span.set_attribute("result", result_val)
|
|
403
|
+
return result
|
|
404
|
+
except Exception as e:
|
|
405
|
+
span.record_exception(e)
|
|
406
|
+
span.set_status(Status(StatusCode.ERROR, str(e)))
|
|
407
|
+
raise
|
|
408
|
+
else:
|
|
409
|
+
with tracer.start_as_current_span(
|
|
410
|
+
span_name,
|
|
411
|
+
kind=kind,
|
|
412
|
+
attributes=attributes,
|
|
413
|
+
) as span:
|
|
414
|
+
return await func(*args, **kwargs)
|
|
415
|
+
|
|
416
|
+
return wrapper # type: ignore
|
|
417
|
+
|
|
418
|
+
return decorator
|
|
419
|
+
|
|
420
|
+
|
|
421
|
+
def _safe_attribute_value(value: Any) -> Optional[Union[str, int, float, bool]]:
|
|
422
|
+
"""
|
|
423
|
+
Convert a value to a safe attribute value for tracing.
|
|
424
|
+
|
|
425
|
+
OpenTelemetry only supports certain types for attributes.
|
|
426
|
+
"""
|
|
427
|
+
if value is None:
|
|
428
|
+
return None
|
|
429
|
+
if isinstance(value, (str, int, float, bool)):
|
|
430
|
+
return value
|
|
431
|
+
if isinstance(value, (list, tuple)):
|
|
432
|
+
if len(value) <= 10: # Limit list size
|
|
433
|
+
return str(value)
|
|
434
|
+
return f"[{len(value)} items]"
|
|
435
|
+
if isinstance(value, dict):
|
|
436
|
+
if len(value) <= 5: # Limit dict size
|
|
437
|
+
return str(value)
|
|
438
|
+
return f"{{{len(value)} items}}"
|
|
439
|
+
# For complex objects, return type and id
|
|
440
|
+
return f"<{type(value).__name__}>"
|