wehub-resource-sync
353 lines
12 KiB
Python
353 lines
12 KiB
Python
import asyncio
|
|
import json
|
|
import random
|
|
from collections.abc import AsyncIterator, Callable
|
|
from typing import Any
|
|
|
|
import json_repair
|
|
|
|
from yuxi.models import select_model
|
|
from yuxi.utils import logger
|
|
|
|
DEFAULT_BENCHMARK_GENERATION_CONCURRENCY = 10
|
|
MAX_BENCHMARK_GENERATION_CONCURRENCY = 20
|
|
DEFAULT_GRAPH_EXPAND_TOP_K = 1
|
|
MAX_GRAPH_EXPAND_TOP_K = 3
|
|
GRAPH_SEED_DECAY = 0.9
|
|
GRAPH_PPR_DAMPING = 0.85
|
|
GRAPH_PPR_MAX_NODES = 10000
|
|
|
|
|
|
async def collect_kb_chunks(kb_instance: Any, kb_id: str) -> list[dict[str, Any]]:
|
|
del kb_instance
|
|
|
|
from yuxi.repositories.knowledge_chunk_repository import KnowledgeChunkRepository
|
|
|
|
return [
|
|
{
|
|
"id": chunk.chunk_id,
|
|
"content": chunk.content or "",
|
|
"file_id": chunk.file_id,
|
|
"chunk_index": chunk.chunk_index,
|
|
"graph_indexed": bool(chunk.graph_indexed),
|
|
"ent_ids": chunk.ent_ids or [],
|
|
"tags": chunk.tags or [],
|
|
"extraction_result": chunk.extraction_result,
|
|
}
|
|
for chunk in await KnowledgeChunkRepository().list_by_kb_id(kb_id)
|
|
]
|
|
|
|
|
|
def clamp_neighbors_count(neighbors_count: int) -> int:
|
|
return min(max(neighbors_count, 0), 10)
|
|
|
|
|
|
def normalize_generation_concurrency_count(value: Any) -> int:
|
|
if value in (None, ""):
|
|
return DEFAULT_BENCHMARK_GENERATION_CONCURRENCY
|
|
return min(max(1, int(value)), MAX_BENCHMARK_GENERATION_CONCURRENCY)
|
|
|
|
|
|
def normalize_graph_expand_top_k(value: Any) -> int:
|
|
if value in (None, ""):
|
|
return DEFAULT_GRAPH_EXPAND_TOP_K
|
|
return min(max(1, int(value)), MAX_GRAPH_EXPAND_TOP_K)
|
|
|
|
|
|
def _chunk_entity_ids(chunk: dict[str, Any]) -> list[str]:
|
|
return [str(entity_id) for entity_id in chunk.get("ent_ids") or [] if entity_id]
|
|
|
|
|
|
def _is_anchor_chunk(candidate: dict[str, Any], anchor_chunk: dict[str, Any]) -> bool:
|
|
metadata = candidate.get("metadata") or {}
|
|
candidate_id = metadata.get("chunk_id")
|
|
if candidate_id is not None and str(candidate_id) == str(anchor_chunk.get("id")):
|
|
return True
|
|
|
|
candidate_file_id = metadata.get("file_id")
|
|
candidate_chunk_index = metadata.get("chunk_index")
|
|
return candidate_file_id == anchor_chunk.get("file_id") and candidate_chunk_index == anchor_chunk.get("chunk_index")
|
|
|
|
|
|
async def select_neighbor_chunks_by_kb_query(
|
|
*, kb_instance: Any, kb_id: str, anchor_chunk: dict[str, Any], neighbors_count: int
|
|
) -> list[dict[str, Any]]:
|
|
if neighbors_count <= 0:
|
|
return []
|
|
|
|
anchor_content = anchor_chunk.get("content", "")
|
|
if not anchor_content:
|
|
return []
|
|
|
|
candidates = await kb_instance.aquery(
|
|
anchor_content,
|
|
kb_id,
|
|
search_mode="vector",
|
|
final_top_k=neighbors_count + 3,
|
|
use_reranker=False,
|
|
similarity_threshold=0.0,
|
|
)
|
|
|
|
chunks = []
|
|
for candidate in candidates:
|
|
if _is_anchor_chunk(candidate, anchor_chunk):
|
|
continue
|
|
|
|
metadata = candidate.get("metadata") or {}
|
|
chunk_id = metadata.get("chunk_id")
|
|
content = candidate.get("content", "")
|
|
if not chunk_id or not content:
|
|
continue
|
|
|
|
chunks.append(
|
|
{
|
|
"id": str(chunk_id),
|
|
"content": content,
|
|
"file_id": metadata.get("file_id"),
|
|
"chunk_index": metadata.get("chunk_index"),
|
|
}
|
|
)
|
|
if len(chunks) >= neighbors_count:
|
|
break
|
|
|
|
return chunks
|
|
|
|
|
|
async def select_graph_enhanced_chunks(
|
|
*,
|
|
kb_id: str,
|
|
anchor_chunk: dict[str, Any],
|
|
chunks_by_id: dict[str, dict[str, Any]],
|
|
context_count: int,
|
|
graph_expand_top_k: int,
|
|
) -> list[dict[str, Any]] | None:
|
|
if context_count <= 1:
|
|
return [anchor_chunk]
|
|
|
|
from yuxi.knowledge.graphs.milvus_graph_service import MilvusGraphService
|
|
|
|
anchor_entity_ids = _chunk_entity_ids(anchor_chunk)
|
|
if not anchor_entity_ids:
|
|
return None
|
|
|
|
graph_service = MilvusGraphService()
|
|
selected = [anchor_chunk]
|
|
selected_ids = {str(anchor_chunk.get("id"))}
|
|
seed_weights = {entity_id: 1.0 for entity_id in anchor_entity_ids}
|
|
round_index = 1
|
|
|
|
while len(selected) < context_count:
|
|
for entity_id in anchor_entity_ids:
|
|
seed_weights[entity_id] = 1.0
|
|
|
|
ranked_chunks = await graph_service.query_and_rank_chunks_by_ppr(
|
|
kb_id,
|
|
seed_weights,
|
|
max_nodes=GRAPH_PPR_MAX_NODES,
|
|
top_k=max(context_count * 5, 20),
|
|
damping=GRAPH_PPR_DAMPING,
|
|
)
|
|
if not ranked_chunks:
|
|
return None
|
|
|
|
new_chunks = []
|
|
for chunk_id, _ in ranked_chunks:
|
|
chunk_id = str(chunk_id)
|
|
if chunk_id in selected_ids:
|
|
continue
|
|
chunk = chunks_by_id.get(chunk_id)
|
|
if chunk is None:
|
|
continue
|
|
new_chunks.append(chunk)
|
|
if len(new_chunks) >= min(graph_expand_top_k, context_count - len(selected)):
|
|
break
|
|
|
|
if not new_chunks:
|
|
return None
|
|
|
|
new_weight = GRAPH_SEED_DECAY**round_index
|
|
for chunk in new_chunks:
|
|
selected.append(chunk)
|
|
selected_ids.add(str(chunk.get("id")))
|
|
for entity_id in _chunk_entity_ids(chunk):
|
|
seed_weights[entity_id] = max(seed_weights.get(entity_id, 0.0), new_weight)
|
|
round_index += 1
|
|
|
|
return selected
|
|
|
|
|
|
def build_benchmark_generation_prompt(ctx_items: list[tuple[str, str]]) -> str:
|
|
context_text = "\n\n".join([f"片段ID={cid}\n{content}" for cid, content in ctx_items])
|
|
return (
|
|
"你将基于以下上下文生成一个可由上下文准确回答的问题与标准答案。"
|
|
"仅返回一个JSON对象，不要包含其他文字。"
|
|
"键为 query、gold_answer、gold_chunk_ids。gold_chunk_ids 必须是上述上下文片段的ID子集。\n\n"
|
|
"上下文：\n" + context_text + "\n"
|
|
)
|
|
|
|
|
|
async def _generate_benchmark_item_once(
|
|
*,
|
|
kb_instance: Any,
|
|
kb_id: str,
|
|
all_chunks: list[dict[str, Any]],
|
|
llm: Any,
|
|
context_count: int,
|
|
generation_mode: str,
|
|
graph_expand_top_k: int,
|
|
chunks_by_id: dict[str, dict[str, Any]],
|
|
) -> dict[str, Any] | None:
|
|
if generation_mode == "graph_enhanced":
|
|
graph_anchor_chunks = [
|
|
chunk for chunk in all_chunks if chunk.get("graph_indexed") is True and _chunk_entity_ids(chunk)
|
|
]
|
|
if not graph_anchor_chunks:
|
|
raise ValueError("No graph indexed chunks with entities found in knowledge base")
|
|
anchor_chunk = graph_anchor_chunks[random.randrange(len(graph_anchor_chunks))]
|
|
ctx_chunks = await select_graph_enhanced_chunks(
|
|
kb_id=kb_id,
|
|
anchor_chunk=anchor_chunk,
|
|
chunks_by_id=chunks_by_id,
|
|
context_count=context_count,
|
|
graph_expand_top_k=graph_expand_top_k,
|
|
)
|
|
if ctx_chunks is None:
|
|
return None
|
|
else:
|
|
anchor_chunk = all_chunks[random.randrange(len(all_chunks))]
|
|
neighbor_chunks = await select_neighbor_chunks_by_kb_query(
|
|
kb_instance=kb_instance,
|
|
kb_id=kb_id,
|
|
anchor_chunk=anchor_chunk,
|
|
neighbors_count=context_count - 1,
|
|
)
|
|
ctx_chunks = [anchor_chunk] + neighbor_chunks
|
|
ctx_items = [(chunk["id"], chunk["content"]) for chunk in ctx_chunks]
|
|
allowed_ids = {cid for cid, _ in ctx_items}
|
|
|
|
try:
|
|
resp = await llm.call(build_benchmark_generation_prompt(ctx_items), False)
|
|
obj = json_repair.loads(resp.content if resp else "")
|
|
query = obj.get("query")
|
|
answer = obj.get("gold_answer")
|
|
gold_ids = obj.get("gold_chunk_ids")
|
|
if not query or not answer or not isinstance(gold_ids, list):
|
|
logger.warning(f"Generated JSON missing fields or invalid format: {obj}")
|
|
return None
|
|
|
|
gold_ids = [str(item) for item in gold_ids if str(item) in allowed_ids]
|
|
if not gold_ids:
|
|
logger.warning("Generated gold_chunk_ids not found in allowed context")
|
|
return None
|
|
|
|
return {"query": query, "gold_chunk_ids": gold_ids, "gold_answer": answer}
|
|
except Exception as e:
|
|
logger.warning(f"Benchmark generation failed for one item: {e}")
|
|
return None
|
|
|
|
|
|
async def iter_generated_benchmark_items(
|
|
*,
|
|
kb_instance: Any,
|
|
kb_id: str,
|
|
count: int,
|
|
neighbors_count: int,
|
|
llm_model_spec: str | None,
|
|
concurrency_count: int = DEFAULT_BENCHMARK_GENERATION_CONCURRENCY,
|
|
generation_mode: str = "vector",
|
|
graph_expand_top_k: int = DEFAULT_GRAPH_EXPAND_TOP_K,
|
|
progress_cb: Callable[[int, str], Any] | None = None,
|
|
cancel_cb: Callable[[], Any] | None = None,
|
|
) -> AsyncIterator[dict[str, Any]]:
|
|
if progress_cb:
|
|
await progress_cb(5, "加载chunks")
|
|
|
|
all_chunks = await collect_kb_chunks(kb_instance, kb_id)
|
|
if not all_chunks:
|
|
raise ValueError("No chunks found in knowledge base")
|
|
chunks_by_id = {str(chunk["id"]): chunk for chunk in all_chunks if chunk.get("id") is not None}
|
|
|
|
if generation_mode not in {"vector", "graph_enhanced"}:
|
|
raise ValueError("Unsupported benchmark generation mode")
|
|
graph_expand_top_k = normalize_graph_expand_top_k(graph_expand_top_k)
|
|
|
|
if progress_cb:
|
|
await progress_cb(15, "准备生成样本")
|
|
|
|
if not llm_model_spec:
|
|
raise ValueError("llm_model_spec 不能为空")
|
|
|
|
llm = select_model(model_spec=llm_model_spec)
|
|
context_count = max(clamp_neighbors_count(neighbors_count), 1)
|
|
max_attempts = max(count * 5, 50)
|
|
worker_count = normalize_generation_concurrency_count(concurrency_count)
|
|
actual_worker_count = min(worker_count, max(count, 1), max_attempts)
|
|
generated = 0
|
|
results: list[tuple[int, dict[str, Any]]] = []
|
|
state_lock = asyncio.Lock()
|
|
queue: asyncio.Queue[int] = asyncio.Queue()
|
|
|
|
for attempt_no in range(max_attempts):
|
|
queue.put_nowait(attempt_no)
|
|
|
|
async def worker() -> None:
|
|
nonlocal generated
|
|
while True:
|
|
if cancel_cb:
|
|
await cancel_cb()
|
|
async with state_lock:
|
|
if generated >= count:
|
|
return
|
|
try:
|
|
attempt_no = queue.get_nowait()
|
|
except asyncio.QueueEmpty:
|
|
return
|
|
try:
|
|
item = await _generate_benchmark_item_once(
|
|
kb_instance=kb_instance,
|
|
kb_id=kb_id,
|
|
all_chunks=all_chunks,
|
|
llm=llm,
|
|
context_count=context_count,
|
|
generation_mode=generation_mode,
|
|
graph_expand_top_k=graph_expand_top_k,
|
|
chunks_by_id=chunks_by_id,
|
|
)
|
|
if item is None:
|
|
continue
|
|
progress = None
|
|
message = None
|
|
async with state_lock:
|
|
if generated >= count:
|
|
continue
|
|
generated += 1
|
|
results.append((attempt_no, item))
|
|
if progress_cb:
|
|
progress = int(99 * generated / max(count, 1))
|
|
message = f"已生成 {generated}/{count}"
|
|
if progress_cb:
|
|
await progress_cb(progress, message)
|
|
finally:
|
|
queue.task_done()
|
|
|
|
workers = [asyncio.create_task(worker()) for _ in range(actual_worker_count)]
|
|
try:
|
|
await asyncio.gather(*workers)
|
|
except asyncio.CancelledError:
|
|
for task in workers:
|
|
task.cancel()
|
|
await asyncio.gather(*workers, return_exceptions=True)
|
|
raise
|
|
except Exception:
|
|
for task in workers:
|
|
task.cancel()
|
|
await asyncio.gather(*workers, return_exceptions=True)
|
|
raise
|
|
|
|
for _, item in sorted(results, key=lambda pair: pair[0]):
|
|
yield item
|
|
|
|
|
|
def dump_benchmark_item(item: dict[str, Any]) -> str:
|
|
return json.dumps(item, ensure_ascii=False, separators=(",", ":")) + "\n"
|