apunkt/trustgraph
1
0
Fork 0
mirror of https://github.com/trustgraph-ai/trustgraph.git synced 2026-07-09 13:22:10 +02:00
Code Issues Projects Releases Packages Wiki Activity Actions 1

feat: make fetch timeout for chunks configurable (#1031)

Browse source 
* tests: processor accepts default and overrides of fetch_chunk_timeout

* feat: configure fetch timeout for chunks from librarian
This commit is contained in:
Keshav 2026-07-07 10:55:33 -04:00 • committed by GitHub
parent 2bdc930b2a
commit f76f2abe08
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 82 additions and 2 deletions
Download patch file Download diff file Expand all files Collapse all files

13
trustgraph-flow/trustgraph/retrieval/document_rag/rag.py
View file

@ -36,6 +36,7 @@ class Processor(FlowProcessor):
fetch_limit = params.get("fetch_limit", 0)
rerank_diversity_mode = params.get("rerank_diversity_mode", "none")
rerank_diversity_lambda = params.get("rerank_diversity_lambda", 0.7)
fetch_chunk_timeout = params.get("fetch_chunk_timeout", 120)
retrieval_mode = params.get("retrieval_mode", "vector")
vector_weight = params.get("vector_weight", 1.0)
keyword_weight = params.get("keyword_weight", 1.0)
@ -47,6 +48,7 @@ class Processor(FlowProcessor):
"fetch_limit": fetch_limit,
"rerank_diversity_mode": rerank_diversity_mode,
"rerank_diversity_lambda": rerank_diversity_lambda,
"fetch_chunk_timeout": fetch_chunk_timeout,
"retrieval_mode": retrieval_mode,
"vector_weight": vector_weight,
"keyword_weight": keyword_weight,
@ -57,6 +59,7 @@ class Processor(FlowProcessor):
self.fetch_limit = fetch_limit
self.rerank_diversity_mode = rerank_diversity_mode
self.rerank_diversity_lambda = rerank_diversity_lambda
self.fetch_chunk_timeout = fetch_chunk_timeout
self.retrieval_mode = retrieval_mode
self.vector_weight = vector_weight
self.keyword_weight = keyword_weight
@ -139,7 +142,7 @@ class Processor(FlowProcessor):
logger.info(f"Handling input {id}...")
async def fetch_chunk(chunk_id, timeout=120):
async def fetch_chunk(chunk_id, timeout=self.fetch_chunk_timeout):
return await flow.librarian.fetch_document_text(
document_id=chunk_id, timeout=timeout,
)
@ -329,6 +332,14 @@ class Processor(FlowProcessor):
help='MMR relevance/diversity tradeoff, higher values prefer relevance'
)
parser.add_argument(
'--fetch-chunk-timeout',
type=int,
default=120,
help='Timeout in seconds for fetching a document chunk from the '
'librarian (default: 120)'
)
parser.add_argument(
'--retrieval-mode',
choices=['vector', 'keyword', 'hybrid'],