Spaces:
Running
Running
| """ | |
| Google Drive Agent – specialized agent with only Drive MCP tools. | |
| Uses MCPServerStdio (local subprocess) with create_static_tool_filter so the | |
| agent sees *only* file-management tools. Zero network overhead. | |
| """ | |
| import asyncio | |
| import logging | |
| from openai import AsyncOpenAI | |
| from agents import Agent, Runner, OpenAIChatCompletionsModel | |
| from agents.model_settings import ModelSettings | |
| try: | |
| from .google_mcp_config import ( | |
| LONGCAT_API_KEY, LONGCAT_BASE_URL, MODEL_NAME, | |
| DRIVE_TOOLS, create_google_mcp_server, USER_GOOGLE_EMAIL, | |
| ) | |
| except ImportError: | |
| from google_mcp_config import ( | |
| LONGCAT_API_KEY, LONGCAT_BASE_URL, MODEL_NAME, | |
| DRIVE_TOOLS, create_google_mcp_server, USER_GOOGLE_EMAIL, | |
| ) | |
| logger = logging.getLogger(__name__) | |
| SYSTEM_PROMPT = """\ | |
| You are a specialized Google Drive assistant. You can search, browse, create, | |
| share, and manage files and folders in Google Drive using the available tools. | |
| Capabilities: | |
| - **Search** files by name, type, content, or owner | |
| - **List** items inside any folder (with pagination) | |
| - **Read** file content (Google Docs, Sheets, text files, etc.) | |
| - **Download** — get direct download URLs for any file | |
| - **Create** files and import external files as Google Docs | |
| - **Update** file metadata (name, description, parent folder) | |
| - **Copy** files to new locations | |
| - **Permissions** — view, share, batch-share, update, and remove | |
| collaborators with specific roles (viewer, commenter, editor, owner) | |
| - **Shareable links** — generate public or restricted links | |
| - **Transfer ownership** of files between users | |
| Rules: | |
| 1. The user's Google email is provided in the query — use it for every | |
| tool call in the `user_google_email` parameter. NEVER ask the user | |
| for their email; it is always supplied. | |
| 2. When sharing files, double-check the email and role before calling the | |
| share tool to avoid accidental exposure. | |
| 3. Use `check_drive_file_public_access` before sharing to verify current | |
| visibility. | |
| 4. For bulk operations (sharing with many users), prefer `batch_share_drive_file`. | |
| 5. After any change, confirm the action with relevant details (file name, | |
| new permissions, etc.). | |
| """ | |
| class GoogleDriveAgent: | |
| """Thin wrapper around the OpenAI Agent SDK wired to Google Drive tools.""" | |
| def __init__(self, model: str = MODEL_NAME): | |
| self.model = model | |
| self._client = AsyncOpenAI( | |
| api_key=LONGCAT_API_KEY, | |
| base_url=LONGCAT_BASE_URL, | |
| timeout=30.0, | |
| ) | |
| # ── factory helpers ────────────────────────────────────────────────── | |
| def _create_mcp_server(self): | |
| """Spawn a local MCP subprocess with only Drive tools loaded.""" | |
| return create_google_mcp_server(service="drive", tool_names=DRIVE_TOOLS) | |
| def _create_agent(self, mcp_server) -> Agent: | |
| return Agent( | |
| name="Google Drive Agent", | |
| instructions=SYSTEM_PROMPT, | |
| mcp_servers=[mcp_server], | |
| model=OpenAIChatCompletionsModel( | |
| model=self.model, | |
| openai_client=self._client, | |
| ), | |
| model_settings=ModelSettings(tool_choice="auto"), | |
| ) | |
| # ── public API ─────────────────────────────────────────────────────── | |
| async def run(self, query: str) -> str: | |
| """Spawn MCP connection, run a single query, then clean up.""" | |
| mcp_server = self._create_mcp_server() | |
| async with mcp_server: | |
| agent = self._create_agent(mcp_server) | |
| logger.info("Google Drive MCP connected – agent ready") | |
| result = await Runner.run(agent, input=query) | |
| return result.final_output | |
| # ─── CLI entry point ────────────────────────────────────────────────────────── | |
| async def main(): | |
| agent = GoogleDriveAgent() | |
| resp = await agent.run( | |
| "List all files in my Drive root folder. My email is user@example.com" | |
| ) | |
| print("Agent Response:\n", resp) | |
| if __name__ == "__main__": | |
| from dotenv import load_dotenv, find_dotenv | |
| load_dotenv(find_dotenv()) | |
| asyncio.run(main()) | |