mirror of
https://github.com/infiniflow/ragflow.git
synced 2026-07-08 12:24:48 +08:00
72 lines
2.9 KiB
Python
72 lines
2.9 KiB
Python
#
|
|
# Copyright 2025 The InfiniFlow Authors. All Rights Reserved.
|
|
#
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License.
|
|
#
|
|
|
|
"""Per-request identifiers forwarded to upstream LLM providers.
|
|
|
|
An agent run (or any LLM-issuing flow) installs the originating ``session_id`` /
|
|
``user_id`` here. The chat model layer reads it and forwards an end-user
|
|
identifier as the OpenAI-standard ``user`` request field, which providers such as
|
|
OpenAI and OpenRouter include in the request body. This lets upstream activity be
|
|
correlated back to the session/user that produced it.
|
|
|
|
The value is a small dict (or ``None`` when no request context is active), e.g.
|
|
``{"session_id": "...", "user_id": "..."}``.
|
|
"""
|
|
|
|
import contextvars
|
|
import logging
|
|
|
|
llm_request_context: contextvars.ContextVar = contextvars.ContextVar("ragflow_llm_request_context", default=None)
|
|
|
|
|
|
def set_llm_request_context(session_id: str | None = None, user_id: str | None = None):
|
|
"""Install the current request identifiers and return the reset token.
|
|
|
|
Pass the returned token to ``reset_llm_request_context`` (typically in a
|
|
``finally`` block) so the value does not leak to later calls in the same task.
|
|
"""
|
|
ctx = {}
|
|
if session_id:
|
|
ctx["session_id"] = str(session_id)[:128]
|
|
if user_id:
|
|
ctx["user_id"] = str(user_id)[:128]
|
|
# Log only presence flags, never the raw identifiers.
|
|
logging.debug("Installing LLM request context (session=%s, user=%s)", bool(session_id), bool(user_id))
|
|
return llm_request_context.set(ctx or None)
|
|
|
|
|
|
def reset_llm_request_context(token) -> None:
|
|
try:
|
|
llm_request_context.reset(token)
|
|
except (ValueError, RuntimeError):
|
|
# The context may be reset from a different context (e.g. an async generator
|
|
# closed on client disconnect -> ValueError) or with an already-consumed
|
|
# token (Python 3.13+ -> RuntimeError); fall back to clearing the value.
|
|
logging.debug("LLM request context reset failed; clearing active context", exc_info=True)
|
|
llm_request_context.set(None)
|
|
|
|
|
|
def current_llm_user() -> str | None:
|
|
"""Return the identifier to forward as the provider ``user`` field.
|
|
|
|
Prefers ``session_id`` (so upstream activity can be traced per chat session),
|
|
falling back to ``user_id``. Returns ``None`` when no context is active.
|
|
"""
|
|
ctx = llm_request_context.get()
|
|
if not ctx:
|
|
return None
|
|
return ctx.get("session_id") or ctx.get("user_id") or None
|