ray-llm

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

4
from opentelemetry.baggage import get_baggage, set_baggage
5
from opentelemetry.context import get_current
6
from opentelemetry.util.types import AttributeValue
7

8
from rayllm.backend.observability.tracing.context import use_context
9

10

11
@contextmanager
12
def baggage(attributes: Mapping[str, AttributeValue]) -> Generator[None, None, None]:
13
    """
14
    A context manager that sets multiple baggage values on the active context.
15
    This can be combined with the BaggageSpanProcessor to add the baggage values
16
    to the span as attributes.
17

18
    Example:
19

20
    with baggage({"foo": "bar"}):
21
        do_something()
22

23
    """
24
    ctx = get_current()
25
    for key, value in attributes.items():
26
        ctx = set_baggage(key, value, context=ctx)
27

28
    with use_context(ctx):
29
        yield
30

31

32
def get(key: str):
33
    return get_baggage(key)
34