changed the nlp to support list and altered the fast API

app.py

@@ -1,7 +1,7 @@
 from fastapi import FastAPI, HTTPException, BackgroundTasks
 from fastapi.middleware.cors import CORSMiddleware
 from pydantic import BaseModel
-from typing import Dict, Optional, List
+from typing import Dict, Optional, List, Any
 import uuid
 from datetime import datetime
 from contextlib import asynccontextmanager
@@ -12,7 +12,6 @@ from models.nlp import NLPModel
 from database.query import DatabaseService
 from database.query_processor import QueryProcessor
 
-
 # Initialize models
 embedding_model = None
 summarization_model = None
@@ -29,7 +28,8 @@ async def lifespan(app: FastAPI):
     db_service = DatabaseService()
     yield
     # Clean up when app stops
-    await db_service.close()
+    if db_service:
+        await db_service.close()
 
 app = FastAPI(
     title="Kairos News API",
@@ -53,29 +53,20 @@ class PostRequest(BaseModel):
     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]
+    completed_at: Optional[datetime] = None
     request: PostRequest
-    result: Optional[Dict]
+    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,
@@ -93,14 +84,7 @@ async def create_job(request: PostRequest, background_tasks: BackgroundTasks):
         db_service
     )
     
-    return {
-        "id": job_id,
-        "status": "processing",
-        "created_at": jobs_db[job_id]["created_at"],
-        "completed_at": None,
-        "request": request,
-        "result": None
-    }
+    return jobs_db[job_id]  # Return the full job object
 
 @app.get("/loading", response_model=JobStatus)
 async def get_job_status(id: str):
@@ -135,7 +119,7 @@ async def process_job(
         jobs_db[job_id].update({
             "status": "completed",
             "completed_at": datetime.now(),
-            "result": result
+            "result": result if result else {"message": "No results found"}
         })
     except Exception as e:
         jobs_db[job_id].update({

models/nlp.py

@@ -11,8 +11,16 @@ class NLPModel:
         return self.extract_entities(text)  # or another default method
     
     def extract_entities(self, text: str):
-        doc = self.nlp(text)
-        return [(ent.text.lower(), ent.label_) for ent in doc.ents]
+        if isinstance(text, list):  # If input is a list of sentences
+            entities = []
+            for sentence in text:
+                doc = self.nlp(sentence)
+                entities.extend([(ent.text.lower(), ent.label_) for ent in doc.ents])
+            return entities
+        else:  # If input is a single string
+            doc = self.nlp(text)
+            return [(ent.text.lower(), ent.label_) for ent in doc.ents]
+    
     
     def tokenize_sentences(self, text: str):
         return nltk.sent_tokenize(text)