Final SaaS Security & Asset Support

backend/src/services/chat_service.py

@@ -1,4 +1,213 @@
 
+# import json
+# from sqlalchemy.ext.asyncio import AsyncSession
+# from sqlalchemy.future import select
+
+# # --- Model Imports ---
+# from backend.src.models.chat import ChatHistory
+# from backend.src.models.integration import UserIntegration
+# from backend.src.models.user import User  # Added User model for Bot Persona
+
+# # --- Dynamic Factory & Tool Imports ---
+# from backend.src.services.llm.factory import get_llm_model
+# from backend.src.services.vector_store.qdrant_adapter import get_vector_store
+# from backend.src.services.security.pii_scrubber import PIIScrubber
+
+# # --- Agents ---
+# from backend.src.services.tools.secure_agent import get_secure_agent 
+# from backend.src.services.tools.nosql_agent import get_nosql_agent 
+# from backend.src.services.tools.cms_agent import get_cms_agent
+
+# # --- Router ---
+# from backend.src.services.routing.semantic_router import SemanticRouter
+
+# # --- LangChain Core ---
+# from langchain_core.messages import HumanMessage, AIMessage
+# from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
+
+# # ==========================================
+# # HELPER FUNCTIONS
+# # ==========================================
+
+# async def get_user_integrations(user_id: str, db: AsyncSession) -> dict:
+#     """Fetches active integrations and filters valid descriptions."""
+#     if not user_id: return {}
+    
+#     query = select(UserIntegration).where(UserIntegration.user_id == user_id, UserIntegration.is_active == True)
+#     result = await db.execute(query)
+#     integrations = result.scalars().all()
+    
+#     settings = {}
+#     for i in integrations:
+#         try:
+#             creds = json.loads(i.credentials)
+#             creds['provider'] = i.provider
+#             creds['schema_map'] = i.schema_map if i.schema_map else {}
+            
+#             # --- STRICT CHECK ---
+#             if i.profile_description:
+#                 creds['description'] = i.profile_description
+            
+#             settings[i.provider] = creds
+#         except (json.JSONDecodeError, TypeError):
+#             continue
+#     return settings
+
+# async def save_chat_to_db(db: AsyncSession, session_id: str, human_msg: str, ai_msg: str, provider: str):
+#     """Saves chat history with PII redaction."""
+#     if not session_id: return
+#     safe_human = PIIScrubber.scrub(human_msg)
+#     safe_ai = PIIScrubber.scrub(ai_msg)
+#     new_chat = ChatHistory(
+#         session_id=session_id, human_message=safe_human, ai_message=safe_ai, provider=provider
+#     )
+#     db.add(new_chat)
+#     await db.commit()
+
+# async def get_chat_history(session_id: str, db: AsyncSession):
+#     """Retrieves past conversation history."""
+#     if not session_id: return []
+#     query = select(ChatHistory).where(ChatHistory.session_id == session_id).order_by(ChatHistory.timestamp.asc())
+#     result = await db.execute(query)
+#     return result.scalars().all()
+
+# async def get_bot_persona(user_id: str, db: AsyncSession):
+#     """Fetches custom Bot Name and Instructions from User table."""
+#     try:
+#         # User ID ko int mein convert karke query karein
+#         stmt = select(User).where(User.id == int(user_id))
+#         result = await db.execute(stmt)
+#         user = result.scalars().first()
+        
+#         if user:
+#             return {
+#                 "name": getattr(user, "bot_name", "OmniAgent"),
+#                 "instruction": getattr(user, "bot_instruction", "You are a helpful AI assistant.")
+#             }
+#     except Exception as e:
+#         print(f"⚠️ Error fetching persona: {e}")
+#         pass
+    
+#     # Fallback Default Persona
+#     return {"name": "OmniAgent", "instruction": "You are a helpful AI assistant."}
+
+# # ==========================================
+# # MAIN CHAT LOGIC
+# # ==========================================
+# async def process_chat(message: str, session_id: str, user_id: str, db: AsyncSession):
+    
+#     # 1. Fetch User Settings & Persona
+#     user_settings = await get_user_integrations(user_id, db)
+#     bot_persona = await get_bot_persona(user_id, db) # <--- Persona Load kiya
+    
+#     # 2. LLM Check
+#     llm_creds = user_settings.get('groq') or user_settings.get('openai')
+#     if not llm_creds:
+#         return "Please configure your AI Model in Settings."
+
+#     # 3. Build Tool Map for Router
+#     tools_map = {}
+#     for provider, config in user_settings.items():
+#         if provider in ['sanity', 'sql', 'mongodb']:
+#             if config.get('description'):
+#                 tools_map[provider] = config['description']
+
+#     # 4. SEMANTIC DECISION (Router)
+#     selected_provider = None
+#     if tools_map:
+#         router = SemanticRouter() # Singleton Instance
+#         selected_provider = router.route(message, tools_map)
+    
+#     response_text = ""
+#     provider_name = "general_chat"
+
+#     # 5. Route to Winner
+#     if selected_provider:
+#         print(f"👉 [Router] Selected Tool: {selected_provider.upper()}")
+#         try:
+#             if selected_provider == 'sanity':
+#                 schema = user_settings['sanity'].get('schema_map', {})
+#                 agent = get_cms_agent(user_id=user_id, schema_map=schema, llm_credentials=llm_creds)
+#                 res = await agent.ainvoke({"input": message})
+#                 response_text = str(res.get('output', ''))
+#                 provider_name = "cms_agent"
+
+#             elif selected_provider == 'sql':
+#                 role = "admin" if user_id == '99' else "customer"
+#                 agent = get_secure_agent(int(user_id), role, user_settings['sql'], llm_credentials=llm_creds)
+#                 res = await agent.ainvoke({"input": message})
+#                 response_text = str(res.get('output', ''))
+#                 provider_name = "sql_agent"
+
+#             elif selected_provider == 'mongodb':
+#                 agent = get_nosql_agent(user_id, user_settings['mongodb'], llm_credentials=llm_creds)
+#                 res = await agent.ainvoke({"input": message})
+#                 response_text = str(res.get('output', ''))
+#                 provider_name = "nosql_agent"
+
+#             # Anti-Hallucination
+#             if not response_text or "error" in response_text.lower():
+#                 print(f"⚠️ [Router] Tool {selected_provider} failed. Triggering Fallback.")
+#                 response_text = "" 
+
+#         except Exception as e:
+#             print(f"❌ [Router] Execution Failed: {e}")
+#             response_text = "" 
+
+#     # 6. Fallback / RAG (Using Custom Persona)
+#     if not response_text:
+#         print("👉 [Router] Fallback to RAG/General Chat...")
+#         try:
+#             llm = get_llm_model(credentials=llm_creds)
+            
+#             # Context from Vector DB
+#             context = ""
+#             if 'qdrant' in user_settings:
+#                 try:
+#                     vector_store = get_vector_store(credentials=user_settings['qdrant'])
+#                     docs = await vector_store.asimilarity_search(message, k=3)
+#                     if docs:
+#                         context = "\n\n".join([d.page_content for d in docs])
+#                 except Exception as e:
+#                     print(f"⚠️ RAG Warning: {e}")
+            
+#             # --- 🔥 DYNAMIC SYSTEM PROMPT ---
+#             system_instruction = f"""
+#             IDENTITY: You are '{bot_persona['name']}'.
+#             MISSION: {bot_persona['instruction']}
+            
+#             CONTEXT FROM KNOWLEDGE BASE:
+#             {context if context else "No specific documents found."}
+            
+#             Answer the user's question based on the context above or your general knowledge if permitted by your mission.
+#             """
+
+#             # History Load
+#             history = await get_chat_history(session_id, db)
+#             formatted_history = []
+#             for chat in history:
+#                 formatted_history.append(HumanMessage(content=chat.human_message))
+#                 if chat.ai_message: formatted_history.append(AIMessage(content=chat.ai_message))
+
+#             # LLM Call
+#             prompt = ChatPromptTemplate.from_messages([
+#                 ("system", system_instruction),
+#                 MessagesPlaceholder(variable_name="chat_history"),
+#                 ("human", "{question}")
+#             ])
+#             chain = prompt | llm
+            
+#             ai_response = await chain.ainvoke({"chat_history": formatted_history, "question": message})
+#             response_text = ai_response.content
+#             provider_name = "rag_fallback"
+
+#         except Exception as e:
+#             print(f"❌ Fallback Error: {e}")
+#             response_text = "I am currently unable to process your request. Please check your AI configuration."
+
+#     # 7. Save to DB
+#     await save_chat_to_db(db, session_id, message, response_text, provider_name)
+#     return response_text
 import json
 from sqlalchemy.ext.asyncio import AsyncSession
 from sqlalchemy.future import select
@@ -6,7 +215,7 @@ from sqlalchemy.future import select
 # --- Model Imports ---
 from backend.src.models.chat import ChatHistory
 from backend.src.models.integration import UserIntegration
-from backend.src.models.user import User  # Added User model for Bot Persona
+from backend.src.models.user import User 
 
 # --- Dynamic Factory & Tool Imports ---
 from backend.src.services.llm.factory import get_llm_model
@@ -44,7 +253,6 @@ async def get_user_integrations(user_id: str, db: AsyncSession) -> dict:
             creds['provider'] = i.provider
             creds['schema_map'] = i.schema_map if i.schema_map else {}
             
-            # --- STRICT CHECK ---
             if i.profile_description:
                 creds['description'] = i.profile_description
             
@@ -74,7 +282,6 @@ async def get_chat_history(session_id: str, db: AsyncSession):
 async def get_bot_persona(user_id: str, db: AsyncSession):
     """Fetches custom Bot Name and Instructions from User table."""
     try:
-        # User ID ko int mein convert karke query karein
         stmt = select(User).where(User.id == int(user_id))
         result = await db.execute(stmt)
         user = result.scalars().first()
@@ -88,17 +295,16 @@ async def get_bot_persona(user_id: str, db: AsyncSession):
         print(f"⚠️ Error fetching persona: {e}")
         pass
     
-    # Fallback Default Persona
     return {"name": "OmniAgent", "instruction": "You are a helpful AI assistant."}
 
 # ==========================================
-# MAIN CHAT LOGIC
+# MAIN CHAT LOGIC (Ultra-Strict Isolated Mode)
 # ==========================================
 async def process_chat(message: str, session_id: str, user_id: str, db: AsyncSession):
     
     # 1. Fetch User Settings & Persona
     user_settings = await get_user_integrations(user_id, db)
-    bot_persona = await get_bot_persona(user_id, db) # <--- Persona Load kiya
+    bot_persona = await get_bot_persona(user_id, db)
     
     # 2. LLM Check
     llm_creds = user_settings.get('groq') or user_settings.get('openai')
@@ -115,13 +321,13 @@ async def process_chat(message: str, session_id: str, user_id: str, db: AsyncSes
     # 4. SEMANTIC DECISION (Router)
     selected_provider = None
     if tools_map:
-        router = SemanticRouter() # Singleton Instance
+        router = SemanticRouter() 
         selected_provider = router.route(message, tools_map)
     
     response_text = ""
     provider_name = "general_chat"
 
-    # 5. Route to Winner
+    # 5. Route to Winner (Agent Execution)
     if selected_provider:
         print(f"👉 [Router] Selected Tool: {selected_provider.upper()}")
         try:
@@ -145,18 +351,16 @@ async def process_chat(message: str, session_id: str, user_id: str, db: AsyncSes
                 response_text = str(res.get('output', ''))
                 provider_name = "nosql_agent"
 
-            # Anti-Hallucination
             if not response_text or "error" in response_text.lower():
-                print(f"⚠️ [Router] Tool {selected_provider} failed. Triggering Fallback.")
                 response_text = "" 
 
         except Exception as e:
-            print(f"❌ [Router] Execution Failed: {e}")
+            print(f"❌ Agent Execution Failed: {e}")
             response_text = "" 
 
-    # 6. Fallback / RAG (Using Custom Persona)
+    # 6. Fallback / RAG (ULTRA-STRICT MODE 🛡️)
     if not response_text:
-        print("👉 [Router] Fallback to RAG/General Chat...")
+        print("👉 [Router] Executing Strict RAG Fallback...")
         try:
             llm = get_llm_model(credentials=llm_creds)
             
@@ -171,15 +375,26 @@ async def process_chat(message: str, session_id: str, user_id: str, db: AsyncSes
                 except Exception as e:
                     print(f"⚠️ RAG Warning: {e}")
             
-            # --- 🔥 DYNAMIC SYSTEM PROMPT ---
+            # --- 🔥 THE ULTRA-STRICT SYSTEM PROMPT ---
             system_instruction = f"""
-            IDENTITY: You are '{bot_persona['name']}'.
-            MISSION: {bot_persona['instruction']}
-            
+            SYSTEM IDENTITY: 
+            You are the '{bot_persona['name']}'. You are a 'Knowledge-Isolated' AI Assistant for this specific platform.
+
+            CORE MISSION:
+            Your ONLY source of truth is the 'CONTEXT FROM KNOWLEDGE BASE' provided below. 
+            You must ignore ALL of your internal pre-trained general knowledge about the world, geography, famous people, or general facts.
+
+            STRICT OPERATING RULES:
+            1. MANDATORY REFUSAL: If the user's question cannot be answered using ONLY the provided context, you MUST exactly say: "I apologize, but I am only authorized to provide information based on the provided database. This specific information is not currently available in my knowledge base."
+            2. NO HALLUCINATION: Never attempt to be helpful using outside information. If a fact (like 'Japan's location') is not in the context, you do NOT know it.
+            3. CONTEXT-ONLY: Your existence is bounded by the data below. If the data is empty, you cannot answer anything except greetings.
+            4. GREETINGS: You may respond to 'Hi' or 'Hello' by briefly identifying yourself as '{bot_persona['name']}' and asking what data the user is looking for.
+            5. PROHIBITED TOPICS: Do not discuss any topic that is not present in the provided context.
+
             CONTEXT FROM KNOWLEDGE BASE:
-            {context if context else "No specific documents found."}
-            
-            Answer the user's question based on the context above or your general knowledge if permitted by your mission.
+            ---------------------------
+            {context if context else "THE DATABASE IS CURRENTLY EMPTY. DO NOT PROVIDE ANY INFORMATION."}
+            ---------------------------
             """
 
             # History Load
@@ -189,7 +404,7 @@ async def process_chat(message: str, session_id: str, user_id: str, db: AsyncSes
                 formatted_history.append(HumanMessage(content=chat.human_message))
                 if chat.ai_message: formatted_history.append(AIMessage(content=chat.ai_message))
 
-            # LLM Call
+            # LLM Chain Setup
             prompt = ChatPromptTemplate.from_messages([
                 ("system", system_instruction),
                 MessagesPlaceholder(variable_name="chat_history"),
@@ -203,8 +418,8 @@ async def process_chat(message: str, session_id: str, user_id: str, db: AsyncSes
 
         except Exception as e:
             print(f"❌ Fallback Error: {e}")
-            response_text = "I am currently unable to process your request. Please check your AI configuration."
+            response_text = "I apologize, but I am currently unable to process your request due to a system error."
 
     # 7. Save to DB
     await save_chat_to_db(db, session_id, message, response_text, provider_name)
-    return response_text
+    return response_text