mirror of
https://github.com/tiennm99/DocsGPT.git
synced 2026-10-19 06:14:42 +00:00
fix(graphrag): give each extraction thread its own LLM
Provider-reported usage is kept on the LLM instance (_last_usage) and claimed by whichever call finishes next. With GRAPHRAG_EXTRACTION_WORKERS > 1 (the default is 8) every extraction thread shared one instance, so a call could claim another call's provider counts while its own fell back to the estimate: token_usage rows, and the cost they bill, could be attributed to the wrong call and summed wrong. Each pool thread now builds its own extraction LLM on first use. The calling thread's instance is still built up front, so a misconfigured model fails the run before any chunk is touched.
This commit is contained in:
1 parent
022cf69b49
commit
a20402a6dc
2 files changed
+75
-4
No files matched your search
@@ -227,6 +227,7 @@ def extract_graph_for_source(
|
||||
source's graph holds after the run — not how many upserts ran, which
|
||||
counts the same entity once per chunk it appears in.
|
||||
"""
|
||||
import threading
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
|
||||
from docsgpt.graphrag.store import GraphStore
|
||||
@@ -246,9 +247,24 @@ def extract_graph_for_source(
|
||||
|
||||
embedding = get_embeddings()
|
||||
|
||||
llm = _build_extraction_llm(
|
||||
_resolve_extraction_model(config), user, request_id
|
||||
)
|
||||
model_id = _resolve_extraction_model(config)
|
||||
# Built here first so a misconfigured model fails the run before any
|
||||
# chunk is touched; this instance serves the calling thread.
|
||||
thread_llm = threading.local()
|
||||
thread_llm.llm = _build_extraction_llm(model_id, user, request_id)
|
||||
|
||||
def _llm():
|
||||
"""This thread's extraction LLM.
|
||||
|
||||
Provider-reported usage is kept on the LLM instance (``_last_usage``)
|
||||
and claimed by whichever call finishes next, so two calls in flight on
|
||||
one instance can bill each other's tokens. Each pool thread therefore
|
||||
builds its own.
|
||||
"""
|
||||
llm = getattr(thread_llm, "llm", None)
|
||||
if llm is None:
|
||||
llm = thread_llm.llm = _build_extraction_llm(model_id, user, request_id)
|
||||
return llm
|
||||
|
||||
node_upserts = 0
|
||||
edges = 0
|
||||
@@ -284,7 +300,7 @@ def extract_graph_for_source(
|
||||
if not text:
|
||||
return chunk_id, "empty", None
|
||||
|
||||
extracted = _extract_chunk(llm, text, chunk_id)
|
||||
extracted = _extract_chunk(_llm(), text, chunk_id)
|
||||
if extracted is None:
|
||||
return chunk_id, "failed", None
|
||||
try:
|
||||
|
||||
Reference in new issue
Block a user