feat(kb): LLM 图谱抽取器重构为固定 Prompt + Schema,支持并发构建与模型参数

- LLM 抽取器收敛为固定 Prompt + 可选 Schema 约束,禁止自定义完整 Prompt
- 图谱构建改用 asyncio worker 队列并发,并发数由 extractor_options.concurrency_count 控制
- 已锁定的图谱配置允许修改同类参数,类型保持锁定
- OpenAIBase 透传 model_params 到流式和非流式调用
This commit is contained in:
Wenjie Zhang 2026-05-17 19:54:45 +08:00
parent 409df5b073
commit d292e53c23
3 changed files with 130 additions and 47 deletions

View File

@ -23,6 +23,10 @@ JSON 格式:
}
"""
SCHEMA_INSTRUCTION = """抽取 Schema 约束:
{schema}
"""
class LLMGraphExtractor(GraphExtractor):
extractor_type = "llm"
@ -30,15 +34,33 @@ class LLMGraphExtractor(GraphExtractor):
def validate_options(self) -> None:
if not self.options.get("model_spec"):
raise ValueError("LLM 抽取器需要 model_spec")
if self.options.get("prompt"):
raise ValueError("LLM 图谱抽取器不支持自定义完整 Prompt请使用 schema 配置抽取约束")
concurrency_count = self.options.get("concurrency_count", 1)
try:
concurrency_count = int(concurrency_count)
except (TypeError, ValueError) as exc:
raise ValueError("LLM 抽取器 concurrency_count 必须是整数") from exc
if concurrency_count < 1 or concurrency_count > 20:
raise ValueError("LLM 抽取器 concurrency_count 必须在 1 到 20 之间")
if self.options.get("model_params") is not None and not isinstance(self.options["model_params"], dict):
raise ValueError("LLM 抽取器 model_params 必须是对象")
async def extract(self, text: str, *, chunk_metadata: dict[str, Any] | None = None) -> dict[str, Any]:
self.validate_options()
model = select_model(model_spec=self.options["model_spec"], timeout=60.0)
model = select_model(
model_spec=self.options["model_spec"],
timeout=60.0,
model_params=self.options.get("model_params") or {},
)
prompt = self._build_prompt(text)
response = await model.call(prompt, stream=False)
parsed = json_repair.loads(response.content if response else "")
return parsed
def _build_prompt(self, text: str) -> str:
extraction_prompt = self.options.get("prompt") or DEFAULT_TRIPLE_EXTRACTION_PROMPT
extraction_prompt = DEFAULT_TRIPLE_EXTRACTION_PROMPT
schema = str(self.options.get("schema") or "").strip()
if schema:
extraction_prompt = f"{extraction_prompt}\n{SCHEMA_INSTRUCTION.format(schema=schema)}"
return f"{extraction_prompt}\n\n文本:\n{text}"

View File

@ -120,17 +120,26 @@ class MilvusGraphService:
kb = await self._get_milvus_kb(db_id)
additional_params = dict(kb.additional_params or {})
existing_config = additional_params.get(GRAPH_CONFIG_KEY) or {}
normalized_extractor_type = (extractor_type or "").lower()
if existing_config.get("locked"):
raise ValueError("图谱抽取配置已锁定")
existing_extractor_type = (existing_config.get("extractor_type") or "").lower()
if normalized_extractor_type != existing_extractor_type:
raise ValueError("图谱抽取器类型已锁定只能修改模型、Schema 等抽取参数")
GraphExtractorFactory.create(extractor_type, extractor_options)
extractor_options = extractor_options or {}
if normalized_extractor_type == "llm" and extractor_options.get("prompt"):
raise ValueError("LLM 图谱抽取器不支持自定义完整 Prompt请使用 schema 配置抽取约束")
GraphExtractorFactory.create(normalized_extractor_type, extractor_options)
config = {
"locked": True,
"extractor_type": extractor_type,
"extractor_type": normalized_extractor_type,
"extractor_options": extractor_options or {},
"created_at": utc_isoformat(),
"created_by": created_by,
"created_at": existing_config.get("created_at") or utc_isoformat(),
"created_by": existing_config.get("created_by") or created_by,
}
if existing_config.get("locked"):
config["updated_at"] = utc_isoformat()
config["updated_by"] = created_by
additional_params[GRAPH_CONFIG_KEY] = config
await self.kb_repo.update(db_id, {"additional_params": additional_params})
return config
@ -138,11 +147,14 @@ class MilvusGraphService:
async def build_pending_chunks(self, db_id: str, *, batch_size: int, context=None) -> dict[str, Any]:
kb = await self._get_milvus_kb(db_id)
config = self._get_locked_config(kb.additional_params or {})
extractor = GraphExtractorFactory.create(config["extractor_type"], config.get("extractor_options") or {})
extractor_options = self._runtime_extractor_options(config)
extractor = GraphExtractorFactory.create(config["extractor_type"], extractor_options)
worker_count = self._get_worker_count(config)
total_pending = await self.chunk_repo.count_graph_pending_by_db_id(db_id)
processed = 0
failed = 0
failed_chunk_ids: set[str] = set()
write_lock = asyncio.Lock()
while True:
if context is not None:
@ -152,48 +164,86 @@ class MilvusGraphService:
if not unprocessed:
break
queue: asyncio.Queue[Any] = asyncio.Queue()
for chunk in unprocessed:
if context is not None:
await context.raise_if_cancelled()
try:
extraction_result = await self._get_chunk_extraction_result(db_id, chunk, extractor)
entities, triples = await asyncio.to_thread(
self.write_chunk_graph,
db_id,
chunk,
extraction_result,
)
await self.graph_repo.upsert_chunk_graph(
db_id=db_id,
file_id=chunk.file_id,
chunk_id=chunk.chunk_id,
entities=entities,
triples=triples,
)
await self.graph_vector_store.insert_missing_graph_records(
db_id=db_id,
embedding_model_spec=kb.embedding_model_spec,
entities=entities,
triples=triples,
)
await self.chunk_repo.mark_graph_indexed(
chunk.chunk_id,
ent_ids=[entity["entity_id"] for entity in entities],
)
processed += 1
except Exception as exc:
logger.error(f"Chunk 图谱构建失败 chunk_id={chunk.chunk_id}: {exc}")
failed_chunk_ids.add(chunk.chunk_id)
failed += 1
queue.put_nowait(chunk)
if context is not None:
completed = processed + failed
progress = 5.0 + min(90.0, completed / max(total_pending, 1) * 90.0)
await context.set_progress(progress, f"图谱构建 {completed}/{total_pending},失败 {failed}")
async def worker() -> None:
nonlocal processed, failed
while True:
if context is not None:
await context.raise_if_cancelled()
try:
chunk = queue.get_nowait()
except asyncio.QueueEmpty:
return
try:
extraction_result = await self._get_chunk_extraction_result(db_id, chunk, extractor)
async with write_lock:
entities, triples = await asyncio.to_thread(
self.write_chunk_graph,
db_id,
chunk,
extraction_result,
)
await self.graph_repo.upsert_chunk_graph(
db_id=db_id,
file_id=chunk.file_id,
chunk_id=chunk.chunk_id,
entities=entities,
triples=triples,
)
await self.graph_vector_store.insert_missing_graph_records(
db_id=db_id,
embedding_model_spec=kb.embedding_model_spec,
entities=entities,
triples=triples,
)
await self.chunk_repo.mark_graph_indexed(
chunk.chunk_id,
ent_ids=[entity["entity_id"] for entity in entities],
)
processed += 1
except Exception as exc:
logger.error(f"Chunk 图谱构建失败 chunk_id={chunk.chunk_id}: {exc}")
failed_chunk_ids.add(chunk.chunk_id)
failed += 1
finally:
queue.task_done()
if context is not None:
completed = processed + failed
progress = 5.0 + min(90.0, completed / max(total_pending, 1) * 90.0)
await context.set_progress(progress, f"图谱构建 {completed}/{total_pending},失败 {failed}")
workers = [asyncio.create_task(worker()) for _ in range(min(worker_count, len(unprocessed)))]
try:
await asyncio.gather(*workers)
except Exception:
for task in workers:
task.cancel()
await asyncio.gather(*workers, return_exceptions=True)
raise
remaining = await self.chunk_repo.count_graph_pending_by_db_id(db_id)
return {"db_id": db_id, "success": processed, "failed": failed, "remaining": remaining}
@staticmethod
def _get_worker_count(config: dict[str, Any]) -> int:
if (config.get("extractor_type") or "").lower() != "llm":
return 1
try:
worker_count = int((config.get("extractor_options") or {}).get("concurrency_count") or 1)
except (TypeError, ValueError):
return 1
return max(1, min(worker_count, 20))
@staticmethod
def _runtime_extractor_options(config: dict[str, Any]) -> dict[str, Any]:
options = dict(config.get("extractor_options") or {})
options.pop("prompt", None)
return options
async def _get_chunk_extraction_result(self, db_id: str, chunk, extractor: GraphExtractor) -> dict[str, Any]:
extractor_type = extractor.extractor_type
if chunk.extraction_result:
@ -434,7 +484,12 @@ class MilvusGraphService:
neo4j_write(self.driver, query)
async def query_nodes(
self, db_id: str | None = None, *, keyword: str = "", max_depth: int = 1, max_nodes: int = 50,
self,
db_id: str | None = None,
*,
keyword: str = "",
max_depth: int = 1,
max_nodes: int = 50,
exclude_chunk: bool = False,
) -> dict[str, Any]:
effective_db_id = db_id or self.db_id
@ -447,7 +502,8 @@ class MilvusGraphService:
with self.driver.session() as session:
result = session.run(
self._build_query(label, keyword, limit, max_depth, exclude_chunk),
keyword=keyword, limit=limit,
keyword=keyword,
limit=limit,
)
return self._process_query_result(result, limit, effective_db_id, exclude_chunk)
except Exception as e:
@ -565,9 +621,11 @@ class MilvusGraphService:
return {
"locked": bool(config.get("locked")),
"extractor_type": config.get("extractor_type"),
"extractor_options": config.get("extractor_options") or {},
"extractor_options": self._runtime_extractor_options(config),
"created_at": config.get("created_at"),
"created_by": config.get("created_by"),
"updated_at": config.get("updated_at"),
"updated_by": config.get("updated_by"),
}
@staticmethod

View File

@ -7,6 +7,7 @@ from yuxi.utils import logger
class OpenAIBase:
def __init__(self, api_key, base_url, model_name, **kwargs):
self.model_params = kwargs.pop("model_params", {}) or {}
self.api_key = api_key
self.base_url = base_url
self.client = AsyncOpenAI(api_key=api_key, base_url=base_url, **kwargs)
@ -46,6 +47,7 @@ class OpenAIBase:
model=self.model_name,
messages=messages,
stream=True,
**self.model_params,
)
async for chunk in response:
if len(chunk.choices) > 0:
@ -56,6 +58,7 @@ class OpenAIBase:
model=self.model_name,
messages=messages,
stream=False,
**self.model_params,
)
return response.choices[0].message