Update app.py

app.py

@@ -2,14 +2,14 @@
 Khalifa University Library RAG Backend
 LangChain + FAISS + FastAPI on Hugging Face Spaces
 
-This app:
-1. Loads scraped library pages from the 'knowledge/' folder
-2. Chunks and embeds them using OpenAI embeddings
-3. Stores in a FAISS vector store
-4. Exposes a /rag endpoint that retrieves relevant chunks and generates grounded answers
+Features:
+- OpenAI embeddings (text-embedding-3-small) for vector search
+- Switchable LLM: ChatGPT or Claude for answer generation
+- FAISS vector store for fast similarity search
 
 Environment variables (set as HF Space Secrets):
-  OPENAI_API_KEY  — for embeddings + LLM
+  OPENAI_API_KEY     — required (embeddings + ChatGPT answers)
+  ANTHROPIC_API_KEY  — optional (enables Claude answers)
 """
 
 import os
@@ -32,11 +32,29 @@ FAISS_INDEX_PATH = "faiss_index"
 CHUNK_SIZE = 800
 CHUNK_OVERLAP = 100
 EMBEDDING_MODEL = "text-embedding-3-small"
-LLM_MODEL = "gpt-4o-mini"
 TOP_K = 5
 
+# ===== RAG PROMPT =====
+RAG_PROMPT = PromptTemplate(
+    input_variables=["context", "question"],
+    template="""You are the Khalifa University Library AI Assistant in Abu Dhabi, UAE.
+KU means Khalifa University, NOT Kuwait University.
+
+Use ONLY the following context from the Khalifa University Library website to answer the question.
+If the context doesn't contain enough information, say "I don't have specific information about this in our library knowledge base" and suggest contacting Ask a Librarian at https://library.ku.ac.ae/AskUs
+
+Always include relevant URLs from the context when available.
+Keep answers concise (2-4 sentences) and helpful.
+
+Context:
+{context}
+
+Question: {question}
+
+Answer:"""
+)
+
 # ===== GLOBAL STATE =====
-qa_chain = None
 vectorstore = None
 
 
@@ -51,7 +69,6 @@ def load_documents():
             with open(filepath, "r", encoding="utf-8") as f:
                 content = f.read()
 
-            # Extract metadata from first two lines
             lines = content.split("\n", 3)
             source = ""
             title = ""
@@ -89,14 +106,12 @@ def build_vectorstore(docs):
 
     embeddings = OpenAIEmbeddings(model=EMBEDDING_MODEL)
 
-    # Try to load existing index first
-    if os.path.exists(FAISS_INDEX_PATH):
+    if os.path.exists(os.path.join(FAISS_INDEX_PATH, "index.faiss")):
         print("Loading existing FAISS index...")
         store = FAISS.load_local(FAISS_INDEX_PATH, embeddings, allow_dangerous_deserialization=True)
         print(f"Loaded FAISS index with {store.index.ntotal} vectors")
         return store
 
-    # Build new index
     print("Building new FAISS index...")
     store = FAISS.from_documents(chunks, embeddings)
     store.save_local(FAISS_INDEX_PATH)
@@ -104,34 +119,35 @@ def build_vectorstore(docs):
     return store
 
 
-def build_chain(store):
-    """Build the LangChain RetrievalQA chain."""
-    llm = ChatOpenAI(model=LLM_MODEL, temperature=0.2, max_tokens=500)
-
-    prompt_template = PromptTemplate(
-        input_variables=["context", "question"],
-        template="""You are the Khalifa University Library AI Assistant in Abu Dhabi, UAE.
-KU means Khalifa University, NOT Kuwait University.
-
-Use ONLY the following context from the Khalifa University Library website to answer the question.
-If the context doesn't contain enough information, say "I don't have specific information about this in our library knowledge base" and suggest contacting Ask a Librarian at https://library.ku.ac.ae/AskUs
-
-Always include relevant URLs from the context when available.
-Keep answers concise (2-4 sentences) and helpful.
-
-Context:
-{context}
-
-Question: {question}
+def get_llm(model_name: str = "gpt"):
+    """Get LLM based on model selection."""
+    if model_name == "claude":
+        api_key = os.environ.get("ANTHROPIC_API_KEY")
+        if not api_key:
+            raise ValueError("ANTHROPIC_API_KEY not configured. Switch to ChatGPT or add the key in Space Secrets.")
+        from langchain_anthropic import ChatAnthropic
+        return ChatAnthropic(
+            model="claude-sonnet-4-20250514",
+            temperature=0.2,
+            max_tokens=500,
+            anthropic_api_key=api_key,
+        )
+    else:
+        return ChatOpenAI(
+            model="gpt-4o-mini",
+            temperature=0.2,
+            max_tokens=500,
+        )
 
-Answer:"""
-    )
 
+def build_chain(store, model_name: str = "gpt"):
+    """Build the LangChain RetrievalQA chain with the selected LLM."""
+    llm = get_llm(model_name)
     chain = RetrievalQA.from_chain_type(
         llm=llm,
         chain_type="stuff",
         retriever=store.as_retriever(search_kwargs={"k": TOP_K}),
-        chain_type_kwargs={"prompt": prompt_template},
+        chain_type_kwargs={"prompt": RAG_PROMPT},
         return_source_documents=True,
     )
     return chain
@@ -140,20 +156,23 @@ Answer:"""
 # ===== STARTUP =====
 @asynccontextmanager
 async def lifespan(app: FastAPI):
-    global qa_chain, vectorstore
+    global vectorstore
     print("=== Starting KU Library RAG Backend ===")
 
     docs = load_documents()
     if docs:
         vectorstore = build_vectorstore(docs)
-        qa_chain = build_chain(vectorstore)
-        print("RAG chain ready!")
+        print("Vector store ready!")
     else:
-        print("WARNING: No knowledge files found. RAG will not work.")
-        print(f"Please add .txt files to the '{KNOWLEDGE_DIR}/' directory.")
+        print("WARNING: No knowledge files found.")
+        print(f"Add .txt files to '{KNOWLEDGE_DIR}/' and restart.")
 
-    yield
+    # Check which LLMs are available
+    has_openai = bool(os.environ.get("OPENAI_API_KEY"))
+    has_claude = bool(os.environ.get("ANTHROPIC_API_KEY"))
+    print(f"LLMs available: ChatGPT={'YES' if has_openai else 'NO'}, Claude={'YES' if has_claude else 'NO'}")
 
+    yield
     print("Shutting down...")
 
 
@@ -162,14 +181,16 @@ app = FastAPI(title="KU Library RAG", lifespan=lifespan)
 
 app.add_middleware(
     CORSMiddleware,
-    allow_origins=["*"],  # Restrict to your domains in production
+    allow_origins=["*"],
     allow_methods=["POST", "GET"],
     allow_headers=["*"],
 )
 
 
+# ===== REQUEST/RESPONSE MODELS =====
 class QueryRequest(BaseModel):
     question: str
+    model: str = "gpt"  # "gpt" or "claude"
     top_k: int = 5
 
 
@@ -181,30 +202,40 @@ class SourceDoc(BaseModel):
 
 class QueryResponse(BaseModel):
     answer: str
+    model_used: str
     sources: list[SourceDoc]
     error: str | None = None
 
 
+# ===== ENDPOINTS =====
 @app.get("/")
 def health():
     return {
         "status": "ok",
-        "rag_ready": qa_chain is not None,
+        "vectorstore_ready": vectorstore is not None,
+        "models": {
+            "gpt": bool(os.environ.get("OPENAI_API_KEY")),
+            "claude": bool(os.environ.get("ANTHROPIC_API_KEY")),
+        },
         "service": "KU Library RAG Backend",
     }
 
 
 @app.post("/rag", response_model=QueryResponse)
 async def rag_query(req: QueryRequest):
-    if not qa_chain:
+    if not vectorstore:
         return QueryResponse(
             answer="RAG system not initialized. Knowledge base may be empty.",
+            model_used=req.model,
             sources=[],
             error="No knowledge files loaded"
         )
 
+    model_name = req.model if req.model in ("gpt", "claude") else "gpt"
+
     try:
-        result = qa_chain.invoke({"query": req.question})
+        chain = build_chain(vectorstore, model_name)
+        result = chain.invoke({"query": req.question})
         answer = result.get("result", "No answer generated.")
         source_docs = result.get("source_documents", [])
 
@@ -222,11 +253,53 @@ async def rag_query(req: QueryRequest):
                     snippet=doc.page_content[:200] + "..."
                 ))
 
-        return QueryResponse(answer=answer, sources=sources)
+        return QueryResponse(
+            answer=answer,
+            model_used=model_name,
+            sources=sources,
+        )
+
+    except ValueError as e:
+        # Model not available — fallback to the other one
+        fallback = "gpt" if model_name == "claude" else "claude"
+        try:
+            chain = build_chain(vectorstore, fallback)
+            result = chain.invoke({"query": req.question})
+            answer = result.get("result", "No answer generated.")
+            source_docs = result.get("source_documents", [])
+
+            sources = []
+            seen = set()
+            for doc in source_docs:
+                src = doc.metadata.get("source", "")
+                title = doc.metadata.get("title", "")
+                key = src or title
+                if key and key not in seen:
+                    seen.add(key)
+                    sources.append(SourceDoc(
+                        title=title,
+                        source=src,
+                        snippet=doc.page_content[:200] + "..."
+                    ))
+
+            return QueryResponse(
+                answer=answer,
+                model_used=fallback,
+                sources=sources,
+                error=f"{model_name} unavailable, used {fallback} instead"
+            )
+        except Exception as e2:
+            return QueryResponse(
+                answer="Both models failed. Please check API keys in Space Secrets.",
+                model_used="none",
+                sources=[],
+                error=str(e2)
+            )
 
     except Exception as e:
         return QueryResponse(
             answer="Sorry, I encountered an error processing your question.",
+            model_used=model_name,
             sources=[],
             error=str(e)
         )
@@ -235,7 +308,7 @@ async def rag_query(req: QueryRequest):
 @app.post("/rebuild")
 async def rebuild_index():
     """Force rebuild the FAISS index from knowledge files."""
-    global qa_chain, vectorstore
+    global vectorstore
     try:
         if os.path.exists(FAISS_INDEX_PATH):
             import shutil
@@ -246,7 +319,6 @@ async def rebuild_index():
             return {"error": "No knowledge files found"}
 
         vectorstore = build_vectorstore(docs)
-        qa_chain = build_chain(vectorstore)
         return {"status": "ok", "chunks": vectorstore.index.ntotal}
     except Exception as e:
-        return {"error": str(e)}
+        return {"error": str(e)}