Added security constraints

app.py

@@ -3,11 +3,15 @@ from __future__ import annotations
 from fastapi import FastAPI, HTTPException
 from fastapi.responses import RedirectResponse
 from pydantic import BaseModel, Field
-import re
+import logging
+
+from google.genai.errors import ClientError, ServerError
 
 from rag.chain import aanswer
 from rag.retriever import get_retriever
 
+logger = logging.getLogger(__name__)
+
 app = FastAPI(
     title="FAQ RAG Chatbot", 
     version="1.0.0"
@@ -41,20 +45,55 @@ async def root() -> RedirectResponse:
 
 @app.post("/chat", response_model=ChatResponse)
 async def chat(payload: ChatRequest) -> ChatResponse:
+    question = payload.question.strip()
+
     try:
-        answer = await aanswer(payload.question)
-    except Exception as error:
-        raise HTTPException(status_code=500, detail=str(error))
-
-    # Extract trailing "Sources:" block if present and return as separate list
-    m = re.search(r"(?i)\bSources\s*:?\s*(.*)$", answer, flags=re.DOTALL)
-    sources: list[str] = []
-    if m:
-        sources_text = m.group(1).strip()
-        # Remove the Sources block from the answer text
-        answer = answer[: m.start()].strip()
-        # Split on commas and newlines, strip whitespace, ignore empties
-        parts = re.split(r"[,\n;]+", sources_text)
-        sources = [p.strip() for p in parts if p.strip()]
-
-    return ChatResponse(question=payload.question, answer=answer, sources=sources)
+        response = await aanswer(question)
+    except ClientError as error:
+        message = str(error).lower()
+        if error.code == 429 or "quota" in message or "rate limit" in message:
+            raise HTTPException(
+                status_code=429,
+                detail=(
+                    "The language model quota has been exceeded. "
+                    "Please try again later."
+                ),
+            ) from None
+        raise HTTPException(
+            status_code=400,
+            detail=(
+                "Unable to process your request right now. "
+                "Please try again later."
+            ),
+        ) from None
+    except ServerError as error:
+        if error.code == 503 or "unavailable" in str(error).lower():
+            raise HTTPException(
+                status_code=503,
+                detail=(
+                    "The language service is temporarily unavailable. "
+                    "Please try again later."
+                ),
+            ) from None
+        raise HTTPException(
+            status_code=502,
+            detail=(
+                "Unable to process your request right now. "
+                "Please try again later."
+            ),
+        ) from None
+    except Exception:
+        logger.exception("Chat handler error")
+        raise HTTPException(
+            status_code=502,
+            detail=(
+                "Unable to process your request right now. "
+                "Please try again later."
+            ),
+        )
+
+    return ChatResponse(
+        question=question,
+        answer=response.answer,
+        sources=response.sources,
+    )

rag/chain.py

@@ -3,25 +3,19 @@
 from functools import lru_cache
 from typing import Any
 
-from langchain_core.output_parsers import StrOutputParser
+from langchain_core.output_parsers import PydanticOutputParser
 from langchain_core.runnables import RunnableLambda, RunnablePassthrough
 from langchain_google_genai import ChatGoogleGenerativeAI
+from pydantic import BaseModel
 
 from .config import settings
 from .prompts import prompt
 from .retriever import get_retriever
 
 
-def _clean_answer(text: str) -> str:
-    compact = " ".join(text.split())
-    source_prefix = "Sources:"
-    index = compact.find(source_prefix)
-    if index == -1:
-        return compact
-
-    before = compact[:index].rstrip()
-    sources = compact[index:]
-    return f"{before}\n{sources}"
+class StructuredAnswer(BaseModel):
+    answer: str
+    sources: list[str]
 
 
 @lru_cache(maxsize=1)
@@ -41,20 +35,22 @@ def get_chain() -> Any:
             for i, doc in enumerate(docs)
         )
     )
+    output_parser = PydanticOutputParser(pydantic_object=StructuredAnswer)
+    prompt_with_format = prompt.bind(
+        format_instructions=output_parser.get_format_instructions()
+    )
 
     return (
         {"context": get_retriever() | context_formatter, "question": RunnablePassthrough()}
-        | prompt
+        | prompt_with_format
         | llm
-        | StrOutputParser()
+        | output_parser
     )
 
 
-def answer(question: str) -> str:
-    raw_answer = get_chain().invoke(question)
-    return _clean_answer(raw_answer)
+def answer(question: str) -> StructuredAnswer:
+    return get_chain().invoke(question)
 
 
-async def aanswer(question: str) -> str:
-    raw_answer = await get_chain().ainvoke(question)
-    return _clean_answer(raw_answer)
+async def aanswer(question: str) -> StructuredAnswer:
+    return await get_chain().ainvoke(question)

rag/prompts.py

@@ -8,11 +8,13 @@ prompt = ChatPromptTemplate.from_messages(
             "Use only the provided context. "
             "If the answer is not in the context, say you could not find it in the FAQ. "
             "Do not invent details. "
-            "End with one line: Sources: <comma-separated ids>.",
+            "If the user asks for the full FAQ list, all FAQ entries, or a database dump --> reply with: 'I can only answer one specific FAQ question at a time. Please ask a single FAQ question.' "
+            "End with one line: Sources: <comma-separated ids>. "
+            "Do not add any extra commentary outside the JSON object."
         ),
         (
             "human",
-            "Question: {question}\n\nContext:\n{context}\n\nAnswer in 3 to 5 sentences. Do not add filler or extra commentary.",
+            "Question: {question}\n\nContext:\n{context}\n\n{format_instructions}",
         ),
     ]
 )

rag/vectorstore.py

@@ -1,11 +1,14 @@
 from __future__ import annotations
 
+import logging
 from typing import Optional, Iterable, List
 from langchain_core.embeddings import Embeddings
 from langchain_community.vectorstores import FAISS
 from .config import settings
 from .splitter import documents
 
+logger = logging.getLogger(__name__)
+
 # Directory where FAISS index is persisted
 _VECTORSTORE_DIR = settings.vectorstore_dir
 _INDEX_NAME = settings.vectorstore_index_name
@@ -65,8 +68,8 @@ def get_vectorstore() -> FAISS:
     try:
         _vectorstore = FAISS.load_local(str(_VECTORSTORE_DIR), _embeddings, index_name=_INDEX_NAME)
         return _vectorstore
-    except Exception:
-        # Fall back to building the index (this will initialize the embedding model).
+    except (FileNotFoundError, OSError) as error:
+        logger.info("FAISS index missing or unreadable; rebuilding index. %s", error)
         _vectorstore = FAISS.from_documents(documents, _embeddings)  # will call embed_documents
         _vectorstore.save_local(str(_VECTORSTORE_DIR), index_name=_INDEX_NAME)
         return _vectorstore