Enable automatic retry for Google GenAI API calls

+7

tests/conftest.py

··· 185 185 for key, value in k.items(): 186 186 setattr(self, key, value) 187 187 188 + class MockHttpRetryOptions: 189 + def __init__(self, **k): 190 + pass 191 + 188 192 genai_mod.types = types.SimpleNamespace( 189 193 GenerateContentConfig=MockGenerateContentConfig, 190 194 Content=MockContent, 191 195 HttpOptions=MockHttpOptions, 196 + HttpRetryOptions=MockHttpRetryOptions, 192 197 ThinkingConfig=MockThinkingConfig, 193 198 ) 194 199 google_mod.genai = genai_mod ··· 408 413 ThinkingConfig=lambda **k: SimpleNamespace(**k), 409 414 Content=lambda **k: SimpleNamespace(**k), 410 415 Part=lambda **k: SimpleNamespace(**k), 416 + HttpOptions=lambda **k: SimpleNamespace(**k), 417 + HttpRetryOptions=lambda **k: SimpleNamespace(**k), 411 418 ) 412 419 google_mod.genai = genai_mod 413 420 monkeypatch.setitem(sys.modules, "google", google_mod)

+195 -22

think/agents.py

··· 37 37 format_segment_times, 38 38 get_muse_configs, 39 39 get_output_path, 40 - load_output_hook, 41 40 load_prompt, 42 41 segment_parse, 43 42 setup_cli, ··· 294 293 return turns 295 294 296 295 296 + # ============================================================================= 297 + # Hook Framework (unified for agents and generators) 298 + # ============================================================================= 299 + 300 + 301 + class HookContext(TypedDict, total=False): 302 + """Context passed to hook functions. 303 + 304 + Provides unified context for both tool-using agents and generators. 305 + Not all fields are present for all modalities. 306 + """ 307 + 308 + # Identity 309 + name: str # Agent/generator name 310 + agent_id: str # Unique agent ID 311 + provider: str # google/anthropic/openai 312 + model: str # Model used 313 + 314 + # Temporal (generators) 315 + day: str # YYYYMMDD 316 + segment: str # Segment key 317 + span: bool # True if span mode 318 + 319 + # Content 320 + prompt: str # Original prompt (agents) or empty (generators) 321 + transcript: str # Clustered transcript (generators only) 322 + 323 + # Output 324 + output_path: str # Where result will be written 325 + output_format: str # 'md' or 'json' 326 + 327 + # Full config 328 + meta: dict # Full frontmatter/config 329 + 330 + 331 + # MUSE_DIR for hook resolution 332 + _MUSE_DIR = Path(__file__).parent.parent / "muse" 333 + 334 + 335 + def load_post_hook(config: dict) -> Callable[[str, HookContext], str | None] | None: 336 + """Load post-processing hook from config if defined. 337 + 338 + Hook config format: {"hook": {"post": "name"}} 339 + Resolution: 340 + - Named: "name" -> muse/{name}.py 341 + - App-qualified: "app:name" -> apps/{app}/muse/{name}.py 342 + - Explicit path: "path/to/hook.py" -> direct path 343 + 344 + Args: 345 + config: Agent/generator config dict 346 + 347 + Returns: 348 + The post_process function from the hook module, or None if no hook. 349 + 350 + Raises: 351 + ValueError: If hook file doesn't define post_process function. 352 + ImportError: If hook file cannot be loaded. 353 + """ 354 + import importlib.util 355 + 356 + hook_config = config.get("hook") 357 + if not hook_config or not isinstance(hook_config, dict): 358 + return None 359 + 360 + post_hook_name = hook_config.get("post") 361 + if not post_hook_name: 362 + return None 363 + 364 + # Resolve hook path 365 + if "/" in post_hook_name or post_hook_name.endswith(".py"): 366 + # Explicit path 367 + hook_path = Path(post_hook_name) 368 + elif ":" in post_hook_name: 369 + # App-qualified: "app:name" -> apps/{app}/muse/{name}.py 370 + app, name = post_hook_name.split(":", 1) 371 + hook_path = Path(__file__).parent.parent / "apps" / app / "muse" / f"{name}.py" 372 + else: 373 + # Named hook: muse/{name}.py 374 + hook_path = _MUSE_DIR / f"{post_hook_name}.py" 375 + 376 + if not hook_path.exists(): 377 + raise ImportError(f"Hook file not found: {hook_path}") 378 + 379 + spec = importlib.util.spec_from_file_location( 380 + f"post_hook_{hook_path.stem}", hook_path 381 + ) 382 + if spec is None or spec.loader is None: 383 + raise ImportError(f"Cannot load hook from {hook_path}") 384 + 385 + module = importlib.util.module_from_spec(spec) 386 + spec.loader.exec_module(module) 387 + 388 + if not hasattr(module, "post_process"): 389 + raise ValueError(f"Hook {hook_path} must define a 'post_process' function") 390 + 391 + process_func = getattr(module, "post_process") 392 + if not callable(process_func): 393 + raise ValueError(f"Hook {hook_path} 'post_process' must be callable") 394 + 395 + return process_func 396 + 397 + 398 + def build_hook_context(config: dict, **extras: Any) -> HookContext: 399 + """Build unified HookContext from config and extra values. 400 + 401 + Args: 402 + config: Agent/generator config dict 403 + **extras: Additional context values (transcript, output_path, etc.) 404 + 405 + Returns: 406 + HookContext with all available fields populated. 407 + """ 408 + context: HookContext = { 409 + "name": config.get("name", ""), 410 + "agent_id": config.get("agent_id", ""), 411 + "provider": config.get("provider", ""), 412 + "model": config.get("model", ""), 413 + "prompt": config.get("prompt", ""), 414 + "output_format": config.get("output", "md"), 415 + "meta": config, 416 + } 417 + 418 + # Add generator-specific fields if present 419 + if "day" in config: 420 + context["day"] = config["day"] 421 + if "segment" in config: 422 + context["segment"] = config["segment"] 423 + 424 + # Merge extras (transcript, output_path, span, etc.) 425 + context.update(extras) 426 + 427 + return context 428 + 429 + 430 + def run_post_hook( 431 + result: str, 432 + context: HookContext, 433 + hook_fn: Callable[[str, HookContext], str | None], 434 + ) -> str: 435 + """Execute post-processing hook and return (potentially transformed) result. 436 + 437 + Args: 438 + result: The LLM-generated output text 439 + context: Hook context with metadata 440 + hook_fn: The post_process function to call 441 + 442 + Returns: 443 + Transformed result if hook returns string, original result otherwise. 444 + """ 445 + try: 446 + hook_result = hook_fn(result, context) 447 + if hook_result is not None: 448 + logging.info("Hook transformed result") 449 + return hook_result 450 + except Exception as exc: 451 + logging.error("Hook failed: %s", exc) 452 + 453 + return result 454 + 455 + 297 456 __all__ = [ 298 457 "ToolStartEvent", 299 458 "ToolEndEvent", ··· 304 463 "ThinkingEvent", 305 464 "GenerateResult", 306 465 "Event", 466 + "HookContext", 307 467 "JSONEventWriter", 308 468 "JSONEventCallback", 309 469 "format_tool_summary", 310 470 "parse_agent_events_to_turns", 471 + "load_post_hook", 472 + "build_hook_context", 473 + "run_post_hook", 311 474 "scan_day", 312 475 "generate_agent_output", 313 476 ] ··· 438 601 client = None 439 602 cache_name = None 440 603 if cache_display_name and provider == "google": 441 - client = genai.Client(api_key=api_key) 604 + client = genai.Client( 605 + api_key=api_key, 606 + http_options=types.HttpOptions( 607 + retry_options=types.HttpRetryOptions() 608 + ), 609 + ) 442 610 cache_name = _get_or_create_cache( 443 611 client, model, cache_display_name, transcript, system_instruction 444 612 ) ··· 688 856 usage_data = gen_result.get("usage") 689 857 690 858 # Run post-processing hook if present 691 - if meta.get("hook_path"): 692 - hook_path = meta["hook_path"] 693 - try: 694 - hook_process = load_output_hook(hook_path) 695 - hook_context = { 696 - "day": day, 697 - "segment": segment, 698 - "span": span_mode, 699 - "name": name, 700 - "output_path": str(output_path), 701 - "meta": dict(meta), 702 - "transcript": markdown, 703 - } 704 - hook_result = hook_process(result, hook_context) 705 - if hook_result is not None: 706 - result = hook_result 707 - logging.info("Hook %s transformed result", hook_path) 708 - except Exception as exc: 709 - logging.error("Hook %s failed: %s", hook_path, exc) 859 + post_hook = load_post_hook(meta) 860 + if post_hook: 861 + hook_context = build_hook_context( 862 + meta, 863 + day=day, 864 + segment=segment, 865 + span=span_mode, 866 + output_path=str(output_path), 867 + transcript=markdown, 868 + ) 869 + result = run_post_hook(result, hook_context, post_hook) 710 870 711 871 # Emit finish event with result (cortex handles file writing) 712 872 finish_event = { ··· 813 973 f"Unknown provider: {provider!r}. Valid providers: {valid}" 814 974 ) 815 975 976 + # Load post hook if configured 977 + post_hook = load_post_hook(config) 978 + 979 + # Create event handler that intercepts finish for hooks 980 + def agent_emit_event(data: Event) -> None: 981 + if post_hook and data.get("event") == "finish": 982 + result = data.get("result", "") 983 + hook_context = build_hook_context(config) 984 + transformed = run_post_hook(result, hook_context, post_hook) 985 + if transformed != result: 986 + data = {**data, "result": transformed} 987 + emit_event(data) 988 + 816 989 # Pass complete config to provider 817 990 await provider_mod.run_agent( 818 991 config=config, 819 - on_event=emit_event, 992 + on_event=agent_emit_event, 820 993 ) 821 994 822 995 else:

+13 -3

think/providers/google.py

··· 75 75 api_key = os.getenv("GOOGLE_API_KEY") 76 76 if not api_key: 77 77 raise ValueError("GOOGLE_API_KEY not found in environment") 78 - client = genai.Client(api_key=api_key) 78 + client = genai.Client( 79 + api_key=api_key, 80 + http_options=types.HttpOptions( 81 + retry_options=types.HttpRetryOptions() 82 + ), 83 + ) 79 84 return client 80 85 81 86 ··· 630 635 ) 631 636 ) 632 637 633 - # Create client 634 - client = genai.Client(api_key=api_key) 638 + # Create client with retry enabled 639 + client = genai.Client( 640 + api_key=api_key, 641 + http_options=types.HttpOptions( 642 + retry_options=types.HttpRetryOptions() 643 + ), 644 + ) 635 645 636 646 # Create fresh chat session 637 647 chat = client.aio.chats.create(

Configure Feed

Configure Feed