Update main.py

main.py

@@ -5,11 +5,15 @@ import json
 from sentence_transformers import SentenceTransformer
 import google.generativeai as genai
 import os
+from dotenv import load_dotenv
 
 # ---------------------
-# Config
+# Startup Config
 # ---------------------
 
+print("Loading environment variables...")
+load_dotenv()
+
 print("Loading songs data...")
 with open("songs.json", encoding="utf-8") as f:
     songs = json.load(f)
@@ -17,28 +21,28 @@ with open("songs.json", encoding="utf-8") as f:
 print("Loading embeddings...")
 embeddings = np.load("song_embeddings_e5_final.npy")
 
-print("Loading model...")
+print("Loading embedding model...")
 model = SentenceTransformer("intfloat/multilingual-e5-large")
 
-# Configure Gemini API
+print("Configuring Gemini API...")
 genai.configure(api_key=os.getenv("GEMINI_API_KEY"))
 gemini_model = genai.GenerativeModel("gemini-2.5-flash")
 
 print("API ready!")
 
 # ---------------------
-# FastAPI app
+# FastAPI App
 # ---------------------
+
 app = FastAPI(
     title="Thirumandiram Search API",
-    description="Semantic search API for Thirumandiram verses with AI summaries",
-    version="1.0.0"
+    description="Semantic search and AI-assisted explanations for Thirumandiram verses",
+    version="2.0.0"
 )
 
-# Allow CORS for your frontend
 app.add_middleware(
     CORSMiddleware,
-    allow_origins=["*"],  # Change to your frontend domain in production
+    allow_origins=["*"],
     allow_methods=["*"],
     allow_headers=["*"],
 )
@@ -46,6 +50,7 @@ app.add_middleware(
 # ---------------------
 # Payiram Mapper
 # ---------------------
+
 def get_payiram(song_number: int) -> str:
     if 1 <= song_number <= 336:
         return "First Payiram"
@@ -68,8 +73,9 @@ def get_payiram(song_number: int) -> str:
     return "Unknown Payiram"
 
 # ---------------------
-# Utility: Search top-k matches
+# Semantic Search
 # ---------------------
+
 def search_songs(query: str, top_k: int = 3):
     query_text = "query: " + query
     query_vec = model.encode([query_text])[0]
@@ -96,15 +102,40 @@ def search_songs(query: str, top_k: int = 3):
 
     return results
 
+# ---------------------
+# Gemini Scope Classifier
+# ---------------------
+
+def is_thirumandiram_scope(query: str) -> bool:
+    prompt = f"""
+You are a strict classifier.
+
+Decide whether the following user query is related to:
+- Thirumandiram
+- Thirumoolar
+- Saivism, Siddha philosophy, Yoga
+- Spiritual concepts explained in Thirumandiram verses
+
+Respond with ONLY:
+YES or NO
+
+If unsure, respond NO.
+
+User query:
+"{query}"
+"""
+    response = gemini_model.generate_content(prompt)
+    return response.text.strip().upper() == "YES"
+
 # ---------------------
 # API Endpoints
 # ---------------------
+
 @app.get("/")
 def root():
-    """API information and available endpoints"""
     return {
         "name": "Thirumandiram Search API",
-        "version": "1.0.0",
+        "version": "2.0.0",
         "endpoints": {
             "search": "/search?q=<query>&top_k=3",
             "chat_search": "/chat_search?q=<query>&top_k=3",
@@ -114,54 +145,75 @@ def root():
     }
 
 @app.get("/health")
-def health_check():
-    """Health check endpoint"""
-    return {"status": "healthy", "model_loaded": model is not None}
+def health():
+    return {
+        "status": "healthy",
+        "embedding_model_loaded": model is not None,
+        "gemini_configured": os.getenv("GEMINI_API_KEY") is not None
+    }
+
+# ---------------------
+# Endpoint 1: Raw Semantic Search
+# ---------------------
 
 @app.get("/search")
-def search(q: str = Query(..., description="Search query in Tamil or English"), 
-           top_k: int = Query(3, ge=1, le=10, description="Number of results to return")):
-    """
-    Semantic search for Thirumandiram verses
-    
-    Returns matching verses with similarity scores
-    """
-    results = search_songs(q, top_k)
-    return {"query": q, "results": results}
+def search(
+    q: str = Query(..., description="Search query in Tamil or English"),
+    top_k: int = Query(3, ge=1, le=10)
+):
+    return {
+        "query": q,
+        "results": search_songs(q, top_k)
+    }
+
+# ---------------------
+# Endpoint 2: Chat Search (Gemini-Gated)
+# ---------------------
 
 @app.get("/chat_search")
-def chat_search(q: str = Query(..., description="Search query in Tamil or English"),
-                top_k: int = Query(3, ge=1, le=10, description="Number of results to return")):
-    """
-    Semantic search with AI-generated summary
-    
-    Returns matching verses along with a summarized explanation using Google Gemini
-    """
+def chat_search(
+    q: str = Query(..., description="Search query in Tamil or English"),
+    top_k: int = Query(3, ge=1, le=10)
+):
+    # STEP 1: Scope check
+    if not is_thirumandiram_scope(q):
+        return {
+            "query": q,
+            "out_of_scope": True,
+            "message": "The query is not within the scope of Thirumandiram.",
+            "summary": None,
+            "results": []
+        }
+
+    # STEP 2: Semantic search
     results = search_songs(q, top_k)
 
-    # Build context
+    # STEP 3: Context building
     context = "\n\n".join(
-        [
-            f"Song {r['song_number']} ({r['payiram']}):\nVerse:\n{r['padal']}\nExplanation:\n{r['vilakam_en']}"
-            for r in results
-        ]
+        f"Song {r['song_number']} ({r['payiram']}):\n"
+        f"Verse:\n{r['padal']}\n"
+        f"Explanation:\n{r['vilakam_en']}"
+        for r in results
     )
 
     prompt = f"""
-You are a helpful assistant. A user searched for: "{q}".
-Here are some relevant verses from Thirumandiram:
+You are a Thirumandiram expert assistant.
+Answer ONLY using Thirumandiram philosophy.
+
+User query:
+"{q}"
 
+Relevant verses:
 {context}
 
-Please summarize these results in a clear, chatbot-friendly way.
-Explain the key ideas and how they relate to the query.
+Explain clearly how these verses address the query.
 """
 
-    # Generate summary using Gemini
     response = gemini_model.generate_content(prompt)
 
     return {
         "query": q,
+        "out_of_scope": False,
         "summary": response.text,
         "results": results
     }