replypier/backend/tests/tools/test_knowledge_search_chunks.py

"""
Copyright 2024-2026 ReplyPier

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.

Search over chunked pages: neighbour expansion, page grouping, the size cap
and "a repeat search returns what has not been shown yet".
"""

from types import SimpleNamespace
from unittest.mock import MagicMock, patch
from uuid import uuid4

import pytest
from agno.document import Document

from app.models.knowledge import SourceType
from app.tools.knowledge_search_byagent import NO_FURTHER_RESULTS, KnowledgeSearchByAgent

PAGE = "https://x.com/pricing"


def _chunk(page: str, n: int, count: int, text: str) -> Document:
    return Document(
        id=f"{page}::{n}",
        name="https://x.com",
        content=text,
        meta_data={"url": page, "page_id": page, "chunk": n, "chunk_count": count},
    )


def _row(doc: Document):
    return SimpleNamespace(id=doc.id, name=doc.name, content=doc.content, meta_data=doc.meta_data)


@pytest.fixture
def tool():
    source = MagicMock()
    source.source = "https://x.com"
    source.source_type = SourceType.WEBSITE
    source.table_name = "d_org"
    source.schema = "ai"
    with patch("app.tools.knowledge_search_byagent.SessionLocal") as session_local, \
         patch("app.tools.knowledge_search_byagent.KnowledgeRepository") as repo_class:
        db = MagicMock()
        session_local.return_value.__enter__.return_value = db
        repo_class.return_value.get_by_agent.return_value = [source]
        tool = KnowledgeSearchByAgent(agent_id=str(uuid4()), org_id=uuid4())
        tool.agent_knowledge = MagicMock()
        tool._db = db
        yield tool


def test_hit_comes_with_page_head_and_neighbours_grouped_and_ordered(tool):
    hit = _chunk(PAGE, 4, 6, "part four")
    tool.agent_knowledge.search.return_value = [hit]
    tool._db.execute.return_value.fetchall.return_value = [
        _row(_chunk(PAGE, 5, 6, "part five")),
        _row(_chunk(PAGE, 1, 6, "Thermofit gloves 39,90 EUR")),
        _row(_chunk(PAGE, 3, 6, "part three")),
    ]

    result = tool.search_knowledge_base("pricing")

    assert result.startswith(f"[WEBSITE - https://x.com | {PAGE}] ")
    order = ["(part 1/6) Thermofit", "(part 3/6) part three", "(part 4/6) part four", "(part 5/6) part five"]
    positions = [result.index(p) for p in order]
    assert positions == sorted(positions)
    ids = tool._db.execute.call_args.args[1]["ids"]
    assert sorted(ids) == [f"{PAGE}::1", f"{PAGE}::3", f"{PAGE}::5"]
    assert tool.seen_chunk_ids == {f"{PAGE}::1", f"{PAGE}::3", f"{PAGE}::4", f"{PAGE}::5"}
    assert tool.collected_sources == [{"name": "https://x.com", "type": "website"}]


def test_page_head_never_displaces_a_matched_page(tool, monkeypatch):
    """Heads are added last and only where they fit: a page that matched the
    query must not be dropped by the cap to make room for another page's head."""
    a, b = "https://x.com/a", "https://x.com/b"
    tool.agent_knowledge.search.return_value = [_chunk(a, 3, 3, "a" * 60), _chunk(b, 3, 3, "b" * 60)]
    tool._db.execute.return_value.fetchall.return_value = [
        _row(_chunk(a, 1, 3, "HEAD-A " + "h" * 60)),
        _row(_chunk(b, 1, 3, "HEAD-B " + "h" * 60)),
    ]
    one_page = len(f"[WEBSITE - https://x.com | {a}] (part 3/3) " + "a" * 60) + 2
    monkeypatch.setattr("app.tools.knowledge_search_byagent.settings.KNOWLEDGE_SEARCH_MAX_CHARS", 2 * one_page + 10)

    result = tool.search_knowledge_base("q")

    assert "a" * 60 in result and "b" * 60 in result  # both matched pages kept
    assert "HEAD-A" not in result and "HEAD-B" not in result  # no room left for heads
    assert f"{a}::1" not in tool.seen_chunk_ids  # unshown head stays searchable


def test_page_head_added_when_it_fits(tool):
    tool.agent_knowledge.search.return_value = [_chunk(PAGE, 5, 5, "tail")]
    tool._db.execute.return_value.fetchall.return_value = [
        _row(_chunk(PAGE, 4, 5, "four")),
        _row(_chunk(PAGE, 1, 5, "Thermofit gloves 39,90 EUR")),
    ]

    result = tool.search_knowledge_base("Wie viel kosten die Handschuhe?")

    assert result.index("39,90 EUR") < result.index("four") < result.index("tail")


def test_page_head_is_not_fetched_twice(tool):
    """A hit on chunk 2 already asks for chunk 1 as its neighbour; a hit on
    chunk 1 needs no head lookup at all."""
    tool.agent_knowledge.search.return_value = [_chunk(PAGE, 2, 4, "two"), _chunk(PAGE, 1, 4, "one")]
    tool._db.execute.return_value.fetchall.return_value = [_row(_chunk(PAGE, 3, 4, "three"))]

    tool.search_knowledge_base("q")

    ids = tool._db.execute.call_args.args[1]["ids"]
    assert sorted(ids) == [f"{PAGE}::3"]


def test_unchunked_row_has_no_neighbour_lookup(tool):
    legacy = Document(id="https://x.com/about", name="https://x.com", content="About us",
                      meta_data={"url": "https://x.com/about", "chunk": 7})
    tool.agent_knowledge.search.return_value = [legacy]

    result = tool.search_knowledge_base("about")

    assert result == "[WEBSITE - https://x.com | https://x.com/about] About us"
    tool._db.execute.assert_not_called()


def test_repeat_search_returns_unseen_chunks_then_exhausts(tool):
    first, second = _chunk(PAGE, 1, 2, "one"), _chunk("https://x.com/faq", 1, 1, "faq")
    tool.agent_knowledge.search.return_value = [first]
    tool._db.execute.return_value.fetchall.return_value = []
    tool.search_knowledge_base("q")

    tool.agent_knowledge.search.return_value = [first, second]
    result = tool.search_knowledge_base("q again")
    assert "faq" in result and "one" not in result
    assert tool.agent_knowledge.search.call_args.kwargs["num_documents"] > 1

    tool.agent_knowledge.search.return_value = [first, second]
    assert tool.search_knowledge_base("q once more") == NO_FURTHER_RESULTS


def test_reset_turn_forgets_seen_and_sources(tool):
    tool.seen_chunk_ids = {"a"}
    tool.collected_sources = [{"name": "n", "type": "t"}]
    tool.reset_turn()
    assert tool.seen_chunk_ids == set() and tool.collected_sources == []


def test_result_cap_drops_whole_pages_never_cuts_a_chunk(tool, monkeypatch):
    monkeypatch.setattr("app.tools.knowledge_search_byagent.settings.KNOWLEDGE_SEARCH_MAX_CHARS", 120)
    pages = [_chunk(f"https://x.com/p{i}", 1, 1, "x" * 80) for i in range(4)]
    tool.agent_knowledge.search.return_value = pages
    tool._db.execute.return_value.fetchall.return_value = []

    result = tool.search_knowledge_base("q")

    assert result.count("[WEBSITE") == 1
    assert result.endswith("x" * 80)
    # Dropped pages were never shown: not seen, not cited, still searchable.
    assert tool.seen_chunk_ids == {"https://x.com/p0::1"}
    tool.agent_knowledge.search.return_value = pages
    assert "p1" in tool.search_knowledge_base("q again")


def test_not_yet_rechunked_whole_page_row_is_bounded(tool, monkeypatch):
    monkeypatch.setattr("app.tools.knowledge_search_byagent.settings.KNOWLEDGE_CHUNK_SIZE", 100)
    legacy = Document(id="https://x.com/huge", name="https://x.com",
                      content=("sentence. " * 10 + "\n") * 200, meta_data={"url": "https://x.com/huge"})
    tool.agent_knowledge.search.return_value = [legacy]

    result = tool.search_knowledge_base("q")

    assert len(result) < 600
    assert result.endswith("(page continues — not yet re-indexed in full)")
    assert "sentence. " in result