Update app.py

add in gemini api keys

app.py

@@ -9,7 +9,7 @@ import shutil
 # ==========================================
 # 1. PAGE CONFIG (MUST BE FIRST)
 # ==========================================
-st.set_page_config(page_title="Bank Loan Agent (SQL)", layout="wide")
+st.set_page_config(page_title="Bank Loan Agent (Multi-Model)", layout="wide")
 
 # Suppress warnings
 warnings.filterwarnings("ignore")
@@ -22,7 +22,12 @@ INDEX_PATH = "faiss_index"
 REQUIRED_PDFS = ["Bank Loan Overall Risk Policy.pdf", "Bank Loan Interest Rate Policy.pdf"]
 
 try:
+    # GROQ IMPORTS
     from langchain_groq import ChatGroq
+    # GOOGLE IMPORTS
+    from langchain_google_genai import ChatGoogleGenerativeAI
+    
+    # SHARED IMPORTS
     from langchain_huggingface import HuggingFaceEmbeddings
     from langchain_community.vectorstores import FAISS
     from langchain_community.callbacks import StreamlitCallbackHandler
@@ -36,13 +41,14 @@ try:
 
 except ImportError as e:
     st.error(f"❌ Critical Import Error: {e}")
+    st.info("💡 Suggestion: Add 'langchain-google-genai' to requirements.txt")
     st.stop()
 
 # ==========================================
 # 3. DATABASE SETUP
 # ==========================================
 def init_db():
-    """Converts CSV files to SQLite DB. Handles errors gracefully."""
+    """Converts CSV files to SQLite DB."""
     if os.path.exists(DB_FILE):
         return
 
@@ -60,7 +66,6 @@ def init_db():
                 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 Exception:
@@ -70,10 +75,8 @@ def init_db():
     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:
@@ -94,7 +97,7 @@ def get_credit_score(user_id: str) -> str:
     row = run_query("SELECT Credit_Score FROM credit_score WHERE ID = ?", (clean_id,))
     if row and not isinstance(row, str):
         return f"Credit Score: {row[0]}"
-    return "User ID not found in Credit DB."
+    return "User ID not found."
 
 @tool
 def get_account_status(user_id: str) -> str:
@@ -106,28 +109,26 @@ def get_account_status(user_id: str) -> str:
     )
     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 in Account DB."
+    return "User ID not found."
 
 @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,))
-
     if row and not isinstance(row, str):
         return f"PR Status: {row[0]}"
-    return "PR Status: False (Record not found)"
+    return "PR Status: False."
 
 # ==========================================
 # 5. STREAMLIT APP UI
 # ==========================================
-st.title("🤖 Multi-Policy Loan Assessor (SQL + RAG)")
+st.title("🤖 Multi-Model Loan Assessor")
 st.markdown("Agent connects to **SQLite Database** and **Persistent Vector Store**")
 
-# Calculate missing PDFs globally so everyone can see it
+# Calculate missing PDFs
 pdfs_missing = [f for f in REQUIRED_PDFS if not os.path.exists(f)]
 
 # --- METRICS FUNCTION ---
@@ -142,58 +143,74 @@ def update_metrics(placeholder):
             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 ---
+# --- SIDEBAR (MULTI-PROVIDER) ---
 with st.sidebar:
     st.header("🔐 Authentication")
     
-    # 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"):
+    # 1. Provider Selector
+    provider_option = st.radio("Select AI Model:", ["Groq (Llama-3)", "Google (Gemini)"])
+    
+    # Initialize State
+    if 'auth_status' not in st.session_state:
+        st.session_state['auth_status'] = False
+        st.session_state['api_key'] = None
+        st.session_state['provider'] = None
+
+    # Reset if user switches provider
+    if st.session_state.get('provider') != provider_option:
+        st.session_state['auth_status'] = False
+        st.session_state['api_key'] = None
+        st.session_state['provider'] = provider_option
+
+    # 2. Authentication Logic
+    if not st.session_state['auth_status']:
+        api_key_input = st.text_input(f"Enter {provider_option} API Key", type="password")
+        
+        if st.button("Validate 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
+                    with st.spinner(f"Contacting {provider_option}..."):
+                        # DYNAMIC VALIDATION
+                        if "Groq" in provider_option:
+                            test_llm = ChatGroq(api_key=api_key_input, model_name="llama-3.3-70b-versatile")
+                        else:
+                            test_llm = ChatGoogleGenerativeAI(google_api_key=api_key_input, model="gemini-1.5-flash")
+                        
+                        test_llm.invoke("Test Connection")
+                        
+                        # Store Success
+                        st.session_state['auth_status'] = True
+                        st.session_state['api_key'] = api_key_input
                         st.success("✅ Valid Key!")
                         time.sleep(0.5)
                         st.rerun()
                 except Exception as e:
-                    st.error(f"❌ Invalid Key: {e}")
+                    st.error(f"❌ Connection Failed: {e}")
     else:
-        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.success(f"✅ {st.session_state['provider']} Active")
+        if st.button("🔴 Logout / Change Model"):
+            st.session_state['auth_status'] = False
+            st.session_state['api_key'] = None
             st.rerun()
 
     st.divider()
-    st.subheader("🛠️ System Maintenance")
+    st.subheader("🛠️ Maintenance")
     
     if st.button("♻️ Rebuild Knowledge Base"):
-        if os.path.exists(INDEX_PATH):
-            shutil.rmtree(INDEX_PATH)
+        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.button("💾 Reload CSVs to DB"):
-        if os.path.exists(DB_FILE):
-            os.remove(DB_FILE)
+    if st.button("💾 Reload CSVs"):
+        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:
@@ -204,18 +221,17 @@ with st.sidebar:
     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']
+if st.session_state.get('auth_status', False):
     
-    # --- RAG SETUP ---
+    # --- RAG SETUP (Shared Embeddings) ---
     @st.cache_resource
     def setup_rag():
-        # Check global variable here
         if pdfs_missing:
             st.error(f"Missing PDFs: {pdfs_missing}")
             st.stop()
             
+        # We use HuggingFace embeddings for BOTH providers to keep the Vector Store compatible
+        # This prevents having to rebuild the index every time you switch models.
         embeddings = HuggingFaceEmbeddings(model_name="all-MiniLM-L6-v2")
         
         if os.path.exists(INDEX_PATH):
@@ -233,11 +249,19 @@ if st.session_state.get('is_key_valid', False):
             vectorstore.save_local(INDEX_PATH)
             return vectorstore.as_retriever()
 
-    with st.spinner("Initializing AI..."):
+    with st.spinner("Initializing Knowledge Base..."):
         retriever = setup_rag()
 
-    llm = ChatGroq(temperature=0, model_name="llama-3.3-70b-versatile")
-    
+    # --- DYNAMIC LLM INSTANTIATION ---
+    current_key = st.session_state['api_key']
+    current_provider = st.session_state['provider']
+
+    if "Groq" in current_provider:
+        llm = ChatGroq(api_key=current_key, temperature=0, model_name="llama-3.3-70b-versatile")
+    else:
+        llm = ChatGoogleGenerativeAI(google_api_key=current_key, temperature=0, model="gemini-1.5-flash")
+
+    # --- AGENT CHAIN ---
     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()}
@@ -260,6 +284,7 @@ if st.session_state.get('is_key_valid', False):
     agent = create_tool_calling_agent(llm, tools, prompt)
     agent_executor = AgentExecutor(agent=agent, tools=tools, verbose=True, return_intermediate_steps=True)
 
+    # --- UI INPUT ---
     col1, col2 = st.columns([1, 2])
     with col1:
         st.subheader("1. Customer Details")
@@ -277,20 +302,16 @@ if st.session_state.get('is_key_valid', False):
     
     with col2:
         if btn:
+            # Build Prompt
             if use_simulation:
                  query = f"""
                 Process Loan for Customer ID: {uid}.
-                *** SIMULATION MODE ***
+                *** SIMULATION MODE ACTIVE ***
                 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.
+                5. Output Final Report table + Justification.
                 """
             else:
                 query = f"""
@@ -298,15 +319,10 @@ if st.session_state.get('is_key_valid', False):
                 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.
+                4. Output Final Report table + Justification.
                 """
             
-            with st.status("🤖 Agent is processing...", expanded=True) as status:
+            with st.status(f"🤖 Agent ({current_provider}) is processing...", expanded=True) as status:
                 st_callback = StreamlitCallbackHandler(st.container())
                 try:
                     start_time = time.time()
@@ -335,5 +351,5 @@ if st.session_state.get('is_key_valid', False):
                         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.")
+elif not st.session_state.get('auth_status', False):
+    st.info("👈 Please select a provider (Groq or Gemini) and validate your key.")