mirror of
https://github.com/AstrBotDevs/AstrBot.git
synced 2026-10-16 06:58:51 +08:00
chore: remove image memory experiment tooling
This commit is contained in:
zenfun
@@ -1,215 +0,0 @@
|
||||
name: Image Memory Lifecycle Validation
|
||||
|
||||
on:
|
||||
push:
|
||||
branches:
|
||||
- feat/optimize-image-memory-lifecycle
|
||||
workflow_dispatch:
|
||||
|
||||
jobs:
|
||||
validate:
|
||||
name: Image and memory validation (${{ matrix.os }})
|
||||
runs-on: ${{ matrix.os }}
|
||||
timeout-minutes: 90
|
||||
strategy:
|
||||
fail-fast: false
|
||||
matrix:
|
||||
os: [ubuntu-latest, windows-latest, macos-latest]
|
||||
|
||||
defaults:
|
||||
run:
|
||||
shell: bash
|
||||
|
||||
steps:
|
||||
- name: Checkout
|
||||
uses: actions/checkout@v7
|
||||
|
||||
- name: Set up Python
|
||||
uses: actions/setup-python@v7
|
||||
with:
|
||||
python-version: '3.12'
|
||||
|
||||
- name: Install uv
|
||||
run: |
|
||||
python -m pip install --upgrade pip
|
||||
python -m pip install uv
|
||||
|
||||
- name: Install dependencies
|
||||
run: uv sync
|
||||
|
||||
- name: Run repository ruff checks
|
||||
run: |
|
||||
uv run ruff format --check .
|
||||
uv run ruff check .
|
||||
|
||||
- name: Run focused tests
|
||||
run: |
|
||||
mkdir -p "$RUNNER_TEMP/image-memory-validation/results"
|
||||
uv run pytest -q \
|
||||
tests/unit/test_image_preparation_budget.py \
|
||||
tests/unit/test_image_media_store.py \
|
||||
tests/test_media_utils.py \
|
||||
tests/test_computer_fs_tools.py \
|
||||
tests/test_tool_loop_agent_runner.py \
|
||||
tests/agent/test_context_manager.py \
|
||||
tests/unit/test_astr_main_agent.py \
|
||||
tests/unit/test_context_image_budget.py \
|
||||
tests/unit/test_image_source_preparation.py \
|
||||
tests/unit/test_image_memory_benchmark.py \
|
||||
tests/unit/test_image_lifecycle_workloads.py \
|
||||
tests/unit/test_image_screenshot_encoding.py \
|
||||
tests/unit/test_image_source_integration.py \
|
||||
tests/unit/test_provider_image_references.py \
|
||||
tests/test_agent_runner_media_resolver.py \
|
||||
tests/unit/test_provider_request_retry.py \
|
||||
tests/unit/test_manage_image_history.py \
|
||||
tests/unit/test_conversation_media_api.py \
|
||||
tests/test_openai_source.py \
|
||||
tests/test_openai_responses_source.py \
|
||||
tests/test_gemini_source.py \
|
||||
tests/test_anthropic_source.py \
|
||||
tests/test_anthropic_kimi_code_provider.py \
|
||||
tests/test_deerflow_agent_runner.py \
|
||||
tests/unit/test_cron_context_compression.py \
|
||||
tests/unit/test_group_chat_context_wiring.py \
|
||||
tests/test_platform_image_format_preservation.py \
|
||||
tests/test_webchat_queue_lifecycle.py \
|
||||
tests/test_webchat_upload_image_format.py \
|
||||
--junitxml="$RUNNER_TEMP/image-memory-validation/results/tests.xml"
|
||||
|
||||
- name: Generate reproducible fixtures outside measured processes
|
||||
run: |
|
||||
mkdir -p "$RUNNER_TEMP/image-memory-validation/fixtures"
|
||||
uv run python scripts/image_memory_bench/generate_fixtures.py \
|
||||
"$RUNNER_TEMP/image-memory-validation/fixtures" \
|
||||
--seeds 7 19 43
|
||||
|
||||
- name: Record versions and fixture manifest
|
||||
run: |
|
||||
mkdir -p "$RUNNER_TEMP/image-memory-validation/results"
|
||||
uv run python scripts/image_memory_bench/baseline_manifest.py \
|
||||
"$RUNNER_TEMP/image-memory-validation/results/manifest.json" \
|
||||
--fixtures "$RUNNER_TEMP/image-memory-validation/fixtures"
|
||||
|
||||
- name: Run single-image memory workloads for every paired seed
|
||||
run: |
|
||||
set +e
|
||||
result="$RUNNER_TEMP/image-memory-validation/results/measurements.jsonl"
|
||||
failures=0
|
||||
for seed in 7 19 43; do
|
||||
while IFS= read -r image; do
|
||||
uv run python scripts/image_memory_bench/measure.py \
|
||||
"$image" "$result" \
|
||||
--repo "$GITHUB_WORKSPACE" \
|
||||
--csv "$RUNNER_TEMP/image-memory-validation/results/measurements.csv" \
|
||||
--workload "$seed-$(basename "$image")"
|
||||
code=$?
|
||||
if [ "$code" -ne 0 ]; then
|
||||
failures=$((failures + 1))
|
||||
fi
|
||||
done < <(find "$RUNNER_TEMP/image-memory-validation/fixtures/$seed" \
|
||||
-type f \( -name '*ordinary*' -o -name '*stress*' \) | sort)
|
||||
done
|
||||
printf '{"measurement_failures":%s}\n' "$failures" \
|
||||
> "$RUNNER_TEMP/image-memory-validation/results/measurement-summary.json"
|
||||
|
||||
- name: Run Python allocation diagnostic workload
|
||||
run: |
|
||||
uv run python scripts/image_memory_bench/measure.py \
|
||||
"$RUNNER_TEMP/image-memory-validation/fixtures/7/png-stress.png" \
|
||||
"$RUNNER_TEMP/image-memory-validation/results/trace-python.jsonl" \
|
||||
--repo "$GITHUB_WORKSPACE" \
|
||||
--trace-python \
|
||||
--csv "$RUNNER_TEMP/image-memory-validation/results/trace-python.csv" \
|
||||
--workload trace-python-png-stress
|
||||
|
||||
- name: Run paired history loading workloads
|
||||
run: |
|
||||
mkdir -p "$RUNNER_TEMP/image-memory-validation/results/history"
|
||||
for seed in 7 19 43; do
|
||||
uv run python scripts/image_memory_bench/history_benchmark.py \
|
||||
"$RUNNER_TEMP/image-memory-validation/fixtures/$seed/manifest.json" \
|
||||
"$RUNNER_TEMP/image-memory-validation/results/history/$seed.jsonl" \
|
||||
--repo "$GITHUB_WORKSPACE" \
|
||||
--images 50 \
|
||||
--keep-turns 25 \
|
||||
--repeats 10 \
|
||||
--cold-repeats 10 \
|
||||
--warm-repeats 10 \
|
||||
--distinct \
|
||||
--csv "$RUNNER_TEMP/image-memory-validation/results/history/$seed.csv" \
|
||||
--rss-jsonl "$RUNNER_TEMP/image-memory-validation/results/history/$seed.rss.jsonl"
|
||||
done
|
||||
|
||||
- name: Run fixed-window concurrent lifecycle and restart workload
|
||||
run: |
|
||||
lifecycle_root="$RUNNER_TEMP/image-memory-validation/lifecycle"
|
||||
mkdir -p "$lifecycle_root"
|
||||
uv run python scripts/image_memory_bench/prepare_lifecycle_fixture.py \
|
||||
"$RUNNER_TEMP/image-memory-validation/fixtures/7/manifest.json" \
|
||||
"$lifecycle_root/input" \
|
||||
--session-count 4 \
|
||||
--history-turns 50
|
||||
lifecycle_manifest="$lifecycle_root/input/lifecycle-manifest.json"
|
||||
database=$(uv run python -c \
|
||||
'import json,sys; print(json.load(open(sys.argv[1]))["database"])' \
|
||||
"$lifecycle_manifest")
|
||||
media_root=$(uv run python -c \
|
||||
'import json,sys; print(json.load(open(sys.argv[1]))["media_root"])' \
|
||||
"$lifecycle_manifest")
|
||||
conversation_ids=$(uv run python -c \
|
||||
'import json,sys; print(" ".join(json.load(open(sys.argv[1]))["conversation_ids"]))' \
|
||||
"$lifecycle_manifest")
|
||||
command=(
|
||||
uv run python scripts/image_memory_bench/lifecycle_workloads.py
|
||||
"$lifecycle_root/result.json"
|
||||
--repo "$GITHUB_WORKSPACE"
|
||||
--database "$database"
|
||||
--media-root "$media_root"
|
||||
--session-count 4
|
||||
--request-count 200
|
||||
--window-turns 25
|
||||
--rss-jsonl "$lifecycle_root/result.rss.jsonl"
|
||||
)
|
||||
for conversation_id in $conversation_ids; do
|
||||
command+=(--conversation-id "$conversation_id")
|
||||
done
|
||||
"${command[@]}"
|
||||
|
||||
- name: Record ablation contract status without fabricating factors
|
||||
run: |
|
||||
uv run python scripts/image_memory_bench/ablation_runner.py \
|
||||
"$RUNNER_TEMP/image-memory-validation/results/ablation-plan.json" \
|
||||
--baseline-repo "$GITHUB_WORKSPACE" \
|
||||
--candidate-repo "$GITHUB_WORKSPACE"
|
||||
|
||||
- name: Upload raw validation artifacts
|
||||
if: always()
|
||||
uses: actions/upload-artifact@v4
|
||||
with:
|
||||
name: image-memory-validation-${{ matrix.os }}
|
||||
path: ${{ runner.temp }}/image-memory-validation
|
||||
if-no-files-found: error
|
||||
|
||||
dashboard:
|
||||
name: Dashboard typecheck
|
||||
runs-on: ubuntu-latest
|
||||
timeout-minutes: 20
|
||||
defaults:
|
||||
run:
|
||||
working-directory: dashboard
|
||||
steps:
|
||||
- name: Checkout
|
||||
uses: actions/checkout@v7
|
||||
- name: Set up Node
|
||||
uses: actions/setup-node@v7
|
||||
with:
|
||||
node-version: 22
|
||||
cache: pnpm
|
||||
cache-dependency-path: dashboard/pnpm-lock.yaml
|
||||
- name: Enable pnpm
|
||||
run: corepack enable
|
||||
- name: Install dashboard dependencies
|
||||
run: pnpm install --frozen-lockfile
|
||||
- name: Run dashboard typecheck
|
||||
run: pnpm typecheck
|
||||
@@ -1,66 +0,0 @@
|
||||
"""Run the real 2^3 image experiment only when production factor flips exist.
|
||||
|
||||
The current checkout does not expose independent A, B, and C experiment
|
||||
switches. This runner therefore refuses to fabricate a matrix and emits an
|
||||
auditable blocked manifest until each factor is supplied by a real checkout.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import itertools
|
||||
import json
|
||||
import subprocess
|
||||
from pathlib import Path
|
||||
|
||||
|
||||
def git_identity(repo: Path) -> dict[str, str]:
|
||||
"""Return commit and dirty diff identity for a checkout."""
|
||||
commit = subprocess.check_output(
|
||||
["git", "rev-parse", "HEAD"], cwd=repo, text=True
|
||||
).strip()
|
||||
diff = subprocess.check_output(["git", "diff", "--binary"], cwd=repo)
|
||||
import hashlib
|
||||
|
||||
return {
|
||||
"commit": commit,
|
||||
"working_tree_diff_sha256": hashlib.sha256(diff).hexdigest(),
|
||||
}
|
||||
|
||||
|
||||
def main() -> int:
|
||||
"""Validate factor provenance and write a blocked or runnable plan."""
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument("output", type=Path)
|
||||
parser.add_argument("--baseline-repo", type=Path, required=True)
|
||||
parser.add_argument("--candidate-repo", type=Path, required=True)
|
||||
parser.add_argument("--factor-entrypoint", type=Path)
|
||||
args = parser.parse_args()
|
||||
factors = {
|
||||
"A": "unavailable: no independent production source-preparation entrypoint supplied",
|
||||
"B": "unavailable: no independent production storage/lazy-materialization entrypoint supplied",
|
||||
"C": "unavailable: no independent production compression-algorithm entrypoint supplied",
|
||||
}
|
||||
entrypoint_ok = (
|
||||
args.factor_entrypoint is not None and args.factor_entrypoint.is_file()
|
||||
)
|
||||
plan = {
|
||||
"status": "blocked" if not entrypoint_ok else "requires_factor_contract",
|
||||
"experiment": "real-2^3-factorial",
|
||||
"factors": factors,
|
||||
"combinations": ["".join(bits) for bits in itertools.product("0A", repeat=0)],
|
||||
"baseline": git_identity(args.baseline_repo),
|
||||
"candidate": git_identity(args.candidate_repo),
|
||||
"warmup_scope": "sdk_send_only; not warm history",
|
||||
"reason": "Do not execute or label combinations until A/B/C independently select real production paths.",
|
||||
}
|
||||
plan["combinations"] = [
|
||||
"".join(bits) for bits in itertools.product("0A", "0B", "0C")
|
||||
]
|
||||
args.output.write_text(json.dumps(plan, indent=2) + "\n", encoding="utf-8")
|
||||
print(json.dumps({"status": plan["status"], "output": str(args.output)}))
|
||||
return 0 if plan["status"] == "blocked" else 2
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
raise SystemExit(main())
|
||||
@@ -1,54 +0,0 @@
|
||||
"""Record the exact code and dependency baseline for image experiments."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import hashlib
|
||||
import importlib.metadata
|
||||
import json
|
||||
import platform
|
||||
import subprocess
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
|
||||
def main() -> int:
|
||||
"""Write a reproducibility manifest without reading application data."""
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument("output", type=Path)
|
||||
parser.add_argument("--fixtures", type=Path)
|
||||
args = parser.parse_args()
|
||||
commit = subprocess.check_output(["git", "rev-parse", "HEAD"], text=True).strip()
|
||||
packages = sorted(
|
||||
f"{dist.metadata['Name']}=={dist.version}"
|
||||
for dist in importlib.metadata.distributions()
|
||||
if dist.metadata.get("Name")
|
||||
)
|
||||
manifest = {
|
||||
"commit": commit,
|
||||
"working_tree_diff_sha256": hashlib.sha256(
|
||||
subprocess.check_output(["git", "diff", "--binary"], cwd=Path.cwd())
|
||||
).hexdigest(),
|
||||
"python": sys.version,
|
||||
"platform": platform.platform(),
|
||||
"packages": packages,
|
||||
}
|
||||
if args.fixtures:
|
||||
files = []
|
||||
for path in sorted(args.fixtures.rglob("*")):
|
||||
if path.is_file() and path.name != "manifest.json":
|
||||
files.append(
|
||||
{
|
||||
"path": str(path.relative_to(args.fixtures)),
|
||||
"bytes": path.stat().st_size,
|
||||
"sha256": hashlib.sha256(path.read_bytes()).hexdigest(),
|
||||
}
|
||||
)
|
||||
manifest["fixtures"] = files
|
||||
args.output.write_text(json.dumps(manifest, indent=2) + "\n", encoding="utf-8")
|
||||
print(json.dumps({"commit": commit, "package_count": len(packages)}))
|
||||
return 0
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
raise SystemExit(main())
|
||||
@@ -1,127 +0,0 @@
|
||||
"""Generate seeded fixtures outside the measured application process."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import hashlib
|
||||
import json
|
||||
import random
|
||||
from pathlib import Path
|
||||
|
||||
from PIL import Image, ImageDraw
|
||||
|
||||
|
||||
def generate(output: Path, seed: int) -> list[dict]:
|
||||
"""Create ten inputs for one paired experiment seed.
|
||||
|
||||
Args:
|
||||
output: New directory for this seed's fixtures.
|
||||
seed: Stable seed for pixel generation.
|
||||
|
||||
Returns:
|
||||
On-disk hashes and decoded image metadata.
|
||||
"""
|
||||
output.mkdir(parents=True, exist_ok=False)
|
||||
rng = random.Random(seed)
|
||||
records = []
|
||||
for kind in ("png", "jpeg", "webp", "gif", "screenshot"):
|
||||
for stress in (False, True):
|
||||
name = f"{kind}-{'stress' if stress else 'ordinary'}"
|
||||
options = {}
|
||||
if kind in ("png", "webp"):
|
||||
size = (1280, 1280) if stress else (320, 240)
|
||||
image = (
|
||||
Image.frombytes("RGBA", size, rng.randbytes(size[0] * size[1] * 4))
|
||||
if stress
|
||||
else Image.new("RGBA", size, (40, 120, 200, 127))
|
||||
)
|
||||
image_format = kind.upper()
|
||||
options = {"lossless": stress} if kind == "webp" else {}
|
||||
elif kind == "jpeg":
|
||||
size = (4000, 3000) if stress else (640, 480)
|
||||
image = Image.frombytes(
|
||||
"RGB", size, rng.randbytes(size[0] * size[1] * 3)
|
||||
)
|
||||
image_format = "JPEG"
|
||||
options = {"quality": 95 if stress else 75}
|
||||
if stress:
|
||||
exif = image.getexif()
|
||||
exif[274] = 6
|
||||
options["exif"] = exif
|
||||
elif kind == "gif":
|
||||
image = Image.new("RGB", (320, 240), (seed % 255, 0, 255))
|
||||
image_format = "GIF"
|
||||
frames = []
|
||||
for frame in range(24 if stress else 3):
|
||||
extra = Image.new("RGB", image.size, (frame * 9, 80, 20))
|
||||
ImageDraw.Draw(extra).rectangle(
|
||||
(frame * 5, 10, frame * 5 + 30, 50), fill="white"
|
||||
)
|
||||
frames.append(extra)
|
||||
options = {
|
||||
"save_all": True,
|
||||
"append_images": frames,
|
||||
"duration": 80,
|
||||
"loop": 0,
|
||||
}
|
||||
else:
|
||||
image = Image.new(
|
||||
"RGB", (3840, 2160) if stress else (1920, 1080), "#202124"
|
||||
)
|
||||
image_format = "PNG"
|
||||
draw = ImageDraw.Draw(image)
|
||||
for row in range(image.height // 20):
|
||||
draw.text(
|
||||
(100, row * 20),
|
||||
f"coordinate (100, {row * 20}) seed={seed} command --verbose "
|
||||
* 4,
|
||||
fill="#e8eaed",
|
||||
)
|
||||
draw.rectangle((20, 20, 70, 70), outline="red", width=3)
|
||||
path = output / f"{name}.{image_format.lower()}"
|
||||
try:
|
||||
image.save(path, image_format, **options)
|
||||
finally:
|
||||
image.close()
|
||||
for frame_image in options.get("append_images", []):
|
||||
frame_image.close()
|
||||
# Reopen the saved bytes: in-memory Image.format and frame counts
|
||||
# do not describe what the encoder actually persisted.
|
||||
with Image.open(path) as saved:
|
||||
record = {
|
||||
"path": path.name,
|
||||
"kind": kind,
|
||||
"stress": stress,
|
||||
"seed": seed,
|
||||
"format": saved.format,
|
||||
"width": saved.width,
|
||||
"height": saved.height,
|
||||
"mode": saved.mode,
|
||||
"frames": getattr(saved, "n_frames", 1),
|
||||
"orientation": saved.getexif().get(274, 1),
|
||||
"has_alpha": saved.mode in ("RGBA", "LA")
|
||||
or "transparency" in saved.info,
|
||||
"bytes": path.stat().st_size,
|
||||
"sha256": hashlib.sha256(path.read_bytes()).hexdigest(),
|
||||
}
|
||||
records.append(record)
|
||||
(output / "manifest.json").write_text(
|
||||
json.dumps(records, indent=2) + "\n", encoding="utf-8"
|
||||
)
|
||||
return records
|
||||
|
||||
|
||||
def main() -> int:
|
||||
"""Generate three paired seeds unless explicitly overridden."""
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument("output", type=Path)
|
||||
parser.add_argument("--seeds", nargs="+", type=int, default=[7, 19, 43])
|
||||
args = parser.parse_args()
|
||||
for seed in args.seeds:
|
||||
records = generate(args.output / str(seed), seed)
|
||||
print(json.dumps({"seed": seed, "count": len(records)}))
|
||||
return 0
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
raise SystemExit(main())
|
||||
@@ -1,461 +0,0 @@
|
||||
"""Measure SQLite history loading and real SDK request bodies for B-only."""
|
||||
|
||||
# The compact benchmark runner keeps subprocess orchestration visibly linear.
|
||||
# ruff: noqa: E701, E702
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import asyncio
|
||||
import base64
|
||||
import csv
|
||||
import hashlib
|
||||
import json
|
||||
import mimetypes
|
||||
import os
|
||||
import subprocess
|
||||
import sys
|
||||
import tempfile
|
||||
import threading
|
||||
import time
|
||||
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
|
||||
from pathlib import Path
|
||||
|
||||
import psutil
|
||||
|
||||
|
||||
async def prepare_database(
|
||||
db_path: Path,
|
||||
media_root: Path,
|
||||
fixture: Path,
|
||||
images: int,
|
||||
distinct: bool,
|
||||
mode: str,
|
||||
) -> None:
|
||||
"""Create one SQLite input outside the measured child."""
|
||||
sys.path.insert(0, str(Path(__file__).resolve().parents[2]))
|
||||
from astrbot.core.db.sqlite import SQLiteDatabase
|
||||
from astrbot.core.utils.image_media_store import (
|
||||
ImageMediaStore,
|
||||
persist_inline_image_refs,
|
||||
)
|
||||
|
||||
rows = json.loads(fixture.read_text())
|
||||
paths = [fixture.parent / row["path"] for row in rows]
|
||||
history = []
|
||||
raw = []
|
||||
for index in range(images):
|
||||
path = paths[index % len(paths)] if distinct else paths[0]
|
||||
# Distinct workloads use separately generated, valid image files listed
|
||||
# by the fixture manifest. Never mutate encoded bytes into invalid images.
|
||||
data = path.read_bytes()
|
||||
mime = mimetypes.guess_type(path.name)[0] or "image/png"
|
||||
raw.append(
|
||||
{
|
||||
"type": "image_url",
|
||||
"image_url": {
|
||||
"url": f"data:{mime};base64,{base64.b64encode(data).decode()}",
|
||||
"detail": "high",
|
||||
},
|
||||
}
|
||||
)
|
||||
for index, part in enumerate(raw):
|
||||
history.extend(
|
||||
[
|
||||
{"role": "user", "content": [part]},
|
||||
{"role": "assistant", "content": f"ack {index}"},
|
||||
]
|
||||
)
|
||||
history.extend(
|
||||
[
|
||||
{"role": "user", "content": "final text"},
|
||||
{"role": "assistant", "content": "final response"},
|
||||
]
|
||||
)
|
||||
if mode == "reference":
|
||||
history = persist_inline_image_refs(history, ImageMediaStore(media_root))
|
||||
db = SQLiteDatabase(str(db_path))
|
||||
await db.initialize()
|
||||
record = await db.create_conversation(
|
||||
user_id="bench", platform_id="bench", content=history, title="benchmark"
|
||||
)
|
||||
(db_path.parent / "conversation-id").write_text(record.conversation_id)
|
||||
await db.engine.dispose()
|
||||
|
||||
|
||||
def count_images(value: object) -> int:
|
||||
"""Count data image URLs in a decoded SDK body."""
|
||||
if isinstance(value, dict):
|
||||
return sum(count_images(item) for item in value.values())
|
||||
if isinstance(value, list):
|
||||
return sum(count_images(item) for item in value)
|
||||
return int(isinstance(value, str) and value.startswith("data:image/"))
|
||||
|
||||
|
||||
async def run_child(args: argparse.Namespace) -> dict:
|
||||
"""Read through SQLite and ConversationManager before selecting context."""
|
||||
sys.path.insert(0, str(args.repo.resolve()))
|
||||
from openai import AsyncOpenAI
|
||||
|
||||
from astrbot.core.agent.context.config import ContextConfig
|
||||
from astrbot.core.agent.context.manager import ContextManager
|
||||
from astrbot.core.agent.message import Message
|
||||
from astrbot.core.conversation_mgr import ConversationManager
|
||||
from astrbot.core.db.sqlite import SQLiteDatabase
|
||||
from astrbot.core.utils.image_media_store import (
|
||||
ImageMediaStore,
|
||||
materialize_image_media_refs,
|
||||
)
|
||||
|
||||
phases = {}
|
||||
|
||||
def mark(phase: str) -> None:
|
||||
with args.marker.open("a", encoding="utf-8") as marker:
|
||||
marker.write(
|
||||
json.dumps({"phase": phase, "monotonic": time.monotonic()}) + "\n"
|
||||
)
|
||||
|
||||
mark("db_and_convert")
|
||||
started = time.perf_counter()
|
||||
db = SQLiteDatabase(str(args.database))
|
||||
await db.initialize()
|
||||
conversation = await ConversationManager(db).get_conversation(
|
||||
"bench", args.conversation_id
|
||||
)
|
||||
phases["db_and_convert_ms"] = (time.perf_counter() - started) * 1000
|
||||
mark("json_load_bind")
|
||||
started = time.perf_counter()
|
||||
history = json.loads(conversation.history)
|
||||
messages = [Message.model_validate(item) for item in history]
|
||||
phases["json_load_bind_ms"] = (time.perf_counter() - started) * 1000
|
||||
mark("selection")
|
||||
started = time.perf_counter()
|
||||
selected = await ContextManager(
|
||||
ContextConfig(enforce_max_turns=args.keep_turns)
|
||||
).process(messages)
|
||||
phases["selection_ms"] = (time.perf_counter() - started) * 1000
|
||||
mark("materialize")
|
||||
selected = await materialize_image_media_refs(
|
||||
selected, ImageMediaStore(args.media_root)
|
||||
)
|
||||
payload = [message.model_dump() for message in selected]
|
||||
mark("sdk")
|
||||
async with AsyncOpenAI(
|
||||
api_key="local", base_url=args.endpoint, max_retries=0
|
||||
) as client:
|
||||
for _ in range(args.warmup):
|
||||
await client.chat.completions.create(model="local", messages=payload)
|
||||
await client.chat.completions.create(model="local", messages=payload)
|
||||
await db.engine.dispose()
|
||||
return {
|
||||
"phases_ms": phases,
|
||||
"selected_messages": len(payload),
|
||||
"selected_images": count_images(payload),
|
||||
}
|
||||
|
||||
|
||||
def main() -> int:
|
||||
"""Prepare paired databases and monitor fresh children."""
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument("fixture", type=Path)
|
||||
parser.add_argument("output", type=Path)
|
||||
parser.add_argument(
|
||||
"--repo", type=Path, default=Path(__file__).resolve().parents[2]
|
||||
)
|
||||
parser.add_argument("--images", type=int, default=50)
|
||||
parser.add_argument("--keep-turns", type=int, default=25)
|
||||
parser.add_argument("--repeats", type=int, default=10)
|
||||
parser.add_argument("--cold-repeats", type=int)
|
||||
parser.add_argument("--warm-repeats", type=int)
|
||||
parser.add_argument("--warmup", type=int, default=3)
|
||||
parser.add_argument("--distinct", action="store_true")
|
||||
parser.add_argument("--csv", type=Path)
|
||||
parser.add_argument("--rss-jsonl", type=Path)
|
||||
parser.add_argument("--child", action="store_true")
|
||||
parser.add_argument("--database", type=Path)
|
||||
parser.add_argument("--media-root", type=Path)
|
||||
parser.add_argument("--history-mode", choices=("inline", "reference"))
|
||||
parser.add_argument("--conversation-id", default="")
|
||||
parser.add_argument("--endpoint")
|
||||
parser.add_argument("--body-file", type=Path)
|
||||
parser.add_argument("--marker", type=Path)
|
||||
args = parser.parse_args()
|
||||
cold_repeats = args.cold_repeats if args.cold_repeats is not None else args.repeats
|
||||
warm_repeats = args.warm_repeats if args.warm_repeats is not None else args.repeats
|
||||
if min(args.repeats, cold_repeats, warm_repeats, args.images, args.keep_turns) < 1:
|
||||
parser.error("repetition, image, and context-window values must be positive")
|
||||
fixture_rows = json.loads(args.fixture.read_text())
|
||||
max_fixture_bytes = max(
|
||||
(
|
||||
row.get("bytes") or row.get("prepared_bytes") or row.get("source_bytes", 0)
|
||||
for row in fixture_rows
|
||||
),
|
||||
default=0,
|
||||
)
|
||||
budget_class = (
|
||||
"over-4MiB-source-stress"
|
||||
if 4 * ((max_fixture_bytes + 2) // 3) > 4 * 1024 * 1024
|
||||
else "representative-within-4MiB-source"
|
||||
)
|
||||
if args.child:
|
||||
if (
|
||||
args.database is None
|
||||
or args.media_root is None
|
||||
or args.endpoint is None
|
||||
or args.body_file is None
|
||||
or args.marker is None
|
||||
or args.history_mode is None
|
||||
):
|
||||
parser.error(
|
||||
"child requires database, media root, endpoint, body file, marker, and history mode"
|
||||
)
|
||||
result = asyncio.run(run_child(args))
|
||||
if args.body_file.exists():
|
||||
metadata = json.loads(args.body_file.read_text())
|
||||
result.update(
|
||||
{
|
||||
"http_json_bytes": metadata["bytes"],
|
||||
"http_body_sha256": metadata["sha256"],
|
||||
"http_images": metadata["images"],
|
||||
}
|
||||
)
|
||||
else:
|
||||
result.update(
|
||||
{
|
||||
"http_json_bytes": None,
|
||||
"http_body_sha256": None,
|
||||
"http_images": 0,
|
||||
}
|
||||
)
|
||||
if not result["selected_images"] or not result["http_images"]:
|
||||
raise RuntimeError("actual SDK body contains no images")
|
||||
args.output.write_text(json.dumps(result) + "\n")
|
||||
return 0
|
||||
args.output.parent.mkdir(parents=True, exist_ok=True)
|
||||
args.output.write_text("", encoding="utf-8")
|
||||
csv_file = args.csv.open("w", newline="", encoding="utf-8") if args.csv else None
|
||||
writer = None
|
||||
received_body = {"path": None}
|
||||
rss_path = args.output.with_name(args.output.stem + ".rss.jsonl")
|
||||
if args.rss_jsonl:
|
||||
rss_path = args.rss_jsonl
|
||||
rss_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
rss_file = rss_path.open("w", encoding="utf-8")
|
||||
|
||||
class Handler(BaseHTTPRequestHandler):
|
||||
def do_POST(self):
|
||||
data = self.rfile.read(int(self.headers.get("Content-Length", "0")))
|
||||
body = json.loads(data)
|
||||
Path(received_body["path"]).write_text(
|
||||
json.dumps(
|
||||
{
|
||||
"bytes": len(data),
|
||||
"sha256": hashlib.sha256(data).hexdigest(),
|
||||
"images": count_images(body),
|
||||
}
|
||||
)
|
||||
)
|
||||
response = b'{"choices":[{"message":{"content":"ok"}}]}'
|
||||
self.send_response(200)
|
||||
self.send_header("Content-Length", str(len(response)))
|
||||
self.end_headers()
|
||||
self.wfile.write(response)
|
||||
|
||||
def log_message(self, *_args):
|
||||
pass
|
||||
|
||||
server = ThreadingHTTPServer(("127.0.0.1", 0), Handler)
|
||||
threading.Thread(target=server.serve_forever, daemon=True).start()
|
||||
endpoint = f"http://127.0.0.1:{server.server_port}/v1"
|
||||
try:
|
||||
for mode in ("inline", "reference"):
|
||||
for repeat in range(cold_repeats + warm_repeats):
|
||||
with tempfile.TemporaryDirectory(
|
||||
prefix="astrbot-history-input-"
|
||||
) as temp:
|
||||
root = Path(temp)
|
||||
db_path = root / f"{mode}.db"
|
||||
media = root / mode
|
||||
asyncio.run(
|
||||
prepare_database(
|
||||
db_path,
|
||||
media,
|
||||
args.fixture,
|
||||
args.images,
|
||||
args.distinct,
|
||||
mode,
|
||||
)
|
||||
)
|
||||
conversation_id = (db_path.parent / "conversation-id").read_text()
|
||||
child_output = root / "child.jsonl"
|
||||
body_file = root / "http-metadata.json"
|
||||
received_body["path"] = body_file
|
||||
warmup = 0 if repeat < cold_repeats else args.warmup
|
||||
command = [
|
||||
sys.executable,
|
||||
str(Path(__file__).resolve()),
|
||||
str(args.fixture),
|
||||
str(child_output),
|
||||
"--child",
|
||||
"--repo",
|
||||
str(args.repo),
|
||||
"--database",
|
||||
str(db_path),
|
||||
"--media-root",
|
||||
str(media),
|
||||
"--conversation-id",
|
||||
conversation_id,
|
||||
"--keep-turns",
|
||||
str(args.keep_turns),
|
||||
"--history-mode",
|
||||
mode,
|
||||
"--warmup",
|
||||
str(warmup),
|
||||
"--endpoint",
|
||||
endpoint,
|
||||
"--body-file",
|
||||
str(body_file),
|
||||
"--marker",
|
||||
str(root / "marker.jsonl"),
|
||||
]
|
||||
proc = subprocess.Popen(
|
||||
command,
|
||||
cwd=args.repo,
|
||||
env={**os.environ, "ASTRBOT_ROOT": str(root)},
|
||||
)
|
||||
monitor = psutil.Process(proc.pid)
|
||||
samples = []
|
||||
started = time.monotonic()
|
||||
limit = min(2 * 1024**3, psutil.virtual_memory().available // 4)
|
||||
stop_reason = None
|
||||
while proc.poll() is None:
|
||||
try:
|
||||
now = time.monotonic()
|
||||
rss = monitor.memory_info().rss
|
||||
private_bytes = None
|
||||
try:
|
||||
private_bytes = monitor.memory_full_info().uss
|
||||
except (AttributeError, psutil.AccessDenied):
|
||||
pass
|
||||
samples.append(
|
||||
{
|
||||
"seconds": now - started,
|
||||
"monotonic": now,
|
||||
"rss": rss,
|
||||
"private_bytes": private_bytes,
|
||||
}
|
||||
)
|
||||
if now - started > 120:
|
||||
stop_reason = "timeout"
|
||||
elif rss > limit:
|
||||
stop_reason = "memory_limit"
|
||||
if stop_reason:
|
||||
proc.kill()
|
||||
break
|
||||
except psutil.NoSuchProcess:
|
||||
break
|
||||
time.sleep(0.01)
|
||||
proc.wait()
|
||||
if child_output.exists():
|
||||
result = json.loads(child_output.read_text())
|
||||
else:
|
||||
result = {
|
||||
"status": "error",
|
||||
"error": stop_reason or "child_failed",
|
||||
}
|
||||
marker_file = root / "marker.jsonl"
|
||||
markers = (
|
||||
[
|
||||
json.loads(line)
|
||||
for line in marker_file.read_text().splitlines()
|
||||
]
|
||||
if marker_file.exists()
|
||||
else []
|
||||
)
|
||||
for sample in samples:
|
||||
eligible = [
|
||||
marker
|
||||
for marker in markers
|
||||
if marker["monotonic"] <= sample["monotonic"]
|
||||
]
|
||||
sample["phase"] = (
|
||||
max(eligible, key=lambda marker: marker["monotonic"])[
|
||||
"phase"
|
||||
]
|
||||
if eligible
|
||||
else "pre-start"
|
||||
)
|
||||
result.update(
|
||||
{
|
||||
"mode": mode,
|
||||
"repeat": repeat,
|
||||
"cold": repeat < cold_repeats,
|
||||
"warmup_mode": "sdk_warmup" if warmup else "cold",
|
||||
"warmup_scope": "sdk_send_only",
|
||||
"workload": f"sqlite-history-{args.images}-{'distinct' if args.distinct else 'same'}",
|
||||
"budget_class": budget_class,
|
||||
"rss_high_water_bytes": max(
|
||||
(sample["rss"] for sample in samples), default=None
|
||||
),
|
||||
"private_high_water_bytes": max(
|
||||
(
|
||||
sample["private_bytes"]
|
||||
for sample in samples
|
||||
if sample["private_bytes"] is not None
|
||||
),
|
||||
default=None,
|
||||
),
|
||||
"stop_reason": stop_reason,
|
||||
"rss_sample_count": len(samples),
|
||||
"rss_jsonl": str(rss_path),
|
||||
"returncode": proc.returncode,
|
||||
"status": (
|
||||
"ok"
|
||||
if proc.returncode == 0 and stop_reason is None
|
||||
else "error"
|
||||
),
|
||||
}
|
||||
)
|
||||
for sample in samples:
|
||||
rss_file.write(
|
||||
json.dumps(
|
||||
{
|
||||
"mode": mode,
|
||||
"repeat": repeat,
|
||||
**sample,
|
||||
}
|
||||
)
|
||||
+ "\n"
|
||||
)
|
||||
rss_file.flush()
|
||||
with args.output.open("a", encoding="utf-8") as stream:
|
||||
stream.write(json.dumps(result) + "\n")
|
||||
if csv_file:
|
||||
fields = [
|
||||
"mode",
|
||||
"repeat",
|
||||
"cold",
|
||||
"workload",
|
||||
"rss_high_water_bytes",
|
||||
"private_high_water_bytes",
|
||||
"rss_sample_count",
|
||||
"status",
|
||||
"returncode",
|
||||
"http_json_bytes",
|
||||
"http_body_sha256",
|
||||
"http_images",
|
||||
]
|
||||
if writer is None:
|
||||
writer = csv.DictWriter(csv_file, fieldnames=fields)
|
||||
writer.writeheader()
|
||||
writer.writerow({field: result.get(field) for field in fields})
|
||||
finally:
|
||||
server.shutdown()
|
||||
rss_file.close()
|
||||
if csv_file:
|
||||
csv_file.close()
|
||||
return 0
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
raise SystemExit(main())
|
||||
@@ -1,371 +0,0 @@
|
||||
"""Run real history lifecycle workloads in an isolated, measured child process."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import asyncio
|
||||
import base64
|
||||
import hashlib
|
||||
import json
|
||||
import subprocess
|
||||
import sys
|
||||
import tempfile
|
||||
import threading
|
||||
import time
|
||||
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
import psutil
|
||||
|
||||
REQUESTS = 200
|
||||
SESSIONS = 4
|
||||
WINDOW_TURNS = 25
|
||||
|
||||
|
||||
def _count_images(value: Any) -> int:
|
||||
if isinstance(value, dict):
|
||||
return sum(_count_images(item) for item in value.values())
|
||||
if isinstance(value, list):
|
||||
return sum(_count_images(item) for item in value)
|
||||
return int(isinstance(value, str) and value.startswith("data:image/"))
|
||||
|
||||
|
||||
def _image_hashes(value: Any) -> list[str]:
|
||||
if isinstance(value, dict):
|
||||
if value.get("type") == "image_url" and isinstance(
|
||||
value.get("image_url"), dict
|
||||
):
|
||||
encoded = value["image_url"].get("url", "").split(",", 1)[-1]
|
||||
try:
|
||||
return [hashlib.sha256(base64.b64decode(encoded)).hexdigest()]
|
||||
except Exception:
|
||||
return []
|
||||
return [item for child in value.values() for item in _image_hashes(child)]
|
||||
if isinstance(value, list):
|
||||
return [item for child in value for item in _image_hashes(child)]
|
||||
return []
|
||||
|
||||
|
||||
async def _child(args: argparse.Namespace) -> dict[str, Any]:
|
||||
"""Execute the production database, context, materialization, and SDK path."""
|
||||
sys.path.insert(0, str(args.repo.resolve()))
|
||||
from openai import AsyncOpenAI
|
||||
|
||||
from astrbot.core.agent.context.config import ContextConfig
|
||||
from astrbot.core.agent.context.manager import ContextManager
|
||||
from astrbot.core.agent.message import Message
|
||||
from astrbot.core.conversation_mgr import ConversationManager
|
||||
from astrbot.core.db.sqlite import SQLiteDatabase
|
||||
from astrbot.core.utils.image_media_store import (
|
||||
ImageMediaStore,
|
||||
materialize_image_media_refs,
|
||||
)
|
||||
|
||||
async def load_payload(
|
||||
database: SQLiteDatabase, conversation_id: str
|
||||
) -> tuple[list, str, list]:
|
||||
conversation = await ConversationManager(database).get_conversation(
|
||||
"bench", conversation_id
|
||||
)
|
||||
messages = [
|
||||
Message.model_validate(item) for item in json.loads(conversation.history)
|
||||
]
|
||||
selected = await ContextManager(
|
||||
ContextConfig(enforce_max_turns=args.window_turns)
|
||||
).process(messages)
|
||||
selected = await materialize_image_media_refs(
|
||||
selected, ImageMediaStore(args.media_root)
|
||||
)
|
||||
payload = [message.model_dump() for message in selected]
|
||||
semantic_hash = hashlib.sha256(
|
||||
json.dumps(payload, sort_keys=True, ensure_ascii=False).encode()
|
||||
).hexdigest()
|
||||
return payload, semantic_hash, messages
|
||||
|
||||
async def session(
|
||||
client: AsyncOpenAI, session_id: str, request_count: int
|
||||
) -> dict[str, Any]:
|
||||
database = SQLiteDatabase(str(args.database))
|
||||
await database.initialize()
|
||||
semantic_hashes: list[str] = []
|
||||
image_counts: list[int] = []
|
||||
image_hash_orders: list[list[str]] = []
|
||||
for turn in range(request_count):
|
||||
_, _, messages = await load_payload(database, session_id)
|
||||
image_refs = [
|
||||
part.model_dump() if hasattr(part, "model_dump") else part
|
||||
for message in messages
|
||||
for part in (
|
||||
message.content if isinstance(message.content, list) else []
|
||||
)
|
||||
if (part.model_dump() if hasattr(part, "model_dump") else part).get(
|
||||
"type"
|
||||
)
|
||||
== "image_media_ref"
|
||||
]
|
||||
if not image_refs:
|
||||
raise RuntimeError(
|
||||
"lifecycle fixture has no persisted image references"
|
||||
)
|
||||
messages.extend(
|
||||
[
|
||||
Message(role="user", content=[image_refs[turn % len(image_refs)]]),
|
||||
Message(role="assistant", content=f"ack {turn}"),
|
||||
]
|
||||
)
|
||||
await ConversationManager(database).update_conversation(
|
||||
"bench",
|
||||
session_id,
|
||||
history=[message.model_dump() for message in messages],
|
||||
)
|
||||
payload, semantic_hash, _ = await load_payload(database, session_id)
|
||||
semantic_hashes.append(semantic_hash)
|
||||
image_counts.append(_count_images(payload))
|
||||
image_hash_orders.append(_image_hashes(payload))
|
||||
if not image_hash_orders[-1] or image_counts[-1] > args.window_turns:
|
||||
raise RuntimeError("active image window is empty or unbounded")
|
||||
await client.chat.completions.create(model="local", messages=payload)
|
||||
del payload
|
||||
await database.engine.dispose()
|
||||
return {
|
||||
"session": session_id,
|
||||
"requests": request_count,
|
||||
"image_counts": image_counts,
|
||||
"image_hash_orders": image_hash_orders,
|
||||
"semantic_hashes": semantic_hashes,
|
||||
}
|
||||
|
||||
if args.restart_check:
|
||||
checks = []
|
||||
database = SQLiteDatabase(str(args.database))
|
||||
await database.initialize()
|
||||
for session_id in args.conversation_ids:
|
||||
payload, semantic_hash, _ = await load_payload(database, session_id)
|
||||
checks.append(
|
||||
{
|
||||
"session": session_id,
|
||||
"images": _count_images(payload),
|
||||
"semantic_hash": semantic_hash,
|
||||
"image_hash_order": _image_hashes(payload),
|
||||
}
|
||||
)
|
||||
await database.engine.dispose()
|
||||
return {"restart_checks": checks}
|
||||
|
||||
async with AsyncOpenAI(
|
||||
api_key="local", base_url=args.endpoint, max_retries=0
|
||||
) as client:
|
||||
session_count = len(args.conversation_ids)
|
||||
base_requests, remainder = divmod(args.request_count, session_count)
|
||||
request_counts = [
|
||||
base_requests + (index < remainder) for index in range(session_count)
|
||||
]
|
||||
results = await asyncio.gather(
|
||||
*(
|
||||
session(client, cid, request_count)
|
||||
for cid, request_count in zip(args.conversation_ids, request_counts)
|
||||
)
|
||||
)
|
||||
return {
|
||||
"sessions": results,
|
||||
"total_requests": sum(request_counts),
|
||||
"request_count_per_session": request_counts,
|
||||
"window_turns": args.window_turns,
|
||||
}
|
||||
|
||||
|
||||
def main() -> int:
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument("output", type=Path)
|
||||
parser.add_argument(
|
||||
"--repo", type=Path, default=Path(__file__).resolve().parents[2]
|
||||
)
|
||||
parser.add_argument("--child", action="store_true")
|
||||
parser.add_argument("--database", type=Path)
|
||||
parser.add_argument("--media-root", type=Path)
|
||||
parser.add_argument(
|
||||
"--conversation-id", dest="conversation_ids", action="append", default=[]
|
||||
)
|
||||
parser.add_argument("--endpoint")
|
||||
parser.add_argument("--rss-jsonl", type=Path)
|
||||
parser.add_argument("--restart-check", action="store_true")
|
||||
parser.add_argument("--session-count", type=int, choices=(1, 4), default=4)
|
||||
parser.add_argument("--request-count", type=int, default=REQUESTS)
|
||||
parser.add_argument("--window-turns", type=int, default=WINDOW_TURNS)
|
||||
args = parser.parse_args()
|
||||
if args.request_count < 1 or args.window_turns < 1:
|
||||
parser.error("request count and window turns must be positive")
|
||||
if args.child:
|
||||
if (
|
||||
args.database is None
|
||||
or args.media_root is None
|
||||
or args.endpoint is None
|
||||
or len(args.conversation_ids) != args.session_count
|
||||
):
|
||||
parser.error(
|
||||
"child requires database, media root, endpoint, and the requested conversation IDs"
|
||||
)
|
||||
result = asyncio.run(_child(args))
|
||||
args.output.write_text(json.dumps(result) + "\n", encoding="utf-8")
|
||||
return 0
|
||||
if len(args.conversation_ids) != args.session_count:
|
||||
parser.error(
|
||||
f"exactly {args.session_count} --conversation-id values are required"
|
||||
)
|
||||
if args.database is None or args.media_root is None:
|
||||
parser.error("--database and --media-root are required")
|
||||
|
||||
requests: list[dict[str, Any]] = []
|
||||
|
||||
class Sink(BaseHTTPRequestHandler):
|
||||
def do_POST(self) -> None:
|
||||
size = int(self.headers.get("Content-Length", "0"))
|
||||
body = self.rfile.read(size)
|
||||
requests.append(
|
||||
{
|
||||
"bytes": len(body),
|
||||
"sha256": hashlib.sha256(body).hexdigest(),
|
||||
"images": _count_images(json.loads(body)),
|
||||
}
|
||||
)
|
||||
response = b'{"choices":[{"message":{"content":"ok"}}]}'
|
||||
self.send_response(200)
|
||||
self.send_header("Content-Length", str(len(response)))
|
||||
self.end_headers()
|
||||
self.wfile.write(response)
|
||||
|
||||
def log_message(self, *_args: Any) -> None:
|
||||
pass
|
||||
|
||||
server = ThreadingHTTPServer(("127.0.0.1", 0), Sink)
|
||||
threading.Thread(target=server.serve_forever, daemon=True).start()
|
||||
with tempfile.TemporaryDirectory(prefix="astrbot-lifecycle-") as run_dir:
|
||||
child_output = Path(run_dir) / "child.json"
|
||||
rss_path = args.rss_jsonl or args.output.with_name(
|
||||
args.output.stem + ".rss.jsonl"
|
||||
)
|
||||
command = [
|
||||
sys.executable,
|
||||
str(Path(__file__).resolve()),
|
||||
str(child_output),
|
||||
"--child",
|
||||
"--repo",
|
||||
str(args.repo.resolve()),
|
||||
"--database",
|
||||
str(args.database.resolve()),
|
||||
"--media-root",
|
||||
str(args.media_root.resolve()),
|
||||
"--endpoint",
|
||||
f"http://127.0.0.1:{server.server_port}/v1",
|
||||
"--session-count",
|
||||
str(args.session_count),
|
||||
"--request-count",
|
||||
str(args.request_count),
|
||||
"--window-turns",
|
||||
str(args.window_turns),
|
||||
]
|
||||
for conversation_id in args.conversation_ids:
|
||||
command.extend(["--conversation-id", conversation_id])
|
||||
started = time.monotonic()
|
||||
process = subprocess.Popen(command, cwd=args.repo)
|
||||
child = psutil.Process(process.pid)
|
||||
stop_reason = None
|
||||
with rss_path.open("w", encoding="utf-8") as rss:
|
||||
while process.poll() is None:
|
||||
current = child.memory_info().rss
|
||||
private_bytes = None
|
||||
try:
|
||||
private_bytes = child.memory_full_info().uss
|
||||
except (AttributeError, psutil.AccessDenied):
|
||||
pass
|
||||
rss.write(
|
||||
json.dumps(
|
||||
{
|
||||
"seconds": time.monotonic() - started,
|
||||
"rss": current,
|
||||
"private_bytes": private_bytes,
|
||||
}
|
||||
)
|
||||
+ "\n"
|
||||
)
|
||||
rss.flush()
|
||||
if (
|
||||
current > min(2 * 1024**3, psutil.virtual_memory().available // 4)
|
||||
or time.monotonic() - started > 120
|
||||
):
|
||||
stop_reason = (
|
||||
"memory_limit"
|
||||
if current
|
||||
> min(2 * 1024**3, psutil.virtual_memory().available // 4)
|
||||
else "timeout"
|
||||
)
|
||||
process.kill()
|
||||
break
|
||||
time.sleep(0.01)
|
||||
return_code = process.wait()
|
||||
server.shutdown()
|
||||
result = (
|
||||
json.loads(child_output.read_text())
|
||||
if child_output.exists()
|
||||
else {"error": "child failed"}
|
||||
)
|
||||
validation_errors = []
|
||||
restart_result = {"skipped": stop_reason or "child_failed"}
|
||||
if stop_reason is None and return_code == 0:
|
||||
restart_output = Path(run_dir) / "restart.json"
|
||||
restart_command = command.copy()
|
||||
restart_command[2] = str(restart_output)
|
||||
restart_command.append("--restart-check")
|
||||
restart_process = subprocess.run(
|
||||
restart_command, cwd=args.repo, check=False
|
||||
)
|
||||
restart_result = (
|
||||
json.loads(restart_output.read_text())
|
||||
if restart_output.exists()
|
||||
else {"error": "restart child failed"}
|
||||
)
|
||||
if restart_result.get("restart_checks"):
|
||||
expected = {
|
||||
item["session"]: item for item in result.get("sessions", [])
|
||||
}
|
||||
for check in restart_result["restart_checks"]:
|
||||
prior = expected.get(check["session"])
|
||||
if (
|
||||
prior is None
|
||||
or check["images"] != prior["image_counts"][-1]
|
||||
or check["image_hash_order"] != prior["image_hash_orders"][-1]
|
||||
):
|
||||
validation_errors.append(
|
||||
"fresh-process restart changed the active image window"
|
||||
)
|
||||
if restart_process.returncode != 0:
|
||||
return_code = restart_process.returncode
|
||||
result.update(
|
||||
{
|
||||
"returncode": return_code,
|
||||
"stop_reason": stop_reason,
|
||||
"request_count": len(requests),
|
||||
"expected_request_count": args.request_count,
|
||||
"requests": requests,
|
||||
"rss_jsonl": str(rss_path),
|
||||
"restart": restart_result,
|
||||
}
|
||||
)
|
||||
if any(
|
||||
request["images"] <= 0 or request["images"] > args.window_turns
|
||||
for request in requests
|
||||
):
|
||||
validation_errors.append("sink observed an invalid active image window")
|
||||
if len(requests) != args.request_count:
|
||||
validation_errors.append("sink did not observe the expected request count")
|
||||
if validation_errors:
|
||||
result["validation_errors"] = validation_errors
|
||||
return_code = return_code or 1
|
||||
result["returncode"] = return_code
|
||||
args.output.write_text(json.dumps(result) + "\n", encoding="utf-8")
|
||||
return return_code
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
raise SystemExit(main())
|
||||
@@ -1,399 +0,0 @@
|
||||
"""Measure real image preparation and SDK requests in isolated processes."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import asyncio
|
||||
import base64
|
||||
import csv
|
||||
import hashlib
|
||||
import json
|
||||
import os
|
||||
import subprocess
|
||||
import sys
|
||||
import tempfile
|
||||
import threading
|
||||
import time
|
||||
import tracemalloc
|
||||
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
|
||||
from pathlib import Path
|
||||
|
||||
import psutil
|
||||
|
||||
|
||||
def _base64_payload_metrics(encoded: str) -> tuple[int, str]:
|
||||
"""Hash an encoded payload in chunks without retaining decoded image bytes.
|
||||
|
||||
Args:
|
||||
encoded: Base64 payload without the data-URI prefix.
|
||||
|
||||
Returns:
|
||||
The decoded byte count and SHA-256 digest.
|
||||
"""
|
||||
digest = hashlib.sha256()
|
||||
decoded_bytes = 0
|
||||
remainder = ""
|
||||
chunk_size = 4 * 1024 * 1024
|
||||
for start in range(0, len(encoded), chunk_size):
|
||||
chunk = remainder + encoded[start : start + chunk_size]
|
||||
usable = len(chunk) - len(chunk) % 4
|
||||
if usable:
|
||||
decoded = base64.b64decode(chunk[:usable])
|
||||
digest.update(decoded)
|
||||
decoded_bytes += len(decoded)
|
||||
remainder = chunk[usable:]
|
||||
if remainder:
|
||||
decoded = base64.b64decode(remainder)
|
||||
digest.update(decoded)
|
||||
decoded_bytes += len(decoded)
|
||||
return decoded_bytes, digest.hexdigest()
|
||||
|
||||
|
||||
async def run_child(args: argparse.Namespace) -> dict:
|
||||
"""Execute production preparation and request assembly in the chosen checkout.
|
||||
|
||||
Args:
|
||||
args: Explicit source, checkout, local server and diagnostic options.
|
||||
|
||||
Returns:
|
||||
Sizes, phase durations and resource counters without image content.
|
||||
"""
|
||||
sys.path.insert(0, str(args.repo.resolve()))
|
||||
from openai import AsyncOpenAI
|
||||
|
||||
from astrbot.core.provider.entities import ProviderRequest
|
||||
|
||||
if args.trace_python:
|
||||
tracemalloc.start()
|
||||
process = psutil.Process()
|
||||
baseline_rss = process.memory_info().rss
|
||||
phases: dict[str, float] = {}
|
||||
stage = "startup"
|
||||
phase_marks = [(time.monotonic(), stage)]
|
||||
failure: dict[str, str] | None = None
|
||||
source_bytes = args.image.stat().st_size
|
||||
source_digest = hashlib.sha256()
|
||||
with args.image.open("rb") as source_stream:
|
||||
for chunk in iter(lambda: source_stream.read(1024 * 1024), b""):
|
||||
source_digest.update(chunk)
|
||||
source_sha256 = source_digest.hexdigest()
|
||||
prepared_bytes: int | None = None
|
||||
prepared_sha256: str | None = None
|
||||
started = time.perf_counter()
|
||||
with tempfile.TemporaryDirectory(prefix="astrbot-image-runtime-") as runtime:
|
||||
os.environ["ASTRBOT_ROOT"] = runtime
|
||||
before_files = set(Path(runtime).rglob("*"))
|
||||
try:
|
||||
stage = "provider_prepare_and_assemble"
|
||||
phase_marks.append((time.monotonic(), stage))
|
||||
tick = time.perf_counter()
|
||||
request = ProviderRequest(
|
||||
prompt="Describe this image.",
|
||||
image_urls=[str(args.image.resolve())],
|
||||
)
|
||||
context = await request.assemble_context()
|
||||
phases["provider_prepare_and_assemble_ms"] = (
|
||||
time.perf_counter() - tick
|
||||
) * 1000
|
||||
image_part = next(
|
||||
part
|
||||
for part in context.get("content", [])
|
||||
if part.get("type") == "image_url"
|
||||
)
|
||||
data_url = image_part["image_url"]["url"]
|
||||
stage = "benchmark_payload_verification"
|
||||
phase_marks.append((time.monotonic(), stage))
|
||||
prepared_bytes, prepared_sha256 = _base64_payload_metrics(
|
||||
data_url.split(",", 1)[1]
|
||||
)
|
||||
stage = "sdk_send"
|
||||
phase_marks.append((time.monotonic(), stage))
|
||||
async with AsyncOpenAI(
|
||||
api_key="local-experiment-only",
|
||||
base_url=args.endpoint,
|
||||
max_retries=0,
|
||||
) as client:
|
||||
for _ in range(args.warmup):
|
||||
await client.chat.completions.create(
|
||||
model="local", messages=[context]
|
||||
)
|
||||
tick = time.perf_counter()
|
||||
for _ in range(args.repeat):
|
||||
await client.chat.completions.create(
|
||||
model="local", messages=[context]
|
||||
)
|
||||
phases["sdk_send_ms"] = (time.perf_counter() - tick) * 1000
|
||||
del context, request
|
||||
except Exception as exc:
|
||||
failure = {"stage": stage, "error_type": type(exc).__name__}
|
||||
finally:
|
||||
residual = [
|
||||
path
|
||||
for path in Path(runtime).rglob("*")
|
||||
if path.is_file() and path not in before_files
|
||||
]
|
||||
result = {
|
||||
"source_bytes": source_bytes,
|
||||
"prepared_bytes": prepared_bytes,
|
||||
"source_sha256": source_sha256,
|
||||
"prepared_sha256": prepared_sha256,
|
||||
"preserved_source_bytes": (
|
||||
prepared_sha256 == source_sha256
|
||||
if prepared_sha256 is not None
|
||||
else None
|
||||
),
|
||||
"base64_bytes": (
|
||||
4 * ((prepared_bytes + 2) // 3)
|
||||
if prepared_bytes is not None
|
||||
else None
|
||||
),
|
||||
"phase_ms": phases,
|
||||
"phase_marks": phase_marks,
|
||||
"last_stage": stage,
|
||||
"failure": failure,
|
||||
"baseline_rss": baseline_rss,
|
||||
"post_request_rss": process.memory_info().rss,
|
||||
"temporary_files_remaining": len(residual),
|
||||
"temporary_bytes_remaining": sum(
|
||||
path.stat().st_size for path in residual
|
||||
),
|
||||
"elapsed_ms": (time.perf_counter() - started) * 1000,
|
||||
}
|
||||
if args.trace_python:
|
||||
result["python_current_bytes"], result["python_peak_bytes"] = (
|
||||
tracemalloc.get_traced_memory()
|
||||
)
|
||||
tracemalloc.stop()
|
||||
if sys.platform != "win32":
|
||||
import resource
|
||||
|
||||
peak = resource.getrusage(resource.RUSAGE_SELF).ru_maxrss
|
||||
result["os_peak_rss_bytes"] = (
|
||||
peak if sys.platform == "darwin" else peak * 1024
|
||||
)
|
||||
return result
|
||||
|
||||
|
||||
def main() -> int:
|
||||
"""Run an isolated request and record measurements or a bounded failure."""
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument("image", type=Path)
|
||||
parser.add_argument("output", type=Path)
|
||||
parser.add_argument(
|
||||
"--repo", type=Path, default=Path(__file__).resolve().parents[2]
|
||||
)
|
||||
parser.add_argument("--trace-python", action="store_true")
|
||||
parser.add_argument("--child", action="store_true")
|
||||
parser.add_argument("--endpoint")
|
||||
parser.add_argument("--repeat", type=int, default=1)
|
||||
parser.add_argument("--warmup", type=int, default=0)
|
||||
parser.add_argument("--csv")
|
||||
parser.add_argument("--workload", default="single-image")
|
||||
args = parser.parse_args()
|
||||
if args.child:
|
||||
try:
|
||||
result = asyncio.run(run_child(args))
|
||||
result["status"] = "error" if result.get("failure") else "ok"
|
||||
except Exception as error:
|
||||
result = {"status": "error", "error_type": type(error).__name__}
|
||||
args.output.write_text(json.dumps(result), encoding="utf-8")
|
||||
return 0 if result["status"] == "ok" else 1
|
||||
|
||||
# The server measures actual SDK bytes without storing the request body.
|
||||
request_sizes = []
|
||||
|
||||
class Handler(BaseHTTPRequestHandler):
|
||||
def do_POST(self):
|
||||
length = int(self.headers.get("Content-Length", "0"))
|
||||
remaining = length
|
||||
while remaining:
|
||||
chunk = self.rfile.read(min(65536, remaining))
|
||||
if not chunk:
|
||||
break
|
||||
remaining -= len(chunk)
|
||||
request_sizes.append(length - remaining)
|
||||
body = b'{"id":"local","object":"chat.completion","created":0,"model":"local","choices":[{"index":0,"message":{"role":"assistant","content":"fixture response"},"finish_reason":"stop"}],"usage":{"prompt_tokens":100,"completion_tokens":2,"total_tokens":102}}'
|
||||
self.send_response(200)
|
||||
self.send_header("Content-Type", "application/json")
|
||||
self.send_header("Content-Length", str(len(body)))
|
||||
self.end_headers()
|
||||
self.wfile.write(body)
|
||||
|
||||
def log_message(self, *args):
|
||||
pass
|
||||
|
||||
server = ThreadingHTTPServer(("127.0.0.1", 0), Handler)
|
||||
server_thread = threading.Thread(target=server.serve_forever, daemon=True)
|
||||
server_thread.start()
|
||||
limit = min(2 * 1024**3, psutil.virtual_memory().available // 4)
|
||||
started = time.monotonic()
|
||||
samples = []
|
||||
stop_reason = None
|
||||
try:
|
||||
with tempfile.TemporaryDirectory(prefix="astrbot-image-measure-") as run_dir:
|
||||
child_result = Path(run_dir) / "child.json"
|
||||
command = [
|
||||
sys.executable,
|
||||
str(Path(__file__).resolve()),
|
||||
str(args.image.resolve()),
|
||||
str(child_result),
|
||||
"--child",
|
||||
"--repo",
|
||||
str(args.repo.resolve()),
|
||||
"--endpoint",
|
||||
f"http://127.0.0.1:{server.server_port}/v1",
|
||||
]
|
||||
if args.repeat != 1:
|
||||
command.extend(["--repeat", str(args.repeat)])
|
||||
if args.warmup:
|
||||
command.extend(["--warmup", str(args.warmup)])
|
||||
command.extend(["--workload", args.workload])
|
||||
if args.trace_python:
|
||||
command.append("--trace-python")
|
||||
env = os.environ.copy()
|
||||
env["ASTRBOT_ROOT"] = str(Path(run_dir) / "runtime")
|
||||
with (
|
||||
(Path(run_dir) / "stdout").open("wb") as out,
|
||||
(Path(run_dir) / "stderr").open("wb") as err,
|
||||
):
|
||||
proc = subprocess.Popen(
|
||||
command, cwd=args.repo, env=env, stdout=out, stderr=err
|
||||
)
|
||||
process = psutil.Process(proc.pid)
|
||||
try:
|
||||
while proc.poll() is None:
|
||||
try:
|
||||
children = process.children(recursive=True)
|
||||
rss = process.memory_info().rss
|
||||
descendants_rss = sum(
|
||||
p.memory_info().rss for p in children if p.is_running()
|
||||
)
|
||||
info = process.memory_info()
|
||||
samples.append(
|
||||
{
|
||||
"seconds": time.monotonic() - started,
|
||||
"rss": rss,
|
||||
"descendants_rss": descendants_rss,
|
||||
"private_bytes": getattr(info, "private", None),
|
||||
}
|
||||
)
|
||||
if rss + descendants_rss > limit:
|
||||
stop_reason = "memory_limit"
|
||||
elif time.monotonic() - started > 120:
|
||||
stop_reason = "timeout"
|
||||
if stop_reason:
|
||||
for child in children:
|
||||
child.kill()
|
||||
proc.kill()
|
||||
break
|
||||
except psutil.NoSuchProcess:
|
||||
break
|
||||
time.sleep(0.01)
|
||||
finally:
|
||||
if proc.poll() is None:
|
||||
proc.kill()
|
||||
proc.wait()
|
||||
result = (
|
||||
json.loads(child_result.read_text())
|
||||
if child_result.exists()
|
||||
else {"status": "error"}
|
||||
)
|
||||
phase_peaks = {}
|
||||
for sample in samples:
|
||||
phase = "process_startup"
|
||||
for timestamp, name in result.get("phase_marks", []):
|
||||
if timestamp > started + sample["seconds"]:
|
||||
break
|
||||
phase = name
|
||||
sample["phase"] = phase
|
||||
phase_peaks[phase] = max(phase_peaks.get(phase, 0), sample["rss"])
|
||||
result.update(
|
||||
{
|
||||
"status": stop_reason or result["status"],
|
||||
"workload": args.workload,
|
||||
"phase_rss_high_water_bytes": phase_peaks,
|
||||
"returncode": proc.returncode,
|
||||
"rss_high_water_bytes": max(
|
||||
(s["rss"] for s in samples), default=None
|
||||
),
|
||||
"private_high_water_bytes": max(
|
||||
(
|
||||
s["private_bytes"]
|
||||
for s in samples
|
||||
if s["private_bytes"] is not None
|
||||
),
|
||||
default=None,
|
||||
),
|
||||
"descendants_high_water_bytes": max(
|
||||
(s["descendants_rss"] for s in samples), default=None
|
||||
),
|
||||
"request_json_bytes": request_sizes,
|
||||
"samples": samples,
|
||||
"trace_python": args.trace_python,
|
||||
"code_commit": subprocess.check_output(
|
||||
["git", "rev-parse", "HEAD"], cwd=args.repo, text=True
|
||||
).strip(),
|
||||
"working_tree_diff_sha256": hashlib.sha256(
|
||||
subprocess.check_output(
|
||||
["git", "diff", "--binary"], cwd=args.repo
|
||||
)
|
||||
).hexdigest(),
|
||||
"benchmark_source_sha256": hashlib.sha256(
|
||||
Path(__file__).read_bytes()
|
||||
).hexdigest(),
|
||||
"monitor_sample_count": len(samples),
|
||||
}
|
||||
)
|
||||
with args.output.open("a", encoding="utf-8") as output:
|
||||
output.write(json.dumps(result) + "\n")
|
||||
if args.csv:
|
||||
fields = [
|
||||
"status",
|
||||
"workload",
|
||||
"repeat",
|
||||
"warmup",
|
||||
"rss_high_water_bytes",
|
||||
"private_high_water_bytes",
|
||||
"post_request_rss",
|
||||
"temporary_files_remaining",
|
||||
"temporary_bytes_remaining",
|
||||
"request_json_bytes",
|
||||
]
|
||||
csv_path = Path(args.csv)
|
||||
write_header = not csv_path.exists()
|
||||
with csv_path.open("a", newline="", encoding="utf-8") as csv_file:
|
||||
writer = csv.DictWriter(csv_file, fieldnames=fields)
|
||||
if write_header:
|
||||
writer.writeheader()
|
||||
writer.writerow(
|
||||
{
|
||||
"status": result.get("status"),
|
||||
"workload": args.workload,
|
||||
"repeat": args.repeat,
|
||||
"warmup": args.warmup,
|
||||
"rss_high_water_bytes": result.get("rss_high_water_bytes"),
|
||||
"private_high_water_bytes": result.get(
|
||||
"private_high_water_bytes"
|
||||
),
|
||||
"post_request_rss": result.get("post_request_rss"),
|
||||
"temporary_files_remaining": result.get(
|
||||
"temporary_files_remaining"
|
||||
),
|
||||
"temporary_bytes_remaining": result.get(
|
||||
"temporary_bytes_remaining"
|
||||
),
|
||||
"request_json_bytes": json.dumps(
|
||||
result.get("request_json_bytes", [])
|
||||
),
|
||||
}
|
||||
)
|
||||
print(json.dumps({k: v for k, v in result.items() if k != "samples"}))
|
||||
return 0 if result["status"] == "ok" else 1
|
||||
finally:
|
||||
server.shutdown()
|
||||
server.server_close()
|
||||
server_thread.join()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
raise SystemExit(main())
|
||||
@@ -1,99 +0,0 @@
|
||||
"""Explicitly externalize or re-inline exported JSONL conversation histories.
|
||||
|
||||
The input is never modified. Inspect the dry run before using --apply. Keep the
|
||||
input alongside the media directory as the rollback backup.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import asyncio
|
||||
import json
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
sys.path.insert(0, str(Path(__file__).resolve().parents[2]))
|
||||
|
||||
from astrbot.core.utils.image_media_store import (
|
||||
ImageMediaStore,
|
||||
materialize_image_media_refs,
|
||||
persist_inline_image_refs,
|
||||
)
|
||||
|
||||
|
||||
async def main() -> int:
|
||||
"""Transform an export into a new file without modifying live histories."""
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument("input", type=Path)
|
||||
parser.add_argument("output", type=Path)
|
||||
parser.add_argument("--media-dir", type=Path, required=True)
|
||||
parser.add_argument(
|
||||
"--mode", choices=["externalize", "inline"], default="externalize"
|
||||
)
|
||||
parser.add_argument("--apply", action="store_true")
|
||||
parser.add_argument("--rollback", action="store_true")
|
||||
args = parser.parse_args()
|
||||
if args.input.resolve() == args.output.resolve() or args.output.exists():
|
||||
parser.error("Output must be a new file distinct from the rollback input")
|
||||
if args.rollback:
|
||||
if not args.apply:
|
||||
parser.error("--rollback requires --apply")
|
||||
# The input is the verified original export; rollback writes that exact
|
||||
# export to a new path, so the original remains the recovery anchor.
|
||||
args.mode = "inline"
|
||||
store = ImageMediaStore(args.media_dir)
|
||||
record_count = 0
|
||||
image_count = 0
|
||||
output = args.output.open("x", encoding="utf-8") if args.apply else None
|
||||
try:
|
||||
with args.input.open(encoding="utf-8") as source:
|
||||
for line in source:
|
||||
record = json.loads(line)
|
||||
raw_history = record.get("history", [])
|
||||
history = (
|
||||
json.loads(raw_history)
|
||||
if isinstance(raw_history, str)
|
||||
else raw_history
|
||||
)
|
||||
if not isinstance(history, list):
|
||||
raise ValueError("History must be a message list")
|
||||
for message in history:
|
||||
parts = message.get("content")
|
||||
if isinstance(parts, list):
|
||||
image_count += sum(
|
||||
isinstance(part, dict)
|
||||
and part.get("type") in {"image_url", "image_media_ref"}
|
||||
for part in parts
|
||||
)
|
||||
if args.apply:
|
||||
if args.mode == "externalize" and not args.rollback:
|
||||
history = persist_inline_image_refs(history, store)
|
||||
else:
|
||||
history = await materialize_image_media_refs(
|
||||
history, store, strict=True
|
||||
)
|
||||
record["history"] = (
|
||||
json.dumps(history, ensure_ascii=False)
|
||||
if isinstance(raw_history, str)
|
||||
else history
|
||||
)
|
||||
output.write(json.dumps(record, ensure_ascii=False) + "\n")
|
||||
record_count += 1
|
||||
except BaseException:
|
||||
if output is not None:
|
||||
output.close()
|
||||
args.output.unlink(missing_ok=True)
|
||||
raise
|
||||
finally:
|
||||
if output is not None:
|
||||
output.close()
|
||||
print(
|
||||
json.dumps(
|
||||
{"records": record_count, "images": image_count, "applied": args.apply}
|
||||
)
|
||||
)
|
||||
return 0
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
raise SystemExit(asyncio.run(main()))
|
||||
@@ -1,131 +0,0 @@
|
||||
"""Prepare a durable multi-session history fixture before measurement."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import asyncio
|
||||
import json
|
||||
from copy import deepcopy
|
||||
from pathlib import Path
|
||||
|
||||
|
||||
async def prepare_fixture(
|
||||
manifest_path: Path,
|
||||
output: Path,
|
||||
session_count: int,
|
||||
history_turns: int,
|
||||
include_stress: bool,
|
||||
) -> dict:
|
||||
"""Create shared media and persisted conversations outside the measured child.
|
||||
|
||||
Args:
|
||||
manifest_path: Fixture manifest containing valid prepared image paths.
|
||||
output: Fresh directory receiving the SQLite database and media objects.
|
||||
session_count: Number of conversations to create.
|
||||
history_turns: Number of image-bearing turns per conversation.
|
||||
include_stress: Whether to include high-entropy stress fixtures.
|
||||
|
||||
Returns:
|
||||
A JSON-compatible manifest consumed by ``lifecycle_workloads.py``.
|
||||
|
||||
Raises:
|
||||
FileExistsError: The output already contains a database.
|
||||
ValueError: The input manifest or requested sizes are invalid.
|
||||
"""
|
||||
if session_count < 1 or history_turns < 1:
|
||||
raise ValueError("session count and history turns must be positive")
|
||||
rows = json.loads(manifest_path.read_text(encoding="utf-8"))
|
||||
if not isinstance(rows, list) or not rows:
|
||||
raise ValueError("fixture manifest must contain at least one image")
|
||||
if not include_stress:
|
||||
rows = [row for row in rows if not row.get("stress")]
|
||||
if not rows:
|
||||
raise ValueError("fixture manifest must contain at least one image")
|
||||
output.mkdir(parents=True, exist_ok=True)
|
||||
database_path = output / "history.db"
|
||||
if database_path.exists():
|
||||
raise FileExistsError(f"refusing to overwrite {database_path}")
|
||||
|
||||
import sys
|
||||
|
||||
sys.path.insert(0, str(Path(__file__).resolve().parents[2]))
|
||||
from astrbot.core.db.sqlite import SQLiteDatabase
|
||||
from astrbot.core.utils.image_media_store import ImageMediaStore
|
||||
|
||||
media_root = output / "media"
|
||||
store = ImageMediaStore(media_root)
|
||||
refs = []
|
||||
for row in rows:
|
||||
path = manifest_path.parent / row["path"]
|
||||
refs.append(store.put(path.read_bytes(), detail="high"))
|
||||
|
||||
history = []
|
||||
for turn in range(history_turns):
|
||||
history.extend(
|
||||
[
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{"type": "text", "text": f"fixture turn {turn}"},
|
||||
refs[turn % len(refs)].model_dump(),
|
||||
],
|
||||
},
|
||||
{"role": "assistant", "content": f"ack {turn}"},
|
||||
]
|
||||
)
|
||||
|
||||
database = SQLiteDatabase(str(database_path))
|
||||
await database.initialize()
|
||||
conversation_ids = []
|
||||
for index in range(session_count):
|
||||
conversation = await database.create_conversation(
|
||||
user_id="bench",
|
||||
platform_id="bench",
|
||||
content=deepcopy(history),
|
||||
title=f"lifecycle fixture {index}",
|
||||
)
|
||||
conversation_ids.append(conversation.conversation_id)
|
||||
await database.engine.dispose()
|
||||
return {
|
||||
"database": str(database_path.resolve()),
|
||||
"media_root": str(media_root.resolve()),
|
||||
"conversation_ids": conversation_ids,
|
||||
"session_count": session_count,
|
||||
"history_turns": history_turns,
|
||||
"media_count": len(refs),
|
||||
}
|
||||
|
||||
|
||||
async def main_async(args: argparse.Namespace) -> int:
|
||||
"""Prepare the fixture and write its machine-readable manifest."""
|
||||
result = await prepare_fixture(
|
||||
args.manifest,
|
||||
args.output,
|
||||
args.session_count,
|
||||
args.history_turns,
|
||||
args.include_stress,
|
||||
)
|
||||
args.output.mkdir(parents=True, exist_ok=True)
|
||||
output_manifest = args.output / "lifecycle-manifest.json"
|
||||
output_manifest.write_text(json.dumps(result, indent=2) + "\n", encoding="utf-8")
|
||||
print(json.dumps(result))
|
||||
return 0
|
||||
|
||||
|
||||
def main() -> int:
|
||||
"""Parse fixture preparation arguments."""
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument("manifest", type=Path)
|
||||
parser.add_argument("output", type=Path)
|
||||
parser.add_argument("--session-count", type=int, default=4)
|
||||
parser.add_argument("--history-turns", type=int, default=50)
|
||||
parser.add_argument(
|
||||
"--include-stress",
|
||||
action="store_true",
|
||||
help="Include high-entropy stress fixtures in the long-running workload.",
|
||||
)
|
||||
return asyncio.run(main_async(parser.parse_args()))
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
raise SystemExit(main())
|
||||
@@ -1,434 +0,0 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Explicit offline migration, repair, rollback, and cleanup for image history."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import base64
|
||||
import hashlib
|
||||
import json
|
||||
import shutil
|
||||
import sqlite3
|
||||
import sys
|
||||
import tempfile
|
||||
from collections.abc import Iterator
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
from astrbot.core.utils.image_media_store import ImageMediaRef, ImageMediaStore
|
||||
|
||||
|
||||
def _digest(path: Path) -> str:
|
||||
digest = hashlib.sha256()
|
||||
with path.open("rb") as stream:
|
||||
for chunk in iter(lambda: stream.read(1024 * 1024), b""):
|
||||
digest.update(chunk)
|
||||
return digest.hexdigest()
|
||||
|
||||
|
||||
def _logical_digest(connection: sqlite3.Connection) -> str:
|
||||
"""Hash every table schema and row without loading the database in memory."""
|
||||
digest = hashlib.sha256()
|
||||
tables = connection.execute(
|
||||
"SELECT name, sql, type FROM sqlite_master WHERE type IN ('table','index','trigger','view') ORDER BY name"
|
||||
)
|
||||
for name, sql, object_type in tables:
|
||||
digest.update(json.dumps([name, sql], ensure_ascii=False).encode())
|
||||
if object_type != "table":
|
||||
continue
|
||||
rows = connection.execute(
|
||||
f'SELECT * FROM "{name.replace(chr(34), chr(34) * 2)}"'
|
||||
)
|
||||
for row in rows:
|
||||
digest.update(
|
||||
json.dumps(
|
||||
list(row), ensure_ascii=False, default=str, separators=(",", ":")
|
||||
).encode()
|
||||
)
|
||||
return digest.hexdigest()
|
||||
|
||||
|
||||
def _manifest_media_name(name: str) -> tuple[str, str]:
|
||||
path = Path(name)
|
||||
if path.name != name or path.suffix not in {".bin", ".json"}:
|
||||
raise ValueError("invalid media manifest filename")
|
||||
media_id = path.stem
|
||||
if len(media_id) != 64 or any(char not in "0123456789abcdef" for char in media_id):
|
||||
raise ValueError("invalid media manifest media id")
|
||||
return media_id, path.suffix
|
||||
|
||||
|
||||
def _inline_parts(value: Any) -> Iterator[tuple[dict[str, Any], str, str]]:
|
||||
if isinstance(value, list):
|
||||
for item in value:
|
||||
yield from _inline_parts(item)
|
||||
elif isinstance(value, dict):
|
||||
if value.get("type") == "image_url" and isinstance(
|
||||
value.get("image_url"), dict
|
||||
):
|
||||
url = value["image_url"].get("url")
|
||||
if isinstance(url, str) and url.startswith("data:image/"):
|
||||
header, encoded = url.split(",", 1)
|
||||
yield value, header[5:].split(";", 1)[0], encoded
|
||||
else:
|
||||
for child in value.values():
|
||||
yield from _inline_parts(child)
|
||||
|
||||
|
||||
def _replace_inline(value: Any, store: ImageMediaStore) -> tuple[Any, int, int]:
|
||||
changed = 0
|
||||
bytes_migrated = 0
|
||||
|
||||
def visit(item: Any) -> Any:
|
||||
nonlocal changed, bytes_migrated
|
||||
if isinstance(item, list):
|
||||
return [visit(child) for child in item]
|
||||
if not isinstance(item, dict):
|
||||
return item
|
||||
if item.get("type") == "image_url" and isinstance(item.get("image_url"), dict):
|
||||
image = item["image_url"]
|
||||
url = image.get("url")
|
||||
if isinstance(url, str) and url.startswith("data:image/"):
|
||||
header, encoded = url.split(",", 1)
|
||||
data = base64.b64decode(encoded, validate=True)
|
||||
ref = store.put(
|
||||
data,
|
||||
header[5:].split(";", 1)[0],
|
||||
image.get("detail"),
|
||||
image.get("id"),
|
||||
)
|
||||
changed += 1
|
||||
bytes_migrated += len(data)
|
||||
return ref.model_dump()
|
||||
return {key: visit(child) for key, child in item.items()}
|
||||
|
||||
return visit(value), changed, bytes_migrated
|
||||
|
||||
|
||||
def _connect(path: Path) -> sqlite3.Connection:
|
||||
connection = sqlite3.connect(path)
|
||||
connection.row_factory = sqlite3.Row
|
||||
connection.execute("PRAGMA foreign_keys=ON")
|
||||
return connection
|
||||
|
||||
|
||||
def _backup(db: Path, media: Path, destination: Path) -> Path:
|
||||
"""Create a consistent database snapshot and a verified media snapshot."""
|
||||
destination.mkdir(parents=True, exist_ok=False)
|
||||
snapshot = sqlite3.connect(destination / db.name)
|
||||
readonly = sqlite3.connect(f"file:{db}?mode=ro", uri=True)
|
||||
try:
|
||||
readonly.backup(snapshot)
|
||||
snapshot.commit()
|
||||
finally:
|
||||
readonly.close()
|
||||
snapshot.close()
|
||||
if media.exists():
|
||||
(destination / "media").mkdir()
|
||||
manifest_media = []
|
||||
for path in media.iterdir():
|
||||
if path.is_symlink() or not path.is_file():
|
||||
raise RuntimeError("media backup refuses symlinks and non-files")
|
||||
target = destination / "media" / path.name
|
||||
shutil.copy2(path, target)
|
||||
manifest_media.append({"name": path.name, "sha256": _digest(target)})
|
||||
else:
|
||||
manifest_media = []
|
||||
manifest = {
|
||||
"database": _digest(destination / db.name),
|
||||
"media": manifest_media,
|
||||
}
|
||||
if (destination / "media").exists():
|
||||
manifest["media"] = sorted(manifest_media, key=lambda item: item["name"])
|
||||
(destination / "manifest.json").write_text(
|
||||
json.dumps(manifest, sort_keys=True) + "\n"
|
||||
)
|
||||
return destination
|
||||
|
||||
|
||||
def _migrate(args: argparse.Namespace) -> dict[str, int]:
|
||||
if args.apply and not args.offline:
|
||||
raise SystemExit("--apply requires --offline: stop all history writes first")
|
||||
db = Path(args.database).resolve()
|
||||
media = Path(args.media).resolve()
|
||||
connection = _connect(db)
|
||||
lock = args.apply
|
||||
try:
|
||||
if lock:
|
||||
connection.execute("BEGIN IMMEDIATE")
|
||||
backup = _backup(
|
||||
db,
|
||||
media,
|
||||
Path(args.backup)
|
||||
if args.backup
|
||||
else Path(
|
||||
tempfile.mkdtemp(prefix="astrbot-image-history-backup-parent-")
|
||||
)
|
||||
/ "snapshot",
|
||||
)
|
||||
query = (
|
||||
"SELECT inner_conversation_id, conversation_id, content FROM conversations"
|
||||
)
|
||||
parameters: tuple[Any, ...] = ()
|
||||
if args.conversation_id:
|
||||
query += (
|
||||
" WHERE conversation_id IN ("
|
||||
+ ",".join("?" for _ in args.conversation_id)
|
||||
+ ")"
|
||||
)
|
||||
parameters = tuple(args.conversation_id)
|
||||
total = changed = bytes_migrated = 0
|
||||
store = ImageMediaStore(media)
|
||||
for row in connection.execute(query, parameters):
|
||||
total += 1
|
||||
try:
|
||||
content = (
|
||||
json.loads(row["content"])
|
||||
if isinstance(row["content"], str)
|
||||
else row["content"]
|
||||
)
|
||||
converted, count, amount = (
|
||||
_replace_inline(content, store)
|
||||
if args.apply
|
||||
else (content, sum(1 for _ in _inline_parts(content)), 0)
|
||||
)
|
||||
except Exception as exc: # noqa: BLE001
|
||||
raise RuntimeError(
|
||||
f"conversation {row['conversation_id']} cannot be parsed or verified"
|
||||
) from exc
|
||||
changed += count
|
||||
bytes_migrated += amount
|
||||
if args.apply and count:
|
||||
connection.execute(
|
||||
"UPDATE conversations SET content=? WHERE inner_conversation_id=?",
|
||||
(
|
||||
json.dumps(converted, ensure_ascii=False),
|
||||
row["inner_conversation_id"],
|
||||
),
|
||||
)
|
||||
if args.apply:
|
||||
post_migration_digest = _logical_digest(connection)
|
||||
connection.commit()
|
||||
manifest_path = backup / "manifest.json"
|
||||
manifest = json.loads(manifest_path.read_text())
|
||||
manifest["post_migration_logical_digest"] = post_migration_digest
|
||||
manifest_path.write_text(json.dumps(manifest, sort_keys=True) + "\n")
|
||||
print(
|
||||
json.dumps(
|
||||
{
|
||||
"backup": str(backup),
|
||||
"records": total,
|
||||
"images": changed,
|
||||
"bytes": bytes_migrated,
|
||||
}
|
||||
)
|
||||
)
|
||||
else:
|
||||
connection.rollback()
|
||||
print(json.dumps({"dry_run": True, "records": total, "images": changed}))
|
||||
return {"records": total, "images": changed}
|
||||
finally:
|
||||
connection.close()
|
||||
|
||||
|
||||
def _cleanup(args: argparse.Namespace) -> None:
|
||||
if args.apply and not args.offline:
|
||||
raise SystemExit("cleanup --apply requires --offline")
|
||||
connection = _connect(Path(args.database).resolve())
|
||||
media = Path(args.media).resolve()
|
||||
try:
|
||||
connection.execute("BEGIN IMMEDIATE")
|
||||
referenced: set[str] = set()
|
||||
for row in connection.execute("SELECT content FROM conversations"):
|
||||
try:
|
||||
content = json.loads(row[0]) if isinstance(row[0], str) else row[0]
|
||||
if not isinstance(content, list):
|
||||
raise ValueError("history content is not a list")
|
||||
for item in content:
|
||||
if not isinstance(item, dict):
|
||||
raise ValueError("history message is not an object")
|
||||
pending = list(content)
|
||||
while pending:
|
||||
part = pending.pop()
|
||||
if isinstance(part, list):
|
||||
pending.extend(part)
|
||||
elif isinstance(part, dict):
|
||||
if part.get("type") == "image_media_ref":
|
||||
parsed = ImageMediaRef(
|
||||
media_id=part["media_id"],
|
||||
mime_type=part["mime_type"],
|
||||
width=part.get("width"),
|
||||
height=part.get("height"),
|
||||
byte_size=part["byte_size"],
|
||||
detail=part.get("detail"),
|
||||
version=part.get("version", 1),
|
||||
image_id=part.get("image_id"),
|
||||
)
|
||||
ImageMediaStore(media).read(parsed, {parsed.media_id})
|
||||
referenced.add(parsed.media_id)
|
||||
else:
|
||||
pending.extend(part.values())
|
||||
except Exception as exc: # noqa: BLE001
|
||||
raise RuntimeError("history parse failed; cleanup aborted") from exc
|
||||
quarantine = Path(
|
||||
args.quarantine or (media.parent / (media.name + ".quarantine"))
|
||||
)
|
||||
if args.apply:
|
||||
quarantine.mkdir(parents=True, exist_ok=True)
|
||||
moved = 0
|
||||
for media_id in sorted(
|
||||
{path.stem for path in media.iterdir() if path.suffix in {".bin", ".json"}}
|
||||
):
|
||||
paths = [media / f"{media_id}.bin", media / f"{media_id}.json"]
|
||||
if any(path.is_symlink() for path in paths if path.exists()):
|
||||
continue
|
||||
if media_id in referenced:
|
||||
continue
|
||||
existing_paths = [path for path in paths if path.exists()]
|
||||
if len(existing_paths) != 2:
|
||||
continue
|
||||
try:
|
||||
metadata = json.loads(paths[1].read_text())
|
||||
if (
|
||||
metadata.get("media_id") != media_id
|
||||
or _digest(paths[0]) != media_id
|
||||
):
|
||||
continue
|
||||
except (OSError, json.JSONDecodeError):
|
||||
continue
|
||||
if args.apply:
|
||||
targets = [quarantine / path.name for path in existing_paths]
|
||||
if any(target.exists() for target in targets):
|
||||
raise RuntimeError(
|
||||
"quarantine target already exists; refusing overwrite"
|
||||
)
|
||||
for path, target in zip(existing_paths, targets):
|
||||
shutil.move(str(path), target)
|
||||
moved += len(existing_paths)
|
||||
if args.apply:
|
||||
connection.commit()
|
||||
else:
|
||||
connection.rollback()
|
||||
print(
|
||||
json.dumps(
|
||||
{
|
||||
"dry_run": not args.apply,
|
||||
"quarantined": moved,
|
||||
"referenced": len(referenced),
|
||||
}
|
||||
)
|
||||
)
|
||||
finally:
|
||||
connection.close()
|
||||
|
||||
|
||||
def _restore(args: argparse.Namespace) -> None:
|
||||
"""Restore only when the database and media are unchanged since migration."""
|
||||
if not args.offline:
|
||||
raise SystemExit("restore requires --offline")
|
||||
backup = Path(args.restore).resolve()
|
||||
db = Path(args.database).resolve()
|
||||
backup_db = backup / db.name
|
||||
manifest_path = backup / "manifest.json"
|
||||
if not backup_db.is_file() or not manifest_path.is_file():
|
||||
raise SystemExit("backup is incomplete")
|
||||
manifest = json.loads(manifest_path.read_text())
|
||||
if _digest(backup_db) != manifest.get("database"):
|
||||
raise SystemExit("restore refused: archived database hash mismatch")
|
||||
archived_media = backup / "media"
|
||||
manifest_names: set[str] = set()
|
||||
for item in manifest.get("media", []):
|
||||
_manifest_media_name(item["name"])
|
||||
checksum = item.get("sha256", "")
|
||||
if (
|
||||
item["name"] in manifest_names
|
||||
or len(checksum) != 64
|
||||
or any(char not in "0123456789abcdef" for char in checksum)
|
||||
):
|
||||
raise SystemExit("restore refused: invalid media manifest")
|
||||
manifest_names.add(item["name"])
|
||||
source = archived_media / item["name"]
|
||||
if (
|
||||
source.is_symlink()
|
||||
or not source.is_file()
|
||||
or _digest(source) != item["sha256"]
|
||||
):
|
||||
raise SystemExit("restore refused: archived media hash mismatch")
|
||||
current = _connect(db)
|
||||
try:
|
||||
if _logical_digest(current) != manifest.get("post_migration_logical_digest"):
|
||||
raise SystemExit("restore refused: database changed since migration")
|
||||
finally:
|
||||
current.close()
|
||||
media = Path(args.media).resolve()
|
||||
for item in manifest.get("media", []):
|
||||
target = media / item["name"]
|
||||
if target.exists() and (
|
||||
target.is_symlink() or _digest(target) != item["sha256"]
|
||||
):
|
||||
raise SystemExit(f"restore refused: conflicting media: {target.name}")
|
||||
rollback = db.with_name(db.name + ".before-restore")
|
||||
if rollback.exists():
|
||||
raise SystemExit("refusing to overwrite an existing pre-restore copy")
|
||||
try:
|
||||
staged_media: list[tuple[Path, Path]] = []
|
||||
staging = Path(
|
||||
tempfile.mkdtemp(prefix="astrbot-restore-media-", dir=media.parent)
|
||||
)
|
||||
for item in manifest.get("media", []):
|
||||
source = archived_media / item["name"]
|
||||
staged = staging / item["name"]
|
||||
shutil.copy2(source, staged)
|
||||
if _digest(staged) != item["sha256"]:
|
||||
raise SystemExit("restore refused: staged media hash mismatch")
|
||||
staged_media.append((staged, media / item["name"]))
|
||||
media.mkdir(parents=True, exist_ok=True)
|
||||
for staged, target in staged_media:
|
||||
if not target.exists():
|
||||
staged.replace(target)
|
||||
|
||||
rollback_connection = sqlite3.connect(rollback)
|
||||
current = _connect(db)
|
||||
try:
|
||||
current.backup(rollback_connection)
|
||||
rollback_connection.commit()
|
||||
finally:
|
||||
current.close()
|
||||
rollback_connection.close()
|
||||
destination = _connect(db)
|
||||
archived = _connect(backup_db)
|
||||
try:
|
||||
archived.backup(destination)
|
||||
destination.commit()
|
||||
finally:
|
||||
archived.close()
|
||||
destination.close()
|
||||
finally:
|
||||
if "staging" in locals():
|
||||
shutil.rmtree(staging, ignore_errors=True)
|
||||
print(json.dumps({"restored": str(db), "pre_restore_copy": str(rollback)}))
|
||||
|
||||
|
||||
def main() -> int:
|
||||
parser = argparse.ArgumentParser(description=__doc__)
|
||||
parser.add_argument("database")
|
||||
parser.add_argument("--media", required=True)
|
||||
parser.add_argument("--conversation-id", action="append")
|
||||
parser.add_argument("--apply", action="store_true")
|
||||
parser.add_argument("--offline", action="store_true")
|
||||
parser.add_argument("--backup")
|
||||
parser.add_argument("--cleanup", action="store_true")
|
||||
parser.add_argument("--quarantine")
|
||||
parser.add_argument("--restore")
|
||||
args = parser.parse_args()
|
||||
if args.restore:
|
||||
_restore(args)
|
||||
elif args.cleanup:
|
||||
_cleanup(args)
|
||||
else:
|
||||
_migrate(args)
|
||||
return 0
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(main())
|
||||
@@ -1,36 +0,0 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import subprocess
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
|
||||
def test_lifecycle_workload_defines_bounded_matrix_and_no_payload_dump():
|
||||
source = Path("scripts/image_memory_bench/lifecycle_workloads.py").read_text(
|
||||
encoding="utf-8"
|
||||
)
|
||||
assert "REQUESTS = 200" in source
|
||||
assert "SESSIONS = 4" in source
|
||||
assert "--request-count" in source
|
||||
assert "--window-turns" in source
|
||||
assert "time.sleep(0.01)" in source
|
||||
assert "120" in source
|
||||
assert "b64encode" not in source
|
||||
assert "print(" not in source
|
||||
|
||||
|
||||
def test_lifecycle_child_requires_real_database_and_reports_wire_metadata(tmp_path):
|
||||
output = tmp_path / "result.jsonl"
|
||||
result = subprocess.run(
|
||||
[
|
||||
sys.executable,
|
||||
"scripts/image_memory_bench/lifecycle_workloads.py",
|
||||
str(output),
|
||||
"--child",
|
||||
],
|
||||
check=False,
|
||||
capture_output=True,
|
||||
text=True,
|
||||
)
|
||||
assert result.returncode != 0
|
||||
assert not output.exists()
|
||||
@@ -1,59 +0,0 @@
|
||||
"""Safety checks for the standalone image benchmark utilities."""
|
||||
|
||||
import json
|
||||
import subprocess
|
||||
import sys
|
||||
|
||||
|
||||
def test_baseline_manifest_records_fixture_fingerprints(tmp_path):
|
||||
fixture = tmp_path / "fixtures" / "7"
|
||||
fixture.mkdir(parents=True)
|
||||
(fixture / "sample.png").write_bytes(b"fixture")
|
||||
output = tmp_path / "manifest.json"
|
||||
subprocess.run(
|
||||
[
|
||||
sys.executable,
|
||||
"scripts/image_memory_bench/baseline_manifest.py",
|
||||
str(output),
|
||||
"--fixtures",
|
||||
str(tmp_path / "fixtures"),
|
||||
],
|
||||
check=True,
|
||||
)
|
||||
manifest = json.loads(output.read_text())
|
||||
assert manifest["fixtures"][0]["sha256"]
|
||||
|
||||
|
||||
def test_migration_dry_run_and_rollback_keep_input(tmp_path):
|
||||
source = tmp_path / "history.jsonl"
|
||||
source.write_text(
|
||||
json.dumps({"history": [{"role": "user", "content": "hello"}]}) + "\n"
|
||||
)
|
||||
output = tmp_path / "migrated.jsonl"
|
||||
media = tmp_path / "media"
|
||||
command = [
|
||||
sys.executable,
|
||||
"scripts/image_memory_bench/migrate_history.py",
|
||||
str(source),
|
||||
str(output),
|
||||
"--media-dir",
|
||||
str(media),
|
||||
]
|
||||
subprocess.run(command, check=True)
|
||||
assert not output.exists()
|
||||
subprocess.run(command + ["--apply"], check=True)
|
||||
rollback = tmp_path / "rollback.jsonl"
|
||||
subprocess.run(
|
||||
[
|
||||
sys.executable,
|
||||
"scripts/image_memory_bench/migrate_history.py",
|
||||
str(output),
|
||||
str(rollback),
|
||||
"--media-dir",
|
||||
str(media),
|
||||
"--rollback",
|
||||
"--apply",
|
||||
],
|
||||
check=True,
|
||||
)
|
||||
assert output.read_bytes() == rollback.read_bytes()
|
||||
@@ -1,354 +0,0 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import base64
|
||||
import importlib.util
|
||||
import json
|
||||
import sqlite3
|
||||
import subprocess
|
||||
import sys
|
||||
|
||||
import pytest
|
||||
from PIL import Image
|
||||
|
||||
|
||||
def _image() -> bytes:
|
||||
import io
|
||||
|
||||
output = io.BytesIO()
|
||||
Image.new("RGB", (3, 2), "red").save(output, "PNG")
|
||||
return output.getvalue()
|
||||
|
||||
|
||||
def _db(path, content):
|
||||
connection = sqlite3.connect(path)
|
||||
connection.execute(
|
||||
"CREATE TABLE conversations (inner_conversation_id INTEGER PRIMARY KEY, conversation_id TEXT, content JSON)"
|
||||
)
|
||||
connection.execute(
|
||||
"INSERT INTO conversations VALUES (1, 'one', ?)", (json.dumps(content),)
|
||||
)
|
||||
connection.commit()
|
||||
connection.close()
|
||||
|
||||
|
||||
def _run(db, media, *args):
|
||||
return subprocess.run(
|
||||
[
|
||||
sys.executable,
|
||||
"scripts/manage_image_history.py",
|
||||
str(db),
|
||||
"--media",
|
||||
str(media),
|
||||
*args,
|
||||
],
|
||||
check=False,
|
||||
capture_output=True,
|
||||
text=True,
|
||||
)
|
||||
|
||||
|
||||
def test_migration_defaults_to_dry_run(tmp_path):
|
||||
data = _image()
|
||||
history = [
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{
|
||||
"type": "image_url",
|
||||
"image_url": {
|
||||
"url": "data:image/png;base64,"
|
||||
+ base64.b64encode(data).decode()
|
||||
},
|
||||
}
|
||||
],
|
||||
}
|
||||
]
|
||||
db = tmp_path / "db.sqlite"
|
||||
media = tmp_path / "media"
|
||||
_db(db, history)
|
||||
result = _run(db, media)
|
||||
assert result.returncode == 0
|
||||
assert not media.exists()
|
||||
assert (
|
||||
json.loads(
|
||||
sqlite3.connect(db)
|
||||
.execute("SELECT content FROM conversations")
|
||||
.fetchone()[0]
|
||||
)
|
||||
== history
|
||||
)
|
||||
|
||||
|
||||
def test_apply_creates_backup_and_references(tmp_path):
|
||||
data = _image()
|
||||
history = [
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{
|
||||
"type": "image_url",
|
||||
"image_url": {
|
||||
"url": "data:image/png;base64,"
|
||||
+ base64.b64encode(data).decode(),
|
||||
"id": "keep",
|
||||
},
|
||||
}
|
||||
],
|
||||
}
|
||||
]
|
||||
db = tmp_path / "db.sqlite"
|
||||
media = tmp_path / "media"
|
||||
backup = tmp_path / "backup"
|
||||
_db(db, history)
|
||||
result = _run(db, media, "--apply", "--offline", "--backup", str(backup))
|
||||
assert result.returncode == 0, result.stderr
|
||||
content = json.loads(
|
||||
sqlite3.connect(db).execute("SELECT content FROM conversations").fetchone()[0]
|
||||
)
|
||||
ref = content[0]["content"][0]
|
||||
assert ref["type"] == "image_media_ref" and ref["image_id"] == "keep"
|
||||
assert (backup / "db.sqlite").exists() and (backup / "manifest.json").exists()
|
||||
|
||||
|
||||
def test_apply_requires_offline_and_cleanup_dry_run_is_read_only(tmp_path):
|
||||
db = tmp_path / "db.sqlite"
|
||||
media = tmp_path / "media"
|
||||
_db(db, [])
|
||||
media.mkdir()
|
||||
(media / "dead.bin").write_bytes(b"x")
|
||||
result = _run(db, media, "--apply")
|
||||
assert result.returncode != 0
|
||||
result = _run(db, media, "--cleanup")
|
||||
assert result.returncode == 0
|
||||
assert (media / "dead.bin").exists()
|
||||
|
||||
|
||||
def test_cleanup_aborts_on_malformed_history(tmp_path):
|
||||
db = tmp_path / "db.sqlite"
|
||||
media = tmp_path / "media"
|
||||
_db(db, {"not": "a history"})
|
||||
media.mkdir()
|
||||
(media / ("a" * 64 + ".bin")).write_bytes(b"orphan")
|
||||
result = _run(db, media, "--cleanup", "--apply", "--offline")
|
||||
assert result.returncode != 0
|
||||
assert (media / ("a" * 64 + ".bin")).exists()
|
||||
|
||||
|
||||
def test_cleanup_preserves_references_nested_in_retained_messages(tmp_path):
|
||||
from astrbot.core.utils.image_media_store import ImageMediaStore
|
||||
|
||||
media = tmp_path / "media"
|
||||
store = ImageMediaStore(media)
|
||||
ref = store.put(_image())
|
||||
db = tmp_path / "db.sqlite"
|
||||
_db(db, [{"role": "tool", "content": [{"resource": ref.model_dump()}]}])
|
||||
result = _run(db, media, "--cleanup", "--apply", "--offline")
|
||||
assert result.returncode == 0, result.stderr
|
||||
assert store.read(ref, {ref.media_id}) == _image()
|
||||
|
||||
|
||||
def test_cleanup_keeps_shared_media_until_last_history_reference_is_removed(tmp_path):
|
||||
from astrbot.core.utils.image_media_store import ImageMediaStore
|
||||
|
||||
media = tmp_path / "media"
|
||||
store = ImageMediaStore(media)
|
||||
ref = store.put(_image())
|
||||
first_history = [{"role": "user", "content": [ref.model_dump()]}]
|
||||
second_history = [{"role": "user", "content": [ref.model_dump()]}]
|
||||
db = tmp_path / "db.sqlite"
|
||||
_db(db, first_history)
|
||||
connection = sqlite3.connect(db)
|
||||
connection.execute(
|
||||
"INSERT INTO conversations VALUES (2, 'two', ?)",
|
||||
(json.dumps(second_history),),
|
||||
)
|
||||
connection.commit()
|
||||
connection.close()
|
||||
|
||||
first_delete = _run(db, media, "--cleanup", "--apply", "--offline")
|
||||
assert first_delete.returncode == 0, first_delete.stderr
|
||||
assert (media / f"{ref.media_id}.bin").exists()
|
||||
|
||||
connection = sqlite3.connect(db)
|
||||
connection.execute("DELETE FROM conversations WHERE inner_conversation_id=1")
|
||||
connection.commit()
|
||||
connection.close()
|
||||
second_delete = _run(db, media, "--cleanup", "--apply", "--offline")
|
||||
assert second_delete.returncode == 0, second_delete.stderr
|
||||
assert (media / f"{ref.media_id}.bin").exists()
|
||||
|
||||
connection = sqlite3.connect(db)
|
||||
connection.execute("DELETE FROM conversations WHERE inner_conversation_id=2")
|
||||
connection.commit()
|
||||
connection.close()
|
||||
last_delete = _run(db, media, "--cleanup", "--apply", "--offline")
|
||||
assert last_delete.returncode == 0, last_delete.stderr
|
||||
assert not (media / f"{ref.media_id}.bin").exists()
|
||||
quarantine = media.parent / "media.quarantine"
|
||||
assert (quarantine / f"{ref.media_id}.bin").exists()
|
||||
|
||||
|
||||
def test_cleanup_leaves_unpaired_and_symlink_objects(tmp_path):
|
||||
db = tmp_path / "db.sqlite"
|
||||
media = tmp_path / "media"
|
||||
_db(db, [])
|
||||
media.mkdir()
|
||||
media_id = "b" * 64
|
||||
(media / f"{media_id}.bin").write_bytes(b"bad")
|
||||
(media / f"{media_id}.json").write_text("{}")
|
||||
(media / ("c" * 64 + ".bin")).write_bytes(b"bad")
|
||||
result = _run(db, media, "--cleanup", "--apply", "--offline")
|
||||
assert result.returncode == 0
|
||||
assert (media / f"{media_id}.bin").exists()
|
||||
|
||||
|
||||
def test_restore_refuses_newer_history_and_wal_is_not_left_stale(tmp_path):
|
||||
data = _image()
|
||||
history = [
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{
|
||||
"type": "image_url",
|
||||
"image_url": {
|
||||
"url": "data:image/png;base64,"
|
||||
+ base64.b64encode(data).decode()
|
||||
},
|
||||
}
|
||||
],
|
||||
}
|
||||
]
|
||||
db = tmp_path / "db.sqlite"
|
||||
media = tmp_path / "media"
|
||||
backup = tmp_path / "backup"
|
||||
_db(db, history)
|
||||
assert (
|
||||
_run(db, media, "--apply", "--offline", "--backup", str(backup)).returncode == 0
|
||||
)
|
||||
connection = sqlite3.connect(db)
|
||||
connection.execute("UPDATE conversations SET content='[]'")
|
||||
connection.execute("CREATE TABLE unrelated (value TEXT)")
|
||||
connection.execute("INSERT INTO unrelated VALUES ('changed')")
|
||||
connection.commit()
|
||||
connection.close()
|
||||
(db.with_name(db.name + "-wal")).write_bytes(b"stale")
|
||||
result = _run(db, media, "--restore", str(backup), "--offline")
|
||||
assert result.returncode != 0
|
||||
assert not (db.with_name(db.name + ".before-restore")).exists()
|
||||
|
||||
|
||||
def test_restore_validates_backup_and_conflicting_media_before_database_write(tmp_path):
|
||||
db = tmp_path / "db.sqlite"
|
||||
media = tmp_path / "media"
|
||||
backup = tmp_path / "backup"
|
||||
_db(db, [])
|
||||
assert (
|
||||
_run(db, media, "--apply", "--offline", "--backup", str(backup)).returncode == 0
|
||||
)
|
||||
manifest = backup / "manifest.json"
|
||||
manifest.write_text(
|
||||
manifest.read_text().replace(
|
||||
'"database":', '"database": "tampered", "ignored":'
|
||||
)
|
||||
)
|
||||
before = db.read_bytes()
|
||||
result = _run(db, media, "--restore", str(backup), "--offline")
|
||||
assert result.returncode != 0
|
||||
assert db.read_bytes() == before
|
||||
|
||||
|
||||
def test_successful_restore_preserves_inline_history_and_detail(tmp_path):
|
||||
data = _image()
|
||||
history = [
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{
|
||||
"type": "image_url",
|
||||
"image_url": {
|
||||
"url": "data:image/png;base64,"
|
||||
+ base64.b64encode(data).decode(),
|
||||
"detail": "high",
|
||||
},
|
||||
}
|
||||
],
|
||||
}
|
||||
]
|
||||
db = tmp_path / "db.sqlite"
|
||||
media = tmp_path / "media"
|
||||
backup = tmp_path / "backup"
|
||||
_db(db, history)
|
||||
assert (
|
||||
_run(db, media, "--apply", "--offline", "--backup", str(backup)).returncode == 0
|
||||
)
|
||||
result = _run(db, media, "--restore", str(backup), "--offline")
|
||||
assert result.returncode == 0, result.stderr
|
||||
restored = json.loads(
|
||||
sqlite3.connect(db).execute("SELECT content FROM conversations").fetchone()[0]
|
||||
)
|
||||
assert restored == history
|
||||
|
||||
|
||||
def test_backup_reads_consistent_open_wal_database(tmp_path):
|
||||
db = tmp_path / "db.sqlite"
|
||||
media = tmp_path / "media"
|
||||
backup = tmp_path / "backup"
|
||||
connection = sqlite3.connect(db)
|
||||
connection.execute("PRAGMA journal_mode=WAL")
|
||||
connection.execute(
|
||||
"CREATE TABLE conversations (inner_conversation_id INTEGER PRIMARY KEY, conversation_id TEXT, content JSON)"
|
||||
)
|
||||
connection.execute("INSERT INTO conversations VALUES (1, 'wal', '[]')")
|
||||
connection.commit()
|
||||
assert (db.with_name(db.name + "-wal")).exists()
|
||||
connection.close()
|
||||
result = _run(db, media, "--apply", "--offline", "--backup", str(backup))
|
||||
assert result.returncode == 0, result.stderr
|
||||
assert (
|
||||
sqlite3.connect(backup / "db.sqlite")
|
||||
.execute("SELECT conversation_id FROM conversations")
|
||||
.fetchone()[0]
|
||||
== "wal"
|
||||
)
|
||||
|
||||
|
||||
def test_restore_replace_failure_leaves_database_and_cleans_stage(
|
||||
tmp_path, monkeypatch
|
||||
):
|
||||
db = tmp_path / "db.sqlite"
|
||||
media = tmp_path / "media"
|
||||
backup = tmp_path / "backup"
|
||||
_db(db, [])
|
||||
from astrbot.core.utils.image_media_store import ImageMediaStore
|
||||
|
||||
ImageMediaStore(media).put(_image(), detail="high")
|
||||
assert (
|
||||
_run(db, media, "--apply", "--offline", "--backup", str(backup)).returncode == 0
|
||||
)
|
||||
for path in media.iterdir():
|
||||
path.unlink()
|
||||
module_spec = importlib.util.spec_from_file_location(
|
||||
"manage_image_history", "scripts/manage_image_history.py"
|
||||
)
|
||||
module = importlib.util.module_from_spec(module_spec)
|
||||
module_spec.loader.exec_module(module)
|
||||
|
||||
def fail_replace(self, target):
|
||||
raise OSError("injected replace failure")
|
||||
|
||||
monkeypatch.setattr(module.Path, "replace", fail_replace)
|
||||
args = type(
|
||||
"Args",
|
||||
(),
|
||||
{
|
||||
"offline": True,
|
||||
"restore": str(backup),
|
||||
"database": str(db),
|
||||
"media": str(media),
|
||||
},
|
||||
)
|
||||
before = db.read_bytes()
|
||||
with pytest.raises(OSError, match="injected replace failure"):
|
||||
module._restore(args)
|
||||
assert db.read_bytes() == before
|
||||
assert not list(tmp_path.glob("astrbot-restore-media-*"))
|
||||
Reference in New Issue
Block a user