cli-jaw 0.1.0

package/skills_ref/deep-research/scripts/research.py

@@ -0,0 +1,692 @@
+#!/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()