Update ki_gen/planner.py

ki_gen/planner.py

@@ -34,17 +34,26 @@ If needed, give me an updated plan to follow this instruction. If your plan alre
     return {"messages" : [output]}
 
 
-def error_chatbot_groq(error,llm_groq,query):
+def error_chatbot_groq(error, model_name, query): # Pass model_name instead of llm_groq object
+    # Switch API key logic...
     if os.environ["GROQ_API_KEY"] == os.getenv("groq_api_key"):
         os.environ["GROQ_API_KEY"] = os.getenv("groq_api_key2")
-    
     elif os.environ["GROQ_API_KEY"] == os.getenv("groq_api_key2"):
         os.environ["GROQ_API_KEY"] = os.getenv("groq_api_key3")
-
     else:
         os.environ["GROQ_API_KEY"] = os.getenv("groq_api_key")
 
-    return {"messages" : [llm_groq.invoke(state["messages"])]}
+    # Re-initialize the model *after* switching the key
+    try:
+        # Use the model_name passed in
+        llm_groq_retry = ChatGroq(model=model_name)
+        # Pass the original query messages
+        return {"messages" : [llm_groq_retry.invoke(query)]}
+    except Exception as retry_error:
+        # Handle potential error during retry
+        print(f"Error during retry: {retry_error}")
+        # Decide what to return or raise here
+        return {"messages": [SystemMessage(content=f"Failed to process after retry: {retry_error}")]}
 
 
 # Wrappers to call LLMs on the state messsages field