Initial release: aircontext
This commit is contained in:
admin
260
scripts/core/compactor.py
Normal file
260
scripts/core/compactor.py
Normal file
@@ -0,0 +1,260 @@
|
||||
"""Compactor — runs out-of-band, mutates JSONL, signals wrapper to resume.
|
||||
|
||||
Invocation (by on_tool_use.py or /aircontext-now):
|
||||
python compactor.py --project <root> --transcript <jsonl> --session-id <sid>
|
||||
|
||||
Lifecycle:
|
||||
1. Acquire single-instance lock (compactor.lock).
|
||||
2. Mark state.compaction_in_progress=true (already set by hook, idempotent).
|
||||
3. Load active chain. If too short, abort.
|
||||
4. Backup JSONL.
|
||||
5. Slice tail (preserved) vs head (to compress).
|
||||
6. Build conversation text + system prompt (rules.md).
|
||||
7. Call LLM backend.
|
||||
8. Append [summary, continuation] messages with parentUuid=null chain head.
|
||||
9. Write state.compaction_ready=true so wrapper restarts claude.
|
||||
"""
|
||||
from __future__ import annotations
|
||||
import argparse
|
||||
import json
|
||||
import os
|
||||
import sys
|
||||
import time
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
_HERE = Path(__file__).resolve().parent.parent
|
||||
if str(_HERE) not in sys.path:
|
||||
sys.path.insert(0, str(_HERE))
|
||||
|
||||
from core.backends.base import build_backend # noqa: E402
|
||||
from core.config_loader import load_config # noqa: E402
|
||||
from core.jsonl_ops import ( # noqa: E402
|
||||
append_isolated_summary,
|
||||
backup_jsonl,
|
||||
find_active_chain,
|
||||
load_messages,
|
||||
rotate_snapshots,
|
||||
)
|
||||
from core.state import StateFile # noqa: E402
|
||||
|
||||
|
||||
MIN_CHAIN_TO_COMPACT = 6
|
||||
|
||||
|
||||
def _serialise_message_for_llm(msg: dict[str, Any]) -> str | None:
|
||||
"""Reduce a JSONL message to plain text the LLM can read.
|
||||
|
||||
Returns None for messages that should be skipped entirely
|
||||
(file-history-snapshot, queue-operation, etc.).
|
||||
"""
|
||||
t = msg.get("type")
|
||||
inner = msg.get("message") or {}
|
||||
if t == "user":
|
||||
content = inner.get("content")
|
||||
if isinstance(content, str):
|
||||
return f"USER: {content}"
|
||||
if isinstance(content, list):
|
||||
parts = [p.get("text", "") if isinstance(p, dict) else str(p) for p in content]
|
||||
return "USER: " + " ".join(p for p in parts if p)
|
||||
if t == "assistant":
|
||||
content = inner.get("content")
|
||||
if isinstance(content, str):
|
||||
return f"ASSISTANT: {content}"
|
||||
if isinstance(content, list):
|
||||
parts: list[str] = []
|
||||
for p in content:
|
||||
if not isinstance(p, dict):
|
||||
continue
|
||||
if p.get("type") == "text":
|
||||
parts.append(p.get("text", ""))
|
||||
elif p.get("type") == "tool_use":
|
||||
parts.append(
|
||||
f"[tool_use {p.get('name')}({json.dumps(p.get('input', {}), ensure_ascii=False)[:300]})]"
|
||||
)
|
||||
elif p.get("type") == "thinking":
|
||||
pass # drop
|
||||
return "ASSISTANT: " + " ".join(s for s in parts if s)
|
||||
if t == "tool_result":
|
||||
# `message.content` for tool_result is the result text
|
||||
content = inner.get("content") if isinstance(inner, dict) else msg.get("content")
|
||||
if isinstance(content, list):
|
||||
content = " ".join(
|
||||
(p.get("text", "") if isinstance(p, dict) else str(p)) for p in content
|
||||
)
|
||||
if not isinstance(content, str):
|
||||
return None
|
||||
return f"TOOL_RESULT: {content}"
|
||||
if t in ("system",):
|
||||
content = inner.get("content")
|
||||
if isinstance(content, str):
|
||||
return f"SYSTEM: {content}"
|
||||
return None
|
||||
|
||||
|
||||
def _truncate_long_tool_results(rendered: list[str], max_lines: int) -> list[str]:
|
||||
out: list[str] = []
|
||||
for r in rendered:
|
||||
if not r.startswith("TOOL_RESULT: "):
|
||||
out.append(r)
|
||||
continue
|
||||
body = r[len("TOOL_RESULT: "):]
|
||||
lines = body.splitlines()
|
||||
if len(lines) <= max_lines:
|
||||
out.append(r)
|
||||
continue
|
||||
head = "\n".join(lines[:50])
|
||||
tail = "\n".join(lines[-50:])
|
||||
out.append(
|
||||
"TOOL_RESULT: "
|
||||
+ head
|
||||
+ f"\n... [{len(lines) - 100} lines omitted by AirContext pre-truncation] ...\n"
|
||||
+ tail
|
||||
)
|
||||
return out
|
||||
|
||||
|
||||
def run(project: Path, transcript: Path, session_id: str) -> int:
|
||||
air = project / "AirContext"
|
||||
cfg_path = air / "config.yaml"
|
||||
state = StateFile(air / "state.json")
|
||||
|
||||
try:
|
||||
cfg = load_config(cfg_path)
|
||||
except Exception as e:
|
||||
print(f"[compactor] config load failed: {e}", file=sys.stderr)
|
||||
state.update(compaction_in_progress=False)
|
||||
return 2
|
||||
|
||||
if state.read().get("paused"):
|
||||
state.update(compaction_in_progress=False)
|
||||
return 0
|
||||
|
||||
# Single-instance lock (best-effort; sufficient for single-user dev environment).
|
||||
lock = air / "compactor.lock"
|
||||
try:
|
||||
fd = os.open(str(lock), os.O_CREAT | os.O_EXCL | os.O_WRONLY)
|
||||
os.write(fd, str(os.getpid()).encode())
|
||||
os.close(fd)
|
||||
except FileExistsError:
|
||||
print("[compactor] another compactor is already running; exiting", file=sys.stderr)
|
||||
return 0
|
||||
|
||||
try:
|
||||
return _run_locked(project, transcript, session_id, cfg, state, air)
|
||||
finally:
|
||||
try:
|
||||
os.unlink(lock)
|
||||
except OSError:
|
||||
pass
|
||||
|
||||
|
||||
def _run_locked(
|
||||
project: Path,
|
||||
transcript: Path,
|
||||
session_id: str,
|
||||
cfg: dict,
|
||||
state: StateFile,
|
||||
air: Path,
|
||||
) -> int:
|
||||
state.update(compaction_in_progress=True, last_session_id=session_id)
|
||||
msgs = load_messages(transcript)
|
||||
chain = find_active_chain(msgs)
|
||||
if len(chain) < MIN_CHAIN_TO_COMPACT:
|
||||
print(
|
||||
f"[compactor] chain too short ({len(chain)}); skipping",
|
||||
file=sys.stderr,
|
||||
)
|
||||
state.update(compaction_in_progress=False)
|
||||
return 0
|
||||
|
||||
comp_cfg = cfg.get("compaction") or {}
|
||||
safety = cfg.get("safety") or {}
|
||||
preserve_tail = int(comp_cfg.get("preserve_tail_messages", 10))
|
||||
drop_lines = int(comp_cfg.get("drop_tool_results_over_lines", 1000))
|
||||
rules_file = comp_cfg.get("rules_file", "rules.md")
|
||||
continuation_prompt = comp_cfg.get("continuation_prompt", "") or ""
|
||||
|
||||
rules_path = air / rules_file
|
||||
rules_text = rules_path.read_text(encoding="utf-8") if rules_path.exists() else ""
|
||||
|
||||
head = chain[:-preserve_tail] if preserve_tail > 0 else chain
|
||||
tail = chain[-preserve_tail:] if preserve_tail > 0 else []
|
||||
|
||||
rendered = [r for r in (_serialise_message_for_llm(m) for m in head) if r]
|
||||
rendered = _truncate_long_tool_results(rendered, drop_lines)
|
||||
conversation_text = "\n\n".join(rendered)
|
||||
if not conversation_text.strip():
|
||||
print("[compactor] nothing to summarise", file=sys.stderr)
|
||||
state.update(compaction_in_progress=False)
|
||||
return 0
|
||||
|
||||
system_prompt = (
|
||||
"You are a context-compression engine for a Claude Code session. "
|
||||
"Apply the user-supplied compression rules below to the conversation "
|
||||
"transcript that follows. Output ONLY the compressed summary text. "
|
||||
"Do not add preamble, headers, or apologies.\n\n"
|
||||
"=== USER COMPRESSION RULES ===\n"
|
||||
f"{rules_text}\n"
|
||||
"=== END RULES ==="
|
||||
)
|
||||
|
||||
backend = build_backend(cfg)
|
||||
print(
|
||||
f"[compactor] calling {cfg['backend']['endpoint']} model={cfg['backend']['model']} "
|
||||
f"head_msgs={len(head)} tail_preserved={len(tail)}",
|
||||
file=sys.stderr,
|
||||
)
|
||||
try:
|
||||
summary_text = backend.summarise(system_prompt, conversation_text)
|
||||
except Exception as e:
|
||||
print(f"[compactor] LLM call failed: {e}", file=sys.stderr)
|
||||
state.update(compaction_in_progress=False)
|
||||
return 3
|
||||
|
||||
if safety.get("dry_run"):
|
||||
print(
|
||||
f"[compactor] dry_run=true; summary length={len(summary_text)}; "
|
||||
"JSONL not modified",
|
||||
file=sys.stderr,
|
||||
)
|
||||
state.update(compaction_in_progress=False, last_compaction_unix=int(time.time()))
|
||||
return 0
|
||||
|
||||
if safety.get("backup", True):
|
||||
backup_jsonl(transcript, air / "snapshots", session_id)
|
||||
|
||||
template = chain[-1] # use the latest message's metadata as template
|
||||
summary_uuid, cont_uuid = append_isolated_summary(
|
||||
transcript,
|
||||
summary_text=summary_text,
|
||||
continuation_text=continuation_prompt or None,
|
||||
template_message=template,
|
||||
)
|
||||
print(
|
||||
f"[compactor] appended summary={summary_uuid} continuation={cont_uuid}",
|
||||
file=sys.stderr,
|
||||
)
|
||||
|
||||
rotate_snapshots(air / "snapshots", int(safety.get("max_snapshots", 50)))
|
||||
|
||||
state.update(
|
||||
compaction_in_progress=False,
|
||||
compaction_ready=True,
|
||||
pending_resume_session_id=session_id,
|
||||
last_compaction_unix=int(time.time()),
|
||||
)
|
||||
return 0
|
||||
|
||||
|
||||
def main() -> int:
|
||||
p = argparse.ArgumentParser()
|
||||
p.add_argument("--project", required=True)
|
||||
p.add_argument("--transcript", required=True)
|
||||
p.add_argument("--session-id", required=True)
|
||||
args = p.parse_args()
|
||||
return run(Path(args.project), Path(args.transcript), args.session_id)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(main())
|
||||
Reference in New Issue
Block a user