aiqa-client 0.1.6__tar.gz → 0.2.0__tar.gz

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.
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: aiqa-client
3
- Version: 0.1.6
3
+ Version: 0.2.0
4
4
  Summary: OpenTelemetry-based Python client for tracing functions and sending traces to the AIQA server
5
5
  Author-email: AIQA <info@aiqa.dev>
6
6
  License: MIT
@@ -4,26 +4,26 @@ Python client for AIQA server - OpenTelemetry tracing decorators.
4
4
 
5
5
  from .tracing import (
6
6
  WithTracing,
7
- flush_spans,
7
+ flush_tracing,
8
8
  shutdown_tracing,
9
9
  set_span_attribute,
10
10
  set_span_name,
11
11
  get_active_span,
12
- provider,
13
- exporter,
12
+ get_provider,
13
+ get_exporter,
14
14
  )
15
15
 
16
- __version__ = "0.1.6"
16
+ __version__ = "0.2.0"
17
17
 
18
18
  __all__ = [
19
19
  "WithTracing",
20
- "flush_spans",
20
+ "flush_tracing",
21
21
  "shutdown_tracing",
22
22
  "set_span_attribute",
23
23
  "set_span_name",
24
24
  "get_active_span",
25
- "provider",
26
- "exporter",
25
+ "get_provider",
26
+ "get_exporter",
27
27
  "__version__",
28
28
  ]
29
29
 
@@ -65,7 +65,7 @@ class AIQASpanExporter(SpanExporter):
65
65
  if not spans:
66
66
  logger.debug("export() called with empty spans list")
67
67
  return SpanExportResult.SUCCESS
68
-
68
+ logger.debug(f"AIQA export() called with {len(spans)} spans")
69
69
  # Serialize and add to buffer
70
70
  with self.buffer_lock:
71
71
  serialized_spans = [self._serialize_span(span) for span in spans]
@@ -0,0 +1,49 @@
1
+ # aiqa/client.py
2
+ import os
3
+ from functools import lru_cache
4
+ from opentelemetry import trace
5
+ from opentelemetry.sdk.trace import TracerProvider
6
+ from opentelemetry.sdk.trace.export import BatchSpanProcessor
7
+
8
+ from .aiqa_exporter import AIQASpanExporter
9
+
10
+ AIQA_TRACER_NAME = "aiqa-tracer"
11
+
12
+ client = {
13
+ "provider": None,
14
+ "exporter": None,
15
+ }
16
+
17
+ @lru_cache(maxsize=1)
18
+ def get_client():
19
+ global client
20
+ _init_tracing()
21
+ # optionally return a richer client object; for now you just need init
22
+ return client
23
+
24
+ def _init_tracing():
25
+ provider = trace.get_tracer_provider()
26
+
27
+ # If it's still the default proxy, install a real SDK provider
28
+ if not isinstance(provider, TracerProvider):
29
+ provider = TracerProvider()
30
+ trace.set_tracer_provider(provider)
31
+
32
+ # Idempotently add your processor
33
+ _attach_aiqa_processor(provider)
34
+ global client
35
+ client["provider"] = provider
36
+
37
+ def _attach_aiqa_processor(provider: TracerProvider):
38
+ # Avoid double-adding if get_client() is called multiple times
39
+ for p in provider._active_span_processor._span_processors:
40
+ if isinstance(getattr(p, "exporter", None), AIQASpanExporter):
41
+ return
42
+
43
+ exporter = AIQASpanExporter(
44
+ server_url=os.getenv("AIQA_SERVER_URL"),
45
+ api_key=os.getenv("AIQA_API_KEY"),
46
+ )
47
+ provider.add_span_processor(BatchSpanProcessor(exporter))
48
+ global client
49
+ client["exporter"] = exporter
@@ -0,0 +1,543 @@
1
+ """
2
+ OpenTelemetry tracing setup and utilities. Initializes tracer provider on import.
3
+ Provides WithTracing decorator to automatically trace function calls.
4
+ """
5
+
6
+ import os
7
+ import json
8
+ import logging
9
+ import inspect
10
+ from typing import Any, Callable, Optional, Dict
11
+ from functools import wraps
12
+ from opentelemetry import trace
13
+ from opentelemetry.sdk.trace import TracerProvider
14
+ from opentelemetry.sdk.trace.export import BatchSpanProcessor
15
+ from opentelemetry.sdk.trace.sampling import ALWAYS_ON
16
+ from opentelemetry.sdk.resources import Resource
17
+ from opentelemetry.semconv.resource import ResourceAttributes
18
+ from opentelemetry.trace import Status, StatusCode
19
+ from opentelemetry.trace.propagation.tracecontext import TraceContextTextMapPropagator
20
+ from .aiqa_exporter import AIQASpanExporter
21
+ from .client import get_client, AIQA_TRACER_NAME
22
+
23
+ logger = logging.getLogger(__name__)
24
+
25
+
26
+ async def flush_tracing() -> None:
27
+ """
28
+ Flush all pending spans to the server.
29
+ Flushes also happen automatically every few seconds. So you only need to call this function
30
+ if you want to flush immediately, e.g. before exiting a process.
31
+
32
+ This flushes both the BatchSpanProcessor and the exporter buffer.
33
+ """
34
+ client = get_client()
35
+ if client.get("provider"):
36
+ client["provider"].force_flush() # Synchronous method
37
+ if client.get("exporter"):
38
+ await client["exporter"].flush()
39
+
40
+
41
+ async def shutdown_tracing() -> None:
42
+ """
43
+ Shutdown the tracer provider and exporter.
44
+ It is not necessary to call this function.
45
+ """
46
+ try:
47
+ client = get_client()
48
+ if client.get("provider"):
49
+ client["provider"].shutdown() # Synchronous method
50
+ if client.get("exporter"):
51
+ await client["exporter"].shutdown() # async method
52
+ except Exception as e:
53
+ logger.error(f"Error shutting down tracing: {e}", exc_info=True)
54
+
55
+
56
+ # Export provider and exporter accessors for advanced usage
57
+ __all__ = ["get_provider", "get_exporter", "flush_tracing", "shutdown_tracing", "WithTracing", "set_span_attribute", "set_span_name", "get_active_span"]
58
+
59
+
60
+ class TracingOptions:
61
+ """Options for WithTracing decorator"""
62
+
63
+ def __init__(
64
+ self,
65
+ name: Optional[str] = None,
66
+ ignore_input: Optional[Any] = None,
67
+ ignore_output: Optional[Any] = None,
68
+ filter_input: Optional[Callable[[Any], Any]] = None,
69
+ filter_output: Optional[Callable[[Any], Any]] = None,
70
+ ):
71
+ self.name = name
72
+ self.ignore_input = ignore_input
73
+ self.ignore_output = ignore_output
74
+ self.filter_input = filter_input
75
+ self.filter_output = filter_output
76
+
77
+
78
+ def _serialize_for_span(value: Any) -> Any:
79
+ """
80
+ Serialize a value for span attributes.
81
+ OpenTelemetry only accepts primitives (bool, str, bytes, int, float) or sequences of those.
82
+ Complex types (dicts, lists, objects) are converted to JSON strings.
83
+ """
84
+ # Keep primitives as is (including None)
85
+ if value is None or isinstance(value, (str, int, float, bool, bytes)):
86
+ return value
87
+
88
+ # For sequences, check if all elements are primitives
89
+ if isinstance(value, (list, tuple)):
90
+ # If all elements are primitives, return as list
91
+ if all(isinstance(item, (str, int, float, bool, bytes, type(None))) for item in value):
92
+ return list(value)
93
+ # Otherwise serialize to JSON string
94
+ try:
95
+ return json.dumps(value)
96
+ except (TypeError, ValueError):
97
+ return str(value)
98
+
99
+ # For dicts and other complex types, serialize to JSON string
100
+ try:
101
+ return json.dumps(value)
102
+ except (TypeError, ValueError):
103
+ # If JSON serialization fails, convert to string
104
+ return str(value)
105
+
106
+
107
+ def _prepare_input(args: tuple, kwargs: dict) -> Any:
108
+ """Prepare input for span attributes."""
109
+ if not args and not kwargs:
110
+ return None
111
+ if len(args) == 1 and not kwargs:
112
+ return _serialize_for_span(args[0])
113
+ # Multiple args or kwargs - combine into dict
114
+ result = {}
115
+ if args:
116
+ result["args"] = [_serialize_for_span(arg) for arg in args]
117
+ if kwargs:
118
+ result["kwargs"] = {k: _serialize_for_span(v) for k, v in kwargs.items()}
119
+ return result
120
+
121
+
122
+ def _prepare_and_filter_input(
123
+ args: tuple,
124
+ kwargs: dict,
125
+ filter_input: Optional[Callable[[Any], Any]],
126
+ ignore_input: Optional[Any],
127
+ ) -> Any:
128
+ """Prepare and filter input for span attributes."""
129
+ input_data = _prepare_input(args, kwargs)
130
+ if filter_input:
131
+ input_data = filter_input(input_data)
132
+ if ignore_input and isinstance(input_data, dict):
133
+ for key in ignore_input:
134
+ if key in input_data:
135
+ del input_data[key]
136
+ return input_data
137
+
138
+
139
+ def _prepare_and_filter_output(
140
+ result: Any,
141
+ filter_output: Optional[Callable[[Any], Any]],
142
+ ignore_output: Optional[Any],
143
+ ) -> Any:
144
+ """Prepare and filter output for span attributes."""
145
+ output_data = result
146
+ if filter_output:
147
+ output_data = filter_output(output_data)
148
+ if ignore_output and isinstance(output_data, dict):
149
+ output_data = output_data.copy()
150
+ for key in ignore_output:
151
+ if key in output_data:
152
+ del output_data[key]
153
+ return output_data
154
+
155
+
156
+ def _handle_span_exception(span: trace.Span, exception: Exception) -> None:
157
+ """Record exception on span and set error status."""
158
+ error = exception if isinstance(exception, Exception) else Exception(str(exception))
159
+ span.record_exception(error)
160
+ span.set_status(Status(StatusCode.ERROR, str(error)))
161
+
162
+
163
+ class TracedGenerator:
164
+ """Wrapper for sync generators that traces iteration."""
165
+
166
+ def __init__(
167
+ self,
168
+ generator: Any,
169
+ span: trace.Span,
170
+ fn_name: str,
171
+ filter_output: Optional[Callable[[Any], Any]],
172
+ ignore_output: Optional[Any],
173
+ context_token: Any,
174
+ ):
175
+ self._generator = generator
176
+ self._span = span
177
+ self._fn_name = fn_name
178
+ self._filter_output = filter_output
179
+ self._ignore_output = ignore_output
180
+ self._context_token = context_token
181
+ self._yielded_values = []
182
+ self._exhausted = False
183
+
184
+ def __iter__(self):
185
+ return self
186
+
187
+ def __next__(self):
188
+ if self._exhausted:
189
+ raise StopIteration
190
+
191
+ try:
192
+ value = next(self._generator)
193
+ self._yielded_values.append(value)
194
+ return value
195
+ except StopIteration:
196
+ self._exhausted = True
197
+ self._finalize_span_success()
198
+ trace.context_api.detach(self._context_token)
199
+ self._span.end()
200
+ raise
201
+ except Exception as exception:
202
+ self._exhausted = True
203
+ _handle_span_exception(self._span, exception)
204
+ trace.context_api.detach(self._context_token)
205
+ self._span.end()
206
+ raise
207
+
208
+ def _finalize_span_success(self):
209
+ """Set output and success status on span."""
210
+ # Record summary of yielded values
211
+ output_data = {
212
+ "type": "generator",
213
+ "yielded_count": len(self._yielded_values),
214
+ }
215
+
216
+ # Optionally include sample values (limit to avoid huge spans)
217
+ if self._yielded_values:
218
+ sample_size = min(10, len(self._yielded_values))
219
+ output_data["sample_values"] = [
220
+ _serialize_for_span(v) for v in self._yielded_values[:sample_size]
221
+ ]
222
+ if len(self._yielded_values) > sample_size:
223
+ output_data["truncated"] = True
224
+
225
+ output_data = _prepare_and_filter_output(output_data, self._filter_output, self._ignore_output)
226
+ self._span.set_attribute("output", _serialize_for_span(output_data))
227
+ self._span.set_status(Status(StatusCode.OK))
228
+
229
+
230
+ class TracedAsyncGenerator:
231
+ """Wrapper for async generators that traces iteration."""
232
+
233
+ def __init__(
234
+ self,
235
+ generator: Any,
236
+ span: trace.Span,
237
+ fn_name: str,
238
+ filter_output: Optional[Callable[[Any], Any]],
239
+ ignore_output: Optional[Any],
240
+ context_token: Any,
241
+ ):
242
+ self._generator = generator
243
+ self._span = span
244
+ self._fn_name = fn_name
245
+ self._filter_output = filter_output
246
+ self._ignore_output = ignore_output
247
+ self._context_token = context_token
248
+ self._yielded_values = []
249
+ self._exhausted = False
250
+
251
+ def __aiter__(self):
252
+ return self
253
+
254
+ async def __anext__(self):
255
+ if self._exhausted:
256
+ raise StopAsyncIteration
257
+
258
+ try:
259
+ value = await self._generator.__anext__()
260
+ self._yielded_values.append(value)
261
+ return value
262
+ except StopAsyncIteration:
263
+ self._exhausted = True
264
+ self._finalize_span_success()
265
+ trace.context_api.detach(self._context_token)
266
+ self._span.end()
267
+ raise
268
+ except Exception as exception:
269
+ self._exhausted = True
270
+ _handle_span_exception(self._span, exception)
271
+ trace.context_api.detach(self._context_token)
272
+ self._span.end()
273
+ raise
274
+
275
+ def _finalize_span_success(self):
276
+ """Set output and success status on span."""
277
+ # Record summary of yielded values
278
+ output_data = {
279
+ "type": "async_generator",
280
+ "yielded_count": len(self._yielded_values),
281
+ }
282
+
283
+ # Optionally include sample values (limit to avoid huge spans)
284
+ if self._yielded_values:
285
+ sample_size = min(10, len(self._yielded_values))
286
+ output_data["sample_values"] = [
287
+ _serialize_for_span(v) for v in self._yielded_values[:sample_size]
288
+ ]
289
+ if len(self._yielded_values) > sample_size:
290
+ output_data["truncated"] = True
291
+
292
+ output_data = _prepare_and_filter_output(output_data, self._filter_output, self._ignore_output)
293
+ self._span.set_attribute("output", _serialize_for_span(output_data))
294
+ self._span.set_status(Status(StatusCode.OK))
295
+
296
+
297
+ def WithTracing(
298
+ func: Optional[Callable] = None,
299
+ *,
300
+ name: Optional[str] = None,
301
+ ignore_input: Optional[Any] = None,
302
+ ignore_output: Optional[Any] = None,
303
+ filter_input: Optional[Callable[[Any], Any]] = None,
304
+ filter_output: Optional[Callable[[Any], Any]] = None,
305
+ ):
306
+ """
307
+ Decorator to automatically create spans for function calls.
308
+ Records input/output as span attributes. Spans are automatically linked via OpenTelemetry context.
309
+
310
+ Works with synchronous functions, asynchronous functions, generator functions, and async generator functions.
311
+
312
+ Args:
313
+ func: The function to trace (when used as @WithTracing)
314
+ name: Optional custom name for the span (defaults to function name)
315
+ ignore_input: Fields to ignore in input (not yet implemented)
316
+ ignore_output: Fields to ignore in output (not yet implemented)
317
+ filter_input: Function to filter/transform input before recording
318
+ filter_output: Function to filter/transform output before recording
319
+
320
+ Example:
321
+ @WithTracing
322
+ def my_function(x, y):
323
+ return x + y
324
+
325
+ @WithTracing
326
+ async def my_async_function(x, y):
327
+ return x + y
328
+
329
+ @WithTracing
330
+ def my_generator(n):
331
+ for i in range(n):
332
+ yield i * 2
333
+
334
+ @WithTracing
335
+ async def my_async_generator(n):
336
+ for i in range(n):
337
+ yield i * 2
338
+
339
+ @WithTracing(name="custom_name")
340
+ def another_function():
341
+ pass
342
+ """
343
+ def decorator(fn: Callable) -> Callable:
344
+ fn_name = name or fn.__name__ or "_"
345
+
346
+ # Check if already traced
347
+ if hasattr(fn, "_is_traced"):
348
+ logger.warning(f"Function {fn_name} is already traced, skipping tracing again")
349
+ return fn
350
+
351
+ is_async = inspect.iscoroutinefunction(fn)
352
+ is_generator = inspect.isgeneratorfunction(fn)
353
+ is_async_generator = inspect.isasyncgenfunction(fn) if hasattr(inspect, 'isasyncgenfunction') else False
354
+
355
+ tracer = trace.get_tracer(AIQA_TRACER_NAME)
356
+
357
+ def _setup_span(span: trace.Span, input_data: Any) -> bool:
358
+ """Setup span with input data. Returns True if span is recording."""
359
+ if not span.is_recording():
360
+ logger.warning(f"Span {fn_name} is not recording - will not be exported")
361
+ return False
362
+
363
+ logger.debug(f"Span {fn_name} is recording, trace_id={format(span.get_span_context().trace_id, '032x')}")
364
+
365
+ if input_data is not None:
366
+ span.set_attribute("input", _serialize_for_span(input_data))
367
+
368
+ trace_id = format(span.get_span_context().trace_id, "032x")
369
+ logger.debug(f"do traceable stuff {fn_name} {trace_id}")
370
+ return True
371
+
372
+ def _finalize_span_success(span: trace.Span, result: Any) -> None:
373
+ """Set output and success status on span."""
374
+ output_data = _prepare_and_filter_output(result, filter_output, ignore_output)
375
+ span.set_attribute("output", _serialize_for_span(output_data))
376
+ span.set_status(Status(StatusCode.OK))
377
+
378
+ def _execute_with_span_sync(executor: Callable[[], Any], input_data: Any) -> Any:
379
+ """Execute sync function within span context, handling input/output and exceptions."""
380
+ with tracer.start_as_current_span(fn_name) as span:
381
+ if not _setup_span(span, input_data):
382
+ return executor()
383
+
384
+ try:
385
+ result = executor()
386
+ _finalize_span_success(span, result)
387
+ return result
388
+ except Exception as exception:
389
+ _handle_span_exception(span, exception)
390
+ raise
391
+
392
+ async def _execute_with_span_async(executor: Callable[[], Any], input_data: Any) -> Any:
393
+ """Execute async function within span context, handling input/output and exceptions."""
394
+ with tracer.start_as_current_span(fn_name) as span:
395
+ if not _setup_span(span, input_data):
396
+ return await executor()
397
+
398
+ try:
399
+ result = await executor()
400
+ _finalize_span_success(span, result)
401
+ logger.debug(f"Span {fn_name} completed successfully, is_recording={span.is_recording()}")
402
+ return result
403
+ except Exception as exception:
404
+ _handle_span_exception(span, exception)
405
+ raise
406
+ finally:
407
+ logger.debug(f"Span {fn_name} context exiting, is_recording={span.is_recording()}")
408
+
409
+ def _execute_generator_sync(executor: Callable[[], Any], input_data: Any) -> Any:
410
+ """Execute sync generator function, returning a traced generator."""
411
+ # Create span but don't use 'with' - span will be closed by TracedGenerator
412
+ span = tracer.start_span(fn_name)
413
+ token = trace.context_api.attach(trace.context_api.set_span_in_context(span))
414
+
415
+ try:
416
+ if not _setup_span(span, input_data):
417
+ generator = executor()
418
+ trace.context_api.detach(token)
419
+ span.end()
420
+ return generator
421
+
422
+ generator = executor()
423
+ return TracedGenerator(generator, span, fn_name, filter_output, ignore_output, token)
424
+ except Exception as exception:
425
+ trace.context_api.detach(token)
426
+ _handle_span_exception(span, exception)
427
+ span.end()
428
+ raise
429
+
430
+ async def _execute_generator_async(executor: Callable[[], Any], input_data: Any) -> Any:
431
+ """Execute async generator function, returning a traced async generator."""
432
+ # Create span but don't use 'with' - span will be closed by TracedAsyncGenerator
433
+ span = tracer.start_span(fn_name)
434
+ token = trace.context_api.attach(trace.context_api.set_span_in_context(span))
435
+
436
+ try:
437
+ if not _setup_span(span, input_data):
438
+ generator = await executor()
439
+ trace.context_api.detach(token)
440
+ span.end()
441
+ return generator
442
+
443
+ generator = await executor()
444
+ return TracedAsyncGenerator(generator, span, fn_name, filter_output, ignore_output, token)
445
+ except Exception as exception:
446
+ trace.context_api.detach(token)
447
+ _handle_span_exception(span, exception)
448
+ span.end()
449
+ raise
450
+
451
+ if is_async_generator:
452
+ @wraps(fn)
453
+ async def async_gen_traced_fn(*args, **kwargs):
454
+ input_data = _prepare_and_filter_input(args, kwargs, filter_input, ignore_input)
455
+ return await _execute_generator_async(
456
+ lambda: fn(*args, **kwargs),
457
+ input_data
458
+ )
459
+
460
+ async_gen_traced_fn._is_traced = True
461
+ logger.debug(f"Function {fn_name} is now traced (async generator)")
462
+ return async_gen_traced_fn
463
+ elif is_generator:
464
+ @wraps(fn)
465
+ def gen_traced_fn(*args, **kwargs):
466
+ input_data = _prepare_and_filter_input(args, kwargs, filter_input, ignore_input)
467
+ return _execute_generator_sync(
468
+ lambda: fn(*args, **kwargs),
469
+ input_data
470
+ )
471
+
472
+ gen_traced_fn._is_traced = True
473
+ logger.debug(f"Function {fn_name} is now traced (generator)")
474
+ return gen_traced_fn
475
+ elif is_async:
476
+ @wraps(fn)
477
+ async def async_traced_fn(*args, **kwargs):
478
+ input_data = _prepare_and_filter_input(args, kwargs, filter_input, ignore_input)
479
+ return await _execute_with_span_async(
480
+ lambda: fn(*args, **kwargs),
481
+ input_data
482
+ )
483
+
484
+ async_traced_fn._is_traced = True
485
+ logger.debug(f"Function {fn_name} is now traced (async)")
486
+ return async_traced_fn
487
+ else:
488
+ @wraps(fn)
489
+ def sync_traced_fn(*args, **kwargs):
490
+ input_data = _prepare_and_filter_input(args, kwargs, filter_input, ignore_input)
491
+ return _execute_with_span_sync(
492
+ lambda: fn(*args, **kwargs),
493
+ input_data
494
+ )
495
+
496
+ sync_traced_fn._is_traced = True
497
+ logger.debug(f"Function {fn_name} is now traced (sync)")
498
+ return sync_traced_fn
499
+
500
+ # Support both @WithTracing and @WithTracing(...) syntax
501
+ if func is None:
502
+ return decorator
503
+ else:
504
+ return decorator(func)
505
+
506
+
507
+ def set_span_attribute(attribute_name: str, attribute_value: Any) -> bool:
508
+ """
509
+ Set an attribute on the active span.
510
+
511
+ Returns:
512
+ True if attribute was set, False if no active span found
513
+ """
514
+ span = trace.get_current_span()
515
+ if span and span.is_recording():
516
+ span.set_attribute(attribute_name, _serialize_for_span(attribute_value))
517
+ return True
518
+ return False
519
+
520
+ def set_span_name(span_name: str) -> bool:
521
+ """
522
+ Set the name of the active span.
523
+ """
524
+ span = trace.get_current_span()
525
+ if span and span.is_recording():
526
+ span.update_name(span_name)
527
+ return True
528
+ return False
529
+
530
+ def get_active_span() -> Optional[trace.Span]:
531
+ """Get the currently active span."""
532
+ return trace.get_current_span()
533
+
534
+ def get_provider() -> Optional[TracerProvider]:
535
+ """Get the tracer provider for advanced usage."""
536
+ client = get_client()
537
+ return client.get("provider")
538
+
539
+ def get_exporter() -> Optional[AIQASpanExporter]:
540
+ """Get the exporter for advanced usage."""
541
+ client = get_client()
542
+ return client.get("exporter")
543
+
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: aiqa-client
3
- Version: 0.1.6
3
+ Version: 0.2.0
4
4
  Summary: OpenTelemetry-based Python client for tracing functions and sending traces to the AIQA server
5
5
  Author-email: AIQA <info@aiqa.dev>
6
6
  License: MIT
@@ -5,6 +5,7 @@ pyproject.toml
5
5
  setup.py
6
6
  aiqa/__init__.py
7
7
  aiqa/aiqa_exporter.py
8
+ aiqa/client.py
8
9
  aiqa/py.typed
9
10
  aiqa/tracing.py
10
11
  aiqa_client.egg-info/PKG-INFO
@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
4
4
 
5
5
  [project]
6
6
  name = "aiqa-client"
7
- version = "0.1.6"
7
+ version = "0.2.0"
8
8
  description = "OpenTelemetry-based Python client for tracing functions and sending traces to the AIQA server"
9
9
  readme = "README.md"
10
10
  requires-python = ">=3.8"
@@ -1,330 +0,0 @@
1
- """
2
- OpenTelemetry tracing setup and utilities. Initializes tracer provider on import.
3
- Provides WithTracing decorator to automatically trace function calls.
4
- """
5
-
6
- import os
7
- import json
8
- import logging
9
- import inspect
10
- from typing import Any, Callable, Optional, Dict
11
- from functools import wraps
12
- from opentelemetry import trace
13
- from opentelemetry.sdk.trace import TracerProvider
14
- from opentelemetry.sdk.trace.export import BatchSpanProcessor
15
- from opentelemetry.sdk.trace.sampling import ALWAYS_ON
16
- from opentelemetry.sdk.resources import Resource
17
- from opentelemetry.semconv.resource import ResourceAttributes
18
- from opentelemetry.trace import Status, StatusCode
19
- from opentelemetry.trace.propagation.tracecontext import TraceContextTextMapPropagator
20
- from .aiqa_exporter import AIQASpanExporter
21
-
22
- logger = logging.getLogger(__name__)
23
-
24
- # Load environment variables
25
- exporter = AIQASpanExporter()
26
-
27
- # Initialize OpenTelemetry
28
- # Use ALWAYS_ON sampler to ensure all spans are recorded (important for FastAPI async contexts)
29
- provider = TracerProvider(
30
- resource=Resource.create(
31
- {
32
- ResourceAttributes.SERVICE_NAME: os.getenv("OTEL_SERVICE_NAME", "aiqa-service"),
33
- }
34
- ),
35
- sampler=ALWAYS_ON
36
- )
37
-
38
- provider.add_span_processor(BatchSpanProcessor(exporter))
39
- trace.set_tracer_provider(provider)
40
-
41
- # Get a tracer instance
42
- tracer = trace.get_tracer("aiqa-tracer")
43
-
44
-
45
- async def flush_spans() -> None:
46
- """
47
- Flush all pending spans to the server.
48
- Flushes also happen automatically every few seconds. So you only need to call this function
49
- if you want to flush immediately, e.g. before exiting a process.
50
-
51
- This flushes both the BatchSpanProcessor and the exporter buffer.
52
- """
53
- provider.force_flush() # Synchronous method
54
- await exporter.flush()
55
-
56
-
57
- async def shutdown_tracing() -> None:
58
- """
59
- Shutdown the tracer provider and exporter.
60
- It is not necessary to call this function.
61
- """
62
- provider.shutdown() # Synchronous method
63
- await exporter.shutdown()
64
-
65
-
66
- # Export provider and exporter for advanced usage
67
- __all__ = ["provider", "exporter", "flush_spans", "shutdown_tracing", "WithTracing", "set_span_attribute", "set_span_name", "get_active_span"]
68
-
69
-
70
- class TracingOptions:
71
- """Options for WithTracing decorator"""
72
-
73
- def __init__(
74
- self,
75
- name: Optional[str] = None,
76
- ignore_input: Optional[Any] = None,
77
- ignore_output: Optional[Any] = None,
78
- filter_input: Optional[Callable[[Any], Any]] = None,
79
- filter_output: Optional[Callable[[Any], Any]] = None,
80
- ):
81
- self.name = name
82
- self.ignore_input = ignore_input
83
- self.ignore_output = ignore_output
84
- self.filter_input = filter_input
85
- self.filter_output = filter_output
86
-
87
-
88
- def _serialize_for_span(value: Any) -> Any:
89
- """
90
- Serialize a value for span attributes.
91
- OpenTelemetry only accepts primitives (bool, str, bytes, int, float) or sequences of those.
92
- Complex types (dicts, lists, objects) are converted to JSON strings.
93
- """
94
- # Keep primitives as is (including None)
95
- if value is None or isinstance(value, (str, int, float, bool, bytes)):
96
- return value
97
-
98
- # For sequences, check if all elements are primitives
99
- if isinstance(value, (list, tuple)):
100
- # If all elements are primitives, return as list
101
- if all(isinstance(item, (str, int, float, bool, bytes, type(None))) for item in value):
102
- return list(value)
103
- # Otherwise serialize to JSON string
104
- try:
105
- return json.dumps(value)
106
- except (TypeError, ValueError):
107
- return str(value)
108
-
109
- # For dicts and other complex types, serialize to JSON string
110
- try:
111
- return json.dumps(value)
112
- except (TypeError, ValueError):
113
- # If JSON serialization fails, convert to string
114
- return str(value)
115
-
116
-
117
- def _prepare_input(args: tuple, kwargs: dict) -> Any:
118
- """Prepare input for span attributes."""
119
- if not args and not kwargs:
120
- return None
121
- if len(args) == 1 and not kwargs:
122
- return _serialize_for_span(args[0])
123
- # Multiple args or kwargs - combine into dict
124
- result = {}
125
- if args:
126
- result["args"] = [_serialize_for_span(arg) for arg in args]
127
- if kwargs:
128
- result["kwargs"] = {k: _serialize_for_span(v) for k, v in kwargs.items()}
129
- return result
130
-
131
-
132
- def WithTracing(
133
- func: Optional[Callable] = None,
134
- *,
135
- name: Optional[str] = None,
136
- ignore_input: Optional[Any] = None,
137
- ignore_output: Optional[Any] = None,
138
- filter_input: Optional[Callable[[Any], Any]] = None,
139
- filter_output: Optional[Callable[[Any], Any]] = None,
140
- ):
141
- """
142
- Decorator to automatically create spans for function calls.
143
- Records input/output as span attributes. Spans are automatically linked via OpenTelemetry context.
144
-
145
- Works with both synchronous and asynchronous functions.
146
-
147
- Args:
148
- func: The function to trace (when used as @WithTracing)
149
- name: Optional custom name for the span (defaults to function name)
150
- ignore_input: Fields to ignore in input (not yet implemented)
151
- ignore_output: Fields to ignore in output (not yet implemented)
152
- filter_input: Function to filter/transform input before recording
153
- filter_output: Function to filter/transform output before recording
154
-
155
- Example:
156
- @WithTracing
157
- def my_function(x, y):
158
- return x + y
159
-
160
- @WithTracing
161
- async def my_async_function(x, y):
162
- return x + y
163
-
164
- @WithTracing(name="custom_name")
165
- def another_function():
166
- pass
167
- """
168
- def decorator(fn: Callable) -> Callable:
169
- fn_name = name or fn.__name__ or "_"
170
-
171
- # Check if already traced
172
- if hasattr(fn, "_is_traced"):
173
- logger.warning(f"Function {fn_name} is already traced, skipping tracing again")
174
- return fn
175
-
176
- is_async = inspect.iscoroutinefunction(fn)
177
-
178
- if is_async:
179
- @wraps(fn)
180
- async def async_traced_fn(*args, **kwargs):
181
- # Prepare input
182
- input_data = _prepare_input(args, kwargs)
183
- if filter_input:
184
- input_data = filter_input(input_data)
185
- if ignore_input and isinstance(input_data, dict):
186
- for key in ignore_input:
187
- if key in input_data:
188
- del input_data[key]
189
-
190
- # Use start_as_current_span to ensure span is recorded by BatchSpanProcessor
191
- with tracer.start_as_current_span(fn_name) as span:
192
- # Check if span is recording - if not, spans won't be exported
193
- if not span.is_recording():
194
- logger.warning(f"Span {fn_name} is not recording - will not be exported")
195
- return await fn(*args, **kwargs)
196
-
197
- logger.debug(f"Span {fn_name} is recording, trace_id={format(span.get_span_context().trace_id, '032x')}")
198
-
199
- if input_data is not None:
200
- # Serialize for span attributes (OpenTelemetry only accepts primitives or JSON strings)
201
- serialized_input = _serialize_for_span(input_data)
202
- span.set_attribute("input", serialized_input)
203
-
204
- try:
205
- # Call the function within the span context
206
- trace_id = format(span.get_span_context().trace_id, "032x")
207
- logger.debug(f"do traceable stuff {fn_name} {trace_id}")
208
-
209
- result = await fn(*args, **kwargs)
210
-
211
- # Prepare output
212
- output_data = result
213
- if filter_output:
214
- output_data = filter_output(output_data)
215
- if ignore_output and isinstance(output_data, dict):
216
- # Make a copy of output_data to avoid modifying the original
217
- output_data = output_data.copy()
218
- for key in ignore_output:
219
- if key in output_data:
220
- del output_data[key]
221
-
222
- span.set_attribute("output", _serialize_for_span(output_data))
223
- span.set_status(Status(StatusCode.OK))
224
-
225
- logger.debug(f"Span {fn_name} completed successfully, is_recording={span.is_recording()}")
226
- return result
227
- except Exception as exception:
228
- error = exception if isinstance(exception, Exception) else Exception(str(exception))
229
- span.record_exception(error)
230
- span.set_status(Status(StatusCode.ERROR, str(error)))
231
- raise
232
- finally:
233
- # Log when span context exits - span should be ended automatically by context manager
234
- logger.debug(f"Span {fn_name} context exiting, is_recording={span.is_recording()}")
235
-
236
- async_traced_fn._is_traced = True
237
- logger.debug(f"Function {fn_name} is now traced (async)")
238
- return async_traced_fn
239
- else:
240
- @wraps(fn)
241
- def sync_traced_fn(*args, **kwargs):
242
- # Prepare input
243
- input_data = _prepare_input(args, kwargs)
244
- if filter_input:
245
- input_data = filter_input(input_data)
246
- if ignore_input and isinstance(input_data, dict):
247
- for key in ignore_input:
248
- if key in input_data:
249
- del input_data[key]
250
-
251
- # Use start_as_current_span to ensure span is recorded by BatchSpanProcessor
252
- with tracer.start_as_current_span(fn_name) as span:
253
- # Check if span is recording - if not, spans won't be exported
254
- if not span.is_recording():
255
- logger.warning(f"Span {fn_name} is not recording - will not be exported")
256
- return fn(*args, **kwargs)
257
-
258
- logger.debug(f"Span {fn_name} is recording, trace_id={format(span.get_span_context().trace_id, '032x')}")
259
-
260
- if input_data is not None:
261
- # Serialize for span attributes (OpenTelemetry only accepts primitives or JSON strings)
262
- serialized_input = _serialize_for_span(input_data)
263
- span.set_attribute("input", serialized_input)
264
-
265
- try:
266
- # Call the function within the span context
267
- trace_id = format(span.get_span_context().trace_id, "032x")
268
- logger.debug(f"do traceable stuff {fn_name} {trace_id}")
269
-
270
- result = fn(*args, **kwargs)
271
-
272
- # Prepare output
273
- output_data = result
274
- if filter_output:
275
- output_data = filter_output(output_data)
276
- if ignore_output and isinstance(output_data, dict):
277
- # Make a copy of output_data to avoid modifying the original
278
- output_data = output_data.copy()
279
- for key in ignore_output:
280
- if key in output_data:
281
- del output_data[key]
282
-
283
- span.set_attribute("output", _serialize_for_span(output_data))
284
- span.set_status(Status(StatusCode.OK))
285
-
286
- return result
287
- except Exception as exception:
288
- error = exception if isinstance(exception, Exception) else Exception(str(exception))
289
- span.record_exception(error)
290
- span.set_status(Status(StatusCode.ERROR, str(error)))
291
- raise
292
-
293
- sync_traced_fn._is_traced = True
294
- logger.debug(f"Function {fn_name} is now traced (sync)")
295
- return sync_traced_fn
296
-
297
- # Support both @WithTracing and @WithTracing(...) syntax
298
- if func is None:
299
- return decorator
300
- else:
301
- return decorator(func)
302
-
303
-
304
- def set_span_attribute(attribute_name: str, attribute_value: Any) -> bool:
305
- """
306
- Set an attribute on the active span.
307
-
308
- Returns:
309
- True if attribute was set, False if no active span found
310
- """
311
- span = trace.get_current_span()
312
- if span and span.is_recording():
313
- span.set_attribute(attribute_name, _serialize_for_span(attribute_value))
314
- return True
315
- return False
316
-
317
- def set_span_name(span_name: str) -> bool:
318
- """
319
- Set the name of the active span.
320
- """
321
- span = trace.get_current_span()
322
- if span and span.is_recording():
323
- span.update_name(span_name)
324
- return True
325
- return False
326
-
327
- def get_active_span() -> Optional[trace.Span]:
328
- """Get the currently active span."""
329
- return trace.get_current_span()
330
-
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes