from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
from transformers import pipeline
from typing import List

app = FastAPI()

class SentenceListPayload(BaseModel):
    sentences: List[str]

# Load the text classification pipeline on startup
try:
    action_item_classifier = pipeline(
        "text-classification",
        model="knkarthick/Action_Items",
        device="cpu",
    )
    print("✅ Action item model loaded successfully")
except Exception as e:
    action_item_classifier = None
    print(f"❌ Error loading action item model: {e}")

@app.post("/classify-action-items")
async def classify_sentences(payload: SentenceListPayload):
    if not action_item_classifier:
        raise HTTPException(status_code=503, detail="Action item model is not available.")
    
    results = action_item_classifier(payload.sentences)
    
    # Filter for sentences classified as action items with a confidence threshold
    action_items = []
    for i, sentence in enumerate(payload.sentences):
        if results[i]['label'] == 'LABEL_1' and results[i]['score'] > 0.8:
            action_items.append({
                "sentence": sentence,
                "confidence": results[i]['score']
            })
            
    return {"action_items": action_items}