773 lines
28 KiB
Python
773 lines
28 KiB
Python
#!/usr/bin/env python3
|
|
"""Incrementally fetch new ASMR works and update web/data assets."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import json
|
|
import mmap
|
|
import os
|
|
import struct
|
|
import sys
|
|
from concurrent.futures import FIRST_COMPLETED, ThreadPoolExecutor, wait
|
|
from datetime import datetime, timezone
|
|
from pathlib import Path
|
|
from types import SimpleNamespace
|
|
from typing import Any, Iterable
|
|
|
|
|
|
ROOT_DIR = Path(__file__).resolve().parents[1]
|
|
if str(ROOT_DIR) not in sys.path:
|
|
sys.path.insert(0, str(ROOT_DIR))
|
|
|
|
from fetch_asmr_works import fetch_page # noqa: E402
|
|
from build_vector_data import embedding_text, normalize_work # noqa: E402
|
|
from build_vector_data_remote import ( # noqa: E402
|
|
embeddings_url,
|
|
normalize_vector,
|
|
request_embeddings_with_retries,
|
|
)
|
|
from filter_asmr_works import ( # noqa: E402
|
|
DEFAULT_EXCLUDE_TAG_IDS,
|
|
DEFAULT_EXCLUDE_TAGS,
|
|
filter_works,
|
|
)
|
|
|
|
|
|
def parse_args() -> argparse.Namespace:
|
|
parser = argparse.ArgumentParser(
|
|
description="Fetch latest ASMR works, filter them, and incrementally update web/data."
|
|
)
|
|
parser.add_argument("--input", default="asmr_works.jsonl", help="Local raw ASMR JSONL file.")
|
|
parser.add_argument(
|
|
"--filtered-output",
|
|
default="asmr_works.filtered.jsonl",
|
|
help="Filtered ASMR JSONL output file.",
|
|
)
|
|
parser.add_argument("--data-dir", default="web/data", help="Directory containing web data assets.")
|
|
parser.add_argument(
|
|
"--skip-web-data",
|
|
action="store_true",
|
|
help="Update raw and filtered JSONL files without reading or writing web/data assets.",
|
|
)
|
|
parser.add_argument("--order", default="create_date", help="ASMR API order parameter.")
|
|
parser.add_argument("--sort", default="desc", help="ASMR API sort parameter.")
|
|
parser.add_argument("--subtitle", default="0", help="ASMR API subtitle parameter.")
|
|
parser.add_argument("--page-size", type=int, default=100, help="ASMR API pageSize parameter.")
|
|
parser.add_argument("--timeout", type=float, default=30.0, help="ASMR API request timeout seconds.")
|
|
parser.add_argument("--retries", type=int, default=3, help="ASMR API retries per failed page.")
|
|
parser.add_argument(
|
|
"--max-pages",
|
|
type=int,
|
|
default=None,
|
|
help="Maximum latest pages to scan. By default, stop when a page has no new works.",
|
|
)
|
|
parser.add_argument(
|
|
"--force",
|
|
action="store_true",
|
|
help="Rebuild filtered/web data even when no new works are found.",
|
|
)
|
|
parser.add_argument("--dry-run", action="store_true", help="Fetch and report without writing files.")
|
|
parser.add_argument(
|
|
"--diff-limit",
|
|
type=int,
|
|
default=30,
|
|
help="Maximum rows to print per dry-run diff section. Use 0 for all rows.",
|
|
)
|
|
parser.add_argument(
|
|
"--exclude-tag",
|
|
action="append",
|
|
default=[],
|
|
help="Additional tag name to remove during filtering.",
|
|
)
|
|
parser.add_argument(
|
|
"--exclude-tag-id",
|
|
action="append",
|
|
type=int,
|
|
default=[],
|
|
help="Additional tag id to remove during filtering.",
|
|
)
|
|
parser.add_argument(
|
|
"--keep-default-exclude-tags",
|
|
action=argparse.BooleanOptionalAction,
|
|
default=True,
|
|
help="Keep default 女性向/乙女向 filtering.",
|
|
)
|
|
parser.add_argument(
|
|
"--api-key-env",
|
|
default="EMBEDDING_API_KEY",
|
|
help="Environment variable that contains the remote embedding API key.",
|
|
)
|
|
parser.add_argument("--embedding-timeout", type=float, default=120.0, help="Embedding request timeout seconds.")
|
|
parser.add_argument("--embedding-retries", type=int, default=3, help="Embedding retries per failed batch.")
|
|
parser.add_argument("--embedding-retry-wait", type=float, default=2.0, help="Initial embedding retry wait seconds.")
|
|
return parser.parse_args()
|
|
|
|
|
|
def clean_text(value: Any) -> str:
|
|
return value.strip() if isinstance(value, str) else ""
|
|
|
|
|
|
def identity_keys(work: dict[str, Any]) -> list[str]:
|
|
keys: list[str] = []
|
|
source_id = clean_text(work.get("source_id")) or clean_text(work.get("sourceId"))
|
|
if source_id:
|
|
keys.append(f"source:{source_id.lower()}")
|
|
work_id = work.get("id")
|
|
if work_id is not None:
|
|
keys.append(f"id:{work_id}")
|
|
return keys
|
|
|
|
|
|
def primary_key(work: dict[str, Any]) -> str:
|
|
keys = identity_keys(work)
|
|
if not keys:
|
|
raise RuntimeError(f"work has no id/source_id: {work!r}")
|
|
return keys[0]
|
|
|
|
|
|
def display_id(work: dict[str, Any]) -> str:
|
|
return clean_text(work.get("source_id")) or clean_text(work.get("sourceId")) or str(work.get("id") or "unknown")
|
|
|
|
|
|
def display_title(work: dict[str, Any]) -> str:
|
|
return clean_text(work.get("title")) or "title unknown"
|
|
|
|
|
|
def truncate(value: str, max_length: int = 90) -> str:
|
|
if len(value) <= max_length:
|
|
return value
|
|
return value[: max_length - 3] + "..."
|
|
|
|
|
|
def format_work(work: dict[str, Any]) -> str:
|
|
date = clean_text(work.get("create_date")) or clean_text(work.get("createDate"))
|
|
suffix = f" / {date}" if date else ""
|
|
return f"{display_id(work)} / {truncate(display_title(work))}{suffix}"
|
|
|
|
|
|
def read_jsonl(path: Path) -> list[dict[str, Any]]:
|
|
works: list[dict[str, Any]] = []
|
|
with path.open(encoding="utf-8") as file:
|
|
for line_no, line in enumerate(file, 1):
|
|
if not line.strip():
|
|
continue
|
|
try:
|
|
value = json.loads(line)
|
|
except json.JSONDecodeError as exc:
|
|
raise SystemExit(f"{path}:{line_no}: invalid JSON: {exc}") from exc
|
|
if not isinstance(value, dict):
|
|
raise SystemExit(f"{path}:{line_no}: expected JSON object")
|
|
works.append(value)
|
|
return works
|
|
|
|
|
|
def write_jsonl_atomic(path: Path, works: Iterable[dict[str, Any]]) -> None:
|
|
path.parent.mkdir(parents=True, exist_ok=True)
|
|
temp_path = path.with_name(path.name + ".tmp")
|
|
try:
|
|
with temp_path.open("w", encoding="utf-8") as file:
|
|
for work in works:
|
|
clean_work = {key: value for key, value in work.items() if key != "__line_no"}
|
|
file.write(json.dumps(clean_work, ensure_ascii=False, separators=(",", ":")))
|
|
file.write("\n")
|
|
temp_path.replace(path)
|
|
except Exception:
|
|
temp_path.unlink(missing_ok=True)
|
|
raise
|
|
|
|
|
|
def write_json_atomic(path: Path, value: Any, *, indent: int | None = None) -> None:
|
|
path.parent.mkdir(parents=True, exist_ok=True)
|
|
temp_path = path.with_name(path.name + ".tmp")
|
|
try:
|
|
with temp_path.open("w", encoding="utf-8") as file:
|
|
json.dump(value, file, ensure_ascii=False, indent=indent, separators=None if indent else (",", ":"))
|
|
if indent:
|
|
file.write("\n")
|
|
temp_path.replace(path)
|
|
except Exception:
|
|
temp_path.unlink(missing_ok=True)
|
|
raise
|
|
|
|
|
|
def fetch_latest_works(args: argparse.Namespace, existing_keys: set[str]) -> tuple[list[dict[str, Any]], dict[str, dict[str, Any]], int]:
|
|
fetch_args = SimpleNamespace(
|
|
order=args.order,
|
|
sort=args.sort,
|
|
subtitle=args.subtitle,
|
|
page_size=args.page_size,
|
|
timeout=args.timeout,
|
|
retries=args.retries,
|
|
)
|
|
|
|
new_works: list[dict[str, Any]] = []
|
|
fetched_existing_by_key: dict[str, dict[str, Any]] = {}
|
|
seen_new_keys: set[str] = set()
|
|
scanned_pages = 0
|
|
page = 1
|
|
|
|
while True:
|
|
if args.max_pages is not None and page > args.max_pages:
|
|
break
|
|
|
|
payload = fetch_page(fetch_args, page)
|
|
scanned_pages += 1
|
|
works = payload.get("works", [])
|
|
if not isinstance(works, list) or not works:
|
|
break
|
|
|
|
page_new = 0
|
|
for item in works:
|
|
if not isinstance(item, dict):
|
|
continue
|
|
keys = identity_keys(item)
|
|
if not keys:
|
|
continue
|
|
if any(key in existing_keys for key in keys):
|
|
for key in keys:
|
|
fetched_existing_by_key[key] = item
|
|
continue
|
|
if any(key in seen_new_keys for key in keys):
|
|
continue
|
|
new_works.append(item)
|
|
seen_new_keys.update(keys)
|
|
page_new += 1
|
|
|
|
print(
|
|
f"Fetched latest page {page}: {page_new} new / {len(works)} works",
|
|
file=sys.stderr,
|
|
)
|
|
if page_new == 0:
|
|
break
|
|
page += 1
|
|
|
|
return new_works, fetched_existing_by_key, scanned_pages
|
|
|
|
|
|
def merge_raw_works(
|
|
existing_works: list[dict[str, Any]],
|
|
new_works: list[dict[str, Any]],
|
|
fetched_existing_by_key: dict[str, dict[str, Any]],
|
|
) -> list[dict[str, Any]]:
|
|
merged: list[dict[str, Any]] = []
|
|
emitted_keys: set[str] = set()
|
|
|
|
for work in new_works:
|
|
keys = identity_keys(work)
|
|
if any(key in emitted_keys for key in keys):
|
|
continue
|
|
merged.append(work)
|
|
emitted_keys.update(keys)
|
|
|
|
for work in existing_works:
|
|
replacement = next((fetched_existing_by_key[key] for key in identity_keys(work) if key in fetched_existing_by_key), None)
|
|
candidate = replacement or work
|
|
keys = identity_keys(candidate)
|
|
if keys and any(key in emitted_keys for key in keys):
|
|
continue
|
|
merged.append(candidate)
|
|
emitted_keys.update(keys)
|
|
|
|
return merged
|
|
|
|
|
|
def numbered_works(works: list[dict[str, Any]]) -> list[dict[str, Any]]:
|
|
result: list[dict[str, Any]] = []
|
|
for index, work in enumerate(works, 1):
|
|
numbered = dict(work)
|
|
numbered["__line_no"] = index
|
|
result.append(numbered)
|
|
return result
|
|
|
|
|
|
def filter_raw_works(args: argparse.Namespace, works: list[dict[str, Any]]) -> tuple[list[dict[str, Any]], dict[str, int]]:
|
|
exclude_tags = set(args.exclude_tag)
|
|
exclude_tag_ids = set(args.exclude_tag_id)
|
|
if args.keep_default_exclude_tags:
|
|
exclude_tags.update(DEFAULT_EXCLUDE_TAGS)
|
|
exclude_tag_ids.update(DEFAULT_EXCLUDE_TAG_IDS)
|
|
return filter_works(numbered_works(works), exclude_tags, exclude_tag_ids)
|
|
|
|
|
|
def read_json(path: Path) -> Any:
|
|
with path.open(encoding="utf-8") as file:
|
|
return json.load(file)
|
|
|
|
|
|
def work_by_key(works: Iterable[dict[str, Any]]) -> dict[str, dict[str, Any]]:
|
|
result: dict[str, dict[str, Any]] = {}
|
|
for work in works:
|
|
for key in identity_keys(work):
|
|
result.setdefault(key, work)
|
|
return result
|
|
|
|
|
|
def unique_by_primary_key(works: Iterable[dict[str, Any]]) -> list[dict[str, Any]]:
|
|
result: list[dict[str, Any]] = []
|
|
seen: set[str] = set()
|
|
for work in works:
|
|
key = primary_key(work)
|
|
if key in seen:
|
|
continue
|
|
seen.add(key)
|
|
result.append(work)
|
|
return result
|
|
|
|
|
|
def raw_diff_reasons(old: dict[str, Any], new: dict[str, Any]) -> list[str]:
|
|
old_normalized = normalize_work(old, 0)
|
|
new_normalized = normalize_work(new, 0)
|
|
checks = (
|
|
("title", old_normalized.get("title"), new_normalized.get("title")),
|
|
("tags", old_normalized.get("tags"), new_normalized.get("tags")),
|
|
("vas", old_normalized.get("vas"), new_normalized.get("vas")),
|
|
("circle", old_normalized.get("circle"), new_normalized.get("circle")),
|
|
("duration", old_normalized.get("duration"), new_normalized.get("duration")),
|
|
("dlCount", old_normalized.get("dlCount"), new_normalized.get("dlCount")),
|
|
("rateAverage", old_normalized.get("rateAverage"), new_normalized.get("rateAverage")),
|
|
("rateCount", old_normalized.get("rateCount"), new_normalized.get("rateCount")),
|
|
("release", old_normalized.get("release"), new_normalized.get("release")),
|
|
("createDate", old_normalized.get("createDate"), new_normalized.get("createDate")),
|
|
)
|
|
return [name for name, old_value, new_value in checks if old_value != new_value]
|
|
|
|
|
|
def updated_existing_diffs(
|
|
existing_works: list[dict[str, Any]],
|
|
fetched_existing_by_key: dict[str, dict[str, Any]],
|
|
) -> list[dict[str, Any]]:
|
|
existing_by_key = work_by_key(existing_works)
|
|
diffs: list[dict[str, Any]] = []
|
|
for new_work in unique_by_primary_key(fetched_existing_by_key.values()):
|
|
old_work = next((existing_by_key[key] for key in identity_keys(new_work) if key in existing_by_key), None)
|
|
if old_work is None:
|
|
continue
|
|
reasons = raw_diff_reasons(old_work, new_work)
|
|
if reasons:
|
|
diffs.append({"work": new_work, "reasons": reasons})
|
|
return diffs
|
|
|
|
|
|
def filtered_change_diffs(
|
|
filtered_path: Path,
|
|
filtered_works: list[dict[str, Any]],
|
|
) -> tuple[list[dict[str, Any]], list[dict[str, Any]]]:
|
|
old_filtered = read_jsonl(filtered_path) if filtered_path.exists() else []
|
|
old_by_key = work_by_key(old_filtered)
|
|
new_by_key = work_by_key(filtered_works)
|
|
|
|
added: list[dict[str, Any]] = []
|
|
removed: list[dict[str, Any]] = []
|
|
emitted_added: set[str] = set()
|
|
emitted_removed: set[str] = set()
|
|
for work in filtered_works:
|
|
key = primary_key(work)
|
|
if key not in old_by_key and key not in emitted_added:
|
|
added.append(work)
|
|
emitted_added.add(key)
|
|
for work in old_filtered:
|
|
key = primary_key(work)
|
|
if key not in new_by_key and key not in emitted_removed:
|
|
removed.append(work)
|
|
emitted_removed.add(key)
|
|
return added, removed
|
|
|
|
|
|
def embedding_diff_reasons(old: dict[str, Any], new: dict[str, Any]) -> list[str]:
|
|
reasons: list[str] = []
|
|
for field in ("title", "tags", "vas"):
|
|
if old.get(field) != new.get(field):
|
|
reasons.append(field)
|
|
return reasons or ["embedding text"]
|
|
|
|
|
|
def jsonl_record_count(path: Path) -> int:
|
|
if not path.exists():
|
|
return -1
|
|
count = 0
|
|
with path.open(encoding="utf-8") as file:
|
|
for line in file:
|
|
if line.strip():
|
|
count += 1
|
|
return count
|
|
|
|
|
|
def web_data_out_of_sync(filtered_path: Path, data_dir: Path) -> bool:
|
|
manifest_path = data_dir / "manifest.json"
|
|
if not manifest_path.exists():
|
|
return True
|
|
try:
|
|
manifest = read_json(manifest_path)
|
|
web_count = int(manifest.get("count", -1))
|
|
except (OSError, ValueError, TypeError, json.JSONDecodeError):
|
|
return True
|
|
return jsonl_record_count(filtered_path) != web_count
|
|
|
|
|
|
def old_embedding_index_by_key(works: list[dict[str, Any]]) -> dict[str, int]:
|
|
index_by_key: dict[str, int] = {}
|
|
for index, work in enumerate(works):
|
|
if not isinstance(work, dict):
|
|
continue
|
|
for key in identity_keys(work):
|
|
index_by_key.setdefault(key, index)
|
|
return index_by_key
|
|
|
|
|
|
def existing_embedding_index(work: dict[str, Any], index_by_key: dict[str, int]) -> int | None:
|
|
for key in identity_keys(work):
|
|
index = index_by_key.get(key)
|
|
if index is not None:
|
|
return index
|
|
return None
|
|
|
|
|
|
def request_missing_embeddings(
|
|
*,
|
|
texts: list[str],
|
|
start_indexes: list[int],
|
|
manifest: dict[str, Any],
|
|
args: argparse.Namespace,
|
|
) -> dict[int, bytes]:
|
|
if not texts:
|
|
return {}
|
|
|
|
base_url = manifest.get("baseUrl")
|
|
model = manifest.get("model")
|
|
if not base_url or not model:
|
|
raise RuntimeError("manifest must contain baseUrl and model for remote embeddings")
|
|
|
|
batch_size = int(manifest.get("batchSize") or 32)
|
|
concurrency = int(manifest.get("concurrency") or 4)
|
|
retry_forever = bool(manifest.get("retryForever", True))
|
|
max_retry_wait = float(manifest.get("maxRetryWait") or 60.0)
|
|
should_normalize = bool(manifest.get("normalized", True))
|
|
url = embeddings_url(str(base_url))
|
|
api_key = os.environ.get(args.api_key_env, "")
|
|
packed_by_index: dict[int, bytes] = {}
|
|
pending = {}
|
|
completed: dict[int, list[list[float]]] = {}
|
|
next_submit = 0
|
|
next_write = 0
|
|
dimensions = int(manifest["dimensions"])
|
|
|
|
def submit_available(executor: ThreadPoolExecutor) -> None:
|
|
nonlocal next_submit
|
|
while next_submit < len(texts) and len(pending) + len(completed) < concurrency:
|
|
start = next_submit
|
|
batch = texts[start : start + batch_size]
|
|
future = executor.submit(
|
|
request_embeddings_with_retries,
|
|
url=url,
|
|
model=str(model),
|
|
inputs=batch,
|
|
api_key=api_key,
|
|
timeout=args.embedding_timeout,
|
|
retries=args.embedding_retries,
|
|
retry_wait=args.embedding_retry_wait,
|
|
max_retry_wait=max_retry_wait,
|
|
retry_forever=retry_forever,
|
|
batch_start=start,
|
|
)
|
|
pending[future] = start
|
|
next_submit += len(batch)
|
|
|
|
def pack_vectors(batch_start: int, vectors: list[list[float]]) -> None:
|
|
for offset, vector in enumerate(vectors):
|
|
if should_normalize:
|
|
vector = normalize_vector(vector)
|
|
if len(vector) != dimensions:
|
|
raise RuntimeError(
|
|
f"embedding dimensions mismatch: expected {dimensions}, got {len(vector)}"
|
|
)
|
|
packed_by_index[start_indexes[batch_start + offset]] = struct.pack(f"<{dimensions}f", *vector)
|
|
|
|
with ThreadPoolExecutor(max_workers=concurrency) as executor:
|
|
submit_available(executor)
|
|
while pending:
|
|
done, _ = wait(pending, return_when=FIRST_COMPLETED)
|
|
for future in done:
|
|
start = pending.pop(future)
|
|
vectors = future.result()
|
|
expected = min(batch_size, len(texts) - start)
|
|
if len(vectors) != expected:
|
|
raise RuntimeError(
|
|
f"batch at {start} returned {len(vectors)} embeddings; expected {expected}"
|
|
)
|
|
completed[start] = vectors
|
|
|
|
while next_write in completed:
|
|
vectors = completed.pop(next_write)
|
|
pack_vectors(next_write, vectors)
|
|
next_write += len(vectors)
|
|
print(f"Embedded {next_write}/{len(texts)} missing works", file=sys.stderr, flush=True)
|
|
|
|
submit_available(executor)
|
|
|
|
return packed_by_index
|
|
|
|
|
|
def write_embeddings_atomic(
|
|
*,
|
|
data_dir: Path,
|
|
final_works: list[dict[str, Any]],
|
|
old_works: list[dict[str, Any]],
|
|
old_index_by_key: dict[str, int],
|
|
old_embeddings_path: Path,
|
|
missing_embeddings: dict[int, bytes],
|
|
dimensions: int,
|
|
) -> None:
|
|
vector_size = dimensions * 4
|
|
expected_size = len(old_works) * vector_size
|
|
actual_size = old_embeddings_path.stat().st_size
|
|
if actual_size != expected_size:
|
|
raise RuntimeError(f"embedding size mismatch: expected {expected_size} bytes, got {actual_size} bytes")
|
|
|
|
temp_path = data_dir / "embeddings.f32.tmp"
|
|
try:
|
|
with old_embeddings_path.open("rb") as old_file, temp_path.open("wb") as new_file:
|
|
with mmap.mmap(old_file.fileno(), 0, access=mmap.ACCESS_READ) as old_map:
|
|
for index, work in enumerate(final_works):
|
|
if index in missing_embeddings:
|
|
new_file.write(missing_embeddings[index])
|
|
continue
|
|
old_index = existing_embedding_index(work, old_index_by_key)
|
|
if old_index is None:
|
|
raise RuntimeError(f"missing embedding for {primary_key(work)}")
|
|
offset = old_index * vector_size
|
|
new_file.write(old_map[offset : offset + vector_size])
|
|
temp_path.replace(data_dir / "embeddings.f32")
|
|
except Exception:
|
|
temp_path.unlink(missing_ok=True)
|
|
raise
|
|
|
|
|
|
def build_incremental_web_data(args: argparse.Namespace, filtered_works: list[dict[str, Any]]) -> dict[str, int]:
|
|
data_dir = Path(args.data_dir)
|
|
manifest_path = data_dir / "manifest.json"
|
|
works_path = data_dir / "works.json"
|
|
embeddings_path = data_dir / "embeddings.f32"
|
|
manifest = read_json(manifest_path)
|
|
if manifest.get("method") != "remote-openai-compatible":
|
|
raise RuntimeError("incremental web/data update currently supports only remote-openai-compatible data")
|
|
|
|
old_works = read_json(works_path)
|
|
if not isinstance(old_works, list):
|
|
raise RuntimeError(f"{works_path} must contain a JSON array")
|
|
|
|
dimensions = int(manifest["dimensions"])
|
|
old_index_by_key = old_embedding_index_by_key(old_works)
|
|
final_works: list[dict[str, Any]] = []
|
|
missing_texts: list[str] = []
|
|
missing_indexes: list[int] = []
|
|
reused = 0
|
|
|
|
for raw_work in filtered_works:
|
|
work = normalize_work(raw_work, len(final_works))
|
|
old_index = existing_embedding_index(work, old_index_by_key)
|
|
if old_index is None or embedding_text(work) != embedding_text(old_works[old_index]):
|
|
missing_indexes.append(len(final_works))
|
|
missing_texts.append(embedding_text(work))
|
|
else:
|
|
reused += 1
|
|
final_works.append(work)
|
|
|
|
missing_embeddings = request_missing_embeddings(
|
|
texts=missing_texts,
|
|
start_indexes=missing_indexes,
|
|
manifest=manifest,
|
|
args=args,
|
|
)
|
|
write_embeddings_atomic(
|
|
data_dir=data_dir,
|
|
final_works=final_works,
|
|
old_works=old_works,
|
|
old_index_by_key=old_index_by_key,
|
|
old_embeddings_path=embeddings_path,
|
|
missing_embeddings=missing_embeddings,
|
|
dimensions=dimensions,
|
|
)
|
|
|
|
new_manifest = dict(manifest)
|
|
new_manifest.update(
|
|
{
|
|
"count": len(final_works),
|
|
"dimensions": dimensions,
|
|
"embeddingFile": "embeddings.f32",
|
|
"worksFile": "works.json",
|
|
"generatedAt": datetime.now(timezone.utc).isoformat(),
|
|
}
|
|
)
|
|
write_json_atomic(works_path, final_works)
|
|
write_json_atomic(manifest_path, new_manifest, indent=2)
|
|
return {"reused": reused, "embedded": len(missing_indexes), "count": len(final_works)}
|
|
|
|
|
|
def dry_run_web_stats(args: argparse.Namespace, filtered_works: list[dict[str, Any]]) -> dict[str, Any]:
|
|
data_dir = Path(args.data_dir)
|
|
manifest = read_json(data_dir / "manifest.json")
|
|
if manifest.get("method") != "remote-openai-compatible":
|
|
raise RuntimeError("incremental web/data update currently supports only remote-openai-compatible data")
|
|
old_works = read_json(data_dir / "works.json")
|
|
old_index_by_key = old_embedding_index_by_key(old_works)
|
|
reused = 0
|
|
embedded = 0
|
|
embedding_diffs: list[dict[str, Any]] = []
|
|
for raw_work in filtered_works:
|
|
work = normalize_work(raw_work, reused + embedded)
|
|
old_index = existing_embedding_index(work, old_index_by_key)
|
|
if old_index is None:
|
|
embedded += 1
|
|
embedding_diffs.append({"work": work, "reason": "new work"})
|
|
elif embedding_text(work) != embedding_text(old_works[old_index]):
|
|
embedded += 1
|
|
reasons = ", ".join(embedding_diff_reasons(old_works[old_index], work))
|
|
embedding_diffs.append({"work": work, "reason": f"changed {reasons}"})
|
|
else:
|
|
reused += 1
|
|
return {"reused": reused, "embedded": embedded, "count": reused + embedded, "embeddingDiffs": embedding_diffs}
|
|
|
|
|
|
def limited_rows(rows: list[Any], limit: int) -> list[Any]:
|
|
return rows if limit == 0 else rows[:limit]
|
|
|
|
|
|
def remaining_count(rows: list[Any], limit: int) -> int:
|
|
return 0 if limit == 0 else max(0, len(rows) - limit)
|
|
|
|
|
|
def print_work_section(title: str, prefix: str, works: list[dict[str, Any]], limit: int) -> None:
|
|
print(f"\n{title}: {len(works)}")
|
|
for work in limited_rows(works, limit):
|
|
print(f" {prefix} {format_work(work)}")
|
|
remaining = remaining_count(works, limit)
|
|
if remaining:
|
|
print(f" ... {remaining} more")
|
|
|
|
|
|
def print_updated_section(diffs: list[dict[str, Any]], limit: int) -> None:
|
|
print(f"\nUpdated existing raw works: {len(diffs)}")
|
|
for diff in limited_rows(diffs, limit):
|
|
print(f" ~ {format_work(diff['work'])} / {', '.join(diff['reasons'])}")
|
|
remaining = remaining_count(diffs, limit)
|
|
if remaining:
|
|
print(f" ... {remaining} more")
|
|
|
|
|
|
def print_embedding_section(diffs: list[dict[str, Any]], limit: int) -> None:
|
|
print(f"\nEmbedding requests: {len(diffs)}")
|
|
for diff in limited_rows(diffs, limit):
|
|
marker = "+" if diff["reason"] == "new work" else "~"
|
|
print(f" {marker} {format_work(diff['work'])} / {diff['reason']}")
|
|
remaining = remaining_count(diffs, limit)
|
|
if remaining:
|
|
print(f" ... {remaining} more")
|
|
|
|
|
|
def print_dry_run_diff(
|
|
*,
|
|
args: argparse.Namespace,
|
|
new_works: list[dict[str, Any]],
|
|
existing_works: list[dict[str, Any]],
|
|
fetched_existing_by_key: dict[str, dict[str, Any]],
|
|
filtered_path: Path,
|
|
filtered_works: list[dict[str, Any]],
|
|
web_stats: dict[str, Any],
|
|
) -> None:
|
|
limit = args.diff_limit
|
|
updated_diffs = updated_existing_diffs(existing_works, fetched_existing_by_key)
|
|
filtered_added, filtered_removed = filtered_change_diffs(filtered_path, filtered_works)
|
|
|
|
print_work_section("New raw works", "+", new_works, limit)
|
|
print_updated_section(updated_diffs, limit)
|
|
print_work_section("Filtered additions", "+", filtered_added, limit)
|
|
print_work_section("Filtered removals", "-", filtered_removed, limit)
|
|
print_embedding_section(web_stats.get("embeddingDiffs", []), limit)
|
|
|
|
|
|
def validate_args(args: argparse.Namespace) -> None:
|
|
if args.page_size <= 0:
|
|
raise SystemExit("--page-size must be greater than 0")
|
|
if args.timeout <= 0:
|
|
raise SystemExit("--timeout must be greater than 0")
|
|
if args.retries < 0:
|
|
raise SystemExit("--retries must be greater than or equal to 0")
|
|
if args.max_pages is not None and args.max_pages <= 0:
|
|
raise SystemExit("--max-pages must be greater than 0")
|
|
if args.diff_limit < 0:
|
|
raise SystemExit("--diff-limit must be greater than or equal to 0")
|
|
if args.embedding_timeout <= 0:
|
|
raise SystemExit("--embedding-timeout must be greater than 0")
|
|
if args.embedding_retries < 0:
|
|
raise SystemExit("--embedding-retries must be greater than or equal to 0")
|
|
if args.embedding_retry_wait < 0:
|
|
raise SystemExit("--embedding-retry-wait must be greater than or equal to 0")
|
|
|
|
|
|
def main() -> int:
|
|
args = parse_args()
|
|
validate_args(args)
|
|
input_path = Path(args.input)
|
|
filtered_path = Path(args.filtered_output)
|
|
if not input_path.exists():
|
|
raise SystemExit(f"input file not found: {input_path}")
|
|
|
|
existing_works = read_jsonl(input_path)
|
|
existing_keys = {key for work in existing_works for key in identity_keys(work)}
|
|
new_works, fetched_existing_by_key, scanned_pages = fetch_latest_works(args, existing_keys)
|
|
|
|
needs_web_data_update = False if args.skip_web_data else web_data_out_of_sync(filtered_path, Path(args.data_dir))
|
|
if not new_works and not args.force and not needs_web_data_update:
|
|
print(f"Scanned pages: {scanned_pages}")
|
|
print("No new works found. update skipped.")
|
|
return 0
|
|
|
|
merged_works = merge_raw_works(existing_works, new_works, fetched_existing_by_key)
|
|
filtered_works, filter_stats = filter_raw_works(args, merged_works)
|
|
web_stats = None
|
|
if args.dry_run:
|
|
web_stats = {"embeddingDiffs": []} if args.skip_web_data else dry_run_web_stats(args, filtered_works)
|
|
|
|
print(f"Scanned pages: {scanned_pages}")
|
|
print(f"New raw works: {len(new_works)}")
|
|
print(f"Raw works after merge: {len(merged_works)}")
|
|
print(f"Filtered works: {filter_stats['output']} ({filter_stats['total_removed']} removed)")
|
|
if web_stats:
|
|
if args.skip_web_data:
|
|
print("web/data would be skipped")
|
|
else:
|
|
print(f"web/data would reuse {web_stats['reused']} embeddings and request {web_stats['embedded']} embeddings")
|
|
print_dry_run_diff(
|
|
args=args,
|
|
new_works=new_works,
|
|
existing_works=existing_works,
|
|
fetched_existing_by_key=fetched_existing_by_key,
|
|
filtered_path=filtered_path,
|
|
filtered_works=filtered_works,
|
|
web_stats=web_stats,
|
|
)
|
|
|
|
if args.dry_run:
|
|
print("Dry run: no files were written.")
|
|
return 0
|
|
|
|
write_jsonl_atomic(input_path, merged_works)
|
|
write_jsonl_atomic(filtered_path, filtered_works)
|
|
print(f"Updated {input_path}")
|
|
print(f"Updated {filtered_path}")
|
|
if args.skip_web_data:
|
|
print("Skipped web/data update")
|
|
return 0
|
|
|
|
web_stats = build_incremental_web_data(args, filtered_works)
|
|
print(
|
|
f"Updated {args.data_dir}: {web_stats['count']} works, "
|
|
f"reused {web_stats['reused']} embeddings, requested {web_stats['embedded']} embeddings"
|
|
)
|
|
return 0
|
|
|
|
|
|
if __name__ == "__main__":
|
|
raise SystemExit(main())
|