ray-llm

1
import contextvars
2
import weakref
3
from contextlib import contextmanager
4

5
from fastapi.datastructures import State
6

7
from rayllm.backend.observability.tracing import baggage
8

9
# Fast api state context.
10
# This may include secrets
11
_fastapi_state_context: contextvars.ContextVar[
12
    weakref.ReferenceType[State]
13
] = contextvars.ContextVar("aviary_fastapi_state")
14

15

16
def set(**kwargs):
17
    return baggage.baggage(kwargs)
18

19

20
def get(key: str):
21
    return baggage.get(key)
22

23

24
@contextmanager
25
def set_fastapi_state(request_state: State):
26
    # Hold a weakref to make sure that we clean up any state once Fastapi request object is garbage collected.
27
    ref = weakref.ref(request_state)
28
    token = _fastapi_state_context.set(ref)
29
    try:
30
        yield
31
    finally:
32
        _fastapi_state_context.reset(token)
33

34

35
def get_fastapi_state():
36
    ctx = _fastapi_state_context.get(None)
37
    if ctx:
38
        return ctx()
39

40

41
def maybe_get_string_field(field: str):
42
    state = get_fastapi_state()
43
    val = getattr(state, field, None)
44
    return val if isinstance(val, str) else None
45