increase retries for query handler

wandb · Jan 6, 2025 · f6e4f7b · f6e4f7b
1 parent 61ddd45
commit f6e4f7b
Showing 1 changed file with 3 additions and 3 deletions.
diff --git a/src/wandbot/rag/query_handler.py b/src/wandbot/rag/query_handler.py
@@ -291,12 +291,12 @@ def _load_chain(self, model: ChatOpenAI) -> Runnable:
         base_query_enhancer = self.prompt | model.with_structured_output(
             EnhancedQuery
         )
-        
+
         # Add retry specifically for validation errors
         query_enhancer_chain = base_query_enhancer.with_retry(
-            retry_if_exception_type=(ValidationError,),
+            retry_if_exception_type=ValidationError,
             wait_exponential_jitter=True,  # Add jitter to prevent thundering herd
-            stop_after_attempt=3  # Try up to 3 times
+            stop_after_attempt=6,
         )
 
         input_chain = RunnableParallel(