forked from fangtianchen/algonotes_rag
313 lines
10 KiB
Python
313 lines
10 KiB
Python
# scripts/query.py
|
||
# CLI backend: query notes — list, info, show raw content, semantic search.
|
||
|
||
from dataclasses import dataclass
|
||
import argparse
|
||
import json
|
||
from pathlib import Path
|
||
|
||
from langchain_core.documents import Document
|
||
from rich.console import Console
|
||
from rich.table import Table
|
||
|
||
from src.logger import setup_logger
|
||
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
|
||
|
||
_file_store = get_file_store()
|
||
_sql_store = get_sql_store()
|
||
_vector_store = get_vector_store()
|
||
|
||
logger = setup_logger("algonotes.cli.query")
|
||
console = Console()
|
||
|
||
|
||
@dataclass
|
||
class NoteInfo:
|
||
"""Note metadata from SQLite."""
|
||
|
||
id: int
|
||
filename: str
|
||
filepath: str
|
||
title: str | None
|
||
tags: str | None
|
||
source_url: str | None
|
||
content_hash: str
|
||
ingested_at: str
|
||
updated_at: str | None
|
||
chunk_count: int
|
||
file_size: int | None
|
||
|
||
@classmethod
|
||
def from_dict(cls, d: dict) -> "NoteInfo":
|
||
return cls(**{k: d[k] for k in cls.__dataclass_fields__ if k in d})
|
||
|
||
|
||
def show_raw(filename: str, lines: int | None = None) -> str | None:
|
||
"""Read and return the raw content of a note file.
|
||
|
||
Args:
|
||
filename: Note filename (e.g. ``"fenwick.md"``).
|
||
lines: If provided, return only the first N lines.
|
||
|
||
Returns:
|
||
File content as string, or ``None`` if the note does not exist.
|
||
"""
|
||
content = _file_store.read(filename)
|
||
if content is None:
|
||
return None
|
||
if lines is not None:
|
||
content = "\n".join(content.splitlines()[:lines])
|
||
return content
|
||
|
||
|
||
def list_notes(tag: str | None = None) -> list[NoteInfo]:
|
||
"""List all notes, optionally filtered by tag.
|
||
|
||
Args:
|
||
tag: If provided, only return notes whose tags contain this keyword.
|
||
|
||
Returns:
|
||
A list of NoteInfo objects.
|
||
"""
|
||
if tag:
|
||
return [NoteInfo.from_dict(d) for d in _sql_store.search_by_tags(tag)]
|
||
return [NoteInfo.from_dict(d) for d in _sql_store.list_all()]
|
||
|
||
|
||
def show_info(filename: str) -> NoteInfo | None:
|
||
"""Get detailed metadata for a note.
|
||
|
||
Args:
|
||
filename: Note filename (e.g. ``"fenwick.md"``).
|
||
|
||
Returns:
|
||
A NoteInfo object, or ``None`` if not found.
|
||
"""
|
||
d = _sql_store.get_by_filename(filename)
|
||
return NoteInfo.from_dict(d) if d else None
|
||
|
||
|
||
def semantic_search(query: str, top_k: int = 5) -> list[Document]:
|
||
"""Semantic search across all notes.
|
||
|
||
Args:
|
||
query: Search query string.
|
||
top_k: Number of results to return.
|
||
|
||
Returns:
|
||
A list of matching Document objects.
|
||
"""
|
||
return _vector_store.search(query, k=top_k)
|
||
|
||
|
||
def build_frontmatter(record: dict) -> str:
|
||
"""Build YAML frontmatter for Obsidian export.
|
||
|
||
Args:
|
||
record: Note metadata dict from SQL.
|
||
|
||
Returns:
|
||
YAML frontmatter string.
|
||
"""
|
||
tags_str = record.get("tags") or ""
|
||
tag_list = [t.strip() for t in tags_str.split(",") if t.strip()]
|
||
|
||
lines = ["---"]
|
||
lines.append(f"title: \"{record.get('title', '')}\"")
|
||
lines.append(f"type: \"{record.get('type') or 'note'}\"")
|
||
lines.append(f"tags: [{', '.join(tag_list)}]")
|
||
lines.append(f"author: \"{record.get('author') or ''}\"")
|
||
lines.append(f"ingested_at: {record.get('ingested_at', '')}")
|
||
if record.get("updated_at"):
|
||
lines.append(f"updated_at: {record['updated_at']}")
|
||
if record.get("source_url"):
|
||
lines.append(f"source_url: \"{record['source_url']}\"")
|
||
lines.append("---")
|
||
return "\n".join(lines)
|
||
|
||
|
||
def export_note(filename: str, output_dir: Path) -> Path:
|
||
"""Export a single note with YAML frontmatter.
|
||
|
||
Args:
|
||
filename: Note filename.
|
||
output_dir: Output directory.
|
||
|
||
Returns:
|
||
Path to exported file.
|
||
|
||
Raises:
|
||
FileNotFoundError: If note not found.
|
||
"""
|
||
record = _sql_store.get_by_filename(filename)
|
||
if not record:
|
||
raise FileNotFoundError(filename)
|
||
content = _file_store.read(filename)
|
||
frontmatter = build_frontmatter(record)
|
||
output_path = output_dir / filename
|
||
output_path.parent.mkdir(parents=True, exist_ok=True)
|
||
output_path.write_text(f"{frontmatter}\n\n{content}", encoding="utf-8")
|
||
return output_path
|
||
|
||
|
||
def export_all(output_dir: Path) -> list[dict]:
|
||
"""Export all notes with YAML frontmatter.
|
||
|
||
Args:
|
||
output_dir: Output directory.
|
||
|
||
Returns:
|
||
List of dicts with filename and path.
|
||
"""
|
||
records = _sql_store.list_all()
|
||
results = []
|
||
for r in records:
|
||
path = export_note(r["filename"], output_dir)
|
||
results.append({"filename": r["filename"], "path": str(path)})
|
||
return results
|
||
|
||
|
||
def print_list(notes: list[NoteInfo]) -> None:
|
||
"""Print notes as a formatted table using rich.
|
||
|
||
Creates a table with columns: ID, filename, tags, ingested_at.
|
||
Uses rich.table.Table for automatic column alignment and styling.
|
||
|
||
Table construction steps:
|
||
1. Create Table instance with title
|
||
2. Add columns with header text and optional style/color
|
||
3. Add rows with cell values
|
||
4. Console.print() renders the table to terminal
|
||
|
||
Args:
|
||
notes: List of NoteInfo objects to display.
|
||
"""
|
||
# 1. 创建表格，设置标题
|
||
table = Table(title="笔记列表")
|
||
|
||
# 2. 添加列：header(列名), style(整列颜色), no_wrap(禁止换行)
|
||
table.add_column("ID", style="cyan", no_wrap=True)
|
||
table.add_column("文件名", style="green")
|
||
table.add_column("标签")
|
||
table.add_column("入库时间", style="dim")
|
||
|
||
# 3. 逐行填充数据
|
||
for note in notes:
|
||
table.add_row(
|
||
str(note.id),
|
||
note.filename,
|
||
note.tags or "-",
|
||
note.ingested_at[:10], # 只显示日期部分
|
||
)
|
||
|
||
# 4. 渲染输出
|
||
console.print(table)
|
||
|
||
|
||
if __name__ == "__main__":
|
||
parser = argparse.ArgumentParser(
|
||
prog="algonotes query",
|
||
description="查询命令：查看笔记、语义搜索、RAG 问答",
|
||
)
|
||
sub = parser.add_subparsers(dest="command", required=True)
|
||
|
||
# ── show ──
|
||
p_show = 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 文件")
|
||
|
||
# ── list ──
|
||
p_list = 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 文件")
|
||
|
||
# ── info ──
|
||
p_info = 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 文件")
|
||
|
||
# ── search ──
|
||
p_search = 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 文件")
|
||
|
||
args = parser.parse_args()
|
||
|
||
# ── dispatch ──
|
||
|
||
if args.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.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.command == "info":
|
||
if not args.filename and not args.id:
|
||
console.print("[red]须指定 filename 或 --id[/red]")
|
||
else:
|
||
note = show_info(args.filename) if args.filename else None
|
||
if note is None:
|
||
console.print(f"[red]笔记不存在: {args.filename}[/red]")
|
||
elif args.json:
|
||
args.json.write_text(
|
||
json.dumps(note.__dict__, ensure_ascii=False, indent=2),
|
||
encoding="utf-8",
|
||
)
|
||
else:
|
||
# key-value 格式输出
|
||
for k, v in note.__dict__.items():
|
||
console.print(f" [cyan]{k}:[/cyan] {v}")
|
||
|
||
elif args.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")
|