참고소스 수정본

참고/guardrails-main/guardrails/stores/context.py

@@ -0,0 +1,48 @@
+from contextvars import ContextVar, copy_context
+from typing import Any, Dict, Literal, Union
+
+
+GUARD_NAME_KEY: Literal["gr.reserved.guard.name"] = "gr.reserved.guard.name"
+DOCUMENT_STORE_KEY: Literal["gr.reserved.document_store"] = "gr.reserved.document_store"
+CALL_KWARGS_KEY: Literal["gr.reserved.call_kwargs"] = "gr.reserved.call_kwargs"
+
+
+def set_guard_name(guard_name: str) -> None:
+    set_context_var(GUARD_NAME_KEY, guard_name)
+
+
+def get_guard_name() -> str:
+    return get_context_var(GUARD_NAME_KEY) or ""
+
+
+def set_call_kwargs(kwargs: Dict[str, Any]) -> None:
+    set_context_var(CALL_KWARGS_KEY, kwargs)
+
+
+def get_call_kwargs() -> Dict[str, Any]:
+    return get_context_var(CALL_KWARGS_KEY) or {}
+
+
+def get_call_kwarg(kwarg_key: str) -> Union[Any, None]:
+    kwargs = get_call_kwargs()
+    return kwargs.get(kwarg_key)
+
+
+def _get_contextvar(key):
+    context = copy_context()
+    context_var = None
+    for c_key in context.keys():
+        if c_key.name == key:
+            context_var = c_key
+            break
+    return context_var
+
+
+def set_context_var(key, value):
+    context_var = _get_contextvar(key) or ContextVar(key)
+    context_var.set(value)
+
+
+def get_context_var(key):
+    context_var = _get_contextvar(key)
+    return context_var.get() if context_var else None