Update agent.py

mod invoke() to invoke_with_retry()

agent.py

@@ -24,6 +24,46 @@ from pydantic import Field
 
 from smolagents import WikipediaSearchTool
 
+def invoke_with_retry(
+    llm: ChatGoogleGenerativeAI,
+    prompt: str,
+    max_retries: int = 5,
+    initial_delay: int = 60
+):
+    """
+    Google Generative AIへのAPI呼び出しを、`ResourceExhausted`エラー時に再試行する関数。
+    
+    Args:
+        llm: ChatGoogleGenerativeAIのインスタンス。
+        prompt: ユーザーからのプロンプト文字列。
+        max_retries: 最大再試行回数。デフォルトは5。
+        initial_delay: 最初の再試行までの待機時間（秒）。デフォルトは60。
+        
+    Returns:
+        成功した場合のAPIレスポンス、失敗した場合はNone。
+    """
+    retries = 0
+    delay = initial_delay
+    
+    while retries < max_retries:
+        try:
+            messages = [HumanMessage(content=prompt)]
+            response = llm.invoke(messages)
+            return response
+        except ResourceExhausted as e:
+            print(f"APIアクセス上限を超えました。待機して再試行します。({retries + 1}/{max_retries})")
+            print(f"エラー詳細: {e}")
+            time.sleep(delay)
+            delay *= 2  # 指数バックオフ
+            retries += 1
+        except Exception as e:
+            # 他の予期せぬエラーに対する処理
+            print(f"予期せぬエラーが発生しました: {e}")
+            break
+            
+    print("最大再試行回数に達しました。API呼び出しに失敗しました。")
+    return None
+
 class SmolagentToolWrapper(BaseTool):
     """Wrapper for smolagents tools to make them compatible with LangChain."""
     
@@ -212,8 +252,9 @@ Please provide a detailed analysis focusing on:
 4. Target audience"""
 
                         # Use the LLM with proper message format
-                        messages = [HumanMessage(content=prompt)]
-                        response = self.llm.invoke(messages)
+                        #messages = [HumanMessage(content=prompt)]
+                        #response = self.llm.invoke(messages)
+                        response = invoke_with_retry(self.llm, prompt)
                         return response.content if hasattr(response, 'content') else str(response)
 
                     except Exception as e:
@@ -241,8 +282,9 @@ Provide a detailed analysis including:
 3. Notable findings
 4. Any mathematical properties (if applicable)"""
 
-            messages = [HumanMessage(content=prompt)]
-            response = self.llm.invoke(messages)
+            #messages = [HumanMessage(content=prompt)]
+            #response = self.llm.invoke(messages)
+            response = invoke_with_retry(self.llm, prompt)
             return response.content if hasattr(response, 'content') else str(response)
 
         except Exception as e:
@@ -262,8 +304,9 @@ Focus on:
 3. Text or numbers (if present)
 4. Overall context and meaning"""
 
-            messages = [HumanMessage(content=prompt)]
-            response = self.llm.invoke(messages)
+            #messages = [HumanMessage(content=prompt)]
+            #response = self.llm.invoke(messages)
+            response = invoke_with_retry(self.llm, prompt)
             return response.content if hasattr(response, 'content') else str(response)
 
         except Exception as e: