33from __future__ import annotations
44
55import json
6+ import os
67import shlex
8+ import time
79from typing import Any , override
810
911from harbor .agents .base import BaseAgent
@@ -24,6 +26,28 @@ def _is_reasoning_model(model: str) -> bool:
2426MAX_TURNS = 32
2527MAX_TOOL_OUTPUT_CHARS = 40_000
2628
29+ # Wall clock, not turn count, is this agent's binding constraint. Most of a turn
30+ # is the sequential `search.py` invocations, each of which boots a fresh
31+ # pyserini/Lucene JVM against the BM25 index (~13s apiece), against ~3-9s of
32+ # inference. Trials therefore run out of clock long before MAX_TURNS, so the
33+ # turn-budget fallback below never fires and harbor instead kills the trial
34+ # mid-search with no answer recorded. This gives that same best-effort safety
35+ # net a deadline to watch. Harbor does not tell an agent its own timeout, so the
36+ # budget is read from the environment; keep it in step with whatever
37+ # --agent-timeout-multiplier the run uses.
38+ AGENT_BUDGET_SEC = float (os .environ .get ("BROWSECOMP_AGENT_BUDGET_SEC" , "900" ))
39+ # Enough to prefill the accumulated history and write a short answer (~8s for a
40+ # 120k-token prompt measured against the target endpoint), plus the answer
41+ # upload and margin for one retry.
42+ FINAL_RESERVE_SEC = float (os .environ .get ("BROWSECOMP_FINAL_RESERVE_SEC" , "120" ))
43+ # A malformed, empty, or length-truncated assistant turn is recoverable; nudge
44+ # rather than abort, and only give up after this many consecutive nudges.
45+ MAX_EMPTY_RETRIES = 3
46+ # The required answer is three short labeled lines, so the forced-final call is
47+ # capped well below the per-turn budget to stop the model from spending its last
48+ # seconds on chain-of-thought and getting truncated a second time.
49+ ANSWER_MAX_TOKENS = 3_000
50+
2751INSTRUCTIONS = """You are a persistent deep-research agent working only against the
2852fixed BrowseComp-Plus corpus. Find the precise answer by issuing focused searches,
2953opening promising documents, reformulating queries, and cross-checking evidence.
@@ -133,13 +157,17 @@ def _trace(self, event: dict[str, Any]) -> None:
133157 file .write (json .dumps (event , ensure_ascii = False , default = str ) + "\n " )
134158
135159 async def _index_command (
136- self , environment : BaseEnvironment , command : str , value : str
160+ self ,
161+ environment : BaseEnvironment ,
162+ command : str ,
163+ value : str ,
164+ timeout_sec : float = 120 ,
137165 ) -> dict [str , Any ]:
138166 flag = "--query" if command == "search" else "--docid"
139167 result = await environment .exec (
140168 f"python3 /opt/browsecomp/search.py { command } { flag } { shlex .quote (value )} " ,
141169 cwd = "/app" ,
142- timeout_sec = 120 ,
170+ timeout_sec = max ( 15 , int ( timeout_sec )) ,
143171 )
144172 if result .return_code != 0 :
145173 return {
@@ -166,7 +194,11 @@ def _usage_value(value: Any, name: str) -> int:
166194 return int (result or 0 )
167195
168196 def _completion_kwargs (
169- self , messages : list [dict [str , Any ]], * , tools : bool = True
197+ self ,
198+ messages : list [dict [str , Any ]],
199+ * ,
200+ tools : bool = True ,
201+ max_tokens : int = 12_000 ,
170202 ) -> dict [str , Any ]:
171203 # Reasoning models replaced max_tokens with max_completion_tokens and
172204 # reject the old name outright ("Unsupported parameter: 'max_tokens'
@@ -181,14 +213,43 @@ def _completion_kwargs(
181213 "model" : self ._api_model ,
182214 "messages" : messages ,
183215 }
184- kwargs [_token_limit_key ] = 12_000
216+ kwargs [_token_limit_key ] = max_tokens
185217 if tools :
186218 kwargs ["tools" ] = TOOLS
187219 # Capability, not provider: see _is_reasoning_model.
188220 if _is_reasoning_model (self ._api_model ):
189221 kwargs ["reasoning_effort" ] = "medium"
190222 return kwargs
191223
224+ async def _complete (
225+ self ,
226+ messages : list [dict [str , Any ]],
227+ * ,
228+ tools : bool = True ,
229+ timeout : float ,
230+ max_retries : int | None = None ,
231+ max_tokens : int = 12_000 ,
232+ ) -> Any :
233+ """One chat completion, validated for shape.
234+
235+ A gateway can answer 200 with a body the SDK does not turn into a
236+ completion (observed once as `'str' object has no attribute 'usage'`),
237+ so anything without usable `.choices` is rejected here rather than
238+ crashing the trial on an attribute access two lines later.
239+ """
240+ options : dict [str , Any ] = {"timeout" : max (30.0 , timeout )}
241+ if max_retries is not None :
242+ options ["max_retries" ] = max_retries
243+ response = await self ._client .with_options (** options ).chat .completions .create (
244+ ** self ._completion_kwargs (messages , tools = tools , max_tokens = max_tokens )
245+ )
246+ choices = getattr (response , "choices" , None )
247+ if not choices or getattr (choices [0 ], "message" , None ) is None :
248+ raise RuntimeError (
249+ f"malformed completion from provider: { type (response ).__name__ } "
250+ )
251+ return response
252+
192253 def _account (self , usage : Any , totals : dict [str , int ]) -> None :
193254 totals ["input" ] += self ._usage_value (usage , "prompt_tokens" )
194255 totals ["output" ] += self ._usage_value (usage , "completion_tokens" )
@@ -211,25 +272,49 @@ async def run(
211272 {"role" : "user" , "content" : instruction },
212273 ]
213274 totals = {"input" : 0 , "output" : 0 , "cached" : 0 }
275+ deadline = time .monotonic () + AGENT_BUDGET_SEC
276+
277+ def remaining () -> float :
278+ return deadline - time .monotonic ()
279+
280+ completed = False
281+ out_of_time = False
282+ empty_turns = 0
283+ last_turn = 0
214284
215285 for turn in range (1 , MAX_TURNS + 1 ):
216- response = await self ._client .chat .completions .create (
217- ** self ._completion_kwargs (messages )
286+ last_turn = turn
287+ # Stop starting new work once only the final-answer reserve is left,
288+ # so the answer is written before harbor's AgentTimeoutError fires.
289+ if remaining () <= FINAL_RESERVE_SEC :
290+ out_of_time = True
291+ self ._trace ({"turn" : turn , "stopping" : "deadline" , "left" : remaining ()})
292+ break
293+ call_started = time .monotonic ()
294+ response = await self ._complete (
295+ messages , timeout = remaining () - FINAL_RESERVE_SEC / 2
218296 )
219- self ._account (response .usage , totals )
297+ llm_sec = time .monotonic () - call_started
298+ self ._account (getattr (response , "usage" , None ), totals )
220299 message = response .choices [0 ].message
300+ finish_reason = getattr (response .choices [0 ], "finish_reason" , None )
221301 calls = message .tool_calls or []
222302 self ._trace (
223303 {
224304 "turn" : turn ,
225305 "content" : message .content ,
306+ "finish_reason" : finish_reason ,
307+ "llm_sec" : round (llm_sec , 2 ),
308+ "left_sec" : round (remaining (), 1 ),
226309 "tool_calls" : [
227310 {"name" : c .function .name , "arguments" : c .function .arguments }
228311 for c in calls
229312 ],
230313 }
231314 )
232315 # Record the assistant turn (with any tool calls) in the history.
316+ # An assistant message carrying neither is rejected by strict
317+ # OpenAI-compatible servers, so it is dropped rather than appended.
233318 assistant : dict [str , Any ] = {"role" : "assistant" }
234319 if message .content :
235320 assistant ["content" ] = message .content
@@ -245,47 +330,88 @@ async def run(
245330 }
246331 for c in calls
247332 ]
248- messages .append (assistant )
333+ if len (assistant ) > 1 :
334+ messages .append (assistant )
249335
250336 if not calls :
251- if (message .content or "" ).strip ():
337+ # `length` means the model was cut off mid-thought, so its
338+ # content is truncated reasoning, not an answer. This target
339+ # writes its chain-of-thought into ordinary `content` (it has no
340+ # separate reasoning_content channel) and can run past the token
341+ # cap, and the truncated ramble was previously submitted verbatim
342+ # as the final answer, which the judge always scores 0. Ask for
343+ # the answer instead of submitting the thinking.
344+ truncated = finish_reason == "length"
345+ if (message .content or "" ).strip () and not truncated :
252346 await self ._submit (environment , message .content )
347+ completed = True
253348 context .metadata = {
254349 "turns" : turn ,
255350 "trace" : "browsecomp-plus-trace.jsonl" ,
256351 }
257352 break
258- raise RuntimeError ("model returned neither a response nor a tool call" )
353+ empty_turns += 1
354+ self ._trace (
355+ {"turn" : turn , "recover" : empty_turns , "truncated" : truncated }
356+ )
357+ if empty_turns >= MAX_EMPTY_RETRIES :
358+ break
359+ messages .append (
360+ {
361+ "role" : "user" ,
362+ "content" : (
363+ "You were cut off mid-thought. Stop reasoning and "
364+ "reply with only the three required labeled lines "
365+ "(Explanation, Exact Answer, Confidence), keeping the "
366+ "explanation to a few sentences."
367+ if truncated
368+ else "Your last message was empty. Either call a tool "
369+ "or give your final response in the required format."
370+ ),
371+ }
372+ )
373+ continue
374+ empty_turns = 0
259375
260376 submitted = False
261377 for call in calls :
262- try :
263- arguments = json .loads (call .function .arguments )
264- if call .function .name == "search" :
265- result = await self ._index_command (
266- environment , "search" , arguments ["query" ]
267- )
268- elif call .function .name == "get_document" :
269- result = await self ._index_command (
270- environment , "get-document" , arguments ["docid" ]
271- )
272- elif call .function .name == "submit_response" :
273- await self ._submit (environment , arguments ["response" ])
274- result = {"submitted" : True }
275- submitted = True
276- else :
277- result = {"error" : f"unknown tool: { call .function .name } " }
278- except (
279- json .JSONDecodeError ,
280- KeyError ,
281- OSError ,
282- TypeError ,
283- ValueError ,
284- ) as error :
285- result = {"error" : str (error )}
286- self ._trace (
287- {"turn" : turn , "tool" : call .function .name , "result" : result }
288- )
378+ # Every tool_call must still get a tool message even once the
379+ # deadline hits, or the forced-final request below is rejected
380+ # for an unanswered tool call. So stub, never skip.
381+ if out_of_time or remaining () <= FINAL_RESERVE_SEC :
382+ if not out_of_time :
383+ out_of_time = True
384+ self ._trace ({"turn" : turn , "stopping" : "deadline-midturn" })
385+ result = {"error" : "search budget exhausted" }
386+ else :
387+ try :
388+ arguments = json .loads (call .function .arguments )
389+ budget = remaining () - FINAL_RESERVE_SEC
390+ if call .function .name == "search" :
391+ result = await self ._index_command (
392+ environment , "search" , arguments ["query" ], budget
393+ )
394+ elif call .function .name == "get_document" :
395+ result = await self ._index_command (
396+ environment , "get-document" , arguments ["docid" ], budget
397+ )
398+ elif call .function .name == "submit_response" :
399+ await self ._submit (environment , arguments ["response" ])
400+ result = {"submitted" : True }
401+ submitted = True
402+ else :
403+ result = {"error" : f"unknown tool: { call .function .name } " }
404+ except (
405+ json .JSONDecodeError ,
406+ KeyError ,
407+ OSError ,
408+ TypeError ,
409+ ValueError ,
410+ ) as error :
411+ result = {"error" : str (error )}
412+ self ._trace (
413+ {"turn" : turn , "tool" : call .function .name , "result" : result }
414+ )
289415 messages .append (
290416 {
291417 "role" : "tool" ,
@@ -298,39 +424,71 @@ async def run(
298424 if submitted :
299425 break
300426 if submitted :
427+ completed = True
301428 context .metadata = {
302429 "turns" : turn ,
303430 "trace" : "browsecomp-plus-trace.jsonl" ,
304431 }
305432 break
306- else :
307- # Turn budget exhausted: force one final tool-free response from what
308- # was retrieved rather than crashing, so the case scores best-effort
433+ if out_of_time :
434+ break
435+
436+ if not completed :
437+ # Budget exhausted (wall clock, turns, or a model that kept coming
438+ # back empty): force one final tool-free response from what was
439+ # retrieved rather than crashing, so the case scores best-effort
309440 # instead of being lost with no answer recorded.
441+ stub = (
442+ "Explanation: no answer could be determined within the search "
443+ "budget.\n Exact Answer: unknown\n Confidence: 0%"
444+ )
310445 messages .append (
311446 {
312447 "role" : "user" ,
313448 "content" : (
314- "You have used your full search budget. Give your single "
315- "best final response now, in the required format, based on "
316- "the documents you have retrieved."
449+ "You have used your full search budget. Do not reason "
450+ "further. Reply with only the three required labeled "
451+ "lines (Explanation, Exact Answer, Confidence), keeping "
452+ "the explanation to a few sentences, giving your single "
453+ "best answer from the documents you have retrieved."
317454 ),
318455 }
319456 )
320- final = await self ._client .chat .completions .create (
321- ** self ._completion_kwargs (messages , tools = False )
322- )
323- self ._account (final .usage , totals )
324- answer = (final .choices [0 ].message .content or "" ).strip () or (
325- "Explanation: no answer could be determined within the search "
326- "budget.\n Exact Answer: unknown\n Confidence: 0%"
327- )
457+ # This call is the last thing standing between a completed trial and
458+ # a lost one, so a failure here degrades to the stub answer instead
459+ # of propagating.
460+ answer = stub
461+ try :
462+ # Bounded retries: the SDK default of 8 could outlast the
463+ # reserve on its own and hand the trial back to the killer.
464+ # A tight token cap keeps the model from spending the whole
465+ # reserve on chain-of-thought and getting truncated again.
466+ final = await self ._complete (
467+ messages ,
468+ tools = False ,
469+ timeout = max (30.0 , remaining () - 20 ),
470+ max_retries = 1 ,
471+ max_tokens = ANSWER_MAX_TOKENS ,
472+ )
473+ self ._account (getattr (final , "usage" , None ), totals )
474+ answer = (final .choices [0 ].message .content or "" ).strip () or stub
475+ except Exception as error : # noqa: BLE001 - last-resort safety net
476+ self ._trace ({"forced_final_error" : f"{ type (error ).__name__ } : { error } " })
328477 await self ._submit (environment , answer )
329- self ._trace ({"turn" : MAX_TURNS , "forced_final_answer" : answer })
478+ self ._trace (
479+ {
480+ "turn" : last_turn ,
481+ "forced_final_answer" : answer ,
482+ "reason" : "deadline" if out_of_time else "turns" ,
483+ }
484+ )
330485 context .metadata = {
331- "turns" : MAX_TURNS ,
486+ "turns" : last_turn ,
332487 "trace" : "browsecomp-plus-trace.jsonl" ,
333488 "forced_final" : True ,
489+ # Lets a post-hoc conservative score treat exactly the trials
490+ # that the pre-fix agent would have lost as zeros.
491+ "forced_final_reason" : "deadline" if out_of_time else "turns" ,
334492 }
335493
336494 context .n_input_tokens = totals ["input" ]
0 commit comments