wehub-resource-sync
127 lines
4.5 KiB
Python
127 lines
4.5 KiB
Python
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)
|