Update app.py

app.py

@@ -1,8 +1,10 @@
-from fastapi import FastAPI, UploadFile, File, HTTPException
 import os
+import json
 import faiss
 import pickle
 import numpy as np
+from fastapi import FastAPI, HTTPException
+from pydantic import BaseModel
 from sentence_transformers import SentenceTransformer
 from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, pipeline
 
@@ -10,13 +12,13 @@ app = FastAPI(title="MetaGPT AI - Local Q&A")
 
 # Config
 DATA_DIR = "data"
-DOCS_DIR = os.path.join(DATA_DIR, "docs")
 INDEX_FILE = os.path.join(DATA_DIR, "index.faiss")
 METADATA_FILE = os.path.join(DATA_DIR, "metadata.pkl")
 CHUNK_SIZE = 500
 CHUNK_OVERLAP = 100
+JSON_FILE = "articles.json"
 
-os.makedirs(DOCS_DIR, exist_ok=True)
+os.makedirs(DATA_DIR, exist_ok=True)
 
 # Load models
 embed_model = SentenceTransformer("sentence-transformers/all-MiniLM-L6-v2")
@@ -25,7 +27,7 @@ tokenizer = AutoTokenizer.from_pretrained(gen_model_name)
 gen_model = AutoModelForSeq2SeqLM.from_pretrained(gen_model_name)
 gen_pipeline = pipeline("text2text-generation", model=gen_model, tokenizer=tokenizer, device=-1)
 
-# Helper: chunk text
+# Helpers
 def chunk_text(text):
     chunks = []
     start = 0
@@ -37,7 +39,6 @@ def chunk_text(text):
         if start >= len(text): break
     return chunks
 
-# Helper: save/load FAISS index
 def save_index(index, metadata):
     faiss.write_index(index, INDEX_FILE)
     with open(METADATA_FILE, "wb") as f:
@@ -51,58 +52,64 @@ def load_index():
         metadata = pickle.load(f)
     return index, metadata
 
-# Endpoint: upload document
-@app.post("/upload")
-async def upload(files: list[UploadFile] = File(...)):
-    index, metadata = load_index()
-    embeddings, metas, texts = [], [], []
+# ---- Build / Rebuild index from JSON ----
+@app.post("/build_index")
+def build_index():
+    if not os.path.exists(JSON_FILE):
+        raise HTTPException(status_code=404, detail=f"{JSON_FILE} not found")
+    
+    with open(JSON_FILE, "r", encoding="utf-8") as f:
+        articles = json.load(f)
 
-    for up in files:
-        content = (await up.read()).decode("utf-8")
+    embeddings, texts, metas = [], [], []
+
+    for art_id, art in enumerate(articles):
+        content = art.get("Continut", "")
+        url = art.get("URL", "")
         chunks = chunk_text(content)
         embs = embed_model.encode(chunks, convert_to_numpy=True)
         embeddings.append(embs)
         texts.extend(chunks)
-        metas.extend([{"source": up.filename, "chunk_id": i} for i in range(len(chunks))])
+        metas.extend([{"source": art.get("Titlu", f"articol_{art_id}"), "url": url, "chunk_id": i} for i in range(len(chunks))])
 
     embeddings = np.vstack(embeddings).astype("float32")
-    if index is None:
-        index = faiss.IndexFlatIP(embeddings.shape[1])
-        faiss.normalize_L2(embeddings)
-        index.add(embeddings)
-        metadata = {"texts": texts, "metas": metas}
-    else:
-        faiss.normalize_L2(embeddings)
-        index.add(embeddings)
-        metadata["texts"].extend(texts)
-        metadata["metas"].extend(metas)
+    faiss.normalize_L2(embeddings)
+    index = faiss.IndexFlatIP(embeddings.shape[1])
+    index.add(embeddings)
+    metadata = {"texts": texts, "metas": metas}
 
     save_index(index, metadata)
-    return {"added_chunks": embeddings.shape[0], "total_chunks": len(metadata["texts"])}
+    return {"status": "ok", "total_chunks": len(texts)}
 
-# Endpoint: ask question
-from pydantic import BaseModel
+# ---- Ask endpoint ----
 class AskRequest(BaseModel):
     question: str
     top_k: int = 4
     max_answer_tokens: int = 256
 
-from fastapi import Depends
 @app.post("/ask")
 def ask(req: AskRequest):
     index, metadata = load_index()
     if index is None:
-        raise HTTPException(status_code=404, detail="No index found. Upload docs first.")
+        raise HTTPException(status_code=404, detail="No index found. Call /build_index first.")
+    
     q_emb = embed_model.encode([req.question], convert_to_numpy=True).astype("float32")
     faiss.normalize_L2(q_emb)
     D, I = index.search(q_emb, req.top_k)
+
     retrieved = [metadata["texts"][i] for i in I[0]]
+    urls = [metadata["metas"][i]["url"] for i in I[0] if "url" in metadata["metas"][i]]
+
     context = "\n\n".join(retrieved)
     prompt = f"Context:\n{context}\n\nQuestion: {req.question}\nAnswer:"
     out = gen_pipeline(prompt, max_length=req.max_answer_tokens, do_sample=False)[0]["generated_text"]
-    return {"answer": out, "sources": [metadata["metas"][i] for i in I[0]]}
 
-# Health check
+    return {
+        "answer": f"{out} Find out more at {', '.join(urls)}",
+        "sources": [metadata["metas"][i] for i in I[0]]
+    }
+
+# ---- Health check ----
 @app.get("/health")
 def health():
     return {"status": "ok"}