File size: 3,609 Bytes
4a24dbd 47ce483 e00c07d 47ce483 a2682b3 dc9275a a2682b3 77370a4 a2682b3 e00c07d a2682b3 47ce483 a2682b3 47ce483 a2682b3 4a24dbd e00c07d 4a24dbd e00c07d 4a24dbd 47ce483 e00c07d 47ce483 4a24dbd 47ce483 a2682b3 47ce483 e00c07d 47ce483 4a24dbd 47ce483 a2682b3 4a24dbd a2682b3 e00c07d a2682b3 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 |
from fastapi import FastAPI, HTTPException, BackgroundTasks
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel
from typing import Dict, Optional, List, Any
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
if db_service:
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 JobStatus(BaseModel):
id: str
status: str # "processing", "completed", "failed"
created_at: datetime
completed_at: Optional[datetime] = None
request: PostRequest
result: Optional[Dict[str, Any]] = None # Flexible result structure
@app.post("/index", response_model=JobStatus)
async def create_job(request: PostRequest, background_tasks: BackgroundTasks):
job_id = str(uuid.uuid4())
jobs_db[job_id] = {
"id": job_id, # Ensure `id` is included
"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 jobs_db[job_id] # Return the full job object
@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 if result else {"message": "No results found"}
})
except Exception as e:
jobs_db[job_id].update({
"status": "failed",
"completed_at": datetime.now(),
"result": {"error": str(e)}
}) |