wehub-resource-sync
387 lines
11 KiB
Python
Executable File
387 lines
11 KiB
Python
Executable File
#!/usr/bin/env python3
|
|
|
|
from __future__ import annotations
|
|
|
|
from collections.abc import Iterable, Iterator
|
|
from dataclasses import dataclass, field
|
|
from pathlib import Path
|
|
import json
|
|
import random
|
|
import time
|
|
from typing import Any, Literal
|
|
|
|
LimitMode = Literal["calls", "events"]
|
|
|
|
DEFAULT_MAX_ITEMS = 50_000
|
|
DEFAULT_SINCE_DAYS = 30
|
|
DEFAULT_MAX_FILES = 500
|
|
DEFAULT_SESSIONS_DIR = Path.home() / ".omp" / "agent" / "sessions"
|
|
|
|
|
|
TOOL_GROUPS: dict[str, tuple[str, ...]] = {
|
|
"edits": ("edit", "ast_edit"),
|
|
"reads": ("read", "grep", "find", "ast_grep", "lsp"),
|
|
"writes": ("edit", "ast_edit", "write"),
|
|
}
|
|
|
|
|
|
@dataclass(slots=True)
|
|
class ToolIOConfig:
|
|
sessions_dir: Path = DEFAULT_SESSIONS_DIR
|
|
since_days: int = DEFAULT_SINCE_DAYS
|
|
max_files: int = DEFAULT_MAX_FILES
|
|
max_items: int = DEFAULT_MAX_ITEMS
|
|
limit_mode: LimitMode = "calls"
|
|
include_unresolved: bool = True
|
|
|
|
|
|
@dataclass(slots=True)
|
|
class ToolCall:
|
|
session_file: Path
|
|
tool_call_id: str
|
|
tool_name: str
|
|
arguments: dict[str, Any]
|
|
assistant_thinking: str | None = None
|
|
assistant_timestamp: str | None = None
|
|
path_hint: str = ""
|
|
|
|
|
|
@dataclass(slots=True)
|
|
class ToolResult:
|
|
tool_call_id: str
|
|
tool_name: str
|
|
is_error: bool
|
|
result_text: str
|
|
details: dict[str, Any] = field(default_factory=dict)
|
|
tool_timestamp: str | None = None
|
|
|
|
|
|
@dataclass(slots=True)
|
|
class ToolInvocation:
|
|
call: ToolCall
|
|
result: ToolResult | None = None
|
|
|
|
@property
|
|
def session_file(self) -> Path:
|
|
return self.call.session_file
|
|
|
|
@property
|
|
def tool_call_id(self) -> str:
|
|
return self.call.tool_call_id
|
|
|
|
@property
|
|
def tool_name(self) -> str:
|
|
return self.call.tool_name
|
|
|
|
@property
|
|
def arguments(self) -> dict[str, Any]:
|
|
return self.call.arguments
|
|
|
|
@property
|
|
def assistant_thinking(self) -> str | None:
|
|
return self.call.assistant_thinking
|
|
|
|
@property
|
|
def assistant_timestamp(self) -> str | None:
|
|
return self.call.assistant_timestamp
|
|
|
|
@property
|
|
def tool_timestamp(self) -> str | None:
|
|
return self.result.tool_timestamp if self.result else None
|
|
|
|
@property
|
|
def path_hint(self) -> str:
|
|
return self.call.path_hint
|
|
|
|
@property
|
|
def has_result(self) -> bool:
|
|
return self.result is not None
|
|
|
|
@property
|
|
def is_error(self) -> bool:
|
|
return bool(self.result and self.result.is_error)
|
|
|
|
@property
|
|
def result_text(self) -> str:
|
|
return self.result.result_text if self.result else ""
|
|
|
|
@property
|
|
def details(self) -> dict[str, Any]:
|
|
return self.result.details if self.result else {}
|
|
|
|
@property
|
|
def diff(self) -> str | None:
|
|
diff = self.details.get("diff")
|
|
return diff if isinstance(diff, str) else None
|
|
|
|
|
|
@dataclass(slots=True)
|
|
class ReservoirSample[T]:
|
|
size: int
|
|
items: list[T] = field(default_factory=list)
|
|
seen: int = 0
|
|
rng: random.Random = field(default_factory=random.Random)
|
|
|
|
def add(self, item: T) -> None:
|
|
if self.size <= 0:
|
|
return
|
|
self.seen += 1
|
|
if len(self.items) < self.size:
|
|
self.items.append(item)
|
|
return
|
|
index = self.rng.randrange(self.seen)
|
|
if index < self.size:
|
|
self.items[index] = item
|
|
|
|
|
|
|
|
def list_recent_session_files(config: ToolIOConfig) -> list[Path]:
|
|
min_mtime = time.time() - config.since_days * 24 * 60 * 60
|
|
candidates: list[tuple[float, Path]] = []
|
|
for session_file in config.sessions_dir.rglob("*.jsonl"):
|
|
try:
|
|
stat = session_file.stat()
|
|
except FileNotFoundError:
|
|
continue
|
|
if stat.st_mtime < min_mtime:
|
|
continue
|
|
candidates.append((stat.st_mtime, session_file))
|
|
candidates.sort(key=lambda entry: entry[0], reverse=True)
|
|
return [entry[1] for entry in candidates[: config.max_files]]
|
|
|
|
|
|
|
|
def iter_tool_invocations(
|
|
tool_names: str | Iterable[str],
|
|
config: ToolIOConfig | None = None,
|
|
) -> Iterator[ToolInvocation]:
|
|
resolved = config or ToolIOConfig()
|
|
wanted = _normalize_tool_names(tool_names)
|
|
seen_items = 0
|
|
|
|
for session_file in list_recent_session_files(resolved):
|
|
pending: dict[str, ToolCall] = {}
|
|
for entry in _iter_session_entries(session_file):
|
|
if entry.get("type") != "message":
|
|
continue
|
|
message = _as_record(entry.get("message"))
|
|
if message is None:
|
|
continue
|
|
|
|
role = message.get("role")
|
|
if role == "assistant":
|
|
content = message.get("content")
|
|
if not isinstance(content, list):
|
|
continue
|
|
thinking = _extract_thinking(content)
|
|
assistant_timestamp = _as_string(entry.get("timestamp"))
|
|
for item in content:
|
|
payload = _as_record(item)
|
|
if payload is None:
|
|
continue
|
|
if payload.get("type") != "toolCall":
|
|
continue
|
|
tool_name = _as_string(payload.get("name"))
|
|
tool_call_id = _as_string(payload.get("id"))
|
|
if tool_name is None or tool_call_id is None or tool_name not in wanted:
|
|
continue
|
|
arguments = _as_record(payload.get("arguments")) or {}
|
|
pending[tool_call_id] = ToolCall(
|
|
session_file=session_file,
|
|
tool_call_id=tool_call_id,
|
|
tool_name=tool_name,
|
|
arguments=arguments,
|
|
assistant_thinking=thinking,
|
|
assistant_timestamp=assistant_timestamp,
|
|
path_hint=extract_path(arguments),
|
|
)
|
|
continue
|
|
|
|
if role != "toolResult":
|
|
continue
|
|
tool_name = _as_string(message.get("toolName"))
|
|
tool_call_id = _as_string(message.get("toolCallId"))
|
|
if tool_name is None or tool_call_id is None or tool_name not in wanted:
|
|
continue
|
|
pending_call = pending.pop(tool_call_id, None)
|
|
if pending_call is None:
|
|
continue
|
|
result = ToolResult(
|
|
tool_call_id=tool_call_id,
|
|
tool_name=tool_name,
|
|
is_error=message.get("isError") is True,
|
|
result_text=extract_result_text(message),
|
|
details=_as_record(message.get("details")) or {},
|
|
tool_timestamp=_as_string(entry.get("timestamp")),
|
|
)
|
|
invocation = ToolInvocation(call=pending_call, result=result)
|
|
seen_items += _event_weight(invocation, resolved.limit_mode)
|
|
yield invocation
|
|
if seen_items >= resolved.max_items:
|
|
return
|
|
|
|
if not resolved.include_unresolved:
|
|
continue
|
|
for pending_call in pending.values():
|
|
invocation = ToolInvocation(call=pending_call)
|
|
seen_items += _event_weight(invocation, resolved.limit_mode)
|
|
yield invocation
|
|
if seen_items >= resolved.max_items:
|
|
return
|
|
|
|
|
|
|
|
def iter_results(stream: Iterable[ToolInvocation]) -> Iterator[ToolInvocation]:
|
|
for invocation in stream:
|
|
if invocation.has_result:
|
|
yield invocation
|
|
|
|
|
|
|
|
def iter_failed(stream: Iterable[ToolInvocation]) -> Iterator[ToolInvocation]:
|
|
for invocation in stream:
|
|
if invocation.is_error:
|
|
yield invocation
|
|
|
|
|
|
|
|
def iter_successful(stream: Iterable[ToolInvocation]) -> Iterator[ToolInvocation]:
|
|
for invocation in stream:
|
|
if invocation.has_result and not invocation.is_error:
|
|
yield invocation
|
|
|
|
|
|
|
|
def iter_with_diff(stream: Iterable[ToolInvocation]) -> Iterator[ToolInvocation]:
|
|
for invocation in stream:
|
|
if invocation.diff:
|
|
yield invocation
|
|
|
|
|
|
|
|
def iter_paths(stream: Iterable[ToolInvocation], *paths: str) -> Iterator[ToolInvocation]:
|
|
wanted = set(paths)
|
|
for invocation in stream:
|
|
if invocation.path_hint in wanted:
|
|
yield invocation
|
|
|
|
|
|
|
|
def take(stream: Iterable[ToolInvocation], limit: int) -> Iterator[ToolInvocation]:
|
|
if limit <= 0:
|
|
return
|
|
remaining = limit
|
|
for invocation in stream:
|
|
if remaining <= 0:
|
|
return
|
|
yield invocation
|
|
remaining -= 1
|
|
|
|
|
|
|
|
def sample_reservoir[T](stream: Iterable[T], size: int, seed: int | None = None) -> list[T]:
|
|
sample: ReservoirSample[T] = ReservoirSample(size=size, rng=random.Random(seed))
|
|
for item in stream:
|
|
sample.add(item)
|
|
return sample.items
|
|
|
|
|
|
|
|
def extract_result_text(message: dict[str, Any] | None) -> str:
|
|
if message is None:
|
|
return ""
|
|
content = message.get("content")
|
|
if not isinstance(content, list):
|
|
return ""
|
|
for item in content:
|
|
payload = _as_record(item)
|
|
if payload is None:
|
|
continue
|
|
if payload.get("type") != "text":
|
|
continue
|
|
text = _as_string(payload.get("text"))
|
|
if text is not None:
|
|
return text
|
|
return ""
|
|
|
|
|
|
|
|
def extract_path(arguments: dict[str, Any]) -> str:
|
|
for key in ("path", "file", "move"):
|
|
value = arguments.get(key)
|
|
if isinstance(value, str):
|
|
return value
|
|
return ""
|
|
|
|
|
|
|
|
def _iter_session_entries(session_file: Path) -> Iterator[dict[str, Any]]:
|
|
with session_file.open("r", encoding="utf-8") as handle:
|
|
for line in handle:
|
|
line = line.strip()
|
|
if not line:
|
|
continue
|
|
try:
|
|
entry = json.loads(line)
|
|
except json.JSONDecodeError:
|
|
continue
|
|
payload = _as_record(entry)
|
|
if payload is not None:
|
|
yield payload
|
|
|
|
|
|
|
|
def _extract_thinking(content: list[Any]) -> str | None:
|
|
for item in content:
|
|
payload = _as_record(item)
|
|
if payload is None:
|
|
continue
|
|
if payload.get("type") != "thinking":
|
|
continue
|
|
thinking = _as_string(payload.get("thinking"))
|
|
if thinking:
|
|
return thinking
|
|
return None
|
|
|
|
|
|
|
|
def resolve_tool_names(*names_or_groups: str) -> tuple[str, ...]:
|
|
ordered: list[str] = []
|
|
seen: set[str] = set()
|
|
for name in names_or_groups:
|
|
expanded = TOOL_GROUPS.get(name, (name,))
|
|
for tool_name in expanded:
|
|
if tool_name in seen:
|
|
continue
|
|
seen.add(tool_name)
|
|
ordered.append(tool_name)
|
|
return tuple(ordered)
|
|
|
|
|
|
def _normalize_tool_names(tool_names: str | Iterable[str]) -> set[str]:
|
|
if isinstance(tool_names, str):
|
|
return set(resolve_tool_names(tool_names))
|
|
ordered: list[str] = []
|
|
for name in tool_names:
|
|
ordered.extend(resolve_tool_names(name))
|
|
return set(ordered)
|
|
|
|
|
|
|
|
def _event_weight(invocation: ToolInvocation, limit_mode: LimitMode) -> int:
|
|
if limit_mode == "calls":
|
|
return 1
|
|
return 2 if invocation.has_result else 1
|
|
|
|
|
|
|
|
def _as_record(value: Any) -> dict[str, Any] | None:
|
|
if not isinstance(value, dict):
|
|
return None
|
|
return value
|
|
|
|
|
|
|
|
def _as_string(value: Any) -> str | None:
|
|
return value if isinstance(value, str) else None
|