251 lines
8.5 KiB
Python
251 lines
8.5 KiB
Python
# scripts/ingest.py
|
|
# CLI backend: import notes into the three-layer storage.
|
|
#
|
|
# Provides three functions:
|
|
# ingest_local — single local .md file
|
|
# ingest_locals — batch import from a directory
|
|
# ingest_web — single web page
|
|
|
|
import argparse
|
|
import json
|
|
import sys
|
|
from dataclasses import dataclass
|
|
from pathlib import Path
|
|
import time
|
|
from typing import List
|
|
|
|
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.loader import load_local, load_web
|
|
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.ingest")
|
|
|
|
|
|
@dataclass
|
|
class IngestResult:
|
|
"""Result of a single note ingestion."""
|
|
|
|
file_name: str
|
|
chunk_count: int
|
|
tags: str
|
|
title: str
|
|
note_id: int
|
|
|
|
|
|
def _save_to_stores(file_name: str, filepath: str, content: str,
|
|
source_url: str | None = None,
|
|
type: str | None = None, author: str | None = None,
|
|
tag: bool = True, verbose: bool = False) -> IngestResult:
|
|
"""Split, tag, and save content into SQLStore and VectorStore.
|
|
|
|
Shared helper used by both ``ingest_local`` and ``ingest_web``.
|
|
|
|
Args:
|
|
file_name: Saved filename (may include timestamp suffix).
|
|
filepath: Absolute path to the saved file.
|
|
content: Text content to process.
|
|
source_url: Source URL for web imports, None for local files.
|
|
type: Note type (note/solution/template).
|
|
author: Note author.
|
|
tag: Extract tags via LLM.
|
|
verbose: Print processing details to stdout.
|
|
|
|
Returns:
|
|
An ``IngestResult`` with file_name, chunk_count, tags, title, note_id.
|
|
"""
|
|
chunks: List[Document] = split_markdown(content)
|
|
logger.info(f"✂️ Splitter: {len(chunks)} chunks")
|
|
|
|
tags = ""
|
|
if tag:
|
|
tags = 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}")
|
|
|
|
note_id = _sql_store.insert(
|
|
filename=file_name,
|
|
filepath=filepath,
|
|
title=title,
|
|
tags=tags,
|
|
source_url=source_url,
|
|
type=type,
|
|
author=author,
|
|
chunk_count=len(chunks),
|
|
)
|
|
logger.info(f"🗄️ SQLStore: id={note_id}")
|
|
|
|
logger.info(f"Start to save {len(chunks)} chunks to VectorStore")
|
|
for chunk in chunks:
|
|
chunk.metadata["tags"] = tags
|
|
doc_ids: List[str] = _vector_store.save(file_name, chunks)
|
|
logger.info(f"🧬 VectorStore: {len(doc_ids)} chunks saved, "
|
|
f"e.g. first chunk id: {doc_ids[0]}")
|
|
|
|
return IngestResult(
|
|
file_name=file_name,
|
|
chunk_count=len(chunks),
|
|
tags=tags,
|
|
title=title,
|
|
note_id=note_id,
|
|
)
|
|
|
|
|
|
def ingest_local(file_path: Path, tag: bool = True,
|
|
verbose: bool = False, type: str | None = None,
|
|
author: str | None = None,
|
|
filename: str | None = None) -> IngestResult:
|
|
"""Import a single local .md file into three-layer storage.
|
|
|
|
Loads from filesystem, splits into chunks, extracts tags (optional),
|
|
then inserts into SQLStore and VectorStore.
|
|
|
|
Args:
|
|
file_path: Path to the local .md file.
|
|
tag: Extract tags via LLM.
|
|
verbose: Print processing details to stdout.
|
|
type: Note type (note/solution/template).
|
|
author: Note author.
|
|
filename: Custom filename override. Defaults to the source file's basename.
|
|
|
|
Returns:
|
|
An ``IngestResult`` with metadata of the imported note.
|
|
"""
|
|
saved_path, content = load_local(
|
|
file_path, _file_store, stream=verbose, filename=filename)
|
|
return _save_to_stores(saved_path.name, str(saved_path.resolve()),
|
|
content, source_url=None, type=type, author=author,
|
|
tag=tag, verbose=verbose)
|
|
|
|
|
|
def ingest_locals(dir_path: Path, tag: bool = True,
|
|
verbose: bool = False, type: str | None = None,
|
|
author: str | None = None) -> List[IngestResult]:
|
|
"""Import all .md files from a directory.
|
|
|
|
Iterates over sorted *.md files and calls ingest_local for each.
|
|
Prints a progress line like ``[1/3] fenwick.md`` for each file.
|
|
|
|
Args:
|
|
dir_path: Directory containing .md files.
|
|
tag: Extract tags via LLM.
|
|
verbose: Print per-file details.
|
|
type: Note type (note/solution/template).
|
|
author: Note author.
|
|
|
|
Returns:
|
|
List of ``IngestResult`` for each imported note.
|
|
"""
|
|
start_time = time.time()
|
|
note_files = sorted(dir_path.glob("*.md"))
|
|
files_number = len(note_files)
|
|
results = []
|
|
for idx, file_path in enumerate(note_files, 1):
|
|
print(f"[{idx}/{files_number}] {file_path.name}")
|
|
results.append(ingest_local(file_path, tag, verbose,
|
|
type=type, author=author))
|
|
print(f"导入完成: {files_number} 个文件", end="")
|
|
elapsed = time.time() - start_time
|
|
if verbose:
|
|
print(f" 耗时: {elapsed:.1f}s")
|
|
return results
|
|
|
|
|
|
def ingest_web(url: str, tag: bool = True,
|
|
verbose: bool = False, type: str | None = None,
|
|
author: str | None = None,
|
|
filename: str | None = None) -> IngestResult:
|
|
"""Import a web page into three-layer storage.
|
|
|
|
Loads via WebBaseLoader, splits into chunks, extracts tags (optional),
|
|
then inserts into SQLStore and VectorStore.
|
|
|
|
Args:
|
|
url: The web page URL to import.
|
|
tag: Extract tags via LLM.
|
|
verbose: Print processing details to stdout.
|
|
type: Note type (note/solution/template).
|
|
author: Note author.
|
|
filename: Custom filename override. Defaults to URL-derived name.
|
|
|
|
Returns:
|
|
An ``IngestResult`` with metadata of the imported note.
|
|
"""
|
|
start_time = time.time()
|
|
if verbose:
|
|
print(" 清洗: ")
|
|
print(" > ", end="")
|
|
saved_path, content = load_web(url, _file_store, stream=verbose, filename=filename)
|
|
result = _save_to_stores(saved_path.name, str(saved_path.resolve()),
|
|
content, source_url=url, type=type, author=author,
|
|
tag=tag, verbose=verbose)
|
|
elapsed = time.time() - start_time
|
|
if verbose:
|
|
print(f" 耗时: {elapsed:.1f}s")
|
|
return result
|
|
|
|
|
|
if __name__ == "__main__":
|
|
parser = argparse.ArgumentParser(
|
|
prog="algonotes ingest",
|
|
description="从本地文件/目录或 URL 导入笔记到三层存储。",
|
|
)
|
|
source = parser.add_mutually_exclusive_group(required=True)
|
|
source.add_argument("-i", "--path", type=Path,
|
|
help="本地文件或目录路径")
|
|
source.add_argument("-u", "--url",
|
|
help="网页 URL 地址")
|
|
|
|
# Must one and no more - source
|
|
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
|
|
|
|
if args.path:
|
|
if args.path.is_file():
|
|
result = ingest_local(args.path,
|
|
tag=tag, verbose=args.verbose)
|
|
elif args.path.is_dir():
|
|
results = ingest_locals(args.path,
|
|
tag=tag, verbose=args.verbose)
|
|
else:
|
|
print(f"❌ 路径不存在: {args.path}")
|
|
sys.exit(1)
|
|
elif args.url:
|
|
result = ingest_web(args.url,
|
|
tag=tag, verbose=args.verbose)
|
|
|
|
# ── Output ──
|
|
if args.path and args.path.is_dir():
|
|
output_data = [r.__dict__ for r in results]
|
|
else:
|
|
output_data = result.__dict__
|
|
|
|
if args.json:
|
|
args.json.write_text(
|
|
json.dumps(output_data, ensure_ascii=False, indent=2),
|
|
encoding="utf-8",
|
|
)
|
|
elif not (args.path and args.path.is_dir()):
|
|
# Single file/URL — print summary to stdout
|
|
print(
|
|
f" ✅ {result.file_name} ({result.chunk_count} chunks, tags: {result.tags})")
|