11import sys
2+ from dataclasses import dataclass
23from typing import Dict , Any , Optional , Union , Tuple
34import requests
45from rich .console import Console
78from .utils import truncate_output
89
910
11+ @dataclass
12+ class RunSummary :
13+ """Brief run summary from execution task response."""
14+
15+ id : str
16+ status : str
17+ name : Optional [str ] = None
18+ require_review : bool = False
19+
20+
21+ @dataclass
22+ class ExecutionTasksResult :
23+ """Result from get_execution_tasks containing tasks and run info."""
24+
25+ tasks : list
26+ run : Optional [RunSummary ] = None
27+
28+
1029def handle_api_error (e : requests .exceptions .HTTPError , console : Console ) -> None :
1130 """Extract and display error messages from API responses in a structured format."""
1231 status = getattr (e .response , "status_code" , None )
@@ -110,6 +129,7 @@ def start_optimization_run(
110129 auth_headers : dict = {},
111130 timeout : Union [int , Tuple [int , int ]] = (10 , 3650 ),
112131 api_keys : Optional [Dict [str , str ]] = None ,
132+ require_review : bool = False ,
113133) -> Optional [Dict [str , Any ]]:
114134 """Start the optimization run."""
115135 with console .status ("[bold green]Starting Optimization..." ):
@@ -128,6 +148,7 @@ def start_optimization_run(
128148 "eval_timeout" : eval_timeout ,
129149 "save_logs" : save_logs ,
130150 "log_dir" : log_dir ,
151+ "require_review" : require_review ,
131152 "metadata" : {"client_name" : "cli" , "client_version" : __pkg_version__ },
132153 }
133154 if api_keys :
@@ -315,3 +336,107 @@ def report_termination(
315336 except Exception as e :
316337 print (f"Warning: Failed to report termination to backend for run { run_id } : { e } " , file = sys .stderr )
317338 return False
339+
340+
341+ # --- Execution Queue API Functions ---
342+
343+
344+ def get_execution_tasks (
345+ run_id : str , auth_headers : dict = {}, timeout : Union [int , Tuple [int , int ]] = (5 , 30 )
346+ ) -> Optional [ExecutionTasksResult ]:
347+ """Poll for ready execution tasks.
348+
349+ Args:
350+ run_id: The run ID to get tasks for.
351+ auth_headers: Authentication headers.
352+ timeout: Request timeout.
353+
354+ Returns:
355+ ExecutionTasksResult with tasks and run summary, or None if request failed.
356+ """
357+ try :
358+ response = requests .get (
359+ f"{ __base_url__ } /execution-tasks/" , params = {"run_id" : run_id }, headers = auth_headers , timeout = timeout
360+ )
361+ response .raise_for_status ()
362+ data = response .json ()
363+
364+ # Extract run summary from top-level run field
365+ run_summary = None
366+ if data .get ("run" ):
367+ run_data = data ["run" ]
368+ run_summary = RunSummary (
369+ id = run_data ["id" ],
370+ status = run_data ["status" ],
371+ name = run_data .get ("name" ),
372+ require_review = run_data .get ("require_review" , False ),
373+ )
374+
375+ return ExecutionTasksResult (tasks = data .get ("tasks" , []), run = run_summary )
376+ except requests .exceptions .HTTPError :
377+ return None
378+ except Exception :
379+ return None
380+
381+
382+ def claim_execution_task (
383+ task_id : str , auth_headers : dict = {}, timeout : Union [int , Tuple [int , int ]] = (5 , 30 )
384+ ) -> Optional [Dict [str , Any ]]:
385+ """Claim an execution task.
386+
387+ Args:
388+ task_id: The task ID to claim.
389+ auth_headers: Authentication headers.
390+ timeout: Request timeout.
391+
392+ Returns:
393+ The claimed task with revision, or None if already claimed or error.
394+ """
395+ try :
396+ response = requests .post (f"{ __base_url__ } /execution-tasks/{ task_id } /claim" , headers = auth_headers , timeout = timeout )
397+ if response .status_code == 409 :
398+ return None # Already claimed
399+ response .raise_for_status ()
400+ return response .json ()
401+ except requests .exceptions .HTTPError :
402+ return None
403+ except Exception :
404+ return None
405+
406+
407+ def submit_execution_result (
408+ run_id : str ,
409+ task_id : str ,
410+ execution_output : str ,
411+ auth_headers : dict = {},
412+ timeout : Union [int , Tuple [int , int ]] = (10 , 3650 ),
413+ api_keys : Optional [Dict [str , str ]] = None ,
414+ ) -> Optional [Dict [str , Any ]]:
415+ """Submit execution result for a task.
416+
417+ Args:
418+ run_id: The run ID.
419+ task_id: The task ID being completed.
420+ execution_output: The execution output to submit.
421+ auth_headers: Authentication headers.
422+ timeout: Request timeout.
423+ api_keys: Optional API keys for LLM providers.
424+
425+ Returns:
426+ The suggest response, or None if request failed.
427+ """
428+ try :
429+ truncated_output = truncate_output (execution_output )
430+ request_json = {"execution_output" : truncated_output , "task_id" : task_id , "metadata" : {}}
431+ if api_keys :
432+ request_json ["api_keys" ] = api_keys
433+
434+ response = requests .post (
435+ f"{ __base_url__ } /runs/{ run_id } /suggest" , json = request_json , headers = auth_headers , timeout = timeout
436+ )
437+ response .raise_for_status ()
438+ return response .json ()
439+ except requests .exceptions .HTTPError :
440+ return None
441+ except Exception :
442+ return None
0 commit comments