Create app.py

app.py

@@ -0,0 +1,226 @@
+import os
+import requests
+import gradio as gr
+import google.generativeai as genai
+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
+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_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) ---
+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.
+
+def get_ai_advisor_chat_history(supabase_user_client: Client, user_id: str):
+    """Fetches the conversation history for the logged-in user."""
+    print(f"TOOL CALL: get_ai_advisor_chat_history for user '{user_id}' from Supabase")
+    try:
+        response = supabase_user_client.table('chat_history').select("*").eq('user_id', user_id).order('timestamp').limit(50).execute()
+        return response.data
+    except Exception as e:
+        return {"error": f"Database query failed: {str(e)}"}
+
+def save_chat_turn(supabase_user_client: Client, user_id: str, session_id: str, user_message: str, assistant_message: str):
+    """Saves the conversation turn to the database for the logged-in user."""
+    print(f"TOOL CALL: Saving chat turn for user '{user_id}' to Supabase")
+    try:
+        supabase_user_client.table('chat_history').insert([
+            {"user_id": user_id, "session_id": session_id, "role": "user", "content": user_message},
+            {"user_id": user_id, "session_id": session_id, "role": "assistant", "content": assistant_message}
+        ]).execute()
+        return {"status": "success"}
+    except Exception as e:
+        return {"error": f"Database insert failed: {str(e)}"}
+
+# --- FastAPI Application with Authentication ---
+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")
+
+@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",
+    }
+    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
+    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!
+        
+    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
+    
+    def get_my_chat_history():
+        """Fetches my most recent conversation history."""
+        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
+    try:
+        # Fetch the user's history to provide context to the model
+        history_data = get_my_chat_history()
+        chat_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)
+        
+        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)
+        
+        return {"response": final_response_text}
+        
+    except Exception as e:
+        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")
+    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())
+
+    with gr.Row():
+        user_id_input = gr.Textbox(label="Enter a User ID to Log In", placeholder="e.g., user-a")
+        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)
+    
+    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
+                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),
+                }
+            else:
+                 return {login_status: gr.update(value=f"<p style='color:red;'>Login failed: {response.text}</p>"), user_id_state: ""}
+        except requests.RequestException as e:
+            return {login_status: gr.update(value=f"<p style='color:red;'>Error connecting to backend: {e}</p>"), user_id_state: ""}
+
+    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."))
+            return "", chat_history
+        
+        try:
+            headers = {"Authorization": f"Bearer {token}"}
+            payload = {"prompt": message, "user_id": user_id, "session_id": session_id}
+            response = requests.post("http://127.0.0.1:7860/agent/invoke", json=payload, headers=headers)
+            
+            if response.status_code == 200:
+                bot_message = response.json()["response"]
+            else:
+                bot_message = f"Error from agent: {response.text}"
+        except requests.RequestException as e:
+            bot_message = f"Error connecting to agent: {e}"
+
+        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])
+    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="/")