Update app.py

fix model error

app.py

@@ -26,14 +26,14 @@ INDEX_PATH = "faiss_index"
 REQUIRED_PDFS = ["Bank Loan Overall Risk Policy.pdf", "Bank Loan Interest Rate Policy.pdf"]
 
 try:
-    # GROQ (Keep LangChain)
+    # GROQ
     from langchain_groq import ChatGroq
     
-    # GOOGLE (Use Raw SDK - Stable)
+    # GOOGLE (Native SDK)
     import google.generativeai as genai
     from google.generativeai.types import HarmCategory, HarmBlockThreshold
     
-    # SHARED UTILS
+    # SHARED
     from langchain_huggingface import HuggingFaceEmbeddings
     from langchain_community.vectorstores import FAISS
     from langchain_community.document_loaders import PyPDFLoader
@@ -71,7 +71,7 @@ def run_query(query, params=()):
             return cursor.fetchone()
     except Exception as e: return f"DB Error: {e}"
 
-# --- DIRECT TOOL FUNCTIONS ---
+# --- TOOLS ---
 def tool_get_credit_score(user_id):
     clean_id = ''.join(filter(str.isdigit, str(user_id)))
     row = run_query("SELECT Credit_Score FROM credit_score WHERE ID = ?", (clean_id,))
@@ -92,7 +92,7 @@ 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. HYBRID AGENT ENGINE (The Solution)
+# 3. HYBRID AGENT (Dynamic Model Loader)
 # ==========================================
 class HybridAgent:
     def __init__(self, provider, api_key, tools_map, rag_chain):
@@ -101,25 +101,47 @@ class HybridAgent:
         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 (Native SDK)
+        # Initialize Gemini with DYNAMIC DISCOVERY
         if "Google" in provider:
             genai.configure(api_key=api_key)
-            # FIX: Use 'gemini-pro' (Stable) instead of 'flash' (404 Error)
-            self.gemini_model = genai.GenerativeModel('gemini-pro') 
+            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):
-        """Switches between LangChain (Groq) and Raw SDK (Gemini)"""
         if "Groq" in self.provider:
             return self.groq_chat.invoke(prompt).content
         else:
-            # Native Google Call
             try:
-                # Safety Settings to prevent blocking
+                # 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,
@@ -155,15 +177,12 @@ Begin!
         logs = []
 
         for i in range(self.max_steps):
-            # 1. Get LLM Response
             response = self.call_llm(history)
             history += response + "\n"
             
-            # 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)
 
@@ -173,7 +192,6 @@ Begin!
                 
                 logs.append((tool_name, val))
                 
-                # Execute
                 result = "Error: Tool not found"
                 if tool_name in self.tools:
                     try: result = self.tools[tool_name](val)
@@ -182,11 +200,8 @@ Begin!
                     try: result = self.rag_chain.invoke(val)
                     except Exception as e: result = f"RAG Error: {e}"
                 
-                # Feed back
-                obs = f"Observation: {result}\n"
-                history += obs
+                history += f"Observation: {result}\n"
             else:
-                # 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"
 
@@ -204,7 +219,6 @@ with st.sidebar:
     
     if 'auth' not in st.session_state: st.session_state.auth = False
     
-    # 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
@@ -213,12 +227,12 @@ with st.sidebar:
         key_in = st.text_input("API Key", type="password")
         if st.button("Validate"):
             try:
-                # Validation Logic
                 if "Groq" in provider_opt:
                     ChatGroq(api_key=key_in).invoke("Hi")
                 else:
                     genai.configure(api_key=key_in)
-                    genai.list_models()
+                    # 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
@@ -238,13 +252,10 @@ with st.sidebar:
         st.rerun()
 
 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()
         
@@ -257,28 +268,20 @@ if st.session_state.auth:
 
     retriever = setup_rag()
 
-    # --- RAG CHAIN FOR TOOLS ---
-    # Use Groq for RAG processing if available (faster), otherwise skip or use simplified
     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])
-        return f"Context from Policy: {ctx}"
+        return "Context: " + "\n".join([d.page_content for d in docs])
 
-    # 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
     }
     
-    # 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")
@@ -295,8 +298,12 @@ if st.session_state.auth:
             q += " Check Policy. Report Risk, Rate, Decision."
 
             with st.status("Agent Working...", expanded=True):
-                ans, logs = agent.run(q)
-                st.write("Done!")
+                try:
+                    ans, logs = agent.run(q)
+                    st.write("Done!")
+                except Exception as e:
+                    st.error(f"Execution Error: {e}")
+                    ans, logs = "Error", []
 
             st.success("### Final Report")
             st.markdown(ans)