247 lines
8.6 KiB
Python
247 lines
8.6 KiB
Python
"""What a conversation's model holds: prompt granules, skills, tools, files touched."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import re
|
|
from collections import Counter
|
|
from pathlib import Path
|
|
from typing import TYPE_CHECKING, Any
|
|
|
|
if TYPE_CHECKING:
|
|
from collections.abc import Iterable, Mapping
|
|
|
|
from beaver_gateway.agents.claude import ClaudeAgent
|
|
from beaver_gateway.conversations.kinds import Kind
|
|
|
|
__all__ = ["bash_paths", "compose", "files_touched", "granules", "skill_sets"]
|
|
|
|
CHARS_PER_TOKEN = 3.2
|
|
READ_TOOLS = frozenset({"Read", "NotebookRead"})
|
|
WRITE_TOOLS = frozenset({"Write", "Edit", "MultiEdit", "NotebookEdit"})
|
|
SEARCH_TOOLS = frozenset({"Glob", "Grep"})
|
|
PATH_KEYS = ("file_path", "notebook_path", "path", "file", "filename")
|
|
_FRONTMATTER = re.compile(r"^---\s*\n(.*?)\n---", re.DOTALL)
|
|
# A shell command names files as quoted strings or bare words; only notes
|
|
# and plain text files count, code and binaries are the agent's own business.
|
|
_SHELL_TOKEN = re.compile(r'"([^"\n]+)"|\'([^\'\n]+)\'|([^\s"\'`;|&()<>]+)')
|
|
_NOTE_SUFFIXES = (".md", ".canvas", ".txt", ".csv", ".json", ".yaml", ".yml")
|
|
_SHELL_CD = re.compile(r'(?:^|&&|;|\|\|)\s*cd\s+("([^"]+)"|\'([^\']+)\'|(\S+))')
|
|
_SHELL_WRITES = re.compile(
|
|
r"write_text|write_bytes|>>?\s*[\"\']?[^\s&|]|\btee\b|\bsed\s+-i|\bmv\b|\brm\b"
|
|
r"|\btouch\b|\bcp\b|open\([^)]*[\"\'][wa]"
|
|
)
|
|
|
|
|
|
def bash_paths(command: str, cwd: Path | None) -> tuple[list[str], Path | None, bool]:
|
|
"""Files a shell command names, where it leaves the shell, whether it writes.
|
|
|
|
Relative names resolve against ``cwd``, or against the directory of a
|
|
leading ``cd`` in the same command. Absolute names stay as they are.
|
|
"""
|
|
here = cwd
|
|
for match in _SHELL_CD.finditer(command):
|
|
target = match.group(2) or match.group(3) or match.group(4) or ""
|
|
target = target.strip()
|
|
if not target:
|
|
continue
|
|
here = Path(target) if target.startswith("/") else (here or Path()) / target
|
|
found: list[str] = []
|
|
for match in _SHELL_TOKEN.finditer(command):
|
|
raw = (match.group(1) or match.group(2) or match.group(3) or "").strip()
|
|
if not raw.lower().endswith(_NOTE_SUFFIXES) or "*" in raw:
|
|
continue
|
|
path = Path(raw) if raw.startswith("/") else (here or Path()) / raw
|
|
text = str(path)
|
|
if text not in found:
|
|
found.append(text)
|
|
return found, here, bool(_SHELL_WRITES.search(command))
|
|
|
|
|
|
def _tokens(chars: int) -> int:
|
|
return round(chars / CHARS_PER_TOKEN)
|
|
|
|
|
|
def _relative(path: str, cwd: Path | None) -> str:
|
|
if cwd is None:
|
|
return path
|
|
try:
|
|
return str(Path(path).relative_to(cwd))
|
|
except ValueError:
|
|
return path
|
|
|
|
|
|
def _tool_uses(entries: Iterable[Mapping[str, Any]]) -> Iterable[tuple[str, Any, str]]:
|
|
for entry in entries:
|
|
message = entry.get("message")
|
|
if entry.get("type") != "assistant" or not isinstance(message, dict):
|
|
continue
|
|
content = message.get("content")
|
|
if not isinstance(content, list):
|
|
continue
|
|
stamp = str(entry.get("timestamp") or "")
|
|
for block in content:
|
|
if isinstance(block, dict) and block.get("type") == "tool_use":
|
|
yield str(block.get("name") or "?"), block.get("input"), stamp
|
|
|
|
|
|
def files_touched(
|
|
entries: Iterable[Mapping[str, Any]], *, cwd: Path | None = None
|
|
) -> tuple[list[dict[str, Any]], dict[str, int]]:
|
|
"""Files named in tool inputs, with how they were touched, plus tool counts."""
|
|
files: dict[str, dict[str, Any]] = {}
|
|
counts: Counter[str] = Counter()
|
|
# The shell keeps its directory between calls; follow it.
|
|
shell_cwd = cwd
|
|
|
|
def touch(raw: str, how: str, stamp: str) -> None:
|
|
key = _relative(raw, cwd)
|
|
row = files.setdefault(
|
|
key, {"path": key, "reads": 0, "writes": 0, "other": 0, "last_at": stamp}
|
|
)
|
|
row[how] += 1
|
|
row["last_at"] = max(row["last_at"], stamp)
|
|
|
|
for name, tool_input, stamp in _tool_uses(entries):
|
|
counts[name] += 1
|
|
if not isinstance(tool_input, dict):
|
|
continue
|
|
if name == "Bash" and isinstance(tool_input.get("command"), str):
|
|
paths, shell_cwd, writes = bash_paths(tool_input["command"], shell_cwd)
|
|
for raw in paths:
|
|
touch(raw, "writes" if writes else "reads", stamp)
|
|
continue
|
|
raw = next(
|
|
(v for k in PATH_KEYS if isinstance(v := tool_input.get(k), str) and v),
|
|
None,
|
|
)
|
|
if raw is None or (name in SEARCH_TOOLS and "file_path" not in tool_input):
|
|
continue
|
|
if name in READ_TOOLS:
|
|
touch(raw, "reads", stamp)
|
|
elif name in WRITE_TOOLS:
|
|
touch(raw, "writes", stamp)
|
|
else:
|
|
touch(raw, "other", stamp)
|
|
ordered = sorted(files.values(), key=lambda r: r["last_at"], reverse=True)
|
|
return ordered, dict(counts)
|
|
|
|
|
|
def granules(agent: ClaudeAgent, kind: Kind) -> dict[str, Any]:
|
|
"""The system prompt as its sources: tag, path, size, token estimate."""
|
|
sources = agent.prompt_for(kind) if agent.serves(kind) else None
|
|
rows: list[dict[str, Any]] = []
|
|
if sources is None:
|
|
chars = len(agent.system_prompt)
|
|
if chars:
|
|
rows.append(
|
|
{
|
|
"tag": None,
|
|
"path": None,
|
|
"bytes": chars,
|
|
"tokens_est": _tokens(chars),
|
|
}
|
|
)
|
|
else:
|
|
for source in sources:
|
|
tag, raw = source if isinstance(source, tuple) else (None, source)
|
|
path = Path(raw)
|
|
try:
|
|
chars = len(path.read_text(encoding="utf-8"))
|
|
except OSError:
|
|
chars = 0
|
|
rows.append(
|
|
{
|
|
"tag": tag,
|
|
"path": _relative(str(path), agent.cwd),
|
|
"bytes": chars,
|
|
"tokens_est": _tokens(chars),
|
|
}
|
|
)
|
|
total = sum(r["bytes"] for r in rows)
|
|
return {"bytes": total, "tokens_est": _tokens(total), "granules": rows}
|
|
|
|
|
|
def _skill_meta(skill_md: Path) -> dict[str, str]:
|
|
try:
|
|
text = skill_md.read_text(encoding="utf-8")
|
|
except OSError:
|
|
return {}
|
|
match = _FRONTMATTER.match(text)
|
|
meta: dict[str, str] = {}
|
|
if match:
|
|
for line in match.group(1).splitlines():
|
|
key, sep, value = line.partition(":")
|
|
if sep:
|
|
meta[key.strip()] = value.strip()
|
|
return meta
|
|
|
|
|
|
def skill_sets(agent: ClaudeAgent, kind: Kind) -> list[dict[str, Any]]:
|
|
"""Skill-set directories the kind loads, with each skill's name and description."""
|
|
out: list[dict[str, Any]] = []
|
|
for directory in agent.skills_for(kind):
|
|
skills = []
|
|
try:
|
|
children = sorted(p for p in directory.iterdir() if p.is_dir())
|
|
except OSError:
|
|
children = []
|
|
for child in children:
|
|
skill_md = child / "SKILL.md"
|
|
if skill_md.is_file():
|
|
meta = _skill_meta(skill_md)
|
|
skills.append(
|
|
{
|
|
"name": meta.get("name") or child.name,
|
|
"description": meta.get("description", ""),
|
|
"path": _relative(str(child), agent.cwd),
|
|
}
|
|
)
|
|
out.append(
|
|
{
|
|
"set": directory.name,
|
|
"path": _relative(str(directory), agent.cwd),
|
|
"skills": skills,
|
|
}
|
|
)
|
|
return out
|
|
|
|
|
|
def compose(
|
|
agent: ClaudeAgent,
|
|
kind: Kind,
|
|
entries: Iterable[Mapping[str, Any]],
|
|
*,
|
|
context_tokens: int,
|
|
) -> dict[str, Any]:
|
|
files, counts = files_touched(entries, cwd=agent.cwd)
|
|
prompt = granules(agent, kind)
|
|
skills = skill_sets(agent, kind)
|
|
skills_chars = sum(
|
|
len(s["name"]) + len(s["description"])
|
|
for group in skills
|
|
for s in group["skills"]
|
|
)
|
|
return {
|
|
"agent": {
|
|
"name": agent.name,
|
|
"model": agent.model,
|
|
"effort": agent.options.effort,
|
|
},
|
|
"kind": kind,
|
|
"context_tokens": context_tokens,
|
|
"prompt": prompt,
|
|
"skills": skills,
|
|
"skills_tokens_est": _tokens(skills_chars),
|
|
"tools": {
|
|
"allowed": list(agent.options.tools) if agent.options.tools else None,
|
|
"disallowed": list(agent.options.disallowed_tools),
|
|
"gateway": list(agent.gateway_tools),
|
|
"mcps": [m.name for m in agent.expose_mcps],
|
|
},
|
|
"files": files,
|
|
"tool_counts": counts,
|
|
"history_tokens_est": max(
|
|
0, context_tokens - prompt["tokens_est"] - _tokens(skills_chars)
|
|
),
|
|
}
|