Update app.py

app.py

@@ -2,30 +2,40 @@ import os
 import json
 import random
 import logging
-from typing import List
+from typing import List, Dict, Any
 import streamlit as st
 from dotenv import load_dotenv
+
+import torch
+import faiss
+import numpy as np
+from sentence_transformers import SentenceTransformer
+from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig
+
 import autogen
 from autogen import AssistantAgent, UserProxyAgent
 
+from langchain.docstore.document import Document
 from langchain_community.embeddings import HuggingFaceEmbeddings
 from langchain_community.vectorstores import Chroma
-from langchain.docstore.document import Document
 
+# --- Load environment ---
 load_dotenv()
 
+# --- Logging ---
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
-st.set_page_config(page_title="IT Support System (RAG)", layout="centered")
 
-# Initialize session state for chat and logs
+st.set_page_config(page_title="AI Help Desk with Hybrid Open Source + Multi-Agent", layout="centered")
+
+# --- Initialize session state ---
 if "chat_history" not in st.session_state:
     st.session_state.chat_history = []
 
 if "workflow_logs" not in st.session_state:
     st.session_state.workflow_logs = []
 
-# Knowledge Base Setup
+# --- Load Knowledge Base ---
 kb_path = os.path.join(os.path.dirname(__file__), 'kb.json')
 with open(kb_path, encoding='utf-8') as f:
     kb_entries = json.load(f)
@@ -34,10 +44,7 @@ docs: List[Document] = []
 for entry in kb_entries:
     docs.append(Document(
         page_content=entry['answer'],
-        metadata={
-            'id': entry.get('id'),
-            'question': entry.get('question')
-        }
+        metadata={'id': entry.get('id'), 'question': entry.get('question')}
     ))
 
 embeddings = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2")
@@ -64,6 +71,7 @@ def escalate_ticket(query: str, analysis: str = "") -> str:
     logger.info(f"Escalating issue with ticket {ticket_id}: {description}")
     return f"Escalated issue. Created ticket {ticket_id}. A support technician will contact you shortly."
 
+# --- OpenAI Multi-agent LLM Config ---
 llm_config = {
     "config_list": [{
         "model": "llama3",
@@ -74,10 +82,8 @@ llm_config = {
     "temperature": 0.5,
 }
 
-master_agent = AssistantAgent(
-    name="Master",
-    llm_config=llm_config,
-    system_message="""
+# --- Define Agents ---
+master_agent = AssistantAgent(name="Master", llm_config=llm_config, system_message="""
 You are the Master Agent that orchestrates the IT support workflow:
 1. First determine if the query is IT-related. If not, provide a direct response explaining your limitations.
 2. For IT-related queries, pass to the Planning Agent for execution plan development
@@ -85,13 +91,9 @@ You are the Master Agent that orchestrates the IT support workflow:
 4. Provide a comprehensive yet concise final response to the user
 
 Only handle one query at a time through the complete workflow.
-"""
-)
+""")
 
-planning_agent = AssistantAgent(
-    name="Planning",
-    llm_config=llm_config,
-    system_message="""
+planning_agent = AssistantAgent(name="Planning", llm_config=llm_config, system_message="""
 You are the Planning Agent responsible for:
 1. Validating if the user query is clear and complete
 2. Refining the query if needed for better processing
@@ -104,13 +106,9 @@ Provide your analysis as a structured output with sections for:
 - Execution Plan
 
 Always end your message with: "Forwarding to Analysis Agent"
-"""
-)
+""")
 
-analysis_agent = AssistantAgent(
-    name="Analysis",
-    llm_config=llm_config,
-    system_message="""
+analysis_agent = AssistantAgent(name="Analysis", llm_config=llm_config, system_message="""
 You are the Analysis Agent responsible for:
 1. Identifying key entities in the user query (devices, software, errors, etc.)
 2. Determining severity level (Low, Medium, High, Critical)
@@ -124,8 +122,7 @@ Provide your analysis as structured output with sections for:
 - Analysis Summary
 
 Always end your message with: "Forwarding to Resolution Agent"
-"""
-)
+""")
 
 resolution_agent = AssistantAgent(
     name="Resolution",
@@ -166,7 +163,59 @@ Always include the ticket ID and expected follow-up timeframe.
     function_map={"escalate_ticket": escalate_ticket}
 )
 
-def handle_it_query(query: str) -> (str, dict):
+# --- Initialize open-source fallback Mistral-7B model for AI fallback ---
+quantization_config = BitsAndBytesConfig(
+    load_in_4bit=True,
+    bnb_4bit_compute_dtype=torch.float16,
+    bnb_4bit_use_double_quant=True,
+    bnb_4bit_quant_type="nf4"
+)
+model_name = "mistralai/Mistral-7B-Instruct-v0.1"
+device = "cuda" if torch.cuda.is_available() else "cpu"
+tokenizer = AutoTokenizer.from_pretrained(model_name)
+mistral_model = AutoModelForCausalLM.from_pretrained(
+    model_name,
+    quantization_config=quantization_config,
+    device_map="auto"
+)
+
+# --- Load FAQ Data for fallback ---
+faq_data = [
+    # Populate with your actual FAQ data or import from file if needed
+    {"question": "What is Onfon Mobile?", "answer": "Onfon Mobile is a device financing company..."},
+    {"question": "How do I get a phone via Onfon Mobile?", "answer": "To get a smartphone, register with *797#..."},
+    # ... Add as needed ...
+]
+
+embed_model = SentenceTransformer("BAAI/bge-large-en")
+faq_embeddings = np.array([embed_model.encode(f["question"], convert_to_numpy=True) for f in faq_data])
+faiss_index = faiss.IndexFlatL2(faq_embeddings.shape[1])
+faiss_index.add(faq_embeddings)
+
+def find_top_faqs(user_input: str, top_k=3):
+    emb = embed_model.encode(user_input, convert_to_numpy=True).reshape(1, -1)
+    _, idxs = faiss_index.search(emb, top_k)
+    return [faq_data[i] for i in idxs[0]]
+
+def generate_ai_faq_answer(user_input: str):
+    top_faqs = find_top_faqs(user_input, 3)
+    system_prompt = (
+        "You are an Onfon Mobile customer support assistant. Be kind and helpful to onfon mobile customers. "
+        "From the following FAQs, provide the best possible answer briefly:\n\n"
+    )
+    faq_context = "\n".join([f"Q: {f['question']}\nA: {f['answer']}" for f in top_faqs])
+    prompt = f"{system_prompt}{faq_context}\n\nUser Question: {user_input}\nAnswer:"
+
+    inputs = tokenizer(prompt, return_tensors="pt", truncation=True).to(device)
+    with torch.no_grad():
+        outputs = mistral_model.generate(**inputs, max_new_tokens=150)
+    answer = tokenizer.decode(outputs[0], skip_special_tokens=True)
+    if "Answer:" in answer:
+        answer = answer.split("Answer:")[-1].strip()
+    return answer
+
+# --- Main Query Handler ---
+def handle_it_query(query: str) -> (str, Dict[str, Any]):
     query = query.strip()
     if not query:
         return "Please enter an IT question or issue.", {}
@@ -236,15 +285,25 @@ def handle_it_query(query: str) -> (str, dict):
         logger.info(f"Final Master Agent Response: {final_response}")
         workflow_logs["final_response"] = final_response
 
+        # Fallback to open-source model if no valid response or error
+        if not final_response or final_response.lower().startswith("an error occurred"):
+            logger.warning("Using fallback open-source Mistral-7B FAQ model due to empty or error response")
+            fallback_response = generate_ai_faq_answer(query)
+            workflow_logs["fallback_response"] = fallback_response
+            final_response = fallback_response
+
         return final_response, workflow_logs
 
     except Exception as e:
         logger.error(f"Error in workflow: {e}", exc_info=True)
-        return f"An error occurred during processing: {str(e)}\n\nPlease try rephrasing your question.", {}
+        fallback_response = generate_ai_faq_answer(query)
+        workflow_logs["fallback_response"] = fallback_response
+        return f"An error occurred during processing, fallback answer provided.\n\n{fallback_response}", workflow_logs
 
+# --- Streamlit UI ---
+st.title("AI Help Desk (Hybrid Multi-Agent + Open Source)")
 
-st.title("AI Help Desk")
-st.write("Ask any IT support question and our multi-agent system will assist you.")
+st.write("Ask any IT support question and our hybrid multi-agent + open-source system will assist you.")
 
 with st.form(key="query_form", clear_on_submit=True):
     user_input = st.text_area("Describe your IT issue:", height=100)
@@ -258,19 +317,19 @@ if submitted:
         with st.spinner("Processing your request through our agent workflow..."):
             response, logs = handle_it_query(user_input)
 
-        # Append to session history to maintain conversation flow
+        # Append conversation to history
         st.session_state.chat_history.append({"user": user_input, "assistant": response})
         if show_logs:
             st.session_state.workflow_logs.append(logs)
 
-# Display conversation history
+# --- Show conversation ---
 st.markdown("## Conversation")
 for chat in st.session_state.chat_history:
     st.markdown(f"**User:** {chat['user']}")
     st.markdown(f"**Assistant:** {chat['assistant']}")
     st.markdown("---")
 
-# Display logs if requested
+# --- Show logs if requested ---
 if show_logs and st.session_state.workflow_logs:
     st.markdown("## Workflow Logs")
     for i, log in enumerate(st.session_state.workflow_logs):