app.py

from fastapi import FastAPI, HTTPException, BackgroundTasks
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel
from typing import Dict, Optional, List
import uuid
from datetime import datetime
from contextlib import asynccontextmanager

from models.embedding import EmbeddingModel
from models.summarization import SummarizationModel
from models.nlp import NLPModel
from database.query import DatabaseService
from database.query_processor import QueryProcessor


# Initialize models
embedding_model = None
summarization_model = None
nlp_model = None
db_service = None

@asynccontextmanager
async def lifespan(app: FastAPI):
    # Load models when app starts
    global embedding_model, summarization_model, nlp_model, db_service
    embedding_model = EmbeddingModel()
    summarization_model = SummarizationModel()
    nlp_model = NLPModel()
    db_service = DatabaseService()
    yield
    # Clean up when app stops
    await db_service.close()

app = FastAPI(
    title="Kairos News API",
    version="1.0",
    lifespan=lifespan
)

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_methods=["*"],
    allow_headers=["*"],
)

# In-memory job storage
jobs_db: Dict[str, Dict] = {}

class PostRequest(BaseModel):
    query: str
    topic: Optional[str] = None
    start_date: Optional[str] = None  # Format: "YYYY-MM-DD"
    end_date: Optional[str] = None    # Format: "YYYY-MM-DD"

class ArticleResult(BaseModel):
    url: str
    content: str
    distance: float
    date: str
    topic: str

class SummaryResult(BaseModel):
    summary: str

class JobStatus(BaseModel):
    id: str
    status: str  # "processing", "completed", "failed"
    created_at: datetime
    completed_at: Optional[datetime]
    request: PostRequest
    result: Optional[Dict]

@app.post("/index", response_model=JobStatus)
async def create_job(request: PostRequest, background_tasks: BackgroundTasks):
    job_id = str(uuid.uuid4())
    
    jobs_db[job_id] = {
        "status": "processing",
        "created_at": datetime.now(),
        "completed_at": None,
        "request": request.dict(),
        "result": None
    }

    background_tasks.add_task(
        process_job,
        job_id,
        request,
        embedding_model,
        summarization_model,
        nlp_model,
        db_service
    )
    
    return {
        "id": job_id,
        "status": "processing",
        "created_at": jobs_db[job_id]["created_at"],
        "completed_at": None,
        "request": request,
        "result": None
    }

@app.get("/loading", response_model=JobStatus)
async def get_job_status(id: str):
    if id not in jobs_db:
        raise HTTPException(status_code=404, detail="Job not found")
    
    return jobs_db[id]

async def process_job(
    job_id: str,
    request: PostRequest,
    embedding_model: EmbeddingModel,
    summarization_model: SummarizationModel,
    nlp_model: NLPModel,
    db_service: DatabaseService
):
    try:
        processor = QueryProcessor(
            embedding_model=embedding_model,
            summarization_model=summarization_model,
            nlp_model=nlp_model,
            db_service=db_service
        )
        
        result = await processor.process(
            query=request.query,
            topic=request.topic,
            start_date=request.start_date,
            end_date=request.end_date
        )
        
        jobs_db[job_id].update({
            "status": "completed",
            "completed_at": datetime.now(),
            "result": result
        })
    except Exception as e:
        jobs_db[job_id].update({
            "status": "failed",
            "completed_at": datetime.now(),
            "result": {"error": str(e)}
        })