lmnr 0.6.21__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.
Files changed (28) hide show
  1. lmnr/__init__.py +0 -4
  2. lmnr/opentelemetry_lib/decorators/__init__.py +38 -28
  3. lmnr/opentelemetry_lib/opentelemetry/instrumentation/anthropic/__init__.py +6 -2
  4. lmnr/opentelemetry_lib/opentelemetry/instrumentation/google_genai/__init__.py +4 -0
  5. lmnr/opentelemetry_lib/opentelemetry/instrumentation/groq/__init__.py +3 -0
  6. lmnr/opentelemetry_lib/opentelemetry/instrumentation/langgraph/__init__.py +16 -16
  7. lmnr/opentelemetry_lib/opentelemetry/instrumentation/openai/shared/chat_wrappers.py +3 -0
  8. lmnr/opentelemetry_lib/opentelemetry/instrumentation/openai/shared/completion_wrappers.py +3 -0
  9. lmnr/opentelemetry_lib/opentelemetry/instrumentation/openai/v1/assistant_wrappers.py +3 -0
  10. lmnr/opentelemetry_lib/opentelemetry/instrumentation/openai/v1/responses_wrappers.py +7 -0
  11. lmnr/opentelemetry_lib/opentelemetry/instrumentation/threading/__init__.py +190 -0
  12. lmnr/opentelemetry_lib/tracing/__init__.py +89 -1
  13. lmnr/opentelemetry_lib/tracing/context.py +109 -0
  14. lmnr/opentelemetry_lib/tracing/processor.py +5 -6
  15. lmnr/opentelemetry_lib/tracing/tracer.py +29 -0
  16. lmnr/sdk/browser/browser_use_otel.py +5 -5
  17. lmnr/sdk/browser/patchright_otel.py +14 -0
  18. lmnr/sdk/browser/playwright_otel.py +32 -6
  19. lmnr/sdk/browser/pw_utils.py +78 -6
  20. lmnr/sdk/client/asynchronous/resources/browser_events.py +1 -0
  21. lmnr/sdk/laminar.py +109 -164
  22. lmnr/sdk/types.py +0 -6
  23. lmnr/version.py +1 -1
  24. {lmnr-0.6.21.dist-info → lmnr-0.7.0.dist-info}/METADATA +3 -2
  25. {lmnr-0.6.21.dist-info → lmnr-0.7.0.dist-info}/RECORD +27 -26
  26. {lmnr-0.6.21.dist-info → lmnr-0.7.0.dist-info}/WHEEL +1 -1
  27. lmnr/opentelemetry_lib/tracing/context_properties.py +0 -65
  28. {lmnr-0.6.21.dist-info → lmnr-0.7.0.dist-info}/entry_points.txt +0 -0
lmnr/__init__.py CHANGED
@@ -9,7 +9,6 @@ from .sdk.types import (
9
9
  HumanEvaluator,
10
10
  RunAgentResponseChunk,
11
11
  StepChunkContent,
12
- TracingLevel,
13
12
  )
14
13
  from .sdk.decorators import observe
15
14
  from .sdk.types import LaminarSpanContext
@@ -18,7 +17,6 @@ from .opentelemetry_lib.tracing.attributes import Attributes
18
17
  from .opentelemetry_lib.tracing.instruments import Instruments
19
18
  from .opentelemetry_lib.tracing.processor import LaminarSpanProcessor
20
19
  from .opentelemetry_lib.tracing.tracer import get_laminar_tracer_provider, get_tracer
21
- from opentelemetry.trace import use_span
22
20
 
23
21
  __all__ = [
24
22
  "AgentOutput",
@@ -36,10 +34,8 @@ __all__ = [
36
34
  "LaminarSpanProcessor",
37
35
  "RunAgentResponseChunk",
38
36
  "StepChunkContent",
39
- "TracingLevel",
40
37
  "get_laminar_tracer_provider",
41
38
  "get_tracer",
42
39
  "evaluate",
43
40
  "observe",
44
- "use_span",
45
41
  ]
@@ -5,13 +5,12 @@ import orjson
5
5
  import types
6
6
  from typing import Any, AsyncGenerator, Callable, Generator, Literal
7
7
 
8
- from opentelemetry import trace
9
8
  from opentelemetry import context as context_api
10
9
  from opentelemetry.trace import Span
11
10
 
12
11
  from lmnr.sdk.utils import get_input_from_func_args, is_method
13
12
  from lmnr.opentelemetry_lib import MAX_MANUAL_SPAN_PAYLOAD_SIZE
14
- from lmnr.opentelemetry_lib.tracing.tracer import get_tracer
13
+ from lmnr.opentelemetry_lib.tracing.tracer import get_tracer_with_context
15
14
  from lmnr.opentelemetry_lib.tracing.attributes import (
16
15
  ASSOCIATION_PROPERTIES,
17
16
  SPAN_INPUT,
@@ -37,6 +36,7 @@ def default_json(o):
37
36
  try:
38
37
  return str(o)
39
38
  except Exception:
39
+ logger.debug("Failed to serialize data to JSON, inner type: %s", type(o))
40
40
  pass
41
41
  return DEFAULT_PLACEHOLDER
42
42
 
@@ -61,8 +61,13 @@ def _setup_span(
61
61
  span_name: str, span_type: str, association_properties: dict[str, Any] | None
62
62
  ):
63
63
  """Set up a span with the given name, type, and association properties."""
64
- with get_tracer() as tracer:
65
- span = tracer.start_span(span_name, attributes={SPAN_TYPE: span_type})
64
+ with get_tracer_with_context() as (tracer, isolated_context):
65
+ # Create span in isolated context
66
+ span = tracer.start_span(
67
+ span_name,
68
+ context=isolated_context,
69
+ attributes={SPAN_TYPE: span_type},
70
+ )
66
71
 
67
72
  if association_properties is not None:
68
73
  for key, value in association_properties.items():
@@ -148,10 +153,10 @@ def _process_output(
148
153
  pass
149
154
 
150
155
 
151
- def _cleanup_span(span: Span, ctx_token):
156
+ def _cleanup_span(span: Span, wrapper: TracerWrapper):
152
157
  """Clean up span and context."""
153
158
  span.end()
154
- context_api.detach(ctx_token)
159
+ wrapper.pop_span_context()
155
160
 
156
161
 
157
162
  def observe_base(
@@ -171,10 +176,11 @@ def observe_base(
171
176
  return fn(*args, **kwargs)
172
177
 
173
178
  span_name = name or fn.__name__
179
+ wrapper = TracerWrapper()
174
180
 
175
181
  span = _setup_span(span_name, span_type, association_properties)
176
- ctx = trace.set_span_in_context(span, context_api.get_current())
177
- ctx_token = context_api.attach(ctx)
182
+ new_context = wrapper.push_span_context(span)
183
+ ctx_token = context_api.attach(new_context)
178
184
 
179
185
  _process_input(
180
186
  span, fn, args, kwargs, ignore_input, ignore_inputs, input_formatter
@@ -184,8 +190,11 @@ def observe_base(
184
190
  res = fn(*args, **kwargs)
185
191
  except Exception as e:
186
192
  _process_exception(span, e)
187
- _cleanup_span(span, ctx_token)
193
+ _cleanup_span(span, wrapper)
188
194
  raise e
195
+ finally:
196
+ # Always restore global context
197
+ context_api.detach(ctx_token)
189
198
 
190
199
  # span will be ended in the generator
191
200
  if isinstance(res, types.GeneratorType):
@@ -201,7 +210,7 @@ def observe_base(
201
210
  return _ahandle_generator(span, ctx_token, res)
202
211
 
203
212
  _process_output(span, res, ignore_output, output_formatter)
204
- _cleanup_span(span, ctx_token)
213
+ _cleanup_span(span, wrapper)
205
214
  return res
206
215
 
207
216
  return wrap
@@ -227,10 +236,11 @@ def async_observe_base(
227
236
  return await fn(*args, **kwargs)
228
237
 
229
238
  span_name = name or fn.__name__
239
+ wrapper = TracerWrapper()
230
240
 
231
241
  span = _setup_span(span_name, span_type, association_properties)
232
- ctx = trace.set_span_in_context(span, context_api.get_current())
233
- ctx_token = context_api.attach(ctx)
242
+ new_context = wrapper.push_span_context(span)
243
+ ctx_token = context_api.attach(new_context)
234
244
 
235
245
  _process_input(
236
246
  span, fn, args, kwargs, ignore_input, ignore_inputs, input_formatter
@@ -240,8 +250,11 @@ def async_observe_base(
240
250
  res = await fn(*args, **kwargs)
241
251
  except Exception as e:
242
252
  _process_exception(span, e)
243
- _cleanup_span(span, ctx_token)
253
+ _cleanup_span(span, wrapper)
244
254
  raise e
255
+ finally:
256
+ # Always restore global context
257
+ context_api.detach(ctx_token)
245
258
 
246
259
  # span will be ended in the generator
247
260
  if isinstance(res, types.AsyncGeneratorType):
@@ -250,7 +263,7 @@ def async_observe_base(
250
263
  return await _ahandle_generator(span, ctx_token, res)
251
264
 
252
265
  _process_output(span, res, ignore_output, output_formatter)
253
- _cleanup_span(span, ctx_token)
266
+ _cleanup_span(span, wrapper)
254
267
  return res
255
268
 
256
269
  return wrap
@@ -258,22 +271,19 @@ def async_observe_base(
258
271
  return decorate
259
272
 
260
273
 
261
- def _handle_generator(span: Span, ctx_token, res: Generator[Any, Any, Any]):
262
- yield from res
263
-
264
- span.end()
265
- if ctx_token is not None:
266
- context_api.detach(ctx_token)
267
-
274
+ def _handle_generator(span: Span, wrapper: TracerWrapper, res: Generator):
275
+ try:
276
+ yield from res
277
+ finally:
278
+ _cleanup_span(span, wrapper)
268
279
 
269
- async def _ahandle_generator(span: Span, ctx_token, res: AsyncGenerator[Any, Any]):
270
- # async with contextlib.aclosing(res) as closing_gen:
271
- async for part in res:
272
- yield part
273
280
 
274
- span.end()
275
- if ctx_token is not None:
276
- context_api.detach(ctx_token)
281
+ async def _ahandle_generator(span: Span, wrapper: TracerWrapper, res: AsyncGenerator):
282
+ try:
283
+ async for part in res:
284
+ yield part
285
+ finally:
286
+ _cleanup_span(span, wrapper)
277
287
 
278
288
 
279
289
  def _process_exception(span: Span, e: Exception):
@@ -30,6 +30,8 @@ from .utils import (
30
30
  should_emit_events,
31
31
  )
32
32
  from .version import __version__
33
+
34
+ from lmnr.opentelemetry_lib.tracing.context import get_current_context
33
35
  from opentelemetry.instrumentation.instrumentor import BaseInstrumentor
34
36
  from opentelemetry.instrumentation.utils import _SUPPRESS_INSTRUMENTATION_KEY, unwrap
35
37
  from opentelemetry.metrics import Counter, Histogram, Meter, get_meter
@@ -396,9 +398,10 @@ def _wrap(
396
398
  name,
397
399
  kind=SpanKind.CLIENT,
398
400
  attributes={
399
- SpanAttributes.LLM_SYSTEM: "Anthropic",
401
+ SpanAttributes.LLM_SYSTEM: "anthropic",
400
402
  SpanAttributes.LLM_REQUEST_TYPE: LLMRequestTypeValues.COMPLETION.value,
401
403
  },
404
+ context=get_current_context(),
402
405
  )
403
406
 
404
407
  _handle_input(span, event_logger, kwargs)
@@ -493,9 +496,10 @@ async def _awrap(
493
496
  name,
494
497
  kind=SpanKind.CLIENT,
495
498
  attributes={
496
- SpanAttributes.LLM_SYSTEM: "Anthropic",
499
+ SpanAttributes.LLM_SYSTEM: "anthropic",
497
500
  SpanAttributes.LLM_REQUEST_TYPE: LLMRequestTypeValues.COMPLETION.value,
498
501
  },
502
+ context=get_current_context(),
499
503
  )
500
504
  await _ahandle_input(span, event_logger, kwargs)
501
505
 
@@ -8,6 +8,8 @@ from typing import AsyncGenerator, Callable, Collection, Generator
8
8
 
9
9
  from google.genai import types
10
10
 
11
+ from lmnr.opentelemetry_lib.tracing.context import get_current_context
12
+
11
13
  from .config import (
12
14
  Config,
13
15
  )
@@ -474,6 +476,7 @@ def _wrap(tracer: Tracer, to_wrap, wrapped, instance, args, kwargs):
474
476
  SpanAttributes.LLM_SYSTEM: "gemini",
475
477
  SpanAttributes.LLM_REQUEST_TYPE: LLMRequestTypeValues.COMPLETION.value,
476
478
  },
479
+ context=get_current_context(),
477
480
  )
478
481
 
479
482
  if span.is_recording():
@@ -509,6 +512,7 @@ async def _awrap(tracer: Tracer, to_wrap, wrapped, instance, args, kwargs):
509
512
  SpanAttributes.LLM_SYSTEM: "gemini",
510
513
  SpanAttributes.LLM_REQUEST_TYPE: LLMRequestTypeValues.COMPLETION.value,
511
514
  },
515
+ context=get_current_context(),
512
516
  )
513
517
 
514
518
  if span.is_recording():
@@ -27,6 +27,7 @@ from .utils import (
27
27
  should_emit_events,
28
28
  )
29
29
  from .version import __version__
30
+ from lmnr.opentelemetry_lib.tracing.context import get_current_context
30
31
  from opentelemetry.instrumentation.instrumentor import BaseInstrumentor
31
32
  from opentelemetry.instrumentation.utils import _SUPPRESS_INSTRUMENTATION_KEY, unwrap
32
33
  from opentelemetry.metrics import Counter, Histogram, Meter, get_meter
@@ -245,6 +246,7 @@ def _wrap(
245
246
  SpanAttributes.LLM_SYSTEM: "Groq",
246
247
  SpanAttributes.LLM_REQUEST_TYPE: LLMRequestTypeValues.COMPLETION.value,
247
248
  },
249
+ context=get_current_context(),
248
250
  )
249
251
 
250
252
  _handle_input(span, kwargs, event_logger)
@@ -327,6 +329,7 @@ async def _awrap(
327
329
  SpanAttributes.LLM_SYSTEM: "Groq",
328
330
  SpanAttributes.LLM_REQUEST_TYPE: LLMRequestTypeValues.COMPLETION.value,
329
331
  },
332
+ context=get_current_context(),
330
333
  )
331
334
 
332
335
  _handle_input(span, kwargs, event_logger)
@@ -12,10 +12,7 @@ from langchain_core.runnables.graph import Graph
12
12
  from opentelemetry.trace import Tracer
13
13
  from wrapt import wrap_function_wrapper
14
14
  from opentelemetry.trace import get_tracer
15
-
16
- from lmnr.opentelemetry_lib.tracing.context_properties import (
17
- update_association_properties,
18
- )
15
+ from opentelemetry.context import get_value, attach, set_value
19
16
 
20
17
  from opentelemetry.instrumentation.instrumentor import BaseInstrumentor
21
18
  from opentelemetry.instrumentation.utils import unwrap
@@ -45,12 +42,13 @@ def wrap_pregel_stream(tracer: Tracer, to_wrap, wrapped, instance, args, kwargs)
45
42
  }
46
43
  for edge in graph.edges
47
44
  ]
48
- update_association_properties(
49
- {
50
- "langgraph.edges": json.dumps(edges),
51
- "langgraph.nodes": json.dumps(nodes),
52
- },
53
- )
45
+ d = {
46
+ "langgraph.edges": json.dumps(edges),
47
+ "langgraph.nodes": json.dumps(nodes),
48
+ }
49
+ association_properties = get_value("lmnr.langgraph.graph") or {}
50
+ association_properties.update(d)
51
+ attach(set_value("lmnr.langgraph.graph", association_properties))
54
52
  return wrapped(*args, **kwargs)
55
53
 
56
54
 
@@ -75,12 +73,14 @@ async def async_wrap_pregel_stream(
75
73
  }
76
74
  for edge in graph.edges
77
75
  ]
78
- update_association_properties(
79
- {
80
- "langgraph.edges": json.dumps(edges),
81
- "langgraph.nodes": json.dumps(nodes),
82
- },
83
- )
76
+
77
+ d = {
78
+ "langgraph.edges": json.dumps(edges),
79
+ "langgraph.nodes": json.dumps(nodes),
80
+ }
81
+ association_properties = get_value("lmnr.langgraph.graph") or {}
82
+ association_properties.update(d)
83
+ attach(set_value("lmnr.langgraph.graph", association_properties))
84
84
 
85
85
  async for item in wrapped(*args, **kwargs):
86
86
  yield item
@@ -39,6 +39,7 @@ from ..utils import (
39
39
  should_emit_events,
40
40
  should_send_prompts,
41
41
  )
42
+ from lmnr.opentelemetry_lib.tracing.context import get_current_context
42
43
  from opentelemetry.instrumentation.utils import _SUPPRESS_INSTRUMENTATION_KEY
43
44
  from opentelemetry.metrics import Counter, Histogram
44
45
  from opentelemetry.semconv.attributes.error_attributes import ERROR_TYPE
@@ -87,6 +88,7 @@ def chat_wrapper(
87
88
  SPAN_NAME,
88
89
  kind=SpanKind.CLIENT,
89
90
  attributes={SpanAttributes.LLM_REQUEST_TYPE: LLM_REQUEST_TYPE.value},
91
+ context=get_current_context(),
90
92
  )
91
93
 
92
94
  run_async(_handle_request(span, kwargs, instance))
@@ -184,6 +186,7 @@ async def achat_wrapper(
184
186
  SPAN_NAME,
185
187
  kind=SpanKind.CLIENT,
186
188
  attributes={SpanAttributes.LLM_REQUEST_TYPE: LLM_REQUEST_TYPE.value},
189
+ context=get_current_context(),
187
190
  )
188
191
 
189
192
  await _handle_request(span, kwargs, instance)
@@ -27,6 +27,7 @@ from ..utils import (
27
27
  should_emit_events,
28
28
  should_send_prompts,
29
29
  )
30
+ from lmnr.opentelemetry_lib.tracing.context import get_current_context
30
31
  from opentelemetry.instrumentation.utils import _SUPPRESS_INSTRUMENTATION_KEY
31
32
  from opentelemetry.semconv.attributes.error_attributes import ERROR_TYPE
32
33
  from opentelemetry.semconv_ai import (
@@ -55,6 +56,7 @@ def completion_wrapper(tracer, wrapped, instance, args, kwargs):
55
56
  SPAN_NAME,
56
57
  kind=SpanKind.CLIENT,
57
58
  attributes={SpanAttributes.LLM_REQUEST_TYPE: LLM_REQUEST_TYPE.value},
59
+ context=get_current_context(),
58
60
  )
59
61
 
60
62
  _handle_request(span, kwargs, instance)
@@ -89,6 +91,7 @@ async def acompletion_wrapper(tracer, wrapped, instance, args, kwargs):
89
91
  name=SPAN_NAME,
90
92
  kind=SpanKind.CLIENT,
91
93
  attributes={SpanAttributes.LLM_REQUEST_TYPE: LLM_REQUEST_TYPE.value},
94
+ context=get_current_context(),
92
95
  )
93
96
 
94
97
  _handle_request(span, kwargs, instance)
@@ -17,6 +17,7 @@ from ..utils import (
17
17
  dont_throw,
18
18
  should_emit_events,
19
19
  )
20
+ from lmnr.opentelemetry_lib.tracing.context import get_current_context
20
21
  from opentelemetry.instrumentation.utils import _SUPPRESS_INSTRUMENTATION_KEY
21
22
  from opentelemetry.semconv.attributes.error_attributes import ERROR_TYPE
22
23
  from opentelemetry.semconv_ai import LLMRequestTypeValues, SpanAttributes
@@ -126,6 +127,7 @@ def messages_list_wrapper(tracer, wrapped, instance, args, kwargs):
126
127
  kind=SpanKind.CLIENT,
127
128
  attributes={SpanAttributes.LLM_REQUEST_TYPE: LLMRequestTypeValues.CHAT.value},
128
129
  start_time=run.get("start_time"),
130
+ context=get_current_context(),
129
131
  )
130
132
 
131
133
  if exception := run.get("exception"):
@@ -250,6 +252,7 @@ def runs_create_and_stream_wrapper(tracer, wrapped, instance, args, kwargs):
250
252
  "openai.assistant.run_stream",
251
253
  kind=SpanKind.CLIENT,
252
254
  attributes={SpanAttributes.LLM_REQUEST_TYPE: LLMRequestTypeValues.CHAT.value},
255
+ context=get_current_context(),
253
256
  )
254
257
 
255
258
  i = 0
@@ -36,6 +36,7 @@ except ImportError:
36
36
  ResponseOutputMessageParam = Dict[str, Any]
37
37
  RESPONSES_AVAILABLE = False
38
38
 
39
+ from lmnr.opentelemetry_lib.tracing.context import get_current_context
39
40
  from openai._legacy_response import LegacyAPIResponse
40
41
  from opentelemetry import context as context_api
41
42
  from opentelemetry.instrumentation.utils import _SUPPRESS_INSTRUMENTATION_KEY
@@ -429,6 +430,7 @@ def responses_get_or_create_wrapper(tracer: Tracer, wrapped, instance, args, kwa
429
430
  start_time=(
430
431
  start_time if traced_data is None else int(traced_data.start_time)
431
432
  ),
433
+ context=get_current_context(),
432
434
  )
433
435
  span.set_attribute(ERROR_TYPE, e.__class__.__name__)
434
436
  span.record_exception(e)
@@ -472,6 +474,7 @@ def responses_get_or_create_wrapper(tracer: Tracer, wrapped, instance, args, kwa
472
474
  SPAN_NAME,
473
475
  kind=SpanKind.CLIENT,
474
476
  start_time=int(traced_data.start_time),
477
+ context=get_current_context(),
475
478
  )
476
479
  set_data_attributes(traced_data, span)
477
480
  span.end()
@@ -523,6 +526,7 @@ async def async_responses_get_or_create_wrapper(
523
526
  start_time=(
524
527
  start_time if traced_data is None else int(traced_data.start_time)
525
528
  ),
529
+ context=get_current_context(),
526
530
  )
527
531
  span.set_attribute(ERROR_TYPE, e.__class__.__name__)
528
532
  span.record_exception(e)
@@ -566,6 +570,7 @@ async def async_responses_get_or_create_wrapper(
566
570
  SPAN_NAME,
567
571
  kind=SpanKind.CLIENT,
568
572
  start_time=int(traced_data.start_time),
573
+ context=get_current_context(),
569
574
  )
570
575
  set_data_attributes(traced_data, span)
571
576
  span.end()
@@ -590,6 +595,7 @@ def responses_cancel_wrapper(tracer: Tracer, wrapped, instance, args, kwargs):
590
595
  kind=SpanKind.CLIENT,
591
596
  start_time=existing_data.start_time,
592
597
  record_exception=True,
598
+ context=get_current_context(),
593
599
  )
594
600
  span.record_exception(Exception("Response cancelled"))
595
601
  set_data_attributes(existing_data, span)
@@ -616,6 +622,7 @@ async def async_responses_cancel_wrapper(
616
622
  kind=SpanKind.CLIENT,
617
623
  start_time=existing_data.start_time,
618
624
  record_exception=True,
625
+ context=get_current_context(),
619
626
  )
620
627
  span.record_exception(Exception("Response cancelled"))
621
628
  set_data_attributes(existing_data, span)
@@ -0,0 +1,190 @@
1
+ # Copyright The OpenTelemetry Authors
2
+ #
3
+ # Licensed under the Apache License, Version 2.0 (the "License");
4
+ # you may not use this file except in compliance with the License.
5
+ # You may obtain a copy of the License at
6
+ #
7
+ # http://www.apache.org/licenses/LICENSE-2.0
8
+ #
9
+ # Unless required by applicable law or agreed to in writing, software
10
+ # distributed under the License is distributed on an "AS IS" BASIS,
11
+ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12
+ # See the License for the specific language governing permissions and
13
+ # limitations under the License.
14
+ """
15
+ Instrument threading to propagate OpenTelemetry context.
16
+
17
+ Copied from opentelemetry-instrumentation-threading at commit:
18
+ ad2fe813abb2ab0b6e25bedeebef5041ca3189f7
19
+ https://github.com/open-telemetry/opentelemetry-python-contrib/blob/ad2fe813abb2ab0b6e25bedeebef5041ca3189f7/instrumentation/opentelemetry-instrumentation-threading/src/opentelemetry/instrumentation/threading/__init__.py
20
+
21
+ Modified to use the Laminar isolated context.
22
+
23
+ Usage
24
+ -----
25
+
26
+ .. code-block:: python
27
+
28
+ from opentelemetry.instrumentation.threading import ThreadingInstrumentor
29
+
30
+ ThreadingInstrumentor().instrument()
31
+
32
+ This library provides instrumentation for the `threading` module to ensure that
33
+ the OpenTelemetry context is propagated across threads. It is important to note
34
+ that this instrumentation does not produce any telemetry data on its own. It
35
+ merely ensures that the context is correctly propagated when threads are used.
36
+
37
+
38
+ When instrumented, new threads created using threading.Thread, threading.Timer,
39
+ or within futures.ThreadPoolExecutor will have the current OpenTelemetry
40
+ context attached, and this context will be re-activated in the thread's
41
+ run method or the executor's worker thread."
42
+ """
43
+
44
+ from __future__ import annotations
45
+
46
+ import threading
47
+ from concurrent import futures
48
+ from typing import TYPE_CHECKING, Any, Callable, Collection
49
+
50
+ from wrapt import (
51
+ wrap_function_wrapper, # type: ignore[reportUnknownVariableType]
52
+ )
53
+
54
+ from lmnr.opentelemetry_lib.tracing.context import (
55
+ get_current_context,
56
+ attach_context,
57
+ detach_context,
58
+ )
59
+ from opentelemetry import context
60
+ from opentelemetry.instrumentation.instrumentor import BaseInstrumentor
61
+ from opentelemetry.instrumentation.utils import unwrap
62
+
63
+ _instruments = ()
64
+
65
+ if TYPE_CHECKING:
66
+ from typing import Protocol, TypeVar
67
+
68
+ R = TypeVar("R")
69
+
70
+ class HasOtelContext(Protocol):
71
+ _otel_context: context.Context
72
+
73
+
74
+ class ThreadingInstrumentor(BaseInstrumentor):
75
+ __WRAPPER_START_METHOD = "start"
76
+ __WRAPPER_RUN_METHOD = "run"
77
+ __WRAPPER_SUBMIT_METHOD = "submit"
78
+
79
+ def instrumentation_dependencies(self) -> Collection[str]:
80
+ return _instruments
81
+
82
+ def _instrument(self, **kwargs: Any):
83
+ self._instrument_thread()
84
+ self._instrument_timer()
85
+ self._instrument_thread_pool()
86
+
87
+ def _uninstrument(self, **kwargs: Any):
88
+ self._uninstrument_thread()
89
+ self._uninstrument_timer()
90
+ self._uninstrument_thread_pool()
91
+
92
+ @staticmethod
93
+ def _instrument_thread():
94
+ wrap_function_wrapper(
95
+ threading.Thread,
96
+ ThreadingInstrumentor.__WRAPPER_START_METHOD,
97
+ ThreadingInstrumentor.__wrap_threading_start,
98
+ )
99
+ wrap_function_wrapper(
100
+ threading.Thread,
101
+ ThreadingInstrumentor.__WRAPPER_RUN_METHOD,
102
+ ThreadingInstrumentor.__wrap_threading_run,
103
+ )
104
+
105
+ @staticmethod
106
+ def _instrument_timer():
107
+ wrap_function_wrapper(
108
+ threading.Timer,
109
+ ThreadingInstrumentor.__WRAPPER_START_METHOD,
110
+ ThreadingInstrumentor.__wrap_threading_start,
111
+ )
112
+ wrap_function_wrapper(
113
+ threading.Timer,
114
+ ThreadingInstrumentor.__WRAPPER_RUN_METHOD,
115
+ ThreadingInstrumentor.__wrap_threading_run,
116
+ )
117
+
118
+ @staticmethod
119
+ def _instrument_thread_pool():
120
+ wrap_function_wrapper(
121
+ futures.ThreadPoolExecutor,
122
+ ThreadingInstrumentor.__WRAPPER_SUBMIT_METHOD,
123
+ ThreadingInstrumentor.__wrap_thread_pool_submit,
124
+ )
125
+
126
+ @staticmethod
127
+ def _uninstrument_thread():
128
+ unwrap(threading.Thread, ThreadingInstrumentor.__WRAPPER_START_METHOD)
129
+ unwrap(threading.Thread, ThreadingInstrumentor.__WRAPPER_RUN_METHOD)
130
+
131
+ @staticmethod
132
+ def _uninstrument_timer():
133
+ unwrap(threading.Timer, ThreadingInstrumentor.__WRAPPER_START_METHOD)
134
+ unwrap(threading.Timer, ThreadingInstrumentor.__WRAPPER_RUN_METHOD)
135
+
136
+ @staticmethod
137
+ def _uninstrument_thread_pool():
138
+ unwrap(
139
+ futures.ThreadPoolExecutor,
140
+ ThreadingInstrumentor.__WRAPPER_SUBMIT_METHOD,
141
+ )
142
+
143
+ @staticmethod
144
+ def __wrap_threading_start(
145
+ call_wrapped: Callable[[], None],
146
+ instance: HasOtelContext,
147
+ args: tuple[()],
148
+ kwargs: dict[str, Any],
149
+ ) -> None:
150
+ instance._otel_context = get_current_context()
151
+ return call_wrapped(*args, **kwargs)
152
+
153
+ @staticmethod
154
+ def __wrap_threading_run(
155
+ call_wrapped: Callable[..., R],
156
+ instance: HasOtelContext,
157
+ args: tuple[Any, ...],
158
+ kwargs: dict[str, Any],
159
+ ) -> R:
160
+ token = None
161
+ try:
162
+ token = attach_context(instance._otel_context)
163
+ return call_wrapped(*args, **kwargs)
164
+ finally:
165
+ if token is not None:
166
+ detach_context(token)
167
+
168
+ @staticmethod
169
+ def __wrap_thread_pool_submit(
170
+ call_wrapped: Callable[..., R],
171
+ instance: futures.ThreadPoolExecutor,
172
+ args: tuple[Callable[..., Any], ...],
173
+ kwargs: dict[str, Any],
174
+ ) -> R:
175
+ # obtain the original function and wrapped kwargs
176
+ original_func = args[0]
177
+ otel_context = get_current_context()
178
+
179
+ def wrapped_func(*func_args: Any, **func_kwargs: Any) -> R:
180
+ token = None
181
+ try:
182
+ token = attach_context(otel_context)
183
+ return original_func(*func_args, **func_kwargs)
184
+ finally:
185
+ if token is not None:
186
+ detach_context(token)
187
+
188
+ # replace the original function with the wrapped function
189
+ new_args: tuple[Callable[..., Any], ...] = (wrapped_func,) + args[1:]
190
+ return call_wrapped(*new_args, **kwargs)