Update src/streamlit_app.py

src/streamlit_app.py

@@ -1,40 +1,86 @@
-import altair as alt
-import numpy as np
-import pandas as pd
 import streamlit as st
+import torch
+import faiss
+import numpy as np
+from sentence_transformers import SentenceTransformer
+from transformers import AutoTokenizer, AutoModelForCausalLM
+from transformers import DistilBertTokenizerFast, DistilBertForSequenceClassification
+import pickle
 
-"""
-# Welcome to Streamlit!
-
-Edit `/streamlit_app.py` to customize this app to your heart's desire :heart:.
-If you have any questions, checkout our [documentation](https://docs.streamlit.io) and [community
-forums](https://discuss.streamlit.io).
+st.set_page_config(page_title="AutoResolve Agent", page_icon="🤖", layout="centered")
 
-In the meantime, below is an example of what you can do with just a few lines of code:
-"""
+st.title("🤖 AutoResolve: IT Support Agent")
+st.markdown("This end-to-end LLM Agent classifies your IT issue, retrieves the relevant enterprise policy, and generates a solution.")
 
-num_points = st.slider("Number of points in spiral", 1, 10000, 1100)
-num_turns = st.slider("Number of turns in spiral", 1, 300, 31)
+# --- 1. Load Models (Cached so they only load once) ---
+@st.cache_resource
+def load_pipeline():
+    # 1. Load DistilBERT Classifier
+    # Note: You must upload your 'autoresolve_distilbert_final' folder to the HF space!
+    distil_tokenizer = DistilBertTokenizerFast.from_pretrained("./autoresolve_distilbert_final")
+    distil_model = DistilBertForSequenceClassification.from_pretrained("./autoresolve_distilbert_final")
+    
+    # 2. Knowledge Base & Retriever
+    kb = [
+        "Refund Policy: Customers are entitled to a full refund within 30 days of purchase. To process, verify the order number and issue the refund to the original payment method.",
+        "Order Tracking: To locate an order, query the shipping database using the 10-digit order number. If the status is 'Dispatched', provide the user with the carrier tracking link.",
+        "Password Recovery: If a user cannot log in, send a secure password reset link to their registered email address. Ensure they check their spam folder.",
+        "Payment Issues: If a transfer or payment fails, verify if the credit card is expired or if the anti-fraud system flagged the transaction. Recommend trying a different payment method."
+    ]
+    embedder = SentenceTransformer('all-MiniLM-L6-v2')
+    kb_embeddings = embedder.encode(kb, convert_to_numpy=True)
+    index = faiss.IndexFlatL2(kb_embeddings.shape[1])
+    index.add(kb_embeddings)
+    
+    # 3. Load Generative LLM (CPU mode)
+    llama_tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen2.5-0.5B-Instruct")
+    llama_model = AutoModelForCausalLM.from_pretrained("Qwen/Qwen2.5-0.5B-Instruct", device_map="cpu")
+    
+    return distil_tokenizer, distil_model, kb, embedder, index, llama_tokenizer, llama_model
 
-indices = np.linspace(0, 1, num_points)
-theta = 2 * np.pi * num_turns * indices
-radius = indices
+with st.spinner("Loading AI Models... (This takes about 60 seconds on initial boot)"):
+    distil_tokenizer, distil_model, knowledge_base, embedder, index, llama_tokenizer, llama_model = load_pipeline()
 
-x = radius * np.cos(theta)
-y = radius * np.sin(theta)
+# Define the intents manually to avoid needing the full dataset for the LabelEncoder
+INTENTS = ['cancel_order', 'change_order', 'change_shipping_address', 'check_cancellation_fee', 'check_invoice', 'check_payment_methods', 'check_refund_policy', 'complaint', 'contact_customer_service', 'contact_human_agent', 'create_account', 'delete_account', 'delivery_options', 'delivery_period', 'edit_account', 'get_invoice', 'get_refund', 'newsletter_subscription', 'payment_issue', 'place_order', 'recover_password', 'registration_problems', 'review', 'set_up_shipping_address', 'switch_account', 'track_order', 'track_refund']
 
-df = pd.DataFrame({
-    "x": x,
-    "y": y,
-    "idx": indices,
-    "rand": np.random.randn(num_points),
-})
+# --- 2. The User Interface ---
+user_query = st.text_input("Describe your IT or Support issue:", placeholder="e.g., am I entitled to a reimbursement?")
 
-st.altair_chart(alt.Chart(df, height=700, width=700)
-    .mark_point(filled=True)
-    .encode(
-        x=alt.X("x", axis=None),
-        y=alt.Y("y", axis=None),
-        color=alt.Color("idx", legend=None, scale=alt.Scale()),
-        size=alt.Size("rand", legend=None, scale=alt.Scale(range=[1, 150])),
-    ))
+if st.button("Submit Ticket"):
+    if user_query:
+        with st.spinner("Processing..."):
+            # Step A: Intent Classification
+            inputs = distil_tokenizer(user_query, return_tensors="pt", truncation=True, padding=True)
+            with torch.no_grad():
+                logits = distil_model(**inputs).logits
+            predicted_class_id = logits.argmax().item()
+            predicted_intent = INTENTS[predicted_class_id]
+            
+            st.success(f"**Intent Classified:** `{predicted_intent}`")
+            
+            # Step B: Retrieval
+            query_vector = embedder.encode([user_query], convert_to_numpy=True)
+            distances, indices = index.search(query_vector, 1)
+            retrieved_doc = knowledge_base[indices[0][0]]
+            
+            st.info(f"**Retrieved Knowledge Base Document:** {retrieved_doc}")
+            
+            # Step C: Generation
+            prompt = f"""<|im_start|>system
+You are AutoResolve, an IT support agent. Answer the user's query using ONLY the provided IT Document. Be polite, concise, and professional.<|im_end|>
+<|im_start|>user
+User Query: {user_query}
+IT Document: {retrieved_doc}<|im_end|>
+<|im_start|>assistant
+"""
+            gen_inputs = llama_tokenizer(prompt, return_tensors="pt")
+            outputs = llama_model.generate(**gen_inputs, max_new_tokens=150, temperature=0.1, pad_token_id=llama_tokenizer.eos_token_id)
+            
+            full_response = llama_tokenizer.decode(outputs[0], skip_special_tokens=True)
+            final_answer = full_response.split("assistant\n")[-1].strip()
+            
+            st.write("### 💬 AutoResolve Agent Response:")
+            st.write(f"> {final_answer}")
+    else:
+        st.warning("Please enter a query first.")