wehub-resource-sync
1861 lines
62 KiB
Python
1861 lines
62 KiB
Python
"""Tests for entity extraction stability after refactoring.
|
||
|
||
Covers:
|
||
- entity_types_guidance injected into prompts (text mode and JSON mode)
|
||
- custom entity_types_guidance via addon_params overrides default
|
||
- ENTITY_TYPES env var raises SystemExit at LightRAG init
|
||
- EntityExtractionResult Pydantic schema used in JSON mode (entity_extraction kwarg)
|
||
- Default entity type guidance constant is present and non-empty
|
||
"""
|
||
|
||
import json
|
||
import os
|
||
import re
|
||
from pathlib import Path
|
||
from unittest.mock import AsyncMock, patch
|
||
|
||
import pytest
|
||
|
||
from lightrag.utils import EmbeddingFunc, 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(
|
||
addon_params: dict | None = None,
|
||
use_json: bool = False,
|
||
max_gleaning: int = 0,
|
||
prompt_profile: dict | None = None,
|
||
) -> dict:
|
||
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": max_gleaning,
|
||
"entity_extract_max_records": 100,
|
||
"entity_extract_max_entities": 40,
|
||
"addon_params": addon_params if addon_params is not None else {},
|
||
"tokenizer": tokenizer,
|
||
"max_extract_input_tokens": 20480,
|
||
"llm_model_max_async": 1,
|
||
"entity_extraction_use_json": use_json,
|
||
"_entity_extraction_prompt_profile": prompt_profile,
|
||
}
|
||
|
||
|
||
def _make_chunks(content: str = "Alice founded Acme Corp in 1990.") -> dict[str, dict]:
|
||
return {
|
||
"chunk-001": {
|
||
"tokens": len(content),
|
||
"content": content,
|
||
"full_doc_id": "doc-001",
|
||
"chunk_order_index": 0,
|
||
}
|
||
}
|
||
|
||
|
||
def _require_yaml() -> None:
|
||
pytest.importorskip("yaml")
|
||
|
||
|
||
def _write_prompt_profile(
|
||
path: Path,
|
||
*,
|
||
guidance: str | None = None,
|
||
text_examples: list[str] | None = None,
|
||
json_examples: list[str] | None = None,
|
||
) -> None:
|
||
lines: list[str] = []
|
||
|
||
def _append_block(key: str, value: str) -> None:
|
||
lines.append(f"{key}: |")
|
||
for line in value.strip("\n").splitlines():
|
||
lines.append(f" {line}")
|
||
|
||
def _append_examples(key: str, values: list[str]) -> None:
|
||
lines.append(f"{key}:")
|
||
for value in values:
|
||
lines.append(" - |")
|
||
for line in value.strip("\n").splitlines():
|
||
lines.append(f" {line}")
|
||
|
||
if guidance is not None:
|
||
_append_block("entity_types_guidance", guidance)
|
||
if text_examples is not None:
|
||
_append_examples("entity_extraction_examples", text_examples)
|
||
if json_examples is not None:
|
||
_append_examples("entity_extraction_json_examples", json_examples)
|
||
|
||
path.write_text("\n".join(lines) + "\n", encoding="utf-8")
|
||
|
||
|
||
def _dummy_embedding_func() -> EmbeddingFunc:
|
||
async def _embed(texts):
|
||
return [[0.0, 0.0, 0.0] for _ in texts]
|
||
|
||
return EmbeddingFunc(embedding_dim=3, func=_embed)
|
||
|
||
|
||
def _patch_prompt_dir(path: Path):
|
||
return patch("lightrag.prompt.get_entity_type_prompt_dir", return_value=path)
|
||
|
||
|
||
def _text_profile_example(label: str) -> str:
|
||
return f"""---Entity Types---
|
||
- ExampleType: Test type
|
||
|
||
---Input Text---
|
||
```
|
||
{label}
|
||
```
|
||
|
||
---Output---
|
||
entity{{tuple_delimiter}}{label}{{tuple_delimiter}}ExampleType{{tuple_delimiter}}{label} description.
|
||
{{completion_delimiter}}"""
|
||
|
||
|
||
def _json_profile_example(label: str) -> str:
|
||
return f"""---Entity Types---
|
||
- ExampleType: Test type
|
||
|
||
---Input Text---
|
||
```
|
||
{label}
|
||
```
|
||
|
||
---Output---
|
||
{{
|
||
"entities": [
|
||
{{"name": "{label}", "type": "ExampleType", "description": "{label} description."}}
|
||
],
|
||
"relationships": []
|
||
}}"""
|
||
|
||
|
||
# --- Minimal valid LLM responses ---
|
||
|
||
_TEXT_MODE_RESPONSE = (
|
||
"entity<|#|>Alice<|#|>Person<|#|>Alice is the founder of Acme Corp."
|
||
"\nentity<|#|>Acme Corp<|#|>Organization<|#|>Acme Corp is a company founded by Alice."
|
||
"\nrelation<|#|>Alice<|#|>Acme Corp<|#|>founded<|#|>Alice founded Acme Corp."
|
||
"\n<|COMPLETE|>"
|
||
)
|
||
|
||
_TEXT_MODE_MISPREFIXED_RELATION_RESPONSE = (
|
||
"entity<|#|>Alice<|#|>Person<|#|>Alice is the founder of Acme Corp."
|
||
"\nentity<|#|>Acme Corp<|#|>Organization<|#|>Acme Corp is a company founded by Alice."
|
||
"\nentity<|#|>Alice<|#|>Acme Corp<|#|>founded<|#|>Alice founded Acme Corp."
|
||
"\n<|COMPLETE|>"
|
||
)
|
||
|
||
_TEXT_MODE_GLEANED_RELATION_RESPONSES = [
|
||
_TEXT_MODE_MISPREFIXED_RELATION_RESPONSE,
|
||
"\nrelation<|#|>Alice<|#|>Acme Corp<|#|>founded<|#|>Alice founded Acme Corp.\n<|COMPLETE|>",
|
||
]
|
||
|
||
_TEXT_MODE_CROSS_PASS_RELATION_RESPONSES = [
|
||
"entity<|#|>Alice<|#|>Person<|#|>Alice founded a company.\n<|COMPLETE|>",
|
||
"entity<|#|>Acme Corp<|#|>Organization<|#|>Acme Corp was founded by Alice."
|
||
"\nrelation<|#|>Alice<|#|>Acme Corp<|#|>founded<|#|>Alice founded Acme Corp.\n<|COMPLETE|>",
|
||
]
|
||
|
||
_JSON_MODE_RESPONSE = json.dumps(
|
||
{
|
||
"entities": [
|
||
{
|
||
"name": "Alice",
|
||
"type": "Person",
|
||
"description": "Alice is the founder of Acme Corp.",
|
||
},
|
||
{
|
||
"name": "Acme Corp",
|
||
"type": "Organization",
|
||
"description": "Acme Corp is a company founded by Alice.",
|
||
},
|
||
],
|
||
"relationships": [
|
||
{
|
||
"source": "Alice",
|
||
"target": "Acme Corp",
|
||
"keywords": "founded",
|
||
"description": "Alice founded Acme Corp.",
|
||
},
|
||
],
|
||
}
|
||
)
|
||
|
||
|
||
class _DummyTextChunksStorage:
|
||
async def get_by_id(self, chunk_id: str):
|
||
return {"file_path": "test.md"}
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# 1. Default entity_types_guidance constant
|
||
# ---------------------------------------------------------------------------
|
||
|
||
|
||
@pytest.mark.offline
|
||
def test_default_entity_types_guidance_exists():
|
||
"""PROMPTS['default_entity_types_guidance'] must be a non-empty string."""
|
||
from lightrag.prompt import PROMPTS
|
||
|
||
guidance = PROMPTS["default_entity_types_guidance"]
|
||
assert isinstance(guidance, str)
|
||
assert len(guidance.strip()) > 0
|
||
|
||
|
||
@pytest.mark.offline
|
||
def test_default_entity_types_guidance_covers_all_types():
|
||
"""Default guidance must mention all 11 canonical entity types."""
|
||
from lightrag.prompt import PROMPTS
|
||
|
||
guidance = PROMPTS["default_entity_types_guidance"]
|
||
expected_types = [
|
||
"Person",
|
||
"Creature",
|
||
"Organization",
|
||
"Location",
|
||
"Event",
|
||
"Concept",
|
||
"Method",
|
||
"Content",
|
||
"Data",
|
||
"Artifact",
|
||
"NaturalObject",
|
||
]
|
||
for t in expected_types:
|
||
assert t in guidance, (
|
||
f"Expected entity type '{t}' missing from default_entity_types_guidance"
|
||
)
|
||
|
||
|
||
@pytest.mark.offline
|
||
def test_builtin_entity_extraction_examples_are_format_only():
|
||
"""Built-in examples must be placeholder templates, not extractable demos.
|
||
|
||
Rather than denylisting specific sample names (brittle: any new concrete
|
||
content with different names would slip through), assert the structural
|
||
shape of a format-only template: no per-example section headers that would
|
||
reintroduce a sample ``---Input Text---`` / ``---Output---`` demo, and every
|
||
data value is an angle-bracket placeholder rather than concrete prose.
|
||
"""
|
||
from lightrag.prompt import PROMPTS
|
||
|
||
section_markers = ("---Input Text---", "---Output---", "---Entity Types---")
|
||
placeholder = re.compile(r"<[^<>]+>")
|
||
tuple_delimiter = PROMPTS["DEFAULT_TUPLE_DELIMITER"]
|
||
completion_delimiter = PROMPTS["DEFAULT_COMPLETION_DELIMITER"]
|
||
|
||
# Text examples: every field after the leading entity/relation tag must be a
|
||
# bare placeholder; concrete sample values would not match.
|
||
for example in PROMPTS["entity_extraction_examples"]:
|
||
for marker in section_markers:
|
||
assert marker not in example
|
||
rendered = example.format(
|
||
tuple_delimiter=tuple_delimiter,
|
||
completion_delimiter=completion_delimiter,
|
||
)
|
||
for line in rendered.splitlines():
|
||
line = line.strip()
|
||
if not line or line == completion_delimiter:
|
||
continue
|
||
tag, *fields = line.split(tuple_delimiter)
|
||
assert tag in {"entity", "relation"}
|
||
assert fields # data rows must carry at least one value field
|
||
for field in fields:
|
||
assert placeholder.fullmatch(field), field
|
||
|
||
# JSON examples: every entity/relationship field value must be a placeholder.
|
||
for example in PROMPTS["entity_extraction_json_examples"]:
|
||
for marker in section_markers:
|
||
assert marker not in example
|
||
parsed = json.loads(example)
|
||
records = parsed["entities"] + parsed["relationships"]
|
||
assert records
|
||
for record in records:
|
||
for value in record.values():
|
||
assert placeholder.fullmatch(value), value
|
||
|
||
|
||
@pytest.mark.offline
|
||
def test_entity_extraction_system_prompts_label_examples_as_format_templates():
|
||
from lightrag.prompt import PROMPTS
|
||
|
||
for prompt_key in (
|
||
"entity_extraction_system_prompt",
|
||
"entity_extraction_json_system_prompt",
|
||
):
|
||
prompt = PROMPTS[prompt_key]
|
||
assert "---Output Format Template---" in prompt
|
||
assert "---Examples---" not in prompt
|
||
assert "output format template only" in prompt
|
||
assert "not source text" in prompt
|
||
assert "must never be used as extraction content" in prompt
|
||
|
||
|
||
@pytest.mark.offline
|
||
def test_text_examples_render_tuple_and_completion_delimiters():
|
||
from lightrag.prompt import PROMPTS
|
||
|
||
rendered = "\n".join(PROMPTS["entity_extraction_examples"]).format(
|
||
tuple_delimiter=PROMPTS["DEFAULT_TUPLE_DELIMITER"],
|
||
completion_delimiter=PROMPTS["DEFAULT_COMPLETION_DELIMITER"],
|
||
)
|
||
|
||
assert (
|
||
"entity<|#|><entity_name><|#|><entity_type><|#|><entity_description>"
|
||
in rendered
|
||
)
|
||
assert (
|
||
"relation<|#|><source_entity><|#|><target_entity><|#|>"
|
||
"<relationship_keywords><|#|><relationship_description>" in rendered
|
||
)
|
||
assert "<|COMPLETE|>" in rendered
|
||
assert "{tuple_delimiter}" not in rendered
|
||
assert "{completion_delimiter}" not in rendered
|
||
|
||
|
||
@pytest.mark.offline
|
||
def test_json_examples_are_parseable_format_templates():
|
||
"""JSON examples must be raw JSON templates with valid endpoint references."""
|
||
from lightrag.prompt import PROMPTS
|
||
|
||
for example in PROMPTS["entity_extraction_json_examples"]:
|
||
parsed = json.loads(example)
|
||
assert set(parsed) == {"entities", "relationships"}
|
||
assert isinstance(parsed["entities"], list)
|
||
assert isinstance(parsed["relationships"], list)
|
||
assert parsed["entities"]
|
||
assert parsed["relationships"]
|
||
|
||
entity_names = {
|
||
entity["name"] for entity in parsed.get("entities", []) if entity
|
||
}
|
||
for relationship in parsed.get("relationships", []):
|
||
assert relationship["source"] in entity_names
|
||
assert relationship["target"] in entity_names
|
||
assert "<entity_name>" in entity_names
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# 2. DEFAULT_ENTITY_TYPES is gone from constants
|
||
# ---------------------------------------------------------------------------
|
||
|
||
|
||
@pytest.mark.offline
|
||
def test_default_entity_types_removed_from_constants():
|
||
"""DEFAULT_ENTITY_TYPES must no longer exist in lightrag.constants."""
|
||
import lightrag.constants as constants
|
||
|
||
assert not hasattr(constants, "DEFAULT_ENTITY_TYPES"), (
|
||
"DEFAULT_ENTITY_TYPES should have been removed from constants.py"
|
||
)
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# 3. ENTITY_TYPES env var raises SystemExit
|
||
# ---------------------------------------------------------------------------
|
||
|
||
|
||
@pytest.mark.offline
|
||
def test_entity_types_env_var_raises_system_exit(tmp_path):
|
||
"""LightRAG.__post_init__ must raise SystemExit when ENTITY_TYPES env var is set."""
|
||
from lightrag import LightRAG
|
||
|
||
with patch.dict(os.environ, {"ENTITY_TYPES": '["Person"]'}):
|
||
with pytest.raises(SystemExit) as exc_info:
|
||
LightRAG(
|
||
working_dir=str(tmp_path),
|
||
llm_model_func=AsyncMock(),
|
||
embedding_func=None,
|
||
)
|
||
assert "ENTITY_TYPES" in str(exc_info.value)
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# 4. Text mode: entity_types_guidance injected into prompt
|
||
# ---------------------------------------------------------------------------
|
||
|
||
|
||
@pytest.mark.offline
|
||
@pytest.mark.asyncio
|
||
async def test_text_mode_default_guidance_injected_into_prompt():
|
||
"""Default entity_types_guidance is passed to LLM system prompt in text mode."""
|
||
from lightrag.operate import extract_entities
|
||
from lightrag.prompt import PROMPTS
|
||
|
||
global_config = _make_global_config(use_json=False)
|
||
llm_func = global_config["llm_model_func"]
|
||
llm_func.return_value = _TEXT_MODE_RESPONSE
|
||
|
||
with patch("lightrag.operate.logger"):
|
||
await extract_entities(
|
||
chunks=_make_chunks(),
|
||
global_config=global_config,
|
||
)
|
||
|
||
# The system prompt passed to the LLM must contain the default guidance
|
||
assert llm_func.await_count >= 1
|
||
call_kwargs = llm_func.call_args_list[0][1]
|
||
system_prompt = call_kwargs.get("system_prompt", "")
|
||
assert PROMPTS["default_entity_types_guidance"] in system_prompt
|
||
assert "must start with `relation`, never `entity`" in system_prompt
|
||
assert "After the last entity row, switch prefixes to `relation`" in system_prompt
|
||
assert "Output at most 100 total rows" in system_prompt
|
||
assert "Output at most 40 entity rows" in system_prompt
|
||
|
||
|
||
@pytest.mark.offline
|
||
@pytest.mark.asyncio
|
||
async def test_text_mode_custom_guidance_overrides_default():
|
||
"""Custom entity_types_guidance in addon_params overrides default."""
|
||
from lightrag.operate import extract_entities
|
||
|
||
custom_guidance = "- Widget: A test widget type"
|
||
global_config = _make_global_config(
|
||
addon_params={"entity_types_guidance": custom_guidance},
|
||
use_json=False,
|
||
)
|
||
llm_func = global_config["llm_model_func"]
|
||
llm_func.return_value = _TEXT_MODE_RESPONSE
|
||
|
||
with patch("lightrag.operate.logger"):
|
||
await extract_entities(
|
||
chunks=_make_chunks(),
|
||
global_config=global_config,
|
||
)
|
||
|
||
call_kwargs = llm_func.call_args_list[0][1]
|
||
system_prompt = call_kwargs.get("system_prompt", "")
|
||
assert custom_guidance in system_prompt
|
||
|
||
|
||
@pytest.mark.offline
|
||
def test_text_continue_prompt_requires_relation_prefix_for_corrections():
|
||
from lightrag.prompt import PROMPTS
|
||
|
||
prompt = PROMPTS["entity_continue_extraction_user_prompt"]
|
||
assert (
|
||
"Any corrected relationship row must be emitted with the literal `relation` prefix"
|
||
in prompt
|
||
)
|
||
assert (
|
||
"output at most {max_total_records} total rows and at most {max_entity_records} entity rows"
|
||
in prompt
|
||
)
|
||
assert (
|
||
"may reference entities that were already extracted correctly in the previous response"
|
||
in prompt
|
||
)
|
||
assert (
|
||
"whose source and target entities are both included in this response"
|
||
not in prompt
|
||
)
|
||
|
||
|
||
@pytest.mark.offline
|
||
def test_text_user_prompt_includes_quantity_limits():
|
||
from lightrag.prompt import PROMPTS
|
||
|
||
prompt = PROMPTS["entity_extraction_user_prompt"]
|
||
assert (
|
||
"output at most {max_total_records} total rows and at most {max_entity_records} entity rows"
|
||
in prompt
|
||
)
|
||
assert (
|
||
"If the row limit is reached, output `{completion_delimiter}` immediately"
|
||
in prompt
|
||
)
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# 5. JSON mode: entity_types_guidance injected + entity_extraction kwarg set
|
||
# ---------------------------------------------------------------------------
|
||
|
||
|
||
@pytest.mark.offline
|
||
@pytest.mark.asyncio
|
||
async def test_rebuild_from_cached_fenced_json_uses_json_parser():
|
||
"""Cached JSON wrapped in markdown fences must not fall back to text parsing."""
|
||
from lightrag import operate
|
||
|
||
fenced_json = f"```json\n{_JSON_MODE_RESPONSE}\n```"
|
||
|
||
with patch(
|
||
"lightrag.operate._process_extraction_result",
|
||
side_effect=AssertionError("text parser should not be used"),
|
||
):
|
||
nodes, edges = await operate._rebuild_from_extraction_result(
|
||
text_chunks_storage=_DummyTextChunksStorage(),
|
||
extraction_result=fenced_json,
|
||
chunk_id="chunk-001",
|
||
timestamp=123,
|
||
)
|
||
|
||
assert set(nodes) == {"Alice", "Acme Corp"}
|
||
assert ("Alice", "Acme Corp") in edges
|
||
assert nodes["Alice"][0]["file_path"] == "test.md"
|
||
|
||
|
||
@pytest.mark.offline
|
||
@pytest.mark.asyncio
|
||
async def test_json_mode_default_guidance_injected_into_prompt():
|
||
"""Default entity_types_guidance is passed to LLM system prompt in JSON mode."""
|
||
from lightrag.operate import extract_entities
|
||
from lightrag.prompt import PROMPTS
|
||
|
||
global_config = _make_global_config(use_json=True)
|
||
llm_func = global_config["llm_model_func"]
|
||
llm_func.return_value = _JSON_MODE_RESPONSE
|
||
|
||
with patch("lightrag.operate.logger"):
|
||
await extract_entities(
|
||
chunks=_make_chunks(),
|
||
global_config=global_config,
|
||
)
|
||
|
||
assert llm_func.await_count >= 1
|
||
call_kwargs = llm_func.call_args_list[0][1]
|
||
system_prompt = call_kwargs.get("system_prompt", "")
|
||
assert PROMPTS["default_entity_types_guidance"] in system_prompt
|
||
assert "Output at most 100 total records" in system_prompt
|
||
assert "Output at most 40 entity objects" in system_prompt
|
||
|
||
|
||
@pytest.mark.offline
|
||
@pytest.mark.asyncio
|
||
async def test_json_mode_entity_extraction_kwarg_passed():
|
||
"""JSON mode must pass response_format={'type':'json_object'} to the LLM function."""
|
||
from lightrag.operate import extract_entities
|
||
|
||
global_config = _make_global_config(use_json=True)
|
||
llm_func = global_config["llm_model_func"]
|
||
llm_func.return_value = _JSON_MODE_RESPONSE
|
||
|
||
with patch("lightrag.operate.logger"):
|
||
await extract_entities(
|
||
chunks=_make_chunks(),
|
||
global_config=global_config,
|
||
)
|
||
|
||
assert llm_func.await_count >= 1
|
||
call_kwargs = llm_func.call_args_list[0][1]
|
||
assert call_kwargs.get("response_format") == {"type": "json_object"}
|
||
assert call_kwargs.get("entity_extraction") is not True
|
||
|
||
|
||
@pytest.mark.offline
|
||
@pytest.mark.asyncio
|
||
async def test_json_mode_custom_guidance_overrides_default():
|
||
"""Custom entity_types_guidance in addon_params overrides default in JSON mode."""
|
||
from lightrag.operate import extract_entities
|
||
|
||
custom_guidance = "- Gadget: A test gadget type"
|
||
global_config = _make_global_config(
|
||
addon_params={"entity_types_guidance": custom_guidance},
|
||
use_json=True,
|
||
)
|
||
llm_func = global_config["llm_model_func"]
|
||
llm_func.return_value = _JSON_MODE_RESPONSE
|
||
|
||
with patch("lightrag.operate.logger"):
|
||
await extract_entities(
|
||
chunks=_make_chunks(),
|
||
global_config=global_config,
|
||
)
|
||
|
||
call_kwargs = llm_func.call_args_list[0][1]
|
||
system_prompt = call_kwargs.get("system_prompt", "")
|
||
assert custom_guidance in system_prompt
|
||
|
||
|
||
@pytest.mark.offline
|
||
def test_json_user_prompt_includes_quantity_limits():
|
||
from lightrag.prompt import PROMPTS
|
||
|
||
prompt = PROMPTS["entity_extraction_json_user_prompt"]
|
||
assert (
|
||
"output at most {max_total_records} total records and at most {max_entity_records} entity objects"
|
||
in prompt
|
||
)
|
||
assert (
|
||
"Only output relationship objects whose `source` and `target` are both included"
|
||
in prompt
|
||
)
|
||
|
||
|
||
@pytest.mark.offline
|
||
def test_json_continue_prompt_includes_quantity_limits():
|
||
from lightrag.prompt import PROMPTS
|
||
|
||
prompt = PROMPTS["entity_continue_extraction_json_user_prompt"]
|
||
assert (
|
||
"output at most {max_total_records} total records and at most {max_entity_records} entity objects"
|
||
in prompt
|
||
)
|
||
assert (
|
||
"may reference entities already extracted correctly in the previous response"
|
||
in prompt
|
||
)
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# 6. Text mode: entity_extraction kwarg NOT passed (only JSON mode uses it)
|
||
# ---------------------------------------------------------------------------
|
||
|
||
|
||
@pytest.mark.offline
|
||
@pytest.mark.asyncio
|
||
async def test_text_mode_no_entity_extraction_kwarg():
|
||
"""Text mode must NOT pass entity_extraction=True to the LLM function."""
|
||
from lightrag.operate import extract_entities
|
||
|
||
global_config = _make_global_config(use_json=False)
|
||
llm_func = global_config["llm_model_func"]
|
||
llm_func.return_value = _TEXT_MODE_RESPONSE
|
||
|
||
with patch("lightrag.operate.logger"):
|
||
await extract_entities(
|
||
chunks=_make_chunks(),
|
||
global_config=global_config,
|
||
)
|
||
|
||
call_kwargs = llm_func.call_args_list[0][1]
|
||
assert not call_kwargs.get("entity_extraction", False)
|
||
|
||
|
||
@pytest.mark.offline
|
||
@pytest.mark.asyncio
|
||
async def test_text_mode_recovers_mis_prefixed_relationship_row():
|
||
from lightrag.operate import extract_entities
|
||
|
||
global_config = _make_global_config(use_json=False)
|
||
llm_func = global_config["llm_model_func"]
|
||
llm_func.return_value = _TEXT_MODE_MISPREFIXED_RELATION_RESPONSE
|
||
|
||
with patch("lightrag.operate.logger"):
|
||
chunk_results = await extract_entities(
|
||
chunks=_make_chunks(),
|
||
global_config=global_config,
|
||
)
|
||
|
||
entities, relationships = chunk_results[0]
|
||
assert len(entities) == 2
|
||
assert len(relationships) == 1
|
||
assert next(iter(relationships.keys())) == ("Alice", "Acme Corp")
|
||
|
||
|
||
@pytest.mark.offline
|
||
@pytest.mark.asyncio
|
||
async def test_text_mode_gleaned_relation_merges_cleanly_after_recovery():
|
||
from lightrag.operate import extract_entities
|
||
|
||
global_config = _make_global_config(use_json=False, max_gleaning=1)
|
||
llm_func = global_config["llm_model_func"]
|
||
llm_func.side_effect = _TEXT_MODE_GLEANED_RELATION_RESPONSES
|
||
|
||
with patch("lightrag.operate.logger"):
|
||
chunk_results = await extract_entities(
|
||
chunks=_make_chunks(),
|
||
global_config=global_config,
|
||
)
|
||
|
||
entities, relationships = chunk_results[0]
|
||
assert len(entities) == 2
|
||
assert len(relationships) == 1
|
||
relation_data = next(iter(relationships.values()))[0]
|
||
assert relation_data["src_id"] == "Alice"
|
||
assert relation_data["tgt_id"] == "Acme Corp"
|
||
|
||
|
||
@pytest.mark.offline
|
||
@pytest.mark.asyncio
|
||
async def test_text_mode_gleaned_relation_can_reference_prior_entity():
|
||
from lightrag.operate import extract_entities
|
||
|
||
global_config = _make_global_config(use_json=False, max_gleaning=1)
|
||
llm_func = global_config["llm_model_func"]
|
||
llm_func.side_effect = _TEXT_MODE_CROSS_PASS_RELATION_RESPONSES
|
||
|
||
with patch("lightrag.operate.logger"):
|
||
chunk_results = await extract_entities(
|
||
chunks=_make_chunks(),
|
||
global_config=global_config,
|
||
)
|
||
|
||
entities, relationships = chunk_results[0]
|
||
assert set(entities.keys()) == {"Alice", "Acme Corp"}
|
||
assert len(relationships) == 1
|
||
relation_data = next(iter(relationships.values()))[0]
|
||
assert relation_data["src_id"] == "Alice"
|
||
assert relation_data["tgt_id"] == "Acme Corp"
|
||
|
||
|
||
@pytest.mark.offline
|
||
def test_addon_params_default_includes_entity_type_prompt_file_env(tmp_path):
|
||
_require_yaml()
|
||
|
||
from lightrag import LightRAG
|
||
|
||
prompt_dir = tmp_path / "entity_type"
|
||
prompt_dir.mkdir()
|
||
_write_prompt_profile(
|
||
prompt_dir / "entity_type_prompt.sample.yml",
|
||
text_examples=[_text_profile_example("Env Default Example")],
|
||
)
|
||
|
||
with patch.dict(
|
||
os.environ,
|
||
{
|
||
"SUMMARY_LANGUAGE": "English",
|
||
"ENTITY_TYPE_PROMPT_FILE": "entity_type_prompt.sample.yml",
|
||
},
|
||
):
|
||
with _patch_prompt_dir(prompt_dir):
|
||
rag = LightRAG(
|
||
working_dir=str(tmp_path / "rag-default-env"),
|
||
llm_model_func=AsyncMock(),
|
||
embedding_func=_dummy_embedding_func(),
|
||
entity_extraction_use_json=False,
|
||
)
|
||
|
||
assert (
|
||
rag.addon_params["entity_type_prompt_file"] == "entity_type_prompt.sample.yml"
|
||
)
|
||
|
||
|
||
@pytest.mark.offline
|
||
@pytest.mark.asyncio
|
||
async def test_text_mode_prompt_file_injects_examples_and_guidance():
|
||
_require_yaml()
|
||
|
||
from lightrag.operate import extract_entities
|
||
|
||
guidance = "- ExampleType: Injected guidance"
|
||
example_label = "Custom Text Example"
|
||
prompt_profile = {
|
||
"entity_types_guidance": guidance,
|
||
"entity_extraction_examples": [_text_profile_example(example_label)],
|
||
"entity_extraction_json_examples": [],
|
||
}
|
||
|
||
global_config = _make_global_config(
|
||
prompt_profile=prompt_profile,
|
||
use_json=False,
|
||
)
|
||
llm_func = global_config["llm_model_func"]
|
||
llm_func.return_value = _TEXT_MODE_RESPONSE
|
||
|
||
with patch("lightrag.operate.logger"):
|
||
await extract_entities(chunks=_make_chunks(), global_config=global_config)
|
||
|
||
call_kwargs = llm_func.call_args_list[0][1]
|
||
system_prompt = call_kwargs.get("system_prompt", "")
|
||
assert guidance in system_prompt
|
||
assert example_label in system_prompt
|
||
|
||
|
||
@pytest.mark.offline
|
||
@pytest.mark.asyncio
|
||
async def test_json_mode_prompt_file_injects_examples_and_guidance():
|
||
_require_yaml()
|
||
|
||
from lightrag.operate import extract_entities
|
||
|
||
guidance = "- ExampleType: Injected JSON guidance"
|
||
example_label = "Custom Json Example"
|
||
prompt_profile = {
|
||
"entity_types_guidance": guidance,
|
||
"entity_extraction_examples": [],
|
||
"entity_extraction_json_examples": [_json_profile_example(example_label)],
|
||
}
|
||
|
||
global_config = _make_global_config(
|
||
prompt_profile=prompt_profile,
|
||
use_json=True,
|
||
)
|
||
llm_func = global_config["llm_model_func"]
|
||
llm_func.return_value = _JSON_MODE_RESPONSE
|
||
|
||
with patch("lightrag.operate.logger"):
|
||
await extract_entities(chunks=_make_chunks(), global_config=global_config)
|
||
|
||
call_kwargs = llm_func.call_args_list[0][1]
|
||
system_prompt = call_kwargs.get("system_prompt", "")
|
||
assert guidance in system_prompt
|
||
assert example_label in system_prompt
|
||
|
||
|
||
@pytest.mark.offline
|
||
@pytest.mark.asyncio
|
||
async def test_prompt_file_guidance_falls_back_to_default_when_missing():
|
||
_require_yaml()
|
||
|
||
from lightrag.operate import extract_entities
|
||
from lightrag.prompt import PROMPTS
|
||
|
||
global_config = _make_global_config(
|
||
prompt_profile={
|
||
"entity_types_guidance": PROMPTS["default_entity_types_guidance"].rstrip(),
|
||
"entity_extraction_examples": [
|
||
_text_profile_example("Fallback Guidance Example")
|
||
],
|
||
"entity_extraction_json_examples": [],
|
||
},
|
||
use_json=False,
|
||
)
|
||
llm_func = global_config["llm_model_func"]
|
||
llm_func.return_value = _TEXT_MODE_RESPONSE
|
||
|
||
with patch("lightrag.operate.logger"):
|
||
await extract_entities(chunks=_make_chunks(), global_config=global_config)
|
||
|
||
call_kwargs = llm_func.call_args_list[0][1]
|
||
system_prompt = call_kwargs.get("system_prompt", "")
|
||
assert PROMPTS["default_entity_types_guidance"] in system_prompt
|
||
|
||
|
||
@pytest.mark.offline
|
||
@pytest.mark.asyncio
|
||
async def test_cached_prompt_profile_supplies_merged_guidance():
|
||
from lightrag.operate import extract_entities
|
||
|
||
merged_guidance = "- ExampleType: Addon override"
|
||
|
||
global_config = _make_global_config(
|
||
prompt_profile={
|
||
"entity_types_guidance": merged_guidance,
|
||
"entity_extraction_examples": [_text_profile_example("Override Example")],
|
||
"entity_extraction_json_examples": [],
|
||
},
|
||
use_json=False,
|
||
)
|
||
llm_func = global_config["llm_model_func"]
|
||
llm_func.return_value = _TEXT_MODE_RESPONSE
|
||
|
||
with patch("lightrag.operate.logger"):
|
||
await extract_entities(chunks=_make_chunks(), global_config=global_config)
|
||
|
||
call_kwargs = llm_func.call_args_list[0][1]
|
||
system_prompt = call_kwargs.get("system_prompt", "")
|
||
assert merged_guidance in system_prompt
|
||
|
||
|
||
@pytest.mark.offline
|
||
def test_text_mode_prompt_file_can_omit_json_examples(tmp_path):
|
||
_require_yaml()
|
||
|
||
from lightrag import LightRAG
|
||
|
||
prompt_dir = tmp_path / "entity_type"
|
||
prompt_dir.mkdir()
|
||
_write_prompt_profile(
|
||
prompt_dir / "text_only.yml",
|
||
text_examples=[_text_profile_example("Text Only Example")],
|
||
)
|
||
|
||
with _patch_prompt_dir(prompt_dir):
|
||
rag = LightRAG(
|
||
working_dir=str(tmp_path / "rag-text"),
|
||
llm_model_func=AsyncMock(),
|
||
embedding_func=_dummy_embedding_func(),
|
||
entity_extraction_use_json=False,
|
||
addon_params={"entity_type_prompt_file": "text_only.yml"},
|
||
)
|
||
|
||
assert rag.addon_params["entity_type_prompt_file"] == "text_only.yml"
|
||
|
||
|
||
@pytest.mark.offline
|
||
def test_json_mode_prompt_file_can_omit_text_examples(tmp_path):
|
||
_require_yaml()
|
||
|
||
from lightrag import LightRAG
|
||
|
||
prompt_dir = tmp_path / "entity_type"
|
||
prompt_dir.mkdir()
|
||
_write_prompt_profile(
|
||
prompt_dir / "json_only.yml",
|
||
json_examples=[_json_profile_example("Json Only Example")],
|
||
)
|
||
|
||
with _patch_prompt_dir(prompt_dir):
|
||
rag = LightRAG(
|
||
working_dir=str(tmp_path / "rag-json"),
|
||
llm_model_func=AsyncMock(),
|
||
embedding_func=_dummy_embedding_func(),
|
||
entity_extraction_use_json=True,
|
||
addon_params={"entity_type_prompt_file": "json_only.yml"},
|
||
)
|
||
|
||
assert rag.addon_params["entity_type_prompt_file"] == "json_only.yml"
|
||
|
||
|
||
@pytest.mark.offline
|
||
def test_text_mode_prompt_file_requires_text_examples(tmp_path):
|
||
_require_yaml()
|
||
|
||
from lightrag import LightRAG
|
||
|
||
prompt_dir = tmp_path / "entity_type"
|
||
prompt_dir.mkdir()
|
||
_write_prompt_profile(
|
||
prompt_dir / "missing_text_examples.yml",
|
||
json_examples=[_json_profile_example("Wrong Mode Only")],
|
||
)
|
||
|
||
with _patch_prompt_dir(prompt_dir):
|
||
with pytest.raises(ValueError) as exc_info:
|
||
LightRAG(
|
||
working_dir=str(tmp_path / "rag-missing-text"),
|
||
llm_model_func=AsyncMock(),
|
||
embedding_func=None,
|
||
entity_extraction_use_json=False,
|
||
addon_params={"entity_type_prompt_file": "missing_text_examples.yml"},
|
||
)
|
||
|
||
assert "entity_extraction_examples" in str(exc_info.value)
|
||
|
||
|
||
@pytest.mark.offline
|
||
def test_json_mode_prompt_file_requires_json_examples(tmp_path):
|
||
_require_yaml()
|
||
|
||
from lightrag import LightRAG
|
||
|
||
prompt_dir = tmp_path / "entity_type"
|
||
prompt_dir.mkdir()
|
||
_write_prompt_profile(
|
||
prompt_dir / "missing_json_examples.yml",
|
||
text_examples=[_text_profile_example("Wrong Mode Only")],
|
||
)
|
||
|
||
with _patch_prompt_dir(prompt_dir):
|
||
with pytest.raises(ValueError) as exc_info:
|
||
LightRAG(
|
||
working_dir=str(tmp_path / "rag-missing-json"),
|
||
llm_model_func=AsyncMock(),
|
||
embedding_func=None,
|
||
entity_extraction_use_json=True,
|
||
addon_params={"entity_type_prompt_file": "missing_json_examples.yml"},
|
||
)
|
||
|
||
assert "entity_extraction_json_examples" in str(exc_info.value)
|
||
|
||
|
||
@pytest.mark.offline
|
||
def test_prompt_file_rejects_directory_segments(tmp_path):
|
||
_require_yaml()
|
||
|
||
from lightrag import LightRAG
|
||
|
||
with pytest.raises(ValueError) as exc_info:
|
||
LightRAG(
|
||
working_dir=str(tmp_path / "rag-bad-path"),
|
||
llm_model_func=AsyncMock(),
|
||
embedding_func=None,
|
||
addon_params={"entity_type_prompt_file": "../outside.yml"},
|
||
)
|
||
|
||
assert "file name only" in str(exc_info.value)
|
||
|
||
|
||
@pytest.mark.offline
|
||
def test_prompt_file_rejects_absolute_paths(tmp_path):
|
||
_require_yaml()
|
||
|
||
from lightrag import LightRAG
|
||
|
||
with pytest.raises(ValueError) as exc_info:
|
||
LightRAG(
|
||
working_dir=str(tmp_path / "rag-abs-path"),
|
||
llm_model_func=AsyncMock(),
|
||
embedding_func=None,
|
||
addon_params={"entity_type_prompt_file": str(tmp_path / "abs.yml")},
|
||
)
|
||
|
||
assert "file name only" in str(exc_info.value)
|
||
|
||
|
||
@pytest.mark.offline
|
||
@pytest.mark.asyncio
|
||
async def test_extract_entities_uses_cached_prompt_profile_without_reloading():
|
||
from lightrag.operate import extract_entities
|
||
|
||
cached_profile = {
|
||
"entity_types_guidance": "- ExampleType: Cached guidance",
|
||
"entity_extraction_examples": [_text_profile_example("Cached Text Example")],
|
||
"entity_extraction_json_examples": [],
|
||
}
|
||
global_config = _make_global_config(use_json=False, prompt_profile=cached_profile)
|
||
llm_func = global_config["llm_model_func"]
|
||
llm_func.return_value = _TEXT_MODE_RESPONSE
|
||
|
||
with patch(
|
||
"lightrag.operate.resolve_entity_extraction_prompt_profile",
|
||
side_effect=AssertionError("should not resolve profile when cache exists"),
|
||
):
|
||
with patch("lightrag.operate.logger"):
|
||
await extract_entities(chunks=_make_chunks(), global_config=global_config)
|
||
await extract_entities(chunks=_make_chunks(), global_config=global_config)
|
||
|
||
system_prompt = llm_func.call_args_list[0][1].get("system_prompt", "")
|
||
assert "Cached Text Example" in system_prompt
|
||
assert "Cached guidance" in system_prompt
|
||
|
||
|
||
@pytest.mark.offline
|
||
def test_sample_prompt_file_matches_builtin_prompt_data():
|
||
_require_yaml()
|
||
|
||
from lightrag.prompt import (
|
||
get_default_entity_extraction_prompt_profile,
|
||
load_entity_extraction_prompt_profile,
|
||
)
|
||
|
||
sample_file = (
|
||
Path(__file__).resolve().parents[2]
|
||
/ "prompts"
|
||
/ "samples"
|
||
/ "entity_type_prompt.sample.yml"
|
||
)
|
||
|
||
loaded_profile = load_entity_extraction_prompt_profile(sample_file)
|
||
assert loaded_profile == get_default_entity_extraction_prompt_profile()
|
||
|
||
|
||
@pytest.mark.offline
|
||
def test_prompt_dir_env_var_overrides_default(tmp_path, monkeypatch):
|
||
_require_yaml()
|
||
|
||
from lightrag.prompt import (
|
||
get_entity_type_prompt_dir,
|
||
resolve_entity_type_prompt_path,
|
||
)
|
||
|
||
monkeypatch.setenv("PROMPT_DIR", str(tmp_path))
|
||
expected_dir = (tmp_path / "entity_type").resolve()
|
||
assert get_entity_type_prompt_dir() == expected_dir
|
||
resolved = resolve_entity_type_prompt_path("custom.yml")
|
||
assert resolved == expected_dir / "custom.yml"
|
||
|
||
|
||
@pytest.mark.offline
|
||
def test_prompt_dir_defaults_to_cwd_relative(tmp_path, monkeypatch):
|
||
_require_yaml()
|
||
|
||
from lightrag.prompt import get_entity_type_prompt_dir
|
||
|
||
monkeypatch.delenv("PROMPT_DIR", raising=False)
|
||
monkeypatch.chdir(tmp_path)
|
||
assert (
|
||
get_entity_type_prompt_dir() == (tmp_path / "prompts" / "entity_type").resolve()
|
||
)
|
||
|
||
|
||
@pytest.mark.offline
|
||
def test_prompt_file_rejects_unsupported_extension(tmp_path):
|
||
_require_yaml()
|
||
|
||
from lightrag import LightRAG
|
||
|
||
with pytest.raises(ValueError, match="'.yml' or '.yaml'"):
|
||
LightRAG(
|
||
working_dir=str(tmp_path / "rag-bad-ext"),
|
||
llm_model_func=AsyncMock(),
|
||
embedding_func=None,
|
||
addon_params={"entity_type_prompt_file": "profile.txt"},
|
||
)
|
||
|
||
|
||
@pytest.mark.offline
|
||
def test_prompt_file_malformed_yaml_raises_valueerror(tmp_path):
|
||
_require_yaml()
|
||
|
||
from lightrag.prompt import load_entity_extraction_prompt_profile
|
||
|
||
bad_file = tmp_path / "broken.yml"
|
||
bad_file.write_text("entity_types_guidance: [unclosed", encoding="utf-8")
|
||
|
||
with pytest.raises(ValueError, match="invalid YAML"):
|
||
load_entity_extraction_prompt_profile(bad_file)
|
||
|
||
|
||
@pytest.mark.offline
|
||
def test_addon_guidance_overrides_file_profile(tmp_path):
|
||
_require_yaml()
|
||
|
||
from lightrag.prompt import resolve_entity_extraction_prompt_profile
|
||
|
||
prompt_dir = tmp_path / "entity_type"
|
||
prompt_dir.mkdir()
|
||
_write_prompt_profile(
|
||
prompt_dir / "profile.yml",
|
||
guidance="- FileType: from file",
|
||
text_examples=[_text_profile_example("Merged Example")],
|
||
)
|
||
|
||
with _patch_prompt_dir(prompt_dir):
|
||
profile = resolve_entity_extraction_prompt_profile(
|
||
addon_params={
|
||
"entity_type_prompt_file": "profile.yml",
|
||
"entity_types_guidance": "- AddonType: from addon_params",
|
||
},
|
||
use_json=False,
|
||
)
|
||
|
||
assert profile["entity_types_guidance"] == "- AddonType: from addon_params"
|
||
# File-provided examples must still be honored.
|
||
assert any(
|
||
"Merged Example" in example for example in profile["entity_extraction_examples"]
|
||
)
|
||
|
||
|
||
@pytest.mark.offline
|
||
def test_explicit_addon_params_still_picks_up_env_defaults(tmp_path, monkeypatch):
|
||
"""Passing addon_params explicitly must not drop env-based defaults."""
|
||
_require_yaml()
|
||
|
||
from lightrag import LightRAG
|
||
|
||
prompt_dir = tmp_path / "entity_type"
|
||
prompt_dir.mkdir()
|
||
_write_prompt_profile(
|
||
prompt_dir / "from_env.yml",
|
||
text_examples=[_text_profile_example("Env Example")],
|
||
)
|
||
|
||
monkeypatch.setenv("ENTITY_TYPE_PROMPT_FILE", "from_env.yml")
|
||
|
||
with _patch_prompt_dir(prompt_dir):
|
||
rag = LightRAG(
|
||
working_dir=str(tmp_path / "rag-env-default"),
|
||
llm_model_func=AsyncMock(),
|
||
embedding_func=_dummy_embedding_func(),
|
||
entity_extraction_use_json=False,
|
||
addon_params={"language": "English"},
|
||
)
|
||
|
||
assert rag.addon_params["entity_type_prompt_file"] == "from_env.yml"
|
||
|
||
|
||
@pytest.mark.offline
|
||
def test_runtime_addon_params_item_update_refreshes_cached_values(tmp_path):
|
||
_require_yaml()
|
||
|
||
from lightrag import LightRAG
|
||
|
||
prompt_dir = tmp_path / "entity_type"
|
||
prompt_dir.mkdir()
|
||
_write_prompt_profile(
|
||
prompt_dir / "initial.yml",
|
||
text_examples=[_text_profile_example("Initial Example")],
|
||
)
|
||
_write_prompt_profile(
|
||
prompt_dir / "updated.yml",
|
||
guidance="- UpdatedType: runtime update",
|
||
text_examples=[_text_profile_example("Updated Example")],
|
||
)
|
||
|
||
with _patch_prompt_dir(prompt_dir):
|
||
rag = LightRAG(
|
||
working_dir=str(tmp_path / "rag-runtime-update"),
|
||
llm_model_func=AsyncMock(),
|
||
embedding_func=_dummy_embedding_func(),
|
||
entity_extraction_use_json=False,
|
||
addon_params={
|
||
"entity_type_prompt_file": "initial.yml",
|
||
"language": "English",
|
||
},
|
||
)
|
||
rag.addon_params["entity_type_prompt_file"] = "updated.yml"
|
||
rag.addon_params["language"] = "French"
|
||
global_config = rag._build_global_config()
|
||
|
||
assert global_config["addon_params"]["language"] == "French"
|
||
assert global_config["_resolved_summary_language"] == "French"
|
||
assert (
|
||
global_config["_entity_extraction_prompt_profile"]["entity_types_guidance"]
|
||
== "- UpdatedType: runtime update"
|
||
)
|
||
assert any(
|
||
"Updated Example" in example
|
||
for example in global_config["_entity_extraction_prompt_profile"][
|
||
"entity_extraction_examples"
|
||
]
|
||
)
|
||
|
||
|
||
@pytest.mark.offline
|
||
def test_runtime_addon_params_replacement_refreshes_cached_values(tmp_path):
|
||
_require_yaml()
|
||
|
||
from lightrag import LightRAG
|
||
|
||
rag = LightRAG(
|
||
working_dir=str(tmp_path / "rag-runtime-replace"),
|
||
llm_model_func=AsyncMock(),
|
||
embedding_func=_dummy_embedding_func(),
|
||
entity_extraction_use_json=False,
|
||
addon_params={"language": "English"},
|
||
)
|
||
|
||
rag.addon_params = {
|
||
"language": "German",
|
||
"entity_types_guidance": "- ReplacementType: runtime replace",
|
||
}
|
||
global_config = rag._build_global_config()
|
||
|
||
assert global_config["addon_params"]["language"] == "German"
|
||
assert global_config["_resolved_summary_language"] == "German"
|
||
assert (
|
||
global_config["_entity_extraction_prompt_profile"]["entity_types_guidance"]
|
||
== "- ReplacementType: runtime replace"
|
||
)
|
||
|
||
|
||
@pytest.mark.offline
|
||
def test_runtime_mode_flip_invalidates_cached_prompt_profile(tmp_path):
|
||
_require_yaml()
|
||
|
||
from lightrag import LightRAG
|
||
|
||
prompt_dir = tmp_path / "entity_type"
|
||
prompt_dir.mkdir()
|
||
_write_prompt_profile(
|
||
prompt_dir / "text_only.yml",
|
||
text_examples=[_text_profile_example("Text Only Example")],
|
||
)
|
||
|
||
with _patch_prompt_dir(prompt_dir):
|
||
rag = LightRAG(
|
||
working_dir=str(tmp_path / "rag-mode-flip"),
|
||
llm_model_func=AsyncMock(),
|
||
embedding_func=_dummy_embedding_func(),
|
||
entity_extraction_use_json=False,
|
||
addon_params={"entity_type_prompt_file": "text_only.yml"},
|
||
)
|
||
|
||
rag._build_global_config()
|
||
rag.entity_extraction_use_json = True
|
||
|
||
with pytest.raises(ValueError) as exc_info:
|
||
rag._build_global_config()
|
||
|
||
assert "entity_extraction_json_examples" in str(exc_info.value)
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Section Context (heading breadcrumb) injection into extraction user prompts
|
||
# ---------------------------------------------------------------------------
|
||
|
||
_SECTION_MARKER = "---Section Context---"
|
||
|
||
|
||
def _render_text_user_prompt(heading_context_block: str) -> str:
|
||
from lightrag.prompt import PROMPTS
|
||
|
||
return PROMPTS["entity_extraction_user_prompt"].format(
|
||
max_total_records=100,
|
||
max_entity_records=40,
|
||
completion_delimiter="<|COMPLETE|>",
|
||
language="English",
|
||
input_text="Alice founded Acme Corp.",
|
||
heading_context_block=heading_context_block,
|
||
)
|
||
|
||
|
||
def _render_json_user_prompt(heading_context_block: str) -> str:
|
||
from lightrag.prompt import PROMPTS
|
||
|
||
return PROMPTS["entity_extraction_json_user_prompt"].format(
|
||
max_total_records=100,
|
||
max_entity_records=40,
|
||
language="English",
|
||
entity_types_guidance="- Person: humans",
|
||
input_text="Alice founded Acme Corp.",
|
||
heading_context_block=heading_context_block,
|
||
)
|
||
|
||
|
||
def _section_block(heading_path: str) -> str:
|
||
from lightrag.prompt import PROMPTS
|
||
|
||
return PROMPTS["entity_extraction_section_context"].format(
|
||
heading_path=heading_path
|
||
)
|
||
|
||
|
||
@pytest.mark.offline
|
||
def test_user_prompts_keep_single_real_input_text_section():
|
||
"""Only the rendered task prompt should carry the real input section marker."""
|
||
|
||
text_markers = [
|
||
line
|
||
for line in _render_text_user_prompt("").splitlines()
|
||
if line == "---Input Text---"
|
||
]
|
||
json_markers = [
|
||
line
|
||
for line in _render_json_user_prompt("").splitlines()
|
||
if line == "---Input Text---"
|
||
]
|
||
|
||
assert len(text_markers) == 1
|
||
assert len(json_markers) == 1
|
||
|
||
|
||
@pytest.mark.offline
|
||
def test_format_heading_context_full_path_includes_current_heading():
|
||
"""The breadcrumb appends the chunk's own heading after the parent chain."""
|
||
from lightrag.chunk_schema import format_heading_context
|
||
|
||
chunk = {
|
||
"content": "...",
|
||
"heading": {
|
||
"level": 2,
|
||
"heading": "Data Collection",
|
||
"parent_headings": ["Methods"],
|
||
},
|
||
}
|
||
assert format_heading_context(chunk) == "Methods → Data Collection"
|
||
|
||
|
||
@pytest.mark.offline
|
||
def test_format_heading_context_empty_when_no_heading():
|
||
"""A chunk without heading info yields an empty breadcrumb (block omitted)."""
|
||
from lightrag.chunk_schema import format_heading_context
|
||
|
||
chunk = {
|
||
"content": "...",
|
||
"tokens": 1,
|
||
"full_doc_id": "d",
|
||
"chunk_order_index": 0,
|
||
}
|
||
assert format_heading_context(chunk) == ""
|
||
|
||
|
||
@pytest.mark.offline
|
||
def test_text_user_prompt_section_context_hidden_and_byte_identical_when_no_heading():
|
||
"""No heading -> the whole `---Section Context---` block disappears and the
|
||
rendered text user prompt is byte-identical to the placeholder-free form."""
|
||
from lightrag.prompt import PROMPTS
|
||
|
||
rendered = _render_text_user_prompt("")
|
||
assert _SECTION_MARKER not in rendered
|
||
|
||
# The placeholder is the ONLY change to this template, so rendering it empty
|
||
# must equal a version with the placeholder physically removed (i.e. the
|
||
# pre-change template). This is the hard no-noise regression guard.
|
||
baseline_template = PROMPTS["entity_extraction_user_prompt"].replace(
|
||
"{heading_context_block}", ""
|
||
)
|
||
baseline = baseline_template.format(
|
||
max_total_records=100,
|
||
max_entity_records=40,
|
||
completion_delimiter="<|COMPLETE|>",
|
||
language="English",
|
||
input_text="Alice founded Acme Corp.",
|
||
)
|
||
assert rendered == baseline
|
||
|
||
|
||
@pytest.mark.offline
|
||
def test_json_user_prompt_section_context_hidden_and_byte_identical_when_no_heading():
|
||
from lightrag.prompt import PROMPTS
|
||
|
||
rendered = _render_json_user_prompt("")
|
||
assert _SECTION_MARKER not in rendered
|
||
|
||
baseline_template = PROMPTS["entity_extraction_json_user_prompt"].replace(
|
||
"{heading_context_block}", ""
|
||
)
|
||
baseline = baseline_template.format(
|
||
max_total_records=100,
|
||
max_entity_records=40,
|
||
language="English",
|
||
entity_types_guidance="- Person: humans",
|
||
input_text="Alice founded Acme Corp.",
|
||
)
|
||
assert rendered == baseline
|
||
|
||
|
||
@pytest.mark.offline
|
||
def test_text_user_prompt_includes_section_context_when_heading_present():
|
||
rendered = _render_text_user_prompt(_section_block("Methods → Data Collection"))
|
||
assert _SECTION_MARKER in rendered
|
||
assert "Methods → Data Collection" in rendered
|
||
# Block sits immediately above the input text section.
|
||
assert "Methods → Data Collection\n\n---Input Text---" in rendered
|
||
|
||
|
||
@pytest.mark.offline
|
||
def test_json_user_prompt_includes_section_context_when_heading_present():
|
||
rendered = _render_json_user_prompt(_section_block("Methods → Data Collection"))
|
||
assert _SECTION_MARKER in rendered
|
||
assert "Methods → Data Collection" in rendered
|
||
assert "Methods → Data Collection\n\n---Input Text---" in rendered
|
||
|
||
|
||
@pytest.mark.offline
|
||
def test_section_context_breadcrumb_is_not_at_line_start():
|
||
"""A heading that looks like a prompt marker must be rendered inline (as
|
||
data), never at the start of a line where it could forge a new section."""
|
||
block = _section_block("---Output---")
|
||
# The breadcrumb follows a label on the same line, so the marker text never
|
||
# begins a line of its own.
|
||
assert "\n---Output---" not in block
|
||
assert "---Output---" in block # still present, just inert/inline
|
||
|
||
|
||
@pytest.mark.offline
|
||
def test_extraction_system_prompts_reference_section_context():
|
||
"""Both system prompts carry the static conditional instruction."""
|
||
from lightrag.prompt import PROMPTS
|
||
|
||
for key in (
|
||
"entity_extraction_system_prompt",
|
||
"entity_extraction_json_system_prompt",
|
||
):
|
||
assert _SECTION_MARKER in PROMPTS[key]
|
||
assert "only as background" in PROMPTS[key]
|
||
|
||
|
||
@pytest.mark.offline
|
||
@pytest.mark.asyncio
|
||
async def test_extract_entities_injects_section_context_for_chunk_with_heading():
|
||
"""End-to-end: a chunk carrying a heading produces a user prompt containing
|
||
its full section breadcrumb; a heading-free chunk does not."""
|
||
from lightrag.operate import extract_entities
|
||
|
||
global_config = _make_global_config(use_json=False)
|
||
llm_func = global_config["llm_model_func"]
|
||
llm_func.return_value = _TEXT_MODE_RESPONSE
|
||
|
||
chunks = {
|
||
"chunk-001": {
|
||
"tokens": 10,
|
||
"content": "Alice founded Acme Corp.",
|
||
"full_doc_id": "doc-001",
|
||
"chunk_order_index": 0,
|
||
"heading": {
|
||
"level": 2,
|
||
"heading": "Data Collection",
|
||
"parent_headings": ["Methods"],
|
||
},
|
||
}
|
||
}
|
||
|
||
with patch("lightrag.operate.logger"):
|
||
await extract_entities(chunks=chunks, global_config=global_config)
|
||
|
||
assert llm_func.await_count >= 1
|
||
user_prompt = llm_func.call_args_list[0][0][0]
|
||
assert _SECTION_MARKER in user_prompt
|
||
assert "Methods → Data Collection\n\n---Input Text---" in user_prompt
|
||
|
||
|
||
@pytest.mark.offline
|
||
@pytest.mark.asyncio
|
||
async def test_extract_entities_omits_section_context_for_chunk_without_heading():
|
||
from lightrag.operate import extract_entities
|
||
|
||
global_config = _make_global_config(use_json=False)
|
||
llm_func = global_config["llm_model_func"]
|
||
llm_func.return_value = _TEXT_MODE_RESPONSE
|
||
|
||
with patch("lightrag.operate.logger"):
|
||
await extract_entities(chunks=_make_chunks(), global_config=global_config)
|
||
|
||
assert llm_func.await_count >= 1
|
||
user_prompt = llm_func.call_args_list[0][0][0]
|
||
assert _SECTION_MARKER not in user_prompt
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Section Context length bounding: per-level char cap + overall token budget
|
||
# ---------------------------------------------------------------------------
|
||
|
||
|
||
@pytest.mark.offline
|
||
def test_format_heading_context_caps_long_level():
|
||
"""A single runaway heading level is truncated to the per-level char cap."""
|
||
from lightrag.chunk_schema import (
|
||
DEFAULT_HEADING_LEVEL_MAX_CHARS,
|
||
format_heading_context,
|
||
)
|
||
|
||
long_title = "A" * (DEFAULT_HEADING_LEVEL_MAX_CHARS + 50)
|
||
chunk = {"heading": {"level": 1, "heading": long_title, "parent_headings": []}}
|
||
|
||
out = format_heading_context(chunk)
|
||
assert out.endswith("…")
|
||
assert len(out) == DEFAULT_HEADING_LEVEL_MAX_CHARS
|
||
|
||
|
||
@pytest.mark.offline
|
||
def test_format_heading_context_per_level_cap_can_be_disabled():
|
||
from lightrag.chunk_schema import format_heading_context
|
||
|
||
long_title = "B" * 300
|
||
chunk = {"heading": {"level": 1, "heading": long_title, "parent_headings": []}}
|
||
|
||
assert format_heading_context(chunk, max_heading_len=0) == long_title
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Query-stage format_parent_headings: same per-level cap + cleaning as extraction
|
||
# ---------------------------------------------------------------------------
|
||
|
||
|
||
@pytest.mark.offline
|
||
def test_format_parent_headings_caps_long_level():
|
||
"""A runaway parent heading is truncated to the per-level char cap, matching
|
||
the extraction-stage format_heading_context."""
|
||
from lightrag.chunk_schema import (
|
||
DEFAULT_HEADING_LEVEL_MAX_CHARS,
|
||
format_parent_headings,
|
||
)
|
||
|
||
long_title = "A" * (DEFAULT_HEADING_LEVEL_MAX_CHARS + 50)
|
||
chunk = {
|
||
"heading": {"level": 2, "heading": "Leaf", "parent_headings": [long_title]}
|
||
}
|
||
|
||
out = format_parent_headings(chunk)
|
||
assert out.endswith("…")
|
||
assert len(out) == DEFAULT_HEADING_LEVEL_MAX_CHARS # only the parent, capped
|
||
|
||
|
||
@pytest.mark.offline
|
||
def test_format_parent_headings_per_level_cap_can_be_disabled():
|
||
from lightrag.chunk_schema import format_parent_headings
|
||
|
||
long_title = "B" * 300
|
||
chunk = {
|
||
"heading": {"level": 2, "heading": "Leaf", "parent_headings": [long_title]}
|
||
}
|
||
|
||
assert format_parent_headings(chunk, max_heading_len=0) == long_title
|
||
|
||
|
||
@pytest.mark.offline
|
||
def test_format_parent_headings_cleaning_matches_extraction():
|
||
"""Parent headings get the same cleaning as extraction: → folded to a space,
|
||
Cc/Cf control chars stripped (shared _clean_heading_text)."""
|
||
from lightrag.chunk_schema import format_parent_headings
|
||
|
||
# chr(0) is a Cc control; chr(0x200B) is ZWSP (Cf) — both stripped. Built
|
||
# via chr() so the source carries no literal invisible characters.
|
||
second_level = "x" + chr(0) + "y" + chr(0x200B) + "z"
|
||
chunk = {
|
||
"heading": {
|
||
"level": 2,
|
||
"heading": "Leaf",
|
||
"parent_headings": ["A→B", second_level],
|
||
}
|
||
}
|
||
# "A→B" -> "A B"; control + format chars removed from the second level.
|
||
assert format_parent_headings(chunk) == "A B → xyz"
|
||
|
||
|
||
@pytest.mark.offline
|
||
def test_format_parent_headings_basic_behavior_preserved():
|
||
"""Existing behavior is unchanged: empty when no heading, normal multi-level
|
||
path joined with the breadcrumb separator."""
|
||
from lightrag.chunk_schema import format_parent_headings
|
||
|
||
assert format_parent_headings({"content": "...", "chunk_order_index": 0}) == ""
|
||
|
||
chunk = {
|
||
"heading": {"level": 2, "heading": "Leaf", "parent_headings": ["h1", "h2"]}
|
||
}
|
||
assert format_parent_headings(chunk) == "h1 → h2" # leaf NOT appended
|
||
|
||
|
||
class _FakeChunksDB:
|
||
"""Minimal text_chunks_db for _attach_content_headings: get_by_ids + config."""
|
||
|
||
def __init__(self, data_by_id: dict, tokenizer):
|
||
self._data = data_by_id
|
||
self.global_config = {"tokenizer": tokenizer}
|
||
|
||
async def get_by_ids(self, ids):
|
||
return [self._data.get(i) for i in ids]
|
||
|
||
|
||
@pytest.mark.offline
|
||
@pytest.mark.asyncio
|
||
async def test_attach_content_headings_token_budgets_deep_path():
|
||
"""A deep heading chain (per-level cap bounds length, not count) is collapsed
|
||
to fit DEFAULT_MAX_SECTION_CONTEXT_TOKENS, mirroring the extraction stage."""
|
||
from lightrag.chunk_schema import HEADING_BREADCRUMB_SEP
|
||
from lightrag.constants import DEFAULT_MAX_SECTION_CONTEXT_TOKENS
|
||
from lightrag.operate import _attach_content_headings
|
||
|
||
tok = Tokenizer("dummy", DummyTokenizer()) # 1 char == 1 token
|
||
deep = [f"Level{i:02d}" for i in range(100)] # well over the token budget
|
||
db = _FakeChunksDB(
|
||
{"c1": {"heading": {"level": 99, "heading": "Leaf", "parent_headings": deep}}},
|
||
tok,
|
||
)
|
||
chunks = [{"chunk_id": "c1"}]
|
||
|
||
await _attach_content_headings(chunks, db)
|
||
|
||
out = chunks[0]["content_headings"]
|
||
assert len(tok.encode(out)) <= DEFAULT_MAX_SECTION_CONTEXT_TOKENS
|
||
# Collapsed to first → … → leaf, so a middle level is gone.
|
||
assert f"{HEADING_BREADCRUMB_SEP}…{HEADING_BREADCRUMB_SEP}" in out
|
||
assert "Level50" not in out
|
||
|
||
|
||
@pytest.mark.offline
|
||
@pytest.mark.asyncio
|
||
async def test_attach_content_headings_keeps_short_path_intact():
|
||
"""A within-budget path is attached unchanged (no token collapsing)."""
|
||
from lightrag.operate import _attach_content_headings
|
||
|
||
tok = Tokenizer("dummy", DummyTokenizer())
|
||
db = _FakeChunksDB(
|
||
{
|
||
"c1": {
|
||
"heading": {
|
||
"level": 2,
|
||
"heading": "Leaf",
|
||
"parent_headings": ["h1", "h2"],
|
||
}
|
||
}
|
||
},
|
||
tok,
|
||
)
|
||
chunks = [{"chunk_id": "c1"}]
|
||
|
||
await _attach_content_headings(chunks, db)
|
||
|
||
assert chunks[0]["content_headings"] == "h1 → h2"
|
||
|
||
|
||
@pytest.mark.offline
|
||
def test_truncate_section_context_noop_within_budget():
|
||
from lightrag.operate import _truncate_section_context
|
||
|
||
tok = Tokenizer("dummy", DummyTokenizer())
|
||
path = "Methods → Data Collection"
|
||
assert _truncate_section_context(path, tok, 256) == path
|
||
|
||
|
||
@pytest.mark.offline
|
||
def test_truncate_section_context_keeps_first_and_last_when_over_budget():
|
||
"""Over budget -> keep first (top-level) + last (leaf) section, elide middle."""
|
||
from lightrag.chunk_schema import HEADING_BREADCRUMB_SEP
|
||
from lightrag.operate import _truncate_section_context
|
||
|
||
tok = Tokenizer("dummy", DummyTokenizer()) # 1 char == 1 token
|
||
levels = [f"Level{i:02d}" for i in range(100)]
|
||
path = HEADING_BREADCRUMB_SEP.join(levels)
|
||
# Budget large enough for the collapsed two-level form (~21 tokens) so the
|
||
# hard-cap backstop does not also fire here.
|
||
budget = 40
|
||
|
||
out = _truncate_section_context(path, tok, budget)
|
||
expected = (
|
||
f"{levels[0]}{HEADING_BREADCRUMB_SEP}…{HEADING_BREADCRUMB_SEP}{levels[-1]}"
|
||
)
|
||
assert out == expected
|
||
assert "Level50" not in out # middle levels are gone
|
||
assert len(tok.encode(out)) <= budget
|
||
|
||
|
||
@pytest.mark.offline
|
||
def test_truncate_section_context_hard_caps_dense_short_path():
|
||
"""A 1-/2-level path that is itself over budget must still be capped
|
||
(not bypassed) — guards token-dense / byte-level tokenizers."""
|
||
from lightrag.chunk_schema import HEADING_BREADCRUMB_SEP
|
||
from lightrag.operate import _truncate_section_context
|
||
|
||
tok = Tokenizer("dummy", DummyTokenizer())
|
||
path = HEADING_BREADCRUMB_SEP.join(["A" * 50, "B" * 50]) # 103 chars/tokens
|
||
budget = 10
|
||
|
||
out = _truncate_section_context(path, tok, budget)
|
||
assert out != path
|
||
assert out.endswith("…")
|
||
assert len(tok.encode(out)) <= budget
|
||
|
||
|
||
@pytest.mark.offline
|
||
def test_truncate_section_context_accounts_for_multitoken_ellipsis():
|
||
"""The hard cap must reserve the tokenizer's actual ellipsis cost."""
|
||
from lightrag.operate import _truncate_section_context
|
||
|
||
class TwoTokenEllipsisTokenizer(TokenizerInterface):
|
||
def encode(self, content: str):
|
||
tokens = []
|
||
for ch in content:
|
||
if ch == "…":
|
||
tokens.extend([0x110000, 0x110001])
|
||
else:
|
||
tokens.append(ord(ch))
|
||
return tokens
|
||
|
||
def decode(self, tokens):
|
||
return "".join(chr(token) for token in tokens if token <= 0x10FFFF)
|
||
|
||
tok = Tokenizer("two-token-ellipsis", TwoTokenEllipsisTokenizer())
|
||
budget = 10
|
||
|
||
out = _truncate_section_context("A" * 20, tok, budget)
|
||
assert out == "A" * 8 + "…"
|
||
assert len(tok.encode(out)) <= budget
|
||
|
||
|
||
@pytest.mark.offline
|
||
def test_truncate_section_context_hard_caps_collapsed_form_when_still_over():
|
||
"""Even the collapsed first→…→leaf form is capped if it still exceeds."""
|
||
from lightrag.chunk_schema import HEADING_BREADCRUMB_SEP
|
||
from lightrag.operate import _truncate_section_context
|
||
|
||
tok = Tokenizer("dummy", DummyTokenizer())
|
||
levels = [f"Level{i:02d}" for i in range(100)]
|
||
path = HEADING_BREADCRUMB_SEP.join(levels)
|
||
budget = 8 # smaller than the ~21-token collapsed form
|
||
|
||
out = _truncate_section_context(path, tok, budget)
|
||
assert out.endswith("…")
|
||
assert len(tok.encode(out)) <= budget
|
||
|
||
|
||
@pytest.mark.offline
|
||
def test_heading_level_cap_below_one_third_of_token_budget():
|
||
"""Invariant guard: collapsed first+leaf must fit the token budget."""
|
||
from lightrag.constants import (
|
||
DEFAULT_HEADING_LEVEL_MAX_CHARS,
|
||
DEFAULT_MAX_SECTION_CONTEXT_TOKENS,
|
||
)
|
||
|
||
assert DEFAULT_HEADING_LEVEL_MAX_CHARS * 3 < DEFAULT_MAX_SECTION_CONTEXT_TOKENS
|
||
|
||
|
||
@pytest.mark.offline
|
||
def test_truncate_section_context_disabled_or_no_tokenizer():
|
||
from lightrag.operate import _truncate_section_context
|
||
|
||
tok = Tokenizer("dummy", DummyTokenizer())
|
||
path = "X" * 1000
|
||
assert _truncate_section_context(path, tok, 0) == path
|
||
assert _truncate_section_context(path, None, 256) == path
|
||
|
||
|
||
@pytest.mark.offline
|
||
@pytest.mark.asyncio
|
||
async def test_extract_entities_bounds_pathological_heading_in_prompt():
|
||
"""A chunk with an absurdly long heading must not inject it verbatim."""
|
||
from lightrag.chunk_schema import DEFAULT_HEADING_LEVEL_MAX_CHARS
|
||
from lightrag.operate import extract_entities
|
||
|
||
global_config = _make_global_config(use_json=False)
|
||
llm_func = global_config["llm_model_func"]
|
||
llm_func.return_value = _TEXT_MODE_RESPONSE
|
||
|
||
long_title = "Z" * 500
|
||
chunks = {
|
||
"chunk-001": {
|
||
"tokens": 10,
|
||
"content": "Alice founded Acme Corp.",
|
||
"full_doc_id": "doc-001",
|
||
"chunk_order_index": 0,
|
||
"heading": {
|
||
"level": 1,
|
||
"heading": long_title,
|
||
"parent_headings": [],
|
||
},
|
||
}
|
||
}
|
||
|
||
with patch("lightrag.operate.logger"):
|
||
await extract_entities(chunks=chunks, global_config=global_config)
|
||
|
||
user_prompt = llm_func.call_args_list[0][0][0]
|
||
assert _SECTION_MARKER in user_prompt
|
||
assert long_title not in user_prompt # full title never reaches the prompt
|
||
assert "Z" * DEFAULT_HEADING_LEVEL_MAX_CHARS not in user_prompt
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Heading text symbol cleaning: → -> space, strip Cc/Cf, preserve everything else
|
||
# ---------------------------------------------------------------------------
|
||
|
||
|
||
@pytest.mark.offline
|
||
def test_clean_heading_text_converts_arrow_to_space():
|
||
"""The breadcrumb separator char must never survive inside one heading."""
|
||
from lightrag.chunk_schema import _clean_heading_text
|
||
|
||
assert _clean_heading_text("A→B") == "A B"
|
||
assert _clean_heading_text("A → B") == "A B"
|
||
|
||
|
||
@pytest.mark.offline
|
||
def test_clean_heading_text_strips_control_and_format_chars():
|
||
"""Cc (NUL, BEL, file/unit separators) and Cf (zero-width marks) are removed."""
|
||
from lightrag.chunk_schema import _clean_heading_text
|
||
|
||
# \x00 (Cc), ZWSP (Cf), BOM (Cf) all vanish.
|
||
assert _clean_heading_text("a\x00bc") == "abc"
|
||
assert _clean_heading_text("x\x07y") == "xy"
|
||
# \x1c-\x1f are Cc but NOT matched by \s — must be stripped, not kept.
|
||
assert _clean_heading_text("p\x1c\x1fq") == "pq"
|
||
|
||
|
||
@pytest.mark.offline
|
||
def test_clean_heading_text_preserves_normal_characters():
|
||
"""CJK / Latin / digits / punctuation are left untouched; only → is folded."""
|
||
from lightrag.chunk_schema import _clean_heading_text
|
||
|
||
assert _clean_heading_text("方法 → 数据采集 (2024)!") == "方法 数据采集 (2024)!"
|
||
# Adjacent CJK never gets a space inserted between characters.
|
||
assert _clean_heading_text("数据采集") == "数据采集"
|
||
|
||
|
||
@pytest.mark.offline
|
||
def test_clean_heading_text_whitespace_collapse_is_last():
|
||
"""Newline/tab still fold to a single space (kept through the strip pass)."""
|
||
from lightrag.chunk_schema import _clean_heading_text
|
||
|
||
assert _clean_heading_text("a\nb\tc") == "a b c"
|
||
# A control char removed between two words must not leave a double space.
|
||
assert _clean_heading_text("a \x00 b") == "a b"
|
||
|
||
|
||
@pytest.mark.offline
|
||
def test_format_heading_context_arrow_in_heading_does_not_forge_level():
|
||
"""A heading containing → is cleaned, so the breadcrumb split stays accurate."""
|
||
from lightrag.chunk_schema import (
|
||
HEADING_BREADCRUMB_SEP,
|
||
format_heading_context,
|
||
)
|
||
|
||
chunk = {
|
||
"heading": {"level": 2, "heading": "C", "parent_headings": ["A→B"]},
|
||
}
|
||
out = format_heading_context(chunk)
|
||
assert out == "A B → C"
|
||
# The breadcrumb still splits into exactly the two real levels.
|
||
assert out.split(HEADING_BREADCRUMB_SEP) == ["A B", "C"]
|