chore: import upstream snapshot with attribution

src/agents/sandbox/memory/phase_one.py

@@ -0,0 +1,126 @@
+from __future__ import annotations
+
+import json
+import re
+from pathlib import Path
+
+from ...run_config import RunConfig
+from ..config import MemoryGenerateConfig
+from ..sandbox_agent import SandboxAgent
+from ..util.token_truncation import TruncationPolicy, truncate_text
+from .interface import RolloutExtractionArtifacts
+from .prompts import (
+    render_rollout_extraction_prompt,
+    render_rollout_extraction_user_prompt,
+)
+
+_ROLLOUT_SLUG_RE = re.compile(r"^[a-z0-9][a-z0-9_-]{0,79}$")
+_ROLLOUT_ID_RE = re.compile(r"^[A-Za-z0-9][A-Za-z0-9._-]{0,127}$")
+_PHASE_ONE_ROLLOUT_TOKEN_LIMIT = 150_000
+_PHASE_ONE_ROLLOUT_OMISSION_MARKER_TEMPLATE = (
+    "\n\n"
+    "[rollout content omitted: this phase-one memory prompt contains a truncated view of "
+    "the saved rollout. original_chars={original_chars}; rendered_chars={rendered_chars}. "
+    "Do not assume the rendered rollout below is complete.]"
+    "\n\n"
+)
+
+
+def normalize_rollout_slug(value: str) -> str:
+    slug = value.strip()
+    if slug.endswith(".md"):
+        slug = slug[:-3]
+    if not _ROLLOUT_SLUG_RE.fullmatch(slug):
+        raise ValueError(f"Invalid rollout_slug: {value!r}")
+    return slug
+
+
+def rollout_id_from_rollout_path(value: str) -> str:
+    rollout_id = Path(Path(value).name.strip()).stem
+    if not rollout_id or not _ROLLOUT_ID_RE.fullmatch(rollout_id):
+        raise ValueError(f"Invalid rollout id for memory: {value!r}")
+    return rollout_id
+
+
+def render_phase_one_prompt(*, rollout_contents: str) -> str:
+    payloads = [json.loads(line) for line in rollout_contents.splitlines() if line.strip()]
+    if not payloads:
+        raise ValueError("rollout_contents must contain at least one JSONL record")
+    payload = payloads[-1]
+    if len(payloads) == 1:
+        terminal_metadata: object = payload.get("terminal_metadata", {})
+    else:
+        terminal_metadata = {
+            "segment_count": len(payloads),
+            "final_terminal_metadata": payload.get("terminal_metadata", {}),
+            "terminal_states": [
+                item.get("terminal_metadata", {}).get("terminal_state", "unknown")
+                for item in payloads
+                if isinstance(item, dict)
+            ],
+        }
+    terminal_metadata_json = json.dumps(
+        terminal_metadata,
+        sort_keys=True,
+        separators=(",", ":"),
+        indent=2,
+    )
+    # TODO: Replace this fixed cap with 70% of the phase-one model's effective
+    # context window once model metadata is available in the SDK.
+    truncated_rollout_contents = truncate_text(
+        rollout_contents,
+        TruncationPolicy.tokens(_PHASE_ONE_ROLLOUT_TOKEN_LIMIT),
+    )
+    if truncated_rollout_contents != rollout_contents:
+        marker = _PHASE_ONE_ROLLOUT_OMISSION_MARKER_TEMPLATE.format(
+            original_chars=len(rollout_contents),
+            rendered_chars=len(truncated_rollout_contents),
+        )
+        truncated_rollout_contents = marker + truncated_rollout_contents
+    return render_rollout_extraction_user_prompt(
+        terminal_metadata_json=terminal_metadata_json,
+        rollout_contents=truncated_rollout_contents,
+    )
+
+
+def validate_rollout_artifacts(artifacts: RolloutExtractionArtifacts) -> bool:
+    if (
+        artifacts.rollout_slug.strip() == ""
+        and artifacts.rollout_summary.strip() == ""
+        and artifacts.raw_memory.strip() == ""
+    ):
+        return False
+    if (
+        not artifacts.rollout_slug.strip()
+        or not artifacts.rollout_summary.strip()
+        or not artifacts.raw_memory.strip()
+    ):
+        raise ValueError("Phase 1 returned partially-empty memory artifacts.")
+    return True
+
+
+async def run_phase_one(
+    *,
+    config: MemoryGenerateConfig,
+    prompt: str,
+    run_config: RunConfig,
+) -> RolloutExtractionArtifacts:
+    from ...run import Runner
+
+    if config.phase_one_model_settings is None:
+        agent = SandboxAgent(
+            name="sandbox-memory-phase-one",
+            instructions=render_rollout_extraction_prompt(extra_prompt=config.extra_prompt),
+            output_type=RolloutExtractionArtifacts,
+            model=config.phase_one_model,
+        )
+    else:
+        agent = SandboxAgent(
+            name="sandbox-memory-phase-one",
+            instructions=render_rollout_extraction_prompt(extra_prompt=config.extra_prompt),
+            output_type=RolloutExtractionArtifacts,
+            model=config.phase_one_model,
+            model_settings=config.phase_one_model_settings,
+        )
+    result = await Runner.run(agent, prompt, run_config=run_config)
+    return result.final_output_as(RolloutExtractionArtifacts, raise_if_incorrect_type=True)