162 lines
5.6 KiB
Python
162 lines
5.6 KiB
Python
from __future__ import annotations
|
|
|
|
from datetime import date
|
|
|
|
import pytest
|
|
|
|
from deerflow.integrations.weknora.local_index.advanced import (
|
|
advanced_search,
|
|
build_statistics,
|
|
enrich_chunk_references,
|
|
expand_related_terms,
|
|
extract_subject_terms,
|
|
parse_date_range,
|
|
)
|
|
from deerflow.persistence.llmwiki_index import MemoryLlmWikiIndexStore
|
|
|
|
|
|
def _page(*, title: str, slug: str, content: str, published_at: str, source_refs=None) -> dict:
|
|
return {
|
|
"id": slug,
|
|
"knowledge_base_mapping_id": "kb",
|
|
"slug": slug,
|
|
"title": title,
|
|
"summary": "",
|
|
"content_md": f"---\npublished_at: {published_at}\n---\n{content}",
|
|
"aliases": [],
|
|
"page_type": "article",
|
|
"status": "published",
|
|
"is_remote_deleted": False,
|
|
"source_refs": source_refs or [],
|
|
}
|
|
|
|
|
|
def test_parses_chinese_date_ranges_and_subject() -> None:
|
|
start, end = parse_date_range(
|
|
"帮我总结9月1号到9月15号的报告，有多少是关于赖清德的？",
|
|
today=date(2026, 9, 20),
|
|
)
|
|
assert start == date(2026, 9, 1)
|
|
assert end == date(2026, 9, 15)
|
|
assert extract_subject_terms("有多少是关于赖清德的？") == ["赖清德"]
|
|
|
|
|
|
def test_expands_first_degree_wiki_relations() -> None:
|
|
page = _page(
|
|
title="赖清德",
|
|
slug="people/lai",
|
|
published_at="2026-09-03",
|
|
content="赖清德与[[民进党]]、[[台湾行政机构|行政机构]]有关。",
|
|
)
|
|
expanded = expand_related_terms(["赖清德"], [page], limit=5)
|
|
assert [item["term"] for item in expanded] == ["民进党", "行政机构"]
|
|
assert all(item["from_page"] == "赖清德" for item in expanded)
|
|
|
|
|
|
def test_statistics_prefers_source_documents_and_applies_exact_dates() -> None:
|
|
pages = [
|
|
_page(
|
|
title="赖清德相关新闻汇总",
|
|
slug="news/lai",
|
|
published_at="2026-09-05",
|
|
content="赖清德出席活动。",
|
|
source_refs=[
|
|
{"document_id": "n1", "title": "新闻一", "published_at": "2026-09-02", "url": "https://example.test/1"},
|
|
{"document_id": "n2", "title": "新闻二", "published_at": "2026-09-18", "url": "https://example.test/2"},
|
|
{"document_id": "n3", "title": "新闻三"},
|
|
],
|
|
)
|
|
]
|
|
payload = build_statistics(
|
|
pages,
|
|
[{"id": "kb", "name": "台湾新闻"}],
|
|
subject_terms=["赖清德"],
|
|
date_start=date(2026, 9, 1),
|
|
date_end=date(2026, 9, 15),
|
|
requested=True,
|
|
)
|
|
assert payload["unit"] == "source_document"
|
|
assert payload["matched_count"] == 1
|
|
assert payload["items"][0]["title"] == "新闻一"
|
|
assert payload["undated_excluded_count"] == 1
|
|
|
|
|
|
class _ChunkClient:
|
|
async def get_chunk(self, chunk_id: str):
|
|
return {
|
|
"id": chunk_id,
|
|
"knowledge_base_id": "remote-kb",
|
|
"knowledge_id": "news-document-1",
|
|
"knowledge_title": "赖清德活动新闻",
|
|
"content": "发布日期：2026年9月6日\n赖清德出席活动。",
|
|
"metadata": {"url": "https://example.test/news"},
|
|
}
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_resolves_only_wiki_cited_chunks_to_news_sources() -> None:
|
|
page = _page(
|
|
title="赖清德",
|
|
slug="people/lai",
|
|
published_at="2026-09-03",
|
|
content="赖清德活动",
|
|
source_refs=["chunk-1"],
|
|
)
|
|
pages, warnings = await enrich_chunk_references(
|
|
[page],
|
|
[{"id": "kb", "weknora_id": "remote-kb"}],
|
|
client=_ChunkClient(),
|
|
)
|
|
assert warnings == []
|
|
resolved = next(item for item in pages[0]["source_refs"] if isinstance(item, dict))
|
|
assert resolved["document_id"] == "news-document-1"
|
|
assert resolved["title"] == "赖清德活动新闻"
|
|
|
|
|
|
class _Retrieval:
|
|
async def search(self, query, mappings, *, top_k_pages, include_drafts):
|
|
return {
|
|
"mode": "hybrid",
|
|
"results": [
|
|
{
|
|
"id": "p1",
|
|
"wiki_page_id": "p1",
|
|
"wiki_slug": "news/lai",
|
|
"knowledge_base_id": "kb",
|
|
"knowledge_base_name": "台湾新闻",
|
|
"title": "赖清德相关新闻",
|
|
"content": "赖清德与民进党",
|
|
"snippet": "赖清德与民进党",
|
|
"score": 0.8,
|
|
"url": "/preview",
|
|
"matched_sections": [],
|
|
}
|
|
],
|
|
"warnings": [],
|
|
"diagnostics": {"query_seen": query},
|
|
}
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_advanced_search_returns_plan_counts_and_references() -> None:
|
|
store = MemoryLlmWikiIndexStore()
|
|
page = _page(
|
|
title="赖清德相关新闻",
|
|
slug="news/lai",
|
|
published_at="2026-09-03",
|
|
content="赖清德与[[民进党]]",
|
|
)
|
|
page.update({"slug_hash": "hash", "content_hash": "content"})
|
|
await store.replace_page_vectors(page=page, vectors=[], fingerprint="fp", sync_id="sync")
|
|
payload = await advanced_search(
|
|
query="9月1号到9月15号有多少关于赖清德的新闻？",
|
|
mappings=[{"id": "kb", "name": "台湾新闻", "description": "每日新闻", "retrieval_settings": None}],
|
|
retrieval_service=_Retrieval(),
|
|
index_store=store,
|
|
top_k=20,
|
|
)
|
|
assert payload["query_plan"]["subject_terms"] == ["赖清德"]
|
|
assert payload["query_plan"]["expanded_terms"][0]["term"] == "民进党"
|
|
assert payload["statistics"]["matched_count"] == 1
|
|
assert payload["references"][0]["reference_id"] == "wiki:kb:news/lai"
|