wehub-resource-sync
414 lines
13 KiB
Python
414 lines
13 KiB
Python
import asyncio
|
|
import os
|
|
from types import SimpleNamespace
|
|
|
|
import pytest
|
|
|
|
os.environ.setdefault("OPENAI_API_KEY", "test-key")
|
|
|
|
from yuxi.knowledge.eval import benchmark_generation
|
|
from yuxi.knowledge.eval.benchmark_generation import (
|
|
build_benchmark_generation_prompt,
|
|
clamp_neighbors_count,
|
|
collect_kb_chunks,
|
|
iter_generated_benchmark_items,
|
|
normalize_generation_concurrency_count,
|
|
select_graph_enhanced_chunks,
|
|
select_neighbor_chunks_by_kb_query,
|
|
)
|
|
|
|
|
|
class FakeKnowledgeBase:
|
|
pass
|
|
|
|
|
|
class FakeGenerationKnowledgeBase:
|
|
def __init__(self, query_results=None):
|
|
self.query_results = query_results or []
|
|
self.query_calls = []
|
|
|
|
async def aquery(self, query_text, kb_id, **kwargs):
|
|
self.query_calls.append({"query_text": query_text, "kb_id": kb_id, **kwargs})
|
|
return self.query_results
|
|
|
|
|
|
class FakeLlm:
|
|
def __init__(self, gold_chunk_id="anchor_chunk"):
|
|
self.gold_chunk_id = gold_chunk_id
|
|
self.prompts = []
|
|
|
|
async def call(self, prompt, stream):
|
|
self.prompts.append(prompt)
|
|
return SimpleNamespace(
|
|
content=('{"query":"问题","gold_answer":"答案","gold_chunk_ids":["' + self.gold_chunk_id + '"]}')
|
|
)
|
|
|
|
|
|
class NoQueryKnowledgeBase(FakeGenerationKnowledgeBase):
|
|
async def aquery(self, query_text, kb_id, **kwargs):
|
|
raise AssertionError("neighbors_count=1 时不应调用 aquery")
|
|
|
|
|
|
class TrackingLlm:
|
|
def __init__(self, content=None, delay=0):
|
|
self.content = content or '{"query":"问题","gold_answer":"答案","gold_chunk_ids":["anchor_chunk"]}'
|
|
self.delay = delay
|
|
self.active_calls = 0
|
|
self.max_active_calls = 0
|
|
self.calls = 0
|
|
|
|
async def call(self, prompt, stream):
|
|
self.calls += 1
|
|
self.active_calls += 1
|
|
self.max_active_calls = max(self.max_active_calls, self.active_calls)
|
|
try:
|
|
if self.delay:
|
|
await asyncio.sleep(self.delay)
|
|
return SimpleNamespace(content=self.content)
|
|
finally:
|
|
self.active_calls -= 1
|
|
|
|
|
|
class FakeGraphGenerationKnowledgeBase(FakeGenerationKnowledgeBase):
|
|
pass
|
|
|
|
|
|
def make_chunk(
|
|
chunk_id: str,
|
|
*,
|
|
kb_id: str = "db_1",
|
|
file_id: str = "file_a",
|
|
content: str = "anchor content",
|
|
chunk_index: int = 0,
|
|
graph_indexed: bool = False,
|
|
ent_ids: list[str] | None = None,
|
|
):
|
|
return SimpleNamespace(
|
|
chunk_id=chunk_id,
|
|
kb_id=kb_id,
|
|
file_id=file_id,
|
|
content=content,
|
|
chunk_index=chunk_index,
|
|
graph_indexed=graph_indexed,
|
|
ent_ids=ent_ids,
|
|
tags=None,
|
|
extraction_result=None,
|
|
)
|
|
|
|
|
|
@pytest.fixture(autouse=True)
|
|
def fake_chunk_repository(monkeypatch):
|
|
class FakeChunkRepository:
|
|
chunks = [make_chunk("anchor_chunk")]
|
|
|
|
async def list_by_kb_id(self, kb_id):
|
|
return [chunk for chunk in self.chunks if chunk.kb_id == kb_id]
|
|
|
|
monkeypatch.setattr(
|
|
"yuxi.repositories.knowledge_chunk_repository.KnowledgeChunkRepository",
|
|
FakeChunkRepository,
|
|
)
|
|
return FakeChunkRepository
|
|
|
|
|
|
def test_clamp_neighbors_count():
|
|
assert clamp_neighbors_count(-1) == 0
|
|
assert clamp_neighbors_count(3) == 3
|
|
assert clamp_neighbors_count(11) == 10
|
|
|
|
|
|
def test_normalize_generation_concurrency_count():
|
|
assert normalize_generation_concurrency_count(None) == 10
|
|
assert normalize_generation_concurrency_count("") == 10
|
|
assert normalize_generation_concurrency_count(0) == 1
|
|
assert normalize_generation_concurrency_count(-5) == 1
|
|
assert normalize_generation_concurrency_count(10000) == 20
|
|
|
|
|
|
def test_build_benchmark_generation_prompt_contains_required_schema():
|
|
prompt = build_benchmark_generation_prompt([("chunk_1", "片段内容")])
|
|
|
|
assert "片段ID=chunk_1" in prompt
|
|
assert "query、gold_answer、gold_chunk_ids" in prompt
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_collect_kb_chunks_filters_kb_id(fake_chunk_repository):
|
|
fake_chunk_repository.chunks = [
|
|
make_chunk("file_a_chunk", content="内容"),
|
|
make_chunk("file_b_chunk", kb_id="db_2", file_id="file_b", content="其他"),
|
|
]
|
|
|
|
chunks = await collect_kb_chunks(FakeKnowledgeBase(), "db_1")
|
|
|
|
assert chunks == [
|
|
{
|
|
"id": "file_a_chunk",
|
|
"content": "内容",
|
|
"file_id": "file_a",
|
|
"chunk_index": 0,
|
|
"graph_indexed": False,
|
|
"ent_ids": [],
|
|
"tags": [],
|
|
"extraction_result": None,
|
|
}
|
|
]
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_iter_generated_benchmark_items_with_one_chunk_does_not_query(monkeypatch):
|
|
fake_llm = FakeLlm()
|
|
monkeypatch.setattr(benchmark_generation, "select_model", lambda model_spec: fake_llm)
|
|
|
|
items = [
|
|
item
|
|
async for item in iter_generated_benchmark_items(
|
|
kb_instance=NoQueryKnowledgeBase(),
|
|
kb_id="db_1",
|
|
count=1,
|
|
neighbors_count=1,
|
|
llm_model_spec="test-provider:test-model",
|
|
)
|
|
]
|
|
|
|
assert items == [{"query": "问题", "gold_chunk_ids": ["anchor_chunk"], "gold_answer": "答案"}]
|
|
assert "片段ID=anchor_chunk" in fake_llm.prompts[0]
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_select_neighbor_chunks_by_kb_query_filters_anchor():
|
|
kb = FakeGenerationKnowledgeBase(
|
|
query_results=[
|
|
{
|
|
"content": "anchor content",
|
|
"metadata": {"chunk_id": "anchor_chunk", "file_id": "file_a", "chunk_index": 0},
|
|
},
|
|
{
|
|
"content": "neighbor content",
|
|
"metadata": {"chunk_id": "neighbor_chunk", "file_id": "file_a", "chunk_index": 1},
|
|
},
|
|
]
|
|
)
|
|
|
|
chunks = await select_neighbor_chunks_by_kb_query(
|
|
kb_instance=kb,
|
|
kb_id="db_1",
|
|
anchor_chunk={"id": "anchor_chunk", "content": "anchor content", "file_id": "file_a", "chunk_index": 0},
|
|
neighbors_count=1,
|
|
)
|
|
|
|
assert chunks == [{"id": "neighbor_chunk", "content": "neighbor content", "file_id": "file_a", "chunk_index": 1}]
|
|
assert kb.query_calls == [
|
|
{
|
|
"query_text": "anchor content",
|
|
"kb_id": "db_1",
|
|
"search_mode": "vector",
|
|
"final_top_k": 4,
|
|
"use_reranker": False,
|
|
"similarity_threshold": 0.0,
|
|
}
|
|
]
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_select_graph_enhanced_chunks_expands_by_ppr_with_anchor_bias(monkeypatch):
|
|
calls = []
|
|
|
|
async def fake_rank(self, kb_id, seed_weights, *, max_nodes, top_k, damping):
|
|
calls.append(dict(seed_weights))
|
|
if len(calls) == 1:
|
|
return [("anchor", 0.9), ("neighbor_1", 0.8)]
|
|
return [("anchor", 0.9), ("neighbor_1", 0.8), ("neighbor_2", 0.7)]
|
|
|
|
monkeypatch.setattr(
|
|
"yuxi.knowledge.graphs.milvus_graph_service.MilvusGraphService.query_and_rank_chunks_by_ppr",
|
|
fake_rank,
|
|
)
|
|
chunks_by_id = {
|
|
"anchor": {"id": "anchor", "content": "anchor", "ent_ids": ["anchor_entity"]},
|
|
"neighbor_1": {"id": "neighbor_1", "content": "neighbor 1", "ent_ids": ["entity_1"]},
|
|
"neighbor_2": {"id": "neighbor_2", "content": "neighbor 2", "ent_ids": ["entity_2"]},
|
|
}
|
|
|
|
chunks = await select_graph_enhanced_chunks(
|
|
kb_id="db_1",
|
|
anchor_chunk=chunks_by_id["anchor"],
|
|
chunks_by_id=chunks_by_id,
|
|
context_count=3,
|
|
graph_expand_top_k=1,
|
|
)
|
|
|
|
assert [chunk["id"] for chunk in chunks] == ["anchor", "neighbor_1", "neighbor_2"]
|
|
assert calls[0] == {"anchor_entity": 1.0}
|
|
assert calls[1]["anchor_entity"] == 1.0
|
|
assert calls[1]["entity_1"] == 0.9
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_iter_generated_benchmark_items_graph_mode_uses_graph_indexed_anchor(monkeypatch, fake_chunk_repository):
|
|
fake_chunk_repository.chunks = [
|
|
make_chunk(
|
|
"vector_anchor",
|
|
content="vector content",
|
|
chunk_index=0,
|
|
graph_indexed=False,
|
|
ent_ids=["vector_entity"],
|
|
),
|
|
make_chunk(
|
|
"graph_anchor",
|
|
content="graph anchor content",
|
|
chunk_index=1,
|
|
graph_indexed=True,
|
|
ent_ids=["anchor_entity"],
|
|
),
|
|
make_chunk(
|
|
"graph_neighbor",
|
|
content="graph neighbor content",
|
|
chunk_index=2,
|
|
graph_indexed=False,
|
|
ent_ids=["neighbor_entity"],
|
|
),
|
|
]
|
|
|
|
async def fake_rank(self, kb_id, seed_weights, *, max_nodes, top_k, damping):
|
|
assert seed_weights["anchor_entity"] == 1.0
|
|
return [("graph_anchor", 0.9), ("graph_neighbor", 0.8)]
|
|
|
|
fake_llm = FakeLlm(gold_chunk_id="graph_neighbor")
|
|
monkeypatch.setattr(benchmark_generation, "select_model", lambda model_spec: fake_llm)
|
|
monkeypatch.setattr(
|
|
"yuxi.knowledge.graphs.milvus_graph_service.MilvusGraphService.query_and_rank_chunks_by_ppr",
|
|
fake_rank,
|
|
)
|
|
kb = FakeGraphGenerationKnowledgeBase()
|
|
|
|
items = [
|
|
item
|
|
async for item in iter_generated_benchmark_items(
|
|
kb_instance=kb,
|
|
kb_id="db_1",
|
|
count=1,
|
|
neighbors_count=2,
|
|
llm_model_spec="test-provider:test-model",
|
|
generation_mode="graph_enhanced",
|
|
)
|
|
]
|
|
|
|
assert items == [{"query": "问题", "gold_chunk_ids": ["graph_neighbor"], "gold_answer": "答案"}]
|
|
assert kb.query_calls == []
|
|
assert "片段ID=graph_anchor" in fake_llm.prompts[0]
|
|
assert "片段ID=graph_neighbor" in fake_llm.prompts[0]
|
|
assert "片段ID=vector_anchor" not in fake_llm.prompts[0]
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_iter_generated_benchmark_items_uses_query_neighbor(monkeypatch):
|
|
fake_llm = FakeLlm(gold_chunk_id="neighbor_chunk")
|
|
monkeypatch.setattr(benchmark_generation, "select_model", lambda model_spec: fake_llm)
|
|
kb = FakeGenerationKnowledgeBase(
|
|
query_results=[
|
|
{
|
|
"content": "neighbor content",
|
|
"metadata": {"chunk_id": "neighbor_chunk", "file_id": "file_a", "chunk_index": 1},
|
|
}
|
|
]
|
|
)
|
|
|
|
items = [
|
|
item
|
|
async for item in iter_generated_benchmark_items(
|
|
kb_instance=kb,
|
|
kb_id="db_1",
|
|
count=1,
|
|
neighbors_count=2,
|
|
llm_model_spec="test-provider:test-model",
|
|
)
|
|
]
|
|
|
|
assert items == [{"query": "问题", "gold_chunk_ids": ["neighbor_chunk"], "gold_answer": "答案"}]
|
|
assert kb.query_calls[0]["query_text"] == "anchor content"
|
|
assert kb.query_calls[0]["search_mode"] == "vector"
|
|
assert "片段ID=neighbor_chunk" in fake_llm.prompts[0]
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_iter_generated_benchmark_items_falls_back_to_anchor_when_query_empty(monkeypatch):
|
|
fake_llm = FakeLlm()
|
|
monkeypatch.setattr(benchmark_generation, "select_model", lambda model_spec: fake_llm)
|
|
|
|
items = [
|
|
item
|
|
async for item in iter_generated_benchmark_items(
|
|
kb_instance=FakeGenerationKnowledgeBase(query_results=[]),
|
|
kb_id="db_1",
|
|
count=1,
|
|
neighbors_count=2,
|
|
llm_model_spec="test-provider:test-model",
|
|
)
|
|
]
|
|
|
|
assert items == [{"query": "问题", "gold_chunk_ids": ["anchor_chunk"], "gold_answer": "答案"}]
|
|
assert "片段ID=anchor_chunk" in fake_llm.prompts[0]
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_iter_generated_benchmark_items_respects_concurrency_count(monkeypatch):
|
|
fake_llm = TrackingLlm(delay=0.01)
|
|
monkeypatch.setattr(benchmark_generation, "select_model", lambda model_spec: fake_llm)
|
|
|
|
items = [
|
|
item
|
|
async for item in iter_generated_benchmark_items(
|
|
kb_instance=NoQueryKnowledgeBase(),
|
|
kb_id="db_1",
|
|
count=4,
|
|
neighbors_count=1,
|
|
concurrency_count=2,
|
|
llm_model_spec="test-provider:test-model",
|
|
)
|
|
]
|
|
|
|
assert len(items) == 4
|
|
assert fake_llm.max_active_calls == 2
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_iter_generated_benchmark_items_returns_at_most_count(monkeypatch):
|
|
fake_llm = TrackingLlm(delay=0.01)
|
|
monkeypatch.setattr(benchmark_generation, "select_model", lambda model_spec: fake_llm)
|
|
|
|
items = [
|
|
item
|
|
async for item in iter_generated_benchmark_items(
|
|
kb_instance=NoQueryKnowledgeBase(),
|
|
kb_id="db_1",
|
|
count=3,
|
|
neighbors_count=1,
|
|
concurrency_count=10,
|
|
llm_model_spec="test-provider:test-model",
|
|
)
|
|
]
|
|
|
|
assert len(items) == 3
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_iter_generated_benchmark_items_stops_at_max_attempts(monkeypatch):
|
|
fake_llm = TrackingLlm(content='{"query":"","gold_answer":"答案","gold_chunk_ids":["anchor_chunk"]}')
|
|
monkeypatch.setattr(benchmark_generation, "select_model", lambda model_spec: fake_llm)
|
|
|
|
items = [
|
|
item
|
|
async for item in iter_generated_benchmark_items(
|
|
kb_instance=NoQueryKnowledgeBase(),
|
|
kb_id="db_1",
|
|
count=2,
|
|
neighbors_count=1,
|
|
concurrency_count=10,
|
|
llm_model_spec="test-provider:test-model",
|
|
)
|
|
]
|
|
|
|
assert items == []
|
|
assert fake_llm.calls == 50
|