Update src/aibom_generator/api.py

src/aibom_generator/api.py

@@ -1,187 +1,195 @@
-from fastapi import FastAPI, HTTPException, Depends, BackgroundTasks, Query, File, UploadFile, Form, Request
-from fastapi.responses import JSONResponse, FileResponse, HTMLResponse, RedirectResponse
-from fastapi.middleware.cors import CORSMiddleware
-from fastapi.staticfiles import StaticFiles
-from fastapi.templating import Jinja2Templates
-from pydantic import BaseModel, Field
-from typing import Optional, Dict, Any, List
-import uvicorn
-import json
 import os
-import sys
+import json
+import logging
 import uuid
 import shutil
-import logging
+import tempfile
+from typing import Dict, List, Optional, Any, Union
 from datetime import datetime
+from pathlib import Path
+
+from fastapi import FastAPI, HTTPException, Depends, BackgroundTasks, Query, File, UploadFile, Form, Request
+from fastapi.responses import JSONResponse, FileResponse, HTMLResponse
+from fastapi.staticfiles import StaticFiles
+from fastapi.templating import Jinja2Templates
+from pydantic import BaseModel, Field
 
-# Set up logging
+# Configure logging
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
-# Add parent directory to path to import generator module
-sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
-
-# Import the AIBOM generator
-try:
-    from aibom_fix.final_generator import AIBOMGenerator
-except ImportError:
-    # If not found, try the mapping directory
-    try:
-        from aibom_mapping.final_generator import AIBOMGenerator
-    except ImportError:
-        # If still not found, use the original generator
-        try:
-            from aibom_fix.generator import AIBOMGenerator
-        except ImportError:
-            try:
-                from generator import AIBOMGenerator
-            except ImportError:
-                # Last resort: try to import from the aibom_generator module
-                try:
-                    from aibom_generator.generator import AIBOMGenerator
-                except ImportError:
-                    raise ImportError("Could not import AIBOMGenerator from any known location")
-
-# Create FastAPI app
-app = FastAPI(
-    title="Aetheris AI SBOM Generator API",
-    description="API for generating CycloneDX JSON AI SBOMs for AI models",
-    version="1.0.0",
-)
-
-# Add CORS middleware
-app.add_middleware(
-    CORSMiddleware,
-    allow_origins=["*"],  # Allow all origins in development
-    allow_credentials=True,
-    allow_methods=["*"],
-    allow_headers=["*"],
-)
-
-# Define output directory - use a directory that will have proper permissions
-# Instead of creating it at the module level, we'll create it in the startup event
-output_dir = "/tmp/aibom_output"  # Using /tmp which should be writable
-
 # Define templates directory
 templates_dir = "templates"
 
 # Initialize templates with a simple path
 templates = Jinja2Templates(directory=templates_dir)
 
-# Create a global generator instance
-generator = AIBOMGenerator(use_best_practices=True)
+# Create app
+app = FastAPI(title="AI SBOM Generator API")
 
-# Define request models
-class GenerateAIBOMRequest(BaseModel):
-    model_id: str = Field(..., description="The Hugging Face model ID (e.g., 'meta-llama/Llama-4-Scout-17B-16E-Instruct')")
-    hf_token: Optional[str] = Field(None, description="Optional Hugging Face API token for accessing private models")
-    include_inference: Optional[bool] = Field(True, description="Whether to use AI inference to enhance the AI SBOM")
-    use_best_practices: Optional[bool] = Field(True, description="Whether to use industry best practices for scoring")
+# Define output directory for generated AIBOMs
+OUTPUT_DIR = "/tmp/aibom_output"
 
-class AIBOMResponse(BaseModel):
-    aibom: Dict[str, Any] = Field(..., description="The generated AIBOM in CycloneDX JSON format")
-    model_id: str = Field(..., description="The model ID for which the AI SBOM was generated")
-    generated_at: str = Field(..., description="Timestamp when the AI SBOM was generated")
-    request_id: str = Field(..., description="Unique ID for this request")
-    download_url: Optional[str] = Field(None, description="URL to download the AIBOM JSON file")
+# Create output directory if it doesn't exist
+os.makedirs(OUTPUT_DIR, exist_ok=True)
 
-class EnhancementReport(BaseModel):
-    ai_enhanced: bool = Field(..., description="Whether AI enhancement was applied")
-    ai_model: Optional[str] = Field(None, description="The AI model used for enhancement, if any")
-    original_score: Dict[str, Any] = Field(..., description="Original completeness score before enhancement")
-    final_score: Dict[str, Any] = Field(..., description="Final completeness score after enhancement")
-    improvement: float = Field(..., description="Score improvement from enhancement")
+# Mount output directory as static files
+app.mount("/output", StaticFiles(directory=OUTPUT_DIR), name="output")
 
-class AIBOMWithReportResponse(AIBOMResponse):
-    enhancement_report: Optional[EnhancementReport] = Field(None, description="Report on AI enhancement results")
+# Define models
+class GenerateRequest(BaseModel):
+    model_id: str
+    include_inference: bool = False
+    use_best_practices: bool = True
 
-class StatusResponse(BaseModel):
-    status: str = Field(..., description="API status")
-    version: str = Field(..., description="API version")
-    generator_version: str = Field(..., description="AIBOM generator version")
+class GenerateWithReportRequest(BaseModel):
+    model_id: str
+    include_inference: bool = False
+    use_best_practices: bool = True
 
 class BatchGenerateRequest(BaseModel):
-    model_ids: List[str] = Field(..., description="List of Hugging Face model IDs to generate AIBOMs for")
-    hf_token: Optional[str] = Field(None, description="Optional Hugging Face API token for accessing private models")
-    include_inference: Optional[bool] = Field(True, description="Whether to use AI inference to enhance the AI SBOM")
-    use_best_practices: Optional[bool] = Field(True, description="Whether to use industry best practices for scoring")
+    model_ids: List[str]
+    include_inference: bool = False
+    use_best_practices: bool = True
 
-class BatchJobResponse(BaseModel):
-    job_id: str = Field(..., description="Unique ID for the batch job")
-    status: str = Field(..., description="Job status (e.g., 'queued', 'processing', 'completed')")
-    model_ids: List[str] = Field(..., description="List of model IDs in the batch")
-    created_at: str = Field(..., description="Timestamp when the job was created")
+class ModelScoreRequest(BaseModel):
+    model_id: str
+
+class StatusResponse(BaseModel):
+    status: str
+    version: str
+    generator_version: str
 
-# In-memory storage for batch jobs
-batch_jobs = {}
+class BatchJobResponse(BaseModel):
+    job_id: str
+    status: str
+    model_ids: List[str]
+    message: str
 
-# Startup event to create output directory
+# Startup event to ensure directories exist
 @app.on_event("startup")
 async def startup_event():
     """Create necessary directories on startup."""
-    os.makedirs(output_dir, exist_ok=True)
-    # Mount the static files directory after it's created
-    app.mount("/output", StaticFiles(directory=output_dir), name="output")
+    os.makedirs(OUTPUT_DIR, exist_ok=True)
+    logger.info(f"Created output directory at {OUTPUT_DIR}")
+
+# Helper function to ensure enhancement report has the expected structure
+def ensure_enhancement_report_structure(enhancement_report):
+    """Ensure the enhancement report has all the expected fields for templates."""
+    if enhancement_report is None:
+        enhancement_report = {}
+    
+    # Ensure ai_enhanced exists
+    if "ai_enhanced" not in enhancement_report:
+        enhancement_report["ai_enhanced"] = False
+    
+    # Ensure ai_model exists
+    if "ai_model" not in enhancement_report:
+        enhancement_report["ai_model"] = "Default AI Model"
+    
+    # Ensure original_score exists and has completeness_score
+    if "original_score" not in enhancement_report:
+        enhancement_report["original_score"] = {"total_score": 0, "completeness_score": 0}
+    elif "completeness_score" not in enhancement_report["original_score"]:
+        enhancement_report["original_score"]["completeness_score"] = enhancement_report["original_score"].get("total_score", 0)
+    
+    # Ensure final_score exists and has completeness_score
+    if "final_score" not in enhancement_report:
+        enhancement_report["final_score"] = {"total_score": 0, "completeness_score": 0}
+    elif "completeness_score" not in enhancement_report["final_score"]:
+        enhancement_report["final_score"]["completeness_score"] = enhancement_report["final_score"].get("total_score", 0)
+    
+    # Ensure improvement exists
+    if "improvement" not in enhancement_report:
+        enhancement_report["improvement"] = 0
+    
+    return enhancement_report
+
+# Helper function to ensure completeness_score has the expected structure
+def ensure_completeness_score_structure(completeness_score):
+    """Ensure the completeness_score has all the expected fields for templates."""
+    if completeness_score is None:
+        completeness_score = {}
+    
+    # Ensure total_score exists
+    if "total_score" not in completeness_score:
+        completeness_score["total_score"] = 0
     
-    # Log environment information for debugging
-    logger.info(f"Starting up API server")
-    logger.info(f"Current working directory: {os.getcwd()}")
-    logger.info(f"Directory contents: {os.listdir('.')}")
-    if os.path.exists(templates_dir):
-        logger.info(f"Templates directory exists")
-        logger.info(f"Templates directory contents: {os.listdir(templates_dir)}")
-    else:
-        logger.warning(f"Templates directory does not exist")
+    # Ensure completeness_score exists
+    if "completeness_score" not in completeness_score:
+        completeness_score["completeness_score"] = completeness_score.get("total_score", 0)
+    
+    # Ensure section_scores exists
+    if "section_scores" not in completeness_score:
+        completeness_score["section_scores"] = {}
+    
+    # Ensure max_scores exists
+    if "max_scores" not in completeness_score:
+        completeness_score["max_scores"] = {}
+    
+    # Ensure field_checklist exists
+    if "field_checklist" not in completeness_score:
+        completeness_score["field_checklist"] = {}
+    
+    # Ensure field_tiers exists
+    if "field_tiers" not in completeness_score:
+        completeness_score["field_tiers"] = {}
+    
+    # Ensure completeness_profile exists
+    if "completeness_profile" not in completeness_score:
+        completeness_score["completeness_profile"] = {
+            "name": "Incomplete",
+            "description": "This AI SBOM needs improvement to meet transparency requirements."
+        }
+    
+    return completeness_score
 
-# Define API endpoints
-# Root route now serves the UI
+# Root route serves the UI
 @app.get("/", response_class=HTMLResponse)
 async def root(request: Request):
     """Serve the web UI interface as the default view."""
     try:
-        logger.info("Root endpoint accessed - serving UI")
-        logger.info(f"Current working directory: {os.getcwd()}")
-        
+        # Check if templates directory exists
         if os.path.exists(templates_dir):
-            logger.info(f"Files in templates dir: {os.listdir(templates_dir)}")
+            logger.info(f"Templates directory exists")
+            logger.info(f"Templates directory contents: {os.listdir(templates_dir)}")
+        else:
+            logger.warning(f"Templates directory does not exist: {templates_dir}")
         
-        # Try using templates first
+        # Try to render the template
         try:
             return templates.TemplateResponse("index.html", {"request": request})
         except Exception as e:
-            logger.error(f"Error using templates: {str(e)}")
+            logger.error(f"Error rendering template: {str(e)}")
             
-            # Fallback to direct file reading
+            # Try direct file reading as fallback
             try:
-                template_path = os.path.join(templates_dir, "index.html")
-                logger.info(f"Trying direct file read from: {template_path}")
+                with open(os.path.join(templates_dir, "index.html"), "r") as f:
+                    html_content = f.read()
+                return HTMLResponse(content=html_content)
+            except Exception as file_e:
+                logger.error(f"Error reading file directly: {str(file_e)}")
                 
-                if os.path.exists(template_path):
-                    with open(template_path, "r") as f:
+                # Last resort: try absolute path
+                try:
+                    with open("/app/templates/index.html", "r") as f:
                         html_content = f.read()
                     return HTMLResponse(content=html_content)
-                else:
-                    logger.error(f"Template file not found at: {template_path}")
-                    
-                    # Try absolute path as last resort
-                    abs_path = os.path.join(os.getcwd(), templates_dir, "index.html")
-                    logger.info(f"Trying absolute path: {abs_path}")
-                    
-                    if os.path.exists(abs_path):
-                        with open(abs_path, "r") as f:
-                            html_content = f.read()
-                        return HTMLResponse(content=html_content)
-                    else:
-                        logger.error(f"Template file not found at absolute path: {abs_path}")
-                        raise HTTPException(status_code=404, detail="UI template not found")
-            except Exception as file_error:
-                logger.error(f"Error reading file directly: {str(file_error)}")
-                raise HTTPException(status_code=500, detail=f"Error serving UI: {str(file_error)}")
-    except Exception as outer_error:
-        logger.error(f"Outer error in UI endpoint: {str(outer_error)}")
-        raise HTTPException(status_code=500, detail=f"Error in UI endpoint: {str(outer_error)}")
+                except Exception as abs_e:
+                    logger.error(f"Error reading file from absolute path: {str(abs_e)}")
+                    raise HTTPException(status_code=500, detail=f"Error in UI endpoint: {str(e)}")
+    except Exception as outer_e:
+        logger.error(f"Outer error in UI endpoint: {str(outer_e)}")
+        raise HTTPException(status_code=500, detail=f"Error in UI endpoint: {str(outer_e)}")
 
-# Status endpoint moved to /api/status
+# UI route for backward compatibility
+@app.get("/ui", response_class=HTMLResponse)
+async def ui(request: Request):
+    """Serve the web UI interface (kept for backward compatibility)."""
+    return await root(request)
+
+# Status endpoint - both at root and /api path for compatibility
+@app.get("/status", response_model=StatusResponse)
 @app.get("/api/status", response_model=StatusResponse)
 async def get_status():
     """Get the API status and version information."""
@@ -191,525 +199,447 @@ async def get_status():
         "generator_version": "1.0.0",
     }
 
-# Keep the /ui endpoint for backward compatibility
-@app.get("/ui", response_class=HTMLResponse)
-async def ui(request: Request):
-    """Serve the web UI interface (kept for backward compatibility)."""
-    return await root(request)
-
-# Form-based generate endpoint for HTML form submissions
+# Form-based generate endpoint for web UI
 @app.post("/generate", response_class=HTMLResponse)
-async def generate_aibom_form(request: Request, model_id: str = Form(...)):
-    """
-    Generate a CycloneDX JSON AI SBOM for a Hugging Face model from form submission.
-    
-    This endpoint handles traditional HTML form submissions and renders the result template.
-    """
+async def generate_form(
+    request: Request,
+    model_id: str = Form(...),
+    include_inference: bool = Form(False),
+    use_best_practices: bool = Form(True)
+):
+    """Generate an AI SBOM from form data and render the result template."""
     try:
-        logger.info(f"Form-based generate endpoint called with model_id: {model_id}")
-        
-        # Create a new generator instance
-        gen = AIBOMGenerator(
-            use_inference=True,  # Default to true for form submissions
-            use_best_practices=True
-        )
-        
-        # Generate a request ID
-        request_id = str(uuid.uuid4())
+        # Import the generator here to avoid circular imports
+        try:
+            from src.aibom_generator.generator import AIBOMGenerator
+        except ImportError:
+            try:
+                from aibom_generator.generator import AIBOMGenerator
+            except ImportError:
+                try:
+                    from generator import AIBOMGenerator
+                except ImportError:
+                    raise ImportError("Could not import AIBOMGenerator. Please check your installation.")
         
-        # Create output file path
-        safe_model_id = model_id.replace("/", "_")
-        output_file = os.path.join(output_dir, f"{safe_model_id}_{request_id}.json")
+        # Create generator instance
+        generator = AIBOMGenerator()
         
-        # Generate the AIBOM
-        aibom = gen.generate_aibom(
+        # Generate AIBOM
+        aibom, enhancement_report = generator.generate(
             model_id=model_id,
-            include_inference=True,
-            use_best_practices=True,
-            output_file=output_file
+            include_inference=include_inference,
+            use_best_practices=use_best_practices
         )
         
-        # Get the enhancement report
-        enhancement_report = gen.get_enhancement_report()
-        
-        # Create a completely new enhancement report structure that matches what the template expects
-        # This is a more direct approach than trying to modify the existing structure
-        template_enhancement_report = {
-            "ai_enhanced": True,
-            "ai_model": "Default AI Model",
-            "original_score": {
-                "total_score": 0,
-                "completeness_score": 0  # Explicitly add completeness_score
-            },
-            "final_score": {
-                "total_score": 0,
-                "completeness_score": 0  # Explicitly add completeness_score
-            },
-            "improvement": 0
-        }
-        
-        # If we have a real enhancement report, copy the values we have
-        if enhancement_report:
-            template_enhancement_report["ai_enhanced"] = enhancement_report.get("ai_enhanced", True)
-            template_enhancement_report["ai_model"] = enhancement_report.get("ai_model", "Default AI Model")
-            
-            if "original_score" in enhancement_report:
-                template_enhancement_report["original_score"]["total_score"] = enhancement_report["original_score"].get("total_score", 0)
-                # Always set completeness_score equal to total_score if not present
-                template_enhancement_report["original_score"]["completeness_score"] = enhancement_report["original_score"].get("completeness_score", 
-                                                                                     enhancement_report["original_score"].get("total_score", 0))
-            
-            if "final_score" in enhancement_report:
-                template_enhancement_report["final_score"]["total_score"] = enhancement_report["final_score"].get("total_score", 0)
-                # Always set completeness_score equal to total_score if not present
-                template_enhancement_report["final_score"]["completeness_score"] = enhancement_report["final_score"].get("completeness_score", 
-                                                                                  enhancement_report["final_score"].get("total_score", 0))
-            
-            template_enhancement_report["improvement"] = enhancement_report.get("improvement", 0)
+        # Save AIBOM to file
+        filename = f"{model_id.replace('/', '_')}_aibom.json"
+        filepath = os.path.join(OUTPUT_DIR, filename)
+        with open(filepath, "w") as f:
+            json.dump(aibom, f, indent=2)
         
         # Create download URL
-        download_url = f"/output/{os.path.basename(output_file)}"
+        download_url = f"/output/{filename}"
         
-        # Add a downloadJSON function for the template
-        download_script = f"""
+        # Create download script
+        download_script = """
         <script>
-        function downloadJSON() {{
-            const a = document.createElement('a');
-            a.href = '{download_url}';
-            a.download = '{safe_model_id}_aibom.json';
-            document.body.appendChild(a);
-            a.click();
-            document.body.removeChild(a);
-        }}
-        
-        // Add tab switching functionality
-        function switchTab(tabId) {{
-            // Hide all tab contents
-            const tabContents = document.getElementsByClassName('tab-content');
-            for (let i = 0; i < tabContents.length; i++) {{
-                tabContents[i].classList.remove('active');
-            }}
-            
-            // Deactivate all tabs
-            const tabs = document.getElementsByClassName('aibom-tab');
-            for (let i = 0; i < tabs.length; i++) {{
-                tabs[i].classList.remove('active');
-            }}
-            
-            // Activate the selected tab and content
-            document.getElementById(tabId).classList.add('active');
-            document.querySelector('.aibom-tab[onclick="switchTab(\\'' + tabId + '\\')"]').classList.add('active');
-        }}
-        
-        // Add collapsible functionality
-        function toggleCollapsible(element) {{
-            element.classList.toggle('active');
-            const content = element.nextElementSibling;
-            if (content.classList.contains('active')) {{
-                content.classList.remove('active');
-            }} else {{
-                content.classList.add('active');
-            }}
-        }}
+            function downloadJSON() {
+                const a = document.createElement('a');
+                a.href = '%s';
+                a.download = '%s';
+                document.body.appendChild(a);
+                a.click();
+                document.body.removeChild(a);
+            }
         </script>
-        """
+        """ % (download_url, filename)
         
-        # Log the enhancement report structure for debugging
-        logger.info(f"Template enhancement report structure: {template_enhancement_report}")
+        # Process enhancement report for template
+        enhancement_report = ensure_enhancement_report_structure(enhancement_report)
+        logger.info(f"Enhancement report structure: {enhancement_report}")
         
-        # Render the result template with our carefully constructed enhancement report
+        # Get completeness score
+        completeness_score = None
+        if hasattr(generator, 'get_completeness_score'):
+            try:
+                completeness_score = generator.get_completeness_score(model_id)
+                completeness_score = ensure_completeness_score_structure(completeness_score)
+                logger.info(f"Completeness score structure: {completeness_score}")
+            except Exception as e:
+                logger.error(f"Error getting completeness score: {str(e)}")
+                completeness_score = ensure_completeness_score_structure(None)
+        else:
+            # Create a default completeness score if the method doesn't exist
+            completeness_score = ensure_completeness_score_structure(None)
+        
+        # Render result template
         return templates.TemplateResponse(
-            "result.html", 
+            "result.html",
             {
                 "request": request,
                 "model_id": model_id,
                 "aibom": aibom,
-                "enhancement_report": template_enhancement_report,
+                "enhancement_report": enhancement_report,
+                "completeness_score": completeness_score,
                 "download_url": download_url,
                 "download_script": download_script
             }
         )
     except Exception as e:
         logger.error(f"Error generating AI SBOM: {str(e)}")
-        # Render the error template
         return templates.TemplateResponse(
-            "error.html", 
+            "error.html",
             {
                 "request": request,
-                "error": str(e)
+                "error_message": f"Error generating AI SBOM: {str(e)}"
             }
         )
 
-# JSON-based generate endpoint for API clients
-@app.post("/api/generate", response_model=AIBOMResponse)
-async def generate_aibom_json(request: GenerateAIBOMRequest):
-    """
-    Generate a CycloneDX JSON AI SBOM for a Hugging Face model.
-    
-    This endpoint takes a JSON request and returns a JSON response for API clients.
-    """
+# JSON API endpoints
+@app.post("/api/generate")
+async def generate_api(request: GenerateRequest):
+    """Generate an AI SBOM and return it as JSON."""
     try:
-        # Create a new generator instance with the provided token if available
-        gen = AIBOMGenerator(
-            hf_token=request.hf_token,
-            use_inference=request.include_inference,
-            use_best_practices=request.use_best_practices
-        )
-        
-        # Generate a request ID
-        request_id = str(uuid.uuid4())
+        # Import the generator here to avoid circular imports
+        try:
+            from src.aibom_generator.generator import AIBOMGenerator
+        except ImportError:
+            try:
+                from aibom_generator.generator import AIBOMGenerator
+            except ImportError:
+                try:
+                    from generator import AIBOMGenerator
+                except ImportError:
+                    raise ImportError("Could not import AIBOMGenerator. Please check your installation.")
         
-        # Create output file path
-        safe_model_id = request.model_id.replace("/", "_")
-        output_file = os.path.join(output_dir, f"{safe_model_id}_{request_id}.json")
+        # Create generator instance
+        generator = AIBOMGenerator()
         
-        # Generate the AIBOM
-        aibom = gen.generate_aibom(
+        # Generate AIBOM
+        aibom, _ = generator.generate(
             model_id=request.model_id,
             include_inference=request.include_inference,
-            use_best_practices=request.use_best_practices,
-            output_file=output_file
+            use_best_practices=request.use_best_practices
         )
         
-        # Create download URL
-        download_url = f"/output/{os.path.basename(output_file)}"
-        
-        # Create response
-        response = {
-            "aibom": aibom,
-            "model_id": request.model_id,
-            "generated_at": datetime.utcnow().isoformat() + "Z",
-            "request_id": request_id,
-            "download_url": download_url
-        }
-        
-        return response
+        return aibom
     except Exception as e:
         raise HTTPException(status_code=500, detail=f"Error generating AI SBOM: {str(e)}")
 
-# JSON-based generate with report endpoint for API clients
-@app.post("/api/generate-with-report", response_model=AIBOMWithReportResponse)
-async def generate_aibom_with_report_json(request: GenerateAIBOMRequest):
-    """
-    Generate a CycloneDX JSON AI SBOM with an enhancement report.
-    
-    This endpoint is similar to /api/generate but also includes a report
-    on the AI enhancement results, including before/after scores.
-    """
+@app.post("/api/generate-with-report")
+async def generate_with_report_api(request: GenerateWithReportRequest):
+    """Generate an AI SBOM with enhancement report and return both as JSON."""
     try:
-        # Create a new generator instance with the provided token if available
-        gen = AIBOMGenerator(
-            hf_token=request.hf_token,
-            use_inference=request.include_inference,
-            use_best_practices=request.use_best_practices
-        )
-        
-        # Generate a request ID
-        request_id = str(uuid.uuid4())
+        # Import the generator here to avoid circular imports
+        try:
+            from src.aibom_generator.generator import AIBOMGenerator
+        except ImportError:
+            try:
+                from aibom_generator.generator import AIBOMGenerator
+            except ImportError:
+                try:
+                    from generator import AIBOMGenerator
+                except ImportError:
+                    raise ImportError("Could not import AIBOMGenerator. Please check your installation.")
         
-        # Create output file path
-        safe_model_id = request.model_id.replace("/", "_")
-        output_file = os.path.join(output_dir, f"{safe_model_id}_{request_id}.json")
+        # Create generator instance
+        generator = AIBOMGenerator()
         
-        # Generate the AIBOM
-        aibom = gen.generate_aibom(
+        # Generate AIBOM
+        aibom, enhancement_report = generator.generate(
             model_id=request.model_id,
             include_inference=request.include_inference,
-            use_best_practices=request.use_best_practices,
-            output_file=output_file
+            use_best_practices=request.use_best_practices
         )
         
-        # Get the enhancement report
-        enhancement_report = gen.get_enhancement_report()
-        
-        # Create a new enhancement report structure for the API response
-        api_enhancement_report = {
-            "ai_enhanced": True,
-            "ai_model": "Default AI Model",
-            "original_score": {
-                "total_score": 0,
-                "completeness_score": 0
-            },
-            "final_score": {
-                "total_score": 0,
-                "completeness_score": 0
-            },
-            "improvement": 0
-        }
+        # Process enhancement report for API response
+        enhancement_report = ensure_enhancement_report_structure(enhancement_report)
         
-        # If we have a real enhancement report, copy the values we have
-        if enhancement_report:
-            api_enhancement_report["ai_enhanced"] = enhancement_report.get("ai_enhanced", True)
-            api_enhancement_report["ai_model"] = enhancement_report.get("ai_model", "Default AI Model")
-            
-            if "original_score" in enhancement_report:
-                api_enhancement_report["original_score"]["total_score"] = enhancement_report["original_score"].get("total_score", 0)
-                api_enhancement_report["original_score"]["completeness_score"] = enhancement_report["original_score"].get("completeness_score", 
-                                                                               enhancement_report["original_score"].get("total_score", 0))
-            
-            if "final_score" in enhancement_report:
-                api_enhancement_report["final_score"]["total_score"] = enhancement_report["final_score"].get("total_score", 0)
-                api_enhancement_report["final_score"]["completeness_score"] = enhancement_report["final_score"].get("completeness_score", 
-                                                                            enhancement_report["final_score"].get("total_score", 0))
-            
-            api_enhancement_report["improvement"] = enhancement_report.get("improvement", 0)
-        
-        # Create download URL
-        download_url = f"/output/{os.path.basename(output_file)}"
-        
-        # Create response
-        response = {
+        return {
             "aibom": aibom,
-            "model_id": request.model_id,
-            "generated_at": datetime.utcnow().isoformat() + "Z",
-            "request_id": request_id,
-            "download_url": download_url,
-            "enhancement_report": api_enhancement_report
+            "enhancement_report": enhancement_report
         }
-        
-        return response
     except Exception as e:
-        raise HTTPException(status_code=500, detail=f"Error generating AIBOM: {str(e)}")
+        raise HTTPException(status_code=500, detail=f"Error generating AI SBOM: {str(e)}")
 
-@app.get("/models/{model_id}/score", response_model=Dict[str, Any])
-async def get_model_score(
-    model_id: str,
-    hf_token: Optional[str] = Query(None, description="Optional Hugging Face API token for accessing private models"),
-    use_best_practices: bool = Query(True, description="Whether to use industry best practices for scoring")
+@app.post("/api/batch-generate")
+async def batch_generate_api(
+    request: BatchGenerateRequest,
+    background_tasks: BackgroundTasks
 ):
-    """
-    Get the completeness score for a model without generating a full AI SBOM.
-    
-    This is a lightweight endpoint that only returns the scoring information.
-    """
+    """Start a batch job to generate multiple AI SBOMs."""
     try:
-        # Create a new generator instance with the provided token if available
-        gen = AIBOMGenerator(
-            hf_token=hf_token,
-            use_inference=False,  # Don't use inference for scoring only
-            use_best_practices=use_best_practices
-        )
-        
-        # Generate the AIBOM (needed to calculate score)
-        aibom = gen.generate_aibom(
-            model_id=model_id,
-            include_inference=False,  # Don't use inference for scoring only
-            use_best_practices=use_best_practices
-        )
-        
-        # Get the enhancement report for the score
-        enhancement_report = gen.get_enhancement_report()
-        
-        # Create a score response with the expected structure
-        score_response = {
-            "total_score": 0,
-            "completeness_score": 0
-        }
-        
-        if enhancement_report and "final_score" in enhancement_report:
-            score_response["total_score"] = enhancement_report["final_score"].get("total_score", 0)
-            score_response["completeness_score"] = enhancement_report["final_score"].get("completeness_score", 
-                                                                                      enhancement_report["final_score"].get("total_score", 0))
-        
-        return score_response
-    except Exception as e:
-        raise HTTPException(status_code=500, detail=f"Error calculating score: {str(e)}")
-
-@app.post("/batch", response_model=BatchJobResponse)
-async def batch_generate(request: BatchGenerateRequest, background_tasks: BackgroundTasks):
-    """
-    Start a batch job to generate AI SBOMs for multiple models.
-    
-    This endpoint queues a background task to generate AI SBOMs for all the
-    specified model IDs and returns a job ID that can be used to check status.
-    """
-    try:
-        # Generate a job ID
+        # Generate job ID
         job_id = str(uuid.uuid4())
         
-        # Create job record
-        job = {
-            "job_id": job_id,
-            "status": "queued",
-            "model_ids": request.model_ids,
-            "created_at": datetime.utcnow().isoformat() + "Z",
-            "completed": 0,
-            "total": len(request.model_ids),
-            "results": {}
-        }
-        
-        # Store job in memory
-        batch_jobs[job_id] = job
-        
-        # Add background task to process the batch
+        # Create job directory
+        job_dir = os.path.join(OUTPUT_DIR, job_id)
+        os.makedirs(job_dir, exist_ok=True)
+        
+        # Create job status file
+        status_file = os.path.join(job_dir, "status.json")
+        with open(status_file, "w") as f:
+            json.dump({
+                "job_id": job_id,
+                "status": "pending",
+                "model_ids": request.model_ids,
+                "completed": 0,
+                "total": len(request.model_ids),
+                "results": {}
+            }, f, indent=2)
+        
+        # Start background task
         background_tasks.add_task(
             process_batch_job,
             job_id=job_id,
             model_ids=request.model_ids,
-            hf_token=request.hf_token,
             include_inference=request.include_inference,
             use_best_practices=request.use_best_practices
         )
         
-        # Return job info
         return {
             "job_id": job_id,
-            "status": "queued",
+            "status": "pending",
             "model_ids": request.model_ids,
-            "created_at": job["created_at"]
+            "message": f"Batch job started. {len(request.model_ids)} models will be processed."
         }
     except Exception as e:
         raise HTTPException(status_code=500, detail=f"Error starting batch job: {str(e)}")
 
-@app.get("/batch/{job_id}", response_model=Dict[str, Any])
-async def get_batch_status(job_id: str):
-    """
-    Get the status of a batch job.
-    
-    This endpoint returns the current status of a batch job, including
-    progress information and results for completed models.
-    """
-    if job_id not in batch_jobs:
-        raise HTTPException(status_code=404, detail="Batch job not found")
-    
-    return batch_jobs[job_id]
+@app.get("/api/batch-status/{job_id}")
+async def batch_status_api(job_id: str):
+    """Get the status of a batch job."""
+    try:
+        # Check if job exists
+        job_dir = os.path.join(OUTPUT_DIR, job_id)
+        if not os.path.exists(job_dir):
+            raise HTTPException(status_code=404, detail=f"Job {job_id} not found")
+        
+        # Read job status
+        status_file = os.path.join(job_dir, "status.json")
+        with open(status_file, "r") as f:
+            status = json.load(f)
+        
+        return status
+    except HTTPException:
+        raise
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=f"Error getting job status: {str(e)}")
+
+@app.get("/api/batch-results/{job_id}")
+async def batch_results_api(job_id: str):
+    """Get the results of a completed batch job."""
+    try:
+        # Check if job exists
+        job_dir = os.path.join(OUTPUT_DIR, job_id)
+        if not os.path.exists(job_dir):
+            raise HTTPException(status_code=404, detail=f"Job {job_id} not found")
+        
+        # Read job status
+        status_file = os.path.join(job_dir, "status.json")
+        with open(status_file, "r") as f:
+            status = json.load(f)
+        
+        # Check if job is completed
+        if status["status"] != "completed":
+            raise HTTPException(status_code=400, detail=f"Job {job_id} is not completed yet")
+        
+        # Read results
+        results = {}
+        for model_id, result_file in status["results"].items():
+            with open(os.path.join(job_dir, result_file), "r") as f:
+                results[model_id] = json.load(f)
+        
+        return results
+    except HTTPException:
+        raise
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=f"Error getting job results: {str(e)}")
+
+@app.get("/api/model-score/{model_id}")
+async def model_score_api(model_id: str):
+    """Get the completeness score for a model."""
+    try:
+        # Import the generator here to avoid circular imports
+        try:
+            from src.aibom_generator.generator import AIBOMGenerator
+        except ImportError:
+            try:
+                from aibom_generator.generator import AIBOMGenerator
+            except ImportError:
+                try:
+                    from generator import AIBOMGenerator
+                except ImportError:
+                    raise ImportError("Could not import AIBOMGenerator. Please check your installation.")
+        
+        # Create generator instance
+        generator = AIBOMGenerator()
+        
+        # Get completeness score
+        if hasattr(generator, 'get_completeness_score'):
+            completeness_score = generator.get_completeness_score(model_id)
+            completeness_score = ensure_completeness_score_structure(completeness_score)
+            return completeness_score
+        else:
+            # Create a default completeness score if the method doesn't exist
+            completeness_score = ensure_completeness_score_structure(None)
+            return completeness_score
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=f"Error getting model score: {str(e)}")
 
-@app.post("/upload-model-card")
-async def upload_model_card(
+@app.post("/api/upload-model-card")
+async def upload_model_card_api(
     model_id: str = Form(...),
-    model_card: UploadFile = File(...),
-    include_inference: bool = Form(True),
-    use_best_practices: bool = Form(True)
+    model_card: UploadFile = File(...)
 ):
-    """
-    Generate an AI SBOM from an uploaded model card file.
-    
-    This endpoint allows users to upload a model card file directly
-    instead of requiring the model to be on Hugging Face.
-    """
+    """Upload a model card file for a model."""
     try:
-        # Create a temporary directory to store the uploaded file
-        temp_dir = os.path.join("/tmp", "aibom_temp")
-        os.makedirs(temp_dir, exist_ok=True)
+        # Create temporary file
+        with tempfile.NamedTemporaryFile(delete=False, suffix=".md") as temp_file:
+            # Write uploaded file to temporary file
+            shutil.copyfileobj(model_card.file, temp_file)
         
-        # Save the uploaded file
-        file_path = os.path.join(temp_dir, model_card.filename)
-        with open(file_path, "wb") as f:
-            shutil.copyfileobj(model_card.file, f)
+        # Import the generator here to avoid circular imports
+        try:
+            from src.aibom_generator.generator import AIBOMGenerator
+        except ImportError:
+            try:
+                from aibom_generator.generator import AIBOMGenerator
+            except ImportError:
+                try:
+                    from generator import AIBOMGenerator
+                except ImportError:
+                    raise ImportError("Could not import AIBOMGenerator. Please check your installation.")
         
-        # TODO: Implement custom model card processing
-        # This would require modifying the AIBOMGenerator to accept a file path
-        # instead of a model ID, which is beyond the scope of this example
+        # Create generator instance
+        generator = AIBOMGenerator()
         
-        # For now, return a placeholder response
-        return {
-            "status": "not_implemented",
-            "message": "Custom model card processing is not yet implemented"
-        }
+        # Process model card
+        if hasattr(generator, 'process_model_card'):
+            result = generator.process_model_card(model_id, temp_file.name)
+            
+            # Clean up temporary file
+            os.unlink(temp_file.name)
+            
+            return result
+        else:
+            # Clean up temporary file
+            os.unlink(temp_file.name)
+            
+            raise HTTPException(status_code=501, detail="Model card processing not implemented")
     except Exception as e:
-        raise HTTPException(status_code=500, detail=f"Error processing uploaded model card: {str(e)}")
+        raise HTTPException(status_code=500, detail=f"Error processing model card: {str(e)}")
 
 @app.get("/download/{filename}")
-async def download_aibom(filename: str):
-    """
-    Download a previously generated AI SBOM file.
-    
-    This endpoint allows downloading AI SBOM files by filename.
-    """
-    file_path = os.path.join(output_dir, filename)
-    if not os.path.exists(file_path):
-        raise HTTPException(status_code=404, detail="File not found")
-    
-    return FileResponse(file_path, media_type="application/json", filename=filename)
+async def download_file(filename: str):
+    """Download a generated AI SBOM file."""
+    try:
+        filepath = os.path.join(OUTPUT_DIR, filename)
+        if not os.path.exists(filepath):
+            raise HTTPException(status_code=404, detail=f"File {filename} not found")
+        
+        return FileResponse(
+            filepath,
+            media_type="application/json",
+            filename=filename
+        )
+    except HTTPException:
+        raise
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=f"Error downloading file: {str(e)}")
 
-# Background task function for batch processing
-async def process_batch_job(job_id: str, model_ids: List[str], hf_token: Optional[str], include_inference: bool, use_best_practices: bool):
+# Background task for batch processing
+async def process_batch_job(
+    job_id: str,
+    model_ids: List[str],
+    include_inference: bool,
+    use_best_practices: bool
+):
     """Process a batch job in the background."""
-    # Update job status
-    batch_jobs[job_id]["status"] = "processing"
-    
-    # Create output directory
-    batch_output_dir = os.path.join(output_dir, job_id)
-    os.makedirs(batch_output_dir, exist_ok=True)
-    
-    # Process each model
-    for model_id in model_ids:
+    try:
+        # Import the generator here to avoid circular imports
         try:
-            # Create a new generator instance
-            gen = AIBOMGenerator(
-                hf_token=hf_token,
-                use_inference=include_inference,
-                use_best_practices=use_best_practices
-            )
-            
-            # Create output file path
-            safe_model_id = model_id.replace("/", "_")
-            output_file = os.path.join(batch_output_dir, f"{safe_model_id}.json")
-            
-            # Generate the AIBOM
-            aibom = gen.generate_aibom(
-                model_id=model_id,
-                include_inference=include_inference,
-                use_best_practices=use_best_practices,
-                output_file=output_file
-            )
-            
-            # Get the enhancement report
-            enhancement_report = gen.get_enhancement_report()
-            
-            # Create a new enhancement report structure for the batch results
-            batch_enhancement_report = {
-                "ai_enhanced": True,
-                "ai_model": "Default AI Model",
-                "original_score": {
-                    "total_score": 0,
-                    "completeness_score": 0
-                },
-                "final_score": {
-                    "total_score": 0,
-                    "completeness_score": 0
-                },
-                "improvement": 0
-            }
-            
-            # If we have a real enhancement report, copy the values we have
-            if enhancement_report:
-                batch_enhancement_report["ai_enhanced"] = enhancement_report.get("ai_enhanced", True)
-                batch_enhancement_report["ai_model"] = enhancement_report.get("ai_model", "Default AI Model")
+            from src.aibom_generator.generator import AIBOMGenerator
+        except ImportError:
+            try:
+                from aibom_generator.generator import AIBOMGenerator
+            except ImportError:
+                try:
+                    from generator import AIBOMGenerator
+                except ImportError:
+                    raise ImportError("Could not import AIBOMGenerator. Please check your installation.")
+        
+        # Create generator instance
+        generator = AIBOMGenerator()
+        
+        # Create job directory
+        job_dir = os.path.join(OUTPUT_DIR, job_id)
+        
+        # Read job status
+        status_file = os.path.join(job_dir, "status.json")
+        with open(status_file, "r") as f:
+            status = json.load(f)
+        
+        # Update status to running
+        status["status"] = "running"
+        with open(status_file, "w") as f:
+            json.dump(status, f, indent=2)
+        
+        # Process each model
+        for i, model_id in enumerate(model_ids):
+            try:
+                # Generate AIBOM
+                aibom, enhancement_report = generator.generate(
+                    model_id=model_id,
+                    include_inference=include_inference,
+                    use_best_practices=use_best_practices
+                )
+                
+                # Process enhancement report for API response
+                enhancement_report = ensure_enhancement_report_structure(enhancement_report)
                 
-                if "original_score" in enhancement_report:
-                    batch_enhancement_report["original_score"]["total_score"] = enhancement_report["original_score"].get("total_score", 0)
-                    batch_enhancement_report["original_score"]["completeness_score"] = enhancement_report["original_score"].get("completeness_score", 
-                                                                                    enhancement_report["original_score"].get("total_score", 0))
+                # Save result to file
+                result_file = f"{model_id.replace('/', '_')}_aibom.json"
+                result_path = os.path.join(job_dir, result_file)
+                with open(result_path, "w") as f:
+                    json.dump({
+                        "aibom": aibom,
+                        "enhancement_report": enhancement_report
+                    }, f, indent=2)
                 
-                if "final_score" in enhancement_report:
-                    batch_enhancement_report["final_score"]["total_score"] = enhancement_report["final_score"].get("total_score", 0)
-                    batch_enhancement_report["final_score"]["completeness_score"] = enhancement_report["final_score"].get("completeness_score", 
-                                                                                 enhancement_report["final_score"].get("total_score", 0))
+                # Update status
+                status["completed"] += 1
+                status["results"][model_id] = result_file
+                with open(status_file, "w") as f:
+                    json.dump(status, f, indent=2)
+            except Exception as e:
+                # Log error
+                logger.error(f"Error processing model {model_id}: {str(e)}")
                 
-                batch_enhancement_report["improvement"] = enhancement_report.get("improvement", 0)
+                # Update status
+                status["completed"] += 1
+                status["results"][model_id] = f"Error: {str(e)}"
+                with open(status_file, "w") as f:
+                    json.dump(status, f, indent=2)
+        
+        # Update status to completed
+        status["status"] = "completed"
+        with open(status_file, "w") as f:
+            json.dump(status, f, indent=2)
+    except Exception as e:
+        # Log error
+        logger.error(f"Error processing batch job {job_id}: {str(e)}")
+        
+        try:
+            # Update status to failed
+            status_file = os.path.join(OUTPUT_DIR, job_id, "status.json")
+            with open(status_file, "r") as f:
+                status = json.load(f)
             
-            # Create download URL
-            download_url = f"/output/{job_id}/{safe_model_id}.json"
+            status["status"] = "failed"
+            status["error"] = str(e)
             
-            # Store result
-            batch_jobs[job_id]["results"][model_id] = {
-                "status": "completed",
-                "download_url": download_url,
-                "enhancement_report": batch_enhancement_report
-            }
-        except Exception as e:
-            # Store error
-            batch_jobs[job_id]["results"][model_id] = {
-                "status": "error",
-                "error": str(e)
-            }
-        
-        # Update progress
-        batch_jobs[job_id]["completed"] += 1
-    
-    # Update job status
-    batch_jobs[job_id]["status"] = "completed"
-
-if __name__ == "__main__":
-    uvicorn.run(app, host="0.0.0.0", port=8000)
+            with open(status_file, "w") as f:
+                json.dump(status, f, indent=2)
+        except Exception as status_e:
+            logger.error(f"Error updating job status: {str(status_e)}")