Files
deep_research/scripts/ground.py
T

69 lines
2.2 KiB
Python

#!/usr/bin/env python3
"""Native web grounding wrapper via ZenMux chat completions.
Use this when you need reproducible, model-native web search (grounding) and
machine-readable citations.
"""
from __future__ import annotations
import argparse
import json
from pathlib import Path
from scripts.lib.zenmux_client import ZenMuxClient, load_secrets
def build_parser() -> argparse.ArgumentParser:
parser = argparse.ArgumentParser(description="Grounded web query via ZenMux")
parser.add_argument("query", help="Question or search prompt")
parser.add_argument("--model", default="google/gemini-3.1-flash-lite-preview")
parser.add_argument("--max-tokens", type=int, default=2400)
parser.add_argument("--temperature", type=float, default=0.2)
parser.add_argument("--json", action="store_true", help="Emit JSON envelope")
parser.add_argument("--log-file", help="Optional JSONL call log path")
parser.add_argument("--system", default=(
"You are a research assistant. Use web grounding when helpful. "
"Return concise facts with explicit source-backed statements."
))
return parser
def main() -> int:
args = build_parser().parse_args()
load_secrets()
log_file = Path(args.log_file) if args.log_file else None
with ZenMuxClient(log_file=log_file) as client:
result = client.chat_complete_with_meta(
model=args.model,
system=args.system,
user=args.query,
temperature=args.temperature,
max_tokens=args.max_tokens,
web_search=True,
web_search_options={},
tag="ground",
)
if args.json:
payload = {
"query": args.query,
"model": args.model,
"content": result["content"],
"citations": result["citations"],
"usage": result["usage"],
}
print(json.dumps(payload, ensure_ascii=False, indent=2))
else:
print(result["content"])
if result["citations"]:
print("\nCitations:")
for idx, url in enumerate(result["citations"], start=1):
print(f"{idx}. {url}")
return 0
if __name__ == "__main__":
raise SystemExit(main())