New agentic features

agent.py

@@ -71,7 +71,6 @@ llm_with_tools = llm.bind_tools(tools)
 class State(TypedDict):
     messages: Annotated[list,add_messages]
 
-graph_builder=StateGraph(State)
 
 # ==================== NODES =======================
 
@@ -93,9 +92,65 @@ def chatbot(state:State):
     ])
     return {"messages":[response]}
 
+def handle_tools(state: State):
+    """
+    Custom tool-execution node that intercepts AUTH_REQUIRED:: sentinels
+    from the Drive tool and surfaces them via LangGraph interrupt instead
+    of letting the agent loop silently.
+    """
+    last_message: AIMessage = state["messages"][-1]
+ 
+    results = []
+    for tool_call in last_message.tool_calls:
+        # ── Run the tool ──────────────────────────────────────────────────────
+        matched_tool = next(
+            (t for t in tools if t.name == tool_call["name"]), None
+        )
+        if matched_tool is None:
+            result_content = f"Unknown tool: {tool_call['name']}"
+        else:
+            result_content = matched_tool.invoke(tool_call["args"])
+ 
+        # ── Auth-gate check ───────────────────────────────────────────────────
+        if isinstance(result_content, str) and result_content.startswith(AUTH_REQUIRED_PREFIX):
+            # Extract the OAuth URL (everything after the sentinel prefix, first line)
+            first_line = result_content.split("\n")[0]
+            auth_url = first_line.removeprefix(AUTH_REQUIRED_PREFIX).strip()
+ 
+            # Interrupt the graph and surface the URL to the front-end
+            # The interrupt value is returned to whoever is streaming the graph.
+            interrupt({
+                "type": "auth_required",
+                "auth_url": auth_url,
+                "message": (
+                    "🔐 Google Drive access is required. "
+                    "Please authenticate by visiting the link below, then retry your request.\n\n"
+                    f"👉 {auth_url}"
+                ),
+            })
+            # After the user resumes (post-OAuth), return a helpful ToolMessage
+            result_content = (
+                "Authentication flow initiated. Once you have completed Google sign-in, "
+                "please repeat your Drive request."
+            )
+ 
+        results.append(
+            ToolMessage(
+                content=result_content,
+                tool_call_id=tool_call["id"],
+            )
+        )
+ 
+    return {"messages": results}
+ 
+
 # ==================== GRAPH =======================
 
 # Adding Node
+memory = MemorySaver()
+
+graph_builder=StateGraph(State)
+
 graph_builder.add_node("chatbot", chatbot)
 
 tool_node = ToolNode(tools=tools)
@@ -116,7 +171,7 @@ graph_builder.add_conditional_edges(
 
 graph_builder.add_edge("tools","chatbot")
 
-graph=graph_builder.compile()
+graph=graph_builder.compile(checkpointer=memory)
 
 # ==================== ENTRY FUNCTION =======================
 

app.py

@@ -1,33 +1,68 @@
 import gradio as gr
-from fastapi import FastAPI, Request
-from google_auth import get_auth_url, fetch_token
-from agent import run_agent
-import json
+from agent import graph
+from oauth_callback import handle_oauth_callback
 
-app = FastAPI()
+# ── Persistent thread so MemorySaver keeps conversation context ───────────────
+THREAD_CONFIG = {"configurable": {"thread_id": "default-thread"}}
 
-@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)
-
-demo = gr.Interface(
-    fn=chat_interface,
-    inputs="text",
-    outputs="text"
-)
-
-demo.launch(server_name="0.0.0.0", server_port=7860)
+def chat(user_message: str, history: list) -> str:
+    """Called by the Gradio ChatInterface on each user message."""
+    events = graph.stream(
+        {"messages": [{"role": "user", "content": user_message}]},
+        config=THREAD_CONFIG,
+        stream_mode="values",
+    )
+ 
+    last_ai_text = ""
+    for event in events:
+        # Check for an interrupt (auth required)
+        if "__interrupt__" in event:
+            interrupt_val = event["__interrupt__"][0].value
+            if interrupt_val.get("type") == "auth_required":
+                return interrupt_val["message"]
+ 
+        msgs = event.get("messages", [])
+        for msg in reversed(msgs):
+            if hasattr(msg, "content") and msg.type == "ai" and not msg.tool_calls:
+                last_ai_text = msg.content
+                break
+ 
+    return last_ai_text or "Done."
+ 
+ 
+# ── OAuth callback endpoint ───────────────────────────────────────────────────
+ 
+def oauth_callback_page(request: gr.Request) -> str:
+    """
+    Gradio page that Google redirects to after the user grants consent.
+    Mount at /oauth/callback in your Space.
+    """
+    params = dict(request.query_params)
+    code  = params.get("code", "")
+    state = params.get("state", "")
+    result = handle_oauth_callback(code, state)
+    if result["success"]:
+        return f"<h2>{result['message']}</h2><p>You can close this tab and return to the chat.</p>"
+    return f"<h2>❌ Authentication failed</h2><p>{result['message']}</p>"
+ 
+ 
+# ── Gradio UI ─────────────────────────────────────────────────────────────────
+ 
+with gr.Blocks(title="AI Agent") as demo:
+    gr.Markdown("## 🤖 AI Agent  |  Email · Google Drive")
+ 
+    with gr.Tab("Chat"):
+        gr.ChatInterface(fn=chat)
+ 
+    # Hidden page — Google redirects here after OAuth
+    with gr.Tab("OAuth Callback", visible=False) as callback_tab:
+        callback_output = gr.HTML()
+ 
+    # Route /oauth/callback → the handler above
+    demo.load(fn=None)   # placeholder; real routing done via gr.mount_gradio_app or FastAPI
+ 
+ 
+# ── For local dev, run directly ───────────────────────────────────────────────
+if __name__ == "__main__":
+    demo.launch(server_name="0.0.0.0", server_port=7860)

drive_tools.py

@@ -1,49 +1,129 @@
-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
+from langchain_core.tools import tool
+from googleapiclient.discovery import build
+from googleapiclient.http import MediaIoBaseDownload
+from googleapiclient.errors import HttpError
+ 
+from supabase_auth import get_token, save_token
+from google_auth_flow import credentials_from_token_dict, get_auth_url
 
-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)
+AUTH_REQUIRED_PREFIX = "AUTH_REQUIRED::"
+DOWNLOAD_DIR = os.getenv("DOWNLOAD_DIR", "/tmp/drive_downloads")
 
-    creds = dict_to_creds(token_data)
-    service = build("drive", "v3", credentials=creds)
+def _drive_service(user_email: str):
+    """Returns an auth Drive service"""
+    token_dict = get_token(user_email)
+    if not token_dict:
+        return None
+    creds = credentials_from_token_dict(token_dict)
+    # Persist refreshed token back to Supabase
+    save_token(user_email, {
+        "token": creds.token,
+        "refresh_token": creds.refresh_token,
+        "token_uri": creds.token_uri,
+        "client_id": creds.client_id,
+        "client_secret": creds.client_secret,
+        "scopes": list(creds.scopes or []),
+        "expiry": creds.expiry.isoformat() if creds.expiry else None,
+    })
+    return build("drive", "v3", credentials=creds)
 
+def _search_files(service, query: str, max_results: int = 5) -> list[dict]:
+    """Full-text search across Drive files."""
     results = service.files().list(
-        q=f"name contains '{file_name}'",
-        fields="files(id, name)"
+        q=f"fullText contains '{query}' and trashed=false",
+        pageSize=max_results,
+        fields="files(id, name, mimeType, webViewLink, size)",
     ).execute()
+    return results.get("files", [])
 
-    items = results.get("files", [])
-
-    if not items:
-        return "No document found."
-
-    file_id = items[0]["id"]
-    request = service.files().get_media(fileId=file_id)
+def _download_file(service, file_id: str, file_name: str, mime_type: str) -> str:
+    """Downloads a file and returns local path."""
+    os.makedirs(DOWNLOAD_DIR, exist_ok=True)
+    local_path = os.path.join(DOWNLOAD_DIR, file_name)
+    
+    # Google Workspace docs must be exported; regular files use get_media
+    export_map = {
+        "application/vnd.google-apps.document":
+            "application/vnd.openxmlformats-officedocument.wordprocessingml.document",
+        "application/vnd.google-apps.spreadsheet":
+            "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet",
+        "application/vnd.google-apps.presentation":
+            "application/vnd.openxmlformats-officedocument.presentationml.presentation",
+    }
 
-    file_path = f"./downloads/{items[0]['name']}"
-    os.makedirs("downloads", exist_ok=True)
-
-    fh = io.FileIO(file_path, "wb")
+    fh = io.BytesIO()
+    if mime_type in export_map:
+        export_mime = export_map[mime_type]
+        request = service.files().export_media(fileId=file_id, mimeType=export_mime)
+        ext_map = {
+            "application/vnd.openxmlformats-officedocument.wordprocessingml.document": ".docx",
+            "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet": ".xlsx",
+            "application/vnd.openxmlformats-officedocument.presentationml.presentation": ".pptx",
+        }
+        local_path += ext_map.get(export_mime, "")
+    else:
+        request = service.files().get_media(fileId=file_id)
+ 
     downloader = MediaIoBaseDownload(fh, request)
-
     done = False
     while not done:
-        status, done = downloader.next_chunk()
+        _, done = downloader.next_chunk()
+ 
+    with open(local_path, "wb") as f:
+        f.write(fh.getvalue())
+ 
+    return local_path
+
+
+
+@tool
+def search_and_download_doc_tool(user_email: str, query: str) -> str:
+    """
+    Searches Google Drive and downloads a document by name.
+    """
 
-    return f"Downloaded to {file_path}"
+    service = _drive_service(user_email)
+    if service is None:
+        auth_url = get_auth_url(state=user_email)          # state carries email through OAuth
+        return (
+            f"{AUTH_REQUIRED_PREFIX}{auth_url}\n"
+            f"User {user_email} is not authenticated with Google Drive. "
+            f"They must visit the URL above to grant access."
+        )
+    try:
+        files = _search_files(service, query)
+    except HttpError as e:
+        return f"Drive search failed: {e}"
+ 
+    if not files:
+        return f"No files found on Google Drive matching '{query}'."
+ 
+    # Pick the first (most relevant) result
+    best = files[0]
+    file_id   = best["id"]
+    file_name = best["name"]
+    mime_type = best["mimeType"]
+    view_link = best.get("webViewLink", "N/A")
+ 
+    other_matches = [f["name"] for f in files[1:]]
+    other_str = (
+        f"\n\nOther matches: {', '.join(other_matches)}" if other_matches else ""
+    )
+ 
+    # ── 3. Download ────────────────────────────────────────────────────────────
+    try:
+        local_path = _download_file(service, file_id, file_name, mime_type)
+        return (
+            f"✅ Found and downloaded '{file_name}'.\n"
+            f"Saved to: {local_path}\n"
+            f"View online: {view_link}"
+            f"{other_str}"
+        )
+    except HttpError as e:
+        return (
+            f"Found '{file_name}' on Drive but download failed: {e}\n"
+            f"View online: {view_link}"
+        )

google_auth.py

@@ -1,60 +0,0 @@
-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"]
-    )

google_auth_flow.py

@@ -0,0 +1,90 @@
+import os
+import json
+from google_auth_oauthlib.flow import Flow
+from google.oauth2.credentials import Credentials
+from google.auth.transport.requests import Request
+from dotenv import load_dotenv
+
+SCOPES = [
+    "https://www.googleapis.com/auth/drive.readonly",
+    "https://www.googleapis.com/auth/userinfo.email",
+    "openid",
+]
+
+load_dotenv()
+ 
+CLIENT_ID = os.getenv("GOOGLE_CLIENT_ID", "")
+CLIENT_SECRET = os.getenv("GOOGLE_CLIENT_SECRET", "")
+REDIRECT_URI = os.getenv("GOOGLE_REDIRECT_URI", "")
+
+def _client_config() -> dict:
+    """Builds the client config dict that google_auth_oauthlib expects."""
+    return {
+        "web": {
+            "client_id": CLIENT_ID,
+            "client_secret": CLIENT_SECRET,
+            "redirect_uris": [REDIRECT_URI],
+            "auth_uri": "https://accounts.google.com/o/oauth2/auth",
+            "token_uri": "https://oauth2.googleapis.com/token",
+        }
+    }
+
+def get_auth_url(state: str | None = None) -> str:
+    """
+    Returns the Google OAuth consent-screen URL to redirect the user to.
+    `state` can carry any context you want back in the callback (e.g. user_email).
+    """
+    flow = Flow.from_client_config(_client_config(), scopes=SCOPES)
+    flow.redirect_uri = REDIRECT_URI
+    auth_url, _ = flow.authorization_url(
+        access_type="offline",   # get refresh_token
+        include_granted_scopes="true",
+        prompt="consent",        # force refresh_token every time during dev
+        state=state or "",
+    )
+    return auth_url
+
+def exchange_code_for_token(code: str) -> dict:
+    """
+    Exchanges an authorization code (from the OAuth callback) for credentials.
+    Returns a JSON-serialisable token dict.
+    """
+    flow = Flow.from_client_config(_client_config(), scopes=SCOPES)
+    flow.redirect_uri = REDIRECT_URI
+    flow.fetch_token(code=code)
+    creds = flow.credentials
+    return _creds_to_dict(creds)
+
+def credentials_from_token_dict(token_dict: dict) -> Credentials:
+    """
+    Re-hydrates a Credentials object from a stored token dict,
+    refreshing automatically if the access token is expired.
+    """
+    creds = Credentials(
+        token=token_dict.get("token"),
+        refresh_token=token_dict.get("refresh_token"),
+        token_uri="https://oauth2.googleapis.com/token",
+        client_id=CLIENT_ID,
+        client_secret=CLIENT_SECRET,
+        scopes=token_dict.get("scopes", SCOPES),
+    )
+    if creds.expired and creds.refresh_token:
+        creds.refresh(Request())
+    return creds
+
+def _creds_to_dict(creds: Credentials) -> dict:
+    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": list(creds.scopes or SCOPES),
+        "expiry": creds.expiry.isoformat() if creds.expiry else None,
+    }
+
+
+
+
+
+

oauth_callback.py

@@ -0,0 +1,41 @@
+"""Handles the google OAuth 2.0 redirect callback"""
+
+from google_auth_flow import exchange_code_for_token
+from supabase_auth import save_token
+
+def handle_oauth_callback(code: str, state: str) -> dict:
+    """
+    Exchanges the OAuth authorization code for tokens and persists them.
+ 
+    Args:
+        code:  The `code` query parameter from the callback URL.
+        state: The `state` parameter — we use it to carry the user's email.
+ 
+    Returns:
+        A dict with keys:
+          - success (bool)
+          - user_email (str)
+          - message (str)
+    """
+    user_email = state  # we set state=user_email when building the auth URL
+ 
+    if not code:
+        return {"success": False, "user_email": user_email, "message": "No authorization code received."}
+    if not user_email:
+        return {"success": False, "user_email": "", "message": "No user email in OAuth state parameter."}
+ 
+    try:
+        token_dict = exchange_code_for_token(code)
+        save_token(user_email, token_dict)
+        return {
+            "success": True,
+            "user_email": user_email,
+            "message": f"✅ Google Drive access granted and token saved for {user_email}. You can now search your Drive.",
+        }
+    except Exception as e:
+        return {
+            "success": False,
+            "user_email": user_email,
+            "message": f"OAuth token exchange failed: {str(e)}",
+        }
+ 

supabase_auth.py

@@ -0,0 +1,52 @@
+"""Handles storing and retrieving Google OAuth tokens from Supabase."""
+
+import os
+import json
+from supabase import create_client, Client
+from dotenv import load_dotenv
+ 
+load_dotenv()
+ 
+SUPABASE_URL: str = os.getenv("SUPABASE_URL", "")
+SUPABASE_KEY: str = os.getenv("SUPABASE_SERVICE_KEY", "")
+
+def _get_client() -> Client:
+    if not SUPABASE_URL or not SUPABASE_KEY:
+        raise EnvironmentError("SUPABASE_URL and SUPABASE_SERVICE_KEY must be set in .env")
+    return create_client(SUPABASE_URL, SUPABASE_KEY)
+
+def get_token(user_email: str) -> dict | None:
+    """
+    Returns the stored Google OAuth token dict for the user, or None if not found.
+    """
+    client = _get_client()
+    result = (
+        client.table(TABLE)
+        .select("token_json")
+        .eq("user_email", user_email)
+        .maybe_single()
+        .execute()
+    )
+    if result.data:
+        return json.loads(result.data["token_json"])
+    return None
+
+def save_token(user_email: str, token_dict: dict) -> None:
+    """
+    Upserts (insert or update) the Google OAuth token for the user.
+    """
+    client = _get_client()
+    client.table(TABLE).upsert(
+        {
+            "user_email": user_email,
+            "token_json": json.dumps(token_dict),
+        }
+    ).execute()
+
+
+def delete_token(user_email: str) -> None:
+    """
+    Removes the stored token (e.g. on logout or revocation).
+    """
+    client = _get_client()
+    client.table(TABLE).delete().eq("user_email", user_email).execute()