Update agents/simple_tools.py

agents/simple_tools.py

@@ -3,6 +3,7 @@ Simplified LangChain tools for FinRyver financial statement generation
 Focus: Notes, Balance Sheet, P&L, Cash Flow generation only
 """
 from langchain_core.tools import tool
+import subprocess
 import os
 import subprocess
 import json
@@ -424,7 +425,6 @@ def generate_llm_notes(file_path: str, note_numbers: str = "", user_api_key: str
         if user_api_key:
             env["OPENROUTER_API_KEY"] = user_api_key
         elif os.getenv("OPENROUTER_API_KEY"):
-            # Fallback to environment variable if no user key provided
             env["OPENROUTER_API_KEY"] = os.getenv("OPENROUTER_API_KEY")
         
         cwd = os.getcwd()
@@ -436,7 +436,8 @@ def generate_llm_notes(file_path: str, note_numbers: str = "", user_api_key: str
             env=env,
             cwd=cwd,
             capture_output=True,
-            text=True
+            text=True,
+            timeout=120  # ADD THIS: 2 minute timeout
         )
 
         if result1.returncode != 0:
@@ -445,22 +446,22 @@ def generate_llm_notes(file_path: str, note_numbers: str = "", user_api_key: str
         # Step 2: Run LLM notes generator
         logger.info("Step 2: Generating notes using LLM")
         if note_numbers:
-            # Generate specific notes
             result2 = subprocess.run(
                 ["python", "notes/llm_notes_generator.py", "specific", note_numbers],
                 env=env,
                 cwd=cwd,
                 capture_output=True,
-                text=True
+                text=True,
+                timeout=600  # ADD THIS: 10 minute timeout for LLM generation
             )
         else:
-            # Generate all notes
             result2 = subprocess.run(
                 ["python", "notes/llm_notes_generator.py", "all", ""],
                 env=env,
                 cwd=cwd,
                 capture_output=True,
-                text=True
+                text=True,
+                timeout=900  # ADD THIS: 15 minute timeout for all notes
             )
 
         if result2.returncode != 0:
@@ -522,27 +523,13 @@ def generate_llm_notes(file_path: str, note_numbers: str = "", user_api_key: str
             except Exception as e:
                 logger.error(f"Error applying UDFs to JSON: {e}")
 
-        # Check if JSON file has content
-        try:
-            with open(input_json, 'r', encoding='utf-8') as f:
-                json_content = json.load(f)
-            if isinstance(json_content, dict) and 'notes' in json_content and not json_content['notes']:
-                logger.warning("JSON file exists but contains no notes")
-        except Exception as e:
-            execution_time = round(time.time() - start_time, 2)
-            return {
-                "status": "error",
-                "error": f"Invalid JSON file generated: {str(e)}",
-                "execution_id": execution_id,
-                "execution_time": execution_time
-            }
-
         result3 = subprocess.run(
             ["python", "notes/llm_notes_excel_converter.py", input_json, output_excel],
             env=env,
             cwd=cwd,
             capture_output=True,
-            text=True
+            text=True,
+            timeout=120  # ADD THIS: 2 minute timeout
         )
 
         if result3.returncode == 0:
@@ -563,6 +550,16 @@ def generate_llm_notes(file_path: str, note_numbers: str = "", user_api_key: str
                 "execution_time": execution_time
             }
 
+    except subprocess.TimeoutExpired as te:  # ADD THIS: Handle timeout
+        execution_time = round(time.time() - start_time, 2)
+        logger.error(f"LLM notes generation timed out after {execution_time}s")
+        return {
+            "status": "error",
+            "error": f"LLM notes generation timed out. The process took longer than expected. Try generating fewer notes or check your API key.",
+            "execution_id": execution_id,
+            "execution_time": execution_time,
+            "timeout": True
+        }
     except Exception as e:
         execution_time = round(time.time() - start_time, 2)
         logger.error(f"LLM notes generation failed: {e}")