lightrag-cn-git/tests/pipeline/test_multimodal_surrounding_context.py

"""Unit tests for the native multimodal surrounding-context extractor.

See ``docs/NativeMultimodalSurroundingContextPlan-zh.md``.

These tests use a 1:1 character-token mapping so the expected token
budgets in each scenario stay obvious without coupling to tiktoken's
BPE.  The helper functions exercised here are pure (no async, no
network), so the suite runs offline.
"""

import json

import pytest

from lightrag.multimodal_context import (
    build_surrounding,
    enrich_sidecars_with_surrounding,
    find_target_span,
    load_chunk_separators,
)
from lightrag.utils import Tokenizer, TokenizerInterface


class _CharTokenizer(TokenizerInterface):
    def encode(self, content: str):
        return [ord(ch) for ch in content]

    def decode(self, tokens):
        return "".join(chr(t) for t in tokens)


def _tokenizer() -> Tokenizer:
    return Tokenizer(model_name="char", tokenizer=_CharTokenizer())


# ---------------------------------------------------------------------------
# Target-tag locator
# ---------------------------------------------------------------------------


@pytest.mark.offline
def test_find_target_span_drawing_in_mixed_content():
    content = (
        "leading text. "
        '<drawing id="im-abcd-0001" format="png" path="img.png" src="img" /> '
        "trailing text."
    )
    span = find_target_span("drawings", "im-abcd-0001", content)
    assert span is not None
    start, end = span
    assert content[start:end].startswith('<drawing id="im-abcd-0001"')
    assert content[start:end].endswith("/>")


@pytest.mark.offline
def test_find_target_span_table_with_id_anywhere_in_attrs():
    # id is not first attribute — locator must still find it.
    content = (
        'before <table format="json" id="tb-abcd-0007">[[1,2],[3,4]]</table> after'
    )
    span = find_target_span("tables", "tb-abcd-0007", content)
    assert span is not None
    snippet = content[span[0] : span[1]]
    assert snippet.endswith("</table>")
    assert 'id="tb-abcd-0007"' in snippet


@pytest.mark.offline
def test_find_target_span_table_cite_marker():
    content = 'before <cite type="table" refid="tb-abcd-0007">表1</cite> after'
    span = find_target_span("tables", "tb-abcd-0007", content)
    assert span is not None
    assert content[span[0] : span[1]].startswith("<cite")


@pytest.mark.offline
def test_find_target_span_equation():
    content = 'A <equation id="eq-abcd-0002" format="latex">x^2</equation> B'
    span = find_target_span("equations", "eq-abcd-0002", content)
    assert span is not None
    assert content[span[0] : span[1]].endswith("</equation>")


@pytest.mark.offline
def test_find_target_span_unknown_id_returns_none():
    content = '<drawing id="im-1" />'
    assert find_target_span("drawings", "im-other", content) is None


# ---------------------------------------------------------------------------
# Drawings & equations surrounding
# ---------------------------------------------------------------------------


@pytest.mark.offline
def test_drawing_surrounding_kept_within_block_only():
    tok = _tokenizer()
    block = (
        "paragraph one ends. paragraph two. "
        '<drawing id="im-1" path="a.png" src="a" /> '
        "paragraph three. paragraph four."
    )
    span = find_target_span("drawings", "im-1", block)
    surr = build_surrounding(
        kind="drawings",
        block_content=block,
        span=span,
        tokenizer=tok,
        leading_max_tokens=2000,
        trailing_max_tokens=2000,
        separators=load_chunk_separators(),
    )
    assert surr["leading"].endswith("paragraph two. ")
    assert surr["trailing"].startswith(" paragraph three.")


@pytest.mark.offline
def test_equation_surrounding_protects_drawing_atom():
    tok = _tokenizer()
    block = (
        '<drawing id="im-prev" path="a.png" src="a" caption="Fig 1" />'
        " intro text. "
        '<equation id="eq-1" format="latex">a+b=c</equation>'
        " conclusion text."
    )
    span = find_target_span("equations", "eq-1", block)
    surr = build_surrounding(
        kind="equations",
        block_content=block,
        span=span,
        tokenizer=tok,
        leading_max_tokens=2000,
        trailing_max_tokens=2000,
        separators=load_chunk_separators(),
    )
    # Parser-internal id/path/src are stripped, but caption survives and
    # the drawing tag stays atomic (not cut in half).
    assert '<drawing caption="Fig 1" />' in surr["leading"]
    assert "/>" in surr["leading"]
    # No half-open drawing/equation tags
    assert surr["leading"].count("<drawing") == surr["leading"].count("/>")


# ---------------------------------------------------------------------------
# Tables surrounding: other tables must be stripped before token counting.
# ---------------------------------------------------------------------------


@pytest.mark.offline
def test_table_surrounding_strips_other_tables_before_counting():
    tok = _tokenizer()
    block = (
        '<table id="tb-other" format="json">[["a","b"],["c","d"]]</table> '
        "narrative text describing the report. "
        '<table id="tb-target" format="json">[["x","y"]]</table>'
        " concluding remarks."
    )
    span = find_target_span("tables", "tb-target", block)
    surr = build_surrounding(
        kind="tables",
        block_content=block,
        span=span,
        tokenizer=tok,
        leading_max_tokens=2000,
        trailing_max_tokens=2000,
        separators=load_chunk_separators(),
    )
    # Sibling table must NOT appear in surrounding.
    assert "<table" not in surr["leading"]
    assert "</table>" not in surr["leading"]
    assert "<table" not in surr["trailing"]
    assert "narrative text" in surr["leading"]
    assert "concluding remarks" in surr["trailing"]


@pytest.mark.offline
def test_table_surrounding_supports_cite_marker_and_strips_sibling_cites():
    tok = _tokenizer()
    block = (
        'prefix <cite type="table" refid="tb-other">表0</cite> '
        'narrative <cite type="table" refid="tb-target">表1</cite> suffix'
    )
    span = find_target_span("tables", "tb-target", block)
    surr = build_surrounding(
        kind="tables",
        block_content=block,
        span=span,
        tokenizer=tok,
        leading_max_tokens=2000,
        trailing_max_tokens=2000,
        separators=load_chunk_separators(),
    )
    assert "tb-other" not in surr["leading"]
    assert "表0" not in surr["leading"]
    assert "narrative " in surr["leading"]
    assert surr["trailing"] == " suffix"


# ---------------------------------------------------------------------------
# Custom CHUNK_R_SEPARATORS via env
# ---------------------------------------------------------------------------


@pytest.mark.offline
def test_chunk_r_separators_env_drives_segment_boundary(monkeypatch):
    # Only the pipe character is a separator: text must split at '|'.
    monkeypatch.setenv("CHUNK_R_SEPARATORS", json.dumps(["|"]))
    seps = load_chunk_separators()
    assert seps == ["|"]
    tok = _tokenizer()
    # 3 segments separated by '|'; budget = 12 chars/tokens; each seg is
    # 10 chars including the trailing '|', so 1 whole segment fits, 2 do not.
    block = 'aaaaaaaaa|bbbbbbbbb|<drawing id="d" />|ccccccccc|ddddddddd'
    span = find_target_span("drawings", "d", block)
    surr = build_surrounding(
        kind="drawings",
        block_content=block,
        span=span,
        tokenizer=tok,
        leading_max_tokens=12,
        trailing_max_tokens=12,
        separators=seps,
    )
    # Leading should end at a '|' boundary (one whole segment), not be
    # char-truncated.
    assert surr["leading"].endswith("|")
    # And contain whole segment closest to target.
    assert "bbbbbbbbb|" in surr["leading"]


# ---------------------------------------------------------------------------
# Char fallback when the closest segment alone exceeds the budget.
# ---------------------------------------------------------------------------


@pytest.mark.offline
def test_oversized_closest_segment_char_truncated():
    tok = _tokenizer()
    # Single huge "segment" (no separator) right before the target.
    big = "X" * 5000
    block = big + '<drawing id="d" />'
    span = find_target_span("drawings", "d", block)
    surr = build_surrounding(
        kind="drawings",
        block_content=block,
        span=span,
        tokenizer=tok,
        leading_max_tokens=200,
        trailing_max_tokens=200,
        separators=load_chunk_separators(),
    )
    assert len(tok.encode(surr["leading"])) <= 200
    assert surr["trailing"] == ""
    # The suffix should be a tail of the X-run.
    assert surr["leading"].endswith("X")


@pytest.mark.offline
def test_oversized_trailing_char_truncated_at_head():
    tok = _tokenizer()
    big = "Y" * 5000
    block = '<drawing id="d" />' + big
    span = find_target_span("drawings", "d", block)
    surr = build_surrounding(
        kind="drawings",
        block_content=block,
        span=span,
        tokenizer=tok,
        leading_max_tokens=200,
        trailing_max_tokens=200,
        separators=load_chunk_separators(),
    )
    assert len(tok.encode(surr["trailing"])) <= 200
    assert surr["trailing"].startswith("Y")


# ---------------------------------------------------------------------------
# Drawings/equations surrounding: JSON / HTML table row trimming.
# ---------------------------------------------------------------------------


@pytest.mark.offline
def test_drawing_surrounding_row_trims_oversized_json_table():
    tok = _tokenizer()
    # 10 rows of repeating cells; whole table is ~> budget.
    rows = [[f"r{i}c0", f"r{i}c1"] for i in range(10)]
    big_table = '<table id="tb-big" format="json">' + json.dumps(rows) + "</table>"
    block = big_table + ' <drawing id="d" />'
    span = find_target_span("drawings", "d", block)
    # Budget chosen so only a few rows of the JSON table fit.
    surr = build_surrounding(
        kind="drawings",
        block_content=block,
        span=span,
        tokenizer=tok,
        leading_max_tokens=80,
        trailing_max_tokens=80,
        separators=load_chunk_separators(),
    )
    # Result must be a complete (smaller) <table>...</table>, contain
    # closing tag, and fit within budget.
    leading = surr["leading"]
    assert "<table " in leading
    assert (
        leading.rstrip().endswith("</table>")
        or leading.rstrip().endswith("</table> ")
        or "</table>" in leading
    )
    assert len(tok.encode(leading)) <= 80
    # Should keep tail rows (closest to target — last rows by index)
    assert "r9c0" in leading
    # Should not include rows from the far side.
    assert "r0c0" not in leading


@pytest.mark.offline
def test_drawing_surrounding_row_trims_oversized_html_table():
    tok = _tokenizer()
    rows_html = "".join(f"<tr><td>r{i}c0</td><td>r{i}c1</td></tr>" for i in range(10))
    body = f"<tbody>{rows_html}</tbody>"
    big_table = f'<table id="tb-h" format="html">{body}</table>'
    block = f'<drawing id="d" /> {big_table}'
    span = find_target_span("drawings", "d", block)
    surr = build_surrounding(
        kind="drawings",
        block_content=block,
        span=span,
        tokenizer=tok,
        leading_max_tokens=120,
        trailing_max_tokens=120,
        separators=load_chunk_separators(),
    )
    trailing = surr["trailing"]
    assert "<table " in trailing
    assert "</table>" in trailing
    assert "<tbody>" in trailing
    assert "</tbody>" in trailing
    assert len(tok.encode(trailing)) <= 120
    # For trailing we keep head rows.
    assert "r0c0" in trailing
    assert "r9c0" not in trailing


@pytest.mark.offline
def test_drawing_surrounding_char_trims_oversized_single_json_row():
    tok = _tokenizer()
    row_text = "A" * 200 + "TAIL"
    big_table = (
        '<table id="tb-big" format="json">'
        + json.dumps([[row_text]], ensure_ascii=False)
        + "</table>"
    )
    block = big_table + '<drawing id="d" />'
    span = find_target_span("drawings", "d", block)
    surr = build_surrounding(
        kind="drawings",
        block_content=block,
        span=span,
        tokenizer=tok,
        leading_max_tokens=90,
        trailing_max_tokens=90,
        separators=load_chunk_separators(),
    )

    leading = surr["leading"]
    assert leading.startswith("<table ")
    assert leading.endswith("</table>")
    assert "TAIL" in leading
    assert len(tok.encode(leading)) <= 90

    body = leading[leading.index(">") + 1 : -len("</table>")]
    parsed = json.loads(body)
    assert isinstance(parsed, list)


@pytest.mark.offline
def test_drawing_surrounding_char_trims_oversized_single_html_row():
    tok = _tokenizer()
    row_text = "HEAD" + "B" * 200
    big_table = (
        '<table id="tb-h" format="html">'
        f"<tbody><tr><td>{row_text}</td></tr></tbody>"
        "</table>"
    )
    block = f'<drawing id="d" />{big_table}'
    span = find_target_span("drawings", "d", block)
    surr = build_surrounding(
        kind="drawings",
        block_content=block,
        span=span,
        tokenizer=tok,
        leading_max_tokens=100,
        trailing_max_tokens=100,
        separators=load_chunk_separators(),
    )

    trailing = surr["trailing"]
    assert trailing.startswith("<table ")
    assert trailing.endswith("</table>")
    assert "<tr><td>" in trailing
    assert "HEAD" in trailing
    assert len(tok.encode(trailing)) <= 100


# ---------------------------------------------------------------------------
# enrich_sidecars_with_surrounding: idempotency + modality gating.
# ---------------------------------------------------------------------------


def _write_blocks(tmp_path, base, blocks):
    blocks_path = tmp_path / f"{base}.blocks.jsonl"
    lines = [json.dumps({"type": "meta", "format": "lightrag"})]
    for b in blocks:
        lines.append(json.dumps(b, ensure_ascii=False))
    blocks_path.write_text("\n".join(lines) + "\n", encoding="utf-8")
    return blocks_path


def _write_sidecar(path, root_key, items):
    path.write_text(
        json.dumps(
            {"version": "1.0", root_key: items},
            ensure_ascii=False,
            indent=2,
        ),
        encoding="utf-8",
    )


@pytest.mark.offline
def test_enrich_only_updates_enabled_modalities(tmp_path):
    base = "doc"
    blockid = "b1"
    content = (
        "intro. "
        '<drawing id="im-1" path="a.png" src="a" />'
        " middle "
        '<table id="tb-1" format="json">[["a"]]</table>'
        " tail "
        '<equation id="eq-1" format="latex">e</equation>'
        " end."
    )
    _write_blocks(
        tmp_path,
        base,
        [
            {
                "type": "content",
                "blockid": blockid,
                "format": "plain_text",
                "content": content,
                "heading": "h",
                "parent_headings": [],
                "level": 1,
            }
        ],
    )
    drawings_path = tmp_path / f"{base}.drawings.json"
    tables_path = tmp_path / f"{base}.tables.json"
    equations_path = tmp_path / f"{base}.equations.json"
    _write_sidecar(
        drawings_path,
        "drawings",
        {"im-1": {"id": "im-1", "blockid": blockid, "heading": "h"}},
    )
    _write_sidecar(
        tables_path,
        "tables",
        {"tb-1": {"id": "tb-1", "blockid": blockid, "heading": "h"}},
    )
    _write_sidecar(
        equations_path,
        "equations",
        {"eq-1": {"id": "eq-1", "blockid": blockid, "heading": "h"}},
    )

    counts = enrich_sidecars_with_surrounding(
        blocks_path=str(tmp_path / f"{base}.blocks.jsonl"),
        enabled_modalities={"drawings"},
        tokenizer=_tokenizer(),
        leading_max_tokens=2000,
        trailing_max_tokens=2000,
    )
    assert counts["drawings"] == 1
    assert counts["tables"] == 0
    assert counts["equations"] == 0

    drawings = json.loads(drawings_path.read_text(encoding="utf-8"))
    tables = json.loads(tables_path.read_text(encoding="utf-8"))
    equations = json.loads(equations_path.read_text(encoding="utf-8"))
    assert "surrounding" in drawings["drawings"]["im-1"]
    assert drawings["drawings"]["im-1"]["surrounding"]["leading"].startswith("intro.")
    assert "surrounding" not in tables["tables"]["tb-1"]
    assert "surrounding" not in equations["equations"]["eq-1"]


@pytest.mark.offline
def test_enrich_runs_even_when_llm_analyze_result_present(tmp_path):
    """Idempotency: existing ``llm_analyze_result`` does not block
    surrounding backfill — we treat the two fields as independent."""
    base = "doc"
    blockid = "b1"
    content = 'prefix. <drawing id="im-1" path="a.png" src="a" /> suffix.'
    _write_blocks(
        tmp_path,
        base,
        [
            {
                "type": "content",
                "blockid": blockid,
                "format": "plain_text",
                "content": content,
                "heading": "h",
                "parent_headings": [],
                "level": 1,
            }
        ],
    )
    drawings_path = tmp_path / f"{base}.drawings.json"
    _write_sidecar(
        drawings_path,
        "drawings",
        {
            "im-1": {
                "id": "im-1",
                "blockid": blockid,
                "heading": "h",
                "llm_analyze_result": {
                    "name": "x",
                    "summary": "",
                    "detail_description": "",
                },
            }
        },
    )

    counts = enrich_sidecars_with_surrounding(
        blocks_path=str(tmp_path / f"{base}.blocks.jsonl"),
        enabled_modalities={"drawings"},
        tokenizer=_tokenizer(),
        leading_max_tokens=2000,
        trailing_max_tokens=2000,
    )
    assert counts["drawings"] == 1
    payload = json.loads(drawings_path.read_text(encoding="utf-8"))
    item = payload["drawings"]["im-1"]
    assert item["llm_analyze_result"]["name"] == "x"  # untouched
    assert item["surrounding"]["leading"].startswith("prefix.")
    assert item["surrounding"]["trailing"].startswith(" suffix.")


@pytest.mark.offline
def test_enrich_does_not_cross_block_boundaries(tmp_path):
    base = "doc"
    block_a = "earlier block content."
    block_b = 'later block. <drawing id="im-1" path="a.png" src="a" /> tail.'
    _write_blocks(
        tmp_path,
        base,
        [
            {
                "type": "content",
                "blockid": "bA",
                "format": "plain_text",
                "content": block_a,
                "heading": "h1",
                "parent_headings": [],
                "level": 1,
            },
            {
                "type": "content",
                "blockid": "bB",
                "format": "plain_text",
                "content": block_b,
                "heading": "h2",
                "parent_headings": [],
                "level": 1,
            },
        ],
    )
    drawings_path = tmp_path / f"{base}.drawings.json"
    _write_sidecar(
        drawings_path,
        "drawings",
        {"im-1": {"id": "im-1", "blockid": "bB", "heading": "h2"}},
    )

    enrich_sidecars_with_surrounding(
        blocks_path=str(tmp_path / f"{base}.blocks.jsonl"),
        enabled_modalities={"drawings"},
        tokenizer=_tokenizer(),
        leading_max_tokens=2000,
        trailing_max_tokens=2000,
    )
    payload = json.loads(drawings_path.read_text(encoding="utf-8"))
    surr = payload["drawings"]["im-1"]["surrounding"]
    # Must come from block B only — content of block A absent.
    assert "earlier block content" not in surr["leading"]
    assert surr["leading"].startswith("later block.")


# ---------------------------------------------------------------------------
# Per-half token budgets via SURROUNDING_LEADING/TRAILING_MAX_TOKENS env vars.
# ---------------------------------------------------------------------------


@pytest.mark.offline
def test_env_var_leading_and_trailing_budgets_apply_independently(
    tmp_path, monkeypatch
):
    # Asymmetric budgets must produce asymmetric leading / trailing sizes.
    monkeypatch.setenv("SURROUNDING_LEADING_MAX_TOKENS", "5")
    monkeypatch.setenv("SURROUNDING_TRAILING_MAX_TOKENS", "20")

    base = "doc"
    blockid = "b1"
    content = "X" * 200 + '<drawing id="im-1" path="a.png" src="a" />' + "Y" * 200
    _write_blocks(
        tmp_path,
        base,
        [
            {
                "type": "content",
                "blockid": blockid,
                "format": "plain_text",
                "content": content,
                "heading": "h",
                "parent_headings": [],
                "level": 1,
            }
        ],
    )
    drawings_path = tmp_path / f"{base}.drawings.json"
    _write_sidecar(
        drawings_path,
        "drawings",
        {"im-1": {"id": "im-1", "blockid": blockid, "heading": "h"}},
    )

    tok = _tokenizer()
    enrich_sidecars_with_surrounding(
        blocks_path=str(tmp_path / f"{base}.blocks.jsonl"),
        enabled_modalities={"drawings"},
        tokenizer=tok,
    )

    surr = json.loads(drawings_path.read_text(encoding="utf-8"))["drawings"]["im-1"][
        "surrounding"
    ]
    assert len(tok.encode(surr["leading"])) <= 5
    assert len(tok.encode(surr["trailing"])) <= 20
    # Trailing is allowed to use its larger budget, so it must be strictly
    # longer than leading here.
    assert len(surr["trailing"]) > len(surr["leading"])


# ---------------------------------------------------------------------------
# Parser-internal markup stripping inside surrounding (mirrors what
# ``strip_internal_multimodal_markup_for_extraction`` does for chunk
# content before entity extraction).  The cleaning happens *before*
# token-budgeted truncation, so the saved budget reflects what the
# LLM actually receives and a truncation point can never land inside
# an unprocessed ``id="…"`` attribute.
# ---------------------------------------------------------------------------


@pytest.mark.offline
def test_surrounding_strips_drawing_id_path_src():
    tok = _tokenizer()
    block = (
        "leading prose. "
        '<drawing id="im-x" path="figs/a.png" src="raw/a.png" caption="Fig 1" />'
        " between. "
        '<equation id="eq-target" format="latex">x=1</equation>'
        " trailing prose."
    )
    span = find_target_span("equations", "eq-target", block)
    surr = build_surrounding(
        kind="equations",
        block_content=block,
        span=span,
        tokenizer=tok,
        leading_max_tokens=2000,
        trailing_max_tokens=2000,
        separators=load_chunk_separators(),
    )
    leading = surr["leading"]
    assert '<drawing caption="Fig 1" />' in leading
    assert 'id="im-x"' not in leading
    assert "path=" not in leading
    assert "src=" not in leading


@pytest.mark.offline
def test_surrounding_strips_table_internal_id():
    tok = _tokenizer()
    block = (
        "prefix. "
        '<table id="tb-x" format="json" caption="Sales">[[1,2],[3,4]]</table>'
        " between. "
        '<drawing id="im-target" caption="Fig 2" />'
        " suffix."
    )
    span = find_target_span("drawings", "im-target", block)
    surr = build_surrounding(
        kind="drawings",
        block_content=block,
        span=span,
        tokenizer=tok,
        leading_max_tokens=2000,
        trailing_max_tokens=2000,
        separators=load_chunk_separators(),
    )
    leading = surr["leading"]
    assert '<table format="json" caption="Sales">[[1,2],[3,4]]</table>' in leading
    assert 'id="tb-x"' not in leading


@pytest.mark.offline
def test_surrounding_strips_cite_refid_keeping_visible_text():
    tok = _tokenizer()
    block = (
        "see "
        '<cite type="table" refid="tb-x">Table 1</cite>'
        " for details. "
        '<drawing id="im-target" caption="Fig 3" />'
        " end."
    )
    span = find_target_span("drawings", "im-target", block)
    surr = build_surrounding(
        kind="drawings",
        block_content=block,
        span=span,
        tokenizer=tok,
        leading_max_tokens=2000,
        trailing_max_tokens=2000,
        separators=load_chunk_separators(),
    )
    leading = surr["leading"]
    # Surrounding path uses keep_cite_tag=True: the cite wrapper survives
    # (so the VLM/LLM can tell "Table 1" is a reference to an external
    # table, not inline prose) but the parser-internal refid is gone.
    assert '<cite type="table">Table 1</cite>' in leading
    assert "refid=" not in leading
    assert "tb-x" not in leading


@pytest.mark.offline
def test_surrounding_keeps_equation_cite_tag_and_strips_refid():
    """In production, equations without LaTeX content emit as
    ``<cite type="equation" refid="eq-…">公式 N</cite>`` rather than a
    full ``<equation>`` tag.  Surrounding must keep the wrapper so the
    multimodal analyzer can recognize the visible label as an external
    referent, not inline prose."""
    tok = _tokenizer()
    block = (
        "see "
        '<cite type="equation" refid="eq-y">公式 2</cite>'
        " above. "
        '<drawing id="im-target" caption="Fig 4" />'
        " end."
    )
    span = find_target_span("drawings", "im-target", block)
    surr = build_surrounding(
        kind="drawings",
        block_content=block,
        span=span,
        tokenizer=tok,
        leading_max_tokens=2000,
        trailing_max_tokens=2000,
        separators=load_chunk_separators(),
    )
    leading = surr["leading"]
    assert '<cite type="equation">公式 2</cite>' in leading
    assert "refid=" not in leading
    assert "eq-y" not in leading


@pytest.mark.offline
def test_strip_happens_before_budget_truncation():
    """Regression guard for the strip-before-truncate ordering.

    Constructs a leading source whose raw form (with id/path/src) exceeds
    the budget while its stripped form fits.  If strip ran *after*
    truncation, the budget would be measured against the bloated raw
    string and the saved surrounding would be cut early (possibly mid-
    attribute, leaving ``id="…`` residue).
    """
    tok = _tokenizer()
    # Raw drawing tag including attrs (~67 chars), stripped form is
    # just '<drawing caption="C" />' (~24 chars).  Budget at 30 sits
    # between the two — raw is too big, stripped fits.
    block = (
        '<drawing id="im-prev" path="some/long/path.png" src="raw/long/path.png"'
        ' caption="C" />'
        '<equation id="eq-1" format="latex">y</equation>'
        " tail."
    )
    span = find_target_span("equations", "eq-1", block)
    surr = build_surrounding(
        kind="equations",
        block_content=block,
        span=span,
        tokenizer=tok,
        leading_max_tokens=30,
        trailing_max_tokens=2000,
        separators=load_chunk_separators(),
    )
    leading = surr["leading"]
    # Whole stripped tag must be present — proves strip ran before
    # the budget gate.
    assert leading == '<drawing caption="C" />'
    # And no parser-internal markers leaked through.
    assert "id=" not in leading
    assert "path=" not in leading
    assert "src=" not in leading


@pytest.mark.offline
def test_enrich_overwrites_surrounding_when_budget_changes(tmp_path):
    """Idempotency: rerunning with a smaller budget overwrites the prior
    surrounding, demonstrating that ``SURROUNDING_LEADING_MAX_TOKENS``
    changes propagate without needing to clear sidecars first."""
    base = "doc"
    blockid = "b1"
    content = "L" * 500 + '<drawing id="im-1" caption="C" />' + "T" * 500
    _write_blocks(
        tmp_path,
        base,
        [
            {
                "type": "content",
                "blockid": blockid,
                "format": "plain_text",
                "content": content,
                "heading": "h",
                "parent_headings": [],
                "level": 1,
            }
        ],
    )
    drawings_path = tmp_path / f"{base}.drawings.json"
    _write_sidecar(
        drawings_path,
        "drawings",
        {"im-1": {"id": "im-1", "blockid": blockid, "heading": "h"}},
    )

    tok = _tokenizer()
    enrich_sidecars_with_surrounding(
        blocks_path=str(tmp_path / f"{base}.blocks.jsonl"),
        enabled_modalities={"drawings"},
        tokenizer=tok,
        leading_max_tokens=300,
        trailing_max_tokens=300,
    )
    first = json.loads(drawings_path.read_text(encoding="utf-8"))["drawings"]["im-1"][
        "surrounding"
    ]
    first_leading_len = len(first["leading"])
    first_trailing_len = len(first["trailing"])

    enrich_sidecars_with_surrounding(
        blocks_path=str(tmp_path / f"{base}.blocks.jsonl"),
        enabled_modalities={"drawings"},
        tokenizer=tok,
        leading_max_tokens=50,
        trailing_max_tokens=50,
    )
    second = json.loads(drawings_path.read_text(encoding="utf-8"))["drawings"]["im-1"][
        "surrounding"
    ]
    # New budget is smaller, so saved surrounding must shrink — proving
    # the previous value was overwritten, not preserved.
    assert len(second["leading"]) < first_leading_len
    assert len(second["trailing"]) < first_trailing_len
    assert len(tok.encode(second["leading"])) <= 50
    assert len(tok.encode(second["trailing"])) <= 50


@pytest.mark.offline
def test_env_var_invalid_value_falls_back_to_default(monkeypatch):
    # An unparseable env value must not crash; it falls back to 2000.
    monkeypatch.setenv("SURROUNDING_LEADING_MAX_TOKENS", "not-a-number")
    monkeypatch.setenv("SURROUNDING_TRAILING_MAX_TOKENS", "not-a-number")
    from lightrag.multimodal_context import (
        DEFAULT_SURROUNDING_MAX_TOKENS,
        _resolve_surrounding_budget,
    )

    leading, trailing = _resolve_surrounding_budget(None, None)
    assert leading == DEFAULT_SURROUNDING_MAX_TOKENS
    assert trailing == DEFAULT_SURROUNDING_MAX_TOKENS