kb-skills/kb_workflow/cli.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
KB Workflow - Knowledge Base Workflow Orchestrator

知识库工作流管道入口。

此模块提供 process_command() 函数作为管道入口，
内部使用 kb_workflow.pipeline 包实现完整的内容处理流程。

使用方式:
    python -m kb_workflow process "https://www.bilibili.com/video/BVxxx"
    python -m kb_workflow process "https://..." -o /path/to/output
    python -m kb_workflow login                    # QR code login
    python -m kb_workflow login --import-cookie cookies.json
    python -m kb_workflow login --status

或者通过包接口调用:
    from kb_workflow import process_command
    result = process_command("https://...")
"""

import argparse
import asyncio
import json
import sys
from pathlib import Path
from typing import Any, Dict, Optional


def process_command(
    url: str,
    output_dir: Optional[str] = None,
    skip_completed: bool = True
) -> Dict[str, Any]:
    """Execute full pipeline for a URL

    Args:
        url: URL to process
        output_dir: Optional output directory for content. Defaults to
            KB_INTERMEDIATE_PATH env var or ~/work/intermediate/analysis/
        skip_completed: Skip already completed steps (default True)

    Returns:
        Dictionary with success status, content_dir, and results
    """
    try:
        from .pipeline.config import get_settings
        from .pipeline.pipeline import Pipeline
        from .pipeline.processors import (
            ASRProcessor, AnalyzeProcessor, ArchiveProcessor,
            DownloadProcessor, ParseProcessor, SRTProcessor
        )
        from .pipeline.schemas import PipelineConfig

        async def run_pipeline():
            # Determine output directory - default to intermediate_path
            settings = get_settings()
            out_dir = Path(output_dir) if output_dir else Path(settings.intermediate_path)

            # Create processors
            processors = [
                ParseProcessor(),
                DownloadProcessor(),
                ASRProcessor(),
                SRTProcessor(),      # B站视频使用
                AnalyzeProcessor(),  # 知乎/文章使用
                ArchiveProcessor(),
            ]

            # Create pipeline config
            config = PipelineConfig(
                steps=["parse", "download", "asr", "srt", "analyze", "archive"],
                output_dir=out_dir,
                skip_completed=skip_completed,
            )

            # Create and run pipeline
            pipeline = Pipeline(
                processors=processors,
                config=config,
                output_dir=out_dir,
            )

            result = await pipeline.run(url)
            return result

        # Run async pipeline
        result = asyncio.run(run_pipeline())

        # Extract archive_file from results if available
        archive_file = None
        if result.results and "archive" in result.results:
            archive_file = result.results["archive"].get("archive_file")

        return {
            "success": result.success,
            "error": result.error_message,
            "content_dir": str(result.content_dir) if result.content_dir else None,
            "archive_file": archive_file,
            "results": result.results,
        }

    except Exception as e:
        return {
            "success": False,
            "error": str(e),
            "content_dir": None,
            "results": {},
        }


def login_command(
    import_cookie: Optional[str] = None,
    status: bool = False
) -> Dict[str, Any]:
    """Execute login command

    Args:
        import_cookie: Path to JSON cookie file
        status: Check current cookie status

    Returns:
        Dictionary with login result
    """
    try:
        from .auth import ZhihuLogin
        from .auth.cookie_manager import CookieManager, parse_json_cookies
        from .pipeline.config import get_settings

        settings = get_settings()
        manager = CookieManager()

        # Check Cookie status
        if status:
            current_cookie = settings.zhihu_cookie.get_secret_value()
            if not current_cookie:
                return {"valid": False, "reason": "No cookie found in ~/.kb_env"}

            async def check_status():
                login = ZhihuLogin(intermediate_path=settings.intermediate_path)
                try:
                    login.browser = login.browser or type('obj', (object,), {
                        'browser': None,
                        'page': None,
                        'start': lambda cookie=None: None,
                        'close': lambda: None
                    })()
                    login.browser.start(cookie=current_cookie)
                    return await login.validate_cookie(current_cookie)
                finally:
                    if login.browser and hasattr(login.browser, 'close'):
                        login.browser.close()

            result = asyncio.run(check_status())
            return result.model_dump()

        # Import Cookie
        if import_cookie:
            cookie_str = parse_json_cookies(import_cookie)
            manager.save_cookie(cookie_str)
            return {"success": True, "message": f"Cookie imported from {import_cookie}"}

        # QR code login (synchronous method)
        login = ZhihuLogin(intermediate_path=settings.intermediate_path)
        result = login.login_by_qrcode()
        if result.success and result.cookie_str:
            manager.save_cookie(result.cookie_str)
        return result.model_dump(mode='json')

    except Exception as e:
        return {"success": False, "error": str(e)}


def main():
    """CLI入口 - 用于命令行调用"""
    parser = argparse.ArgumentParser(
        description='KB Workflow - Knowledge Base Workflow Orchestrator',
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog='''
使用示例:
  # 处理URL (通过完整管道)
  python -m kb_workflow process "https://www.bilibili.com/video/BVxxx"

  # 指定输出目录
  python -m kb_workflow process "https://..." -o /path/to/output

  # 强制重新执行（忽略已完成状态）
  python -m kb_workflow process "https://..." --no-skip

  # 知乎登录 (二维码扫码)
  python -m kb_workflow login

  # 导入 Cookie
  python -m kb_workflow login --import-cookie cookies.json

  # 检查 Cookie 状态
  python -m kb_workflow login --status
        '''
    )

    # Create subparsers for commands
    subparsers = parser.add_subparsers(dest='command', help='Available commands')

    # Process subcommand
    process_parser = subparsers.add_parser('process', help='Process URL through full pipeline')
    process_parser.add_argument('url', type=str, help='URL to process')
    process_parser.add_argument('--output-dir', '-o', type=str, help='Output directory for content')
    process_parser.add_argument('--skip-completed', action='store_true', default=True,
                                help='Skip completed steps (default: True)')
    process_parser.add_argument('--no-skip', action='store_false', dest='skip_completed',
                                help='Re-run all steps')

    # Login subcommand
    login_parser = subparsers.add_parser('login', help='Zhihu login authentication')
    login_parser.add_argument('--import-cookie', '-i', type=str,
                              help='Import cookie from JSON file')
    login_parser.add_argument('--status', action='store_true',
                              help='Check current cookie status')

    args = parser.parse_args()

    # Handle process subcommand
    if args.command == 'process':
        result = process_command(
            url=args.url,
            output_dir=args.output_dir,
            skip_completed=args.skip_completed
        )
        print(json.dumps(result, ensure_ascii=False, indent=2))
    # Handle login subcommand
    elif args.command == 'login':
        result = login_command(
            import_cookie=args.import_cookie,
            status=args.status
        )
        print(json.dumps(result, ensure_ascii=False, indent=2))
    else:
        parser.print_help()


if __name__ == '__main__':
    main()