Update agents/simple_tools.py

agents/simple_tools.py

@@ -77,7 +77,7 @@ def generate_notes_full_pipeline_from_path(file_path: str, note_numbers: str = "
     
 
 @tool
-def generate_balance_sheet(file_path: str) -> Dict[str, Any]:
+def generate_balance_sheet(file_path: str, **kwargs) -> Dict[str, Any]:
     """
     Generate balance sheet from trial balance file using complete pipeline
     Args:
@@ -93,6 +93,10 @@ def generate_balance_sheet(file_path: str) -> Dict[str, Any]:
         env = os.environ.copy()
         if os.getenv("OPENROUTER_API_KEY"):
             env["OPENROUTER_API_KEY"] = os.getenv("OPENROUTER_API_KEY")
+        # Override with user-provided API key if available
+        user_api_key = kwargs.get("user_api_key")
+        if user_api_key:
+            env["OPENROUTER_API_KEY"] = user_api_key
         env["INPUT_FILE"] = "data/clean_financial_data_bs.json"
         cwd = os.getcwd()
         
@@ -136,22 +140,43 @@ def generate_balance_sheet(file_path: str) -> Dict[str, Any]:
         )
         
         if result3.returncode == 0:
-            # Check for output files in data/output directory
-            output_files = []
-            if os.path.exists("data/output"):
-                output_files = [f for f in os.listdir("data/output") if f.endswith('.xlsx')]
+            # Parse the output file path from stdout
+            output_file_path = None
+            if result3.stdout:
+                for line in result3.stdout.strip().split('\n'):
+                    if line.startswith('Output file:'):
+                        output_file_path = line.split('Output file:', 1)[1].strip()
+                        break
             
-            execution_time = round(time.time() - start_time, 2)
+            # Fallback: check for output files in data/output directory
+            if not output_file_path or not os.path.exists(output_file_path):
+                output_dir = "data/output"
+                os.makedirs(output_dir, exist_ok=True)
+                output_files = []
+                if os.path.exists(output_dir):
+                    output_files = [f for f in os.listdir(output_dir) if f.endswith('.xlsx') and f.startswith('balance_sheet')]
+                    # Sort by modification time, get the most recent
+                    if output_files:
+                        output_files.sort(key=lambda x: os.path.getmtime(os.path.join(output_dir, x)), reverse=True)
+                        output_file_path = os.path.join(output_dir, output_files[0])
             
-  
-            return {
-                "status": "success",
-                "message": "Balance sheet generated successfully",
-                "output_path": "data/output/",
-                "output_files": output_files,
-                "execution_id": execution_id,
-                "execution_time": execution_time
-            }
+            if output_file_path and os.path.exists(output_file_path):
+                execution_time = round(time.time() - start_time, 2)
+                return {
+                    "status": "success",
+                    "message": "Balance sheet generated successfully",
+                    "output_path": output_file_path,
+                    "execution_id": execution_id,
+                    "execution_time": execution_time
+                }
+            else:
+                execution_time = round(time.time() - start_time, 2)
+                return {
+                    "status": "error",
+                    "error": "Balance sheet generation completed but output file not found",
+                    "execution_id": execution_id,
+                    "execution_time": execution_time
+                }
         else:
             execution_time = round(time.time() - start_time, 2)