Update app.py

app.py

@@ -12,10 +12,8 @@ from agents.langgraph import run_workflow
 from agents.rlhf_workflows import run_rlhf_workflow
 from agents.rlhf_routes import rlhf_router
 
-# Load environment variables from .env file
 load_dotenv()
 
-# Configure logging for the application
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger("financial_notes_api")
 
@@ -26,10 +24,8 @@ app = FastAPI(
     version="1.0.0"
 )
 
-# Initialize Interactive Feedback Manager
 feedback_manager = InteractiveFeedbackManager()
 
-# Include RLHF routes
 app.include_router(rlhf_router)
 @app.on_event("startup")
 async def startup_event():
@@ -43,26 +39,25 @@ router = APIRouter()
 
 
 @router.post("/notes-llm")
-async def notes_llm_route(file: UploadFile = File(...), use_rlhf: bool = Query(False)):
-    """Generate AI-powered financial notes using Generator-Validator pattern with Interactive Feedback"""
+async def notes_llm_route(
+    file: UploadFile = File(...), 
+    use_rlhf: bool = Query(False),
+    user_api_key: Optional[str] = Form(None)
+):
     file_path = f"data/input/{file.filename}"
     os.makedirs("data/input", exist_ok=True)
     with open(file_path, "wb") as buffer:
         shutil.copyfileobj(file.file, buffer)
 
     try:
-        # Create Generator-Validator pipeline for LLM notes
-        pipeline = create_notes_pipeline(use_rlhf=use_rlhf)
+        pipeline = create_notes_pipeline(use_rlhf=use_rlhf, user_api_key=user_api_key)
 
-        # Process through pipeline
         generation_result, validation_result = pipeline.process(file_path)
 
-        # Log processing summary
         summary = pipeline.get_processing_summary()
         logger.info(f"LLM Notes Pipeline Summary: {summary}")
 
         if generation_result.success and validation_result.is_valid:
-            # Create interactive feedback session
             session_id = feedback_manager.create_session(file_path)
 
             response = FileResponse(
@@ -70,22 +65,19 @@ async def notes_llm_route(file: UploadFile = File(...), use_rlhf: bool = Query(F
                 filename=os.path.basename(generation_result.output_path)
             )
 
-            # Add comprehensive metadata to headers
             response.headers["X-Generation-Method"] = "llm"
             response.headers["X-Validation-Score"] = str(validation_result.score)
             response.headers["X-Attempts-Made"] = str(generation_result.metadata.get("attempt", 1))
             response.headers["X-Execution-ID"] = generation_result.metadata.get("execution_id", "")
-            response.headers["X-Session-ID"] = session_id  # New: Interactive session ID
-            response.headers["X-Interactive-Enabled"] = "true"  # New: Indicates interactive mode available
+            response.headers["X-Session-ID"] = session_id
+            response.headers["X-Interactive-Enabled"] = "true"
 
-            # Add RLHF metadata if available
             if use_rlhf and "rlhf_metadata" in generation_result.metadata:
                 rlhf_data = generation_result.metadata["rlhf_metadata"]
                 response.headers["X-RLHF-Statement-ID"] = str(rlhf_data.get("statement_id", ""))
                 response.headers["X-RLHF-Quality-Score"] = str(rlhf_data.get("predicted_quality", ""))
                 response.headers["X-RLHF-Confidence"] = str(rlhf_data.get("confidence_score", ""))
 
-            # Add validation feedback
             if validation_result.feedback:
                 response.headers["X-Validation-Feedback"] = json.dumps(validation_result.feedback)
 
@@ -100,6 +92,9 @@ async def notes_llm_route(file: UploadFile = File(...), use_rlhf: bool = Query(F
             }
             raise HTTPException(status_code=500, detail=json.dumps(error_detail))
 
+    except ValueError as ve:
+        logger.error(f"API key error: {ve}")
+        raise HTTPException(status_code=400, detail=str(ve))
     except Exception as e:
         logger.error(f"LLM Notes pipeline failed: {e}")
         raise HTTPException(status_code=500, detail=f"Pipeline processing failed: {str(e)}")
@@ -111,9 +106,7 @@ async def submit_feedback(
     feedback_text: str = Form(...),
     feedback_type: str = Form(..., regex="^(text|numeric|formula|suggestion)$")
 ):
-    """Submit user feedback for iterative improvement"""
     try:
-        # Add feedback and generate UDF
         udf_version = feedback_manager.add_feedback(session_id, feedback_text, feedback_type)
 
         if udf_version is None:
@@ -134,7 +127,6 @@ async def submit_feedback(
 
 @router.post("/notes-llm/approve")
 async def approve_session(session_id: str = Form(...)):
-    """Approve the current session and set final UDF"""
     try:
         success = feedback_manager.approve_session(session_id)
 
@@ -158,7 +150,6 @@ async def approve_session(session_id: str = Form(...)):
 
 @router.get("/notes-llm/session/{session_id}")
 async def get_session_info(session_id: str):
-    """Get session information and feedback history"""
     try:
         session = feedback_manager.get_session(session_id)
 
@@ -194,9 +185,9 @@ async def get_session_info(session_id: str):
 @router.post("/notes-llm/generate")
 async def generate_with_feedback(
     session_id: str = Form(...),
-    file: UploadFile = File(...)
+    file: UploadFile = File(...),
+    user_api_key: Optional[str] = Form(None)
 ):
-    """Generate notes using accumulated feedback and UDFs"""
     try:
         session = feedback_manager.get_session(session_id)
 
@@ -206,16 +197,13 @@ async def generate_with_feedback(
         if session.status != 'active':
             raise HTTPException(status_code=400, detail=f"Session is {session.status}")
 
-        # Save uploaded file
         file_path = f"data/input/{file.filename}"
         os.makedirs("data/input", exist_ok=True)
         with open(file_path, "wb") as buffer:
             shutil.copyfileobj(file.file, buffer)
 
-        # Create pipeline with feedback integration
-        pipeline = create_notes_pipeline(use_rlhf=False)
+        pipeline = create_notes_pipeline(use_rlhf=False, user_api_key=user_api_key)
 
-        # Prepare feedback context for the generator
         udfs_to_apply = []
         if session.final_udf:
             udfs_to_apply.append(session.final_udf)
@@ -224,7 +212,7 @@ async def generate_with_feedback(
             
         feedback_context = {
             'session_id': session_id,
-            'udfs': udfs_to_apply,  # Pass final UDF if available, otherwise archived UDFs
+            'udfs': udfs_to_apply,
             'feedback_history': [
                 {
                     'text': f.feedback_text,
@@ -235,7 +223,6 @@ async def generate_with_feedback(
             'current_iteration': session.current_iteration
         }
 
-        # Process through pipeline with feedback context
         generation_result, validation_result = pipeline.process(file_path, feedback_context=feedback_context)
 
         if generation_result.success and validation_result.is_valid:
@@ -244,13 +231,11 @@ async def generate_with_feedback(
                 filename=os.path.basename(generation_result.output_path)
             )
 
-            # Add session and feedback metadata
             response.headers["X-Session-ID"] = session_id
             response.headers["X-Iteration"] = str(session.current_iteration)
             response.headers["X-Feedbacks-Applied"] = str(len(session.feedback_history))
             response.headers["X-UDFs-Archived"] = str(len(session.archived_udfs))
 
-            # Add generation metadata
             response.headers["X-Generation-Method"] = "llm_with_feedback"
             response.headers["X-Validation-Score"] = str(validation_result.score)
             response.headers["X-Execution-ID"] = generation_result.metadata.get("execution_id", "")
@@ -268,6 +253,9 @@ async def generate_with_feedback(
 
     except HTTPException:
         raise
+    except ValueError as ve:
+        logger.error(f"API key error: {ve}")
+        raise HTTPException(status_code=400, detail=str(ve))
     except Exception as e:
         logger.error(f"Feedback-based generation failed: {e}")
         raise HTTPException(status_code=500, detail=f"Generation failed: {str(e)}")
@@ -276,23 +264,18 @@ async def generate_with_feedback(
 
 @router.post("/notes")
 async def notes_route(file: UploadFile = File(...)):
-    """Generate financial notes directly from uploaded file"""
     try:
-        # Save uploaded file
         file_path = f"data/input/{file.filename}"
         os.makedirs("data/input", exist_ok=True)
         with open(file_path, "wb") as buffer:
             shutil.copyfileobj(file.file, buffer)
 
-        # Generate notes directly
         result = generate_notes_full_pipeline_from_path(file_path)
 
         if result["status"] == "success":
-            # Return the generated Excel file
             output_path = result["output_xlsx_path"]
             return FileResponse(output_path, filename=os.path.basename(output_path))
 
-        # If generation failed, raise HTTP exception
         raise HTTPException(status_code=500, detail=result.get("error", "Notes generation failed"))
 
     except Exception as e:
@@ -306,7 +289,6 @@ async def pnl_route(file: UploadFile = File(...), use_rlhf: bool = Query(False))
     with open(file_path, "wb") as buffer:
         shutil.copyfileobj(file.file, buffer)
     
-    # Choose workflow based on RLHF preference
     if use_rlhf:
         result = run_rlhf_workflow(file_path, "pnl")
     else:
@@ -315,7 +297,6 @@ async def pnl_route(file: UploadFile = File(...), use_rlhf: bool = Query(False))
     if result["status"] == "success":
         response = FileResponse(result["result"].get("output_path", "data/pnl_statement.xlsx"), filename=os.path.basename(result["result"].get("output_path", "data/pnl_statement.xlsx")))
         
-        # Add RLHF metadata to headers if available
         if "rlhf_metadata" in result.get("result", {}):
             rlhf_data = result["result"]["rlhf_metadata"]
             response.headers["X-RLHF-Statement-ID"] = str(rlhf_data.get("statement_id", ""))
@@ -332,17 +313,14 @@ async def bs_route(file: UploadFile = File(...), use_rlhf: bool = Query(False)):
     with open(file_path, "wb") as buffer:
         shutil.copyfileobj(file.file, buffer)
     
-    # Choose workflow based on RLHF preference
     if use_rlhf:
         result = run_rlhf_workflow(file_path, "bs")
     else:
         result = run_workflow(file_path, "bs")
     
     if result["status"] == "success":
-        # Use first xlsx file in output dir if present
         output_file = result["result"].get("output_path")
         if not output_file or not os.path.isfile(output_file):
-            # Try to find the first .xlsx file in data/output/ (ensure it's a file, not a directory)
             output_dir = "data/output/"
             xlsx_files = [f for f in os.listdir(output_dir) if f.endswith('.xlsx') and os.path.isfile(os.path.join(output_dir, f))]
             if xlsx_files:
@@ -352,7 +330,6 @@ async def bs_route(file: UploadFile = File(...), use_rlhf: bool = Query(False)):
         
         response = FileResponse(output_file, filename=os.path.basename(output_file))
         
-        # Add RLHF metadata to headers if available
         if "rlhf_metadata" in result.get("result", {}):
             rlhf_data = result["result"]["rlhf_metadata"]
             response.headers["X-RLHF-Statement-ID"] = str(rlhf_data.get("statement_id", ""))
@@ -370,7 +347,6 @@ async def cf_route(file: UploadFile = File(...), use_rlhf: bool = Query(False)):
     with open(file_path, "wb") as buffer:
         shutil.copyfileobj(file.file, buffer)
     
-    # Choose workflow based on RLHF preference
     if use_rlhf:
         result = run_rlhf_workflow(file_path, "cf")
     else:
@@ -379,7 +355,6 @@ async def cf_route(file: UploadFile = File(...), use_rlhf: bool = Query(False)):
     if result["status"] == "success":
         response = FileResponse(result["result"].get("output_path", "data/cash_flow_statements.xlsx"), filename=os.path.basename(result["result"].get("output_path", "data/cash_flow_statements.xlsx")))
         
-        # Add RLHF metadata to headers if available
         if "rlhf_metadata" in result.get("result", {}):
             rlhf_data = result["result"]["rlhf_metadata"]
             response.headers["X-RLHF-Statement-ID"] = str(rlhf_data.get("statement_id", ""))
@@ -392,4 +367,4 @@ app.include_router(router)
 
 if __name__ == "__main__":
     import uvicorn
-    uvicorn.run(app, host="0.0.0.0", port=7860)
+    uvicorn.run(app, host="0.0.0.0", port=7860)