693 lines
24 KiB
Python
693 lines
24 KiB
Python
#!/usr/bin/env python3
|
|
"""
|
|
Gemini Deep Research Skill - Execute autonomous multi-step research tasks.
|
|
|
|
Uses Google's Deep Research Agent for comprehensive, cited research reports.
|
|
Research tasks take 2-10 minutes but produce detailed analysis.
|
|
"""
|
|
|
|
import argparse
|
|
import asyncio
|
|
import json
|
|
import os
|
|
import re
|
|
import sys
|
|
from datetime import datetime, timezone
|
|
from pathlib import Path
|
|
from typing import Any, AsyncIterator, Callable, Dict, List, Optional
|
|
|
|
try:
|
|
import httpx
|
|
except ImportError:
|
|
print("Error: httpx not installed. Run: pip install httpx", file=sys.stderr)
|
|
sys.exit(1)
|
|
|
|
try:
|
|
from dotenv import load_dotenv
|
|
load_dotenv()
|
|
except ImportError:
|
|
pass # dotenv is optional
|
|
|
|
|
|
class DeepResearchError(Exception):
|
|
"""Custom exception for Deep Research API errors."""
|
|
pass
|
|
|
|
|
|
class HistoryManager:
|
|
"""Manage local research history cache."""
|
|
|
|
def __init__(self, cache_dir: Optional[str] = None):
|
|
default_dir = os.path.expanduser("~/.cache/deep-research")
|
|
self.cache_dir = Path(cache_dir or os.getenv("DEEP_RESEARCH_CACHE_DIR", default_dir))
|
|
self.history_file = self.cache_dir / "history.json"
|
|
self._ensure_cache_dir()
|
|
|
|
def _ensure_cache_dir(self):
|
|
"""Create cache directory if it doesn't exist."""
|
|
self.cache_dir.mkdir(parents=True, exist_ok=True)
|
|
|
|
def _load_history(self) -> Dict:
|
|
"""Load history from file."""
|
|
if self.history_file.exists():
|
|
try:
|
|
return json.loads(self.history_file.read_text())
|
|
except (json.JSONDecodeError, IOError):
|
|
return {"interactions": []}
|
|
return {"interactions": []}
|
|
|
|
def _save_history(self, history: Dict):
|
|
"""Save history to file."""
|
|
self.history_file.write_text(json.dumps(history, indent=2))
|
|
|
|
def add_interaction(self, interaction_id: str, query: str, status: str = "started"):
|
|
"""Add or update an interaction in history."""
|
|
history = self._load_history()
|
|
|
|
# Check if interaction already exists
|
|
for item in history["interactions"]:
|
|
if item["id"] == interaction_id:
|
|
item["status"] = status
|
|
item["updated_at"] = datetime.now(timezone.utc).isoformat().replace("+00:00", "Z")
|
|
self._save_history(history)
|
|
return
|
|
|
|
# Add new interaction
|
|
history["interactions"].insert(0, {
|
|
"id": interaction_id,
|
|
"query": query[:200] + "..." if len(query) > 200 else query,
|
|
"started_at": datetime.now(timezone.utc).isoformat().replace("+00:00", "Z"),
|
|
"status": status
|
|
})
|
|
|
|
# Keep only last 50 interactions
|
|
history["interactions"] = history["interactions"][:50]
|
|
self._save_history(history)
|
|
|
|
def update_status(self, interaction_id: str, status: str):
|
|
"""Update the status of an interaction."""
|
|
history = self._load_history()
|
|
for item in history["interactions"]:
|
|
if item["id"] == interaction_id:
|
|
item["status"] = status
|
|
item["updated_at"] = datetime.now(timezone.utc).isoformat().replace("+00:00", "Z")
|
|
if status == "completed":
|
|
item["completed_at"] = datetime.now(timezone.utc).isoformat().replace("+00:00", "Z")
|
|
break
|
|
self._save_history(history)
|
|
|
|
def get_recent(self, limit: int = 10) -> List[Dict]:
|
|
"""Get recent interactions."""
|
|
history = self._load_history()
|
|
return history["interactions"][:limit]
|
|
|
|
def get_interaction(self, interaction_id: str) -> Optional[Dict]:
|
|
"""Get a specific interaction by ID."""
|
|
history = self._load_history()
|
|
for item in history["interactions"]:
|
|
if item["id"] == interaction_id:
|
|
return item
|
|
return None
|
|
|
|
|
|
class DeepResearchClient:
|
|
"""Client for Gemini Deep Research API."""
|
|
|
|
BASE_URL = "https://generativelanguage.googleapis.com/v1beta/interactions"
|
|
AGENT = "deep-research-pro-preview-12-2025"
|
|
|
|
def __init__(self, api_key: Optional[str] = None):
|
|
self.api_key = api_key or os.getenv("GEMINI_API_KEY")
|
|
if not self.api_key:
|
|
raise DeepResearchError(
|
|
"GEMINI_API_KEY not set. Set it in .env or environment variables."
|
|
)
|
|
self._client: Optional[httpx.AsyncClient] = None
|
|
self.timeout = int(os.getenv("DEEP_RESEARCH_TIMEOUT", "600"))
|
|
self.poll_interval = int(os.getenv("DEEP_RESEARCH_POLL_INTERVAL", "10"))
|
|
self.history = HistoryManager()
|
|
|
|
async def _get_client(self) -> httpx.AsyncClient:
|
|
"""Get or create the HTTP client."""
|
|
if self._client is None or self._client.is_closed:
|
|
self._client = httpx.AsyncClient(timeout=60.0)
|
|
return self._client
|
|
|
|
async def close(self):
|
|
"""Close the HTTP client."""
|
|
if self._client and not self._client.is_closed:
|
|
await self._client.aclose()
|
|
|
|
def _build_prompt(self, query: str, format_spec: Optional[str] = None) -> str:
|
|
"""Build the research prompt with optional format specification."""
|
|
prompt = query
|
|
|
|
if format_spec:
|
|
prompt += f"\n\nFormat the output with the following structure:\n{format_spec}"
|
|
|
|
return prompt
|
|
|
|
async def start_research(
|
|
self,
|
|
query: str,
|
|
format_spec: Optional[str] = None,
|
|
previous_interaction_id: Optional[str] = None
|
|
) -> str:
|
|
"""
|
|
Start a research task.
|
|
|
|
Args:
|
|
query: The research query
|
|
format_spec: Optional output format specification
|
|
previous_interaction_id: Optional ID to continue from previous research
|
|
|
|
Returns:
|
|
Interaction ID for polling results
|
|
"""
|
|
prompt = self._build_prompt(query, format_spec)
|
|
client = await self._get_client()
|
|
|
|
payload: Dict[str, Any] = {
|
|
"input": prompt,
|
|
"agent": self.AGENT,
|
|
"background": True
|
|
}
|
|
|
|
if previous_interaction_id:
|
|
payload["previous_interaction_id"] = previous_interaction_id
|
|
|
|
try:
|
|
response = await client.post(
|
|
self.BASE_URL,
|
|
headers={
|
|
"Content-Type": "application/json",
|
|
"x-goog-api-key": str(self.api_key)
|
|
},
|
|
json=payload
|
|
)
|
|
|
|
if response.status_code != 200:
|
|
error_text = response.text
|
|
raise DeepResearchError(f"API error {response.status_code}: {error_text}")
|
|
|
|
data = response.json()
|
|
interaction_id = data.get("id") or data.get("name", "").split("/")[-1]
|
|
|
|
if not interaction_id:
|
|
raise DeepResearchError("No interaction ID returned from API")
|
|
|
|
# Track in history
|
|
self.history.add_interaction(interaction_id, query, "started")
|
|
|
|
return interaction_id
|
|
|
|
except httpx.HTTPError as e:
|
|
raise DeepResearchError(f"HTTP error: {e}")
|
|
|
|
async def get_status(self, interaction_id: str) -> Dict[str, Any]:
|
|
"""
|
|
Get the status of a research task.
|
|
|
|
Returns:
|
|
Dict with 'status' and optionally 'result' or 'error'
|
|
"""
|
|
client = await self._get_client()
|
|
|
|
try:
|
|
response = await client.get(
|
|
f"{self.BASE_URL}/{interaction_id}",
|
|
headers={"x-goog-api-key": str(self.api_key)},
|
|
timeout=30.0
|
|
)
|
|
|
|
if response.status_code != 200:
|
|
return {"status": "error", "error": f"API error: {response.status_code}"}
|
|
|
|
data = response.json()
|
|
status = data.get("status", "unknown")
|
|
|
|
if status == "completed":
|
|
outputs = data.get("outputs", [])
|
|
if outputs:
|
|
text = outputs[-1].get("text", "")
|
|
return {"status": "completed", "result": text, "raw": data}
|
|
return {"status": "completed", "result": None, "raw": data}
|
|
elif status == "failed":
|
|
return {"status": "failed", "error": data.get("error", "Unknown error")}
|
|
else:
|
|
return {"status": status, "raw": data}
|
|
|
|
except httpx.HTTPError as e:
|
|
return {"status": "error", "error": str(e)}
|
|
|
|
async def wait_for_completion(
|
|
self,
|
|
interaction_id: str,
|
|
timeout: Optional[int] = None,
|
|
poll_interval: Optional[int] = None,
|
|
progress_callback: Optional[Callable[[int, float, str], None]] = None
|
|
) -> Dict[str, Any]:
|
|
"""
|
|
Wait for research to complete with polling.
|
|
|
|
Args:
|
|
interaction_id: The interaction ID to poll
|
|
timeout: Maximum wait time in seconds
|
|
poll_interval: Seconds between polls
|
|
progress_callback: Optional callback for progress updates
|
|
|
|
Returns:
|
|
Dict with research result or error
|
|
"""
|
|
timeout = timeout or self.timeout
|
|
poll_interval = poll_interval or self.poll_interval
|
|
start_time = asyncio.get_event_loop().time()
|
|
poll_count = 0
|
|
|
|
while True:
|
|
elapsed = asyncio.get_event_loop().time() - start_time
|
|
if elapsed > timeout:
|
|
self.history.update_status(interaction_id, "timeout")
|
|
return {"status": "timeout", "error": f"Research timed out after {timeout}s"}
|
|
|
|
result = await self.get_status(interaction_id)
|
|
poll_count += 1
|
|
|
|
if progress_callback:
|
|
progress_callback(poll_count, elapsed, result.get("status", "unknown"))
|
|
|
|
if result["status"] == "completed":
|
|
self.history.update_status(interaction_id, "completed")
|
|
return result
|
|
elif result["status"] in ["failed", "error"]:
|
|
self.history.update_status(interaction_id, "failed")
|
|
return result
|
|
|
|
await asyncio.sleep(poll_interval)
|
|
|
|
async def stream_research(
|
|
self,
|
|
query: str,
|
|
format_spec: Optional[str] = None
|
|
) -> AsyncIterator[Dict[str, Any]]:
|
|
"""
|
|
Stream research progress in real-time.
|
|
|
|
Yields:
|
|
Dict with event type and content
|
|
"""
|
|
prompt = self._build_prompt(query, format_spec)
|
|
client = await self._get_client()
|
|
|
|
try:
|
|
async with client.stream(
|
|
"POST",
|
|
f"{self.BASE_URL}?alt=sse",
|
|
headers={
|
|
"Content-Type": "application/json",
|
|
"x-goog-api-key": str(self.api_key)
|
|
},
|
|
json={
|
|
"input": prompt,
|
|
"agent": self.AGENT,
|
|
"background": True,
|
|
"stream": True,
|
|
"agent_config": {
|
|
"type": "deep-research",
|
|
"thinking_summaries": "auto"
|
|
}
|
|
},
|
|
timeout=None # No timeout for streaming
|
|
) as response:
|
|
interaction_id = None
|
|
buffer = ""
|
|
|
|
async for chunk in response.aiter_text():
|
|
buffer += chunk
|
|
|
|
# Parse SSE events
|
|
while "\n\n" in buffer:
|
|
event_str, buffer = buffer.split("\n\n", 1)
|
|
|
|
if not event_str.strip():
|
|
continue
|
|
|
|
# Parse SSE format
|
|
data_line = None
|
|
for line in event_str.split("\n"):
|
|
if line.startswith("data: "):
|
|
data_line = line[6:]
|
|
break
|
|
|
|
if not data_line:
|
|
continue
|
|
|
|
try:
|
|
event = json.loads(data_line)
|
|
except json.JSONDecodeError:
|
|
continue
|
|
|
|
event_type = event.get("event_type", "")
|
|
|
|
if event_type == "interaction.start":
|
|
interaction_id = event.get("interaction", {}).get("id")
|
|
if interaction_id:
|
|
self.history.add_interaction(interaction_id, query, "streaming")
|
|
yield {"type": "start", "interaction_id": interaction_id}
|
|
|
|
elif event_type == "content.delta":
|
|
delta = event.get("delta", {})
|
|
delta_type = delta.get("type")
|
|
|
|
if delta_type == "text":
|
|
yield {"type": "text", "content": delta.get("text", "")}
|
|
elif delta_type == "thought_summary":
|
|
content = delta.get("content", {})
|
|
yield {"type": "thought", "content": content.get("text", "")}
|
|
|
|
elif event_type == "interaction.complete":
|
|
if interaction_id:
|
|
self.history.update_status(interaction_id, "completed")
|
|
yield {"type": "complete"}
|
|
|
|
elif event_type == "error":
|
|
if interaction_id:
|
|
self.history.update_status(interaction_id, "failed")
|
|
yield {"type": "error", "error": event.get("error", "Unknown error")}
|
|
|
|
except httpx.HTTPError as e:
|
|
yield {"type": "error", "error": str(e)}
|
|
|
|
def parse_result(self, text: str) -> Optional[Dict]:
|
|
"""
|
|
Parse research result text into structured data if JSON.
|
|
|
|
Args:
|
|
text: Raw text output from Gemini
|
|
|
|
Returns:
|
|
Parsed JSON data or None if not JSON
|
|
"""
|
|
if not text:
|
|
return None
|
|
|
|
# Try direct JSON parsing
|
|
try:
|
|
return json.loads(text)
|
|
except json.JSONDecodeError:
|
|
pass
|
|
|
|
# Try extracting JSON from markdown code blocks
|
|
json_match = re.search(r'```(?:json)?\s*(\{.*?\})\s*```', text, re.DOTALL)
|
|
if json_match:
|
|
try:
|
|
return json.loads(json_match.group(1))
|
|
except json.JSONDecodeError:
|
|
pass
|
|
|
|
# Try finding JSON object pattern
|
|
json_match = re.search(r'\{[^{}]*"[^"]+"\s*:[^{}]*\}', text, re.DOTALL)
|
|
if json_match:
|
|
try:
|
|
return json.loads(json_match.group(0))
|
|
except json.JSONDecodeError:
|
|
pass
|
|
|
|
return None
|
|
|
|
|
|
def print_progress(poll_count: int, elapsed: float, status: str):
|
|
"""Print progress update during polling."""
|
|
mins = int(elapsed // 60)
|
|
secs = int(elapsed % 60)
|
|
print(f"\r[{mins:02d}:{secs:02d}] Poll #{poll_count} - Status: {status}", end="", flush=True)
|
|
|
|
|
|
async def cmd_research(args):
|
|
"""Execute a research query."""
|
|
client = DeepResearchClient()
|
|
|
|
try:
|
|
if args.stream:
|
|
# Streaming mode
|
|
print(f"Starting streaming research...\n")
|
|
full_text = ""
|
|
|
|
async for event in client.stream_research(args.query, args.format):
|
|
if event["type"] == "start":
|
|
print(f"Interaction ID: {event['interaction_id']}\n")
|
|
elif event["type"] == "thought":
|
|
print(f"\n[Thinking] {event['content']}\n", file=sys.stderr)
|
|
elif event["type"] == "text":
|
|
print(event["content"], end="", flush=True)
|
|
full_text += event["content"]
|
|
elif event["type"] == "complete":
|
|
print("\n\n[Research Complete]")
|
|
elif event["type"] == "error":
|
|
print(f"\n[Error] {event['error']}", file=sys.stderr)
|
|
return 1
|
|
|
|
# Output JSON if requested
|
|
if args.json and full_text:
|
|
parsed = client.parse_result(full_text)
|
|
if parsed:
|
|
print("\n\n--- Parsed JSON ---")
|
|
print(json.dumps(parsed, indent=2))
|
|
else:
|
|
# Polling mode
|
|
previous_id = args.continue_from if hasattr(args, 'continue_from') else None
|
|
|
|
print(f"Starting research task...")
|
|
interaction_id = await client.start_research(
|
|
args.query,
|
|
args.format,
|
|
previous_id
|
|
)
|
|
print(f"Interaction ID: {interaction_id}")
|
|
print(f"Estimated time: 2-10 minutes\n")
|
|
|
|
if args.no_wait:
|
|
print(f"Research started. Check status with: --status {interaction_id}")
|
|
return 0
|
|
|
|
print("Waiting for completion (Ctrl+C to cancel)...")
|
|
result = await client.wait_for_completion(
|
|
interaction_id,
|
|
progress_callback=print_progress
|
|
)
|
|
print() # New line after progress
|
|
|
|
if result["status"] == "completed":
|
|
text = result.get("result", "")
|
|
|
|
if args.json:
|
|
parsed = client.parse_result(text)
|
|
if parsed:
|
|
print(json.dumps(parsed, indent=2))
|
|
else:
|
|
print(json.dumps({"text": text}, indent=2))
|
|
elif args.raw:
|
|
print(json.dumps(result.get("raw", {}), indent=2))
|
|
else:
|
|
print("\n--- Research Result ---\n")
|
|
print(text)
|
|
|
|
return 0
|
|
else:
|
|
print(f"\nResearch failed: {result.get('error', 'Unknown error')}")
|
|
return 1
|
|
|
|
except DeepResearchError as e:
|
|
print(f"Error: {e}", file=sys.stderr)
|
|
return 1
|
|
except KeyboardInterrupt:
|
|
print("\n\nCancelled by user.")
|
|
return 130
|
|
finally:
|
|
await client.close()
|
|
|
|
|
|
async def cmd_status(args):
|
|
"""Check status of a research task."""
|
|
client = DeepResearchClient()
|
|
|
|
try:
|
|
result = await client.get_status(args.interaction_id)
|
|
|
|
if args.json:
|
|
print(json.dumps(result, indent=2))
|
|
else:
|
|
print(f"Status: {result['status']}")
|
|
|
|
if result["status"] == "completed":
|
|
text = result.get("result", "")
|
|
if text:
|
|
print(f"\n--- Result Preview ---\n{text[:500]}...")
|
|
elif result["status"] in ["failed", "error"]:
|
|
print(f"Error: {result.get('error', 'Unknown')}")
|
|
|
|
return 0 if result["status"] != "error" else 1
|
|
|
|
except DeepResearchError as e:
|
|
print(f"Error: {e}", file=sys.stderr)
|
|
return 1
|
|
finally:
|
|
await client.close()
|
|
|
|
|
|
async def cmd_wait(args):
|
|
"""Wait for a research task to complete."""
|
|
client = DeepResearchClient()
|
|
|
|
try:
|
|
print(f"Waiting for {args.interaction_id}...")
|
|
result = await client.wait_for_completion(
|
|
args.interaction_id,
|
|
progress_callback=print_progress
|
|
)
|
|
print()
|
|
|
|
if result["status"] == "completed":
|
|
text = result.get("result", "")
|
|
|
|
if args.json:
|
|
parsed = client.parse_result(text)
|
|
if parsed:
|
|
print(json.dumps(parsed, indent=2))
|
|
else:
|
|
print(json.dumps({"text": text}, indent=2))
|
|
else:
|
|
print("\n--- Research Result ---\n")
|
|
print(text)
|
|
|
|
return 0
|
|
else:
|
|
print(f"Research failed: {result.get('error', 'Unknown error')}")
|
|
return 1
|
|
|
|
except DeepResearchError as e:
|
|
print(f"Error: {e}", file=sys.stderr)
|
|
return 1
|
|
except KeyboardInterrupt:
|
|
print("\n\nCancelled by user.")
|
|
return 130
|
|
finally:
|
|
await client.close()
|
|
|
|
|
|
async def cmd_list(args):
|
|
"""List recent research tasks."""
|
|
history = HistoryManager()
|
|
interactions = history.get_recent(args.limit)
|
|
|
|
if not interactions:
|
|
print("No research history found.")
|
|
return 0
|
|
|
|
if args.json:
|
|
print(json.dumps(interactions, indent=2))
|
|
else:
|
|
print(f"Recent research tasks (last {len(interactions)}):\n")
|
|
for item in interactions:
|
|
status_icon = {
|
|
"completed": "[ok]",
|
|
"failed": "[!!]",
|
|
"started": "[..]",
|
|
"streaming": "[>>]",
|
|
"timeout": "[to]"
|
|
}.get(item.get("status", ""), "[??]")
|
|
|
|
print(f"{status_icon} {item['id'][:12]}...")
|
|
print(f" Query: {item['query'][:60]}{'...' if len(item['query']) > 60 else ''}")
|
|
print(f" Started: {item.get('started_at', 'N/A')}")
|
|
if item.get("completed_at"):
|
|
print(f" Completed: {item['completed_at']}")
|
|
print()
|
|
|
|
return 0
|
|
|
|
|
|
def main():
|
|
parser = argparse.ArgumentParser(
|
|
description="Gemini Deep Research - Autonomous multi-step research agent",
|
|
formatter_class=argparse.RawDescriptionHelpFormatter,
|
|
epilog="""
|
|
Examples:
|
|
# Start a research task and wait for results
|
|
%(prog)s --query "Research the history of Kubernetes"
|
|
|
|
# With structured output format
|
|
%(prog)s --query "Compare Python web frameworks" \\
|
|
--format "1. Executive Summary\\n2. Comparison Table\\n3. Recommendations"
|
|
|
|
# Stream progress in real-time
|
|
%(prog)s --query "Analyze EV battery market" --stream
|
|
|
|
# Start without waiting
|
|
%(prog)s --query "Research topic" --no-wait
|
|
|
|
# Check status of running research
|
|
%(prog)s --status abc123
|
|
|
|
# Continue from previous research
|
|
%(prog)s --query "Elaborate on point 2" --continue abc123
|
|
|
|
# List recent research
|
|
%(prog)s --list
|
|
|
|
Note: Research tasks typically take 2-10 minutes and cost $2-5 per task.
|
|
"""
|
|
)
|
|
|
|
# Main commands (mutually exclusive)
|
|
cmd_group = parser.add_mutually_exclusive_group(required=True)
|
|
cmd_group.add_argument("--query", "-q", help="Research query to execute")
|
|
cmd_group.add_argument("--status", "-s", dest="interaction_id", metavar="ID",
|
|
help="Check status of a research task")
|
|
cmd_group.add_argument("--wait", "-w", dest="wait_id", metavar="ID",
|
|
help="Wait for a research task to complete")
|
|
cmd_group.add_argument("--list", "-l", action="store_true",
|
|
help="List recent research tasks")
|
|
|
|
# Query options
|
|
parser.add_argument("--format", "-f", metavar="SPEC",
|
|
help="Output format specification")
|
|
parser.add_argument("--continue", dest="continue_from", metavar="ID",
|
|
help="Continue from previous research interaction")
|
|
parser.add_argument("--stream", action="store_true",
|
|
help="Stream progress in real-time")
|
|
parser.add_argument("--no-wait", action="store_true",
|
|
help="Start research without waiting for completion")
|
|
|
|
# Output options
|
|
parser.add_argument("--json", "-j", action="store_true",
|
|
help="Output results as JSON")
|
|
parser.add_argument("--raw", "-r", action="store_true",
|
|
help="Output raw API response")
|
|
|
|
# List options
|
|
parser.add_argument("--limit", type=int, default=10,
|
|
help="Number of recent tasks to show (default: 10)")
|
|
|
|
args = parser.parse_args()
|
|
|
|
# Route to appropriate command
|
|
if args.query:
|
|
exit_code = asyncio.run(cmd_research(args))
|
|
elif args.interaction_id:
|
|
exit_code = asyncio.run(cmd_status(args))
|
|
elif args.wait_id:
|
|
args.interaction_id = args.wait_id
|
|
exit_code = asyncio.run(cmd_wait(args))
|
|
elif args.list:
|
|
exit_code = asyncio.run(cmd_list(args))
|
|
else:
|
|
parser.print_help()
|
|
exit_code = 1
|
|
|
|
sys.exit(exit_code)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|