feat: add agentic pd hybrid benchmark prototype

src/agentic_pd_hybrid/trace_profiles.py

@@ -0,0 +1,127 @@
+from __future__ import annotations
+
+import json
+from collections import defaultdict
+from dataclasses import asdict, dataclass
+from math import ceil
+from pathlib import Path
+
+from agentic_pd_hybrid.trace import TraceRequest, load_trace
+
+
+BLOCK_TOKEN_BUDGET = 24
+
+
+@dataclass(frozen=True)
+class NormalizeTraceLengthsConfig:
+    trace_path: Path
+    output_path: Path
+    initial_input_length: int = 10_000
+    append_input_length: int = 1_000
+    output_length: int = 1_000
+    max_requests: int | None = None
+
+
+@dataclass(frozen=True)
+class NormalizeTraceLengthsSummary:
+    input_trace_path: str
+    output_trace_path: str
+    request_count: int
+    session_count: int
+    multi_turn_session_count: int
+    initial_input_length: int
+    append_input_length: int
+    output_length: int
+    max_turns_per_session: int
+    max_input_length: int
+
+
+def normalize_trace_lengths(
+    config: NormalizeTraceLengthsConfig,
+) -> NormalizeTraceLengthsSummary:
+    if config.initial_input_length < 0:
+        raise ValueError("initial_input_length must be >= 0")
+    if config.append_input_length < 0:
+        raise ValueError("append_input_length must be >= 0")
+    if config.output_length < 0:
+        raise ValueError("output_length must be >= 0")
+
+    requests = load_trace(config.trace_path, request_limit=config.max_requests)
+    sessions: dict[str, list[TraceRequest]] = defaultdict(list)
+    for request in requests:
+        sessions[request.session_id].append(request)
+
+    normalized_records: list[dict[str, object]] = []
+    max_turns_per_session = 0
+    max_input_length = 0
+
+    for session_idx, session_id in enumerate(sorted(sessions, key=_session_sort_key)):
+        session_requests = sorted(
+            sessions[session_id],
+            key=lambda request: (request.timestamp_s, request.turn_id, request.chat_id),
+        )
+        max_turns_per_session = max(max_turns_per_session, len(session_requests))
+        base_block_count = ceil(config.initial_input_length / BLOCK_TOKEN_BUDGET)
+        base_hash_ids = [
+            _hash_id_for(session_idx=session_idx, block_idx=block_idx)
+            for block_idx in range(base_block_count)
+        ]
+
+        for turn_idx, request in enumerate(session_requests):
+            input_length = config.initial_input_length + turn_idx * (
+                config.append_input_length + config.output_length
+            )
+            total_block_count = ceil(input_length / BLOCK_TOKEN_BUDGET)
+            hash_ids = base_hash_ids + [
+                _hash_id_for(
+                    session_idx=session_idx,
+                    block_idx=base_block_count + append_block_idx,
+                )
+                for append_block_idx in range(max(0, total_block_count - base_block_count))
+            ]
+            max_input_length = max(max_input_length, input_length)
+            normalized_records.append(
+                {
+                    "chat_id": request.chat_id,
+                    "parent_chat_id": request.parent_chat_id,
+                    "timestamp": request.timestamp_s,
+                    "input_length": input_length,
+                    "output_length": config.output_length,
+                    "type": request.request_type,
+                    "turn": request.turn_id,
+                    "hash_ids": hash_ids,
+                }
+            )
+
+    normalized_records.sort(key=lambda item: float(item["timestamp"]))
+    config.output_path.parent.mkdir(parents=True, exist_ok=True)
+    with config.output_path.open("w", encoding="utf-8") as handle:
+        for record in normalized_records:
+            handle.write(json.dumps(record, sort_keys=True) + "\n")
+
+    summary = NormalizeTraceLengthsSummary(
+        input_trace_path=str(config.trace_path),
+        output_trace_path=str(config.output_path),
+        request_count=len(normalized_records),
+        session_count=len(sessions),
+        multi_turn_session_count=sum(
+            1 for session_requests in sessions.values() if len(session_requests) > 1
+        ),
+        initial_input_length=config.initial_input_length,
+        append_input_length=config.append_input_length,
+        output_length=config.output_length,
+        max_turns_per_session=max_turns_per_session,
+        max_input_length=max_input_length,
+    )
+    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
+
+
+def _hash_id_for(*, session_idx: int, block_idx: int) -> int:
+    return session_idx * 1_000_000 + block_idx
+
+
+def _session_sort_key(session_id: str) -> tuple[int, str]:
+    return (0, session_id) if session_id.isdigit() else (1, session_id)