algonotes_rag/scripts/update.py

# scripts/update.py
# CLI backend: update existing notes in three-layer storage.
#
# Supports:
#   - Local note update (requires --file)
#   - URL note re-crawl (from stored source_url)

import argparse
import json
import time
from dataclasses import dataclass
from pathlib import Path

from langchain_core.documents import Document

from src.store.file_store import get_file_store
from src.store.sql_store import get_sql_store
from src.store.vector_store import get_vector_store
from src.logger import setup_logger
from src.ingestion.cleaner import clean_text
from src.ingestion.splitter import split_markdown
from src.ingestion.tagger import extract_tags

_file_store = get_file_store()
_sql_store = get_sql_store()
_vector_store = get_vector_store()

logger = setup_logger("algonotes.cli.update")


@dataclass
class UpdateResult:
    """Result of a note update."""

    file_name: str
    chunk_count: int
    tags: str
    title: str
    note_id: int
    success: bool


@dataclass
class MetadataResult:
    """Result of a metadata update."""

    filename: str
    success: bool
    fields_updated: list[str]


def update_metadata(filename: str, title: str | None = None,
                    tags: str | None = None, author: str | None = None,
                    type: str | None = None) -> MetadataResult:
    """Update note metadata without reprocessing content.

    Args:
        filename: The note filename to update.
        title: New title.
        tags: New tags (comma-separated).
        author: New author.
        type: New note type.

    Returns:
        A MetadataResult indicating success and updated fields.
    """
    record = _sql_store.get_by_filename(filename)
    if not record:
        return MetadataResult(filename, False, [])

    _sql_store.update_metadata(
        record["id"], title=title, tags=tags, author=author, type=type)

    vec_fields = {}
    if tags is not None:
        vec_fields["tags"] = tags
    if vec_fields:
        _vector_store.update_metadata_fields(filename, vec_fields)

    updated = [k for k, v in {"title": title, "tags": tags,
                               "author": author, "type": type}.items()
               if v is not None]
    return MetadataResult(filename, True, updated)


def update_note(file_name: str, file_path: Path | None = None,
                tag: bool = True,
                verbose: bool = False, type: str | None = None,
                author: str | None = None) -> UpdateResult:
    """Update a note across all three storage layers.

    For local notes: provide file_path to update content.
    For URL notes: omit file_path to re-crawl from stored URL.

    Args:
        file_name: The note filename to update (e.g. "fenwick.md").
        file_path: Path to new local file (optional).
        tag: Extract tags via LLM.
        verbose: Print processing details.
        type: Note type (note/solution/template).
        author: Note author.

    Returns:
        An UpdateResult indicating success and chunk count.
    """
    sql_record = _sql_store.get_by_filename(file_name)
    if not sql_record:
        logger.warning(f"Note not found: {file_name}")
        return UpdateResult(file_name, 0, "", "", 0, False)

    source_url = sql_record.get("source_url")

    if file_path:
        content = Path(file_path).read_text(encoding="utf-8")
    elif source_url:
        if verbose:
            print("  清洗: ")
            print("  > ", end="")
        from langchain_community.document_loaders import WebBaseLoader
        from src.ingestion.loader import _HEADERS
        docs = WebBaseLoader(source_url, header_template=_HEADERS).load()
        content = docs[0].page_content
        content = clean_text(content, verbose)
        _file_store.update(file_name, content)
    else:
        logger.error(f"Cannot update {file_name}: no --file and no source_url")
        return UpdateResult(file_name, 0, "", "", 0, False)

    new_filepath = str(_file_store._resolve(file_name).resolve())

    _vector_store.delete(file_name)

    chunks: list[Document] = split_markdown(content)

    tags_str = ""
    if tag:
        tags_str = extract_tags(file_name, _file_store)

    title = chunks[0].metadata.get("header_h1", file_name)
    if verbose:
        print(f"  标题: {title}")
        print(f"  分块: {len(chunks)}")
        print(f"  打 tag: {tags_str}")

    _sql_store.update(
        sql_record["id"],
        filepath=new_filepath,
        title=title,
        tags=tags_str,
        type=type,
        author=author,
        chunk_count=len(chunks),
    )

    for chunk in chunks:
        chunk.metadata["tags"] = tags_str
    doc_ids: list[str] = _vector_store.save(file_name, chunks)
    if verbose:
        print(f"  向量: {len(doc_ids)} chunks")

    logger.info(f"Updated {file_name} ({len(chunks)} chunks)")
    return UpdateResult(file_name, len(chunks), tags_str, title, sql_record["id"], True)


if __name__ == "__main__":
    parser = argparse.ArgumentParser(
        prog="algonotes update",
        description="更新已有笔记的内容。",
    )

    parser.add_argument("filename", help="笔记文件名")
    parser.add_argument("--file", type=Path, help="新的本地文件路径")
    parser.add_argument("--no-tag", action="store_true", help="跳过打 tag")
    parser.add_argument("-v", "--verbose", action="store_true", help="详细输出")
    parser.add_argument("--json", type=Path, metavar="FILE",
                        help="输出结果到 JSON 文件")

    args = parser.parse_args()
    tag = not args.no_tag

    sql_record = _sql_store.get_by_filename(args.filename)
    if not sql_record:
        print(f"❌ 笔记不存在: {args.filename}")
        exit(1)

    if not args.file and not sql_record.get("source_url"):
        print(f"❌ 笔记 {args.filename} 无 source_url，请使用 --file 指定新文件")
        exit(1)

    result = update_note(
        args.filename,
        file_path=args.file,
        tag=tag,
        verbose=args.verbose,
    )

    if result.success:
        print(f"Update success: {result.file_name} ({result.chunk_count} chunks)")
    else:
        print(f"Update failed: {result.file_name}")

    if args.json:
        args.json.write_text(
            json.dumps(result.__dict__, ensure_ascii=False, indent=2),
            encoding="utf-8",
        )