Update agents/simple_tools.py

agents/simple_tools.py

@@ -77,11 +77,12 @@ def generate_notes_full_pipeline_from_path(file_path: str, note_numbers: str = "
     
 
 @tool
-def generate_balance_sheet(file_path: str, **kwargs) -> Dict[str, Any]:
+def generate_balance_sheet(file_path: str, user_api_key: str = None, **kwargs) -> Dict[str, Any]:
     """
     Generate balance sheet from trial balance file using complete pipeline
     Args:
         file_path: Path to trial balance Excel file
+        user_api_key: OpenRouter API key for LLM calls
     """
     execution_id = str(uuid.uuid4())[:8]
     start_time = time.time()
@@ -91,12 +92,19 @@ def generate_balance_sheet(file_path: str, **kwargs) -> Dict[str, Any]:
     try:
         # Use the complete BS pipeline from the existing endpoint
         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")
+        
+        # Get user-provided API key from parameter
+        logger.info(f"generate_balance_sheet: user_api_key received = {bool(user_api_key)}")
         if user_api_key:
             env["OPENROUTER_API_KEY"] = user_api_key
+            logger.info("generate_balance_sheet: Setting OPENROUTER_API_KEY from 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")
+            logger.info("generate_balance_sheet: Using OPENROUTER_API_KEY from environment")
+        else:
+            logger.warning("generate_balance_sheet: No API key available!")
+        
         env["INPUT_FILE"] = "data/clean_financial_data_bs.json"
         cwd = os.getcwd()
         
@@ -397,12 +405,13 @@ def generate_cash_flow_statement(file_path: str) -> Dict[str, Any]:
         }
 
 @tool
-def generate_llm_notes(file_path: str, note_numbers: str = "", **kwargs) -> Dict[str, Any]:
+def generate_llm_notes(file_path: str, note_numbers: str = "", user_api_key: str = None, **kwargs) -> Dict[str, Any]:
     """
     Generate notes using LLM-based approach (FlexibleFinancialNoteGenerator)
     Args:
         file_path: Path to trial balance Excel file
         note_numbers: Optional comma-separated note numbers to generate
+        user_api_key: OpenRouter API key for LLM calls
     Returns:
         dict with status, output_xlsx_path, and error if any
     """
@@ -412,12 +421,14 @@ def generate_llm_notes(file_path: str, note_numbers: str = "", **kwargs) -> Dict
     try:
         # Use the complete LLM notes pipeline from existing scripts
         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")
+        
+        # Get user-provided API key from parameter
         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()
 
         # Step 1: Run LLM notes data processor