Update app.py

app.py

@@ -1,37 +1,34 @@
 import os
 import requests
 import gradio as gr
-import google.generativeai as genai
+from google import genai # CHANGED: Modern library import
+from google.generativeai.types import HarmCategory, HarmBlockThreshold # For safety settings
 from fastapi import FastAPI, HTTPException, Depends
 from fastapi.security import OAuth2PasswordBearer
 from pydantic import BaseModel
 from typing import List, Dict, Any
 import uvicorn
 from supabase import create_client, Client
-import jwt # You'll need to add 'PyJWT' to requirements.txt
+import jwt
 from datetime import datetime, timedelta
 
 # --- Configuration ---
 GOOGLE_API_KEY = os.getenv("GOOGLE_API_KEY")
 SUPABASE_URL = os.getenv("SUPABASE_URL")
-SUPABASE_KEY = os.getenv("SUPABASE_KEY") # This is the public 'anon' key
-# NEW: Add the service_role key to your HF Space secrets
+SUPABASE_KEY = os.getenv("SUPABASE_KEY")
 SUPABASE_SERVICE_KEY = os.getenv("SUPABASE_SERVICE_KEY")
-# NEW: A secret for signing our own JWTs
 JWT_SECRET = os.getenv("JWT_SECRET", "a-strong-secret-key-for-local-testing")
 
-# --- Initialize Supabase Admin Client (for user authentication) ---
+# --- Initialize Supabase Admin Client ---
 try:
-    # We use the service key to create an admin client that can manage users
     supabase_admin: Client = create_client(SUPABASE_URL, SUPABASE_SERVICE_KEY)
     print("Successfully created Supabase admin client.")
 except Exception as e:
     supabase_admin = None
     print(f"Warning: Supabase service key invalid. User authentication will fail. Error: {e}")
 
-# --- Tool Definitions (Unchanged, but now they will use a user-specific client) ---
-# We will create the supabase client on-the-fly for each request using the user's token.
-# So we don't need a global client here. The tool functions will accept the client as an argument.
+# --- Tool Definitions (No changes here, these are still correct) ---
+# NOTE: In a real app, these would call your mock API. For this example, we keep the Supabase calls.
 
 def get_ai_advisor_chat_history(supabase_user_client: Client, user_id: str):
     """Fetches the conversation history for the logged-in user."""
@@ -58,91 +55,77 @@ def save_chat_turn(supabase_user_client: Client, user_id: str, session_id: str,
 app = FastAPI()
 
 class LoginRequest(BaseModel):
-    # In a real app, this would be email/password. Here we use a simple user_id.
     user_id: str
 
 class AgentRequest(BaseModel):
     prompt: str
     user_id: str
     session_id: str = "default-session"
-    # We no longer need to pass history, as it will be fetched by the tool
-    # based on the authenticated user.
 
-# This is our security scheme
-oauth2_scheme = OAuth2PasswordBearer(tokenUrl="token")
+oauth2_scheme = OAuth2PasswordBearer(tokenUrl="login") # Corrected tokenUrl
 
 @app.post("/login")
 async def login(request: LoginRequest):
-    """
-    Simulates a login. Creates a custom JWT that Supabase will trust.
-    This endpoint uses the ADMIN client to generate a token for a user.
-    """
     if not supabase_admin:
         raise HTTPException(status_code=500, detail="Authentication service is not configured.")
-    
-    # In a real app, you would verify a password here.
-    # For this demo, we'll just create a token for the given user_id.
     user_id = request.user_id
-    
-    # Create a custom JWT token that Supabase can understand.
-    # This token asserts the user's identity.
-    payload = {
-        "sub": user_id, # The user's UUID
-        "aud": "authenticated",
-        "exp": datetime.utcnow() + timedelta(hours=1), # Token expires in 1 hour
-        "role": "authenticated",
-    }
+    payload = { "sub": user_id, "aud": "authenticated", "exp": datetime.utcnow() + timedelta(hours=1), "role": "authenticated", }
     access_token = jwt.encode(payload, JWT_SECRET, algorithm="HS256")
     return {"access_token": access_token, "token_type": "bearer"}
 
-
 @app.post("/agent/invoke")
 async def invoke_gemini_agent(request: AgentRequest, token: str = Depends(oauth2_scheme)):
-    """
-    This endpoint is now secure. It requires a valid token.
-    """
-    # 1. Authenticate the user and create a user-specific Supabase client
+    # 1. Authenticate and create user-specific Supabase client
     try:
-        # Decode the token to get the user_id
         decoded_token = jwt.decode(token, JWT_SECRET, algorithms=["HS256"], audience="authenticated")
         user_id = decoded_token.get("sub")
         if not user_id or user_id != request.user_id:
             raise HTTPException(status_code=401, detail="Invalid token or user mismatch.")
         
-        # Create a Supabase client that is authenticated as THIS user
         supabase_user_client = create_client(SUPABASE_URL, SUPABASE_KEY)
-        supabase_user_client.auth.set_session(access_token=token, refresh_token="dummy") # Important!
-        
+        supabase_user_client.auth.set_session(access_token=token, refresh_token=token) # Use token for both
     except jwt.PyJWTError as e:
         raise HTTPException(status_code=401, detail=f"Invalid token: {e}")
 
-    # 2. Define tools that can access the user-specific client
-    # We define them here so they have access to the authenticated client and user_id
-    
+    # 2. Define tools available for this specific request
     def get_my_chat_history():
-        """Fetches my most recent conversation history."""
+        """Fetches my most recent conversation history to provide context."""
         return get_ai_advisor_chat_history(supabase_user_client, user_id)
-        
-    # The Gemini model itself doesn't need to know the user_id, it just calls the tool.
-    # We handle passing the id inside the function definition.
-    tools_for_gemini = [get_my_chat_history]
     
-    # 3. Fetch history and run the Gemini agent
+    tools_for_gemini = [get_my_chat_history]
+
+    # 3. Fetch history and run the Gemini agent using the MODERN library structure
     try:
-        # Fetch the user's history to provide context to the model
+        # CHANGED: Initialize the modern client
+        client = genai.Client(api_key=GOOGLE_API_KEY)
+        
         history_data = get_my_chat_history()
-        chat_history = []
+        
+        # CHANGED: Build the conversation history in the new format
+        conversation_history = []
         if isinstance(history_data, list):
             for turn in history_data:
-                chat_history.append({"role": turn['role'], "parts": [turn['content']]})
-
-        genai.configure(api_key=GOOGLE_API_KEY)
-        model = genai.GenerativeModel(model_name='gemini-pro', tools=tools_for_gemini)
-        chat = model.start_chat(history=chat_history)
+                role = "user" if turn['role'] == 'user' else "model"
+                conversation_history.append({'role': role, 'parts': [{'text': turn['content']}]})
+        
+        # Add the current user prompt
+        conversation_history.append({'role': 'user', 'parts': [{'text': request.prompt}]})
+
+        # CHANGED: Call the model with the full history
+        model = client.models.get('gemini-pro')
+        response = model.generate_content(
+            conversation_history,
+            tools=tools_for_gemini,
+            safety_settings={ # Added safety settings for robustness
+                HarmCategory.HARM_CATEGORY_DANGEROUS_CONTENT: HarmBlockThreshold.BLOCK_NONE,
+                HarmCategory.HARM_CATEGORY_HARASSMENT: HarmBlockThreshold.BLOCK_NONE,
+                HarmCategory.HARM_CATEGORY_HATE_SPEECH: HarmBlockThreshold.BLOCK_NONE,
+                HarmCategory.HARM_CATEGORY_SEXUALLY_EXPLICIT: HarmBlockThreshold.BLOCK_NONE,
+            }
+        )
+
+        final_response_text = response.text
         
-        response = chat.send_message(request.prompt)
-        final_response_text = " ".join([part.text for part in response.parts if hasattr(part, 'text')])
-
         # 4. Save the new conversation turn to the database
         save_chat_turn(supabase_user_client, user_id, request.session_id, request.prompt, final_response_text)
         
@@ -152,15 +135,11 @@ async def invoke_gemini_agent(request: AgentRequest, token: str = Depends(oauth2
         print(f"Error in agent invocation: {e}")
         raise HTTPException(status_code=500, detail=str(e))
 
-
-# --- Gradio UI (The Face) ---
-# The UI needs to be updated to handle the login flow.
-
-with gr.Blocks() as demo:
-    gr.Markdown("# 💎 Secure, Multi-User AI Agent")
+# --- Gradio UI (The Face) - No changes needed here ---
+with gr.Blocks(theme=gr.themes.Soft()) as demo:
+    gr.Markdown("# 💎 Secure, Multi-User AI Agent (Upgraded)")
     gr.Markdown("First, 'log in' with a User ID (e.g., 'user-a', 'user-b'). This will generate a secure token. Then, you can chat with the agent, and your conversation will be saved securely to your own user record in Supabase.")
 
-    # State variables to hold session info
     user_id_state = gr.State("")
     token_state = gr.State("")
     session_id_state = gr.State("session-" + os.urandom(8).hex())
@@ -170,27 +149,20 @@ with gr.Blocks() as demo:
         login_button = gr.Button("Login")
         login_status = gr.Markdown("")
 
-    chatbot = gr.Chatbot(label="Agent Chat", visible=False)
-    msg_input = gr.Textbox(label="Your Message", visible=False)
-    clear_button = gr.Button("Clear Chat", visible=False)
+    chatbot = gr.Chatbot(label="Agent Chat", visible=False, height=500)
+    msg_input = gr.Textbox(label="Your Message", visible=False, show_label=False, placeholder="Type your message here...")
     
     def login_fn(user_id):
         if not user_id:
             return {login_status: gr.update(value="<p style='color:red;'>Please enter a User ID.</p>"), user_id_state: ""}
         try:
-            # Call our own FastAPI /login endpoint
             response = requests.post("http://127.0.0.1:7860/login", json={"user_id": user_id})
             if response.status_code == 200:
                 token = response.json()["access_token"]
-                status_html = f"<p style='color:green;'>Logged in as: {user_id}</p>"
-                # After successful login, show the chat interface
+                status_html = f"<p style='color:green;'>Logged in as: {user_id}. You can start chatting now.</p>"
                 return {
-                    login_status: gr.update(value=status_html),
-                    user_id_state: user_id,
-                    token_state: token,
-                    chatbot: gr.update(visible=True),
-                    msg_input: gr.update(visible=True),
-                    clear_button: gr.update(visible=True),
+                    login_status: gr.update(value=status_html), user_id_state: user_id, token_state: token,
+                    chatbot: gr.update(visible=True), msg_input: gr.update(visible=True),
                 }
             else:
                  return {login_status: gr.update(value=f"<p style='color:red;'>Login failed: {response.text}</p>"), user_id_state: ""}
@@ -199,9 +171,8 @@ with gr.Blocks() as demo:
 
     def chat_fn(message, chat_history, user_id, token, session_id):
         if not token:
-            chat_history.append((message, "Error: You are not logged in. Please enter a User ID and click Login."))
+            chat_history.append((message, "Error: You are not logged in."))
             return "", chat_history
-        
         try:
             headers = {"Authorization": f"Bearer {token}"}
             payload = {"prompt": message, "user_id": user_id, "session_id": session_id}
@@ -217,10 +188,7 @@ with gr.Blocks() as demo:
         chat_history.append((message, bot_message))
         return "", chat_history
 
-    login_button.click(login_fn, inputs=[user_id_input], outputs=[login_status, user_id_state, token_state, chatbot, msg_input, clear_button])
+    login_button.click(login_fn, inputs=[user_id_input], outputs=[login_status, user_id_state, token_state, chatbot, msg_input])
     msg_input.submit(chat_fn, [msg_input, chatbot, user_id_state, token_state, session_id_state], [msg_input, chatbot])
-    clear_button.click(lambda: None, None, chatbot, queue=False)
-
 
-# Mount the Gradio app onto the FastAPI server
 app = gr.mount_gradio_app(app, demo, path="/")