Update app.py

app.py

@@ -6,7 +6,7 @@ from fastapi import FastAPI, HTTPException, Depends
 from fastapi.security import OAuth2PasswordBearer
 from pydantic import BaseModel
 from typing import List, Dict, Any
-import uvicorn # Ensure uvicorn is imported
+import uvicorn
 from supabase import create_client, Client
 import jwt
 from datetime import datetime, timedelta
@@ -46,34 +46,11 @@ def save_chat_turn(supabase_user_client: Client, user_id: str, session_id: str,
     except Exception as e:
         return {"error": f"Database insert failed: {str(e)}"}
 
-# --- FastAPI Application ---
-app = FastAPI()
-
-class LoginRequest(BaseModel):
-    user_id: str
-
-class AgentRequest(BaseModel):
-    prompt: str
-    user_id: str
-    session_id: str = "default-session"
-
-oauth2_scheme = OAuth2PasswordBearer(tokenUrl="login")
-
-@app.post("/login")
-async def login(request: LoginRequest):
-    if not supabase_admin:
-        raise HTTPException(status_code=500, detail="Authentication service is not configured.")
-    user_id = request.user_id
-    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)):
+# --- Core Agent Logic (Refactored) ---
+async def _run_agent_logic(prompt: str, user_id: str, session_id: str, token: str):
     try:
         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:
+        if not decoded_token.get("sub") or decoded_token.get("sub") != user_id:
             raise HTTPException(status_code=401, detail="Invalid token or user mismatch.")
         
         supabase_user_client = create_client(SUPABASE_URL, SUPABASE_KEY)
@@ -82,7 +59,6 @@ async def invoke_gemini_agent(request: AgentRequest, token: str = Depends(oauth2
         raise HTTPException(status_code=401, detail=f"Invalid token: {e}")
 
     def get_my_chat_history():
-        """Fetches my most recent conversation history to provide context."""
         return get_ai_advisor_chat_history(supabase_user_client, user_id)
     
     tools_for_gemini = [get_my_chat_history]
@@ -96,75 +72,89 @@ async def invoke_gemini_agent(request: AgentRequest, token: str = Depends(oauth2
                 role = "user" if turn['role'] == 'user' else "model"
                 conversation_history.append({'role': role, 'parts': [{'text': turn['content']}]})
         
-        conversation_history.append({'role': 'user', 'parts': [{'text': request.prompt}]})
+        conversation_history.append({'role': 'user', 'parts': [{'text': prompt}]})
         
         model = genai.GenerativeModel('gemini-pro')
         response = model.generate_content(
-            conversation_history,
-            tools=tools_for_gemini,
-            safety_settings={
-                'HARM_CATEGORY_DANGEROUS_CONTENT': 'BLOCK_NONE', 'HARM_CATEGORY_HARASSMENT': 'BLOCK_NONE',
-                'HARM_CATEGORY_HATE_SPEECH': 'BLOCK_NONE', 'HARM_CATEGORY_SEXUALLY_EXPLICIT': 'BLOCK_NONE',
-            }
+            conversation_history, tools=tools_for_gemini,
+            safety_settings={'HARM_CATEGORY_DANGEROUS_CONTENT': 'BLOCK_NONE', 'HARM_CATEGORY_HARASSMENT': 'BLOCK_NONE', 'HARM_CATEGORY_HATE_SPEECH': 'BLOCK_NONE', 'HARM_CATEGORY_SEXUALLY_EXPLICIT': 'BLOCK_NONE'}
         )
         final_response_text = response.text
-        save_chat_turn(supabase_user_client, user_id, request.session_id, request.prompt, final_response_text)
-        return {"response": final_response_text}
+        save_chat_turn(supabase_user_client, user_id, session_id, prompt, final_response_text)
+        return final_response_text
     except Exception as e:
         print(f"Error in agent invocation: {e}")
-        raise HTTPException(status_code=500, detail=str(e))
+        # Return error as a string to be displayed in chat
+        return f"An error occurred while processing your request: {e}"
+
+# --- FastAPI Endpoints ---
+app = FastAPI()
+
+class LoginRequest(BaseModel):
+    user_id: str
+
+class AgentRequest(BaseModel):
+    prompt: str
+    user_id: str
+    session_id: str = "default-session"
+
+oauth2_scheme = OAuth2PasswordBearer(tokenUrl="login")
+
+@app.post("/login")
+async def login(request: LoginRequest):
+    if not supabase_admin: raise HTTPException(status_code=500, detail="Auth service not configured.")
+    user_id = request.user_id
+    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)):
+    response_text = await _run_agent_logic(request.prompt, request.user_id, request.session_id, token)
+    return {"response": response_text}
 
 # --- Gradio UI ---
 with gr.Blocks() 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.")
+    gr.Markdown("# 💎 Secure, Multi-User AI Agent (Live)")
+    gr.Markdown("Log in with a User ID (e.g., 'user-a'). This creates a secure token. Your chat history is then saved to your own user record in Supabase.")
 
     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")
+        user_id_input = gr.Textbox(label="Enter 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, height=500)
     msg_input = gr.Textbox(label="Your Message", visible=False, show_label=False, placeholder="Type your message here...")
     
-    def login_fn(user_id):
+    # CHANGED: login_fn now calls the login function directly
+    async 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: ""}
+            return {login_status: gr.update(value="<p style='color:red;'>Please enter a User ID.</p>")}
         try:
-            # For HF Spaces, the app runs on a specific port, but we call it internally via its name
-            response = requests.post("http://localhost: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}. 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),
-                }
-            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):
+            # Direct async call, no requests
+            token_data = await login(LoginRequest(user_id=user_id))
+            token = token_data["access_token"]
+            status_html = f"<p style='color:green;'>Logged in as: {user_id}. You can start chatting.</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),
+            }
+        except Exception as e:
+            return {login_status: gr.update(value=f"<p style='color:red;'>Login failed: {e}</p>")}
+
+    # CHANGED: chat_fn now calls the agent logic function directly
+    async def chat_fn(message, chat_history, user_id, token, session_id):
         if not token:
             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}
-            response = requests.post("http://localhost: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}"
-
+        
+        # Direct async call, no requests or headers
+        bot_message = await _run_agent_logic(message, user_id, session_id, token)
+        
         chat_history.append((message, bot_message))
         return "", chat_history
 
@@ -173,6 +163,5 @@ with gr.Blocks() as demo:
 
 app = gr.mount_gradio_app(app, demo, path="/")
 
-# --- ADD THIS BLOCK AT THE VERY END ---
 if __name__ == "__main__":
     uvicorn.run(app, host="0.0.0.0", port=7860)