Update kig_core/llm_interface.py

kig_core/llm_interface.py

@@ -1,10 +1,12 @@
 import os
+import re
 import time
+import logging
+
 from langchain_google_genai import ChatGoogleGenerativeAI
 from langchain_openai import ChatOpenAI
 from langchain_core.language_models.chat_models import BaseChatModel
 from .config import settings
-import logging
 
 logger = logging.getLogger(__name__)
 
@@ -55,15 +57,22 @@ def get_llm(model_name: str) -> BaseChatModel:
         logger.error(f"Unsupported model provider for model name: {model_name}")
         raise ValueError(f"Model '{model_name}' is not supported or configuration is missing.")
 
-def invoke_llm(var,parameters):
+def invoke_llm(var, parameters):
     try:
         return var.invoke(parameters)
     except Exception as e:
-        print(f"Error during .invoke : {e} \nwaiting 60 seconds")
-        time.sleep(60)
+        # Try to extract retry_delay seconds from the error message string
+        match = re.search(r'retry_delay\s*{\s*seconds:\s*(\d+)', str(e))
+        if match:
+            retry_seconds = int(match.group(1)) + 1  # Add 1 second buffer
+        else:
+            retry_seconds = 60  # fallback to 60 seconds if not found
+
+        print(f"Error during .invoke : {e} \nwaiting {retry_seconds} seconds")
+        time.sleep(retry_seconds)
         print("Waiting is finished")
         return var.invoke(parameters)
-        
+
 # Example usage (could be called from other modules)
 # main_llm = get_llm(settings.main_llm_model)
 # eval_llm = get_llm(settings.eval_llm_model)