synth-ai 0.2.14__py3-none-any.whl → 0.4.1__py3-none-any.whl

examples/task_apps/enron/task_app/grpo_enron.py

@@ -1,906 +0,0 @@
-"""Task App configuration for the GRPO Enron email QA example."""
-
-from __future__ import annotations
-
-import contextlib
-import json
-import logging
-import os
-import time
-from dataclasses import dataclass
-from datetime import datetime, timezone
-from pathlib import Path
-from typing import Any, Iterable, Sequence
-from uuid import UUID, uuid4
-
-from datasets import load_dataset
-import httpx
-
-from fastapi import HTTPException
-
-from synth_ai.environments.examples.enron.environment import EnronEnvironment
-from synth_ai.environments.examples.enron.taskset import (
-    EnronTaskInstance,
-    EnronTaskInstanceMetadata,
-)
-from synth_ai.environments.tasks.core import (
-    Impetus,
-    Intent,
-    SplitInfo,
-    TaskInstanceSet,
-)
-from synth_ai.task.apps import ModalDeploymentConfig, TaskAppEntry, register_task_app
-from synth_ai.task.contracts import (
-    RolloutMetrics,
-    RolloutRequest,
-    RolloutResponse,
-    RolloutStep,
-    RolloutTrajectory,
-    TaskInfo,
-)
-from synth_ai.task.datasets import TaskDatasetRegistry, TaskDatasetSpec
-from synth_ai.task.rubrics import load_rubric
-from synth_ai.task.server import ProxyConfig, RubricBundle, TaskAppConfig
-from synth_ai.task.tracing_utils import (
-    build_tracer_factory,
-    resolve_sft_output_dir,
-    resolve_tracing_db_url,
-    tracing_env_enabled,
-)
-from synth_ai.tracing_v3.session_tracer import SessionTracer
-from synth_ai.environments.environment.tools import EnvToolCall
-
-logger = logging.getLogger(__name__)
-
-_HERE = Path(__file__).resolve()
-REPO_ROOT = _HERE.parents[4]
-
-DATASET_SPEC = TaskDatasetSpec(
-    id="enron_email_qa",
-    name="Enron Email QA",
-    version="1.0.0",
-    splits=["train", "test"],
-    default_split="train",
-    description="Question answering over a sample of Enron emails.",
-)
-
-HF_DATASET_ID = "corbt/enron_emails_sample_questions"
-HF_CACHE_DIR = os.path.join(
-    os.getenv("ENRON_DATASET_CACHE_DIR", str(REPO_ROOT / ".cache" / "hf-datasets"))
-)
-
-TOOLS = ["search_emails", "read_email", "answer_question", "terminate"]
-GROQ_CHAT_URL = "https://api.groq.com/openai/v1/chat/completions"
-DEFAULT_GROQ_MODEL = "qwen/qwen3-32b"
-ENRON_SYSTEM_PROMPT = (
-    "You are an Enron investigations analyst. Answer the user's question by reading emails. "
-    "You can call tools to search the corpus, read specific messages, and submit a final answer. "
-    "Use the tools deliberately, gather evidence before answering, and when confident call "
-    "answer_question with your final answer. If you cannot find the answer after thorough search, "
-    "answer_question with your best attempt noting uncertainty."
-)
-
-
-def _simplify(obj: Any) -> Any:
-    if isinstance(obj, (str, int, float, bool)) or obj is None:
-        return obj
-    if isinstance(obj, dict):
-        return {str(k): _simplify(v) for k, v in obj.items()}
-    if isinstance(obj, (list, tuple, set)):
-        return [_simplify(v) for v in obj]
-    return str(obj)
-
-
-def _render_search_results(results: list[dict[str, Any]]) -> str:
-    if not results:
-        return "No search results."
-    lines = []
-    for item in results[:5]:
-        message_id = item.get("message_id") or item.get("id") or "<unknown>"
-        snippet = (item.get("snippet") or item.get("snip") or "").strip()
-        lines.append(f"- {message_id}: {snippet[:280]}")
-    return "\n".join(lines)
-
-
-def _render_email(email: dict[str, Any] | None) -> str:
-    if not email:
-        return "No email loaded."
-    subject = email.get("subject", "<no subject>")
-    from_addr = email.get("from_address") or email.get("from_addr") or "<unknown>"
-    date = email.get("date", "<unknown date>")
-    snippet = (email.get("body") or "")[:600]
-    return f"Subject: {subject}\nFrom: {from_addr}\nDate: {date}\nBody Preview:\n{snippet}"
-
-
-def _render_observation(obs: dict[str, Any]) -> str:
-    lines = [
-        f"Question: {obs.get('question', '')}",
-        f"Already answered: {bool(obs.get('already_answered'))}",
-        f"Available tools: {', '.join(obs.get('tools') or [])}",
-        f"Inbox address: {obs.get('inbox_address', '<unknown>')}",
-        f"Reward Δ: {obs.get('reward_last', 0)}   Total Reward: {obs.get('total_reward', 0)}",
-    ]
-    tool_error = obs.get("tool_error")
-    if tool_error:
-        lines.append(f"Last tool error: {tool_error}")
-    search_results = obs.get("search_results") or []
-    if search_results:
-        lines.append("Search Results:")
-        lines.append(_render_search_results(search_results))
-    email = obs.get("email")
-    if email:
-        lines.append("Email Content:")
-        lines.append(_render_email(email))
-    gold = obs.get("gold_answer")
-    if gold and obs.get("terminated"):
-        lines.append(f"Gold Answer: {gold}")
-    return "\n".join(lines)
-
-
-def _conversation_message(role: str, content: Any, **metadata: Any) -> dict[str, Any]:
-    if isinstance(content, (dict, list)):
-        rendered = json.dumps(_simplify(content), ensure_ascii=False)
-    else:
-        rendered = str(content)
-    message: dict[str, Any] = {"role": role, "content": rendered}
-    message.update({k: v for k, v in metadata.items() if v is not None})
-    return message
-
-
-def _build_trace_payload_enron(
-    run_id: str,
-    request: RolloutRequest,
-    steps: list[RolloutStep],
-    metrics: RolloutMetrics,
-    *,
-    provider: str,
-    model: str,
-    conversation: list[dict[str, Any]],
-    metadata: dict[str, Any] | None = None,
-) -> dict[str, Any]:
-    created_at = datetime.now(timezone.utc)
-    event_time = time.time()
-    session_steps: list[dict[str, Any]] = []
-    event_history: list[dict[str, Any]] = []
-    markov_history: list[dict[str, Any]] = []
-    for msg in conversation:
-        event_time += 0.005
-        markov_history.append(
-            {
-                "content": {"text": msg.get("content", "")},
-                "message_type": msg.get("role", "system"),
-                "time_record": {"event_time": event_time},
-                "metadata": _simplify({k: v for k, v in msg.items() if k not in {"role", "content"}}),
-            }
-        )
-
-    session_trace = {
-        "session_id": run_id,
-        "created_at": created_at.isoformat(),
-        "metadata": {
-            "task": "enron_email_qa",
-            "provider": provider,
-            "model": model,
-            "policy": _simplify(request.policy.model_dump() if request.policy else {}),
-            "env": _simplify(request.env.model_dump() if request.env else {}),
-            **(_simplify(metadata or {})),
-        },
-        "session_time_steps": session_steps,
-        "event_history": event_history,
-        "markov_blanket_message_history": markov_history,
-    }
-
-    return {
-        "version": 3,
-        "session_trace": session_trace,
-        "run_id": run_id,
-        "policy_id": request.policy.policy_id or request.policy.policy_name,
-        "reward": metrics.mean_return,
-        "episode_returns": metrics.episode_returns,
-        "mean_return": metrics.mean_return,
-        "num_steps": metrics.num_steps,
-    }
-
-
-async def _call_groq_chat(
-    client: httpx.AsyncClient,
-    api_key: str,
-    payload: dict[str, Any],
-) -> tuple[dict[str, Any], dict[str, Any]]:
-    response = await client.post(
-        GROQ_CHAT_URL,
-        json=payload,
-        headers={"Authorization": f"Bearer {api_key}"},
-    )
-    if response.status_code >= 400:
-        try:
-            body = response.json()
-        except Exception:
-            body = {"raw": response.text}
-        detail = {
-            "status": response.status_code,
-            "body": body,
-            "headers": dict(response.headers),
-        }
-        raise HTTPException(status_code=response.status_code, detail=detail)
-    data = response.json()
-    return data, {
-        "status": response.status_code,
-        "headers": dict(response.headers),
-        "body": data,
-    }
-
-
-def _load_taskset_blocking() -> TaskInstanceSet:
-    """Build the Enron taskset synchronously."""
-
-    cache_dir = Path(HF_CACHE_DIR)
-    cache_dir.mkdir(parents=True, exist_ok=True)
-
-    ds_train = load_dataset(HF_DATASET_ID, split="train", cache_dir=cache_dir)
-    ds_test = load_dataset(HF_DATASET_ID, split="test", cache_dir=cache_dir)
-
-    def _metadata_from_row(row: dict[str, Any], split: str) -> EnronTaskInstance:
-        question = str(row.get("question") or "").strip()
-        answer = str(row.get("answer") or "").strip()
-        message_ids = row.get("message_ids") or []
-        if not isinstance(message_ids, list):
-            message_ids = list(message_ids)
-        impetus = Impetus(instructions=question)
-        intent = Intent(
-            rubric={"goal": "Answer the question using the Enron emails."},
-            gold_trajectories=None,
-            gold_state_diff={"answer": answer},
-        )
-        metadata = EnronTaskInstanceMetadata(
-            split=split,
-            email_count=len(message_ids),
-            message_ids=message_ids,
-        )
-        return EnronTaskInstance(
-            id=uuid4(),
-            impetus=impetus,
-            intent=intent,
-            metadata=metadata,
-            is_reproducible=True,
-            initial_engine_snapshot=row,
-        )
-
-    train_instances = [_metadata_from_row(r, "train") for r in ds_train]
-    test_instances = [_metadata_from_row(r, "test") for r in ds_test]
-
-    split_info = SplitInfo(
-        val_instance_ids=set(),
-        test_instance_ids={inst.id for inst in test_instances},
-        _is_split_defined=True,
-    )
-
-    return TaskInstanceSet(
-        name="Enron-QA",
-        description="QA over Enron email dataset sample.",
-        instances=train_instances + test_instances,
-        split_info=split_info,
-    )
-
-
-def _safe_uuid(value: Any) -> UUID:
-    if isinstance(value, UUID):
-        return value
-    try:
-        return UUID(str(value))
-    except Exception:
-        return UUID(int=0)
-
-
-@dataclass
-class EnronDataset:
-    spec: TaskDatasetSpec
-
-    def __post_init__(self) -> None:
-        self._taskset = _load_taskset_blocking()
-        self.instances: list[EnronTaskInstance] = list(self._taskset.instances)
-        self.instance_ids = [str(_safe_uuid(inst.id)) for inst in self.instances]
-        self.default_seed = 0
-        self.seed_min = 0
-        self.seed_max = max(len(self.instances) - 1, 0)
-
-    def describe(self) -> dict[str, Any]:
-        return {
-            **self.spec.model_dump(),
-            "instance_count": len(self.instances),
-            "instance_ids": self.instance_ids[:50],
-        }
-
-    def instance_by_seed(self, seed: int | None) -> EnronTaskInstance:
-        if not self.instances:
-            raise ValueError("Enron dataset is empty.")
-        if seed is None:
-            index = 0
-        else:
-            index = int(seed) % len(self.instances)
-        return self.instances[index]
-
-
-def build_dataset() -> tuple[TaskDatasetRegistry, EnronDataset]:
-    registry = TaskDatasetRegistry()
-    dataset = EnronDataset(DATASET_SPEC)
-    registry.register(DATASET_SPEC, lambda _spec: dataset, cache=True)
-    return registry, dataset
-
-
-def _base_task_info(dataset: EnronDataset) -> TaskInfo:
-    return TaskInfo(
-        task={"id": "enron_email_qa", "name": "Enron Email QA", "version": "1.0.0"},
-        environment="enron",
-        action_space={
-            "type": "tool_calls",
-            "tools": TOOLS,
-            "description": "Tool-assisted QA workflow over an email corpus.",
-        },
-        observation={
-            "summary": "Text observations describing the question, tool status, and last reward.",
-            "format": "text",
-        },
-        dataset={**dataset.describe(), "default_seed": dataset.default_seed},
-        rubric={
-            "version": "1",
-            "criteria_count": 1,
-            "source": "inline",
-            "aggregation": "weighted_sum",
-        },
-        inference={
-            "supports_proxy": False,
-            "endpoints": {},
-            "tool": {"name": "enron_tools", "parallel_tool_calls": False},
-        },
-        limits={"max_ops": 0, "max_time_s": 900},
-    )
-
-
-OUTCOME_RUBRIC = load_rubric(
-    {
-        "version": "1",
-        "goal_text": "Provide the correct answer to the question using the Enron emails.",
-        "aggregation": "weighted_sum",
-        "criteria": [
-            {
-                "id": "accuracy",
-                "description": "Final answer matches the gold answer.",
-                "weight": 1.0,
-            }
-        ],
-    }
-)
-
-EVENTS_RUBRIC = load_rubric(
-    {
-        "version": "1",
-        "goal_text": "Encourage efficient use of tools when exploring the corpus.",
-        "aggregation": "weighted_sum",
-        "criteria": [
-            {
-                "id": "tool_use",
-                "description": "Use search, read, and answer tools deliberately.",
-                "weight": 1.0,
-            }
-        ],
-    }
-)
-
-
-def describe_taskset(dataset: EnronDataset) -> dict[str, Any]:
-    return dataset.describe()
-
-
-def provide_task_instances(
-    dataset: EnronDataset, base_info: TaskInfo, seeds: Sequence[int]
-) -> Iterable[TaskInfo]:
-    infos: list[TaskInfo] = []
-    base_observation = getattr(base_info, "observation", None)
-    if hasattr(base_observation, "model_dump"):
-        observation_template = base_observation.model_dump()
-    elif isinstance(base_observation, dict):
-        observation_template = dict(base_observation)
-    else:
-        observation_template = {}
-
-    for seed in seeds:
-        instance = dataset.instance_by_seed(seed)
-        metadata = instance.metadata
-        meta_dict = {
-            "split": getattr(metadata, "split", None),
-            "email_count": getattr(metadata, "email_count", None),
-            "message_ids": getattr(metadata, "message_ids", None),
-        }
-        infos.append(
-            TaskInfo(
-                task=base_info.task,
-                environment=base_info.environment,
-                action_space=base_info.action_space,
-                observation={
-                    **observation_template,
-                    "question": instance.impetus.instructions,
-                },
-                dataset={
-                    **base_info.dataset.model_dump(),
-                    "instance_id": str(_safe_uuid(instance.id)),
-                    "metadata": meta_dict,
-                },
-                rubric=base_info.rubric,
-                inference=base_info.inference,
-                limits=base_info.limits,
-            )
-        )
-    return infos
-
-
-def _ensure_dataset_from_state(fastapi_request, fallback: EnronDataset) -> EnronDataset:
-    if fastapi_request is None:
-        return fallback
-    dataset = getattr(getattr(fastapi_request, "app", None), "state", None)
-    candidate = getattr(dataset, "dataset", None)
-    return candidate or fallback
-
-
-def _normalise_observation(value: Any) -> dict[str, Any]:
-    if isinstance(value, dict):
-        return value
-    if hasattr(value, "observation"):
-        obs = getattr(value, "observation")
-        if isinstance(obs, dict):
-            return obs
-        return {"text": str(obs)}
-    return {"text": str(value)}
-
-
-async def rollout_executor(request: RolloutRequest, fastapi_request) -> RolloutResponse:
-    policy_cfg = dict(request.policy.config or {})
-    provider = str(policy_cfg.get("provider") or "").strip().lower()
-    if provider == "groq":
-        return await _rollout_with_groq(request, fastapi_request, policy_cfg)
-
-    # Fallback: return initial observation but include minimal trace payload
-    dataset = _ensure_dataset_from_state(fastapi_request, RUNTIME_DATASET)
-    env_seed = getattr(request.env, "seed", None) if request and request.env else None
-    instance = dataset.instance_by_seed(env_seed)
-    env = EnronEnvironment(task_instance=instance)
-    env.custom_obs = None
-    try:
-        initial_observation = await env.initialize()
-    finally:
-        with contextlib.suppress(Exception):
-            await env.terminate()
-
-    obs_dict = _normalise_observation(initial_observation)
-    step = RolloutStep(
-        obs=obs_dict,
-        tool_calls=[],
-        reward=0.0,
-        done=True,
-        truncated=None,
-        info={"note": "No rollout executed; provider unset."},
-    )
-    # No inference_url for noop policy
-    trajectory = RolloutTrajectory(
-        env_id=request.env.env_id or "enron",
-        policy_id=request.policy.policy_id or request.policy.policy_name or "noop-policy",
-        steps=[step],
-        final={"observation": obs_dict},
-        length=1,
-        inference_url=None,  # NEW: No inference for noop policy
-        decision_samples=None,
-    )
-    metrics = RolloutMetrics(
-        episode_returns=[0.0],
-        mean_return=0.0,
-        num_steps=1,
-        num_episodes=1,
-        outcome_score=None,
-        events_score=None,
-        details={"note": "Provider not configured; returning initial state."},
-    )
-    trace_payload = _build_trace_payload_enron(
-        request.run_id,
-        request,
-        [step],
-        metrics,
-        provider="local",
-        model=policy_cfg.get("model") or "noop",
-        conversation=[
-            _conversation_message("system", ENRON_SYSTEM_PROMPT),
-            _conversation_message("user", _render_observation(obs_dict)),
-        ],
-        metadata={"mode": "noop"},
-    )
-    return RolloutResponse(
-        run_id=request.run_id,
-        trajectories=[trajectory],
-        branches={},
-        metrics=metrics,
-        aborted=False,
-        ops_executed=0,
-        trace=trace_payload,
-    )
-
-
-def _prepare_tool_call(
-    tool_name: str,
-    raw_args: dict[str, Any],
-    current_obs: dict[str, Any],
-) -> EnvToolCall:
-    if tool_name == "search_emails":
-        keywords = raw_args.get("keywords")
-        if isinstance(keywords, str):
-            keywords = [k.strip() for k in keywords.split(",") if k.strip()]
-        if not isinstance(keywords, list) or not keywords:
-            raise ValueError("search_emails requires a non-empty list of keywords.")
-        inbox = raw_args.get("inbox") or current_obs.get("inbox_address") or "investigator@enron.com"
-        args = {
-            "inbox": str(inbox),
-            "keywords": [str(k) for k in keywords],
-            "from_addr": raw_args.get("from_addr"),
-            "to_addr": raw_args.get("to_addr"),
-            "sent_after": raw_args.get("sent_after"),
-            "sent_before": raw_args.get("sent_before"),
-            "max_results": int(raw_args.get("max_results") or 5),
-        }
-        return EnvToolCall(tool="search_emails", args=args)
-
-    if tool_name == "read_email":
-        message_id = raw_args.get("message_id")
-        if not message_id:
-            raise ValueError("read_email requires 'message_id'.")
-        return EnvToolCall(tool="read_email", args={"message_id": str(message_id)})
-
-    if tool_name == "answer_question":
-        answer = raw_args.get("answer")
-        if not isinstance(answer, str) or not answer.strip():
-            raise ValueError("answer_question requires a non-empty 'answer'.")
-        return EnvToolCall(tool="answer_question", args={"answer": answer.strip()})
-
-    if tool_name == "terminate":
-        return EnvToolCall(tool="terminate", args={})
-
-    raise ValueError(f"Unsupported tool '{tool_name}'")
-
-
-async def _rollout_with_groq(
-    request: RolloutRequest,
-    fastapi_request,
-    config: dict[str, Any],
-) -> RolloutResponse:
-    api_key = os.getenv("GROQ_API_KEY")
-    if not api_key:
-        raise HTTPException(
-            status_code=503,
-            detail="GROQ_API_KEY environment variable is required for Groq rollouts.",
-        )
-
-    dataset = _ensure_dataset_from_state(fastapi_request, RUNTIME_DATASET)
-    env_seed = getattr(request.env, "seed", None) if request and request.env else None
-    instance = dataset.instance_by_seed(env_seed)
-    env = EnronEnvironment(task_instance=instance)
-    env.custom_obs = None
-
-    metadata_extra = {
-        "split": getattr(instance.metadata, "split", None),
-        "email_count": getattr(instance.metadata, "email_count", None),
-        "message_ids": list(getattr(instance.metadata, "message_ids", []))[:10],
-    }
-
-    model = config.get("model") or DEFAULT_GROQ_MODEL
-    temperature = float(config.get("temperature", 0.2) or 0.2)
-    top_p = float(config.get("top_p", 0.8) or 0.8)
-    max_tokens = int(config.get("max_tokens", 768) or 768)
-    max_turns = int(config.get("max_turns", config.get("max_steps", 12)) or 12)
-
-    tool_schemas = [
-        {
-            "type": "function",
-            "function": {
-                "name": "search_emails",
-                "description": "Search the Enron corpus for emails matching keywords.",
-                "parameters": {
-                    "type": "object",
-                    "properties": {
-                        "inbox": {"type": "string", "description": "Email address performing the search."},
-                        "keywords": {
-                            "type": "array",
-                            "items": {"type": "string"},
-                            "minItems": 1,
-                            "description": "Keywords to include in the search.",
-                        },
-                        "from_addr": {"type": "string"},
-                        "to_addr": {"type": "string"},
-                        "sent_after": {"type": "string", "description": "YYYY-MM-DD"},
-                        "sent_before": {"type": "string", "description": "YYYY-MM-DD"},
-                        "max_results": {"type": "integer", "minimum": 1, "maximum": 10},
-                    },
-                    "required": ["keywords"],
-                    "additionalProperties": False,
-                },
-            },
-        },
-        {
-            "type": "function",
-            "function": {
-                "name": "read_email",
-                "description": "Read the full contents of an email by message_id.",
-                "parameters": {
-                    "type": "object",
-                    "properties": {"message_id": {"type": "string"}},
-                    "required": ["message_id"],
-                    "additionalProperties": False,
-                },
-            },
-        },
-        {
-            "type": "function",
-            "function": {
-                "name": "answer_question",
-                "description": "Submit the final answer to the investigation question.",
-                "parameters": {
-                    "type": "object",
-                    "properties": {"answer": {"type": "string"}},
-                    "required": ["answer"],
-                    "additionalProperties": False,
-                },
-            },
-        },
-        {
-            "type": "function",
-            "function": {
-                "name": "terminate",
-                "description": "Terminate the investigation without answering.",
-                "parameters": {"type": "object", "properties": {}, "additionalProperties": False},
-            },
-        },
-    ]
-
-    steps: list[RolloutStep] = []
-    conversation: list[dict[str, Any]] = []
-    executed = 0
-    try:
-        observation = await env.initialize()
-        obs_dict = _normalise_observation(observation)
-        conversation.append(_conversation_message("system", ENRON_SYSTEM_PROMPT))
-        conversation.append(_conversation_message("user", _render_observation(obs_dict)))
-
-        async with httpx.AsyncClient(timeout=httpx.Timeout(60.0)) as client:
-            for turn in range(max_turns):
-                payload = {
-                    "model": model,
-                    "messages": conversation,
-                    "temperature": temperature,
-                    "top_p": top_p,
-                    "max_tokens": max_tokens,
-                    "tools": tool_schemas,
-                    "tool_choice": "auto",
-                }
-                vendor_attempts: list[dict[str, Any]] = []
-                response, response_meta = await _call_groq_chat(client, api_key, payload)
-                vendor_attempts.append({"request": payload, "response": response_meta})
-
-                choices = response.get("choices") or []
-                if not choices:
-                    break
-                message = choices[0].get("message") or {}
-                tool_calls = message.get("tool_calls") or []
-                assistant_msg_meta = {"tool_calls": _simplify(tool_calls)} if tool_calls else {}
-                conversation.append(
-                    _conversation_message("assistant", message.get("content") or "", **assistant_msg_meta)
-                )
-
-                tool_call_records: list[dict[str, Any]] = []
-                step_reward = 0.0
-                done = False
-                truncated = False
-
-                if not tool_calls:
-                    final_answer = (message.get("content") or "").strip()
-                    if final_answer:
-                        env_call = EnvToolCall(tool="answer_question", args={"answer": final_answer})
-                        observation = await env.step(env_call)
-                        executed += 1
-                        obs_dict = _normalise_observation(observation)
-                        step_reward += float(obs_dict.get("reward_last") or 0.0)
-                        done = bool(obs_dict.get("terminated"))
-                        truncated = bool(obs_dict.get("truncated"))
-                        tool_call_records.append({"tool": "answer_question", "args": env_call.args})
-                        conversation.append(
-                            _conversation_message(
-                                "tool",
-                                {"result": "answer_submitted", "observation": obs_dict},
-                                name="answer_question",
-                            )
-                        )
-                    else:
-                        break
-                else:
-                    for call in tool_calls:
-                        func = call.get("function") or {}
-                        name = func.get("name")
-                        raw_args = func.get("arguments")
-                        if isinstance(raw_args, str):
-                            try:
-                                parsed_args = json.loads(raw_args)
-                            except json.JSONDecodeError:
-                                parsed_args = {}
-                        elif isinstance(raw_args, dict):
-                            parsed_args = raw_args
-                        else:
-                            parsed_args = {}
-
-                        env_call = _prepare_tool_call(name, parsed_args, obs_dict)
-                        observation = await env.step(env_call)
-                        executed += 1
-                        obs_dict = _normalise_observation(observation)
-                        reward_delta = float(obs_dict.get("reward_last") or 0.0)
-                        step_reward += reward_delta
-                        done = bool(obs_dict.get("terminated"))
-                        truncated = bool(obs_dict.get("truncated"))
-                        tool_call_records.append({"tool": env_call.tool, "args": env_call.args})
-                        conversation.append(
-                            _conversation_message(
-                                "tool",
-                                {
-                                    "tool": env_call.tool,
-                                    "args": env_call.args,
-                                    "reward_delta": reward_delta,
-                                    "observation": obs_dict,
-                                },
-                                name=env_call.tool,
-                                tool_call_id=call.get("id"),
-                            )
-                        )
-                        if done or truncated:
-                            break
-
-                conversation.append(_conversation_message("user", _render_observation(obs_dict)))
-
-                step = RolloutStep(
-                    obs=obs_dict,
-                    tool_calls=tool_call_records,
-                    reward=step_reward,
-                    done=done,
-                    truncated=truncated if truncated else None,
-                    info={
-                        "provider": "groq",
-                        "model": model,
-                        "vendor_attempts": vendor_attempts,
-                        "turn": turn,
-                    },
-                )
-                steps.append(step)
-
-                if done or truncated:
-                    break
-    finally:
-        with contextlib.suppress(Exception):
-            await env.terminate()
-
-    if steps:
-        final_obs = steps[-1].obs
-        total_reward = float(final_obs.get("total_reward") or 0.0)
-    else:
-        total_reward = 0.0
-
-    metrics = RolloutMetrics(
-        episode_returns=[total_reward],
-        mean_return=total_reward if steps else 0.0,
-        num_steps=len(steps),
-        num_episodes=1,
-        outcome_score=None,
-        events_score=None,
-        details={"provider": "groq", "model": model},
-    )
-    inference_url_groq = "https://api.groq.com/openai/v1/chat/completions"
-    
-    trajectory = RolloutTrajectory(
-        env_id=request.env.env_id or "enron",
-        policy_id=request.policy.policy_id or request.policy.policy_name or "enron-groq",
-        steps=steps,
-        final={"observation": steps[-1].obs if steps else {}},
-        length=len(steps),
-        inference_url=inference_url_groq,  # NEW: Required for trace correlation
-        decision_samples=None,
-    )
-    trace_payload = _build_trace_payload_enron(
-        request.run_id,
-        request,
-        steps,
-        metrics,
-        provider="groq",
-        model=model,
-        conversation=conversation,
-        metadata=metadata_extra,
-    )
-    return RolloutResponse(
-        run_id=request.run_id,
-        trajectories=[trajectory],
-        branches={},
-        metrics=metrics,
-        aborted=False,
-        ops_executed=executed,
-        trace=trace_payload,
-    )
-
-
-RUNTIME_DATASET: EnronDataset
-registry, RUNTIME_DATASET = build_dataset()
-BASE_INFO = _base_task_info(RUNTIME_DATASET)
-
-
-def build_config() -> TaskAppConfig:
-    tracing_enabled = tracing_env_enabled()
-    tracing_db_url = resolve_tracing_db_url()
-    tracer_factory = build_tracer_factory(
-        SessionTracer, enabled=tracing_enabled, db_url=tracing_db_url
-    )
-    sft_output_dir = resolve_sft_output_dir()
-
-    app_state: dict[str, Any] = {
-        "dataset": RUNTIME_DATASET,
-        "allowed_environments": ["enron"],
-        "tracing_enabled": tracing_enabled,
-    }
-    if tracer_factory is not None:
-        app_state["session_tracer_factory"] = tracer_factory
-    if sft_output_dir:
-        app_state["sft_output_dir"] = sft_output_dir
-
-    if tracing_enabled:
-        logger.info("[enron:tracing] enabled (db=%s)", tracing_db_url or "default")
-    else:
-        logger.info("[enron:tracing] disabled")
-    if sft_output_dir:
-        logger.info("[enron:sft] writing JSONL to %s", sft_output_dir)
-
-    config = TaskAppConfig(
-        app_id="grpo-enron",
-        name="GRPO Enron Email QA Task App",
-        description="Tool-assisted QA environment over Enron emails with GRPO-compatible endpoints.",
-        base_task_info=BASE_INFO,
-        describe_taskset=lambda: describe_taskset(RUNTIME_DATASET),
-        provide_task_instances=lambda seeds: provide_task_instances(RUNTIME_DATASET, BASE_INFO, seeds),
-        rollout=rollout_executor,
-        dataset_registry=registry,
-        rubrics=RubricBundle(outcome=OUTCOME_RUBRIC, events=EVENTS_RUBRIC),
-        proxy=ProxyConfig(enable_openai=False, enable_groq=False),
-        routers=(),
-        app_state=app_state,
-        cors_origins=["*"],
-    )
-    return config
-
-
-register_task_app(
-    entry=TaskAppEntry(
-        app_id="grpo-enron",
-        description="Enron email QA task app with rollout metadata endpoints.",
-        config_factory=build_config,
-        aliases=("enron", "enron-task"),
-        env_files=(str(REPO_ROOT / "backend" / ".env.dev"),),
-        modal=ModalDeploymentConfig(
-            app_name="grpo-enron-task-app",
-            python_version="3.11",
-            pip_packages=(
-                "fastapi>=0.100.0",
-                "uvicorn>=0.23.0",
-                "pydantic>=2.0.0",
-                "httpx>=0.24.0",
-                "python-dotenv>=1.0.1",
-                "datasets>=2.10.0",
-            ),
-            extra_local_dirs=(
-                (str(REPO_ROOT), "/opt/synth_ai_repo"),
-                (str(REPO_ROOT / "synth_ai"), "/opt/synth_ai_repo/synth_ai"),
-                (str(_HERE.parent), "/opt/synth_ai_repo/examples/task_apps/enron/task_app"),
-            ),
-            secret_names=("groq-api-key", "openai-api-key"),
-            memory=8192,
-            cpu=2.0,
-            max_containers=4,
-        ),
-    )
-)