Spaces:
Running
Running
| from __future__ import annotations | |
| import os | |
| import time | |
| import httpx | |
| import web_tools | |
| from content.roles import looks_like_refusal | |
| from oplog import log_event | |
| from scheduler import ProviderError, call_slot | |
| # Provider-agnostic research agent: a ReAct think -> act -> observe loop around a | |
| # single model. The model calls web_tools by emitting one ACTION line; we execute | |
| # it, feed back an OBSERVATION, and loop until it writes a final answer (no ACTION) | |
| # or hits the step/budget cap. Reasoning params ride in extra_body; the streamed | |
| # thinking trace is surfaced via usage["reasoning_content"]. Never raises - mirrors | |
| # jobs.call_once's result shape (+ steps/searches/tool_calls). | |
| RESEARCH_MAX_STEPS = int(os.environ.get("RESEARCH_MAX_STEPS", "8")) | |
| async def research_call(client: httpx.AsyncClient, picked, system_prompt: str, *, | |
| user_msg: str, max_tokens: int, timeout_s: float, | |
| extra_body: dict | None = None, max_steps: int | None = None, | |
| system_in_user: bool = False) -> dict: | |
| max_steps = RESEARCH_MAX_STEPS if max_steps is None else max_steps | |
| t0 = time.monotonic() | |
| prompt_full = system_prompt + web_tools.TOOLS_PROTOCOL | |
| if system_in_user: | |
| # hosts that strip custom system prompts (e.g. GitHub Phi-4) get the whole | |
| # protocol folded into the first user message instead. | |
| messages = [{"role": "user", "content": f"{prompt_full}\n\n---\n\n{user_msg}"}] | |
| else: | |
| messages = [ | |
| {"role": "system", "content": prompt_full}, | |
| {"role": "user", "content": user_msg}, | |
| ] | |
| steps = searches = fetches = 0 | |
| reasoning_all: list[str] = [] | |
| in_tok = out_tok = 0 | |
| deadline = t0 + timeout_s | |
| try: | |
| for step in range(max_steps + 1): | |
| final_turn = step == max_steps | |
| msgs = messages | |
| if final_turn: | |
| msgs = messages + [{"role": "user", | |
| "content": "Tool budget reached. Write your FINAL answer now, no ACTION lines."}] | |
| remaining = max(15.0, deadline - time.monotonic()) | |
| content, usage = await call_slot(client, picked, messages=msgs, | |
| max_tokens=max_tokens, timeout_s=remaining, | |
| extra_body=extra_body) | |
| in_tok += int(usage.get("prompt_tokens") or 0) | |
| out_tok += int(usage.get("completion_tokens") or 0) | |
| if usage.get("reasoning_content"): | |
| reasoning_all.append(usage["reasoning_content"]) | |
| action = None if final_turn else web_tools.parse_action(content) | |
| if action is None: | |
| text = (content or "").strip() | |
| if looks_like_refusal(text): | |
| return _result(False, "refusal", error=f"refusal: {text[:160]!r}", | |
| steps=steps, searches=searches, tool_calls=searches + fetches, | |
| t0=t0, in_tok=in_tok, out_tok=out_tok, reasoning=reasoning_all) | |
| return _result(True, "ok", output=text, steps=steps, searches=searches, | |
| tool_calls=searches + fetches, t0=t0, in_tok=in_tok, | |
| out_tok=out_tok, reasoning=reasoning_all) | |
| tool, args = action | |
| steps += 1 | |
| if tool == "web_search": | |
| searches += 1 | |
| obs = await web_tools.web_search(str(args.get("query", "")), http_client=client) | |
| else: | |
| fetches += 1 | |
| obs = await web_tools.web_fetch(str(args.get("url", "")), http_client=client) | |
| log_event("research_step", slot=picked.who, step=steps, tool=tool, | |
| arg=str(args)[:120], obs_chars=len(obs)) | |
| messages = messages + [ | |
| {"role": "assistant", "content": content}, | |
| {"role": "user", "content": f"OBSERVATION:\n{obs}"}, | |
| ] | |
| return _result(False, "empty", error="no final answer", steps=steps, searches=searches, | |
| tool_calls=searches + fetches, t0=t0, in_tok=in_tok, out_tok=out_tok, | |
| reasoning=reasoning_all) | |
| except ProviderError as e: | |
| msg = str(e) | |
| return _result(False, msg.split(":", 1)[0], error=msg[:300], steps=steps, | |
| searches=searches, tool_calls=searches + fetches, t0=t0, | |
| in_tok=in_tok, out_tok=out_tok, reasoning=reasoning_all) | |
| except Exception as e: # noqa: BLE001 | |
| return _result(False, "exc", error=f"{type(e).__name__}: {str(e)[:180]}", steps=steps, | |
| searches=searches, tool_calls=searches + fetches, t0=t0, | |
| in_tok=in_tok, out_tok=out_tok, reasoning=reasoning_all) | |
| def _result(ok, code, *, output=None, error=None, steps, searches, tool_calls, t0, | |
| in_tok, out_tok, reasoning) -> dict: | |
| usage = {"prompt_tokens": in_tok, "completion_tokens": out_tok} | |
| if reasoning: | |
| usage["reasoning_content"] = "\n\n".join(reasoning) | |
| res = {"ok": ok, "code": code, "usage": usage, | |
| "elapsed_s": round(time.monotonic() - t0, 1), | |
| "steps": steps, "searches": searches, "tool_calls": tool_calls} | |
| if output is not None: | |
| res["output"] = output | |
| if error is not None: | |
| res["error"] = error | |
| return res | |