Agent new agent features

app.py

@@ -1,38 +1,33 @@
 import gradio as gr
-from indexer import run_email_agent
-
-with gr.Blocks(css="""
-    body {background-color: #fffaf0;}  /* Light creme background */
-    .gr-textbox {min-height: 200px;}  /* Make textboxes taller */
-""") as demo:
-    gr.Markdown("## Tool-Based AI Email Agent")
-
-    user_input = gr.Textbox(
-        label="User Query",
-        placeholder="Type your question here...",
-        lines=4  # make input box slightly bigger
-    )
-    submit_btn = gr.Button("Run Agent")
-
-    output = gr.Textbox(
-        label="Agent Response",
-        placeholder="The agent will respond here...",
-        lines=10  # bigger response box
-    )
-
-    human_box = gr.Textbox(
-        label="Human Assistance Required",
-        visible=False
-    )
-
-    def handle_query(text):
-        response = run_email_agent(text)
-        return response
-
-    submit_btn.click(
-        handle_query,
-        inputs=user_input,
-        outputs=output
-    )
-
-demo.launch()
+from fastapi import FastAPI, Request
+from google_auth import get_auth_url, fetch_token
+from agent import run_agent
+import json
+
+app = FastAPI()
+
+@app.get("/login")
+def login():
+    auth_url, state = get_auth_url()
+    return {"auth_url": auth_url}
+
+@app.get("/auth/callback")
+async def callback(request: Request):
+    code = request.query_params.get("code")
+    token_data = fetch_token(code)
+
+    with open("user_token.json", "w") as f:
+        json.dump(token_data, f)
+
+    return {"status": "Login successful. You may close this tab."}
+
+def chat_interface(message):
+    return run_agent(message)
+
+gradio_ui = gr.Interface(
+    fn=chat_interface,
+    inputs="text",
+    outputs="text"
+)
+
+app = gr.mount_gradio_app(app, gradio_ui, path="/")

drive_tools.py

@@ -0,0 +1,49 @@
+from langchain_core.tools import tool
+from googleapiclient.discovery import build
+from googleapiclient.http import MediaIoBaseDownload
+from google_auth import dict_to_creds
+import io
+import os
+import json
+
+TOKEN_STORE = "user_token.json"
+
+@tool
+def search_and_download_doc_tool(file_name: str) -> str:
+    """
+    Searches Google Drive and downloads a document by name.
+    """
+
+    if not os.path.exists(TOKEN_STORE):
+        return "User not authenticated. Please login first."
+
+    with open(TOKEN_STORE, "r") as f:
+        token_data = json.load(f)
+
+    creds = dict_to_creds(token_data)
+    service = build("drive", "v3", credentials=creds)
+
+    results = service.files().list(
+        q=f"name contains '{file_name}'",
+        fields="files(id, name)"
+    ).execute()
+
+    items = results.get("files", [])
+
+    if not items:
+        return "No document found."
+
+    file_id = items[0]["id"]
+    request = service.files().get_media(fileId=file_id)
+
+    file_path = f"./downloads/{items[0]['name']}"
+    os.makedirs("downloads", exist_ok=True)
+
+    fh = io.FileIO(file_path, "wb")
+    downloader = MediaIoBaseDownload(fh, request)
+
+    done = False
+    while not done:
+        status, done = downloader.next_chunk()
+
+    return f"Downloaded to {file_path}"

google_auth.py

@@ -0,0 +1,60 @@
+import os
+import json
+from google_auth_oauthlib.flow import Flow
+from google.oauth2.credentials import Credentials
+
+SCOPES = ['https://www.googleapis.com/auth/drive.readonly']
+
+REDIRECT_URI = os.getenv("REDIRECT_URI")
+
+def create_flow():
+    client_config = {
+        "web": {
+            "client_id": os.getenv("GOOGLE_CLIENT_ID"),
+            "client_secret": os.getenv("GOOGLE_CLIENT_SECRET"),
+            "auth_uri": "https://accounts.google.com/o/oauth2/auth",
+            "token_uri": "https://oauth2.googleapis.com/token"
+        }
+    }
+
+    flow = Flow.from_client_config(
+        client_config,
+        scopes=SCOPES,
+        redirect_uri=REDIRECT_URI
+    )
+
+    return flow
+
+def get_auth_url():
+    flow = create_flow()
+    auth_url, state = flow.authorization_url(
+        access_type='offline',
+        include_granted_scopes='true'
+    )
+    return auth_url, state
+
+def fetch_token(code):
+    flow = create_flow()
+    flow.fetch_token(code=code)
+    creds = flow.credentials
+    return creds_to_dict(creds)
+
+def creds_to_dict(creds):
+    return {
+        "token": creds.token,
+        "refresh_token": creds.refresh_token,
+        "token_uri": creds.token_uri,
+        "client_id": creds.client_id,
+        "client_secret": creds.client_secret,
+        "scopes": creds.scopes
+    }
+
+def dict_to_creds(data):
+    return Credentials(
+        token=data["token"],
+        refresh_token=data["refresh_token"],
+        token_uri=data["token_uri"],
+        client_id=data["client_id"],
+        client_secret=data["client_secret"],
+        scopes=data["scopes"]
+    )

indexer.py

@@ -14,6 +14,7 @@ from langchain_tavily import TavilySearch
 from langchain_groq import ChatGroq
 from langgraph.checkpoint.memory import MemorySaver
 from langgraph.types import Command, interrupt
+from drive_tools import search_and_download_doc_tool
 import os
 
 
@@ -63,7 +64,7 @@ def send_email_tool(to_email: str, subject: str, body: str) -> str:
     except Exception as e:
         return f"Failed to send email: {str(e)}"
     
-tools = [send_email_tool]
+tools = [search_and_download_doc_tool, send_email_tool]
 llm_with_tools = llm.bind_tools(tools)
 
 # ==================== STATE =======================
@@ -77,12 +78,16 @@ graph_builder=StateGraph(State)
 def chatbot(state:State):
     response = llm_with_tools.invoke([
     SystemMessage(content="""
-    You are an email assistant.
-    When the user wants to send an email:
-    - Extract recipient email
-    - Generate subject
-    - Generate professional body
-    - ALWAYS call send_email_tool
+    You are an AI assistant with access to tools.
+
+    Available tools:
+    1. send_email_tool → Use when the user wants to send an email.
+    2. search_and_download_doc_tool → Use when the user wants to find or download a document from Google Drive.
+
+    Rules:
+    - Always call the appropriate tool when the request requires action.
+    - Do NOT respond with plain text if an action is required.
+    - After tool execution, summarize the result for the user.
     """),
         *state["messages"]
     ])
@@ -115,15 +120,11 @@ graph=graph_builder.compile()
 
 # ==================== ENTRY FUNCTION =======================
 
-def run_email_agent(user_input: str):
+def run_agent(user_input: str):
     result = graph.invoke({
         "messages": [HumanMessage(content=user_input)]
     })
-
-    final_message = result["messages"][-1].content
-    return final_message
-
-
+    return result["messages"][-1].content
 
 # Start
 # LLM + promt -> Chatbot 

requirements.txt

@@ -3,6 +3,7 @@ langchain>=0.3.0
 langsmith
 python-dotenv
 langchain-groq
+gradio
 pandas
 numpy
 scikit-learn
@@ -17,4 +18,9 @@ langchain-core>=0.3.0
 langchain-community>=0.3.0
 langchain-openai
 langchain_tavily>=0.1.0
-langchain_huggingface>=0.1.0
+langchain_huggingface>=0.1.0
+fastapi
+google-auth 
+google-auth-oauthlib 
+google-auth-httplib2 
+google-api-python-client