samvish commited on
Commit
984961b
·
verified ·
1 Parent(s): 515e2a7

Create app.py

Browse files
Files changed (1) hide show
  1. app.py +143 -0
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...")