MuhammadSheraza002/hardware-scanner-api

1import os2import json3import base644import io5from datetime import datetime6from typing import List, Optional7from fastapi import FastAPI, HTTPException8from fastapi.middleware.cors import CORSMiddleware9from fastapi.responses import StreamingResponse10from pydantic import BaseModel11from groq import Groq12from openpyxl import Workbook13from dotenv import load_dotenv14 15load_dotenv()16 17app = FastAPI(title="Hardware Inventory Scanner API")18 19# CORS - Allow all origins for development, restrict in production20app.add_middleware(21    CORSMiddleware,22    allow_origins=["*"],23    allow_credentials=True,24    allow_methods=["*"],25    allow_headers=["*"],26)27 28# In-memory session storage (for demo; use Redis/DB in production)29sessions = {}30 31# --- Models ---32class ImageRequest(BaseModel):33    session_id: str34    image_base64: str  # Base64 encoded image data (without data:image prefix)35 36class HardwareData(BaseModel):37    capacity: str = "N/A"38    generation: str = "N/A"39    brand: str = "N/A"40    speed: str = "N/A"41    timestamp: Optional[str] = None42 43class ProcessResponse(BaseModel):44    success: bool45    data: Optional[HardwareData] = None46    error: Optional[str] = None47    scan_count: int = 048 49class SessionStats(BaseModel):50    session_id: str51    scan_count: int52    started_at: str53 54# --- Groq Vision Processing ---55def extract_hardware_info(image_base64: str) -> dict:56    """Process image with Groq Vision API to extract hardware info"""57    api_key = os.environ.get("GROQ_API_KEY")58    if not api_key:59        raise HTTPException(status_code=500, detail="GROQ_API_KEY not configured")60    61    client = Groq(api_key=api_key)62    63    prompt = """64    Extract the following information from this hardware label in JSON format:65    - capacity (e.g., 8GB, 16GB, 256GB)66    - generation (DDR3, DDR4, DDR5)67    - brand68    - speed (bus speed in MHz, e.g., 2133, 2400, 2666, 3200)69    If any field is missing, set it to "N/A". Return ONLY the JSON object.70    """71    72    try:73        completion = client.chat.completions.create(74            model="meta-llama/llama-4-scout-17b-16e-instruct",75            messages=[76                {77                    "role": "user",78                    "content": [79                        {"type": "text", "text": prompt},80                        {81                            "type": "image_url",82                            "image_url": {"url": f"data:image/jpeg;base64,{image_base64}"}83                        },84                    ],85                }86            ],87            response_format={"type": "json_object"}88        )89        return json.loads(completion.choices[0].message.content)90    except Exception as e:91        raise HTTPException(status_code=500, detail=str(e))92 93# --- API Endpoints ---94@app.get("/")95def health_check():96    return {"status": "healthy", "service": "Hardware Inventory Scanner API"}97 98@app.post("/api/start-session")99def start_session():100    """Start a new scanning session"""101    session_id = datetime.now().strftime("%Y%m%d_%H%M%S")102    sessions[session_id] = {103        "items": [],104        "started_at": datetime.now().isoformat()105    }106    return {"session_id": session_id, "message": "Session started"}107 108@app.post("/api/process-image", response_model=ProcessResponse)109def process_image(request: ImageRequest):110    """Process a captured image and extract hardware information"""111    if request.session_id not in sessions:112        # Auto-create session if not exists113        sessions[request.session_id] = {114            "items": [],115            "started_at": datetime.now().isoformat()116        }117    118    try:119        # Extract hardware info from image120        result = extract_hardware_info(request.image_base64)121        122        # Add timestamp123        result["timestamp"] = datetime.now().strftime("%Y-%m-%d %H:%M:%S")124        125        # Store in session126        sessions[request.session_id]["items"].append(result)127        128        return ProcessResponse(129            success=True,130            data=HardwareData(**result),131            scan_count=len(sessions[request.session_id]["items"])132        )133    except HTTPException as e:134        return ProcessResponse(success=False, error=e.detail)135    except Exception as e:136        return ProcessResponse(success=False, error=str(e))137 138@app.get("/api/session/{session_id}")139def get_session(session_id: str):140    """Get session data"""141    if session_id not in sessions:142        raise HTTPException(status_code=404, detail="Session not found")143    144    session = sessions[session_id]145    return {146        "session_id": session_id,147        "items": session["items"],148        "scan_count": len(session["items"]),149        "started_at": session["started_at"]150    }151 152@app.get("/api/export/{session_id}")153def export_session(session_id: str):154    """Export session data as Excel file"""155    if session_id not in sessions:156        raise HTTPException(status_code=404, detail="Session not found")157    158    session = sessions[session_id]159    items = session["items"]160    161    if not items:162        raise HTTPException(status_code=400, detail="No items to export")163    164    # Create Excel workbook165    wb = Workbook()166    ws = wb.active167    ws.title = "Hardware Inventory"168    169    # Header row170    headers = ["#", "Brand", "Capacity", "Generation", "Speed (MHz)", "Scanned At"]171    for col, header in enumerate(headers, 1):172        ws.cell(row=1, column=col, value=header)173        ws.cell(row=1, column=col).font = ws.cell(row=1, column=col).font.copy(bold=True)174    175    # Data rows176    for idx, item in enumerate(items, 1):177        ws.cell(row=idx + 1, column=1, value=idx)178        ws.cell(row=idx + 1, column=2, value=item.get("brand", "N/A"))179        ws.cell(row=idx + 1, column=3, value=item.get("capacity", "N/A"))180        ws.cell(row=idx + 1, column=4, value=item.get("generation", "N/A"))181        ws.cell(row=idx + 1, column=5, value=item.get("speed", "N/A"))182        ws.cell(row=idx + 1, column=6, value=item.get("timestamp", "N/A"))183    184    # Auto-adjust column widths185    for col in ws.columns:186        max_length = 0187        column = col[0].column_letter188        for cell in col:189            try:190                if len(str(cell.value)) > max_length:191                    max_length = len(str(cell.value))192            except:193                pass194        adjusted_width = (max_length + 2)195        ws.column_dimensions[column].width = adjusted_width196    197    # Save to bytes buffer198    buffer = io.BytesIO()199    wb.save(buffer)200    buffer.seek(0)201    202    filename = f"hardware_inventory_{session_id}.xlsx"203    204    return StreamingResponse(205        buffer,206        media_type="application/vnd.openxmlformats-officedocument.spreadsheetml.sheet",207        headers={"Content-Disposition": f"attachment; filename={filename}"}208    )209 210@app.delete("/api/session/{session_id}")211def end_session(session_id: str):212    """End and clean up a session"""213    if session_id in sessions:214        del sessions[session_id]215    return {"message": "Session ended", "session_id": session_id}216 217if __name__ == "__main__":218    import uvicorn219    uvicorn.run(app, host="0.0.0.0", port=8000)220