Feat: add citation option to agent and enlarge the timeouts. (#9484)

### What problem does this PR solve?

#9422

### Type of change

- [x] New Feature (non-breaking change which adds functionality)
This commit is contained in:
Kevin Hu
2025-08-15 10:05:01 +08:00
committed by GitHub
parent 562349eb02
commit 2114e966d8
6 changed files with 22 additions and 21 deletions

View File

@ -57,20 +57,22 @@ async def run_graphrag(
):
chunks.append(d["content_with_weight"])
subgraph = await generate_subgraph(
LightKGExt
if "method" not in row["kb_parser_config"].get("graphrag", {}) or row["kb_parser_config"]["graphrag"]["method"] != "general"
else GeneralKGExt,
tenant_id,
kb_id,
doc_id,
chunks,
language,
row["kb_parser_config"]["graphrag"].get("entity_types", []),
chat_model,
embedding_model,
callback,
)
with trio.fail_after(len(chunks)*60):
subgraph = await generate_subgraph(
LightKGExt
if "method" not in row["kb_parser_config"].get("graphrag", {}) or row["kb_parser_config"]["graphrag"]["method"] != "general"
else GeneralKGExt,
tenant_id,
kb_id,
doc_id,
chunks,
language,
row["kb_parser_config"]["graphrag"].get("entity_types", []),
chat_model,
embedding_model,
callback,
)
if not subgraph:
return
@ -125,7 +127,6 @@ async def run_graphrag(
return
@timeout(60*60, 1)
async def generate_subgraph(
extractor: Extractor,
tenant_id: str,