Update app.py

fix the model error

app.py

@@ -2,7 +2,6 @@ import streamlit as st
 import pandas as pd
 import os
 import warnings
-import time
 import sqlite3
 import shutil
 import asyncio
@@ -27,9 +26,13 @@ INDEX_PATH = "faiss_index"
 REQUIRED_PDFS = ["Bank Loan Overall Risk Policy.pdf", "Bank Loan Interest Rate Policy.pdf"]
 
 try:
+    # GROQ (Keep LangChain)
     from langchain_groq import ChatGroq
-    from langchain_google_genai import ChatGoogleGenerativeAI
+    
+    # GOOGLE (Use Raw SDK - More Stable)
     import google.generativeai as genai
+    
+    # SHARED UTILS
     from langchain_huggingface import HuggingFaceEmbeddings
     from langchain_community.vectorstores import FAISS
     from langchain_community.document_loaders import PyPDFLoader
@@ -37,13 +40,12 @@ try:
     from langchain_core.prompts import PromptTemplate
     from langchain_core.runnables import RunnablePassthrough
     from langchain_core.output_parsers import StrOutputParser
-    from langchain_core.tools import tool
 except ImportError as e:
     st.error(f"❌ Import Error: {e}")
     st.stop()
 
 # ==========================================
-# 2. DATABASE & TOOLS SETUP
+# 2. DATABASE & TOOLS
 # ==========================================
 def init_db():
     if os.path.exists(DB_FILE): return
@@ -68,15 +70,13 @@ def run_query(query, params=()):
             return cursor.fetchone()
     except Exception as e: return f"DB Error: {e}"
 
-# --- TOOL FUNCTIONS (Pure Python) ---
+# --- DIRECT TOOL FUNCTIONS ---
 def tool_get_credit_score(user_id):
-    """Input: User ID. Returns 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."
 
 def tool_get_account_status(user_id):
-    """Input: User ID. Returns Name, Nationality, Status."""
     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,))
     if row and not isinstance(row, str):
@@ -84,7 +84,6 @@ def tool_get_account_status(user_id):
     return "User ID not found."
 
 def tool_check_pr_status(user_id):
-    """Input: User ID. Returns 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()):
@@ -92,194 +91,220 @@ def tool_check_pr_status(user_id):
     return f"PR Status: {row[0]}" if (row and not isinstance(row, str)) else "PR Status: False."
 
 # ==========================================
-# 3. MANUAL AGENT ENGINE (The Fix)
+# 3. HYBRID AGENT ENGINE (The Solution)
 # ==========================================
-class ManualReActAgent:
-    def __init__(self, llm, tools_map, rag_chain):
-        self.llm = llm
+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 = 6
+        self.max_steps = 8
+        
+        # Initialize Groq here (Reusable)
+        if "Groq" in provider:
+            self.groq_chat = ChatGroq(api_key=api_key, model_name="llama-3.3-70b-versatile", temperature=0)
+        
+        # Initialize Gemini Config
+        if "Google" in provider:
+            genai.configure(api_key=api_key)
+            # Use Flash - it's faster and smarter for tools
+            self.gemini_model = genai.GenerativeModel('gemini-1.5-flash') 
+
+    def call_llm(self, prompt):
+        """Switches between LangChain (Groq) and Raw SDK (Gemini)"""
+        if "Groq" in self.provider:
+            return self.groq_chat.invoke(prompt).content
+        else:
+            # Native Google Call - Bypasses LangChain errors
+            try:
+                response = self.gemini_model.generate_content(prompt)
+                return response.text
+            except Exception as e:
+                return f"Gemini Error: {str(e)}"
 
     def run(self, query):
-        """Runs the ReAct loop manually to avoid Library Parsing Errors."""
-        
-        # 1. DEFINE PROMPT
         tool_desc = "\n".join([f"- {name}: {func.__doc__}" for name, func in self.tools.items()])
-        system_prompt = f"""You are a Loan Risk Officer. Answer the question using the tools below.
+        
+        history = f"""You are a Loan Officer. Solve this request: "{query}"
 
-TOOLS:
+TOOLS AVAILABLE:
 {tool_desc}
-- consult_policy_doc: Consult policy PDF for risk rules. Input: a question string.
+- consult_policy_doc: Search PDF policies. Input: question string.
 
-FORMAT:
+RULES:
+1. You run in a loop. OUTPUT ONLY ONE STEP AT A TIME.
+2. Format:
 Thought: <reasoning>
 Action: <tool_name>
-Action Input: <input>
+Action Input: <input_string>
 Observation: <result>
-... (repeat)
-Final Answer: <answer>
+...
+Final Answer: <the full report>
 
 Begin!
-Question: {query}
 """
-        history = system_prompt
         logs = []
 
-        # 2. LOOP
         for i in range(self.max_steps):
-            # A. Call LLM
-            response = self.llm.invoke(history).content
+            # 1. Get LLM Response
+            response = self.call_llm(history)
             history += response + "\n"
             
-            # B. Parse "Action"
+            # 2. Check for Final Answer
+            if "Final Answer:" in response:
+                return response.split("Final Answer:")[-1].strip(), logs
+
+            # 3. Parse Tool Call
             action_match = re.search(r"Action:\s*(.+)", response)
             input_match = re.search(r"Action Input:\s*(.+)", response)
-            
-            # C. Check for Final Answer (Stop Condition)
-            if "Final Answer:" in response:
-                final_ans = response.split("Final Answer:")[-1].strip()
-                return final_ans, logs
 
-            # D. Execute Tool
             if action_match and input_match:
                 tool_name = action_match.group(1).strip()
-                tool_input = input_match.group(1).strip()
-                
-                # Strip quotes if present
-                tool_input = tool_input.strip('"').strip("'")
+                val = input_match.group(1).strip().strip('"').strip("'")
                 
-                logs.append((tool_name, tool_input))
+                logs.append((tool_name, val))
                 
                 # Execute
-                observation = f"Error: Tool {tool_name} not found."
+                result = "Error: Tool not found"
                 if tool_name in self.tools:
-                    try:
-                        observation = self.tools[tool_name](tool_input)
-                    except Exception as e:
-                        observation = f"Tool Error: {e}"
+                    try: result = self.tools[tool_name](val)
+                    except Exception as e: result = f"Error: {e}"
                 elif tool_name == "consult_policy_doc":
-                    try:
-                        observation = self.rag_chain.invoke(tool_input)
-                    except Exception as e:
-                        observation = f"RAG Error: {e}"
+                    try: result = self.rag_chain.invoke(val)
+                    except Exception as e: result = f"RAG Error: {e}"
                 
-                obs_str = f"Observation: {observation}\n"
-                history += obs_str
+                # Feed back
+                obs = f"Observation: {result}\n"
+                history += obs
             else:
-                # If LLM didn't output an action but didn't finish, force it
-                if i == self.max_steps - 1:
-                    return response, logs
-                history += "Observation: Please continue. If you have the answer, say 'Final Answer:'.\n"
+                # Force agent to continue if it stops early
+                if i == self.max_steps - 1: return response, logs
+                history += "Observation: Please continue. Use 'Final Answer:' when done.\n"
 
         return "Agent timed out.", logs
 
 # ==========================================
-# 4. UI & SETUP
+# 4. UI & LOGIC
 # ==========================================
 st.title("🤖 Multi-Model Loan Assessor")
 pdfs_missing = [f for f in REQUIRED_PDFS if not os.path.exists(f)]
 
 with st.sidebar:
     st.header("🔐 Authentication")
-    provider = st.radio("Model:", ["Groq (Llama-3)", "Google (Gemini)"])
+    provider_opt = st.radio("Model:", ["Groq (Llama-3)", "Google (Gemini)"])
     
-    if 'api_key' not in st.session_state: st.session_state.api_key = None
+    if 'auth' not in st.session_state: st.session_state.auth = False
     
-    key_input = st.text_input("API Key", type="password")
-    if st.button("Set Key"):
-        st.session_state.api_key = key_input
-        st.success("Key Set!")
-        st.rerun()
+    # Reset if provider changes
+    if st.session_state.get('last_provider') != provider_opt:
+        st.session_state.auth = False
+        st.session_state.last_provider = provider_opt
 
-    if st.button("♻️ Reset"):
+    if not st.session_state.auth:
+        key_in = st.text_input("API Key", type="password")
+        if st.button("Validate"):
+            try:
+                # Simple Validation
+                if "Groq" in provider_opt:
+                    ChatGroq(api_key=key_in).invoke("Hi")
+                else:
+                    genai.configure(api_key=key_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}")
+    else:
+        st.success("Active")
+        if st.button("Logout"): 
+            st.session_state.auth = False
+            st.rerun()
+
+    if st.button("♻️ Reset DB"):
         if os.path.exists(INDEX_PATH): shutil.rmtree(INDEX_PATH)
         st.cache_resource.clear()
         st.rerun()
 
-if st.session_state.api_key:
+if st.session_state.auth:
     # --- RAG SETUP ---
     @st.cache_resource
     def setup_rag():
         if pdfs_missing: return None
+        # Always use HuggingFace embeddings (Free, Fast, Compatible)
         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()
-        documents = []
-        for f in REQUIRED_PDFS: documents.extend(PyPDFLoader(f).load())
-        splits = CharacterTextSplitter(chunk_size=600, chunk_overlap=50).split_documents(documents)
+        
+        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()
 
-    with st.spinner("Loading AI..."):
-        retriever = setup_rag()
+    retriever = setup_rag()
 
-    # --- LLM SETUP ---
-    if "Groq" in provider:
-        llm = ChatGroq(api_key=st.session_state.api_key, temperature=0, model_name="llama-3.3-70b-versatile")
-    else:
-        # Using Gemini 1.5 Flash with REST transport
-        llm = ChatGoogleGenerativeAI(
-            google_api_key=st.session_state.api_key, 
-            temperature=0, 
-            model="gemini-1.5-flash",
-            transport="rest"
-        )
-
-    # --- RAG CHAIN ---
-    rag_chain = (
-        {"context": retriever | (lambda d: "\n".join([x.page_content for x in d])), "question": RunnablePassthrough()}
-        | PromptTemplate.from_template("Info: {context}\nQ: {question}\nA:") 
-        | llm 
-        | StrOutputParser()
-    )
-
-    # --- AGENT INSTANCE ---
-    tools_map = {
+    # --- RAG CHAIN FOR TOOLS ---
+    # We use a separate Groq LLM for the RAG lookup to ensure it's fast/stable
+    # regardless of the main agent choice.
+    rag_llm = ChatGroq(api_key=st.session_state.key, model_name="llama-3.3-70b-versatile") if "Groq" in provider_opt else None
+    
+    # Simple RAG Chain
+    def query_rag(q):
+        if not retriever: return "No PDFs found."
+        docs = retriever.invoke(q)
+        ctx = "\n".join([d.page_content for d in docs])
+        # If using Gemini, we format prompt manually for RAG too
+        return f"Context from Policy: {ctx}"
+
+    # Agent Tools Map
+    tools = {
         "get_credit_score": tool_get_credit_score,
         "get_account_status": tool_get_account_status,
         "check_pr_status": tool_check_pr_status
     }
-    agent = ManualReActAgent(llm, tools_map, rag_chain)
+    
+    # Initialize Hybrid Agent
+    # For RAG, we pass a simple lambda that calls our query_rag function
+    rag_lambda = type('RAG', (object,), {"invoke": lambda self, x: query_rag(x)})()
+    
+    agent = HybridAgent(provider_opt, st.session_state.key, tools, rag_lambda)
 
     # --- UI ---
     col1, col2 = st.columns([1, 2])
     with col1:
         uid = st.text_input("Customer ID", "1111")
         use_sim = st.checkbox("Simulation Mode")
-        sim_score = st.slider("Score", 300, 900, 450) if use_sim else 0
-        sim_status = st.selectbox("Status", ["good-standing", "closed", "delinquent"]) if use_sim else ""
-        btn = st.button("Assess Risk", type="primary")
+        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")
 
     with col2:
         if btn:
-            query = f"Process Loan for ID {uid}. "
-            if use_sim: query += f"SIMULATION: Score {sim_score}, Status '{sim_status}'. Do NOT query DB for score/status."
-            else: query += "Query DB for all info."
-            query += " Check policies. Report Risk, Rate, and Decision."
-
-            with st.status(f"🤖 {provider} Agent Running...", expanded=True):
-                st.write("Thinking...")
-                try:
-                    # Run Manual Loop
-                    final_res, logs = agent.run(query)
-                    st.write("✅ Done!")
-                except Exception as e:
-                    st.error(f"Error: {e}")
-                    final_res = "Failed."
-                    logs = []
-
-            st.success("### 📋 Report")
-            st.markdown(final_res)
+            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):
+                ans, logs = agent.run(q)
+                st.write("Done!")
+
+            st.success("### Final Report")
+            st.markdown(ans)
             
             with st.expander("Trace"):
-                for tool_name, tool_in in logs:
-                    st.markdown(f"**Tool:** `{tool_name}` | **Input:** `{tool_in}`")
+                for t, i in logs: st.write(f"**{t}**: {i}")
             
             if not use_sim:
                 st.divider()
-                st.text_area("✉️ Email Draft", value=llm.invoke(f"Draft email for: {final_res}").content)
+                with st.expander("Draft Email"):
+                    st.text_area("Content", value=agent.call_llm(f"Draft email for: {ans}"))
 
 else:
-    st.info("👈 Enter API Key")
+    st.info("👈 Login Required")