Update app.py

app.py

@@ -17,6 +17,25 @@ import boto3
 from io import BytesIO
 import base64
 from langchain.chains import RetrievalQA
+
+
+class AskRequest(BaseModel):
+    question: str
+
+class ImageRequest(BaseModel):
+    question: str
+    answer: str
+
+class VideoRequest(BaseModel):
+    question: str
+
+class SourceRequest(BaseModel):
+    question: str
+
+class InsightRequest(BaseModel):
+    question: str
+    answer: str
+
 # ------------------- FastAPI Init -------------------
 app = FastAPI()
 
@@ -350,132 +369,89 @@ async def ask(payload: AskRequest):
 
 
 @app.post("/get_images")
-async def get_images(question: str = Form(...), answer: str = Form(...)):
-    """
-    Detects people (PERSON) and brands/companies (ORG) from both the question and answer.
-    Returns both local DB-matched base64 images and web photos from DuckDuckGo.
-
-    Output:
-    {
-        "images": [
-            {
-                "name": "Gopal Asthana",
-                "local_photos": ["data:image/jpeg;base64,..."],
-                "web_photos": ["https://..."]
-            },
-            ...
-        ]
-    }
-    """
-    try:
-        # 🔍 NER over question and answer separately
-        doc_q = nlp(question.strip())
-        doc_a = nlp(answer.strip())
-
-        entities = [
-            ent.text.strip()
-            for ent in list(doc_q.ents) + list(doc_a.ents)
-            if ent.label_ in {"PERSON", "ORG"}
-        ]
+async def get_images(payload: ImageRequest):
+    question = payload.question
+    answer = payload.answer
+
+    doc_q = nlp(question.strip())
+    doc_a = nlp(answer.strip())
+    entities = [
+        ent.text.strip()
+        for ent in list(doc_q.ents) + list(doc_a.ents)
+        if ent.label_ in {"PERSON", "ORG"}
+    ]
 
-        # Fallback: Use n-gram tokens if spaCy fails
-        if not entities and len(answer.strip().split()) >= 2:
-            words = answer.strip().split()
-            entities = [' '.join(words[i:i+2]) for i in range(len(words)-1)]
-
-        seen = set()
-        results = []
-
-        for name in entities:
-            norm_key = name.lower().replace(" ", "")
-            if norm_key in seen:
-                continue
-            seen.add(norm_key)
-
-            # ----- Local image base64 from DB -----
-            local_photos = []
-            try:
-                conn = psycopg2.connect(
-                    host="45.194.46.142",
-                    port=5432,
-                    dbname="IMDL_DWH_DEV",
-                    user="imdl_dwh_da_user",
-                    password="ImagesDWH@2025"
-                )
-                cursor = conn.cursor()
-                cursor.execute("""
-                    SELECT image_data FROM tb_fact_image_uploads
-                    WHERE LOWER(title) ILIKE %s
-                      AND is_active = true AND is_deleted = false
-                    LIMIT 1
-                """, (f"%{name.lower()}%",))
-                row = cursor.fetchone()
-                cursor.close()
-                conn.close()
-
-                if row and row[0]:
-                    image_bytes = BytesIO(row[0]).getvalue()
-                    image_base64 = base64.b64encode(image_bytes).decode("utf-8")
-                    local_photos = [f"data:image/jpeg;base64,{image_base64}"]
-
-            except Exception as e:
-                logging.warning(f"⚠️ Local photo DB error for '{name}': {e}")
-
-            # ----- Web fallback -----
-            web_photos = []
-            try:
-                with DDGS() as ddgs:
-                    for r in ddgs.images(name, max_results=3):
-                        web_photos.append(r["image"])
-            except Exception as e:
-                logging.warning(f"⚠️ DuckDuckGo search failed for '{name}': {e}")
-
-            results.append({
-                "name": name,
-                "local_photos": local_photos,
-                "web_photos": web_photos
-            })
-
-        return {"images": results}
+    if not entities and len(answer.strip().split()) >= 2:
+        words = answer.strip().split()
+        entities = [' '.join(words[i:i+2]) for i in range(len(words)-1)]
 
-    except Exception as e:
-        return JSONResponse(status_code=500, content={"error": f"Image detection failed: {e}"})
+    seen, results = set(), []
+
+    for name in entities:
+        norm_key = name.lower().replace(" ", "")
+        if norm_key in seen:
+            continue
+        seen.add(norm_key)
+
+        # Local DB base64 image
+        local_photos = []
+        try:
+            conn = psycopg2.connect(**DB_CONFIG)
+            cursor = conn.cursor()
+            cursor.execute("""
+                SELECT image_data FROM tb_fact_image_uploads
+                WHERE LOWER(title) ILIKE %s
+                AND is_active = true AND is_deleted = false
+                LIMIT 1
+            """, (f"%{name.lower()}%",))
+            row = cursor.fetchone()
+            cursor.close()
+            conn.close()
+            if row and row[0]:
+                image_bytes = BytesIO(row[0]).getvalue()
+                image_base64 = base64.b64encode(image_bytes).decode("utf-8")
+                local_photos = [f"data:image/jpeg;base64,{image_base64}"]
+        except Exception as e:
+            logging.warning(f"Local photo DB error for '{name}': {e}")
 
+        # Web fallback
+        web_photos = []
+        try:
+            with DDGS() as ddgs:
+                for r in ddgs.images(name, max_results=3):
+                    web_photos.append(r["image"])
+        except Exception as e:
+            logging.warning(f"DuckDuckGo failed for '{name}': {e}")
+
+        results.append({
+            "name": name,
+            "local_photos": local_photos,
+            "web_photos": web_photos
+        })
+
+    return {"images": results}
 
 
 @app.post("/get_videos")
-async def get_videos(question: str = Form(...)):
+async def get_videos(payload: VideoRequest):
     try:
-        return {"videos": await asyncio.to_thread(fetch_youtube_videos.invoke, {"query": query})}
+        return {
+            "videos": await asyncio.to_thread(fetch_youtube_videos.invoke, {"question": payload.question})
+        }
     except Exception as e:
         return JSONResponse(status_code=500, content={"error": f"Video fetch failed: {e}"})
 
 
 @app.post("/get_sources")
-async def get_sources(question: str = Form(...)):
-    """
-    Retrieve source metadata (filename, page, snippet, signed S3 URL) 
-    from the vector store using FAISS similarity search.
+async def get_sources(payload: SourceRequest):
+    question = payload.question
 
-    Returns:
-    {
-        "sources": [
-            {
-                "source": "RetailSummit2024.pdf",
-                "page": "3",
-                "snippet": "Gopal Asthana presented innovations in fashion...",
-                "signed_url": "https://..."
-            },
-            ...
-        ]
-    }
-    """
-    try:
-        def clean(text):
-            text = re.sub(r'\[([^\]]+)\]\((https?://[^\)]+)\)', r'\1', text)
-            text = re.sub(r'https?://\S+', '', text)
-            return text.strip()
+    def clean(text):
+        text = re.sub(r'\[([^\]]+)\]\((https?://[^\)]+)\)', r'\1', text)
+        text = re.sub(r'https?://\S+', '', text)
+        return text.strip()
 
+    try:
         if not vector_store:
             return JSONResponse(status_code=500, content={"error": "FAISS vector store not initialized."})
 
@@ -483,7 +459,8 @@ async def get_sources(question: str = Form(...)):
         if not docs:
             return {"sources": []}
 
-        seen = {}
+        seen, sources = {}, []
+
         for doc in docs:
             source = doc.metadata.get("file_name", doc.metadata.get("source", "Unknown"))
             page = str(doc.metadata.get("page", "1"))
@@ -500,21 +477,22 @@ async def get_sources(question: str = Form(...)):
         return {"sources": list(seen.values())}
 
     except Exception as e:
-        logging.error(f"❌ Error in get_sources: {e}")
+        logging.error(f"Error in get_sources: {e}")
         return JSONResponse(status_code=500, content={"error": f"Failed to fetch sources: {e}"})
 
 
 
 @app.post("/get_insights")
-async def get_insights(answer: str = Form(...), question: str = Form(...)):
+async def get_insights(payload: InsightRequest):
     try:
-        people = await asyncio.to_thread(detect_people_and_images.invoke, {"answer": answer})
-        docs = await asyncio.to_thread(retrieve_documents.invoke, {"query": question})
+        people = await asyncio.to_thread(detect_people_and_images.invoke, {"answer": payload.answer})
+        docs = await asyncio.to_thread(retrieve_documents.invoke, {"query": payload.question})
         insights = []
         if people:
             insights.extend([f"{p['name']} was detected." for p in people])
-        if docs:
-            insights.append(f"{len(docs)} documents used. Top: {docs[0]['source']}")
+        if isinstance(docs, str) and docs.strip():
+            insights.append("Document-based answer was successfully generated.")
         return {"insights": insights}
     except Exception as e:
-        return JSONResponse(status_code=500, content={"error": f"Insight gen failed: {e}"})
+        return JSONResponse(status_code=500, content={"error": f"Insight generation failed: {e}"})
+