188 lines
7.6 KiB
Python
188 lines
7.6 KiB
Python
"""
|
|
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
|