296 lines
10 KiB
Python
296 lines
10 KiB
Python
from __future__ import annotations
|
|
|
|
import hashlib
|
|
import json
|
|
from collections import defaultdict
|
|
from dataclasses import asdict, dataclass
|
|
from pathlib import Path
|
|
from typing import Literal
|
|
|
|
from agentic_pd_hybrid.trace import TraceRequest, load_trace
|
|
|
|
|
|
SampleProfile = Literal["default", "small-append"]
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class SessionSampleConfig:
|
|
trace_path: Path
|
|
output_path: Path
|
|
target_duration_s: float = 600.0
|
|
start_time_s: float = 0.0
|
|
session_sample_rate: float = 1.0
|
|
min_turns: int = 1
|
|
max_requests: int | None = None
|
|
profile: SampleProfile = "default"
|
|
min_initial_input_tokens: int | None = None
|
|
max_initial_input_tokens: int | None = None
|
|
max_append_input_tokens: int | None = None
|
|
max_output_tokens: int | None = None
|
|
min_overlap_ratio: float | None = None
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class SessionSampleSummary:
|
|
input_trace_path: str
|
|
output_trace_path: str
|
|
request_count: int
|
|
session_count: int
|
|
multi_turn_session_count: int
|
|
start_time_s: float
|
|
end_time_s: float
|
|
sampled_duration_s: float
|
|
session_sample_rate: float
|
|
min_turns: int
|
|
profile: str
|
|
min_initial_input_tokens: int | None
|
|
max_initial_input_tokens: int | None
|
|
max_append_input_tokens: int | None
|
|
max_output_tokens: int | None
|
|
min_overlap_ratio: float | None
|
|
mean_append_input_tokens: float | None
|
|
mean_turn_overlap_ratio: float | None
|
|
|
|
|
|
def sample_trace_sessions(config: SessionSampleConfig) -> SessionSampleSummary:
|
|
requests = load_trace(config.trace_path)
|
|
sessions: dict[str, list[TraceRequest]] = defaultdict(list)
|
|
for request in requests:
|
|
sessions[request.session_id].append(request)
|
|
|
|
filters = _resolve_filters(config)
|
|
eligible_sessions = {
|
|
session_id: session_requests
|
|
for session_id, session_requests in sessions.items()
|
|
if len(session_requests) >= filters.min_turns
|
|
and _session_matches_filters(session_requests, filters)
|
|
and _keep_session(session_id, config.session_sample_rate)
|
|
}
|
|
ordered_sessions = sorted(
|
|
eligible_sessions.values(),
|
|
key=lambda session_requests: session_requests[0].timestamp_s,
|
|
)
|
|
|
|
selected_requests: list[TraceRequest] = []
|
|
sampled_start: float | None = None
|
|
sampled_end: float | None = None
|
|
for session_requests in ordered_sessions:
|
|
session_first = session_requests[0].timestamp_s
|
|
if session_first < config.start_time_s:
|
|
continue
|
|
|
|
if sampled_start is None:
|
|
sampled_start = session_first
|
|
|
|
selected_requests.extend(session_requests)
|
|
sampled_end = max(request.timestamp_s for request in session_requests)
|
|
|
|
if config.max_requests is not None and len(selected_requests) >= config.max_requests:
|
|
break
|
|
if sampled_end - sampled_start >= config.target_duration_s:
|
|
break
|
|
|
|
selected_requests.sort(key=lambda request: request.timestamp_s)
|
|
if config.max_requests is not None:
|
|
selected_requests = selected_requests[: config.max_requests]
|
|
|
|
if not selected_requests:
|
|
raise ValueError("Sampling produced no requests; adjust the sampling arguments")
|
|
|
|
config.output_path.parent.mkdir(parents=True, exist_ok=True)
|
|
with config.output_path.open("w", encoding="utf-8") as handle:
|
|
for request in selected_requests:
|
|
payload = {
|
|
"request_id": request.request_id,
|
|
"session_id": request.session_id,
|
|
"chat_id": request.chat_id,
|
|
"parent_chat_id": request.parent_chat_id,
|
|
"timestamp": request.timestamp_s,
|
|
"input_length": request.input_length,
|
|
"output_length": request.output_length,
|
|
"type": request.request_type,
|
|
"turn": request.turn_id,
|
|
"hash_ids": list(request.hash_ids),
|
|
}
|
|
handle.write(json.dumps(payload, sort_keys=True) + "\n")
|
|
|
|
selected_session_ids = {request.session_id for request in selected_requests}
|
|
selected_session_requests = [
|
|
eligible_sessions[session_id] for session_id in selected_session_ids
|
|
]
|
|
append_lengths = [
|
|
length
|
|
for session_requests in selected_session_requests
|
|
for length in _turn_append_lengths(session_requests)
|
|
]
|
|
overlap_ratios = [
|
|
ratio
|
|
for session_requests in selected_session_requests
|
|
for ratio in _turn_overlap_ratios(session_requests)
|
|
]
|
|
summary = SessionSampleSummary(
|
|
input_trace_path=str(config.trace_path),
|
|
output_trace_path=str(config.output_path),
|
|
request_count=len(selected_requests),
|
|
session_count=len(selected_session_ids),
|
|
multi_turn_session_count=sum(
|
|
1
|
|
for session_id in selected_session_ids
|
|
if len(eligible_sessions[session_id]) > 1
|
|
),
|
|
start_time_s=selected_requests[0].timestamp_s,
|
|
end_time_s=selected_requests[-1].timestamp_s,
|
|
sampled_duration_s=selected_requests[-1].timestamp_s
|
|
- selected_requests[0].timestamp_s,
|
|
session_sample_rate=config.session_sample_rate,
|
|
min_turns=filters.min_turns,
|
|
profile=config.profile,
|
|
min_initial_input_tokens=filters.min_initial_input_tokens,
|
|
max_initial_input_tokens=filters.max_initial_input_tokens,
|
|
max_append_input_tokens=filters.max_append_input_tokens,
|
|
max_output_tokens=filters.max_output_tokens,
|
|
min_overlap_ratio=filters.min_overlap_ratio,
|
|
mean_append_input_tokens=_mean(append_lengths),
|
|
mean_turn_overlap_ratio=_mean(overlap_ratios),
|
|
)
|
|
summary_path = config.output_path.with_suffix(config.output_path.suffix + ".summary.json")
|
|
with summary_path.open("w", encoding="utf-8") as handle:
|
|
json.dump(asdict(summary), handle, indent=2, sort_keys=True)
|
|
return summary
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class _ResolvedFilters:
|
|
min_turns: int
|
|
min_initial_input_tokens: int | None
|
|
max_initial_input_tokens: int | None
|
|
max_append_input_tokens: int | None
|
|
max_output_tokens: int | None
|
|
min_overlap_ratio: float | None
|
|
|
|
|
|
def _resolve_filters(config: SessionSampleConfig) -> _ResolvedFilters:
|
|
if config.profile == "default":
|
|
return _ResolvedFilters(
|
|
min_turns=config.min_turns,
|
|
min_initial_input_tokens=config.min_initial_input_tokens,
|
|
max_initial_input_tokens=config.max_initial_input_tokens,
|
|
max_append_input_tokens=config.max_append_input_tokens,
|
|
max_output_tokens=config.max_output_tokens,
|
|
min_overlap_ratio=config.min_overlap_ratio,
|
|
)
|
|
|
|
if config.profile != "small-append":
|
|
raise ValueError(f"Unsupported sample profile: {config.profile}")
|
|
|
|
return _ResolvedFilters(
|
|
min_turns=max(config.min_turns, 2),
|
|
min_initial_input_tokens=(
|
|
2048
|
|
if config.min_initial_input_tokens is None
|
|
else config.min_initial_input_tokens
|
|
),
|
|
max_initial_input_tokens=(
|
|
16000
|
|
if config.max_initial_input_tokens is None
|
|
else config.max_initial_input_tokens
|
|
),
|
|
max_append_input_tokens=(
|
|
2048
|
|
if config.max_append_input_tokens is None
|
|
else config.max_append_input_tokens
|
|
),
|
|
max_output_tokens=(
|
|
2048 if config.max_output_tokens is None else config.max_output_tokens
|
|
),
|
|
min_overlap_ratio=(
|
|
0.75 if config.min_overlap_ratio is None else config.min_overlap_ratio
|
|
),
|
|
)
|
|
|
|
|
|
def _session_matches_filters(
|
|
session_requests: list[TraceRequest],
|
|
filters: _ResolvedFilters,
|
|
) -> bool:
|
|
ordered = sorted(
|
|
session_requests,
|
|
key=lambda request: (request.timestamp_s, request.turn_id, request.chat_id),
|
|
)
|
|
if not ordered:
|
|
return False
|
|
|
|
initial = ordered[0]
|
|
if (
|
|
filters.min_initial_input_tokens is not None
|
|
and initial.input_length < filters.min_initial_input_tokens
|
|
):
|
|
return False
|
|
if (
|
|
filters.max_initial_input_tokens is not None
|
|
and initial.input_length > filters.max_initial_input_tokens
|
|
):
|
|
return False
|
|
if filters.max_output_tokens is not None and any(
|
|
request.output_length > filters.max_output_tokens for request in ordered
|
|
):
|
|
return False
|
|
|
|
append_lengths = _turn_append_lengths(ordered)
|
|
if filters.max_append_input_tokens is not None and any(
|
|
append_length <= 0 or append_length > filters.max_append_input_tokens
|
|
for append_length in append_lengths
|
|
):
|
|
return False
|
|
|
|
overlap_ratios = _turn_overlap_ratios(ordered)
|
|
if filters.min_overlap_ratio is not None and any(
|
|
overlap_ratio < filters.min_overlap_ratio for overlap_ratio in overlap_ratios
|
|
):
|
|
return False
|
|
|
|
return True
|
|
|
|
|
|
def _turn_append_lengths(session_requests: list[TraceRequest]) -> list[int]:
|
|
ordered = sorted(
|
|
session_requests,
|
|
key=lambda request: (request.timestamp_s, request.turn_id, request.chat_id),
|
|
)
|
|
return [
|
|
current.input_length - (previous.input_length + previous.output_length)
|
|
for previous, current in zip(ordered, ordered[1:], strict=False)
|
|
]
|
|
|
|
|
|
def _turn_overlap_ratios(session_requests: list[TraceRequest]) -> list[float]:
|
|
ordered = sorted(
|
|
session_requests,
|
|
key=lambda request: (request.timestamp_s, request.turn_id, request.chat_id),
|
|
)
|
|
ratios: list[float] = []
|
|
for previous, current in zip(ordered, ordered[1:], strict=False):
|
|
if not current.hash_ids:
|
|
ratios.append(0.0)
|
|
continue
|
|
previous_blocks = set(previous.hash_ids)
|
|
overlap = sum(1 for block in current.hash_ids if block in previous_blocks)
|
|
ratios.append(overlap / len(current.hash_ids))
|
|
return ratios
|
|
|
|
|
|
def _mean(values: list[int] | list[float]) -> float | None:
|
|
if not values:
|
|
return None
|
|
return sum(values) / len(values)
|
|
|
|
|
|
def _keep_session(session_id: str, sample_rate: float) -> bool:
|
|
if sample_rate >= 1.0:
|
|
return True
|
|
if sample_rate <= 0.0:
|
|
return False
|
|
digest = hashlib.blake2b(session_id.encode("utf-8"), digest_size=8).digest()
|
|
bucket = int.from_bytes(digest, byteorder="big", signed=False) / 2**64
|
|
return bucket < sample_rate
|