Update app.py

revert to only using groq

app.py

@@ -2,319 +2,338 @@ import streamlit as st
 import pandas as pd
 import os
 import warnings
+import time
 import sqlite3
 import shutil
-import asyncio
-import re
 
 # ==========================================
-# 0. ASYNC FIX
+# 1. PAGE CONFIG (MUST BE FIRST)
 # ==========================================
-try:
-    asyncio.get_running_loop()
-except RuntimeError:
-    asyncio.set_event_loop(asyncio.new_event_loop())
+st.set_page_config(page_title="Bank Loan Agent (SQL)", layout="wide")
 
-# ==========================================
-# 1. CONFIG & IMPORTS
-# ==========================================
-st.set_page_config(page_title="Bank Loan Agent", layout="wide")
+# Suppress warnings
 warnings.filterwarnings("ignore")
 
+# ==========================================
+# 2. GLOBAL CONSTANTS & IMPORTS
+# ==========================================
 DB_FILE = "bank.db"
 INDEX_PATH = "faiss_index"
 REQUIRED_PDFS = ["Bank Loan Overall Risk Policy.pdf", "Bank Loan Interest Rate Policy.pdf"]
 
 try:
-    # GROQ
     from langchain_groq import ChatGroq
-    
-    # GOOGLE (Native SDK)
-    import google.generativeai as genai
-    from google.generativeai.types import HarmCategory, HarmBlockThreshold
-    
-    # SHARED
     from langchain_huggingface import HuggingFaceEmbeddings
     from langchain_community.vectorstores import FAISS
+    from langchain_community.callbacks import StreamlitCallbackHandler
     from langchain_community.document_loaders import PyPDFLoader
     from langchain_text_splitters import CharacterTextSplitter
-    from langchain_core.prompts import PromptTemplate
+    from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
     from langchain_core.runnables import RunnablePassthrough
     from langchain_core.output_parsers import StrOutputParser
+    from langchain_core.tools import tool
+    from langchain.agents import AgentExecutor, create_tool_calling_agent
+
 except ImportError as e:
-    st.error(f"❌ Import Error: {e}")
+    st.error(f"❌ Critical Import Error: {e}")
     st.stop()
 
 # ==========================================
-# 2. DATABASE & TOOLS
+# 3. DATABASE SETUP
 # ==========================================
 def init_db():
-    if os.path.exists(DB_FILE): return
+    """Converts CSV files to SQLite DB. Handles errors gracefully."""
+    if os.path.exists(DB_FILE):
+        return
+
     conn = sqlite3.connect(DB_FILE)
-    csv_files = {"credit_score": "credit_score.csv", "account_status": "account_status.csv", "pr_status": "pr_status.csv"}
+    csv_files = {
+        "credit_score": "credit_score.csv",
+        "account_status": "account_status.csv",
+        "pr_status": "pr_status.csv"
+    }
+
     try:
         for table, file in csv_files.items():
             if os.path.exists(file):
                 df = pd.read_csv(file)
                 df.columns = [c.strip() for c in df.columns] 
-                if 'ID' in df.columns: df['ID'] = df['ID'].astype(str)
-                try: df.to_sql(table, conn, if_exists='replace', index=False)
-                except: pass 
-    finally: conn.close()
+                if 'ID' in df.columns:
+                    df['ID'] = df['ID'].astype(str)
+                
+                try:
+                    df.to_sql(table, conn, if_exists='replace', index=False)
+                except Exception:
+                    pass 
+    except Exception as e:
+        st.error(f"DB Init Error: {e}")
+    finally:
+        conn.close()
+
+# Initialize DB on startup
 init_db()
 
+# Helper for SQL tools
 def run_query(query, params=()):
     try:
         with sqlite3.connect(DB_FILE) as conn:
             cursor = conn.cursor()
             cursor.execute(query, params)
             return cursor.fetchone()
-    except Exception as e: return f"DB Error: {e}"
+    except Exception as e:
+        return f"DB Error: {e}"
 
-# --- TOOLS ---
-def tool_get_credit_score(user_id):
+# ==========================================
+# 4. DEFINE TOOLS
+# ==========================================
+
+@tool
+def get_credit_score(user_id: str) -> str:
+    """Queries SQL DB for Credit Score."""
     clean_id = ''.join(filter(str.isdigit, str(user_id)))
     row = run_query("SELECT Credit_Score FROM credit_score WHERE ID = ?", (clean_id,))
-    return f"Credit Score: {row[0]}" if (row and not isinstance(row, str)) else "User ID not found."
+    if row and not isinstance(row, str):
+        return f"Credit Score: {row[0]}"
+    return "User ID not found in Credit DB."
 
-def tool_get_account_status(user_id):
+@tool
+def get_account_status(user_id: str) -> str:
+    """Queries SQL DB for Name, Nationality, Status, and Email."""
     clean_id = ''.join(filter(str.isdigit, str(user_id)))
-    row = run_query("SELECT Name, Nationality, Account_Status, Email FROM account_status WHERE ID = ?", (clean_id,))
+    row = run_query(
+        "SELECT Name, Nationality, Account_Status, Email FROM account_status WHERE ID = ?", 
+        (clean_id,)
+    )
     if row and not isinstance(row, str):
         return f"Customer Name: {row[0]}, Nationality: {row[1]}, Status: {row[2]}, Email: {row[3]}"
-    return "User ID not found."
+    return "User ID not found in Account DB."
 
-def tool_check_pr_status(user_id):
+@tool
+def check_pr_status(user_id: str) -> str:
+    """Queries SQL DB for PR Status."""
     clean_id = ''.join(filter(str.isdigit, str(user_id)))
     row = run_query("SELECT PR_Status FROM pr_status WHERE ID = ?", (clean_id,))
+    
     if not row or (isinstance(row, str) and "no such column" in row.lower()):
          row = run_query("SELECT Is_PR FROM pr_status WHERE ID = ?", (clean_id,))
-    return f"PR Status: {row[0]}" if (row and not isinstance(row, str)) else "PR Status: False."
-
-# ==========================================
-# 3. HYBRID AGENT (Dynamic Model Loader)
-# ==========================================
-class HybridAgent:
-    def __init__(self, provider, api_key, tools_map, rag_chain):
-        self.provider = provider
-        self.api_key = api_key
-        self.tools = tools_map
-        self.rag_chain = rag_chain
-        self.max_steps = 8
-        self.gemini_model = None
-        
-        # Initialize Groq
-        if "Groq" in provider:
-            self.groq_chat = ChatGroq(api_key=api_key, model_name="llama-3.3-70b-versatile", temperature=0)
-        
-        # Initialize Gemini with DYNAMIC DISCOVERY
-        if "Google" in provider:
-            genai.configure(api_key=api_key)
-            self.gemini_model = self._find_best_gemini_model()
-
-    def _find_best_gemini_model(self):
-        """Auto-detects which Gemini model is actually available to avoid 404s."""
-        try:
-            available_models = [m.name for m in genai.list_models() if 'generateContent' in m.supported_generation_methods]
-            
-            # Priority 1: Flash (Fastest)
-            for m in available_models:
-                if "flash" in m and "1.5" in m: return genai.GenerativeModel(m)
-            
-            # Priority 2: Pro 1.5
-            for m in available_models:
-                if "pro" in m and "1.5" in m: return genai.GenerativeModel(m)
-            
-            # Priority 3: Pro 1.0 / Standard
-            for m in available_models:
-                if "gemini-pro" in m: return genai.GenerativeModel(m)
-                
-            # Fallback: Just take the first one
-            if available_models: return genai.GenerativeModel(available_models[0])
-            
-            return genai.GenerativeModel('gemini-pro') # Blind hope
-        except:
-            return genai.GenerativeModel('gemini-1.5-flash') # Default
-
-    def call_llm(self, prompt):
-        if "Groq" in self.provider:
-            return self.groq_chat.invoke(prompt).content
-        else:
-            try:
-                # Disable safety to prevent "list index out of range" errors
-                safety = {
-                    HarmCategory.HARM_CATEGORY_DANGEROUS_CONTENT: HarmBlockThreshold.BLOCK_NONE,
-                    HarmCategory.HARM_CATEGORY_HATE_SPEECH: HarmBlockThreshold.BLOCK_NONE,
-                    HarmCategory.HARM_CATEGORY_HARASSMENT: HarmBlockThreshold.BLOCK_NONE,
-                    HarmCategory.HARM_CATEGORY_SEXUALLY_EXPLICIT: HarmBlockThreshold.BLOCK_NONE,
-                }
-                response = self.gemini_model.generate_content(prompt, safety_settings=safety)
-                return response.text
-            except Exception as e:
-                return f"Gemini Error: {str(e)}"
-
-    def run(self, query):
-        tool_desc = "\n".join([f"- {name}: {func.__doc__}" for name, func in self.tools.items()])
-        
-        history = f"""You are a Loan Officer. Solve this request: "{query}"
-
-TOOLS AVAILABLE:
-{tool_desc}
-- consult_policy_doc: Search PDF policies. Input: question string.
-
-RULES:
-1. You run in a loop. OUTPUT ONLY ONE STEP AT A TIME.
-2. Format:
-Thought: <reasoning>
-Action: <tool_name>
-Action Input: <input_string>
-Observation: <result>
-...
-Final Answer: <the full report>
-
-Begin!
-"""
-        logs = []
-
-        for i in range(self.max_steps):
-            response = self.call_llm(history)
-            history += response + "\n"
-            
-            if "Final Answer:" in response:
-                return response.split("Final Answer:")[-1].strip(), logs
-
-            action_match = re.search(r"Action:\s*(.+)", response)
-            input_match = re.search(r"Action Input:\s*(.+)", response)
-
-            if action_match and input_match:
-                tool_name = action_match.group(1).strip()
-                val = input_match.group(1).strip().strip('"').strip("'")
-                
-                logs.append((tool_name, val))
-                
-                result = "Error: Tool not found"
-                if tool_name in self.tools:
-                    try: result = self.tools[tool_name](val)
-                    except Exception as e: result = f"Error: {e}"
-                elif tool_name == "consult_policy_doc":
-                    try: result = self.rag_chain.invoke(val)
-                    except Exception as e: result = f"RAG Error: {e}"
-                
-                history += f"Observation: {result}\n"
-            else:
-                if i == self.max_steps - 1: return response, logs
-                history += "Observation: Please continue. Use 'Final Answer:' when done.\n"
 
-        return "Agent timed out.", logs
+    if row and not isinstance(row, str):
+        return f"PR Status: {row[0]}"
+    return "PR Status: False (Record not found)"
 
 # ==========================================
-# 4. UI & LOGIC
+# 5. STREAMLIT APP UI
 # ==========================================
-st.title("🤖 Multi-Model Loan Assessor")
+st.title("🤖 Multi-Policy Loan Assessor (SQL + RAG)")
+st.markdown("Agent connects to **SQLite Database** and **Persistent Vector Store**")
+
+# Calculate missing PDFs globally so everyone can see it
 pdfs_missing = [f for f in REQUIRED_PDFS if not os.path.exists(f)]
 
+# --- METRICS FUNCTION ---
+def update_metrics(placeholder):
+    manual_time = 15 * 60  
+    if 'execution_time' in st.session_state:
+        ai_time = st.session_state.execution_time
+        time_saved = manual_time - ai_time
+        saved_pct = (time_saved / manual_time) * 100
+        with placeholder.container():
+            col_kpi1, col_kpi2 = st.columns(2)
+            col_kpi1.metric("AI Processing", f"{ai_time:.1f}s")
+            col_kpi2.metric("Time Saved", f"{time_saved/60:.1f} min", delta=f"{saved_pct:.1f}% faster")
+
+# --- SIDEBAR ---
 with st.sidebar:
     st.header("🔐 Authentication")
-    provider_opt = st.radio("Model:", ["Groq (Llama-3)", "Google (Gemini)"])
     
-    if 'auth' not in st.session_state: st.session_state.auth = False
-    
-    if st.session_state.get('last_provider') != provider_opt:
-        st.session_state.auth = False
-        st.session_state.last_provider = provider_opt
-
-    if not st.session_state.auth:
-        key_in = st.text_input("API Key", type="password")
-        if st.button("Validate"):
-            try:
-                if "Groq" in provider_opt:
-                    ChatGroq(api_key=key_in).invoke("Hi")
-                else:
-                    genai.configure(api_key=key_in)
-                    # Quick list check to validate key
-                    [m.name for m in genai.list_models()] 
-                
-                st.session_state.auth = True
-                st.session_state.key = key_in
-                st.success("Valid!")
-                st.rerun()
-            except Exception as e:
-                st.error(f"Invalid: {e}")
+    # Initialize Session State for Key
+    if 'is_key_valid' not in st.session_state:
+        st.session_state['is_key_valid'] = False
+
+    # MANUAL ENTRY ONLY (No Secret Check)
+    if not st.session_state['is_key_valid']:
+        api_key_input = st.text_input("Enter Groq API Key", type="password", key="input_key")
+        if st.button("Validate API Key"):
+            if not api_key_input:
+                st.error("⚠️ Please enter a key.")
+            else:
+                try:
+                    with st.spinner("Validating..."):
+                        test_llm = ChatGroq(api_key=api_key_input, model_name="llama-3.3-70b-versatile")
+                        test_llm.invoke("Test")
+                        st.session_state['groq_api_key'] = api_key_input
+                        st.session_state['is_key_valid'] = True
+                        st.success("✅ Valid Key!")
+                        time.sleep(0.5)
+                        st.rerun()
+                except Exception as e:
+                    st.error(f"❌ Invalid Key: {e}")
     else:
-        st.success("Active")
-        if st.button("Logout"): 
-            st.session_state.auth = False
+        st.success("✅ API Key Active")
+        if st.button("🔴 Reset Key"):
+            st.session_state['is_key_valid'] = False
+            st.session_state['groq_api_key'] = None
             st.rerun()
 
-    if st.button("♻️ Reset DB"):
-        if os.path.exists(INDEX_PATH): shutil.rmtree(INDEX_PATH)
+    st.divider()
+    st.subheader("🛠️ System Maintenance")
+    
+    if st.button("♻️ Rebuild Knowledge Base"):
+        if os.path.exists(INDEX_PATH):
+            shutil.rmtree(INDEX_PATH)
         st.cache_resource.clear()
+        st.success("Cache cleared.")
+        time.sleep(1)
         st.rerun()
 
-if st.session_state.auth:
+    if st.button("💾 Reload CSVs to DB"):
+        if os.path.exists(DB_FILE):
+            os.remove(DB_FILE)
+        init_db()
+        st.success("Database refreshed.")
+        
+    st.divider()
+    
+    if os.path.exists(DB_FILE) and not pdfs_missing:
+        st.success("✅ System Ready")
+    else:
+        st.warning(f"⚠️ Missing: {pdfs_missing}")
+
+    st.header("📊 Metrics")
+    metrics_placeholder = st.empty()
+    update_metrics(metrics_placeholder)
+
+# --- MAIN LOGIC ---
+if st.session_state.get('is_key_valid', False):
+    
+    os.environ["GROQ_API_KEY"] = st.session_state['groq_api_key']
+    
+    # --- RAG SETUP ---
     @st.cache_resource
     def setup_rag():
-        if pdfs_missing: return None
+        # Check global variable here
+        if pdfs_missing:
+            st.error(f"Missing PDFs: {pdfs_missing}")
+            st.stop()
+            
         embeddings = HuggingFaceEmbeddings(model_name="all-MiniLM-L6-v2")
+        
         if os.path.exists(INDEX_PATH):
             return FAISS.load_local(INDEX_PATH, embeddings, allow_dangerous_deserialization=True).as_retriever()
-        
-        docs = []
-        for f in REQUIRED_PDFS: docs.extend(PyPDFLoader(f).load())
-        splits = CharacterTextSplitter(chunk_size=600, chunk_overlap=50).split_documents(docs)
-        vectorstore = FAISS.from_documents(splits, embeddings)
-        vectorstore.save_local(INDEX_PATH)
-        return vectorstore.as_retriever()
-
-    retriever = setup_rag()
-
-    def query_rag(q):
-        if not retriever: return "No PDFs found."
-        docs = retriever.invoke(q)
-        return "Context: " + "\n".join([d.page_content for d in docs])
-
-    tools = {
-        "get_credit_score": tool_get_credit_score,
-        "get_account_status": tool_get_account_status,
-        "check_pr_status": tool_check_pr_status
-    }
+        else:
+            documents = []
+            for pdf_file in REQUIRED_PDFS:
+                loader = PyPDFLoader(pdf_file)
+                documents.extend(loader.load())
+            
+            text_splitter = CharacterTextSplitter(chunk_size=600, chunk_overlap=50)
+            final_docs = text_splitter.split_documents(documents)
+            
+            vectorstore = FAISS.from_documents(final_docs, embeddings)
+            vectorstore.save_local(INDEX_PATH)
+            return vectorstore.as_retriever()
+
+    with st.spinner("Initializing AI..."):
+        retriever = setup_rag()
+
+    llm = ChatGroq(temperature=0, model_name="llama-3.3-70b-versatile")
+    
+    rag_prompt = ChatPromptTemplate.from_template("Answer based on context:\n{context}\nQuestion: {question}")
+    rag_chain = (
+        {"context": retriever | (lambda d: "\n".join([x.page_content for x in d])), "question": RunnablePassthrough()}
+        | rag_prompt | llm | StrOutputParser()
+    )
+    
+    @tool
+    def consult_policy_doc(query: str) -> str:
+        """Consults Policy Documents for Risk Rules."""
+        return rag_chain.invoke(query)
+
+    tools = [get_credit_score, get_account_status, check_pr_status, consult_policy_doc]
+
+    prompt = ChatPromptTemplate.from_messages([
+        ("system", "You are a Loan Risk Officer. Query SQL DB for customer info. Consult Policy Docs for rules."),
+        ("human", "{input}"),
+        MessagesPlaceholder(variable_name="agent_scratchpad"),
+    ])
     
-    rag_lambda = type('RAG', (object,), {"invoke": lambda self, x: query_rag(x)})()
-    agent = HybridAgent(provider_opt, st.session_state.key, tools, rag_lambda)
+    agent = create_tool_calling_agent(llm, tools, prompt)
+    agent_executor = AgentExecutor(agent=agent, tools=tools, verbose=True, return_intermediate_steps=True)
 
     col1, col2 = st.columns([1, 2])
     with col1:
+        st.subheader("1. Customer Details")
         uid = st.text_input("Customer ID", "1111")
-        use_sim = st.checkbox("Simulation Mode")
-        s_score = st.slider("Score", 300, 900, 450) if use_sim else 0
-        s_status = st.selectbox("Status", ["good-standing", "closed", "delinquent"]) if use_sim else ""
-        btn = st.button("Assess")
-
+        use_simulation = st.checkbox("Simulation Mode")
+        
+        sim_score = 650
+        sim_status = "good-standing"
+        if use_simulation:
+            sim_score = st.slider("Sim Credit Score", 300, 900, 450)
+            sim_status = st.selectbox("Sim Status", ["good-standing", "closed", "delinquent"])
+
+        st.divider()
+        btn = st.button("Assess Loan Risk", type="primary")
+    
     with col2:
         if btn:
-            q = f"Process Loan ID {uid}. "
-            if use_sim: q += f"SIMULATION: Score {s_score}, Status '{s_status}'. Skip DB for those."
-            else: q += "Query DB for all data."
-            q += " Check Policy. Report Risk, Rate, Decision."
-
-            with st.status("Agent Working...", expanded=True):
+            if use_simulation:
+                 query = f"""
+                Process Loan for Customer ID: {uid}.
+                *** SIMULATION MODE ***
+                1. DO NOT query 'get_credit_score' or 'account_status' for Score/Status.
+                2. USE: Score: {sim_score}, Status: {sim_status}
+                3. Query 'get_account_status' ONLY for Name/Nationality.
+                4. Consult Policy Docs for risk/rates.
+                5. Provide a Final Recommendation Report that MUST include:
+                   - Customer Name, ID, Email
+                   - Risk Level, Interest Rate
+                   - Final Decision (Approve/Reject)
+                   - Justification for Decision (Cite specific PDF policies)
+                   - Format in a clear markdown table.
+                """
+            else:
+                query = f"""
+                Process Loan for Customer ID: {uid}.
+                1. Query SQL tools for Name, Email, Nationality, Status, Score.
+                2. IF Nationality is 'Singaporean', SKIP 'check_pr_status'.
+                3. Consult Policy Docs for risk/rates.
+                4.  Provide a Final Recommendation Report that MUST include:
+                   - Customer Name, ID, Email
+                   - Risk Level, Interest Rate
+                   - Final Decision (Approve/Reject)
+                   - Justification for Decision (Cite specific PDF policies)
+                   - Format in a clear markdown table.
+                """
+            
+            with st.status("🤖 Agent is processing...", expanded=True) as status:
+                st_callback = StreamlitCallbackHandler(st.container())
                 try:
-                    ans, logs = agent.run(q)
-                    st.write("Done!")
+                    start_time = time.time()
+                    res = agent_executor.invoke({"input": query}, {"callbacks": [st_callback]})
+                    end_time = time.time()
+                    st.session_state.execution_time = end_time - start_time
+                    update_metrics(metrics_placeholder)
+                    status.update(label="✅ Complete!", state="complete", expanded=False)
                 except Exception as e:
-                    st.error(f"Execution Error: {e}")
-                    ans, logs = "Error", []
+                    st.error(f"Error: {e}")
+                    st.stop()
 
-            st.success("### Final Report")
-            st.markdown(ans)
-            
-            with st.expander("Trace"):
-                for t, i in logs: st.write(f"**{t}**: {i}")
+            st.success("### 📋 Final Recommendation")
+            st.markdown(res['output'])
             
-            if not use_sim:
-                st.divider()
-                with st.expander("Draft Email"):
-                    st.text_area("Content", value=agent.call_llm(f"Draft email for: {ans}"))
+            with st.expander("🔍 Detailed Trace"):
+                steps = res.get("intermediate_steps", [])
+                for i, (action, observation) in enumerate(steps):
+                    st.markdown(f"**Step {i+1}:** Tool `{action.tool}` | Output: `{observation}`")
 
-else:
-    st.info("👈 Login Required")
+            if not use_simulation:
+                st.divider()
+                with st.expander("✉️ Draft Email"):
+                    email_prompt = f"Write a formal email based on this decision: {res['output']}"
+                    with st.spinner("Drafting..."):
+                        email_draft = llm.invoke(email_prompt).content
+                    st.text_area("Email Draft", value=email_draft, height=200)
+
+elif not st.session_state.get('is_key_valid', False):
+    st.info("👈 Please validate your Groq API Key.")