diff --git a/src/openjarvis/agents/hybrid/_base.py b/src/openjarvis/agents/hybrid/_base.py index 5a334551..e9c32840 100644 --- a/src/openjarvis/agents/hybrid/_base.py +++ b/src/openjarvis/agents/hybrid/_base.py @@ -73,12 +73,19 @@ WEB_SEARCH_COST_PER_CALL = 0.01 # $0.01/call number — kept as a separate constant so it can drift. OPENAI_WEB_SEARCH_COST_PER_CALL = 0.01 -# Gemini Google-Search grounding: billed at $35 per 1000 grounded -# *requests* (2025-12 public list price for the Grounding-with-Google-Search -# tool, charged once per request that uses the tool regardless of how many -# internal queries it issues). We charge per grounded request, not per -# `web_search_queries` entry. -GEMINI_SEARCH_COST_PER_CALL = 0.035 +# Gemini 3 Google-Search grounding: billed at $14 per 1000 search queries. +# `_call_gemini_agent` reports the model's `web_search_queries`, so this is +# charged per query, not per outer generate_content request. +GEMINI_SEARCH_COST_PER_CALL = 0.014 + +# Tavily Search, advanced depth: 2 API credits per search request at $0.008 +# per credit on the public pay-as-you-go plan. WebSearchTool captures actual +# credits when Tavily returns usage metadata; this is the fallback estimate. +TAVILY_SEARCH_COST_PER_CREDIT = 0.008 +TAVILY_ADVANCED_SEARCH_CREDITS = 2 +TAVILY_SEARCH_COST_PER_CALL = ( + TAVILY_SEARCH_COST_PER_CREDIT * TAVILY_ADVANCED_SEARCH_CREDITS +) ANTHROPIC_WEB_SEARCH_TOOL = { "type": "web_search_20250305", @@ -101,6 +108,40 @@ def build_web_search_tool(max_uses: int = 8) -> Dict[str, Any]: } +def tavily_search_context( + query: str, + *, + max_results: int = 5, +) -> Dict[str, Any]: + """Run OpenJarvis WebSearchTool and return accounting-friendly metadata.""" + from openjarvis.tools.web_search import WebSearchTool + + tool = WebSearchTool(max_results=max_results) + res = tool.execute(query=query, max_results=max_results) + meta = dict(res.metadata or {}) + engine = str(meta.get("engine") or "unknown") + credits = 0 + cost_usd = 0.0 + if engine == "tavily": + try: + credits = int(meta.get("credits") or TAVILY_ADVANCED_SEARCH_CREDITS) + except (TypeError, ValueError): + credits = TAVILY_ADVANCED_SEARCH_CREDITS + cost_usd = credits * TAVILY_SEARCH_COST_PER_CREDIT + text = res.content or "" + if not res.success and not text: + text = "(no search results)" + return { + "text": text, + "success": bool(res.success), + "engine": engine, + "credits": credits, + "cost_usd": cost_usd, + "n_searches": 1 if (query or "").strip() else 0, + "error": None if res.success else text, + } + + def web_search_cfg(method_cfg: Optional[Dict[str, Any]]) -> Tuple[bool, int]: """Parse ``method_cfg.web_search = { enabled, max_uses }``. @@ -1322,6 +1363,9 @@ __all__ = [ "LocalCloudAgent", "NO_TEMP_PREFIXES", "OPENAI_WEB_SEARCH_COST_PER_CALL", + "TAVILY_ADVANCED_SEARCH_CREDITS", + "TAVILY_SEARCH_COST_PER_CALL", + "TAVILY_SEARCH_COST_PER_CREDIT", "WEB_SEARCH_COST_PER_CALL", "_bump_cloud_calls", "_bump_local_calls", @@ -1329,5 +1373,6 @@ __all__ = [ "estimate_cost", "is_gpt5_family", "supports_temperature", + "tavily_search_context", "web_search_cfg", ] diff --git a/src/openjarvis/agents/hybrid/_prices.py b/src/openjarvis/agents/hybrid/_prices.py index c146496b..2b518e50 100644 --- a/src/openjarvis/agents/hybrid/_prices.py +++ b/src/openjarvis/agents/hybrid/_prices.py @@ -15,13 +15,16 @@ PRICES: dict[str, tuple[float, float]] = { "claude-sonnet-4-6": (3.00, 15.0), "claude-haiku-4-5": (1.00, 5.00), "claude-haiku-4-5-20251001": (1.00, 5.00), + "gpt-5.5": (5.00, 30.0), "gpt-5": (1.25, 10.0), "gpt-5-mini": (0.25, 2.00), "gpt-5-mini-2025-08-07": (0.25, 2.00), "gpt-4o": (0.15, 0.60), - # Gemini Developer API prices (USD per 1M tokens), 2025-12 list price. - # 2.5 Pro uses tiered pricing (>200K context = $2.50/$15); we charge the - # low-context tier since GAIA / SWE-bench prompts stay well under 200K. + # Gemini Developer API prices (USD per 1M tokens). Pro models use tiered + # pricing above 200K prompt tokens; GAIA prompts stay under that tier, so + # charge the low-context standard rate. + "gemini-3.1-pro-preview": (2.00, 12.0), + "gemini-3.1-pro-preview-customtools": (2.00, 12.0), "gemini-2.5-pro": (1.25, 10.0), "gemini-2.5-flash": (0.30, 2.50), "gemini-2.5-flash-lite": (0.10, 0.40), @@ -61,7 +64,11 @@ def is_reasoning_model(model: str) -> bool: before emitting visible answer text. At max_tokens=4096 these silently truncate with empty answers on GAIA (26/100 GPT-5, 18/100 Gemini Pro).""" m = (model or "").lower() - return is_gpt5_family(model) or "gemini-2.5-pro" in m + return ( + is_gpt5_family(model) + or "gemini-2.5-pro" in m + or "gemini-3.1-pro" in m + ) def default_max_output_tokens(model: str) -> int: diff --git a/src/openjarvis/agents/hybrid/advisors.py b/src/openjarvis/agents/hybrid/advisors.py index 307f3fc3..6e8c9f3e 100644 --- a/src/openjarvis/agents/hybrid/advisors.py +++ b/src/openjarvis/agents/hybrid/advisors.py @@ -34,6 +34,7 @@ from openjarvis.agents.hybrid._base import ( WEB_SEARCH_COST_PER_CALL, LocalCloudAgent, build_web_search_tool, + tavily_search_context, web_search_cfg, ) from openjarvis.agents.hybrid.mini_swe_agent import ( @@ -135,7 +136,12 @@ class AdvisorsAgent(LocalCloudAgent): advisor_temperature = float(cfg.get("advisor_temperature", 0.2)) ws_enabled, ws_max_uses = web_search_cfg(cfg) - if ws_enabled and self._cloud_endpoint not in _SEARCH_CAPABLE_ENDPOINTS: + search_backend = str(cfg.get("search_backend", "provider")).lower() + if ( + ws_enabled + and search_backend != "tavily" + and self._cloud_endpoint not in _SEARCH_CAPABLE_ENDPOINTS + ): raise ValueError( f"web_search.enabled=true but cloud_endpoint={self._cloud_endpoint!r}; " "server-side web_search is wired for anthropic / openai / gemini " @@ -146,19 +152,23 @@ class AdvisorsAgent(LocalCloudAgent): use_ws = ws_enabled gaia_max_turns = int(cfg.get("gaia_max_turns", 8)) n_searches_total = 0 + search_cost_total = 0.0 # 1. Initial executor pass — advisor (Qwen) doesn't get tools; # only the cloud executor passes do. With web_search on, dispatch # to the search-capable agent loop for the configured provider. if use_ws: - initial_resp, e1_in, e1_out, n_s1, e1_turns = self._executor_search( + (initial_resp, e1_in, e1_out, n_s1, e1_turns, + e1_search_cost) = self._executor_search( user=f"Question:\n{question}", system=EXECUTOR_INITIAL_SYS, max_tokens=executor_max_tokens, ws_max_uses=ws_max_uses, max_turns=gaia_max_turns, + query=question, ) n_searches_total += n_s1 + search_cost_total += e1_search_cost else: initial_resp, e1_in, e1_out = self._call_cloud( user=f"Question:\n{question}", @@ -196,14 +206,17 @@ class AdvisorsAgent(LocalCloudAgent): f"answer-format rules." ) if use_ws: - final_answer, e2_in, e2_out, n_s2, e2_turns = self._executor_search( + (final_answer, e2_in, e2_out, n_s2, e2_turns, + e2_search_cost) = self._executor_search( user=final_user, system=EXECUTOR_FINAL_SYS, max_tokens=executor_max_tokens, ws_max_uses=ws_max_uses, max_turns=gaia_max_turns, + query=question, ) n_searches_total += n_s2 + search_cost_total += e2_search_cost else: final_answer, e2_in, e2_out = self._call_cloud( user=final_user, @@ -216,7 +229,10 @@ class AdvisorsAgent(LocalCloudAgent): tokens_local = adv_in + adv_out tokens_cloud = e1_in + e1_out + e2_in + e2_out cost = self.cost_usd(self._cloud_model, e1_in + e2_in, e1_out + e2_out) - cost += n_searches_total * _search_cost_per_call(self._cloud_endpoint) + if search_backend == "tavily": + cost += search_cost_total + else: + cost += n_searches_total * _search_cost_per_call(self._cloud_endpoint) meta: Dict[str, Any] = { "tokens_local": tokens_local, @@ -233,7 +249,9 @@ class AdvisorsAgent(LocalCloudAgent): "initial_response": initial_resp, "advisor_feedback": advisor_text, "web_search_enabled": use_ws, + "search_backend": search_backend, "n_web_searches": n_searches_total, + "search_cost_usd": search_cost_total, "note": "inference-only advisor (untrained); lower bound on the technique.", }, } @@ -251,16 +269,40 @@ class AdvisorsAgent(LocalCloudAgent): max_tokens: int, ws_max_uses: int, max_turns: int, - ) -> Tuple[str, int, int, int, int]: + query: Optional[str] = None, + ) -> Tuple[str, int, int, int, int, float]: """Run a search-capable executor pass for the configured cloud. Dispatches by ``self._cloud_endpoint`` to the matching ``_base`` - agent loop. Returns the shared 5-tuple ``(text, p_tok, c_tok, - n_searches, turns)``. The endpoint is assumed already validated - against ``_SEARCH_CAPABLE_ENDPOINTS`` by the caller. + agent loop, or through Tavily when ``method_cfg.search_backend`` is + ``"tavily"``. Returns ``(text, p_tok, c_tok, n_searches, turns, + search_cost_usd)``. """ + if str(self._cfg.get("search_backend", "provider")).lower() == "tavily": + res = tavily_search_context( + query or user, + max_results=int(self._cfg.get("tavily_max_results", 5)), + ) + grounded_user = ( + f"Web search results:\n{res['text']}\n\n" + f"Using the search results above, answer this request:\n{user}" + ) + text, p, c = self._call_cloud( + user=grounded_user, + system=system, + max_tokens=max_tokens, + temperature=0.0, + ) + return ( + text, + p, + c, + int(res["n_searches"]), + 1, + float(res["cost_usd"]), + ) if self._cloud_endpoint == "anthropic": - return self._call_anthropic_agent( + text, p, c, n_searches, turns = self._call_anthropic_agent( self._cloud_model, user=user, system=system, @@ -269,8 +311,9 @@ class AdvisorsAgent(LocalCloudAgent): tools=[build_web_search_tool(ws_max_uses)], max_turns=max_turns, ) + return text, p, c, n_searches, turns, 0.0 if self._cloud_endpoint == "openai": - return self._call_openai_agent( + text, p, c, n_searches, turns = self._call_openai_agent( self._cloud_model, user=user, system=system, @@ -278,8 +321,9 @@ class AdvisorsAgent(LocalCloudAgent): temperature=0.0, max_turns=max_turns, ) + return text, p, c, n_searches, turns, 0.0 if self._cloud_endpoint == "gemini": - return self._call_gemini_agent( + text, p, c, n_searches, turns = self._call_gemini_agent( self._cloud_model, user=user, system=system, @@ -287,6 +331,7 @@ class AdvisorsAgent(LocalCloudAgent): temperature=0.0, max_turns=max_turns, ) + return text, p, c, n_searches, turns, 0.0 # Genuinely unsupported (openrouter / vllm / unknown). The caller # guard should have caught this; raise defensively. raise ValueError( diff --git a/src/openjarvis/agents/hybrid/conductor.py b/src/openjarvis/agents/hybrid/conductor.py index 46dbb1bc..e326059f 100644 --- a/src/openjarvis/agents/hybrid/conductor.py +++ b/src/openjarvis/agents/hybrid/conductor.py @@ -46,6 +46,7 @@ from openjarvis.agents.hybrid._base import ( WEB_SEARCH_COST_PER_CALL, LocalCloudAgent, build_web_search_tool, + tavily_search_context, web_search_cfg, ) from openjarvis.agents.hybrid._prices import ( @@ -456,8 +457,14 @@ def _format_worker_pool(workers: List[Dict[str, Any]]) -> str: ) -def _search_capable_indices(workers: List[Dict[str, Any]]) -> List[int]: +def _search_capable_indices( + workers: List[Dict[str, Any]], + *, + search_backend: str = "provider", +) -> List[int]: """Indices of workers whose endpoint can run server-side web search.""" + if search_backend == "tavily": + return [w["id"] for w in workers] return [ w["id"] for w in workers if (w.get("endpoint") or "openai").lower() @@ -470,6 +477,7 @@ def _build_conductor_prompt( workers: List[Dict[str, Any]], *, web_search_enabled: bool = False, + search_backend: str = "provider", ) -> str: """Build the planner prompt. @@ -485,12 +493,16 @@ def _build_conductor_prompt( ) if not web_search_enabled: return base - capable = _search_capable_indices(workers) + capable = _search_capable_indices(workers, search_backend=search_backend) if capable: cap_str = ", ".join(str(i) for i in capable) + if search_backend == "tavily": + capability = "External Tavily search results will be prepended to worker prompts" + else: + capability = "Only these model indices can perform live web search" constraint = ( "\n\nWEB SEARCH CONSTRAINT:\n" - f"Only these model indices can perform live web search: [{cap_str}]. " + f"{capability}: [{cap_str}]. " "Any step that needs to look up facts, current events, or other " "information not reliably known from memory MUST be routed to one " "of those indices. Steps routed to any other model can only use " @@ -550,8 +562,8 @@ def _call_worker( *, web_search_tool: Optional[Dict[str, Any]] = None, web_search_max_uses: int = 8, -) -> Tuple[str, int, int, bool, int]: - """Returns (text, p_tok, c_tok, is_local, n_web_searches). +) -> Tuple[str, int, int, bool, int, float]: + """Returns (text, p_tok, c_tok, is_local, n_web_searches, extra_cost). ``web_search_tool``: a truthy marker that web_search is enabled for this run. When set AND the worker endpoint is search-capable @@ -565,6 +577,22 @@ def _call_worker( max_tok = int(cfg.get("worker_max_tokens", 4096)) temp = float(cfg.get("worker_temperature", 0.2)) use_ws = web_search_tool is not None + search_backend = str(cfg.get("search_backend", "provider")).lower() + extra_cost = 0.0 + if use_ws and search_backend == "tavily": + res = tavily_search_context( + prompt, + max_results=int(cfg.get("tavily_max_results", 5)), + ) + prompt = ( + f"Web search results:\n{res['text']}\n\n" + f"Using the search results above, answer this request:\n{prompt}" + ) + extra_cost = float(res["cost_usd"]) + use_ws = False + tavily_searches = int(res["n_searches"]) + else: + tavily_searches = 0 if ep == "vllm": text, p, c = LocalCloudAgent._call_vllm( @@ -575,7 +603,7 @@ def _call_worker( temperature=temp, enable_thinking=False, ) - return text, p, c, True, 0 + return text, p, c, True, tavily_searches, extra_cost if ep == "openai": if use_ws: text, p, c, n_searches, _ = LocalCloudAgent._call_openai_agent( @@ -584,14 +612,14 @@ def _call_worker( max_tokens=max_tok, temperature=(1.0 if is_gpt5_family(worker["model"]) else temp), ) - return text, p, c, False, n_searches + return text, p, c, False, n_searches, 0.0 text, p, c = LocalCloudAgent._call_openai( worker["model"], user=prompt, max_tokens=max_tok, temperature=(1.0 if is_gpt5_family(worker["model"]) else temp), ) - return text, p, c, False, 0 + return text, p, c, False, tavily_searches, extra_cost if ep == "openrouter": # OpenRouter is OpenAI-compatible; the helper handles the # base_url + OPENROUTER_API_KEY plumbing. No server-side web @@ -607,7 +635,7 @@ def _call_worker( temperature=temp, extra_body=extra_body if isinstance(extra_body, dict) else None, ) - return text, p, c, False, 0 + return text, p, c, False, tavily_searches, extra_cost if ep == "anthropic": eff_temp = temp if supports_temperature(worker["model"]) else 0.0 anthropic_kwargs: Dict[str, Any] = dict( @@ -620,7 +648,7 @@ def _call_worker( text, p, c, n_searches = LocalCloudAgent._call_anthropic( worker["model"], **anthropic_kwargs ) - return text, p, c, False, n_searches + return text, p, c, False, n_searches or tavily_searches, extra_cost if ep == "gemini": # Gemini Developer API via google-genai. With web_search on, route # through the Google-Search-grounded agent loop; otherwise plain @@ -632,14 +660,14 @@ def _call_worker( max_tokens=max_tok, temperature=temp, ) - return text, p, c, False, n_searches + return text, p, c, False, n_searches, 0.0 text, p, c = LocalCloudAgent._call_gemini( worker["model"], user=prompt, max_tokens=max_tok, temperature=temp, ) - return text, p, c, False, 0 + return text, p, c, False, tavily_searches, extra_cost raise ValueError(f"unsupported worker endpoint: {ep!r}") @@ -672,7 +700,7 @@ def _swe_worker_step( # backbones today (the loop's tool-call format is Anthropic- or # OpenAI-via-vllm-shaped only). Fall back to one-shot for those — # SWE-bench-wise they were already weak; this preserves behavior. - text, p, c, is_local, n_searches = _call_worker(worker, prompt, cfg) + text, p, c, is_local, n_searches, _extra = _call_worker(worker, prompt, cfg) return text, p, c, is_local, n_searches, 0 out = run_swe_agent_loop( task, @@ -753,13 +781,17 @@ class ConductorAgent(LocalCloudAgent): and bool(task_meta_early.get("base_commit")) ) ws_enabled, ws_max_uses = web_search_cfg(cfg) + search_backend = str(cfg.get("search_backend", "provider")).lower() planner_ws = ws_enabled and not swe_mode_early # 1. Plan — when web_search is on (GAIA), the prompt names which # worker indices can actually search, so the planner routes # research steps to a search-capable worker. user = _build_conductor_prompt( - question, workers, web_search_enabled=planner_ws, + question, + workers, + web_search_enabled=planner_ws, + search_backend=search_backend, ) plan_text, p_in, p_out = self._call_cloud( user=user, @@ -833,7 +865,7 @@ class ConductorAgent(LocalCloudAgent): # memory. Fail loud instead of degrading silently. # ``ws_enabled`` / ``ws_max_uses`` computed up front for the planner # constraint — reuse them here. - if ws_enabled and not swe_mode: + if ws_enabled and search_backend != "tavily" and not swe_mode: search_workers = [ w for w in workers if (w.get("endpoint") or "openai").lower() @@ -894,7 +926,7 @@ class ConductorAgent(LocalCloudAgent): # may legitimately not need search; see Task-3 planner # constraint that tries to prevent this upfront). if ( - ws_enabled and not swe_mode + ws_enabled and search_backend != "tavily" and not swe_mode and worker_ep not in _SEARCH_CAPABLE_WORKER_ENDPOINTS ): self.record_trace_event({ @@ -911,6 +943,7 @@ class ConductorAgent(LocalCloudAgent): ), }) + extra_cost = 0.0 if swe_mode: text, w_in, w_out, is_local, n_searches, bash_turns = ( _swe_worker_step( @@ -919,7 +952,9 @@ class ConductorAgent(LocalCloudAgent): ) tool_calls += bash_turns else: - text, w_in, w_out, is_local, n_searches = _call_worker( + ( + text, w_in, w_out, is_local, n_searches, extra_cost + ) = _call_worker( worker, prompt, cfg, web_search_tool=ws_tool, web_search_max_uses=ws_max_uses, @@ -930,7 +965,10 @@ class ConductorAgent(LocalCloudAgent): else: tokens_cloud += w_in + w_out cost += self.cost_usd(worker["model"], w_in, w_out) - cost += n_searches * _worker_search_cost_per_call(worker_ep) + if search_backend != "tavily": + cost += n_searches * _worker_search_cost_per_call(worker_ep) + if search_backend == "tavily": + cost += extra_cost n_web_searches_total += n_searches tool_calls += n_searches steps.append({ @@ -981,6 +1019,7 @@ class ConductorAgent(LocalCloudAgent): "plan": plan, "fallback_used": fallback_used, "web_search_enabled": ws_enabled, + "search_backend": search_backend, "n_web_searches": n_web_searches_total, "parse_attempts": parse_attempts, "workers": [ diff --git a/src/openjarvis/agents/hybrid/minions.py b/src/openjarvis/agents/hybrid/minions.py index da5ab9b6..a5902017 100644 --- a/src/openjarvis/agents/hybrid/minions.py +++ b/src/openjarvis/agents/hybrid/minions.py @@ -48,12 +48,13 @@ from openjarvis.agents.hybrid._base import ( WEB_SEARCH_COST_PER_CALL, LocalCloudAgent, build_web_search_tool, + tavily_search_context, web_search_cfg, ) from openjarvis.agents.hybrid._openai_retry import ( patch_openai_globally as _patch_openai_globally, ) -from openjarvis.agents.hybrid._prices import NO_TEMP_PREFIXES +from openjarvis.agents.hybrid._prices import NO_TEMP_PREFIXES, default_max_output_tokens from openjarvis.agents.hybrid.mini_swe_agent import run_swe_agent_loop from openjarvis.core.registry import AgentRegistry @@ -362,6 +363,8 @@ def _prefetch_context( cloud_endpoint: str, cloud_model: str, max_uses: int = 8, + search_backend: str = "provider", + tavily_max_results: int = 5, ) -> Dict[str, Any]: """Use Anthropic web_search to fetch real source material the worker can read. @@ -375,6 +378,22 @@ def _prefetch_context( out: Dict[str, Any] = { "text": "", "tokens": 0, "cost_usd": 0.0, "n_searches": 0, } + if search_backend == "tavily": + try: + res = tavily_search_context(question, max_results=tavily_max_results) + out.update( + text=res["text"], + cost_usd=float(res["cost_usd"]), + n_searches=int(res["n_searches"]), + tokens=0, + engine=res.get("engine"), + credits=res.get("credits"), + ) + if res.get("error"): + out["error"] = res["error"] + except Exception as e: + out["error"] = f"{type(e).__name__}: {e}" + return out if cloud_endpoint != "anthropic" or not (question or "").strip(): return out try: @@ -498,18 +517,22 @@ class MinionsAgent(LocalCloudAgent): max_tokens=cfg.get("worker_max_tokens", 4096), local=True, ) + cloud_max_tokens = int( + cfg.get("cloud_max_tokens") + or default_max_output_tokens(self._cloud_model) + ) if self._cloud_endpoint == "openai": cloud_client = OpenAIClient( model_name=self._cloud_model, temperature=0.0, - max_tokens=4096, + max_tokens=cloud_max_tokens, ) elif self._cloud_endpoint == "anthropic": # Temperature stripping is handled by the global patch above for Opus 4.7+. cloud_client = AnthropicClient( model_name=self._cloud_model, temperature=0.0, - max_tokens=4096, + max_tokens=cloud_max_tokens, ) elif self._cloud_endpoint == "gemini": # The vendored Minion library already special-cases GeminiClient @@ -520,7 +543,7 @@ class MinionsAgent(LocalCloudAgent): cloud_client = GeminiClient( model_name=self._cloud_model, temperature=0.0, - max_tokens=4096, + max_tokens=cloud_max_tokens, ) else: raise ValueError(f"unsupported cloud endpoint: {self._cloud_endpoint!r}") @@ -560,6 +583,8 @@ class MinionsAgent(LocalCloudAgent): self._cloud_endpoint, self._cloud_model, max_uses=ws_max_uses, + search_backend=str(cfg.get("search_backend", "provider")).lower(), + tavily_max_results=int(cfg.get("tavily_max_results", 5)), ) if prefetch.get("text"): diff --git a/src/openjarvis/agents/hybrid/skillorchestra.py b/src/openjarvis/agents/hybrid/skillorchestra.py index 8b179ea7..510275bd 100644 --- a/src/openjarvis/agents/hybrid/skillorchestra.py +++ b/src/openjarvis/agents/hybrid/skillorchestra.py @@ -149,6 +149,17 @@ def _build_router_schema(agent_ids: List[str]) -> Dict[str, Any]: } +def _openai_response_format(schema: Dict[str, Any]) -> Dict[str, Any]: + return { + "type": "json_schema", + "json_schema": { + "name": "skillorchestra_route", + "schema": schema["format"]["schema"], + "strict": True, + }, + } + + def _parse_router_json(text: str) -> Dict[str, Any]: s = (text or "").strip() try: @@ -196,6 +207,70 @@ class SkillOrchestraAgent(LocalCloudAgent): agent_id = "skillorchestra" + def _route_call( + self, + *, + question: str, + router_sys: str, + router_schema: Dict[str, Any], + router_max: int, + ) -> Tuple[str, int, int]: + user = f"Question:\n{question}" + if self._cloud_endpoint == "anthropic": + kwargs: Dict[str, Any] = { + "user": user, + "system": router_sys, + "max_tokens": router_max, + "output_config": router_schema, + } + if supports_temperature(self._cloud_model): + kwargs["temperature"] = 0.0 + text, r_in, r_out, _ = self._call_anthropic( + self._cloud_model, + **kwargs, + ) + return text, r_in, r_out + if self._cloud_endpoint == "openai": + return self._call_openai( + self._cloud_model, + user=user, + system=router_sys, + max_tokens=router_max, + temperature=0.0, + response_format=_openai_response_format(router_schema), + ) + if self._cloud_endpoint == "gemini": + return self._call_gemini( + self._cloud_model, + user=user, + system=router_sys, + max_tokens=router_max, + temperature=0.0, + ) + raise ValueError( + f"SkillOrchestra router unsupported cloud_endpoint={self._cloud_endpoint!r}" + ) + + def _executor_call( + self, + *, + question: str, + max_tokens: int, + ) -> Tuple[str, int, int]: + if self._cloud_endpoint == "anthropic": + text, w_in, w_out, _ = self._call_anthropic( + self._cloud_model, + user=question, + max_tokens=max_tokens, + temperature=0.0, + ) + return text, w_in, w_out + return self._call_cloud( + user=question, + max_tokens=max_tokens, + temperature=0.0, + ) + def _is_soft_failure(self, exc: BaseException) -> Optional[str]: # Empty/unbalanced router JSON — treat as soft failure to match the # hybrid adapter's behavior (matches `err=1` rows in the n=30 cell). @@ -220,33 +295,13 @@ class SkillOrchestraAgent(LocalCloudAgent): router_sys = _build_router_sys(competence, cost) router_schema = _build_router_schema(agent_ids) - # 1. Route — Anthropic only (output_config schema is Anthropic-specific - # in the hybrid adapter). If you need OpenAI routing, swap the prompt - # to JSON-mode and bypass output_config. - if self._cloud_endpoint != "anthropic": - raise ValueError( - "SkillOrchestra router requires cloud_endpoint='anthropic'; " - f"got {self._cloud_endpoint!r}" - ) router_max = int(cfg.get("router_max_tokens", 1024)) - # Strip temperature for Opus 4.7+; Anthropic's output_config does the schema. - if supports_temperature(self._cloud_model): - router_text, r_in, r_out, _ = self._call_anthropic( - self._cloud_model, - user=f"Question:\n{question}", - system=router_sys, - max_tokens=router_max, - temperature=0.0, - output_config=router_schema, - ) - else: - router_text, r_in, r_out, _ = self._call_anthropic( - self._cloud_model, - user=f"Question:\n{question}", - system=router_sys, - max_tokens=router_max, - output_config=router_schema, - ) + router_text, r_in, r_out = self._route_call( + question=question, + router_sys=router_sys, + router_schema=router_schema, + router_max=router_max, + ) decision = _parse_router_json(router_text) skill_weights: Dict[str, float] = decision.get("skill_weights") or {} @@ -329,11 +384,9 @@ class SkillOrchestraAgent(LocalCloudAgent): tokens_cloud += out["tokens_in"] + out["tokens_out"] run_cost += out["cost_usd"] else: - ans, w_in, w_out, _ = self._call_anthropic( - self._cloud_model, - user=question, + ans, w_in, w_out = self._executor_call( + question=question, max_tokens=int(cfg.get("cloud_max_tokens", 4096)), - temperature=0.0, ) tokens_cloud += w_in + w_out run_cost += self.cost_usd(self._cloud_model, w_in, w_out) diff --git a/src/openjarvis/agents/hybrid/skillorchestra/orchestrator.py b/src/openjarvis/agents/hybrid/skillorchestra/orchestrator.py index a221dbc9..8d252066 100644 --- a/src/openjarvis/agents/hybrid/skillorchestra/orchestrator.py +++ b/src/openjarvis/agents/hybrid/skillorchestra/orchestrator.py @@ -31,7 +31,14 @@ from .stage_router import ( get_routing_strategy, parse_skill_analysis, ) -from .tools import anthropic_tools, openai_tools, run_answer, run_code, run_search +from .tools import ( + anthropic_tools, + gemini_tools, + openai_tools, + run_answer, + run_code, + run_search, +) # tool name -> routing stage (stage_router uses "reasoning" for code). _TOOL_STAGE = { @@ -115,10 +122,49 @@ def _orchestrate_step( u = resp.usage p = getattr(u, "prompt_tokens", 0) if u else 0 c = getattr(u, "completion_tokens", 0) if u else 0 + elif endpoint == "gemini": + from google import genai + from google.genai import types + + client = genai.Client( + http_options=types.HttpOptions(timeout=600_000) + ) + cfg = types.GenerateContentConfig( + temperature=1.0, + max_output_tokens=max_tokens, + tools=[types.Tool(function_declarations=gemini_tools())], + ) + resp = client.models.generate_content( + model=model, + contents=user, + config=cfg, + ) + text = (resp.text or "") if hasattr(resp, "text") else "" + tool_calls = [] + try: + parts = resp.candidates[0].content.parts or [] + except Exception: # noqa: BLE001 + parts = [] + for part in parts: + fc = getattr(part, "function_call", None) + if fc is None: + continue + name = getattr(fc, "name", None) + if not isinstance(name, str) or not name: + continue + args = getattr(fc, "args", None) or {} + try: + args = dict(args) + except Exception: # noqa: BLE001 + args = {} + tool_calls.append({"name": name, "input": args}) + um = getattr(resp, "usage_metadata", None) + p = int(getattr(um, "prompt_token_count", 0) or 0) if um else 0 + c = int(getattr(um, "candidates_token_count", 0) or 0) if um else 0 else: raise ValueError( f"orchestrator endpoint {endpoint!r} unsupported — route the " - "orchestrator through anthropic/openai (set method_cfg." + "orchestrator through anthropic/openai/gemini (set method_cfg." "orchestrator_endpoint)." ) @@ -192,6 +238,8 @@ def run_orchestrator( code_timeout = int(cfg.get("code_timeout_s", 60)) answer_max_tokens = int(cfg.get("answer_max_tokens", 40000)) ws_max_uses = int(cfg.get("web_search_max_uses", 5)) + search_backend = str(cfg.get("search_backend", "provider")).lower() + tavily_max_results = int(cfg.get("tavily_max_results", 5)) # The orchestrator model: a fixed model per run (the original's # MODEL_NAME). Defaults to the cell's cloud model when that endpoint @@ -203,7 +251,7 @@ def run_orchestrator( orch_model = (cfg.get("orchestrator_model") or cfg.get("router_model") or agent._cloud_model) - if orch_endpoint not in ("anthropic", "openai"): + if orch_endpoint not in ("anthropic", "openai", "gemini"): orch_endpoint, orch_model = "anthropic", "claude-opus-4-7" orch_max_tokens = int(cfg.get("orchestrator_max_tokens", 4096)) @@ -306,6 +354,8 @@ def run_orchestrator( res = run_search( agent, spec, context_str=context_str, problem=problem, retriever_url=retriever_url, web_search_max_uses=ws_max_uses, + search_backend=search_backend, + tavily_max_results=tavily_max_results, ) docs = res["search_results_data"] joined = "\n---\n".join(d for d in docs if d)[:char_cap] diff --git a/src/openjarvis/agents/hybrid/skillorchestra/tools.py b/src/openjarvis/agents/hybrid/skillorchestra/tools.py index 59a0ad2f..8401b443 100644 --- a/src/openjarvis/agents/hybrid/skillorchestra/tools.py +++ b/src/openjarvis/agents/hybrid/skillorchestra/tools.py @@ -27,6 +27,7 @@ from .._base import ( OPENAI_WEB_SEARCH_COST_PER_CALL, WEB_SEARCH_COST_PER_CALL, build_web_search_tool, + tavily_search_context, ) from .pool import ModelSpec, call_alias @@ -110,6 +111,26 @@ def openai_tools() -> List[Dict[str, Any]]: return out +def gemini_tools() -> List[Dict[str, Any]]: + """The 3 orchestrator tools in Gemini function-declaration shape.""" + out = [] + for name, desc in ( + ("search", _SEARCH_DESC), + ("enhance_reasoning", _CODE_DESC), + ("answer", _ANSWER_DESC), + ): + out.append({ + "name": name, + "description": desc, + "parameters": { + "type": "object", + "properties": {"model": _model_prop(name)}, + "required": ["model"], + }, + }) + return out + + # --------------------------------------------------------------------------- # enhance_reasoning / code — eval_frames.py:659-812 # --------------------------------------------------------------------------- @@ -256,6 +277,8 @@ def run_search( retriever_url: Optional[str] = None, topk: int = 150, web_search_max_uses: int = 5, + search_backend: str = "provider", + tavily_max_results: int = 5, ) -> Dict[str, Any]: """Write a search query with ``spec``, then retrieve documents. @@ -283,7 +306,12 @@ def run_search( contents: List[str] = [] search_uses = 0 - if retriever_url: + if search_backend == "tavily": + res = tavily_search_context(query, max_results=tavily_max_results) + contents.append(res["text"]) + search_uses = int(res["n_searches"]) + cost += float(res["cost_usd"]) + elif retriever_url: # Faithful path — the original FAISS retriever service. import requests diff --git a/src/openjarvis/agents/hybrid/toolorchestra.py b/src/openjarvis/agents/hybrid/toolorchestra.py index d5ebd78d..50ae93f6 100644 --- a/src/openjarvis/agents/hybrid/toolorchestra.py +++ b/src/openjarvis/agents/hybrid/toolorchestra.py @@ -16,7 +16,8 @@ Two modes, gated by ``method_cfg.orchestrator_mode``: (``answer-1``, ``reasoner-2``, ``search-3``, …) is mapped to a real backend through ``EXPERT_MODEL_MAPPING`` — by default the frontier Anthropic worker for `*-1` slots, gpt-5-mini for `*-2`, local Qwen - for `*-3`. Search routes to the Anthropic server-side web_search. + for `*-3`. Search routes to the configured provider's server-side + web-search helper when available. We do NOT reproduce the upstream Tavily / FAISS-wiki retriever, the code-interpreter sandbox, or the multi-vLLM mix (Llama-3.3-70B, @@ -43,8 +44,8 @@ Prompted-mode pipeline: prompt; fallback to strongest worker on parse failure. Workers come from ``cfg["workers"]`` or a sensible default pool (local -Qwen if vLLM up, plus a web-search tool via Anthropic, Opus 4.7, -gpt-5-mini). +Qwen if vLLM up, plus provider-native web search, the configured frontier +cloud model, and gpt-5-mini). """ from __future__ import annotations @@ -59,8 +60,11 @@ from typing import Any, Dict, List, Optional, Tuple from openjarvis.agents._stubs import AgentContext from openjarvis.agents.hybrid._base import ( ANTHROPIC_WEB_SEARCH_TOOL, + GEMINI_SEARCH_COST_PER_CALL, + OPENAI_WEB_SEARCH_COST_PER_CALL, WEB_SEARCH_COST_PER_CALL, LocalCloudAgent, + tavily_search_context, ) from openjarvis.agents.hybrid._prices import ( PRICES, @@ -197,10 +201,23 @@ def _expert_for(slot: str, local_model: Optional[str], cost tier for mid OpenAI calls) - `*-3` (local tier) -> local vLLM (`local_model`) - `answer-math-*` -> same tiers as the numeric suffix - - `search-*` -> always the Anthropic web_search tool (the - upstream uses Tavily; we have web_search) + - `search-*` -> provider-native web search when the cloud + endpoint supports it; otherwise Anthropic """ if slot.startswith("search"): + ep = (cloud_endpoint or "anthropic").lower() + if ep == "openai": + return { + "name": f"search:{slot}", + "type": "openai-web-search", + "model": cloud_model, + } + if ep == "gemini": + return { + "name": f"search:{slot}", + "type": "gemini-web-search", + "model": cloud_model, + } return { "name": f"search:{slot}", "type": "anthropic-web-search", @@ -340,21 +357,18 @@ def _paper_expert_for( # ---- Tavily + Modal helpers ------------------------------------------------- -def _call_tavily_search(query: str, max_results: int = 5) -> Tuple[str, int, int]: - """One-shot Tavily search. Returns (text, p_tok=0, c_tok=0). +def _call_tavily_search( + query: str, + max_results: int = 5, +) -> Tuple[str, int, int, float, int]: + """One-shot Tavily search. Returns (text, p_tok=0, c_tok=0, cost, uses). Token counts are reported as zero (no LLM was billed); the OpenJarvis accounting layer separately tallies tool-call counts. Falls back to DuckDuckGo if Tavily is unreachable (see ``WebSearchTool``). """ - from openjarvis.tools.web_search import WebSearchTool - - tool = WebSearchTool(max_results=max_results) - res = tool.execute(query=query, max_results=max_results) - text = res.content or "" - if not res.success and not text: - text = "(no results)" - return text, 0, 0 + res = tavily_search_context(query, max_results=max_results) + return res["text"], 0, 0, float(res["cost_usd"]), int(res["n_searches"]) _MODAL_APP_NAME = "openjarvis-toolorchestra-sandbox" @@ -674,13 +688,25 @@ def _default_pool( "concise extraction, formatting, arithmetic on given data." ), }) + if ep == "openai": + search_type = "openai-web-search" + search_model = cloud_model + search_desc = "OpenAI hosted web search on the configured frontier model." + elif ep == "gemini": + search_type = "gemini-web-search" + search_model = cloud_model + search_desc = "Gemini Google Search grounding on the configured frontier model." + else: + search_type = "anthropic-web-search" + search_model = _DEFAULT_WEB_SEARCH_MODEL + search_desc = "Anthropic server-side web_search." pool.append({ "id": len(pool), "name": "web-search", - "type": "anthropic-web-search", - "model": "claude-haiku-4-5", + "type": search_type, + "model": search_model, "description": ( - "Anthropic server-side web_search. Use for facts that need a lookup " + f"{search_desc} Use for facts that need a lookup " "(recent events, rare names/dates, niche sources). Returns a digest." ), }) @@ -717,8 +743,13 @@ def _default_pool( # `modal-python` — One-shot Python exec in a fresh Modal Sandbox (the # paper's "Python sandbox" inside `enhance_reasoning`). _TOOLORCH_VALID_TYPES = ( - "vllm", "openai", "anthropic", "anthropic-web-search", "gemini", - "tavily-search", "openrouter", "modal-python", + "vllm", "openai", "anthropic", "anthropic-web-search", + "openai-web-search", "gemini", "gemini-web-search", "tavily-search", + "openrouter", "modal-python", +) +_TOOLORCH_SEARCH_TYPES = ( + "anthropic-web-search", "openai-web-search", "gemini-web-search", + "tavily-search", ) # Default model used when an `anthropic-web-search` entry omits `model`. @@ -739,10 +770,12 @@ def _resolve_worker_pool( the override is absent. Each user-supplied entry must be a dict with keys ``id``, ``name``, - ``type``, and (for non-search types) ``model``. ``type`` must be one - of ``vllm`` / ``openai`` / ``anthropic`` / ``anthropic-web-search``. - ``anthropic-web-search`` entries may omit ``model`` — it defaults to - ``claude-haiku-4-5``. + ``type``, and (for non-search types) ``model``. Search worker types are + ``anthropic-web-search``, ``openai-web-search``, ``gemini-web-search``, + and ``tavily-search``. ``anthropic-web-search`` entries may omit + ``model`` — it defaults to ``claude-haiku-4-5``. OpenAI and Gemini + search workers default to the configured cloud model. Tavily does not + require a model. Substitution: ``model = "$local"`` (or ``""``) resolves to ``local_model``; ``model = "$cloud"`` / ``""`` to ``cloud_model``. @@ -804,14 +837,24 @@ def _resolve_worker_pool( elif isinstance(model, str) and model in ("$cloud", ""): model = cloud_model entry["model"] = model - if wtype == "anthropic-web-search": + if wtype in _TOOLORCH_SEARCH_TYPES: if model in (None, ""): - model = _DEFAULT_WEB_SEARCH_MODEL + if wtype == "anthropic-web-search": + model = _DEFAULT_WEB_SEARCH_MODEL + elif wtype in ("openai-web-search", "gemini-web-search"): + model = cloud_model + else: + model = wtype entry["model"] = model elif not isinstance(model, str): raise ValueError( f"Invalid worker_pool entry [{wid}]: 'model' must be a string when set" ) + if wtype in ("openai-web-search", "gemini-web-search") and model not in PRICES: + raise ValueError( + f"Invalid worker_pool entry [{wid}]: model {model!r} " + f"is not in PRICES (known: {sorted(PRICES)})" + ) # Search workers don't satisfy the "needs a solver" requirement. else: if not isinstance(model, str) or not model: @@ -843,7 +886,7 @@ def _resolve_worker_pool( if not has_non_search: raise ValueError( "Invalid worker_pool entry [-]: worker_pool must contain at least " - "one non-search worker (vllm / openai / anthropic)" + "one non-search worker (vllm / openai / anthropic / gemini)" ) return resolved @@ -910,13 +953,31 @@ def _call_worker( ) extra = n_searches * WEB_SEARCH_COST_PER_CALL return text, p, c, False, extra, n_searches + if wtype == "openai-web-search": + eff_temp = 1.0 if is_gpt5_family(worker["model"]) else temp + text, p, c, n_searches, _ = LocalCloudAgent._call_openai_agent( + worker["model"], + user=prompt, + max_tokens=max(max_tok, 16384) if is_gpt5_family(worker["model"]) else max_tok, + temperature=eff_temp, + ) + extra = n_searches * OPENAI_WEB_SEARCH_COST_PER_CALL + return text, p, c, False, extra, n_searches + if wtype == "gemini-web-search": + text, p, c, n_searches, _ = LocalCloudAgent._call_gemini_agent( + worker["model"], + user=prompt, + max_tokens=max_tok, + temperature=temp, + ) + extra = n_searches * GEMINI_SEARCH_COST_PER_CALL + return text, p, c, False, extra, n_searches if wtype == "tavily-search": - # Tavily costs are flat per call; charge `WEB_SEARCH_COST_PER_CALL` - # for parity with the Anthropic web-search worker. One call = one - # "n_search" for accounting. max_results = int(cfg.get("tavily_max_results", 5)) - text, p, c = _call_tavily_search(str(prompt), max_results=max_results) - return text, p, c, False, WEB_SEARCH_COST_PER_CALL, 1 + text, p, c, extra, n_searches = _call_tavily_search( + str(prompt), max_results=max_results, + ) + return text, p, c, False, extra, n_searches if wtype == "openrouter": text, p, c = LocalCloudAgent._call_openrouter( worker["model"], @@ -951,7 +1012,7 @@ def _swe_call_worker( caller can surface ``tool_calls`` per row. Fallbacks to one-shot workers return 0 bash turns (no agent loop ran).""" wtype = worker.get("type", "openai") - if wtype == "anthropic-web-search": + if wtype in _TOOLORCH_SEARCH_TYPES: # Search workers stay one-shot. text, p, c, is_local, extra, n_searches = _call_worker(worker, prompt, cfg) return text, p, c, is_local, extra, n_searches, 0 @@ -1197,7 +1258,8 @@ class ToolOrchestraAgent(LocalCloudAgent): # Search workers are excluded — they answer fact-lookup # questions, not synthesis. non_search = [ - w for w in workers if w.get("type") != "anthropic-web-search" + w for w in workers + if w.get("type") not in _TOOLORCH_SEARCH_TYPES ] or workers worker = max( non_search, diff --git a/src/openjarvis/tools/web_search.py b/src/openjarvis/tools/web_search.py index a90d0cec..3d71f5de 100644 --- a/src/openjarvis/tools/web_search.py +++ b/src/openjarvis/tools/web_search.py @@ -165,7 +165,10 @@ class WebSearchTool(BaseTool): client = TavilyClient(api_key=self._api_key) response = client.search( - query, max_results=max_results, search_depth="advanced" + query, + max_results=max_results, + search_depth="advanced", + include_usage=True, ) results = response.get("results", []) formatted_parts = [] @@ -182,7 +185,11 @@ class WebSearchTool(BaseTool): tool_name="web_search", content=formatted or "No results found.", success=True, - metadata={"num_results": len(results), "engine": "tavily"}, + metadata={ + "num_results": len(results), + "engine": "tavily", + "credits": (response.get("usage") or {}).get("credits"), + }, ) except Exception as exc: logger.debug( diff --git a/tests/tools/test_web_search.py b/tests/tools/test_web_search.py index 570c3019..5ada6770 100644 --- a/tests/tools/test_web_search.py +++ b/tests/tools/test_web_search.py @@ -171,7 +171,7 @@ class TestWebSearchTool: tool = WebSearchTool(api_key="test-key", max_results=3) tool.execute(query="test", max_results=7) mock_client.search.assert_called_once_with( - "test", max_results=7, search_depth="advanced" + "test", max_results=7, search_depth="advanced", include_usage=True ) def test_to_openai_function(self):