"""Tool-driven inference for the metric tracker RL environment.""" from __future__ import annotations import asyncio import json import os import sys import textwrap from dataclasses import dataclass, field from typing import Any from openenv.core.containers.runtime.providers import LocalDockerProvider from openai import APIStatusError, OpenAI from websockets.exceptions import ConnectionClosedError from metric_tracker_rl import DEFAULT_TASK_ORDER, MetricTrackerRlAction, MetricTrackerRlEnv, get_task_spec from metric_tracker_rl.analysis_tools import available_analysis_methods from metric_tracker_rl.models import ( MetricSubmissionRow, MetricTrackerRlObservation, PayloadGeneratorMethod, ) IMAGE_NAME = (os.getenv("IMAGE_NAME") or "metric_tracker_rl:latest").strip() API_KEY = os.getenv("HF_TOKEN") or os.getenv("OPENAI_API_KEY") or os.getenv("API_KEY") API_BASE_URL = ( os.getenv("API_BASE_URL") or os.getenv("OPENAI_BASE_URL") or "https://router.huggingface.co/v1" ) MODEL_NAME = os.getenv("MODEL_NAME") or os.getenv("OPENAI_MODEL") or "Qwen/Qwen2.5-72B-Instruct" BASE_URL = os.getenv("BASE_URL") TASK_NAME = os.getenv("MetricTrackerRl_TASK", "multi_task_agent_baseline") BENCHMARK = os.getenv("MetricTrackerRl_BENCHMARK", "metric_tracker_rl") TEMPERATURE = float(os.getenv("TEMPERATURE", "0")) MAX_TOKENS = min(int(os.getenv("MAX_TOKENS", "1000")), 4096) MAX_TOOL_ROUNDS = int(os.getenv("MAX_TOOL_ROUNDS", "16")) CONNECT_TIMEOUT_S = float(os.getenv("OPENENV_CONNECT_TIMEOUT_S", "30")) MESSAGE_TIMEOUT_S = float(os.getenv("OPENENV_MESSAGE_TIMEOUT_S", "180")) DOCKER_WAIT_TIMEOUT_S = float(os.getenv("OPENENV_DOCKER_WAIT_TIMEOUT_S", "120")) TASK_RETRY_COUNT = int(os.getenv("OPENENV_TASK_RETRY_COUNT", "1")) SCORE_EPSILON = float(os.getenv("OPENENV_SCORE_EPSILON", "0.000001")) SYSTEM_PROMPT = textwrap.dedent( """ You are solving a multi-anomaly analytics benchmark with tool use. Rules: - Use only the shared safe analysis methods. - Do not request full hidden answers or assume direct access to ground truth. - Prefer declarative payload generators over manual row construction. - Start from the default reset observation only. - Start by trying `get_rate_spike_from_median_rows` across different conversion metrics to learn which metrics produce useful anomaly rows. - Compare candidate metrics, then refine with raw-data inspection and median/std methods only when needed. - Prefer: task_overview -> get_rate_spike_from_median_rows on several conversion metrics -> compare useful results -> payload_generator -> submit_payload_generator. - Keep notes brief and factual. """ ).strip() @dataclass class ToolRuntimeState: """Mutable state shared across tool calls.""" method_log: list[dict[str, Any]] = field(default_factory=list) last_preview: dict[str, Any] | None = None rewards: list[float] = field(default_factory=list) steps: int = 0 def log_start(task: str, env: str, model: str) -> None: print(f"[START] task={task} env={env} model={model}", flush=True) def log_method(tool_name: str, arguments: dict[str, Any], note: str) -> None: return None def log_payload_generator_methods(tool_name: str, generator_methods: list[dict[str, Any]]) -> None: return None def log_step(step: int, action: str, reward: float, done: bool, error: str | None) -> None: error_val = error if error else "null" clipped_reward = bounded_log_reward(reward) print( f"[STEP] step={step} action={action} reward={clipped_reward:.2f} done={str(done).lower()} error={error_val}", flush=True, ) def bounded_task_score(score: float) -> float: """Clamp task scores to the open interval (0, 1).""" return min(1.0 - SCORE_EPSILON, max(SCORE_EPSILON, score)) def bounded_log_reward(reward: float) -> float: """Clamp logged scores and rewards to [0.01, 0.99].""" return min(0.99, max(0.01, reward)) def log_end(success: bool, steps: int, score: float, rewards: list[float]) -> None: clipped_score = bounded_log_reward(score) rewards_str = ",".join(f"{bounded_log_reward(reward):.2f}" for reward in rewards) print( f"[END] success={str(success).lower()} steps={steps} score={clipped_score:.2f} rewards={rewards_str}", flush=True, ) def tool_schemas() -> list[dict[str, Any]]: """OpenAI-compatible tool definitions.""" shared_schemas = [] for spec in available_analysis_methods(): properties = {} required = [] if spec.name in {"rows_for_date", "hourly_rows_for_date", "detect_funnel_break", "check_impossible_counts"}: properties = {"date": {"type": "string"}} required = ["date"] elif spec.name in {"compare_rate_to_median", "compare_count_to_median"}: properties = { "date": {"type": "string"}, "entity_name": {"type": "string"}, } required = ["date", "entity_name"] elif spec.name == "list_suspicious_dates": properties = {"limit": {"type": "integer", "default": 10}} elif spec.name == "preview_submission": properties = { "rows": { "type": "array", "items": {"type": "object"}, } } elif spec.name == "show_raw_data": properties = {"limit": {"type": "integer", "default": 5}} elif spec.name in {"get_metric_median", "get_metric_std_dev_from_median"}: properties = { "metric_name": {"type": "string"}, "metric_names": {"type": "array", "items": {"type": "string"}}, } elif spec.name == "get_rows_with_abs_diff_from_median_gt": properties = { "metric_name": {"type": "string"}, "metric_names": {"type": "array", "items": {"type": "string"}}, "threshold": {"type": "number"}, } required = ["threshold"] elif spec.name in { "get_median_filter_rows", "get_rate_drop_from_median_rows", "get_rate_spike_from_median_rows", "get_absolute_drop_in_event_count_rows", "get_absolute_spike_in_event_count_rows", }: properties = { "metric_name": {"type": "string"}, "metric_names": {"type": "array", "items": {"type": "string"}}, "threshold_multiplier": {"type": "number"}, } required = ["threshold_multiplier"] elif spec.name in { "get_funnel_break_rows", "get_hourly_traffic_mix_shift_rows", "get_instrumentation_data_quality_issue_rows", }: properties = { "threshold_multiplier": {"type": "number"}, } required = ["threshold_multiplier"] elif spec.name == "payload_generator": properties = { "generator_methods": { "type": "array", "items": {"type": "object"}, } } required = ["generator_methods"] shared_schemas.append( { "type": "function", "function": { "name": spec.name, "description": spec.description, "parameters": { "type": "object", "properties": properties, "required": required, "additionalProperties": False, }, }, } ) shared_schemas.append( { "type": "function", "function": { "name": "submit_payload_generator", "description": "Submit declarative payload generator methods for environment-side payload generation and grading.", "parameters": { "type": "object", "properties": { "generator_methods": { "type": "array", "items": {"type": "object"}, } }, "required": ["generator_methods"], "additionalProperties": False, }, }, } ) shared_schemas.append( { "type": "function", "function": { "name": "submit_solution", "description": "Submit the final anomaly payload to the environment.", "parameters": { "type": "object", "properties": { "rows": { "type": "array", "items": {"type": "object"}, } }, "required": ["rows"], "additionalProperties": False, }, }, } ) return shared_schemas def build_initial_user_prompt(observation: MetricTrackerRlObservation) -> str: return textwrap.dedent( f""" Solve the RL environment with tools. Initial observation: {json.dumps(observation.model_dump(exclude={"debug"}), indent=2)} Prefer building a payload generator first, then submit it. Start by calling `get_rate_spike_from_median_rows` on several conversion metrics and see which ones return useful anomaly rows. If a metric returns nothing or low-signal rows, try another metric. For funnel, hourly mix, or data-quality tasks, use the family-specific generator methods instead. Final payload rows use: `date`, `entity_type`, `entity_name`, `anomaly_type`, `detection_method`, `baseline_value`, `observed_value`, `delta_value`, `severity`. Supported generator method example: `{{"method_name":"get_rate_spike_from_median_rows","threshold_multiplier":2.0}}` or `{{"method_name":"get_rate_spike_from_median_rows","metric_names":["app_open_to_order_placed","app_open_to_payment_successful"],"threshold_multiplier":2.0}}` Use shared analysis methods only. Prefer `submit_payload_generator` over `submit_solution`. """ ).strip() def create_chat_completion(client: OpenAI, **kwargs): try: return client.chat.completions.create(**kwargs) except APIStatusError as exc: if exc.status_code == 402: raise RuntimeError( "The configured inference provider rejected the request with HTTP 402. " "Your Hugging Face router credits are depleted. Update `.env.inference` " "with a working provider/key, or switch `API_BASE_URL`/`MODEL_NAME`." ) from exc raise def decode_arguments(raw_arguments: str | None) -> dict[str, Any]: if not raw_arguments: return {} return json.loads(raw_arguments) def preview_text(text: str, limit: int = 220) -> str: return text.replace("\n", " ")[:limit] def format_action(tool_name: str, arguments: dict[str, Any]) -> str: if not arguments: return f"{tool_name}()" return preview_text(f"{tool_name}({json.dumps(arguments, sort_keys=True)})") def step_error(result: Any) -> str | None: message = getattr(result.observation, "message", None) return message if result.observation.status == "error" and message else None def record_step( runtime_state: ToolRuntimeState, action: str, result: Any, ) -> None: reward = float(result.reward or 0.0) runtime_state.steps += 1 runtime_state.rewards.append(reward) log_step( step=runtime_state.steps, action=action, reward=reward, done=bool(result.done), error=step_error(result), ) async def connect_env() -> MetricTrackerRlEnv: if BASE_URL: client = MetricTrackerRlEnv( base_url=BASE_URL, connect_timeout_s=CONNECT_TIMEOUT_S, message_timeout_s=MESSAGE_TIMEOUT_S, ) return await client.connect() provider = LocalDockerProvider() base_url = provider.start_container(IMAGE_NAME) provider.wait_for_ready(base_url, timeout_s=DOCKER_WAIT_TIMEOUT_S) client = MetricTrackerRlEnv( base_url=base_url, connect_timeout_s=CONNECT_TIMEOUT_S, message_timeout_s=MESSAGE_TIMEOUT_S, provider=provider, ) return await client.connect() async def execute_tool_call( env: MetricTrackerRlEnv, observation: MetricTrackerRlObservation, runtime_state: ToolRuntimeState, tool_name: str, arguments: dict[str, Any], ) -> tuple[dict[str, Any], Any | None, MetricTrackerRlObservation]: """Execute one model-requested tool locally.""" action = format_action(tool_name, arguments) if tool_name == "submit_payload_generator": methods = [ PayloadGeneratorMethod(**item) for item in arguments.get("generator_methods", []) ] runtime_state.method_log.append( { "tool_name": tool_name, "arguments": arguments, "generator_methods": [item.model_dump() for item in methods], "note": _tool_note(tool_name, arguments), } ) result = await env.step(MetricTrackerRlAction(payload_generators=methods)) record_step(runtime_state, action, result) return ( { "status": result.observation.status, "message": result.observation.message, "reward": result.reward, "done": result.done, "generated_rows": [row.model_dump() for row in result.observation.generated_rows], "submission_issues": [issue.model_dump() for issue in result.observation.submission_issues], "reward_breakdown": ( result.observation.reward_breakdown.model_dump() if result.observation.reward_breakdown else None ), }, result, result.observation, ) if tool_name == "submit_solution": rows = [MetricSubmissionRow(**row) for row in arguments.get("rows", [])] result = await env.step(MetricTrackerRlAction(classifications=rows)) record_step(runtime_state, action, result) return ( { "status": result.observation.status, "message": result.observation.message, "reward": result.reward, "done": result.done, "reward_breakdown": ( result.observation.reward_breakdown.model_dump() if result.observation.reward_breakdown else None ), "issue_count": len(result.observation.submission_issues), "correct_row_count": result.observation.correct_row_count, }, result, result.observation, ) result = await env.step( MetricTrackerRlAction( analysis_method=tool_name, analysis_args=arguments, ) ) record_step(runtime_state, action, result) output = result.observation.analysis_result or { "method": tool_name, "arguments": arguments, "result": None, } log_arguments = { "tool_name": tool_name, "arguments": arguments, "note": _tool_note(tool_name, arguments), } if tool_name == "payload_generator": log_arguments["generator_methods"] = arguments.get("generator_methods", []) runtime_state.method_log.append( log_arguments ) if tool_name == "preview_submission": runtime_state.last_preview = output return output, None, result.observation def _tool_note(tool_name: str, arguments: dict[str, Any]) -> str: notes = { "task_overview": "bootstrap the task and payload schema", "list_dates": "confirm the date range", "list_entities": "confirm valid entities", "rows_for_date": "inspect daily counts on one date", "hourly_rows_for_date": "inspect hourly traffic shape", "compare_rate_to_median": "check a conversion-rate anomaly against median baseline", "compare_count_to_median": "check an absolute count anomaly against median baseline", "detect_funnel_break": "test whether a funnel step is broken", "check_impossible_counts": "test for instrumentation or impossible count issues", "list_suspicious_dates": "prioritize dates worth deeper inspection", "preview_submission": "validate payload structure before submit", "show_raw_data": "inspect daily aggregate rows in head() form", "get_metric_median": "measure a baseline median for one metric", "get_metric_std_dev_from_median": "measure metric spread around the median", "get_rows_with_abs_diff_from_median_gt": "inspect dates outside a chosen absolute-difference threshold", "get_median_filter_rows": "generate candidate anomaly rows using median and std-from-median filtering", "get_rate_drop_from_median_rows": "generate candidate conversion-rate drop rows using median and std-from-median filtering", "get_rate_spike_from_median_rows": "generate candidate conversion-rate spike rows using median and std-from-median filtering", "get_absolute_drop_in_event_count_rows": "generate candidate event-count drop rows using median and std-from-median filtering", "get_absolute_spike_in_event_count_rows": "generate candidate event-count spike rows using median and std-from-median filtering", "get_funnel_break_rows": "generate candidate funnel-break rows across funnel steps", "get_hourly_traffic_mix_shift_rows": "generate candidate hourly traffic mix shift rows across dates", "get_instrumentation_data_quality_issue_rows": "generate candidate impossible-count or instrumentation-issue rows across dates", "payload_generator": "merge multiple generator methods into one candidate payload", "submit_payload_generator": "submit generator methods for environment-side generation and grading", } return notes.get(tool_name, f"run {tool_name} with {arguments}") async def run_agent_loop( client: OpenAI, env: MetricTrackerRlEnv, observation: MetricTrackerRlObservation, ) -> tuple[Any, str, int, list[dict[str, Any]], ToolRuntimeState]: """Run a tool-calling loop until the env is solved or the round limit is hit.""" runtime_state = ToolRuntimeState() current_observation = observation messages: list[dict[str, Any]] = [ {"role": "system", "content": SYSTEM_PROMPT}, {"role": "user", "content": build_initial_user_prompt(current_observation)}, ] last_result = None final_text = "" tool_rounds = 0 for _ in range(MAX_TOOL_ROUNDS): completion = create_chat_completion( client, model=MODEL_NAME, messages=messages, tools=tool_schemas(), tool_choice="auto", temperature=TEMPERATURE, max_tokens=MAX_TOKENS, stream=False, ) message = completion.choices[0].message assistant_payload: dict[str, Any] = { "role": "assistant", "content": message.content or "", } if message.tool_calls: assistant_payload["tool_calls"] = [ { "id": tool_call.id, "type": tool_call.type, "function": { "name": tool_call.function.name, "arguments": tool_call.function.arguments, }, } for tool_call in message.tool_calls ] messages.append(assistant_payload) if not message.tool_calls: final_text = (message.content or "").strip() break tool_rounds += 1 for tool_call in message.tool_calls: tool_name = tool_call.function.name arguments = decode_arguments(tool_call.function.arguments) if tool_name != "submit_solution": log_method(tool_name, arguments, _tool_note(tool_name, arguments)) if tool_name in {"payload_generator", "submit_payload_generator"}: log_payload_generator_methods( tool_name, arguments.get("generator_methods", []), ) tool_output, maybe_result, current_observation = await execute_tool_call( env, current_observation, runtime_state, tool_name, arguments, ) messages.append( { "role": "tool", "tool_call_id": tool_call.id, "content": json.dumps(tool_output), } ) if maybe_result is not None: last_result = maybe_result if last_result is not None: completion = create_chat_completion( client, model=MODEL_NAME, messages=messages, temperature=TEMPERATURE, max_tokens=MAX_TOKENS, stream=False, ) final_text = (completion.choices[0].message.content or "").strip() break return last_result, final_text, tool_rounds, runtime_state.method_log, runtime_state async def run_single_task( client: OpenAI, env: MetricTrackerRlEnv, task_id: str, ) -> dict[str, Any]: """Run one named benchmark task and return a reproducible summary.""" task_spec = get_task_spec(task_id) log_start(task=task_spec.task_id, env=BENCHMARK, model=MODEL_NAME) reset_result = await env.reset(task_id=task_spec.task_id) final_result, final_text, tool_rounds, method_log, runtime_state = await run_agent_loop( client, env, reset_result.observation, ) if final_result is None: raise RuntimeError(f"The model never submitted a graded action for task `{task_spec.task_id}`.") reward = float(final_result.reward or 0.0) task_score = bounded_task_score(reward) success = bool(final_result.done and reward >= 0.999999) return { "task_id": task_spec.task_id, "difficulty": task_spec.difficulty, "objective": task_spec.objective, "grader_name": task_spec.grader_name, "score": task_score, "normalized_score": task_score, "done": final_result.done, "success": success, "final_status": final_result.observation.status, "final_message": final_result.observation.message, "issue_count": len(final_result.observation.submission_issues), "correct_row_count": final_result.observation.correct_row_count, "expected_row_count": final_result.observation.expected_row_count, "tool_rounds": tool_rounds, "assistant_summary": final_text, "steps": runtime_state.steps, "rewards": runtime_state.rewards, "reward_breakdown": ( final_result.observation.reward_breakdown.model_dump() if final_result.observation.reward_breakdown else None ), } async def run_single_task_with_retries( client: OpenAI, task_id: str, ) -> dict[str, Any]: """Run one task with a fresh env connection and bounded reconnect retries.""" attempts = TASK_RETRY_COUNT + 1 last_error: Exception | None = None for attempt in range(1, attempts + 1): env = None success = False steps = 0 score = 0.0 rewards: list[float] = [] try: env = await connect_env() summary = await run_single_task(client, env, task_id) success = bool(summary["success"]) steps = int(summary["steps"]) score = float(summary["score"]) rewards = list(summary["rewards"]) return summary except (ConnectionClosedError, ConnectionError, TimeoutError, OSError) as exc: last_error = exc print( ( f"[WARN] task_id={task_id} attempt={attempt}/{attempts} " f"env_connection_error={type(exc).__name__}: {exc}" ), flush=True, file=sys.stderr, ) if attempt >= attempts: raise finally: try: if env is not None: await env.close() except Exception: pass if env is not None: log_end(success=success, steps=steps, score=score, rewards=rewards) assert last_error is not None raise last_error async def main() -> None: if not API_KEY: raise RuntimeError("Set OPENAI_API_KEY, HF_TOKEN, or API_KEY.") client = OpenAI(base_url=API_BASE_URL, api_key=API_KEY) task_summaries: list[dict[str, Any]] = [] for task_id in DEFAULT_TASK_ORDER: task_summaries.append(await run_single_task_with_retries(client, task_id)) if __name__ == "__main__": asyncio.run(main())