added source in response

app.py

@@ -3,7 +3,7 @@ import json
 from fastapi import FastAPI
 from pydantic import BaseModel
 from transformers import AutoTokenizer, AutoModelForSequenceClassification, pipeline
-
+from typing import Optional
 # --- 1. Load Model and Tokenizer ---
 
 # Define the path to your trained model
@@ -32,26 +32,32 @@ print("Classification pipeline loaded successfully.")
 
 # --- 2. Load Knowledge Base (Answers) ---
 
-intent_to_answer_map = {}
+
+# We will now store a dictionary for each intent
+intent_details_map = {}
 print(f"Loading knowledge base from: {KNOWLEDGE_BASE_PATH}")
 try:
     with open(KNOWLEDGE_BASE_PATH, 'r', encoding='utf-8') as f:
         knowledge_base_data = json.load(f)
     
-    # Create a simple lookup map: Intent -> Answer
-    # This assumes the answer is the same for all questions of the same intent.
-    intent_to_answer_map = {
-        item['intent']: item['answer'] for item in knowledge_base_data
-    }
-    print(f"Knowledge base loaded with {len(intent_to_answer_map)} intent-to-answer mappings.")
+    # Create a lookup map: Intent -> {answer, source}
+    # We only need to add each intent once, since the answer/source is the same
+    for item in knowledge_base_data:
+        if item['intent'] not in intent_details_map:
+            intent_details_map[item['intent']] = {
+                "answer": item.get('answer', 'No answer found.'),
+                "source": item.get('source', 'No source found.')
+            }
+            
+    print(f"Knowledge base loaded with {len(intent_details_map)} intent-to-detail mappings.")
 
 except FileNotFoundError:
     print(f"CRITICAL ERROR: Knowledge base file not found at {KNOWLEDGE_BASE_PATH}")
-    # In a real app, you might want to exit if the KB can't be loaded
 except Exception as e:
     print(f"Error loading knowledge base: {e}")
 
 
+
 # --- 3. Initialize FastAPI App ---
 
 app = FastAPI(
@@ -62,17 +68,21 @@ app = FastAPI(
 
 # --- 4. Define Request and Response Models ---
 
+# This is what the API will return
+
 # This is what the user must send in their POST request
 class Query(BaseModel):
     text: str
 
-# This is what the API will return
+# This is what the API will return (now includes 'source')
 class PredictionResponse(BaseModel):
     query: str
     predicted_intent: str
     confidence_score: float
     answer: str
+    source: Optional[str] # Make source optional in case it's missing
 
+ 
 # --- 5. Define API Endpoints ---
 
 @app.get("/")
@@ -81,7 +91,7 @@ def read_root():
     Root endpoint for health check.
     """
     return {"status": "API is running",
-            "message": "Post to /predict with a 'text' field to get an intent and answer."}
+            "message": "Post to /predict with a 'text' field to get an intent and answer and source."}
 
 
 @app.post("/predict", response_model=PredictionResponse)
@@ -96,15 +106,19 @@ def predict_intent(query: Query):
     predicted_intent = model_result['label']
     confidence_score = model_result['score']
     
-    # 2. Retrieve answer from our knowledge base
-    fallback_answer = "Could not find a specific answer for this intent. Please rephrase your question or contact a legal professional for advice."
-    
-    answer = intent_to_answer_map.get(predicted_intent, fallback_answer)
+    # 2. Retrieve details (answer and source) from our knowledge base
+    fallback_details = {
+        "answer": "Could not find a specific answer for this intent. Please rephrase your question or contact a legal professional for advice.",
+        "source": None
+    }
+    details = intent_details_map.get(predicted_intent, fallback_details)
     
     # 3. Return the combined response
     return {
         "query": query.text,
         "predicted_intent": predicted_intent,
         "confidence_score": confidence_score,
-        "answer": answer
+        "answer": details.get('answer'),
+        "source": details.get('source')
     }
+