feat: add RAG_RERANKING_BATCH_SIZE configuration option

Add configurable reranker batch size (env var RAG_RERANKING_BATCH_SIZE,
default 32) following the same pattern as RAG_EMBEDDING_BATCH_SIZE.

- config.py: PersistentConfig for RAG_RERANKING_BATCH_SIZE
- main.py: import, state init, pass to get_reranking_function
- colbert.py: accept batch_size param in predict() (was hardcoded 32)
- utils.py: get_reranking_function passes batch_size at call time
- retrieval.py: expose in config GET/POST endpoints and ConfigForm
- Documents.svelte: add Reranking Batch Size input in admin settings

Closes #23730
This commit is contained in:
Timothy Jaeryang Baek
2026-04-17 08:35:45 +09:00
parent 70a6a24f14
commit 4d2f189810
6 changed files with 40 additions and 5 deletions
+9
View File
@@ -487,6 +487,7 @@ async def get_rag_config(request: Request, user=Depends(get_admin_user)):
# Reranking settings
'RAG_RERANKING_MODEL': request.app.state.config.RAG_RERANKING_MODEL,
'RAG_RERANKING_ENGINE': request.app.state.config.RAG_RERANKING_ENGINE,
'RAG_RERANKING_BATCH_SIZE': request.app.state.config.RAG_RERANKING_BATCH_SIZE,
'RAG_EXTERNAL_RERANKER_URL': request.app.state.config.RAG_EXTERNAL_RERANKER_URL,
'RAG_EXTERNAL_RERANKER_API_KEY': request.app.state.config.RAG_EXTERNAL_RERANKER_API_KEY,
'RAG_EXTERNAL_RERANKER_TIMEOUT': request.app.state.config.RAG_EXTERNAL_RERANKER_TIMEOUT,
@@ -694,6 +695,7 @@ class ConfigForm(BaseModel):
# Reranking settings
RAG_RERANKING_MODEL: Optional[str] = None
RAG_RERANKING_ENGINE: Optional[str] = None
RAG_RERANKING_BATCH_SIZE: Optional[int] = None
RAG_EXTERNAL_RERANKER_URL: Optional[str] = None
RAG_EXTERNAL_RERANKER_API_KEY: Optional[str] = None
RAG_EXTERNAL_RERANKER_TIMEOUT: Optional[str] = None
@@ -940,6 +942,12 @@ async def update_rag_config(request: Request, form_data: ConfigForm, user=Depend
else request.app.state.config.RAG_EXTERNAL_RERANKER_TIMEOUT
)
request.app.state.config.RAG_RERANKING_BATCH_SIZE = (
form_data.RAG_RERANKING_BATCH_SIZE
if form_data.RAG_RERANKING_BATCH_SIZE is not None
else request.app.state.config.RAG_RERANKING_BATCH_SIZE
)
log.info(
f'Updating reranking model: {request.app.state.config.RAG_RERANKING_MODEL} to {form_data.RAG_RERANKING_MODEL}'
)
@@ -967,6 +975,7 @@ async def update_rag_config(request: Request, form_data: ConfigForm, user=Depend
request.app.state.config.RAG_RERANKING_ENGINE,
request.app.state.config.RAG_RERANKING_MODEL,
request.app.state.rf,
reranking_batch_size=request.app.state.config.RAG_RERANKING_BATCH_SIZE,
)
except Exception as e:
log.error(f'Error loading reranking model: {e}')