Update app/app.py

app/app.py

@@ -3,27 +3,29 @@ import json
 import asyncio
 import logging
 import uuid
-import re
 from fastapi import FastAPI, HTTPException, Request
 from pydantic import BaseModel
-from typing import Optional
+from typing import Optional, Dict
 from llama_cpp import Llama
-from typing import Optional, Dict, List
 
 # Correctly reference the module within the 'app' package
 from app.policy_vector_db import PolicyVectorDB, ensure_db_populated
 
+
 # -----------------------------
 # ✅ Logging Configuration
 # -----------------------------
 logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - [%(request_id)s] - %(message)s')
 
+
 class RequestIdAdapter(logging.LoggerAdapter):
     def process(self, msg, kwargs):
-        return '[%s] %s' % (self.extra['request_id'], msg), kwargs
+        return '[%s] %s' % (self.extra.get('request_id', 'N/A'), msg), kwargs
+
 
 logger = logging.getLogger("app")
 
+
 # -----------------------------
 # ✅ Configuration
 # -----------------------------
@@ -34,11 +36,14 @@ LLM_TIMEOUT_SECONDS = int(os.getenv("LLM_TIMEOUT_SECONDS", "90"))
 RELEVANCE_THRESHOLD = float(os.getenv("RELEVANCE_THRESHOLD", "0.3"))
 TOP_K_SEARCH = int(os.getenv("TOP_K_SEARCH", "3"))
 TOP_K_CONTEXT = int(os.getenv("TOP_K_CONTEXT", "1"))
+LLM_THREADS = int(os.getenv("LLM_THREADS", "4"))  # configurable threads
+
 
 # -----------------------------
 # ✅ Initialize FastAPI App
 # -----------------------------
-app = FastAPI(title="NEEPCO DoP RAG Chatbot", version="2.1.0")
+app = FastAPI(title="NEEPCO DoP RAG Chatbot", version="2.1.1")
+
 
 @app.middleware("http")
 async def add_request_id(request: Request, call_next):
@@ -48,6 +53,7 @@ async def add_request_id(request: Request, call_next):
     response.headers["X-Request-ID"] = request_id
     return response
 
+
 # -----------------------------
 # ✅ Vector DB and Data Initialization
 # -----------------------------
@@ -71,19 +77,20 @@ except Exception as e:
     db = None
     db_ready = False
 
+
 # -----------------------------
-# ✅ Load TinyLlama GGUF Model with Improved Settings
+# ✅ Load TinyLlama GGUF Model with Safer Generation
 # -----------------------------
 logger.info(f"Loading GGUF model from: {MODEL_PATH}")
 try:
     llm = Llama(
         model_path=MODEL_PATH,
         n_ctx=2048,
-        n_threads=2,  # Increased threads for better performance
-        n_batch=256,  # Reduced batch size for stability
+        n_threads=LLM_THREADS,
+        n_batch=256,
         use_mlock=True,
         verbose=False,
-        seed=42  # Added seed for reproducible results
+        seed=42
     )
     logger.info("GGUF model loaded successfully.")
     model_ready = True
@@ -92,18 +99,21 @@ except Exception as e:
     llm = None
     model_ready = False
 
+
 # -----------------------------
 # ✅ API Schemas
 # -----------------------------
 class Query(BaseModel):
     question: str
 
+
 class AdvancedQuery(BaseModel):
     question: str
     section_filter: Optional[str] = None
     chunk_type_filter: Optional[str] = None
     top_k: Optional[int] = None
 
+
 class Feedback(BaseModel):
     request_id: str
     question: str
@@ -112,20 +122,22 @@ class Feedback(BaseModel):
     feedback: str
     comment: str | None = None
 
+
 # -----------------------------
 # ✅ Helper Functions
 # -----------------------------
 def get_logger_adapter(request: Request):
     return RequestIdAdapter(logger, {'request_id': getattr(request.state, 'request_id', 'N/A')})
 
+
 def get_chunk_priority(chunk: Dict) -> int:
     """Assign priority to different chunk types for better context selection"""
     priority_order = [
-        'approval_authority', 
-        'delegation_summary', 
-        'requirement', 
-        'method_specific', 
-        'board_approval', 
+        'approval_authority',
+        'delegation_summary',
+        'requirement',
+        'method_specific',
+        'board_approval',
         'financial_concurrence',
         'composition'
     ]
@@ -135,120 +147,103 @@ def get_chunk_priority(chunk: Dict) -> int:
     except ValueError:
         return len(priority_order)  # Lower priority for unknown types
 
+
 def detect_filters(question_lower: str) -> tuple:
     """Detect section and chunk type filters from user question"""
     section_filter = None
     chunk_type_filter = None
-    
+
     # Section keyword mapping
     section_keywords = {
         "annexure": "Annexure A",
-        "financial concurrence": "Financial Concurrence", 
+        "financial concurrence": "Financial Concurrence",
         "guidelines": "Guidelines",
         "section 1": "I", "section i": "I",
         "section 2": "II", "section ii": "II",
-        "section 3": "III", "section iii": "III", 
+        "section 3": "III", "section iii": "III",
         "section 4": "IV", "section iv": "IV"
     }
-    
+
     # Chunk type keyword mapping
     chunk_type_keywords = {
         "approval": "approval_authority",
         "delegation": "delegation_summary",
-        "requirement": "requirement", 
+        "requirement": "requirement",
         "method": "method_specific",
         "board": "board_approval",
         "committee": "composition"
     }
-    
-    # Check for section filters
+
     for keyword, section in section_keywords.items():
         if keyword in question_lower:
             section_filter = section
             break
-    
-    # Check for chunk type filters
+
     for keyword, chunk_type in chunk_type_keywords.items():
         if keyword in question_lower:
             chunk_type_filter = chunk_type
             break
-    
+
     return section_filter, chunk_type_filter
 
+
 def clean_llm_response(raw_response: str) -> str:
-    """Clean and validate LLM response"""
+    """Simplified cleaner to avoid over-trimming."""
     if not raw_response:
         return ""
-    
-    # Remove common unwanted patterns
-    cleaned = raw_response.strip()
-    
-    # Remove incomplete sentences at the end
-    if cleaned and not cleaned.endswith(('.', '!', '?', ':', '|')):
-        # Find the last complete sentence
-        sentences = re.split(r'[.!?]', cleaned)
-        if len(sentences) > 1:
-            cleaned = '.'.join(sentences[:-1]) + '.'
-    
-    return cleaned
+    return raw_response.strip()
+
 
 async def generate_llm_response(prompt: str, request_id: str, adapter: RequestIdAdapter):
-    """Improved LLM response generation with better error handling"""
+    """LLM response generation with safer stops and robust extraction."""
     loop = asyncio.get_running_loop()
-    
-    # Multiple generation attempts with different parameters
+
+    # Use plain completion configs without fragile stop tokens
     generation_configs = [
-        {
-            "max_tokens": 512,
-            "temperature": 0.1,
-            "top_p": 0.9,
-            "repeat_penalty": 1.1,
-            "stop": ["</s>", "[INST]", "[/INST]", "Question:", "Context:", "###"]
-        },
-        {
-            "max_tokens": 256,
-            "temperature": 0.3,
-            "top_p": 0.8,
-            "repeat_penalty": 1.2,
-            "stop": ["</s>", "\n\n", "Question:", "Context:"]
-        },
-        {
-            "max_tokens": 128,
-            "temperature": 0.5,
-            "top_p": 0.7,
-            "repeat_penalty": 1.15,
-            "stop": ["</s>"]
-        }
+        {"max_tokens": 512, "temperature": 0.2, "top_p": 0.9, "repeat_penalty": 1.1, "stop": []},
+        {"max_tokens": 384, "temperature": 0.3, "top_p": 0.9, "repeat_penalty": 1.1, "stop": []},
+        {"max_tokens": 256, "temperature": 0.4, "top_p": 0.9, "repeat_penalty": 1.1, "stop": []},
     ]
-    
+
     for attempt, config in enumerate(generation_configs, 1):
         try:
             adapter.info(f"LLM generation attempt {attempt}/{len(generation_configs)} with config: {config}")
-            
+
             response = await loop.run_in_executor(
                 None,
                 lambda: llm(prompt, echo=False, **config)
             )
-            
-            raw_answer = response["choices"][0]["text"]
+
+            # Debug: log a truncated snapshot of the raw response
+            try:
+                adapter.info(f"Raw LLM response object (truncated): {json.dumps(response)[:1200]}")
+            except Exception:
+                pass
+
+            raw_answer = ""
+            if isinstance(response, dict) and "choices" in response and response["choices"]:
+                choice = response["choices"][0]
+                if isinstance(choice, dict):
+                    raw_answer = choice.get("text") or choice.get("message", {}).get("content", "") or ""
+
             cleaned_answer = clean_llm_response(raw_answer)
-            
             adapter.info(f"Attempt {attempt} - Raw response length: {len(raw_answer)}, Cleaned length: {len(cleaned_answer)}")
-            
-            if cleaned_answer and len(cleaned_answer.strip()) > 10:  # Minimum meaningful response
+
+            # Accept concise answers
+            if cleaned_answer and len(cleaned_answer.strip()) > 3:
                 adapter.info(f"Successful generation on attempt {attempt}")
                 return cleaned_answer
             else:
                 adapter.warning(f"Attempt {attempt} produced insufficient response: '{cleaned_answer}'")
-                
+
         except Exception as e:
-            adapter.error(f"Attempt {attempt} failed: {e}")
+            adapter.error(f"Attempt {attempt} failed: {e}", exc_info=True)
             continue
-    
-    # If all attempts fail, return a fallback message
+
     adapter.error("All LLM generation attempts failed")
     raise ValueError("Unable to generate a meaningful response after multiple attempts")
 
+
 # -----------------------------
 # ✅ Endpoints
 # -----------------------------
@@ -256,6 +251,7 @@ async def generate_llm_response(prompt: str, request_id: str, adapter: RequestId
 async def root():
     return {"status": "✅ Server is running."}
 
+
 @app.get("/health")
 async def health_check():
     status = {
@@ -268,6 +264,7 @@ async def health_check():
         raise HTTPException(status_code=503, detail=status)
     return status
 
+
 @app.post("/chat")
 async def chat(query: Query, request: Request):
     adapter = get_logger_adapter(request)
@@ -298,19 +295,19 @@ async def chat(query: Query, request: Request):
 
     # 1. Enhanced Search with potential filtering
     section_filter, chunk_type_filter = detect_filters(question_lower)
-    
+
     if section_filter or chunk_type_filter:
         adapter.info(f"Detected filters - section: '{section_filter}', chunk_type: '{chunk_type_filter}'")
         search_results = db.search_with_filters(
-            query.question, 
+            query.question,
             top_k=TOP_K_SEARCH,
             section_filter=section_filter,
             chunk_type_filter=chunk_type_filter
         )
-        adapter.info(f"Used filtered search")
+        adapter.info("Used filtered search")
     else:
         search_results = db.search(query.question, top_k=TOP_K_SEARCH)
-        adapter.info(f"Used regular search")
+        adapter.info("Used regular search")
 
     if not search_results:
         adapter.warning("No relevant context found in vector DB.")
@@ -325,7 +322,7 @@ async def chat(query: Query, request: Request):
     chunk_types = [result['metadata'].get('chunk_type', 'unknown') for result in search_results]
     sections = [result['metadata'].get('section', 'unknown') for result in search_results]
     scores = [f"{result['relevance_score']:.4f}" for result in search_results]
-    
+
     adapter.info(f"Found {len(search_results)} relevant chunks")
     adapter.info(f"Chunk types: {chunk_types}")
     adapter.info(f"Sections: {sections}")
@@ -333,8 +330,6 @@ async def chat(query: Query, request: Request):
 
     # 3. Prioritize chunk types for better context selection
     prioritized_results = sorted(search_results, key=lambda x: (get_chunk_priority(x), -x['relevance_score']))
-    
-    # Log prioritization results
     prioritized_types = [result['metadata'].get('chunk_type', 'unknown') for result in prioritized_results]
     adapter.info(f"Prioritized chunk types order: {prioritized_types}")
 
@@ -353,19 +348,25 @@ async def chat(query: Query, request: Request):
             'score': f"{result['relevance_score']:.4f}"
         }
         context_metadata.append(context_info)
-    
-    adapter.info(f"Selected context metadata: {context_metadata}")
 
-    # 6. Build Improved Prompt for TinyLlama
-    prompt = f"""[INST] You are a helpful assistant for NEEPCO's Delegation of Powers policy. Answer the question using only the provided context.
-
-Context: {context}
-
-Question: {query.question}
+    adapter.info(f"Selected context metadata: {context_metadata}")
 
-Provide a clear, direct answer based only on the context above. If the context doesn't contain the information, say "The provided policy context does not contain information on this topic."
+    # 6. Build Plain Completion Prompt (no [INST] tags)
+    prompt = (
+        "You are a helpful assistant for NEEPCO's Delegation of Powers policy. "
+        "Answer the question using only the provided context.\n\n"
+        f"Context:\n{context}\n\n"
+        f"Question:\n{query.question}\n\n"
+        "Provide a clear, direct answer based only on the context above. If the context doesn't contain the information, "
+        "say \"The provided policy context does not contain information on this topic.\"\n\n"
+        "Answer:\n"
+    )
 
-Answer: [/INST]"""
+    # Optional: log a short preview of the prompt to debug future issues (safe/truncated)
+    try:
+        adapter.info(f"Prompt preview (first 400 chars): {prompt[:400].replace(chr(10),' ')}")
+    except Exception:
+        pass
 
     # 7. Generate Response
     answer = "An error occurred while processing your request."
@@ -375,7 +376,7 @@ Answer: [/INST]"""
             generate_llm_response(prompt, request.state.request_id, adapter),
             timeout=LLM_TIMEOUT_SECONDS
         )
-        
+
         adapter.info(f"LLM generation successful. Response length: {len(raw_answer)}")
 
         # --- POST-PROCESSING LOGIC ---
@@ -394,7 +395,7 @@ Answer: [/INST]"""
         adapter.error(f"An unexpected error occurred during LLM generation: {e}", exc_info=True)
         answer = "Sorry, an unexpected error occurred while generating a response."
 
-    adapter.info(f"Final answer prepared. Returning to client.")
+    adapter.info("Final answer prepared. Returning to client.")
 
     return {
         "request_id": request.state.request_id,
@@ -403,23 +404,24 @@ Answer: [/INST]"""
         "answer": answer
     }
 
+
 @app.post("/advanced_search")
 async def advanced_search(query: AdvancedQuery, request: Request):
     """Advanced search endpoint with explicit filters"""
     adapter = get_logger_adapter(request)
-    
+
     if not db_ready:
         raise HTTPException(status_code=503, detail="Database not ready")
-    
+
     adapter.info(f"Advanced search: question='{query.question}', section='{query.section_filter}', chunk_type='{query.chunk_type_filter}'")
-    
+
     search_results = db.search_with_filters(
         query.question,
         top_k=query.top_k or TOP_K_SEARCH,
         section_filter=query.section_filter,
         chunk_type_filter=query.chunk_type_filter
     )
-    
+
     return {
         "request_id": request.state.request_id,
         "query": query.question,
@@ -437,6 +439,7 @@ async def advanced_search(query: AdvancedQuery, request: Request):
         ]
     }
 
+
 @app.post("/feedback")
 async def collect_feedback(feedback: Feedback, request: Request):
     adapter = get_logger_adapter(request)
@@ -451,4 +454,4 @@ async def collect_feedback(feedback: Feedback, request: Request):
     }
 
     adapter.info(json.dumps(feedback_log))
-    return {"status": "✅ Feedback recorded. Thank you!"}
+    return {"status": "✅ Feedback recorded. Thank you!"}