kb-skills/kb_workflow/kb_workflow.py

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

知识库工作流管道入口。

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

使用方式:
    python kb_workflow.py process "https://www.bilibili.com/video/BVxxx"
    python kb_workflow.py process "https://..." -o /path/to/output
"""

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 kb_workflow.pipeline.config import get_settings
        from kb_workflow.pipeline.pipeline import Pipeline
        from kb_workflow.pipeline.processors import (
            ASRProcessor, AnalyzeProcessor, ArchiveProcessor,
            DownloadProcessor, ParseProcessor
        )
        from kb_workflow.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(),
                AnalyzeProcessor(),
                ArchiveProcessor(),
            ]

            # Create pipeline config
            config = PipelineConfig(
                steps=["parse", "download", "asr", "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 main():
    """CLI入口 - 用于命令行调用"""
    parser = argparse.ArgumentParser(
        description='KB Workflow - Knowledge Base Workflow Orchestrator',
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog='''
使用示例:
  # 处理URL (通过完整管道)
  python kb_workflow.py process "https://www.bilibili.com/video/BVxxx"

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

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

    # 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')

    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))
    else:
        parser.print_help()


if __name__ == '__main__':
    main()