query_materials.py

#!/usr/bin/env python3
"""
Get 笔记素材查询工具
快速查询知识库中的素材列表
"""

import os
import sys
import argparse
from mcp_tools.getnote import GetNoteAPI


def display_material(material, index):
    """格式化显示单个素材"""
    print(f"\n{'='*70}")
    print(f"【素材 {index}】")
    print(f"{'='*70}")

    # 标题
    title = material.get('title', '无标题')
    print(f"\n📌 标题: {title}")

    # 元信息
    print(f"\n📊 元信息:")
    print(f"   • 类型: {material.get('type', 'unknown')}")
    print(f"   • 相关度: {material.get('score', 0):.4f}")
    print(f"   • 来源: {material.get('source', 'Get笔记')}")

    # 内容预览
    content = material.get('content', '')
    if content:
        print(f"\n📝 内容预览:")
        # 显示前300字符
        preview = content[:300]
        # 去除多余的换行
        preview = ' '.join(preview.split())
        print(f"   {preview}...")

    # ID（用于进一步查询）
    if material.get('id'):
        print(f"\n🔑 ID: {material.get('id')}")


def search_materials(query, category=None, limit=5, verbose=False):
    """搜索素材"""
    try:
        # 创建客户端
        client = GetNoteAPI()

        # 显示搜索信息
        print(f"\n{'='*70}")
        print(f"🔍 搜索查询")
        print(f"{'='*70}")
        print(f"主题: {query}")
        if category:
            print(f"分类: {category}")
        print(f"返回数量: {limit}")
        print(f"{'='*70}\n")

        # 获取素材
        materials = client.get_materials_for_topic(
            topic=query,
            category=category,
            limit=limit
        )

        if not materials:
            print("\n❌ 未找到相关素材")
            print("\n💡 建议:")
            print("   1. 尝试使用更通用的关键词")
            print("   2. 检查知识库中是否有相关内容")
            print("   3. 尝试不指定分类参数\n")
            return

        # 显示结果摘要
        print(f"\n✅ 找到 {len(materials)} 条相关素材\n")

        # 显示每个素材
        for idx, material in enumerate(materials, 1):
            if verbose:
                display_material(material, idx)
            else:
                # 简洁模式
                print(f"【{idx}】 {material.get('title', '无标题')}")
                print(f"     类型: {material.get('type')} | 相关度: {material.get('score', 0):.4f}")
                content_preview = material.get('content', '')[:100].replace('\n', ' ')
                print(f"     预览: {content_preview}...\n")

        # 提示
        if not verbose:
            print(f"\n💡 使用 -v 参数查看详细信息")

        print(f"\n{'='*70}\n")

    except Exception as e:
        print(f"\n❌ 查询失败: {e}")
        import traceback
        if verbose:
            traceback.print_exc()


def main():
    parser = argparse.ArgumentParser(
        description='Get 笔记素材查询工具',
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
示例:
  # 基本查询
  python query_materials.py "如何提升工作效率"

  # 指定分类
  python query_materials.py "AI技术趋势" -c 科技

  # 返回更多结果
  python query_materials.py "时间管理" -n 10

  # 查看详细信息
  python query_materials.py "职场沟通技巧" -c 职场 -v
        """
    )

    parser.add_argument('query', help='搜索关键词或主题')
    parser.add_argument('-c', '--category', help='分类（职场/生活/科技等）')
    parser.add_argument('-n', '--limit', type=int, default=5,
                       help='返回结果数量（默认5）')
    parser.add_argument('-v', '--verbose', action='store_true',
                       help='显示详细信息')

    # Get 笔记配置
    parser.add_argument('--api-key', help='Get 笔记 API Key')
    parser.add_argument('--kb-id', help='知识库 ID')

    args = parser.parse_args()

    # 设置环境变量（如果提供了参数）
    if args.api_key:
        os.environ['GETNOTE_API_KEY'] = args.api_key
    if args.kb_id:
        os.environ['GETNOTE_KB_ID'] = args.kb_id

    # 执行搜索
    search_materials(
        query=args.query,
        category=args.category,
        limit=args.limit,
        verbose=args.verbose
    )


if __name__ == "__main__":
    main()