chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,178 @@
|
||||
"""Tests for entity extraction gleaning token limit guard."""
|
||||
|
||||
import logging
|
||||
from unittest.mock import AsyncMock
|
||||
|
||||
import pytest
|
||||
|
||||
from lightrag.utils import Tokenizer, TokenizerInterface
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def _propagate_lightrag_logger(monkeypatch):
|
||||
"""``lightrag.utils.logger`` sets ``propagate = False`` to avoid noisy
|
||||
test output; restore propagation locally so ``caplog`` can capture
|
||||
WARNING records emitted from inside ``lightrag.operate``."""
|
||||
monkeypatch.setattr(logging.getLogger("lightrag"), "propagate", True)
|
||||
|
||||
|
||||
class DummyTokenizer(TokenizerInterface):
|
||||
"""Simple 1:1 character-to-token mapping for testing."""
|
||||
|
||||
def encode(self, content: str):
|
||||
return [ord(ch) for ch in content]
|
||||
|
||||
def decode(self, tokens):
|
||||
return "".join(chr(token) for token in tokens)
|
||||
|
||||
|
||||
def _make_global_config(
|
||||
entity_extract_max_gleaning: int = 1,
|
||||
) -> dict:
|
||||
"""Build a minimal global_config dict for extract_entities."""
|
||||
tokenizer = Tokenizer("dummy", DummyTokenizer())
|
||||
extract_func = AsyncMock(return_value="")
|
||||
return {
|
||||
"llm_model_func": extract_func,
|
||||
"role_llm_funcs": {
|
||||
"extract": extract_func,
|
||||
"keyword": extract_func,
|
||||
"query": extract_func,
|
||||
"vlm": extract_func,
|
||||
},
|
||||
"entity_extract_max_gleaning": entity_extract_max_gleaning,
|
||||
"entity_extract_max_records": 100,
|
||||
"entity_extract_max_entities": 40,
|
||||
"addon_params": {},
|
||||
"tokenizer": tokenizer,
|
||||
"llm_model_max_async": 1,
|
||||
}
|
||||
|
||||
|
||||
# Minimal valid extraction result that _process_extraction_result can parse
|
||||
_EXTRACTION_RESULT = (
|
||||
"(entity<|#|>TEST_ENTITY<|#|>CONCEPT<|#|>A test entity)<|COMPLETE|>"
|
||||
)
|
||||
|
||||
|
||||
def _make_chunks(content: str = "Test content.") -> dict[str, dict]:
|
||||
return {
|
||||
"chunk-001": {
|
||||
"tokens": len(content),
|
||||
"content": content,
|
||||
"full_doc_id": "doc-001",
|
||||
"chunk_order_index": 0,
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.offline
|
||||
@pytest.mark.asyncio
|
||||
async def test_gleaning_skipped_when_tokens_exceed_limit(
|
||||
monkeypatch, caplog, _propagate_lightrag_logger
|
||||
):
|
||||
"""Gleaning must be skipped (with a WARNING) when the projected
|
||||
gleaning input — system + history(user+assistant) + continue prompt —
|
||||
exceeds ``MAX_EXTRACT_INPUT_TOKENS``. This prevents
|
||||
``context_length_exceeded`` errors from the LLM provider on the second
|
||||
round when the initial response was long.
|
||||
"""
|
||||
from lightrag.operate import extract_entities
|
||||
|
||||
# 10 tokens cannot fit any realistic prompt — guard must trip.
|
||||
monkeypatch.setenv("MAX_EXTRACT_INPUT_TOKENS", "10")
|
||||
|
||||
global_config = _make_global_config(entity_extract_max_gleaning=1)
|
||||
llm_func = global_config["llm_model_func"]
|
||||
llm_func.return_value = _EXTRACTION_RESULT
|
||||
|
||||
with caplog.at_level("WARNING", logger="lightrag"):
|
||||
await extract_entities(
|
||||
chunks=_make_chunks(),
|
||||
global_config=global_config,
|
||||
)
|
||||
|
||||
# Only the initial extraction round ran; gleaning was skipped.
|
||||
assert llm_func.await_count == 1
|
||||
|
||||
warnings_emitted = [
|
||||
rec.getMessage()
|
||||
for rec in caplog.records
|
||||
if rec.levelname == "WARNING"
|
||||
and rec.getMessage().startswith("Gleaning stopped for chunk chunk-001:")
|
||||
]
|
||||
assert warnings_emitted, (
|
||||
"expected a WARNING log explaining gleaning was skipped due to "
|
||||
"token limit; got: "
|
||||
f"{[r.getMessage() for r in caplog.records]}"
|
||||
)
|
||||
# Message must surface both the measured token count and the limit so
|
||||
# operators can size MAX_EXTRACT_INPUT_TOKENS appropriately.
|
||||
msg = warnings_emitted[0]
|
||||
assert "exceeded limit (10)" in msg
|
||||
assert "Input tokens (" in msg
|
||||
|
||||
|
||||
@pytest.mark.offline
|
||||
@pytest.mark.asyncio
|
||||
async def test_gleaning_proceeds_when_tokens_within_limit(monkeypatch):
|
||||
"""Gleaning runs normally when the projected input fits the cap."""
|
||||
from lightrag.operate import extract_entities
|
||||
|
||||
monkeypatch.setenv("MAX_EXTRACT_INPUT_TOKENS", "999999")
|
||||
|
||||
global_config = _make_global_config(entity_extract_max_gleaning=1)
|
||||
llm_func = global_config["llm_model_func"]
|
||||
llm_func.return_value = _EXTRACTION_RESULT
|
||||
|
||||
await extract_entities(
|
||||
chunks=_make_chunks(),
|
||||
global_config=global_config,
|
||||
)
|
||||
|
||||
# Both rounds run: initial extraction + one gleaning pass.
|
||||
assert llm_func.await_count == 2
|
||||
|
||||
|
||||
@pytest.mark.offline
|
||||
@pytest.mark.asyncio
|
||||
async def test_no_gleaning_when_max_gleaning_zero(monkeypatch):
|
||||
"""``entity_extract_max_gleaning=0`` disables gleaning regardless of
|
||||
token budget — the guard is downstream of the feature flag."""
|
||||
from lightrag.operate import extract_entities
|
||||
|
||||
monkeypatch.setenv("MAX_EXTRACT_INPUT_TOKENS", "999999")
|
||||
|
||||
global_config = _make_global_config(entity_extract_max_gleaning=0)
|
||||
llm_func = global_config["llm_model_func"]
|
||||
llm_func.return_value = _EXTRACTION_RESULT
|
||||
|
||||
await extract_entities(
|
||||
chunks=_make_chunks(),
|
||||
global_config=global_config,
|
||||
)
|
||||
|
||||
assert llm_func.await_count == 1
|
||||
|
||||
|
||||
@pytest.mark.offline
|
||||
@pytest.mark.asyncio
|
||||
async def test_gleaning_guard_disabled_when_max_tokens_zero(monkeypatch):
|
||||
"""Setting ``MAX_EXTRACT_INPUT_TOKENS=0`` opts out of the guard so
|
||||
gleaning always runs regardless of input size — useful for callers
|
||||
whose provider has no hard input ceiling."""
|
||||
from lightrag.operate import extract_entities
|
||||
|
||||
monkeypatch.setenv("MAX_EXTRACT_INPUT_TOKENS", "0")
|
||||
|
||||
global_config = _make_global_config(entity_extract_max_gleaning=1)
|
||||
llm_func = global_config["llm_model_func"]
|
||||
llm_func.return_value = _EXTRACTION_RESULT
|
||||
|
||||
await extract_entities(
|
||||
chunks=_make_chunks(),
|
||||
global_config=global_config,
|
||||
)
|
||||
|
||||
# Guard disabled → gleaning still runs even with tight projected input.
|
||||
assert llm_func.await_count == 2
|
||||
@@ -0,0 +1,270 @@
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import AsyncMock, patch
|
||||
|
||||
import pytest
|
||||
|
||||
from lightrag.llm.lmdeploy import lmdeploy_model_if_cache
|
||||
from lightrag.llm.lollms import lollms_model_complete, lollms_model_if_cache
|
||||
from lightrag.llm.ollama import _ollama_model_if_cache, ollama_model_complete
|
||||
|
||||
|
||||
@pytest.mark.offline
|
||||
@pytest.mark.asyncio
|
||||
async def test_ollama_response_format_forwards_to_inner():
|
||||
hashing_kv = SimpleNamespace(global_config={"llm_model_name": "ollama-model"})
|
||||
|
||||
with patch(
|
||||
"lightrag.llm.ollama._ollama_model_if_cache",
|
||||
AsyncMock(return_value="{}"),
|
||||
) as mocked_complete:
|
||||
await ollama_model_complete(
|
||||
prompt="hello",
|
||||
hashing_kv=hashing_kv,
|
||||
response_format={"type": "json_object"},
|
||||
)
|
||||
|
||||
assert mocked_complete.await_args.kwargs["response_format"] == {
|
||||
"type": "json_object"
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.offline
|
||||
@pytest.mark.asyncio
|
||||
async def test_ollama_legacy_keyword_extraction_emits_deprecation_warning():
|
||||
"""_ollama_model_if_cache is the canonical emission site for the shim."""
|
||||
captured_kwargs = {}
|
||||
|
||||
class FakeAsyncClient:
|
||||
def __init__(self, *args, **kwargs):
|
||||
self._client = SimpleNamespace(aclose=AsyncMock())
|
||||
|
||||
async def chat(self, **kwargs):
|
||||
captured_kwargs.update(kwargs)
|
||||
return {"message": {"content": "{}"}}
|
||||
|
||||
with patch("lightrag.llm.ollama.ollama.AsyncClient", FakeAsyncClient):
|
||||
with pytest.warns(DeprecationWarning):
|
||||
await _ollama_model_if_cache(
|
||||
model="ollama-model",
|
||||
prompt="hello",
|
||||
keyword_extraction=True,
|
||||
)
|
||||
|
||||
assert captured_kwargs["format"] == "json"
|
||||
assert "keyword_extraction" not in captured_kwargs
|
||||
assert "response_format" not in captured_kwargs
|
||||
|
||||
|
||||
@pytest.mark.offline
|
||||
@pytest.mark.asyncio
|
||||
async def test_ollama_complete_forwards_legacy_flag_downstream():
|
||||
"""ollama_model_complete is a pure forwarder; the shim fires inside _if_cache."""
|
||||
hashing_kv = SimpleNamespace(global_config={"llm_model_name": "ollama-model"})
|
||||
|
||||
with patch(
|
||||
"lightrag.llm.ollama._ollama_model_if_cache",
|
||||
AsyncMock(return_value="{}"),
|
||||
) as mocked_complete:
|
||||
await ollama_model_complete(
|
||||
prompt="hello",
|
||||
hashing_kv=hashing_kv,
|
||||
keyword_extraction=True,
|
||||
)
|
||||
|
||||
assert mocked_complete.await_args.kwargs.get("keyword_extraction") is True
|
||||
|
||||
|
||||
@pytest.mark.offline
|
||||
@pytest.mark.asyncio
|
||||
async def test_ollama_translates_json_object_response_format_to_native_format():
|
||||
captured_kwargs = {}
|
||||
|
||||
class FakeAsyncClient:
|
||||
def __init__(self, *args, **kwargs):
|
||||
self._client = SimpleNamespace(aclose=AsyncMock())
|
||||
|
||||
async def chat(self, **kwargs):
|
||||
captured_kwargs.update(kwargs)
|
||||
return {"message": {"content": "{}"}}
|
||||
|
||||
with patch("lightrag.llm.ollama.ollama.AsyncClient", FakeAsyncClient):
|
||||
result = await _ollama_model_if_cache(
|
||||
model="ollama-model",
|
||||
prompt="hello",
|
||||
response_format={"type": "json_object"},
|
||||
)
|
||||
|
||||
assert result == "{}"
|
||||
assert captured_kwargs["format"] == "json"
|
||||
assert "response_format" not in captured_kwargs
|
||||
|
||||
|
||||
@pytest.mark.offline
|
||||
@pytest.mark.asyncio
|
||||
async def test_ollama_unwraps_openai_json_schema_response_format():
|
||||
captured_kwargs = {}
|
||||
schema = {
|
||||
"type": "object",
|
||||
"properties": {"answer": {"type": "string"}},
|
||||
"required": ["answer"],
|
||||
}
|
||||
|
||||
class FakeAsyncClient:
|
||||
def __init__(self, *args, **kwargs):
|
||||
self._client = SimpleNamespace(aclose=AsyncMock())
|
||||
|
||||
async def chat(self, **kwargs):
|
||||
captured_kwargs.update(kwargs)
|
||||
return {"message": {"content": "{}"}}
|
||||
|
||||
with patch("lightrag.llm.ollama.ollama.AsyncClient", FakeAsyncClient):
|
||||
result = await _ollama_model_if_cache(
|
||||
model="ollama-model",
|
||||
prompt="hello",
|
||||
response_format={
|
||||
"type": "json_schema",
|
||||
"json_schema": {"name": "answer_payload", "schema": schema},
|
||||
},
|
||||
)
|
||||
|
||||
assert result == "{}"
|
||||
assert captured_kwargs["format"] == schema
|
||||
assert "response_format" not in captured_kwargs
|
||||
|
||||
|
||||
@pytest.mark.offline
|
||||
@pytest.mark.asyncio
|
||||
async def test_lollms_if_cache_strips_response_format_before_request():
|
||||
"""lollms_model_if_cache drops response_format; lollms has no JSON mode."""
|
||||
captured_requests = []
|
||||
|
||||
class FakeResponse:
|
||||
def __init__(self):
|
||||
pass
|
||||
|
||||
async def __aenter__(self):
|
||||
return self
|
||||
|
||||
async def __aexit__(self, *exc_info):
|
||||
return False
|
||||
|
||||
async def text(self):
|
||||
return "{}"
|
||||
|
||||
class FakeSession:
|
||||
def __init__(self, *args, **kwargs):
|
||||
pass
|
||||
|
||||
async def __aenter__(self):
|
||||
return self
|
||||
|
||||
async def __aexit__(self, *exc_info):
|
||||
return False
|
||||
|
||||
def post(self, url, json):
|
||||
captured_requests.append(json)
|
||||
return FakeResponse()
|
||||
|
||||
with patch("lightrag.llm.lollms.aiohttp.ClientSession", FakeSession):
|
||||
result = await lollms_model_if_cache(
|
||||
model="lollms-model",
|
||||
prompt="hello",
|
||||
response_format={"type": "json_object"},
|
||||
)
|
||||
|
||||
assert result == "{}"
|
||||
assert captured_requests
|
||||
assert "response_format" not in captured_requests[0]
|
||||
|
||||
|
||||
@pytest.mark.offline
|
||||
@pytest.mark.asyncio
|
||||
async def test_lollms_if_cache_emits_deprecation_warning():
|
||||
class FakeResponse:
|
||||
async def __aenter__(self):
|
||||
return self
|
||||
|
||||
async def __aexit__(self, *exc_info):
|
||||
return False
|
||||
|
||||
async def text(self):
|
||||
return "{}"
|
||||
|
||||
class FakeSession:
|
||||
def __init__(self, *args, **kwargs):
|
||||
pass
|
||||
|
||||
async def __aenter__(self):
|
||||
return self
|
||||
|
||||
async def __aexit__(self, *exc_info):
|
||||
return False
|
||||
|
||||
def post(self, url, json):
|
||||
return FakeResponse()
|
||||
|
||||
with patch("lightrag.llm.lollms.aiohttp.ClientSession", FakeSession):
|
||||
with pytest.warns(DeprecationWarning):
|
||||
await lollms_model_if_cache(
|
||||
model="lollms-model",
|
||||
prompt="hello",
|
||||
keyword_extraction=True,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.offline
|
||||
@pytest.mark.asyncio
|
||||
async def test_lollms_complete_forwards_legacy_flag_downstream():
|
||||
hashing_kv = SimpleNamespace(global_config={"llm_model_name": "lollms-model"})
|
||||
|
||||
with patch(
|
||||
"lightrag.llm.lollms.lollms_model_if_cache",
|
||||
AsyncMock(return_value="{}"),
|
||||
) as mocked_complete:
|
||||
await lollms_model_complete(
|
||||
prompt="hello",
|
||||
hashing_kv=hashing_kv,
|
||||
keyword_extraction=True,
|
||||
)
|
||||
|
||||
assert mocked_complete.await_args.kwargs.get("keyword_extraction") is True
|
||||
|
||||
|
||||
@pytest.mark.offline
|
||||
@pytest.mark.asyncio
|
||||
async def test_lmdeploy_strips_response_format_before_generation_config(monkeypatch):
|
||||
captured_gen_config_kwargs = {}
|
||||
|
||||
class FakeGenerationConfig:
|
||||
def __init__(self, **kwargs):
|
||||
captured_gen_config_kwargs.update(kwargs)
|
||||
|
||||
class FakeVersion:
|
||||
def __lt__(self, other):
|
||||
return False
|
||||
|
||||
async def fake_generate(*_args, **_kwargs):
|
||||
yield SimpleNamespace(response="{}")
|
||||
|
||||
monkeypatch.setattr(
|
||||
"lightrag.llm.lmdeploy.initialize_lmdeploy_pipeline",
|
||||
lambda **_kwargs: SimpleNamespace(generate=fake_generate),
|
||||
)
|
||||
|
||||
import sys
|
||||
|
||||
sys.modules["lmdeploy"] = SimpleNamespace(
|
||||
__version__="0.6.0",
|
||||
version_info=FakeVersion(),
|
||||
GenerationConfig=FakeGenerationConfig,
|
||||
)
|
||||
|
||||
result = await lmdeploy_model_if_cache(
|
||||
model="lmdeploy-model",
|
||||
prompt="hello",
|
||||
response_format={"type": "json_object"},
|
||||
)
|
||||
|
||||
assert result == "{}"
|
||||
assert "response_format" not in captured_gen_config_kwargs
|
||||
assert "keyword_extraction" not in captured_gen_config_kwargs
|
||||
@@ -0,0 +1,155 @@
|
||||
import pytest
|
||||
from unittest.mock import patch
|
||||
|
||||
from lightrag.base import QueryParam
|
||||
from lightrag.operate import _parse_keywords_payload, extract_keywords_only
|
||||
|
||||
|
||||
class _FakeKeywordModel:
|
||||
def model_dump(self):
|
||||
return {
|
||||
"high_level_keywords": ["AI"],
|
||||
"low_level_keywords": ["RAG", "Graph"],
|
||||
}
|
||||
|
||||
|
||||
class _FakeTokenizer:
|
||||
def encode(self, content: str) -> list[int]:
|
||||
return [ord(ch) for ch in content]
|
||||
|
||||
|
||||
class _FakeKVStorage:
|
||||
def __init__(self):
|
||||
self.global_config = {"enable_llm_cache": True}
|
||||
self._store = {}
|
||||
|
||||
async def get_by_id(self, key):
|
||||
return self._store.get(key)
|
||||
|
||||
async def upsert(self, entries):
|
||||
self._store.update(entries)
|
||||
|
||||
|
||||
def _keyword_global_config(
|
||||
model: str, binding: str = "openai", keyword_func=None
|
||||
) -> dict:
|
||||
return {
|
||||
"addon_params": {"language": "en"},
|
||||
"tokenizer": _FakeTokenizer(),
|
||||
"role_llm_funcs": {"keyword": keyword_func} if keyword_func else {},
|
||||
"llm_cache_identities": {
|
||||
"keyword": {
|
||||
"role": "keyword",
|
||||
"binding": binding,
|
||||
"model": model,
|
||||
"host": "https://api.example.com/v1",
|
||||
}
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.offline
|
||||
def test_parse_keywords_payload_accepts_model_like_objects():
|
||||
is_valid, hl_keywords, ll_keywords = _parse_keywords_payload(_FakeKeywordModel())
|
||||
|
||||
assert is_valid is True
|
||||
assert hl_keywords == ["AI"]
|
||||
assert ll_keywords == ["RAG", "Graph"]
|
||||
|
||||
|
||||
@pytest.mark.offline
|
||||
def test_parse_keywords_payload_extracts_json_from_wrapped_text():
|
||||
result = """
|
||||
analysis first
|
||||
{"high_level_keywords":"AI, Agents","low_level_keywords":["RAG","LightRAG"]}
|
||||
trailing note
|
||||
"""
|
||||
|
||||
is_valid, hl_keywords, ll_keywords = _parse_keywords_payload(result)
|
||||
|
||||
assert is_valid is True
|
||||
assert hl_keywords == ["AI", "Agents"]
|
||||
assert ll_keywords == ["RAG", "LightRAG"]
|
||||
|
||||
|
||||
@pytest.mark.offline
|
||||
def test_parse_keywords_payload_warns_when_json_repair_is_used():
|
||||
broken_result = (
|
||||
'{"high_level_keywords":"AI, Agents","low_level_keywords":["RAG","LightRAG"]'
|
||||
)
|
||||
|
||||
with patch("lightrag.operate.logger.warning") as mocked_warning:
|
||||
is_valid, hl_keywords, ll_keywords = _parse_keywords_payload(broken_result)
|
||||
|
||||
assert is_valid is True
|
||||
assert hl_keywords == ["AI", "Agents"]
|
||||
assert ll_keywords == ["RAG", "LightRAG"]
|
||||
mocked_warning.assert_called_once()
|
||||
assert (
|
||||
"Keyword extraction response required JSON repair"
|
||||
in mocked_warning.call_args[0][0]
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.offline
|
||||
@pytest.mark.asyncio
|
||||
async def test_extract_keywords_only_accepts_empty_keyword_cache_without_requery():
|
||||
async def should_not_run(*_args, **_kwargs):
|
||||
raise AssertionError(
|
||||
"keyword LLM should not be called on a valid empty cache hit"
|
||||
)
|
||||
|
||||
param = QueryParam()
|
||||
global_config = _keyword_global_config("model-a", keyword_func=should_not_run)
|
||||
|
||||
with patch(
|
||||
"lightrag.operate.handle_cache",
|
||||
return_value=('{"high_level_keywords":[],"low_level_keywords":[]}', None),
|
||||
):
|
||||
hl_keywords, ll_keywords = await extract_keywords_only(
|
||||
"hello",
|
||||
param,
|
||||
global_config,
|
||||
hashing_kv=None,
|
||||
)
|
||||
|
||||
assert hl_keywords == []
|
||||
assert ll_keywords == []
|
||||
|
||||
|
||||
@pytest.mark.offline
|
||||
@pytest.mark.asyncio
|
||||
async def test_extract_keywords_only_partitions_cache_by_keyword_llm_identity():
|
||||
cache = _FakeKVStorage()
|
||||
calls = 0
|
||||
|
||||
async def keyword_model(*_args, **_kwargs):
|
||||
nonlocal calls
|
||||
calls += 1
|
||||
return (
|
||||
'{"high_level_keywords":["model-'
|
||||
+ str(calls)
|
||||
+ '"],"low_level_keywords":["rag"]}'
|
||||
)
|
||||
|
||||
param = QueryParam()
|
||||
|
||||
first_hl, first_ll = await extract_keywords_only(
|
||||
"same query",
|
||||
param,
|
||||
_keyword_global_config("model-a", keyword_func=keyword_model),
|
||||
hashing_kv=cache,
|
||||
)
|
||||
second_hl, second_ll = await extract_keywords_only(
|
||||
"same query",
|
||||
param,
|
||||
_keyword_global_config("model-b", keyword_func=keyword_model),
|
||||
hashing_kv=cache,
|
||||
)
|
||||
|
||||
assert first_hl == ["model-1"]
|
||||
assert first_ll == ["rag"]
|
||||
assert second_hl == ["model-2"]
|
||||
assert second_ll == ["rag"]
|
||||
assert calls == 2
|
||||
assert len(cache._store) == 2
|
||||
@@ -0,0 +1,69 @@
|
||||
import json
|
||||
import re
|
||||
|
||||
import pytest
|
||||
|
||||
from lightrag.prompt import PROMPTS
|
||||
|
||||
|
||||
@pytest.mark.offline
|
||||
def test_keywords_extraction_prompt_template_formats_with_literal_json_braces():
|
||||
rendered = PROMPTS["keywords_extraction"].format(
|
||||
query="hello",
|
||||
examples="example",
|
||||
language="en",
|
||||
)
|
||||
|
||||
assert "first character of your response must be `{`" in rendered
|
||||
assert '{"high_level_keywords": [], "low_level_keywords": []}' in rendered
|
||||
|
||||
|
||||
@pytest.mark.offline
|
||||
def test_keywords_extraction_examples_are_format_only():
|
||||
"""Keyword examples must be placeholder-only JSON templates, not sample demos.
|
||||
|
||||
Rather than denylisting specific sample queries (brittle: generic words like
|
||||
"education" would both false-match unrelated content and let new samples slip
|
||||
through), assert the structural shape: no ``Query:``/``Output:`` demo framing,
|
||||
and every keyword is an angle-bracket placeholder.
|
||||
"""
|
||||
placeholder = re.compile(r"<[^<>]+>")
|
||||
|
||||
for example in PROMPTS["keywords_extraction_examples"]:
|
||||
assert "Query:" not in example
|
||||
assert "Output:" not in example
|
||||
|
||||
parsed = json.loads(example)
|
||||
assert set(parsed) == {"high_level_keywords", "low_level_keywords"}
|
||||
keywords = parsed["high_level_keywords"] + parsed["low_level_keywords"]
|
||||
assert keywords
|
||||
for keyword in keywords:
|
||||
assert placeholder.fullmatch(keyword), keyword
|
||||
|
||||
|
||||
@pytest.mark.offline
|
||||
def test_keywords_extraction_prompt_labels_template_as_not_source_text():
|
||||
prompt = PROMPTS["keywords_extraction"]
|
||||
|
||||
assert "---Output Format Template---" in prompt
|
||||
assert "---Examples---" not in prompt
|
||||
assert "Apple Inc." not in prompt
|
||||
assert "output JSON format template only" in prompt
|
||||
assert "not source text" in prompt
|
||||
assert "must never be used as keyword extraction content" in prompt
|
||||
assert (
|
||||
"derived only from the `User Query` in the `---Real Data---` section" in prompt
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.offline
|
||||
def test_keywords_extraction_prompt_keeps_single_real_user_query_section():
|
||||
rendered = PROMPTS["keywords_extraction"].format(
|
||||
query="How did LightRAG improve retrieval?",
|
||||
examples="\n".join(PROMPTS["keywords_extraction_examples"]),
|
||||
language="English",
|
||||
)
|
||||
|
||||
assert rendered.count("User Query:") == 1
|
||||
assert "User Query: How did LightRAG improve retrieval?" in rendered
|
||||
assert "User Query:" not in "\n".join(PROMPTS["keywords_extraction_examples"])
|
||||
@@ -0,0 +1,115 @@
|
||||
"""Regression tests for the map-reduce chunking path of
|
||||
``_handle_entity_relation_summary``.
|
||||
|
||||
The map phase tokenizes every description to decide (a) whether the list
|
||||
needs splitting and (b) where to cut the chunks. Historically the same
|
||||
description was encoded twice per map-reduce iteration — once to sum the
|
||||
total token count and again while building the chunks — which doubled the
|
||||
tiktoken BPE cost on every entity/edge merge. The counts from the first
|
||||
pass are now memoized and reused by the second.
|
||||
|
||||
These tests guard:
|
||||
|
||||
* the chunking output is unchanged (behavioral), and
|
||||
* each description is encoded exactly once per map-reduce iteration
|
||||
(the perf property the fix introduced).
|
||||
"""
|
||||
|
||||
from unittest.mock import AsyncMock
|
||||
|
||||
import pytest
|
||||
|
||||
from lightrag.utils import Tokenizer, TokenizerInterface
|
||||
|
||||
|
||||
class _CountingTokenizer(TokenizerInterface):
|
||||
"""1:1 character-to-token mapping that records every ``encode`` input."""
|
||||
|
||||
def __init__(self) -> None:
|
||||
self.encoded: list[str] = []
|
||||
|
||||
def encode(self, content: str):
|
||||
self.encoded.append(content)
|
||||
return [ord(ch) for ch in content]
|
||||
|
||||
def decode(self, tokens):
|
||||
return "".join(chr(token) for token in tokens)
|
||||
|
||||
|
||||
def _make_chunking_config(mock_extract) -> tuple[dict, _CountingTokenizer]:
|
||||
"""``global_config`` whose ``summary_context_size`` forces the map-reduce
|
||||
chunking path (4 five-token descriptions vs. a 10-token window)."""
|
||||
counter = _CountingTokenizer()
|
||||
tokenizer = Tokenizer("dummy", counter)
|
||||
global_config = {
|
||||
"role_llm_funcs": {"extract": mock_extract},
|
||||
"addon_params": {},
|
||||
"_resolved_summary_language": "English",
|
||||
"summary_length_recommended": 100,
|
||||
# 4 descriptions x 5 tokens = 20 tokens > 10-token window → chunking
|
||||
"summary_context_size": 10,
|
||||
"summary_max_tokens": 100_000,
|
||||
"force_llm_summary_on_merge": 10,
|
||||
"tokenizer": tokenizer,
|
||||
}
|
||||
return global_config, counter
|
||||
|
||||
|
||||
@pytest.mark.offline
|
||||
@pytest.mark.asyncio
|
||||
async def test_map_reduce_chunking_is_preserved():
|
||||
"""Forcing the chunking path must still split the descriptions into groups,
|
||||
summarize each via the LLM, and join the results — i.e. memoizing the
|
||||
token counts must not change the chunk boundaries or the final output."""
|
||||
from lightrag.operate import _handle_entity_relation_summary
|
||||
|
||||
# Distinct summaries per reduce call so the joined result reflects the
|
||||
# number of chunks produced.
|
||||
mock_extract = AsyncMock(side_effect=["SUM_ONE", "SUM_TWO", "SUM_THREE"])
|
||||
global_config, _ = _make_chunking_config(mock_extract)
|
||||
|
||||
description, llm_was_used = await _handle_entity_relation_summary(
|
||||
"Entity",
|
||||
"TEST_ENTITY",
|
||||
["alpha", "beta", "gamma", "delta"],
|
||||
"<SEP>",
|
||||
global_config,
|
||||
llm_response_cache=None,
|
||||
)
|
||||
|
||||
# Reduce phase ran (the LLM was used to summarize at least one chunk).
|
||||
assert llm_was_used is True
|
||||
# 4 five-token descriptions under a 10-token window → 2 chunks → 2 summaries
|
||||
# joined in the next iteration (len == 2 takes the no-LLM join exit).
|
||||
assert description == "SUM_ONE<SEP>SUM_TWO"
|
||||
|
||||
|
||||
@pytest.mark.offline
|
||||
@pytest.mark.asyncio
|
||||
async def test_map_phase_encodes_each_description_once():
|
||||
"""The map phase must tokenize each description exactly once per
|
||||
map-reduce iteration. Pre-fix the chunk-building pass re-encoded every
|
||||
description, so each of the 4 inputs appeared 2x in the first iteration;
|
||||
post-fix it appears exactly once (the second pass reuses the counts)."""
|
||||
from lightrag.operate import _handle_entity_relation_summary
|
||||
|
||||
mock_extract = AsyncMock(side_effect=["SUM_ONE", "SUM_TWO", "SUM_THREE"])
|
||||
global_config, counter = _make_chunking_config(mock_extract)
|
||||
inputs = ["alpha", "beta", "gamma", "delta"]
|
||||
|
||||
await _handle_entity_relation_summary(
|
||||
"Entity",
|
||||
"TEST_ENTITY",
|
||||
inputs,
|
||||
"<SEP>",
|
||||
global_config,
|
||||
llm_response_cache=None,
|
||||
)
|
||||
|
||||
# In the first (chunking) iteration each input description must be encoded
|
||||
# exactly once by the map phase. Pre-fix this was 2 per description.
|
||||
for desc in inputs:
|
||||
assert counter.encoded.count(desc) == 1, (
|
||||
f"description {desc!r} was encoded {counter.encoded.count(desc)} "
|
||||
f"time(s) by the map phase; expected exactly 1 (memoized counts)"
|
||||
)
|
||||
@@ -0,0 +1,160 @@
|
||||
"""Regression tests: every merge-stage description outcome must be
|
||||
sanitized before it lands on graph nodes/edges.
|
||||
|
||||
Extraction-time descriptions are cleaned by
|
||||
``sanitize_and_normalize_extracted_text``, but merge-stage descriptions
|
||||
can bypass that: LLM re-summaries, descriptions read back from existing
|
||||
graph nodes, and multimodal entity descriptions injected straight from
|
||||
chunk content. Any of them carrying control characters / surrogates
|
||||
(e.g. ``\\frac`` decoded as ``\\x0c`` + ``rac`` from unescaped LaTeX in
|
||||
VLM JSON) would break GraphML (XML) serialization on the next NetworkX
|
||||
flush with::
|
||||
|
||||
ValueError: All strings must be XML compatible: Unicode or ASCII,
|
||||
no NULL bytes or control characters
|
||||
|
||||
``_handle_entity_relation_summary`` (single / join outcomes) and
|
||||
``_summarize_descriptions`` (LLM outcome) now sanitize at every exit.
|
||||
"""
|
||||
|
||||
from unittest.mock import AsyncMock
|
||||
|
||||
import networkx as nx
|
||||
import pytest
|
||||
|
||||
from lightrag.utils import Tokenizer, TokenizerInterface
|
||||
|
||||
|
||||
class DummyTokenizer(TokenizerInterface):
|
||||
"""Simple 1:1 character-to-token mapping for testing."""
|
||||
|
||||
def encode(self, content: str):
|
||||
return [ord(ch) for ch in content]
|
||||
|
||||
def decode(self, tokens):
|
||||
return "".join(chr(token) for token in tokens)
|
||||
|
||||
|
||||
def _make_global_config(summary_return: str) -> dict:
|
||||
"""Minimal global_config for ``_summarize_descriptions`` whose ``extract``
|
||||
role LLM returns ``summary_return`` verbatim."""
|
||||
tokenizer = Tokenizer("dummy", DummyTokenizer())
|
||||
extract_func = AsyncMock(return_value=summary_return)
|
||||
return {
|
||||
"role_llm_funcs": {"extract": extract_func},
|
||||
"addon_params": {},
|
||||
"_resolved_summary_language": "English",
|
||||
"summary_length_recommended": 100,
|
||||
"summary_context_size": 10_000,
|
||||
"tokenizer": tokenizer,
|
||||
}
|
||||
|
||||
|
||||
# A description that the LLM "echoed" from a dirty source: NULL byte, bell,
|
||||
# unit-separator and DEL are all illegal in XML 1.0; \t and \n are legal and
|
||||
# must be preserved.
|
||||
_DIRTY_SUMMARY = "Clean text\x00with\x07control\x1fchars\x7f\tand\nnewlines"
|
||||
_EXPECTED_CLEAN = "Clean textwithcontrolchars\tand\nnewlines"
|
||||
|
||||
|
||||
@pytest.mark.offline
|
||||
@pytest.mark.asyncio
|
||||
async def test_summarize_descriptions_strips_control_chars():
|
||||
"""The LLM summary path must remove XML-incompatible control characters
|
||||
while preserving legal whitespace."""
|
||||
from lightrag.operate import _summarize_descriptions
|
||||
|
||||
global_config = _make_global_config(_DIRTY_SUMMARY)
|
||||
|
||||
summary = await _summarize_descriptions(
|
||||
"Entity",
|
||||
"TEST_ENTITY",
|
||||
["first description", "second description"],
|
||||
global_config,
|
||||
llm_response_cache=None,
|
||||
)
|
||||
|
||||
assert summary == _EXPECTED_CLEAN
|
||||
assert "\x00" not in summary
|
||||
assert "\x07" not in summary
|
||||
assert "\x1f" not in summary
|
||||
assert "\x7f" not in summary
|
||||
|
||||
|
||||
@pytest.mark.offline
|
||||
@pytest.mark.asyncio
|
||||
async def test_single_description_early_return_is_sanitized():
|
||||
"""The ``len(description_list) == 1`` early return skips the LLM
|
||||
entirely — this is exactly the path a dirty multimodal entity
|
||||
description (chunk content reused verbatim) takes on first insert."""
|
||||
from lightrag.operate import _handle_entity_relation_summary
|
||||
|
||||
description, llm_was_used = await _handle_entity_relation_summary(
|
||||
"Entity",
|
||||
"TEST_ENTITY",
|
||||
[_DIRTY_SUMMARY],
|
||||
"<SEP>",
|
||||
# Early return happens before any config key is read.
|
||||
global_config={},
|
||||
llm_response_cache=None,
|
||||
)
|
||||
|
||||
assert llm_was_used is False
|
||||
assert description == _EXPECTED_CLEAN
|
||||
|
||||
|
||||
@pytest.mark.offline
|
||||
@pytest.mark.asyncio
|
||||
async def test_join_without_llm_is_sanitized():
|
||||
"""The no-LLM join outcome must also sanitize: descriptions read back
|
||||
from pre-existing (dirty) graph nodes re-enter the merge here."""
|
||||
from lightrag.operate import _handle_entity_relation_summary
|
||||
|
||||
tokenizer = Tokenizer("dummy", DummyTokenizer())
|
||||
global_config = {
|
||||
"tokenizer": tokenizer,
|
||||
"summary_context_size": 100_000,
|
||||
"summary_max_tokens": 100_000,
|
||||
"force_llm_summary_on_merge": 10,
|
||||
}
|
||||
|
||||
description, llm_was_used = await _handle_entity_relation_summary(
|
||||
"Entity",
|
||||
"TEST_ENTITY",
|
||||
[_DIRTY_SUMMARY, "clean second description"],
|
||||
"<SEP>",
|
||||
global_config,
|
||||
llm_response_cache=None,
|
||||
)
|
||||
|
||||
assert llm_was_used is False
|
||||
assert description == f"{_EXPECTED_CLEAN}<SEP>clean second description"
|
||||
|
||||
|
||||
@pytest.mark.offline
|
||||
@pytest.mark.asyncio
|
||||
async def test_summarized_description_is_graphml_writable(tmp_path):
|
||||
"""End-to-end guard: a node whose description came from the summary path
|
||||
must serialize to GraphML without the original ValueError."""
|
||||
from lightrag.operate import _summarize_descriptions
|
||||
|
||||
global_config = _make_global_config(_DIRTY_SUMMARY)
|
||||
|
||||
description = await _summarize_descriptions(
|
||||
"Entity",
|
||||
"TEST_ENTITY",
|
||||
["first description", "second description"],
|
||||
global_config,
|
||||
llm_response_cache=None,
|
||||
)
|
||||
|
||||
graph = nx.Graph()
|
||||
graph.add_node("TEST_ENTITY", description=description)
|
||||
|
||||
# Pre-fix, write_graphml raised:
|
||||
# "All strings must be XML compatible: ... no NULL bytes or control characters"
|
||||
out = tmp_path / "graph.graphml"
|
||||
nx.write_graphml(graph, out)
|
||||
|
||||
reloaded = nx.read_graphml(out)
|
||||
assert reloaded.nodes["TEST_ENTITY"]["description"] == _EXPECTED_CLEAN
|
||||
Reference in New Issue
Block a user