Spaces:
Sleeping
Sleeping
Create app.py
Browse files
app.py
ADDED
@@ -0,0 +1,143 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
import os
|
2 |
+
from fastapi import FastAPI, HTTPException
|
3 |
+
from pydantic import BaseModel # Used to define expected request data
|
4 |
+
from dotenv import load_dotenv # To load our secret API keys from .env file
|
5 |
+
|
6 |
+
# --- LlamaIndex and Composio Imports ---
|
7 |
+
from composio_llamaindex import ComposioToolSet, Action # Composio tools for LlamaIndex
|
8 |
+
from llama_index.core.agent import FunctionCallingAgentWorker # The agent framework
|
9 |
+
from llama_index.core.llms import ChatMessage # Standard message format
|
10 |
+
from llama_index.llms.gemini import Gemini # The Gemini LLM integration
|
11 |
+
|
12 |
+
# --- Load API Keys Safely ---
|
13 |
+
load_dotenv() # Load variables from the .env file
|
14 |
+
gemini_api_key = os.getenv("GEMINI_API_KEY")
|
15 |
+
composio_api_key = os.getenv("COMPOSIO_API_KEY")
|
16 |
+
|
17 |
+
# Basic check to ensure keys are loaded
|
18 |
+
if not gemini_api_key:
|
19 |
+
raise ValueError("GEMINI_API_KEY not found. Make sure it's set in your .env file.")
|
20 |
+
if not composio_api_key:
|
21 |
+
raise ValueError("COMPOSIO_API_KEY not found. Make sure it's set in your .env file.")
|
22 |
+
|
23 |
+
# --- Global Agent Variable ---
|
24 |
+
# We'll set this up when the app starts
|
25 |
+
agent = None
|
26 |
+
|
27 |
+
# --- Function to Setup the Agent ---
|
28 |
+
def setup_gmail_agent():
|
29 |
+
"""Initializes the LlamaIndex agent with Gemini and Composio tools."""
|
30 |
+
global agent # Allow modification of the global 'agent' variable
|
31 |
+
try:
|
32 |
+
print("Initializing LLM (Gemini)...")
|
33 |
+
# Initialize the Gemini Language Model
|
34 |
+
# You might need to specify a model, e.g., "models/gemini-pro" if default isn't sufficient
|
35 |
+
llm = Gemini(api_key=gemini_api_key)
|
36 |
+
|
37 |
+
print("Initializing Composio ToolSet...")
|
38 |
+
# Initialize Composio ToolSet - it securely connects to your Gmail via the setup you did
|
39 |
+
composio_toolset = ComposioToolSet(api_key=composio_api_key)
|
40 |
+
|
41 |
+
print("Fetching Gmail tools from Composio...")
|
42 |
+
# Get specific Gmail tools (Actions) you want the agent to use
|
43 |
+
gmail_tools = composio_toolset.get_tools(actions=[
|
44 |
+
Action.GMAIL_SEND_EMAIL,
|
45 |
+
# Action.GMAIL_CREATE_EMAIL_DRAFT, # You can uncomment to add more actions
|
46 |
+
Action.GMAIL_REPLY_TO_THREAD,
|
47 |
+
Action.GMAIL_FETCH_MESSAGE_BY_THREAD_ID,
|
48 |
+
Action.GMAIL_LIST_THREADS, # Good for getting recent conversations
|
49 |
+
# Action.GMAIL_FETCH_EMAILS, # More generic fetch if needed
|
50 |
+
])
|
51 |
+
print(f"Loaded {len(gmail_tools)} Gmail tools.")
|
52 |
+
|
53 |
+
# --- Define the Agent's "Personality" or Instructions ---
|
54 |
+
system_prompt_content = """You are a helpful Gmail assistant.
|
55 |
+
Your tasks are to:
|
56 |
+
1. Read and list recent email threads when asked.
|
57 |
+
2. Read the content of a specific email/thread when given a thread ID.
|
58 |
+
3. Reply to email threads when requested, using the provided thread ID and message body.
|
59 |
+
4. Send new emails when requested, using the recipient, subject, and body provided.
|
60 |
+
Use the available tools precisely based on the user's request.
|
61 |
+
Be concise in your confirmation messages.
|
62 |
+
If you need a 'thread_id' to reply or fetch a specific message, and the user hasn't provided one, ask them to list threads first to get the ID.
|
63 |
+
"""
|
64 |
+
|
65 |
+
prefix_messages = [ChatMessage(role="system", content=system_prompt_content)]
|
66 |
+
|
67 |
+
print("Creating the Function Calling Agent Worker...")
|
68 |
+
# Create the agent worker - this combines the LLM, tools, and instructions
|
69 |
+
agent_worker = FunctionCallingAgentWorker(
|
70 |
+
tools=gmail_tools,
|
71 |
+
llm=llm,
|
72 |
+
prefix_messages=prefix_messages,
|
73 |
+
max_function_calls=10, # Max tool uses per query
|
74 |
+
allow_parallel_tool_calls=False, # Process tools one by one
|
75 |
+
verbose=True # Print agent's thinking process (helpful for debugging)
|
76 |
+
)
|
77 |
+
agent = agent_worker.as_agent() # Make the worker usable as an agent
|
78 |
+
print("--- Gmail Agent Initialized Successfully! ---")
|
79 |
+
|
80 |
+
except Exception as e:
|
81 |
+
print(f"---!!! ERROR Initializing Agent: {e} !!!---")
|
82 |
+
agent = None # Ensure agent is None if setup fails
|
83 |
+
|
84 |
+
# --- FastAPI Application Setup ---
|
85 |
+
print("Setting up FastAPI app...")
|
86 |
+
app = FastAPI(
|
87 |
+
title="Gmail AI Agent API (Gemini + Composio)",
|
88 |
+
description="Backend API for a Gmail assistant powered by Gemini, LlamaIndex, and Composio.",
|
89 |
+
version="1.0.0",
|
90 |
+
)
|
91 |
+
|
92 |
+
# --- Define Expected Input Data Structure ---
|
93 |
+
class QueryRequest(BaseModel):
|
94 |
+
query: str # We expect requests to have a single field named "query"
|
95 |
+
|
96 |
+
# --- Run Agent Setup When FastAPI Starts ---
|
97 |
+
@app.on_event("startup")
|
98 |
+
async def startup_event():
|
99 |
+
"""This function runs automatically when the FastAPI server starts."""
|
100 |
+
print("Application startup: Initializing Gmail agent...")
|
101 |
+
setup_gmail_agent()
|
102 |
+
if agent is None:
|
103 |
+
print("WARNING: Agent did not initialize successfully on startup.")
|
104 |
+
|
105 |
+
# --- Health Check Endpoint ---
|
106 |
+
@app.get("/health")
|
107 |
+
async def health_check():
|
108 |
+
"""A simple endpoint to check if the server is running and agent is initialized."""
|
109 |
+
if agent:
|
110 |
+
return {"status": "healthy", "agent_initialized": True}
|
111 |
+
else:
|
112 |
+
# Return unhealthy status if agent setup failed
|
113 |
+
return {"status": "unhealthy", "agent_initialized": False, "message": "Agent initialization failed during startup."}
|
114 |
+
|
115 |
+
# --- Main Interaction Endpoint ---
|
116 |
+
@app.post("/interact") # We use POST because the user is sending data (their query)
|
117 |
+
async def process_query(request: QueryRequest):
|
118 |
+
"""Receives a user query, processes it with the agent, and returns the response."""
|
119 |
+
print(f"Received query: {request.query}")
|
120 |
+
if agent is None:
|
121 |
+
# If agent isn't ready, return an error
|
122 |
+
raise HTTPException(status_code=503, detail="Agent not initialized or initialization failed. Check server logs.")
|
123 |
+
|
124 |
+
try:
|
125 |
+
# --- Core Logic: Use the Agent ---
|
126 |
+
# Use agent.achat for asynchronous execution (better performance for web servers)
|
127 |
+
agent_response = await agent.achat(request.query)
|
128 |
+
print(f"Agent response: {agent_response}")
|
129 |
+
|
130 |
+
# Return the agent's response
|
131 |
+
return {"status": "success", "response": str(agent_response)} # Convert response to string
|
132 |
+
|
133 |
+
except Exception as e:
|
134 |
+
# Handle errors during processing
|
135 |
+
print(f"---!!! ERROR processing query '{request.query}': {e} !!!---")
|
136 |
+
raise HTTPException(status_code=500, detail=f"Error processing your request: {str(e)}")
|
137 |
+
|
138 |
+
# --- Root endpoint (optional, just to show the app is running) ---
|
139 |
+
@app.get("/")
|
140 |
+
async def root():
|
141 |
+
return {"message": "Gmail Agent API is running. Use the /interact endpoint to send queries."}
|
142 |
+
|
143 |
+
print("FastAPI app setup complete. Waiting for server to start...")
|