Update agents/generator_validator.py

agents/generator_validator.py

@@ -433,9 +433,10 @@ class BaseValidator(ABC):
 class LLMNotesGenerator(BaseGenerator):
     """Generator for AI-powered financial notes"""
 
-    def __init__(self, max_attempts: int = 3, use_rlhf: bool = False):
+    def __init__(self, max_attempts: int = 3, use_rlhf: bool = False, user_api_key: Optional[str] = None):
         super().__init__(max_attempts)
         self.use_rlhf = use_rlhf
+        self.user_api_key = user_api_key
 
     def generate(self, file_path: str, **kwargs) -> GenerationResult:
         """Generate notes using AI/LLM approach with feedback integration"""
@@ -451,10 +452,10 @@ class LLMNotesGenerator(BaseGenerator):
             # Choose workflow based on RLHF preference
             if self.use_rlhf:
                 from agents.rlhf_workflows import run_rlhf_workflow
-                result = run_rlhf_workflow(file_path, "notes-llm")
+                result = run_rlhf_workflow(file_path, "notes-llm", user_api_key=self.user_api_key)
             else:
                 from agents.langgraph import run_workflow
-                result = run_workflow(file_path, "notes-llm", feedback_context=feedback_context)
+                result = run_workflow(file_path, "notes-llm", feedback_context=feedback_context, user_api_key=self.user_api_key)
 
             if result["status"] == "success":
                 # UDFs are now applied in generate_llm_notes function before Excel conversion
@@ -544,84 +545,6 @@ class LLMNotesGenerator(BaseGenerator):
         # 2. Retry with different parameters
         # 3. Use fallback models
 
-        if not self.use_rlhf and "quality" in str(feedback).lower():
-            # If quality issues and not using RLHF, try RLHF
-            logger.info("Switching to RLHF for better quality")
-            original_rlhf = self.use_rlhf
-            self.use_rlhf = True
-            result = self.generate(previous_result.data.get("file_path") if previous_result.data else None)
-            self.use_rlhf = original_rlhf  # Reset for future calls
-            return result
-        else:
-            # Otherwise, just retry
-            return self.generate(previous_result.data.get("file_path") if previous_result.data else None)
-    """Generator for AI-powered financial notes"""
-
-    def __init__(self, max_attempts: int = 3, use_rlhf: bool = False):
-        super().__init__(max_attempts)
-        self.use_rlhf = use_rlhf
-
-    def generate(self, file_path: str, **kwargs) -> GenerationResult:
-        """Generate notes using AI/LLM approach"""
-        try:
-            self.attempts_made += 1
-            execution_id = f"notes_llm_{datetime.now().strftime('%Y%m%d_%H%M%S')}_{self.attempts_made}"
-
-            # Choose workflow based on RLHF preference
-            if self.use_rlhf:
-                from agents.rlhf_workflows import run_rlhf_workflow
-                result = run_rlhf_workflow(file_path, "notes-llm")
-            else:
-                from agents.langgraph import run_workflow
-                result = run_workflow(file_path, "notes-llm")
-
-            if result["status"] == "success":
-                return GenerationResult(
-                    success=True,
-                    output_path=result["result"]["output_xlsx_path"],
-                    data=result["result"],
-                    error=None,
-                    metadata={
-                        "execution_id": execution_id,
-                        "generation_method": "llm",
-                        "use_rlhf": self.use_rlhf,
-                        "attempt": self.attempts_made,
-                        "rlhf_metadata": result["result"].get("rlhf_metadata", {})
-                    }
-                )
-            else:
-                return GenerationResult(
-                    success=False,
-                    output_path=None,
-                    data=None,
-                    error=result.get("error", "Unknown error"),
-                    metadata={
-                        "execution_id": execution_id,
-                        "generation_method": "llm",
-                        "use_rlhf": self.use_rlhf,
-                        "attempt": self.attempts_made
-                    }
-                )
-
-        except Exception as e:
-            logger.error(f"LLM Notes generation failed: {e}")
-            return GenerationResult(
-                success=False,
-                output_path=None,
-                data=None,
-                error=str(e),
-                metadata={
-                    "execution_id": f"error_{datetime.now().strftime('%Y%m%d_%H%M%S')}",
-                    "generation_method": "llm",
-                    "use_rlhf": self.use_rlhf,
-                    "attempt": self.attempts_made
-                }
-            )
-
-    def refine(self, previous_result: GenerationResult, feedback: List[str]) -> GenerationResult:
-        """Refine LLM notes generation based on feedback"""
-        logger.info(f"Refining LLM notes generation with feedback: {feedback}")
-
         if not self.use_rlhf and "quality" in str(feedback).lower():
             # If quality issues and not using RLHF, try RLHF
             logger.info("Switching to RLHF for better quality")
@@ -752,9 +675,9 @@ class GeneratorValidatorPipeline:
             "validation_criteria": self.validator.get_validation_criteria()
         }
 
-def create_notes_pipeline(use_rlhf: bool = False) -> GeneratorValidatorPipeline:
+def create_notes_pipeline(use_rlhf: bool = False, user_api_key: Optional[str] = None) -> GeneratorValidatorPipeline:
     """Factory function to create LLM-based pipeline for notes generation"""
-    generator = LLMNotesGenerator(use_rlhf=use_rlhf)
+    generator = LLMNotesGenerator(use_rlhf=use_rlhf, user_api_key=user_api_key)
     validator = NotesValidator()
 
     return GeneratorValidatorPipeline(generator, validator)