aiqa-client 0.5.2__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.
aiqa/tracing.py CHANGED
@@ -1,51 +1,24 @@
1
1
  """
2
- OpenTelemetry tracing setup and utilities. Initializes tracer provider on import.
3
- Provides WithTracing decorator to automatically trace function calls.
2
+ OpenTelemetry tracing decorator. Provides WithTracing decorator to automatically trace function calls.
4
3
  """
5
4
 
6
- import json
7
5
  import logging
8
6
  import inspect
9
- import os
10
- import copy
11
- import requests
7
+ import fnmatch
12
8
  from typing import Any, Callable, Optional, List
13
9
  from functools import wraps
14
10
  from opentelemetry import trace
15
- from opentelemetry.sdk.trace import TracerProvider
16
- from opentelemetry.trace import Status, StatusCode, SpanContext, TraceFlags
17
- from opentelemetry.propagate import inject, extract
18
- from .client import get_aiqa_client, get_component_tag, set_component_tag as _set_component_tag, get_aiqa_tracer
19
- from .constants import AIQA_TRACER_NAME, LOG_TAG
11
+ from opentelemetry.trace import Status, StatusCode
12
+
13
+ from .client import get_aiqa_client, get_component_tag, get_aiqa_tracer
14
+ from .constants import LOG_TAG
20
15
  from .object_serialiser import serialize_for_span
21
- from .http_utils import build_headers, get_server_url, get_api_key
22
16
  from .tracing_llm_utils import _extract_and_set_token_usage, _extract_and_set_provider_and_model
23
17
 
24
18
  logger = logging.getLogger(LOG_TAG)
25
19
 
26
20
 
27
- async def flush_tracing() -> None:
28
- """
29
- Flush all pending spans to the server.
30
- Flushes also happen automatically every few seconds. So you only need to call this function
31
- if you want to flush immediately, e.g. before exiting a process.
32
- A common use is if you are tracing unit tests or experiment runs.
33
-
34
- This flushes the BatchSpanProcessor (OTLP exporter doesn't have a separate flush method).
35
- """
36
- client = get_aiqa_client()
37
- if client.provider:
38
- client.provider.force_flush() # Synchronous method
39
-
40
-
41
- # Export provider and exporter accessors for advanced usage
42
-
43
- __all__ = [
44
- "flush_tracing", "WithTracing",
45
- "set_span_attribute", "set_span_name", "get_active_span",
46
- "get_active_trace_id", "get_span_id", "create_span_from_trace_id", "inject_trace_context", "extract_trace_context",
47
- "set_conversation_id", "set_component_tag", "set_token_usage", "set_provider_and_model", "get_span", "submit_feedback"
48
- ]
21
+ __all__ = ["WithTracing"]
49
22
 
50
23
 
51
24
  class TracingOptions:
@@ -73,20 +46,25 @@ class TracingOptions:
73
46
  descriptive names.
74
47
 
75
48
  ignore_input: Iterable of keys (e.g., list, set) to exclude from
76
- input data when recording span attributes. Only applies when
77
- input is a dictionary. For example, use `["password", "api_key"]`
78
- to exclude sensitive fields from being traced.
49
+ input data when recording span attributes. Applies after filter_input if both are set.
50
+ Supports "self" and simple wildcards (e.g., `"_*"`
51
+ matches `"_apple"`, `"_fruit"`). The pattern `"_*"` is applied by default
52
+ to filter properties starting with '_' in nested objects.
79
53
 
80
54
  ignore_output: Iterable of keys (e.g., list, set) to exclude from
81
55
  output data when recording span attributes. Only applies when
82
- output is a dictionary. Useful for excluding large or sensitive
83
- fields from traces.
56
+ output is a dictionary. Supports simple wildcards (e.g., `"_*"`
57
+ matches `"_apple"`, `"_fruit"`). The pattern `"_*"` is applied by default
58
+ to filter properties starting with '_' in nested objects. Useful for excluding
59
+ large or sensitive fields from traces.
84
60
 
85
- filter_input: Callable function that receives the prepared input data
86
- and returns a filtered/transformed version to be recorded in the
87
- span. The function should accept one argument (the input data)
88
- and return the transformed data. This is applied before
89
- ignore_input filtering.
61
+ filter_input: Callable function that receives the same arguments as the
62
+ decorated function (*args, **kwargs) and returns a filtered/transformed
63
+ version to be recorded in the span. This allows you to extract specific
64
+ properties from any kind of object, including `self` for methods.
65
+ The function receives the exact same inputs as the decorated function,
66
+ including `self` for bound methods. Returns a dict or any value that
67
+ will be converted to a dict. Applied before ignore_input filtering.
90
68
 
91
69
  filter_output: Callable function that receives the output data and
92
70
  returns a filtered/transformed version to be recorded in the span.
@@ -107,6 +85,22 @@ class TracingOptions:
107
85
  )
108
86
  def process_data(items):
109
87
  return items
88
+
89
+ # Extract properties from self in a method
90
+ class ExperimentRunner:
91
+ def __init__(self, dataset_id, experiment_id):
92
+ self.dataset_id = dataset_id
93
+ self.experiment_id = experiment_id
94
+
95
+ @WithTracing(
96
+ filter_input=lambda self, example: {
97
+ "dataset": self.dataset_id,
98
+ "experiment": self.experiment_id,
99
+ "example": example.id if hasattr(example, 'id') else None
100
+ }
101
+ )
102
+ def run_example(self, example):
103
+ return self.process(example)
110
104
  """
111
105
  self.name = name
112
106
  self.ignore_input = ignore_input
@@ -115,6 +109,23 @@ class TracingOptions:
115
109
  self.filter_output = filter_output
116
110
 
117
111
 
112
+ def _matches_ignore_pattern(key: str, ignore_patterns: List[str]) -> bool:
113
+ """
114
+ Check if a key matches any pattern in the ignore list.
115
+ Supports simple wildcards (e.g., "_*" matches "_apple", "_fruit").
116
+ """
117
+ for pattern in ignore_patterns:
118
+ if "*" in pattern or "?" in pattern:
119
+ # Use fnmatch for wildcard matching
120
+ if fnmatch.fnmatch(key, pattern):
121
+ return True
122
+ else:
123
+ # Exact match for non-wildcard patterns
124
+ if key == pattern:
125
+ return True
126
+ return False
127
+
128
+
118
129
  def _prepare_input(args: tuple, kwargs: dict, sig: Optional[inspect.Signature] = None) -> Any:
119
130
  """Prepare input for span attributes.
120
131
  Converts args and kwargs into a unified dict structure using function signature when available.
@@ -157,6 +168,89 @@ def _prepare_input(args: tuple, kwargs: dict, sig: Optional[inspect.Signature] =
157
168
  return result
158
169
 
159
170
 
171
+ def _apply_ignore_patterns(
172
+ data_dict: dict,
173
+ ignore_patterns: Optional[List[str]],
174
+ recursive: bool = True,
175
+ max_depth: int = 100,
176
+ current_depth: int = 0
177
+ ) -> dict:
178
+ """
179
+ Apply ignore patterns to a dict, optionally recursively.
180
+ Supports string keys, wildcard patterns (*), and list of patterns.
181
+ Used for both ignore_input and ignore_output.
182
+
183
+ Args:
184
+ data_dict: Dictionary to filter (may contain nested dictionaries)
185
+ ignore_patterns: List of patterns to exclude (e.g., ["self", "_*", "password"])
186
+ recursive: Whether to apply patterns recursively to nested dictionaries
187
+ max_depth: Maximum recursion depth to prevent infinite loops (default: 100)
188
+ current_depth: Current recursion depth (internal use)
189
+
190
+ Returns:
191
+ Filtered dictionary with matching keys removed
192
+ """
193
+ if not isinstance(data_dict, dict):
194
+ return data_dict
195
+
196
+ # Safety check: prevent infinite loops from extremely deep nesting
197
+ if current_depth >= max_depth:
198
+ logger.warning(
199
+ f"_apply_ignore_patterns: max depth {max_depth} reached, "
200
+ f"stopping recursion to prevent infinite loop"
201
+ )
202
+ return data_dict
203
+
204
+ # If no patterns, return copy (no filtering needed, even if recursive=True)
205
+ if not ignore_patterns:
206
+ return data_dict.copy()
207
+
208
+ result = {}
209
+ for key, value in data_dict.items():
210
+ # Skip keys that match ignore patterns
211
+ if _matches_ignore_pattern(key, ignore_patterns):
212
+ continue
213
+
214
+ # Recursively process nested dictionaries if recursive=True
215
+ if recursive and isinstance(value, dict):
216
+ result[key] = _apply_ignore_patterns(
217
+ value, ignore_patterns, recursive, max_depth, current_depth + 1
218
+ )
219
+ else:
220
+ result[key] = value
221
+
222
+ return result
223
+
224
+
225
+ def _merge_with_default_ignore_patterns(
226
+ ignore_patterns: Optional[List[str]],
227
+ client: Optional[Any] = None
228
+ ) -> List[str]:
229
+ """
230
+ Merge user-provided ignore patterns with client's default ignore patterns.
231
+
232
+ Args:
233
+ ignore_patterns: Optional list of user-provided patterns
234
+ client: Optional client instance (to avoid repeated get_aiqa_client() calls)
235
+
236
+ Returns:
237
+ List of patterns including client's default ignore patterns
238
+ """
239
+ if client is None:
240
+ client = get_aiqa_client()
241
+ default_patterns = client.default_ignore_patterns
242
+
243
+ if ignore_patterns is None:
244
+ return default_patterns.copy() if default_patterns else []
245
+
246
+ # Merge patterns, avoiding duplicates
247
+ merged = list(default_patterns)
248
+ for pattern in ignore_patterns:
249
+ if pattern not in merged:
250
+ merged.append(pattern)
251
+ return merged
252
+
253
+
160
254
  def _prepare_and_filter_input(
161
255
  args: tuple,
162
256
  kwargs: dict,
@@ -165,44 +259,73 @@ def _prepare_and_filter_input(
165
259
  sig: Optional[inspect.Signature] = None,
166
260
  ) -> Any:
167
261
  """
168
- Prepare and filter input for span attributes - applies the user's filter_input and ignore_input.
169
- Converts all args to a dict using function signature when available.
262
+ Prepare and filter input for span attributes.
263
+
264
+ Process flow:
265
+ 1. Apply filter_input to args, kwargs (receives same inputs as decorated function, including self)
266
+ 2. Convert into dict ready for span.attributes.input
267
+ 3. Apply ignore_input to the dict (supports string, wildcard, and list patterns)
268
+ Client's default ignore patterns are automatically merged with ignore_input.
269
+
270
+ Args:
271
+ args: Positional arguments (including self for bound methods)
272
+ kwargs: Keyword arguments
273
+ filter_input: Optional function to filter/transform args and kwargs before conversion.
274
+ Receives *args, **kwargs with the same signature as the function being decorated,
275
+ including `self` for bound methods. This allows extracting properties from any object.
276
+ ignore_input: Optional list of keys/patterns to exclude from the final dict.
277
+ If "self" is in ignore_input, it will be removed from the final dict but filter_input
278
+ still receives it. Client's default ignore patterns are automatically merged.
279
+ sig: Optional function signature for proper arg name resolution
280
+
281
+ Returns:
282
+ Prepared input data (dict, list, or other) ready for span.attributes.input
170
283
  """
171
- # Handle "self" in ignore_input by skipping the first argument
172
- filtered_args = args
173
- filtered_kwargs = kwargs.copy() if kwargs else {}
174
- filtered_ignore_input = ignore_input
175
- filtered_sig = sig
176
- if ignore_input and "self" in ignore_input:
177
- # Remove "self" from ignore_input list (we'll handle it specially)
178
- filtered_ignore_input = [key for key in ignore_input if key != "self"]
179
- # Skip first arg if it exists (typically self for bound methods)
180
- if args:
181
- filtered_args = args[1:]
182
- # Also remove "self" from kwargs if present
183
- if "self" in filtered_kwargs:
184
- del filtered_kwargs["self"]
185
- # Adjust signature to remove "self" parameter if present
186
- # This is needed because we removed self from args, so signature binding will fail otherwise
187
- if filtered_sig is not None:
188
- params = list(filtered_sig.parameters.values())
189
- if params and params[0].name == "self":
190
- filtered_sig = filtered_sig.replace(parameters=params[1:])
191
- # turn args, kwargs into one "nice" object (now always a dict when signature is available)
192
- input_data = _prepare_input(filtered_args, filtered_kwargs, filtered_sig)
193
- if filter_input and input_data is not None:
194
- input_data = filter_input(input_data)
195
- if filtered_ignore_input and len(filtered_ignore_input) > 0:
196
- if not isinstance(input_data, dict):
197
- logger.warning(f"_prepare_and_filter_input: skip: ignore_input is set beyond 'self': {filtered_ignore_input} but input_data is not a dict: {type(input_data)}")
284
+ # Step 1: Apply filter_input to args, kwargs (same inputs as decorated function, including self)
285
+ if filter_input:
286
+ # filter_input receives the exact same args/kwargs as the decorated function
287
+ # This allows it to access self and extract properties from any object
288
+ try:
289
+ filtered_result = filter_input(*args, **kwargs)
290
+ except TypeError:
291
+ # Fallback: backward compatibility - convert to dict first
292
+ temp_dict = _prepare_input(args, kwargs, sig)
293
+ filtered_result = filter_input(temp_dict)
294
+
295
+ # Step 2: Convert filter_input result into dict ready for span.attributes.input
296
+ if isinstance(filtered_result, dict):
297
+ input_data = filtered_result
198
298
  else:
199
- for key in filtered_ignore_input:
200
- if key in input_data:
201
- del input_data[key]
202
- # Also handle case where input_data is just self (single value, not dict)
203
- # If we filtered out self and there are no remaining args/kwargs, return None
204
- if ignore_input and "self" in ignore_input and not filtered_args and not filtered_kwargs:
205
- return None
299
+ # Convert filter_input result to dict using signature
300
+ # Use original sig (not filtered) since filter_input received all args including self
301
+ input_data = _prepare_input(
302
+ (filtered_result,) if not isinstance(filtered_result, tuple) else filtered_result,
303
+ {},
304
+ sig
305
+ )
306
+ else:
307
+ # Step 2: Convert into dict ready for span.attributes.input
308
+ input_data = _prepare_input(args, kwargs, sig)
309
+
310
+ # Step 3: Apply ignore_input to the dict (removes "self" from final dict if specified)
311
+ # Merge with client's default ignore patterns
312
+ client = get_aiqa_client()
313
+ merged_ignore_input = _merge_with_default_ignore_patterns(ignore_input, client)
314
+ should_ignore_self = "self" in merged_ignore_input
315
+
316
+ if isinstance(input_data, dict):
317
+ input_data = _apply_ignore_patterns(
318
+ input_data,
319
+ merged_ignore_input,
320
+ recursive=client.ignore_recursive
321
+ )
322
+ # Handle case where we removed self and there are no remaining args/kwargs
323
+ if should_ignore_self and not input_data:
324
+ return None
325
+ elif merged_ignore_input:
326
+ # Warn if ignore patterns are set but input_data is not a dict
327
+ logger.warning(f"_prepare_and_filter_input: skip: ignore patterns are set but input_data is not a dict: {type(input_data)}")
328
+
206
329
  return input_data
207
330
 
208
331
 
@@ -211,17 +334,30 @@ def _filter_and_serialize_output(
211
334
  filter_output: Optional[Callable[[Any], Any]],
212
335
  ignore_output: Optional[List[str]],
213
336
  ) -> Any:
214
- """Filter and serialize output for span attributes."""
337
+ """
338
+ Filter and serialize output for span attributes.
339
+ Client's default ignore patterns are automatically merged with ignore_output.
340
+ """
215
341
  output_data = result
216
342
  if filter_output:
217
343
  if isinstance(output_data, dict):
218
344
  output_data = output_data.copy() # copy to provide shallow protection against the user accidentally mutating the output with filter_output
219
345
  output_data = filter_output(output_data)
220
- if ignore_output and isinstance(output_data, dict):
221
- output_data = output_data.copy()
222
- for key in ignore_output:
223
- if key in output_data:
224
- del output_data[key]
346
+
347
+ # Apply ignore_output patterns (supports key, wildcard, and list patterns)
348
+ # Merge with client's default ignore patterns
349
+ client = get_aiqa_client()
350
+ merged_ignore_output = _merge_with_default_ignore_patterns(ignore_output, client)
351
+
352
+ if isinstance(output_data, dict):
353
+ output_data = _apply_ignore_patterns(
354
+ output_data,
355
+ merged_ignore_output,
356
+ recursive=client.ignore_recursive
357
+ )
358
+ elif merged_ignore_output:
359
+ # Warn if ignore patterns are set but output_data is not a dict
360
+ logger.warning(f"_filter_and_serialize_output: skip: ignore patterns are set but output_data is not a dict: {type(output_data)}")
225
361
 
226
362
  # Serialize immediately to create immutable result (removes mutable structures)
227
363
  return serialize_for_span(output_data)
@@ -438,13 +574,24 @@ def WithTracing(
438
574
  func: The function to trace (when used as @WithTracing)
439
575
  name: Optional custom name for the span (defaults to function name)
440
576
  ignore_input: List of keys to exclude from input data when recording span attributes.
441
- Only applies when input is a dictionary. For example, use ["password", "api_key"]
442
- to exclude sensitive fields from being traced.
577
+ self is handled as "self"
578
+ Supports simple wildcards (e.g., "_*"
579
+ matches "_apple", "_fruit"). The pattern "_*" is applied by default
580
+ to filter properties starting with '_' in nested objects. For example, use
581
+ ["password", "api_key"] to exclude additional sensitive fields from being traced.
443
582
  ignore_output: List of keys to exclude from output data when recording span attributes.
444
- Only applies when output is a dictionary. Useful for excluding large or sensitive
445
- fields from traces.
446
- filter_input: Function to filter/transform input before recording
447
- filter_output: Function to filter/transform output before recording
583
+ Only applies when output is a dictionary. Supports simple wildcards (e.g., "_*"
584
+ matches "_apple", "_fruit"). The pattern "_*" is applied by default
585
+ to filter properties starting with '_' in nested objects. Useful for excluding
586
+ large or sensitive fields from traces.
587
+ filter_input: Function to filter/transform input before recording.
588
+ Receives the same arguments as the decorated function (*args, **kwargs),
589
+ including `self` for bound methods. This allows you to extract specific
590
+ properties from any kind of object. For example, to extract `dataset_id`
591
+ from `self` in a method: `filter_input=lambda self, x: {"dataset": self.dataset_id, "x": x}`.
592
+ Returns a dict or any value (will be converted to dict). Applied before ignore_input.
593
+ filter_output: Function to filter/transform output before recording.
594
+ Receives the output value and returns a filtered/transformed version.
448
595
 
449
596
  Example:
450
597
  @WithTracing
@@ -454,6 +601,17 @@ def WithTracing(
454
601
  @WithTracing(name="custom_name")
455
602
  def another_function():
456
603
  pass
604
+
605
+ # Extract properties from self in a method
606
+ class MyClass:
607
+ def __init__(self, dataset_id):
608
+ self.dataset_id = dataset_id
609
+
610
+ @WithTracing(
611
+ filter_input=lambda self, x: {"dataset": self.dataset_id, "x": x}
612
+ )
613
+ def process(self, x):
614
+ return x * 2
457
615
  """
458
616
  def decorator(fn: Callable) -> Callable:
459
617
  fn_name = name or fn.__name__ or "_"
@@ -598,7 +756,8 @@ def WithTracing(
598
756
  # This is called lazily when the function runs, not at decorator definition time
599
757
  client = get_aiqa_client()
600
758
  if not client.enabled:
601
- return await executor()
759
+ # executor() returns an async generator object, not a coroutine, so don't await it
760
+ return executor()
602
761
 
603
762
  # Get tracer after initialization (lazy)
604
763
  tracer = get_aiqa_tracer()
@@ -677,476 +836,4 @@ def WithTracing(
677
836
  return decorator(func)
678
837
 
679
838
 
680
- def set_span_attribute(attribute_name: str, attribute_value: Any) -> bool:
681
- """
682
- Set an attribute on the active span.
683
-
684
- Returns:
685
- True if attribute was set, False if no active span found
686
- """
687
- span = trace.get_current_span()
688
- if span and span.is_recording():
689
- span.set_attribute(attribute_name, serialize_for_span(attribute_value))
690
- return True
691
- return False
692
-
693
- def set_span_name(span_name: str) -> bool:
694
- """
695
- Set the name of the active span.
696
- """
697
- span = trace.get_current_span()
698
- if span and span.is_recording():
699
- span.update_name(span_name)
700
- return True
701
- return False
702
-
703
- def get_active_span() -> Optional[trace.Span]:
704
- """Get the currently active span."""
705
- return trace.get_current_span()
706
-
707
-
708
- def set_conversation_id(conversation_id: str) -> bool:
709
- """
710
- Naturally a conversation might span several traces.
711
- Set the gen_ai.conversation.id attribute on the active span.
712
- This allows you to group multiple traces together that are part of the same conversation.
713
- See https://opentelemetry.io/docs/specs/semconv/gen-ai/gen-ai-events/ for more details.
714
-
715
- Args:
716
- conversation_id: A unique identifier for the conversation (e.g., user session ID, chat ID, etc.)
717
-
718
- Returns:
719
- True if gen_ai.conversation.id was set, False if no active span found
720
-
721
- Example:
722
- from aiqa import WithTracing, set_conversation_id
723
-
724
- @WithTracing
725
- def handle_user_request(user_id: str, request: dict):
726
- # Set conversation ID to group all traces for this user session
727
- set_conversation_id(f"user_{user_id}_session_{request.get('session_id')}")
728
- # ... rest of function
729
- """
730
- return set_span_attribute("gen_ai.conversation.id", conversation_id)
731
-
732
-
733
- def set_token_usage(
734
- input_tokens: Optional[int] = None,
735
- output_tokens: Optional[int] = None,
736
- total_tokens: Optional[int] = None,
737
- ) -> bool:
738
- """
739
- Set token usage attributes on the active span using OpenTelemetry semantic conventions for gen_ai.
740
- This allows you to explicitly record token usage information.
741
- AIQA tracing will automatically detect and set token usage from standard OpenAI-like API responses.
742
- See https://opentelemetry.io/docs/specs/semconv/gen-ai/gen-ai-spans/ for more details.
743
-
744
- Args:
745
- input_tokens: Number of input tokens used (maps to gen_ai.usage.input_tokens)
746
- output_tokens: Number of output tokens generated (maps to gen_ai.usage.output_tokens)
747
- total_tokens: Total number of tokens used (maps to gen_ai.usage.total_tokens)
748
-
749
- Returns:
750
- True if at least one token usage attribute was set, False if no active span found
751
-
752
- Example:
753
- from aiqa import WithTracing, set_token_usage
754
-
755
- @WithTracing
756
- def call_llm(prompt: str):
757
- response = openai_client.chat.completions.create(...)
758
- # Explicitly set token usage
759
- set_token_usage(
760
- input_tokens=response.usage.prompt_tokens,
761
- output_tokens=response.usage.completion_tokens,
762
- total_tokens=response.usage.total_tokens
763
- )
764
- return response
765
- """
766
- span = trace.get_current_span()
767
- if not span or not span.is_recording():
768
- return False
769
-
770
- set_count = 0
771
- try:
772
- if input_tokens is not None:
773
- span.set_attribute("gen_ai.usage.input_tokens", input_tokens)
774
- set_count += 1
775
- if output_tokens is not None:
776
- span.set_attribute("gen_ai.usage.output_tokens", output_tokens)
777
- set_count += 1
778
- if total_tokens is not None:
779
- span.set_attribute("gen_ai.usage.total_tokens", total_tokens)
780
- set_count += 1
781
- except Exception as e:
782
- logger.warning(f"Failed to set token usage attributes: {e}")
783
- return False
784
-
785
- return set_count > 0
786
-
787
-
788
- def set_provider_and_model(
789
- provider: Optional[str] = None,
790
- model: Optional[str] = None,
791
- ) -> bool:
792
- """
793
- Set provider and model attributes on the active span using OpenTelemetry semantic conventions for gen_ai.
794
- This allows you to explicitly record provider and model information.
795
- AIQA tracing will automatically detect and set provider/model from standard API responses.
796
- See https://opentelemetry.io/docs/specs/semconv/gen-ai/gen-ai-spans/ for more details.
797
-
798
- Args:
799
- provider: Name of the AI provider (e.g., "openai", "anthropic", "google") (maps to gen_ai.provider.name)
800
- model: Name of the model used (e.g., "gpt-4", "claude-3-5-sonnet") (maps to gen_ai.request.model)
801
-
802
- Returns:
803
- True if at least one attribute was set, False if no active span found
804
-
805
- Example:
806
- from aiqa import WithTracing, set_provider_and_model
807
-
808
- @WithTracing
809
- def call_llm(prompt: str):
810
- response = openai_client.chat.completions.create(...)
811
- # Explicitly set provider and model
812
- set_provider_and_model(
813
- provider="openai",
814
- model=response.model
815
- )
816
- return response
817
- """
818
- span = trace.get_current_span()
819
- if not span or not span.is_recording():
820
- return False
821
-
822
- set_count = 0
823
- try:
824
- if provider is not None:
825
- span.set_attribute("gen_ai.provider.name", str(provider))
826
- set_count += 1
827
- if model is not None:
828
- span.set_attribute("gen_ai.request.model", str(model))
829
- set_count += 1
830
- except Exception as e:
831
- logger.warning(f"Failed to set provider/model attributes: {e}")
832
- return False
833
-
834
- return set_count > 0
835
-
836
-
837
- def set_component_tag(tag: str) -> None:
838
- """
839
- Set the component tag that will be added to all spans created by AIQA.
840
- This can also be set via the AIQA_COMPONENT_TAG environment variable.
841
- The component tag allows you to identify which component/system generated the spans.
842
-
843
- Note: Initialization is automatic when WithTracing is first used. You can also call
844
- get_aiqa_client() explicitly if needed.
845
- the client and load environment variables.
846
-
847
- Args:
848
- tag: A component identifier (e.g., "mynamespace.mysystem", "backend.api", etc.)
849
-
850
- Example:
851
- from aiqa import get_aiqa_client, set_component_tag, WithTracing
852
-
853
- # Initialize client (loads env vars including AIQA_COMPONENT_TAG)
854
- get_aiqa_client()
855
-
856
- # Or set component tag programmatically (overrides env var)
857
- set_component_tag("mynamespace.mysystem")
858
-
859
- @WithTracing
860
- def my_function():
861
- pass
862
- """
863
- _set_component_tag(tag)
864
-
865
-
866
- def get_active_trace_id() -> Optional[str]:
867
- """
868
- Get the current trace ID as a hexadecimal string (32 characters).
869
-
870
- Returns:
871
- The trace ID as a hex string, or None if no active span exists.
872
-
873
- Example:
874
- trace_id = get_active_trace_id()
875
- # Pass trace_id to another service/agent
876
- # e.g., include in HTTP headers, message queue metadata, etc.
877
- # Within a single thread, OpenTelemetry normally does this for you.
878
- """
879
- span = trace.get_current_span()
880
- if span and span.get_span_context().is_valid:
881
- return format(span.get_span_context().trace_id, "032x")
882
- return None
883
-
884
-
885
- def get_span_id() -> Optional[str]:
886
- """
887
- Get the current span ID as a hexadecimal string (16 characters).
888
-
889
- Returns:
890
- The span ID as a hex string, or None if no active span exists.
891
-
892
- Example:
893
- span_id = get_span_id()
894
- # Can be used to create child spans in other services
895
- """
896
- span = trace.get_current_span()
897
- if span and span.get_span_context().is_valid:
898
- return format(span.get_span_context().span_id, "016x")
899
- return None
900
-
901
-
902
- def create_span_from_trace_id(
903
- trace_id: str,
904
- parent_span_id: Optional[str] = None,
905
- span_name: str = "continued_span",
906
- ) -> trace.Span:
907
- """
908
- Create a new span that continues from an existing trace ID.
909
- This is useful for linking traces across different services or agents.
910
-
911
- Args:
912
- trace_id: The trace ID as a hexadecimal string (32 characters)
913
- parent_span_id: Optional parent span ID as a hexadecimal string (16 characters).
914
- If provided, the new span will be a child of this span.
915
- span_name: Name for the new span (default: "continued_span")
916
-
917
- Returns:
918
- A new span that continues the trace. Use it in a context manager or call end() manually.
919
-
920
- Example:
921
- # In service A: get trace ID
922
- trace_id = get_active_trace_id()
923
- span_id = get_span_id()
924
-
925
- # Send to service B (e.g., via HTTP, message queue, etc.)
926
- # ...
927
-
928
- # In service B: continue the trace
929
- with create_span_from_trace_id(trace_id, parent_span_id=span_id, span_name="service_b_operation"):
930
- # Your code here
931
- pass
932
- """
933
- try:
934
- # Parse trace ID from hex string
935
- trace_id_int = int(trace_id, 16)
936
-
937
- # Parse parent span ID if provided
938
- parent_span_id_int = None
939
- if parent_span_id:
940
- parent_span_id_int = int(parent_span_id, 16)
941
-
942
- # Create a parent span context
943
- parent_span_context = SpanContext(
944
- trace_id=trace_id_int,
945
- span_id=parent_span_id_int if parent_span_id_int else 0,
946
- is_remote=True,
947
- trace_flags=TraceFlags(0x01), # SAMPLED flag
948
- )
949
-
950
- # Create a context with this span context as the parent
951
- from opentelemetry.trace import set_span_in_context
952
- parent_context = set_span_in_context(trace.NonRecordingSpan(parent_span_context))
953
-
954
- # Ensure initialization before creating span
955
- get_aiqa_client()
956
- # Start a new span in this context (it will be a child of the parent span)
957
- tracer = get_aiqa_tracer()
958
- span = tracer.start_span(span_name, context=parent_context)
959
-
960
- # Set component tag if configured
961
- component_tag = get_component_tag()
962
- if component_tag:
963
- span.set_attribute("gen_ai.component.id", component_tag)
964
-
965
- return span
966
- except (ValueError, AttributeError) as e:
967
- logger.error(f"Error creating span from trace_id: {e}")
968
- # Ensure initialization before creating span
969
- get_aiqa_client()
970
- # Fallback: create a new span
971
- tracer = get_aiqa_tracer()
972
- span = tracer.start_span(span_name)
973
- component_tag = get_component_tag()
974
- if component_tag:
975
- span.set_attribute("gen_ai.component.id", component_tag)
976
- return span
977
-
978
-
979
- def inject_trace_context(carrier: dict) -> None:
980
- """
981
- Inject the current trace context into a carrier (e.g., HTTP headers).
982
- This allows you to pass trace context to another service.
983
-
984
- Args:
985
- carrier: Dictionary to inject trace context into (e.g., HTTP headers dict)
986
-
987
- Example:
988
- import requests
989
-
990
- headers = {}
991
- inject_trace_context(headers)
992
- response = requests.get("http://other-service/api", headers=headers)
993
- """
994
- try:
995
- inject(carrier)
996
- except Exception as e:
997
- logger.warning(f"Error injecting trace context: {e}")
998
-
999
-
1000
- def extract_trace_context(carrier: dict) -> Any:
1001
- """
1002
- Extract trace context from a carrier (e.g., HTTP headers).
1003
- Use this to continue a trace that was started in another service.
1004
-
1005
- Args:
1006
- carrier: Dictionary containing trace context (e.g., HTTP headers dict)
1007
-
1008
- Returns:
1009
- A context object that can be used with trace.use_span() or tracer.start_span()
1010
-
1011
- Example:
1012
- from opentelemetry.trace import use_span
1013
-
1014
- # Extract context from incoming request headers
1015
- ctx = extract_trace_context(request.headers)
1016
-
1017
- # Use the context to create a span
1018
- with use_span(ctx):
1019
- # Your code here
1020
- pass
1021
-
1022
- # Or create a span with the context
1023
- tracer = get_aiqa_tracer()
1024
- with tracer.start_as_current_span("operation", context=ctx):
1025
- # Your code here
1026
- pass
1027
- """
1028
- try:
1029
- return extract(carrier)
1030
- except Exception as e:
1031
- logger.warning(f"Error extracting trace context: {e}")
1032
- return None
1033
-
1034
-
1035
- def get_span(span_id: str, organisation_id: Optional[str] = None, exclude: Optional[List[str]] = None) -> Optional[dict]:
1036
- """
1037
- Get a span by its ID from the AIQA server.
1038
-
1039
- Expected usage is: re-playing a specific function call in a unit test (either a developer debugging an issue, or as part of a test suite).
1040
-
1041
- Args:
1042
- span_id: The span ID as a hexadecimal string (16 characters) or client span ID
1043
- organisation_id: Optional organisation ID. If not provided, will try to get from
1044
- AIQA_ORGANISATION_ID environment variable. The organisation is typically
1045
- extracted from the API key during authentication, but the API requires it
1046
- as a query parameter.
1047
- exclude: Optional list of fields to exclude from the span data. By default this function WILL return 'attributes' (often large).
1048
-
1049
- Returns:
1050
- The span data as a dictionary, or None if not found
1051
-
1052
- Example:
1053
- from aiqa import get_span
1054
-
1055
- span = get_span('abc123...')
1056
- if span:
1057
- print(f"Found span: {span['name']}")
1058
- my_function(**span['input'])
1059
- """
1060
- server_url = get_server_url()
1061
- api_key = get_api_key()
1062
- org_id = organisation_id or os.getenv("AIQA_ORGANISATION_ID", "")
1063
-
1064
- # Check if server_url is the default (meaning AIQA_SERVER_URL was not set)
1065
- if not os.getenv("AIQA_SERVER_URL"):
1066
- raise ValueError("AIQA_SERVER_URL is not set. Cannot retrieve span.")
1067
- if not org_id:
1068
- raise ValueError("Organisation ID is required. Provide it as parameter or set AIQA_ORGANISATION_ID environment variable.")
1069
- if not api_key:
1070
- raise ValueError("API key is required. Set AIQA_API_KEY environment variable.")
1071
-
1072
- # Try both spanId and clientSpanId queries
1073
- for query_field in ["spanId", "clientSpanId"]:
1074
- url = f"{server_url}/span"
1075
- params = {
1076
- "q": f"{query_field}:{span_id}",
1077
- "organisation": org_id,
1078
- "limit": "1",
1079
- "exclude": ",".join(exclude) if exclude else None,
1080
- "fields": "*" if not exclude else None,
1081
- }
1082
-
1083
- headers = build_headers(api_key)
1084
-
1085
- response = requests.get(url, params=params, headers=headers)
1086
- if response.status_code == 200:
1087
- result = response.json()
1088
- hits = result.get("hits", [])
1089
- if hits and len(hits) > 0:
1090
- return hits[0]
1091
- elif response.status_code == 404:
1092
- # Try next query field
1093
- continue
1094
- else:
1095
- error_text = response.text
1096
- raise ValueError(f"Failed to get span: {response.status_code} - {error_text[:500]}")
1097
- # not found
1098
- return None
1099
-
1100
-
1101
- async def submit_feedback(
1102
- trace_id: str,
1103
- thumbs_up: Optional[bool] = None,
1104
- comment: Optional[str] = None,
1105
- ) -> None:
1106
- """
1107
- Submit feedback for a trace by creating a new span with the same trace ID.
1108
- This allows you to add feedback (thumbs-up, thumbs-down, comment) to a trace after it has completed.
1109
-
1110
- Args:
1111
- trace_id: The trace ID as a hexadecimal string (32 characters)
1112
- thumbs_up: True for positive feedback, False for negative feedback, None for neutral
1113
- comment: Optional text comment
1114
-
1115
- Example:
1116
- from aiqa import submit_feedback
1117
-
1118
- # Submit positive feedback
1119
- await submit_feedback('abc123...', thumbs_up=True, comment='Great response!')
1120
-
1121
- # Submit negative feedback
1122
- await submit_feedback('abc123...', thumbs_up=False, comment='Incorrect answer')
1123
- """
1124
- if not trace_id or len(trace_id) != 32:
1125
- raise ValueError('Invalid trace ID: must be 32 hexadecimal characters')
1126
-
1127
- # Create a span for feedback with the same trace ID
1128
- span = create_span_from_trace_id(trace_id, span_name='feedback')
1129
-
1130
- try:
1131
- # Set feedback attributes
1132
- if thumbs_up is not None:
1133
- span.set_attribute('feedback.thumbs_up', thumbs_up)
1134
- span.set_attribute('feedback.type', 'positive' if thumbs_up else 'negative')
1135
- else:
1136
- span.set_attribute('feedback.type', 'neutral')
1137
-
1138
- if comment:
1139
- span.set_attribute('feedback.comment', comment)
1140
-
1141
- # Mark as feedback span
1142
- span.set_attribute('aiqa.span_type', 'feedback')
1143
-
1144
- # End the span
1145
- span.end()
1146
-
1147
- # Flush to ensure it's sent immediately
1148
- await flush_tracing()
1149
- except Exception as e:
1150
- span.end()
1151
- raise e
1152
839