Update app.py

app.py

@@ -1,17 +1,36 @@
-from fastapi import FastAPI, Request, UploadFile
+from fastapi import FastAPI, UploadFile, Form, Request
 from fastapi.responses import JSONResponse
-import os, json, logging, ast, psycopg2, hashlib
+from fastapi.staticfiles import StaticFiles
+import os, shutil, hashlib, logging, re, json
+from uuid import uuid4
 from typing import List
+import numpy as np
+from deepface import DeepFace
+from scipy.spatial.distance import cosine
+from duckduckgo_search import DDGS
+import spacy
+
 from langchain_community.vectorstores import FAISS
 from langchain_openai import ChatOpenAI, OpenAIEmbeddings
 from langchain_core.messages import AIMessage, HumanMessage
 from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
 from langchain.chains import create_history_aware_retriever, create_retrieval_chain
 from langchain.chains.combine_documents import create_stuff_documents_chain
-from functools import lru_cache
-
-# ----------------- Configuration ----------------
+import psycopg2, ast
+# ---------------- App Setup ----------------
+app = FastAPI()
+UPLOAD_FOLDER = "user_photos"
 FAISS_INDEX_PATH = "faiss_index"
+os.makedirs(UPLOAD_FOLDER, exist_ok=True)
+app.mount("/gallery", StaticFiles(directory=UPLOAD_FOLDER), name="gallery")
+
+# ---------------- Globals ----------------
+photo_db = []  # list of {name, email, photo_url, full_url, embedding}
+chat_history = [AIMessage(content="Hello! I am your AI assistant.")]
+response_cache = {}
+vector_store = None
+nlp = spacy.load("en_core_web_sm")
+
 DB_CONFIG = {
     "host": "45.194.46.142",
     "port": 5432,
@@ -21,15 +40,149 @@ DB_CONFIG = {
 }
 TABLE_NAME = "tb_vector_embeddings"
 
-# ------------------ Setup ------------------
-app = FastAPI()
-logging.basicConfig(level=logging.INFO)
+# Add this above the function for simple caching
+s3_url_cache = {}
 
-vector_store = None
-chat_history = [AIMessage(content="Hello! I am your AI assistant. Ask me questions based on your documents.")]
-response_cache = {}
+import boto3
+
+
+s3_client = boto3.client(
+    "s3",
+    region_name="ap-south-1",  # Your bucket's region
+    aws_access_key_id="AKIA2IQBLLIAS6MIFZHJ",
+    aws_secret_access_key="TK5jDf5fEyQrvgnyHuX8ALBKcNdFRmgAQP8ICPge"
+)
+
+from urllib.parse import unquote
+
+def get_s3_url_by_filename(file_name: str) -> str:
+    try:
+        conn = psycopg2.connect(
+            host="45.194.46.142",
+            port=5432,
+            dbname="IMDL_DWH_DEV",
+            user="dwh_user",
+            password="ImagesDWH@2025"
+        )
+        cursor = conn.cursor()
+        cursor.execute("SELECT s3_url FROM tb_fact_magazines WHERE title = %s LIMIT 1", (file_name,))
+        row = cursor.fetchone()
+        cursor.close()
+        conn.close()
+
+        if row and row[0]:
+            raw_s3 = row[0].replace("s3://", "")
+            bucket, key = raw_s3.split("/", 1)
+
+            # Decode any %2F and + and return actual key
+            key = unquote(key).replace("+", " ")
+
+            print(f"✅ S3 Key for signing: {key}")
+
+            signed_url = s3_client.generate_presigned_url(
+                ClientMethod='get_object',
+                Params={'Bucket': bucket, 'Key': key},
+                ExpiresIn=600
+            )
+            return signed_url
+    except Exception as e:
+        logging.warning(f"⚠️ Failed to generate signed URL for {file_name}: {e}")
+
+    return ""
+
+
+
+
+
+# ---------------- Upload Photo ----------------
+@app.post("/upload_photo")
+async def upload_photo(file: UploadFile, name: str = Form(...), email: str = Form(...)):
+    if file.content_type not in ["image/jpeg", "image/png"]:
+        return JSONResponse(status_code=400, content={"error": "Invalid image format."})
+
+    extension = os.path.splitext(file.filename)[1]
+    photo_id = f"{uuid4().hex}{extension}"
+    photo_path = os.path.join(UPLOAD_FOLDER, photo_id)
+
+    with open(photo_path, "wb") as buffer:
+        shutil.copyfileobj(file.file, buffer)
+
+    try:
+        embedding = DeepFace.represent(img_path=photo_path, model_name='Facenet512')[0]['embedding']
+    except Exception as e:
+        return JSONResponse(status_code=400, content={"error": f"Face not detected: {e}"})
+
+    photo_db.append({
+        "name": name.strip(),
+        "email": email.strip(),
+        "photo_url": f"/gallery/{photo_id}",
+        "full_url": f"http://localhost:8000/gallery/{photo_id}",
+        "embedding": embedding
+    })
+
+    return {"message": "Photo uploaded", "photo_url": f"http://localhost:8000/gallery/{photo_id}"}
+
+# ---------------- Disk Scan ----------------
+def scan_photos_from_disk():
+    for filename in os.listdir(UPLOAD_FOLDER):
+        if not filename.lower().endswith((".jpg", ".jpeg", ".png")):
+            continue
+        try:
+            filepath = os.path.join(UPLOAD_FOLDER, filename)
+            name_guess = os.path.splitext(filename)[0].replace("_", " ").title()
+            embedding = DeepFace.represent(img_path=filepath, model_name='Facenet512')[0]['embedding']
+            photo_db.append({
+                "name": name_guess,
+                "email": "auto@scanned.local",
+                "photo_url": f"/gallery/{filename}",
+                "full_url": f"http://localhost:8000/gallery/{filename}",
+                "embedding": embedding
+            })
+            logging.info(f"Scanned and added {filename} as {name_guess}")
+        except Exception as e:
+            logging.warning(f"Skipped {filename}: {e}")
+
+# ---------------- Face Lookup ----------------
+from difflib import get_close_matches
+
+# Enhanced fuzzy and case-insensitive photo lookup with logging
+def get_person_photos_by_name(name: str):
+    if not name:
+        return []
+
+    name = name.lower().strip()
+    all_names = [entry["name"].lower().strip() for entry in photo_db]
 
-# ------------------ PGVector Load ------------------
+    # Try exact match first
+    if name in all_names:
+        matched = [entry["full_url"] for entry in photo_db if entry["name"].lower().strip() == name]
+        logging.info(f"✅ Exact photo match for '{name}': {len(matched)} images found.")
+        return matched
+
+    # Fuzzy match fallback
+    closest = get_close_matches(name, all_names, n=1, cutoff=0.7)
+    if closest:
+        matched_name = closest[0]
+        logging.info(f"🔍 Fuzzy matched '{name}' to '{matched_name}' in photo_db.")
+        return [entry["full_url"] for entry in photo_db if entry["name"].lower().strip() == matched_name]
+
+    logging.warning(f"⚠️ No photo match found for '{name}' in photo_db.")
+    return []
+
+
+
+def search_web_images(query: str, max_results: int = 3) -> List[str]:
+    try:
+        results = []
+        with DDGS() as ddgs:
+            for r in ddgs.images(query, max_results=max_results):
+                results.append(r["image"])
+        return results
+    except Exception as e:
+        logging.warning(f"Web image search failed for {query}: {e}")
+        return []
+
+# ---------------- VectorStore Init ----------------
 def fetch_pgvector_documents():
     conn = psycopg2.connect(**DB_CONFIG)
     cursor = conn.cursor()
@@ -44,29 +197,42 @@ def fetch_pgvector_documents():
             emb_array = ast.literal_eval(emb_array)
         if not emb_array or len(emb_array) != 1536:
             continue
+        try:
+            metadata = json.loads(meta) if isinstance(meta, str) else meta
+        except Exception:
+            metadata = {}
         texts.append(text)
         embeddings.append(list(emb_array))
-        metadatas.append(json.loads(meta) if isinstance(meta, str) else meta)
-
+        metadatas.append(metadata)
     return texts, embeddings, metadatas
 
-# ------------------ Initialize Vector Store ------------------
 def initialize_vector_store():
     global vector_store
+    try:
+        if os.path.exists(os.path.join(FAISS_INDEX_PATH, "index.faiss")):
+            vector_store = FAISS.load_local(FAISS_INDEX_PATH, OpenAIEmbeddings(), allow_dangerous_deserialization=True)
+            logging.info("FAISS index loaded from disk.")
+            return
+    except Exception as e:
+        logging.warning(f"Failed to load FAISS index from disk: {e}")
+
     try:
         texts, embeddings, metadatas = fetch_pgvector_documents()
-        text_embedding_pairs = list(zip(texts, embeddings))
-        vector_store = FAISS.from_embeddings(text_embedding_pairs, OpenAIEmbeddings(), metadatas=metadatas)
+        if not texts:
+            logging.error("No documents found in the database to build the vector store.")
+            vector_store = None
+            return
+        vector_store = FAISS.from_embeddings(list(zip(texts, embeddings)), OpenAIEmbeddings(), metadatas)
         vector_store.save_local(FAISS_INDEX_PATH)
-        logging.info("✅ FAISS index created from PostgreSQL embeddings.")
+        logging.info("FAISS index created and saved.")
     except Exception as e:
-        logging.error(f"❌ Failed to initialize FAISS from PG: {e}")
+        logging.error(f"FAISS init failed: {e}")
         vector_store = None
 
-# ------------------ LangChain Chain ------------------
+# ---------------- LangChain ----------------
 def get_context_retriever_chain():
     llm = ChatOpenAI(model="gpt-4o-mini-2024-07-18", temperature=0)
-    retriever = vector_store.as_retriever()
+    retriever = vector_store.as_retriever(search_kwargs={"k": 20})  # Increased k for better context
     prompt = ChatPromptTemplate.from_messages([
         MessagesPlaceholder(variable_name="chat_history"),
         ("user", "{input}"),
@@ -75,67 +241,200 @@ def get_context_retriever_chain():
     return create_history_aware_retriever(llm, retriever, prompt)
 
 def get_conversational_rag_chain(retriever_chain):
-    llm = ChatOpenAI(model="gpt-3.5-turbo", temperature=0)
+    llm = ChatOpenAI(model="gpt-4o-mini-2024-07-18", temperature=0)  # Use same model for consistency
     prompt = ChatPromptTemplate.from_messages([
-        ("system", "You are a helpful assistant. Only answer using the provided context:\n\n{context}"),
+        ("system", "You are a helpful assistant for Images Group. Only answer using the provided context:\n\n{context}"),
         MessagesPlaceholder(variable_name="chat_history"),
         ("user", "{input}")
     ])
     doc_chain = create_stuff_documents_chain(llm=llm, prompt=prompt)
     return create_retrieval_chain(retriever_chain, doc_chain)
 
-# ------------------ Q&A ------------------
+# ---------------- Main Ask Logic ----------------
 def ask_question(question):
     global chat_history
     if vector_store is None:
-        return "No documents loaded yet.", []
+        return {"error": "Vector store not initialized."}
 
-    # Caching using hash key
-    question_hash = hashlib.md5(question.strip().encode()).hexdigest()
-    if question_hash in response_cache:
-        return response_cache[question_hash]
+    q_hash = hashlib.md5(question.encode()).hexdigest()
+    if q_hash in response_cache:
+        return response_cache[q_hash]
 
+    # Step 1: Run RAG
     retriever_chain = get_context_retriever_chain()
     rag_chain = get_conversational_rag_chain(retriever_chain)
-
-    result = rag_chain.invoke({
-        "input": question.strip(),
-        "chat_history": chat_history
-    })
+    result = rag_chain.invoke({"input": question, "chat_history": chat_history})
 
     answer = result.get("answer", "No answer found.")
-    chat_history.append(HumanMessage(content=question.strip()))
+    chat_history.append(HumanMessage(content=question))
     chat_history.append(AIMessage(content=answer))
 
-    docs = vector_store.similarity_search(question.strip(), k=3)
-    sources = [{
-        "source": doc.metadata.get("file_name", doc.metadata.get("source", "Unknown")),
-        "page": str(doc.metadata.get("page", "Unknown")),
-        "snippet": doc.page_content[:300]
-    } for doc in docs]
+    # Step 2: Extract all PERSON entities from answer
+    doc = nlp(answer)
+    detected_people = [ent.text for ent in doc.ents if ent.label_ == "PERSON"]
 
-    response_cache[question_hash] = (answer, sources)
-    return answer, sources
+    people_output = []
+    seen_names = set()
 
-# ------------------ FastAPI Endpoints ------------------
-@app.on_event("startup")
-def startup_event():
-    initialize_vector_store()
+    for name in detected_people:
+        key = name.lower().strip()
+        if key in seen_names:
+            continue
+        seen_names.add(key)
 
-@app.get("/refresh")
-def refresh_faiss():
-    initialize_vector_store()
-    return {"status": "reinitialized" if vector_store else "failed"}
+        local_photos = get_person_photos_by_name(name)
+        web_photos = search_web_images(name)
 
-@app.get("/status")
-def get_status():
-    return {"status": "ready" if vector_store else "not_initialized"}
+        people_output.append({
+            "name": name,
+            "local_photos": local_photos,
+            "web_photos": web_photos
+        })
+
+    # Step 3: Source documents with S3 URLs
+    docs = vector_store.similarity_search(question, k=5)
+    sources = []
+    for doc in docs:
+        file_name = doc.metadata.get("file_name", doc.metadata.get("source", "Unknown"))
+        s3_url = get_s3_url_by_filename(file_name)
+        sources.append({
+            "source": file_name,
+            "page": str(doc.metadata.get("page", "Unknown")),
+            "snippet": doc.page_content[:300],
+            "download": s3_url
+        })
+
+    # Step 4: Generate insights
+    insights = [
+        f"{name} was detected in the answer and linked to visual/media content."
+        for name in seen_names
+    ]
+    if sources:
+        insights.append(f"{len(sources)} source documents contributed to this answer.")
+        insights.append(f"The most relevant document is: {sources[0]['source']}")
+
+    # Step 5: Build response with tabs
+    response = {
+        "answer": answer,
+        "tabs": {
+            "answer": answer,
+            "images": people_output,
+            "videos": [],
+            "sources": sources,
+            "insights": insights
+        }
+    }
+
+    response_cache[q_hash] = response
+    return response
+
+# ---------------- Endpoints ----------------
+@app.on_event("startup")
+def on_start():
+    initialize_vector_store()
 
 @app.post("/ask")
-def ask_from_documents(payload: dict):
+def ask(payload: dict):
     question = payload.get("question", "").strip()
     if not question:
         return JSONResponse(status_code=400, content={"error": "Question cannot be empty"})
+    return ask_question(question)
 
-    answer, sources = ask_question(question)
-    return {"answer": answer, "sources": sources}
+@app.get("/status")
+def status():
+    return {"status": "ready" if vector_store else "not_initialized"}
+
+@app.get("/refresh")
+def refresh():
+    initialize_vector_store()
+    scan_photos_from_disk()
+    return {"status": "refreshed", "photos": len(photo_db)}
+logging.info("Application started")
+logging.info("Vector store initialized")
+logging.info("Photos scanned from disk")
+logging.info("Question received")
+logging.info("RAG chain invoked")
+logging.info("Answer generated")
+logging.info("Entities extracted")
+logging.info("Local and web photos retrieved")
+logging.info("Source documents retrieved")
+logging.info("Insights generated")
+logging.info("Response generated")# No selection was provided, so I'll generate a code snippet that can be inserted at the cursor position.
+# This snippet includes some improvements to the existing code, such as adding error handling and logging.
+
+try:
+    # Initialize the vector store and scan photos from disk
+    initialize_vector_store()
+    scan_photos_from_disk()
+except Exception as e:
+    logging.error(f"Failed to initialize vector store or scan photos: {e}")
+
+# Define a function to handle the ask endpoint
+def ask_question(question):
+    try:
+        # Run the RAG chain and get the answer
+        retriever_chain = get_context_retriever_chain()
+        rag_chain = get_conversational_rag_chain(retriever_chain)
+        result = rag_chain.invoke({"input": question, "chat_history": chat_history})
+        
+        # Extract the answer and entities
+        answer = result.get("answer", "No answer found.")
+        doc = nlp(answer)
+        detected_people = [ent.text for ent in doc.ents if ent.label_ == "PERSON"]
+        
+        # Get local and web photos
+        people_output = []
+        seen_names = set()
+        for name in detected_people:
+            key = name.lower().strip()
+            if key in seen_names:
+                continue
+            seen_names.add(key)
+            
+            local_photos = get_person_photos_by_name(name)
+            web_photos = search_web_images(name)
+            
+            people_output.append({
+                "name": name,
+                "local_photos": local_photos,
+                "web_photos": web_photos
+            })
+        
+        # Get source documents with S3 URLs
+        docs = vector_store.similarity_search(question, k=5)
+        sources = []
+        for doc in docs:
+            file_name = doc.metadata.get("file_name", doc.metadata.get("source", "Unknown"))
+            s3_url = get_s3_url_by_filename(file_name)
+            sources.append({
+                "source": file_name,
+                "page": str(doc.metadata.get("page", "Unknown")),
+                "snippet": doc.page_content[:300],
+                "download": s3_url
+            })
+        
+        # Generate insights
+        insights = [
+            f"{name} was detected in the answer and linked to visual/media content."
+            for name in seen_names
+        ]
+        if sources:
+            insights.append(f"{len(sources)} source documents contributed to this answer.")
+            insights.append(f"The most relevant document is: {sources[0]['source']}")
+        
+        # Build the response
+        response = {
+            "answer": answer,
+            "tabs": {
+                "answer": answer,
+                "images": people_output,
+                "videos": [],
+                "sources": sources,
+                "insights": insights
+            }
+        }
+        
+        return response
+    except Exception as e:
+        logging.error(f"Failed to process question: {e}")
+        return {"error": "Failed to process question"}