Update app.py

app.py

@@ -32,33 +32,36 @@ def load_page(url: str) -> str:
 
 # Initialize Gemini client
 client = genai.Client(api_key=os.environ.get("GEMINI_API_KEY"))
-MODEL = "gemini-2.5-pro-exp-03-25"
+MODEL = "gemini-2.0-flash"
 
-TOOLS = [
-    types.Tool(
-        function_declarations=[
-            types.FunctionDeclaration(
-                name="load_page",
-                description="Load webpage content as markdown",
-                parameters={
-                    "type": "object",
-                    "properties": {
-                        "url": {"type": "string", "description": "Full URL to load"}
-                    },
-                    "required": ["url"]
-                }
-            )
-        ]
-    ),
-    types.Tool(google_search=types.GoogleSearch()),
-    types.Tool(code_execution=types.ToolCodeExecution())
-]
+TOOL_CONFIGS = {
+    "web": [
+        types.Tool(
+            function_declarations=[
+                types.FunctionDeclaration(
+                    name="load_page",
+                    description="Load webpage content as markdown",
+                    parameters={
+                        "type": "object",
+                        "properties": {
+                            "url": {"type": "string", "description": "Full URL to load"}
+                        },
+                        "required": ["url"]
+                    }
+                )
+            ]
+        ),
+        types.Tool(google_search=types.GoogleSearch())
+    ],
+    "code": [
+        types.Tool(code_execution=types.ToolCodeExecution())
+    ]
+}
 
-SYSTEM_INSTRUCTION = """You are an AI assistant with:
-1. Web browsing capabilities
-2. Code execution for calculations
-3. Data analysis skills
-Use the most appropriate tool for each query."""
+SYSTEM_INSTRUCTION = """You are an AI assistant. Choose appropriate tools based on the query:
+- Use web tools for real-time information
+- Use code execution for calculations/analysis
+- Never combine search with code execution"""
 
 def format_response(parts):
     """Format response parts with proper Markdown formatting"""
@@ -72,66 +75,36 @@ def format_response(parts):
             formatted.append(f"**Result**:\n```\n{part.code_execution_result.output}\n```")
     return "\n\n".join(formatted)
 
+def select_tools(query: str) -> list:
+    """Select appropriate tools based on query content"""
+    if any(keyword in query.lower() for keyword in ["plot", "calculate", "simulate", "algorithm"]):
+        return TOOL_CONFIGS["code"]
+    return TOOL_CONFIGS["web"]
+
 def generate_response(user_input):
-    full_response = ""
-    chat = client.chats.create(
-        model=MODEL,
-        config=types.GenerateContentConfig(
-            temperature=0.7,
-            tools=TOOLS,
-            system_instruction=SYSTEM_INSTRUCTION
+    try:
+        if not user_input.strip():
+            raise ValueError("Please enter a valid query")
+            
+        selected_tools = select_tools(user_input)
+        
+        response = client.models.generate_content(
+            model=MODEL,
+            contents=[user_input],
+            config=types.GenerateContentConfig(
+                temperature=0.7,
+                tools=selected_tools,
+                system_instruction=SYSTEM_INSTRUCTION
+            )
         )
-    )
-    
-    # Initial request
-    response = chat.send_message(user_input)
-    
-    # Process all response parts
-    response_parts = []
-    for part in response.candidates[0].content.parts:
-        response_parts.append(part)
-        full_response = format_response(response_parts)
-        yield full_response
         
-        # Handle function calls
-        if part.function_call:
-            fn = part.function_call
-            if fn.name == "load_page":
-                result = load_page(**fn.args)
-                chat.send_message(
-                    types.Content(
-                        parts=[
-                            types.Part(
-                                function_response=types.FunctionResponse(
-                                    name=fn.name,
-                                    id=fn.id,
-                                    response={"result": result}
-                                )
-                            )
-                        ]
-                    )
-                )
-                # Get final response after tool execution
-                final_response = chat.send_message("")
-                for final_part in final_response.candidates[0].content.parts:
-                    response_parts.append(final_part)
-                    full_response = format_response(response_parts)
-                    yield full_response
+        return format_response(response.candidates[0].content.parts)
+        
+    except Exception as e:
+        return f"Error: {str(e)}"
 
 # Create Gradio interface
-with gr.Blocks(
-    title="Gemini AI Assistant",
-    css=""".markdown-output {
-        padding: 20px;
-        border-radius: 5px;
-        background: #f9f9f9;
-    }
-    .markdown-output code {
-        background: #f3f3f3;
-        padding: 2px 5px;
-        border-radius: 3px;
-    }"""
-) as demo:
+with gr.Blocks(title="Gemini AI Assistant") as demo:
     gr.Markdown("# 🚀 Gemini AI Assistant")
     gr.Markdown("Web • Code • Data Analysis")
     
@@ -157,8 +130,7 @@ with gr.Blocks(
     submit_btn.click(
         fn=generate_response,
         inputs=input_box,
-        outputs=output_box,
-        queue=True
+        outputs=output_box
     )
     
     clear_btn.click(