#!/usr/bin/env python3 """Tavily search helper for OpenClaw. Goals: - Stable: prefer Tavily REST API. - Secret-safe: never prints API key. - Frequent-use friendly: - Reads Tavily key from Bitwarden item `openclaw/tavily` when possible. - Auto-reads a local `.bw_session` file (if present & fresh) to avoid manual env exports. - Lightweight cache + monthly usage counter to reduce calls and avoid quota surprises. Outputs JSON. """ from __future__ import annotations import argparse import calendar import hashlib import json import os import subprocess import sys import time from dataclasses import dataclass from pathlib import Path from typing import Any, Dict, Optional from urllib import request STATE_DIR = Path(os.environ.get("OPENCLAW_STATE_DIR", str(Path.home() / ".openclaw"))) BW_SESSION_FILE = Path("/home/kuhnn/.openclaw/.bw_session") DEFAULT_CACHE_DIR = STATE_DIR / "cache" / "tavily" DEFAULT_USAGE_PATH = STATE_DIR / "tavily-usage.json" TAVILY_ENDPOINT = "https://api.tavily.com/search" @dataclass class Usage: year: int month: int count: int def _now_ym() -> tuple[int, int]: t = time.gmtime() return t.tm_year, t.tm_mon def load_usage(path: Path) -> Usage: y, m = _now_ym() if not path.exists(): return Usage(y, m, 0) try: obj = json.loads(path.read_text(encoding="utf-8")) if obj.get("year") == y and obj.get("month") == m: return Usage(y, m, int(obj.get("count", 0))) return Usage(y, m, 0) except Exception: return Usage(y, m, 0) def bump_usage(path: Path, inc: int = 1) -> Usage: path.parent.mkdir(parents=True, exist_ok=True) u = load_usage(path) u.count += inc path.write_text(json.dumps({"year": u.year, "month": u.month, "count": u.count}), encoding="utf-8") return u def _sha1(s: str) -> str: return hashlib.sha1(s.encode("utf-8")).hexdigest() def cache_get(cache_dir: Path, key: str, ttl_sec: int) -> Optional[Dict[str, Any]]: p = cache_dir / f"{_sha1(key)}.json" if not p.exists(): return None try: st = p.stat() if (time.time() - st.st_mtime) > ttl_sec: return None return json.loads(p.read_text(encoding="utf-8")) except Exception: return None def cache_set(cache_dir: Path, key: str, value: Dict[str, Any]) -> None: cache_dir.mkdir(parents=True, exist_ok=True) p = cache_dir / f"{_sha1(key)}.json" p.write_text(json.dumps(value, ensure_ascii=False), encoding="utf-8") def read_bw_session_if_fresh(max_age_sec: int = 3600) -> Optional[str]: try: if not BW_SESSION_FILE.exists(): return None age = time.time() - BW_SESSION_FILE.stat().st_mtime if age > max_age_sec: return None s = BW_SESSION_FILE.read_text(encoding="utf-8").strip() return s or None except Exception: return None def bw_get_tavily_key(bw_session: str) -> Optional[str]: """Read Tavily key from Bitwarden CLI. Expects item name contains `openclaw/tavily` and stores API key in login.password. """ env = os.environ.copy() env["BW_SESSION"] = bw_session try: out = subprocess.check_output( ["bw", "list", "items", "--search", "openclaw/tavily"], text=True, env=env, stderr=subprocess.DEVNULL, ) items = json.loads(out) if not items: return None item = items[0] login = item.get("login") or {} key = login.get("password") return key or None except Exception: return None def resolve_api_key() -> Optional[str]: # 1) explicit env k = os.environ.get("TAVILY_API_KEY") if k: return k.strip() or None # 2) Bitwarden via BW_SESSION env bw = os.environ.get("BW_SESSION") if bw: k2 = bw_get_tavily_key(bw) if k2: return k2 # 3) auto-read local .bw_session if fresh bw2 = read_bw_session_if_fresh() if bw2: k3 = bw_get_tavily_key(bw2) if k3: return k3 return None def tavily_search(api_key: str, query: str, max_results: int, depth: str, include_answer: bool, include_images: bool) -> Dict[str, Any]: payload = { "api_key": api_key, "query": query, "max_results": max_results, "search_depth": depth, "include_answer": include_answer, "include_images": include_images, "include_raw_content": False, } data = json.dumps(payload).encode("utf-8") req = request.Request( TAVILY_ENDPOINT, data=data, headers={"Content-Type": "application/json"}, method="POST", ) with request.urlopen(req, timeout=45) as resp: body = resp.read().decode("utf-8", errors="replace") # Tavily returns JSON return json.loads(body) def main() -> int: ap = argparse.ArgumentParser() ap.add_argument("--query", required=True) ap.add_argument("--max-results", type=int, default=5) ap.add_argument("--depth", choices=["basic", "advanced"], default="basic") ap.add_argument("--include-answer", action="store_true") ap.add_argument("--include-images", action="store_true") ap.add_argument("--cache-ttl-sec", type=int, default=3600) ap.add_argument("--cache-dir", default=str(DEFAULT_CACHE_DIR)) ap.add_argument("--usage-path", default=str(DEFAULT_USAGE_PATH)) ap.add_argument("--monthly-cap", type=int, default=2000) args = ap.parse_args() api_key = resolve_api_key() if not api_key: print( "ERROR: missing Tavily API key. Store it in Bitwarden item 'openclaw/tavily' (login.password) and ensure .bw_session is present & fresh, or set TAVILY_API_KEY.", file=sys.stderr, ) return 2 cache_dir = Path(args.cache_dir) usage_path = Path(args.usage_path) # monthly cap guard (best-effort; local counter) u = load_usage(usage_path) if u.count >= args.monthly_cap: print( json.dumps( { "error": "monthly_cap_reached", "message": f"Local counter indicates Tavily monthly cap reached ({u.count}/{args.monthly_cap}).", "results": [], }, ensure_ascii=False, ) ) return 0 cache_key = json.dumps( { "q": args.query, "max": args.max_results, "depth": args.depth, "ans": bool(args.include_answer), "img": bool(args.include_images), }, sort_keys=True, ensure_ascii=False, ) cached = cache_get(cache_dir, cache_key, args.cache_ttl_sec) if cached is not None: cached["_cache"] = {"hit": True, "ttl_sec": args.cache_ttl_sec} print(json.dumps(cached, ensure_ascii=False)) return 0 try: out = tavily_search( api_key=api_key, query=args.query, max_results=args.max_results, depth=args.depth, include_answer=bool(args.include_answer), include_images=bool(args.include_images), ) except Exception as e: # Do not leak secrets. print( json.dumps( { "error": "tavily_request_failed", "message": str(e), "results": [], }, ensure_ascii=False, ) ) return 0 bump_usage(usage_path, 1) out["_cache"] = {"hit": False, "ttl_sec": args.cache_ttl_sec} out["_usage"] = {"year": _now_ym()[0], "month": _now_ym()[1]} cache_set(cache_dir, cache_key, out) print(json.dumps(out, ensure_ascii=False)) return 0 if __name__ == "__main__": raise SystemExit(main())