ray-llm

1
from contextlib import contextmanager
2
from typing import Generator, Optional
3

4
from opentelemetry import context
5
from opentelemetry.context.context import Context
6

7
# Sentinel value representing "use the current context"
8
CURRENT_CONTEXT = Context()
9

10

11
@contextmanager
12
def use_context(
13
    parent_context: Optional[Context] = CURRENT_CONTEXT,
14
) -> Generator[None, None, None]:
15
    if parent_context is CURRENT_CONTEXT:
16
        # Do nothing, default behavior of opentelemetry is to use the current context
17
        yield
18
    else:
19
        new_context = parent_context if parent_context is not None else Context()
20
        token = context.attach(new_context)
21
        try:
22
            yield
23
        finally:
24
            context.detach(token)
25