mirror of
https://github.com/infiniflow/ragflow.git
synced 2026-07-09 04:44:49 +08:00
139 lines
4.9 KiB
Python
139 lines
4.9 KiB
Python
#
|
|
# Copyright 2026 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.
|
|
#
|
|
"""Unit tests for the LLM request-context user-forwarding precedence rules.
|
|
|
|
Covers the behaviour flagged in review: (1) session_id is preferred over
|
|
user_id, (2) an explicit caller-supplied ``user`` overrides the context value,
|
|
and (3) a caller-supplied empty ``user`` suppresses forwarding entirely.
|
|
"""
|
|
import types
|
|
|
|
import pytest
|
|
|
|
from common.llm_request_context import (
|
|
current_llm_user,
|
|
llm_request_context,
|
|
reset_llm_request_context,
|
|
set_llm_request_context,
|
|
)
|
|
|
|
|
|
@pytest.fixture(autouse=True)
|
|
def _isolate_context():
|
|
"""Ensure every test starts and ends with no active request context."""
|
|
token = llm_request_context.set(None)
|
|
yield
|
|
try:
|
|
llm_request_context.reset(token)
|
|
except ValueError:
|
|
llm_request_context.set(None)
|
|
|
|
|
|
@pytest.mark.p2
|
|
class TestCurrentLlmUser:
|
|
def test_no_active_context_returns_none(self):
|
|
assert current_llm_user() is None
|
|
|
|
def test_session_id_preferred_over_user_id(self):
|
|
token = set_llm_request_context(session_id="sess-1", user_id="user-1")
|
|
try:
|
|
assert current_llm_user() == "sess-1"
|
|
finally:
|
|
reset_llm_request_context(token)
|
|
|
|
def test_falls_back_to_user_id_without_session(self):
|
|
token = set_llm_request_context(session_id=None, user_id="user-1")
|
|
try:
|
|
assert current_llm_user() == "user-1"
|
|
finally:
|
|
reset_llm_request_context(token)
|
|
|
|
def test_empty_identifiers_return_none(self):
|
|
token = set_llm_request_context(session_id=None, user_id=None)
|
|
try:
|
|
assert current_llm_user() is None
|
|
finally:
|
|
reset_llm_request_context(token)
|
|
|
|
def test_identifier_is_truncated_to_128_chars(self):
|
|
token = set_llm_request_context(session_id="s" * 200)
|
|
try:
|
|
assert current_llm_user() == "s" * 128
|
|
finally:
|
|
reset_llm_request_context(token)
|
|
|
|
|
|
@pytest.mark.p2
|
|
class TestSetResetLifecycle:
|
|
def test_reset_restores_previous_context(self):
|
|
outer = set_llm_request_context(session_id="outer")
|
|
try:
|
|
inner = set_llm_request_context(session_id="inner")
|
|
assert current_llm_user() == "inner"
|
|
reset_llm_request_context(inner)
|
|
assert current_llm_user() == "outer"
|
|
finally:
|
|
reset_llm_request_context(outer)
|
|
|
|
def test_reset_with_stale_token_does_not_raise(self):
|
|
token = set_llm_request_context(session_id="x")
|
|
reset_llm_request_context(token)
|
|
# Resetting again with the now-stale token must fall back, not raise
|
|
# (mirrors an async generator closed from a different context).
|
|
reset_llm_request_context(token)
|
|
assert current_llm_user() is None
|
|
|
|
|
|
@pytest.mark.p2
|
|
class TestCompletionArgsUserPrecedence:
|
|
"""Exercise the provider chokepoint: LiteLLMBase._construct_completion_args."""
|
|
|
|
def _construct(self, **kwargs):
|
|
chat_model = pytest.importorskip("rag.llm.chat_model")
|
|
# provider="" keeps every provider-specific branch inert, so only the
|
|
# generic completion_args + user-forwarding path is exercised.
|
|
fake = types.SimpleNamespace(model_name="m", api_key="k", max_retries=0, provider="")
|
|
return chat_model.LiteLLMBase._construct_completion_args(fake, [], False, False, **kwargs)
|
|
|
|
def test_context_user_applied_when_caller_omits_it(self):
|
|
token = set_llm_request_context(session_id="sess-9", user_id="user-9")
|
|
try:
|
|
args = self._construct()
|
|
assert args["user"] == "sess-9"
|
|
finally:
|
|
reset_llm_request_context(token)
|
|
|
|
def test_caller_user_overrides_context(self):
|
|
token = set_llm_request_context(session_id="sess-9")
|
|
try:
|
|
args = self._construct(user="explicit-caller")
|
|
assert args["user"] == "explicit-caller"
|
|
finally:
|
|
reset_llm_request_context(token)
|
|
|
|
def test_caller_empty_user_suppresses_forwarding(self):
|
|
token = set_llm_request_context(session_id="sess-9")
|
|
try:
|
|
args = self._construct(user="")
|
|
# Key presence (not truthiness) is honoured: the empty string wins.
|
|
assert args["user"] == ""
|
|
finally:
|
|
reset_llm_request_context(token)
|
|
|
|
def test_no_context_leaves_user_unset(self):
|
|
args = self._construct()
|
|
assert "user" not in args
|