forked from fangtianchen/algonotes_rag
446 lines
16 KiB
Python
446 lines
16 KiB
Python
# scripts/cli.py
|
||
# Unified CLI entry point for `algonotes` command.
|
||
|
||
import argparse
|
||
import json
|
||
import sys
|
||
from pathlib import Path
|
||
|
||
|
||
def _setup_cli_logging():
|
||
"""Configure logging for CLI use.
|
||
|
||
Must be called before any module with logger calls is imported.
|
||
Parses sys.argv to check for ``--verbose`` before argparse (which is
|
||
fine because we only need the flag's presence, not its value).
|
||
"""
|
||
console = "-v" in sys.argv or "--verbose" in sys.argv
|
||
from src.logger import setup_logger
|
||
setup_logger("algonotes", console=console)
|
||
|
||
|
||
def _ensure_utf8_stdout():
|
||
"""Reconfigure stdout encoding to UTF-8 if running on Windows with GBK.
|
||
|
||
Agent responses may contain emoji or other characters not representable
|
||
in the system's default codepage (e.g. cp936/GBK on Chinese Windows).
|
||
"""
|
||
try:
|
||
if sys.stdout.encoding and sys.stdout.encoding.lower() not in (
|
||
"utf-8", "utf8",
|
||
):
|
||
sys.stdout.reconfigure(encoding="utf-8")
|
||
except Exception:
|
||
pass
|
||
|
||
|
||
def main():
|
||
_setup_cli_logging()
|
||
_ensure_utf8_stdout()
|
||
parser = argparse.ArgumentParser(
|
||
prog="algonotes",
|
||
description="AlgoNotes RAG CLI — 算法竞赛笔记管理工具",
|
||
)
|
||
sub = parser.add_subparsers(dest="command", required=True)
|
||
|
||
# ── ingest ──
|
||
p_ingest = sub.add_parser("ingest", help="导入笔记")
|
||
source = p_ingest.add_mutually_exclusive_group(required=True)
|
||
source.add_argument("-i", "--path", type=Path,
|
||
help="本地文件或目录路径")
|
||
source.add_argument("-u", "--url",
|
||
help="网页 URL 地址")
|
||
p_ingest.add_argument("--no-tag", action="store_true",
|
||
help="跳过打 tag")
|
||
p_ingest.add_argument("--type", choices=["note", "solution", "template"],
|
||
default="note", help="笔记类型")
|
||
p_ingest.add_argument("--author", help="笔记作者")
|
||
p_ingest.add_argument("-v", "--verbose", action="store_true",
|
||
help="详细输出")
|
||
p_ingest.add_argument("--json", type=Path, metavar="FILE",
|
||
help="输出结果到 JSON 文件")
|
||
|
||
# ── update ──
|
||
p_update = sub.add_parser("update", help="更新笔记")
|
||
u_sub = p_update.add_subparsers(dest="update_command", required=True)
|
||
|
||
p_content = u_sub.add_parser("content", help="更新笔记内容（重新处理文件）")
|
||
p_content.add_argument("filename", help="笔记文件名")
|
||
p_content.add_argument("--file", type=Path, help="新的本地文件路径")
|
||
p_content.add_argument("--no-tag", action="store_true",
|
||
help="跳过打 tag")
|
||
p_content.add_argument("-v", "--verbose", action="store_true",
|
||
help="详细输出")
|
||
p_content.add_argument("--json", type=Path, metavar="FILE",
|
||
help="输出结果到 JSON 文件")
|
||
|
||
p_meta = u_sub.add_parser("metadata", help="更新笔记元数据（不修改内容）")
|
||
p_meta.add_argument("filename", help="笔记文件名")
|
||
p_meta.add_argument("--title", help="新标题")
|
||
p_meta.add_argument("--tags", help="新标签（逗号分隔）")
|
||
p_meta.add_argument("--author", help="新作者")
|
||
p_meta.add_argument("--type", choices=["note", "solution", "template"],
|
||
help="笔记类型")
|
||
p_meta.add_argument("--json", type=Path, metavar="FILE",
|
||
help="输出结果到 JSON 文件")
|
||
|
||
# ── delete ──
|
||
p_delete = sub.add_parser("delete", help="删除笔记")
|
||
p_delete.add_argument("filename", help="笔记文件名")
|
||
p_delete.add_argument("--force", action="store_true",
|
||
help="跳过确认提示")
|
||
p_delete.add_argument("--json", type=Path, metavar="FILE",
|
||
help="输出结果到 JSON 文件")
|
||
|
||
# ── query ──
|
||
p_query = sub.add_parser("query", help="查询笔记")
|
||
q_sub = p_query.add_subparsers(dest="query_command", required=True)
|
||
|
||
p_show = q_sub.add_parser("show", help="查看笔记原文")
|
||
p_show.add_argument("filename", help="笔记文件名")
|
||
p_show.add_argument("--lines", type=int, help="只显示前 N 行")
|
||
p_show.add_argument("--json", type=Path, metavar="FILE",
|
||
help="输出到 JSON 文件")
|
||
|
||
p_list = q_sub.add_parser("list", help="列出笔记")
|
||
p_list.add_argument("--tag", help="按标签筛选")
|
||
p_list.add_argument("-v", "--verbose", action="store_true",
|
||
help="显示详细信息")
|
||
p_list.add_argument("--json", type=Path, metavar="FILE",
|
||
help="输出到 JSON 文件")
|
||
|
||
p_info = q_sub.add_parser("info", help="查看笔记详情")
|
||
p_info.add_argument("filename", nargs="?", help="笔记文件名")
|
||
p_info.add_argument("--id", type=int, help="笔记 ID")
|
||
p_info.add_argument("--json", type=Path, metavar="FILE",
|
||
help="输出到 JSON 文件")
|
||
|
||
p_search = q_sub.add_parser("search", help="语义搜索")
|
||
p_search.add_argument("query", help="搜索关键词")
|
||
p_search.add_argument("--top-k", type=int, default=5,
|
||
help="返回结果数量")
|
||
p_search.add_argument("--json", type=Path, metavar="FILE",
|
||
help="输出到 JSON 文件")
|
||
|
||
p_export = q_sub.add_parser("export",
|
||
help="导出笔记（YAML frontmatter，兼容 Obsidian）")
|
||
p_export.add_argument("filename", nargs="?", help="笔记文件名")
|
||
p_export.add_argument("--all", action="store_true", help="导出所有笔记")
|
||
p_export.add_argument("-o", "--output", type=Path,
|
||
default=Path("./export"), help="输出目录")
|
||
p_export.add_argument("--json", type=Path, metavar="FILE",
|
||
help="输出到 JSON 文件")
|
||
|
||
p_ask = q_sub.add_parser("ask", help="RAG 问答（单次，直接输出答案）")
|
||
p_ask.add_argument("question", help="问题")
|
||
p_ask.add_argument("--top-k", type=int, default=5,
|
||
help="检索结果数量（默认 5）")
|
||
p_ask.add_argument("--stream", action="store_true",
|
||
help="流式输出答案")
|
||
p_ask.add_argument("--json", type=Path, metavar="FILE",
|
||
help="输出到 JSON 文件")
|
||
|
||
# ── chat ──
|
||
p_chat = sub.add_parser("chat", help="交互式问答")
|
||
p_chat.add_argument("--thread-id", default="default",
|
||
help="会话标识符，用于区分不同对话")
|
||
|
||
# ── mcp ──
|
||
p_mcp = sub.add_parser("mcp", help="启动 MCP 服务器")
|
||
p_mcp.add_argument("--host", default="127.0.0.1",
|
||
help="监听地址（默认 127.0.0.1）")
|
||
p_mcp.add_argument("--port", type=int, default=8000,
|
||
help="监听端口（默认 8000）")
|
||
|
||
# ── dispatch ──
|
||
args = parser.parse_args()
|
||
|
||
if args.command == "ingest":
|
||
_run_ingest(args)
|
||
elif args.command == "update":
|
||
_run_update(args)
|
||
elif args.command == "delete":
|
||
_run_delete(args)
|
||
elif args.command == "query":
|
||
_run_query(args)
|
||
elif args.command == "chat":
|
||
_run_chat(args)
|
||
elif args.command == "mcp":
|
||
_run_mcp(args)
|
||
def _run_ingest(args):
|
||
from scripts.ingest import ingest_local, ingest_locals, ingest_web
|
||
|
||
tag = not args.no_tag
|
||
note_type = args.type
|
||
author = args.author
|
||
|
||
if args.path:
|
||
if args.path.is_file():
|
||
result = ingest_local(args.path,
|
||
tag=tag, verbose=args.verbose,
|
||
type=note_type, author=author)
|
||
elif args.path.is_dir():
|
||
results = ingest_locals(args.path,
|
||
tag=tag, verbose=args.verbose,
|
||
type=note_type, author=author)
|
||
else:
|
||
print(f"❌ 路径不存在: {args.path}")
|
||
sys.exit(1)
|
||
elif args.url:
|
||
result = ingest_web(args.url,
|
||
tag=tag, verbose=args.verbose,
|
||
type=note_type, author=author)
|
||
|
||
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()):
|
||
print(
|
||
f" ✅ {result.file_name} ({result.chunk_count} chunks, tags: {result.tags})")
|
||
|
||
|
||
def _run_update(args):
|
||
if args.update_command == "content":
|
||
_run_update_content(args)
|
||
elif args.update_command == "metadata":
|
||
_run_update_metadata(args)
|
||
|
||
|
||
def _run_update_content(args):
|
||
from scripts.update import update_note
|
||
from src.store.sql_store import get_sql_store
|
||
|
||
sql_store = get_sql_store()
|
||
tag = not args.no_tag
|
||
|
||
sql_record = sql_store.get_by_filename(args.filename)
|
||
if not sql_record:
|
||
print(f"❌ 笔记不存在: {args.filename}")
|
||
sys.exit(1)
|
||
|
||
if not args.file and not sql_record.get("source_url"):
|
||
print(f"❌ 笔记 {args.filename} 无 source_url，请使用 --file 指定新文件")
|
||
sys.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",
|
||
)
|
||
|
||
|
||
def _run_update_metadata(args):
|
||
from scripts.update import update_metadata
|
||
|
||
result = update_metadata(
|
||
args.filename,
|
||
title=args.title,
|
||
tags=args.tags,
|
||
author=args.author,
|
||
type=args.type,
|
||
)
|
||
|
||
if result.success:
|
||
print(f"Metadata updated: {result.filename} ({', '.join(result.fields_updated)})")
|
||
else:
|
||
print(f"Metadata update failed: {result.filename}")
|
||
|
||
if args.json:
|
||
args.json.write_text(
|
||
json.dumps({"filename": result.filename,
|
||
"success": result.success,
|
||
"fields_updated": result.fields_updated},
|
||
ensure_ascii=False, indent=2),
|
||
encoding="utf-8",
|
||
)
|
||
|
||
|
||
def _run_delete(args):
|
||
from scripts.delete import delete_note
|
||
|
||
if not args.force:
|
||
ipt = input(f"请确认删除 {args.filename}（键入 stop 以撤销操作）: ")
|
||
if ipt.lower() == "stop":
|
||
sys.exit(0)
|
||
|
||
result = delete_note(args.filename)
|
||
|
||
if result.success:
|
||
print(f"成功删除 {result.file_name}")
|
||
else:
|
||
print(f"{result.file_name} 文件不存在，无法进行删除")
|
||
|
||
if args.json:
|
||
args.json.write_text(
|
||
json.dumps(result.__dict__, ensure_ascii=False, indent=2),
|
||
encoding="utf-8",
|
||
)
|
||
|
||
|
||
def _run_query(args):
|
||
from scripts.query import (show_raw, list_notes, show_info,
|
||
semantic_search, print_list, console)
|
||
import json as _json
|
||
|
||
if args.query_command == "show":
|
||
content = show_raw(args.filename, args.lines)
|
||
if content is None:
|
||
console.print(f"[red]笔记不存在: {args.filename}[/red]")
|
||
elif args.json:
|
||
line_count = len(content.splitlines())
|
||
args.json.write_text(
|
||
_json.dumps({"filename": args.filename, "content": content,
|
||
"line_count": line_count},
|
||
ensure_ascii=False, indent=2),
|
||
encoding="utf-8",
|
||
)
|
||
else:
|
||
console.print(content)
|
||
|
||
elif args.query_command == "list":
|
||
notes = list_notes(args.tag)
|
||
if args.json:
|
||
args.json.write_text(
|
||
_json.dumps([n.__dict__ for n in notes],
|
||
ensure_ascii=False, indent=2),
|
||
encoding="utf-8",
|
||
)
|
||
else:
|
||
print_list(notes)
|
||
|
||
elif args.query_command == "info":
|
||
if not args.filename and not args.id:
|
||
console.print("[red]须指定 filename 或 --id[/red]")
|
||
else:
|
||
if args.id:
|
||
from src.store.sql_store import get_sql_store
|
||
record = get_sql_store().get(args.id)
|
||
if record:
|
||
note = show_info(record["filename"]) if record.get("filename") else None
|
||
else:
|
||
note = None
|
||
else:
|
||
note = show_info(args.filename)
|
||
if note is None:
|
||
console.print(f"[red]笔记不存在: {args.filename or f'id={args.id}'}[/red]")
|
||
elif args.json:
|
||
args.json.write_text(
|
||
_json.dumps(note.__dict__, ensure_ascii=False, indent=2),
|
||
encoding="utf-8",
|
||
)
|
||
else:
|
||
for k, v in note.__dict__.items():
|
||
console.print(f" [cyan]{k}:[/cyan] {v}")
|
||
|
||
elif args.query_command == "search":
|
||
docs = semantic_search(args.query, args.top_k)
|
||
if not docs:
|
||
console.print("[yellow]无匹配结果[/yellow]")
|
||
elif args.json:
|
||
args.json.write_text(
|
||
_json.dumps({"query": args.query,
|
||
"results": [{"source": d.metadata.get("source"),
|
||
"content": d.page_content,
|
||
"metadata": {k: v for k, v in d.metadata.items()
|
||
if k != "source"}}
|
||
for d in docs]},
|
||
ensure_ascii=False, indent=2),
|
||
encoding="utf-8",
|
||
)
|
||
else:
|
||
console.print(f"找到 {len(docs)} 条相关结果：\n")
|
||
for i, doc in enumerate(docs, 1):
|
||
console.print(
|
||
f"[cyan][{i}][/cyan] 来自: {doc.metadata.get('source', '?')}")
|
||
console.print(f" {doc.page_content[:120]}...\n")
|
||
|
||
elif args.query_command == "export":
|
||
_run_query_export(args)
|
||
elif args.query_command == "ask":
|
||
_run_query_ask(args)
|
||
|
||
|
||
def _run_query_export(args):
|
||
from scripts.query import export_note, export_all
|
||
import json as _json
|
||
|
||
try:
|
||
if args.all:
|
||
results = export_all(args.output)
|
||
if args.json:
|
||
args.json.write_text(
|
||
_json.dumps(results, ensure_ascii=False, indent=2),
|
||
encoding="utf-8",
|
||
)
|
||
else:
|
||
print(f"Exported {len(results)} notes to {args.output}")
|
||
elif args.filename:
|
||
path = export_note(args.filename, args.output)
|
||
if args.json:
|
||
args.json.write_text(
|
||
_json.dumps({"filename": args.filename,
|
||
"path": str(path)},
|
||
ensure_ascii=False, indent=2),
|
||
encoding="utf-8",
|
||
)
|
||
else:
|
||
print(f"Exported {args.filename} to {path}")
|
||
else:
|
||
print("请指定 filename 或 --all")
|
||
except FileNotFoundError as e:
|
||
print(f"❌ 笔记不存在: {e}")
|
||
except Exception as e:
|
||
print(f"❌ 导出失败: {e}")
|
||
|
||
|
||
def _run_query_ask(args):
|
||
from scripts.rag import ask_question
|
||
import json as _json
|
||
|
||
result = ask_question(args.question, stream=args.stream)
|
||
|
||
if not result.success:
|
||
print(f"❌ 问答失败: {result.error}")
|
||
sys.exit(1)
|
||
|
||
if args.json:
|
||
args.json.write_text(
|
||
_json.dumps(result.dict, ensure_ascii=False, indent=2),
|
||
encoding="utf-8",
|
||
)
|
||
elif not args.stream:
|
||
print(result.answer)
|
||
|
||
|
||
def _run_chat(args):
|
||
from scripts.chat import chat_loop
|
||
chat_loop(args.thread_id)
|
||
|
||
|
||
def _run_mcp(args):
|
||
from src.mcp.server import main as mcp_main
|
||
mcp_main(host=args.host, port=args.port)
|
||
|
||
|
||
if __name__ == "__main__":
|
||
main()
|