自动同步: 2026-03-18 02:10
This commit is contained in:
@@ -0,0 +1,26 @@
|
||||
"""Processors module for KB Pipeline
|
||||
|
||||
This module contains processors for various content transformation tasks:
|
||||
- BaseProcessor: Abstract base class for all processors
|
||||
- ParseProcessor: URL parsing for various platforms
|
||||
- DownloadProcessor: Metadata and content downloading
|
||||
- ASRProcessor: Audio-to-text conversion (subtitle download + Whisper.cpp)
|
||||
- AnalyzeProcessor: AI-powered content analysis using DeepAnalyzer
|
||||
- ArchiveProcessor: Copy analysis reports to archive directory
|
||||
"""
|
||||
|
||||
from kb_workflow.pipeline.processors.base import BaseProcessor
|
||||
from kb_workflow.pipeline.processors.parse import ParseProcessor
|
||||
from kb_workflow.pipeline.processors.download import DownloadProcessor
|
||||
from kb_workflow.pipeline.processors.asr import ASRProcessor
|
||||
from kb_workflow.pipeline.processors.analyze import AnalyzeProcessor
|
||||
from kb_workflow.pipeline.processors.archive import ArchiveProcessor
|
||||
|
||||
__all__ = [
|
||||
"BaseProcessor",
|
||||
"ParseProcessor",
|
||||
"DownloadProcessor",
|
||||
"ASRProcessor",
|
||||
"AnalyzeProcessor",
|
||||
"ArchiveProcessor",
|
||||
]
|
||||
@@ -0,0 +1,313 @@
|
||||
"""AnalyzeProcessor - AI-powered content analysis
|
||||
|
||||
Unified processor for generating AI analysis reports from content.
|
||||
Wraps existing DeepAnalyzer for integration into the pipeline.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
from pathlib import Path
|
||||
from typing import Optional, Dict, Any
|
||||
|
||||
from kb_workflow.pipeline.config import get_settings
|
||||
from kb_workflow.pipeline.schemas import AnalyzeInput, AnalyzeOutput
|
||||
from kb_workflow.pipeline.state import WorkflowStateManager
|
||||
from kb_workflow.pipeline.processors.base import BaseProcessor
|
||||
from kb_workflow.deep_analyzer import DeepAnalyzer, SimpleAIClient, generate_frontmatter
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class AnalyzeProcessor(BaseProcessor):
|
||||
"""Analyze 处理器 - AI 深度内容分析
|
||||
|
||||
处理流程：
|
||||
1. 检查步骤是否已完成（幂等执行）
|
||||
2. 加载元数据和内容文件
|
||||
3. 初始化 AI 客户端
|
||||
4. 调用 DeepAnalyzer 生成分析报告
|
||||
5. 生成 frontmatter 并保存 Markdown 文件
|
||||
6. 记录处理状态到 workflow_state.json
|
||||
|
||||
Attributes:
|
||||
STEP_NAME: 步骤名称，用于状态记录
|
||||
settings: KB 配置
|
||||
state_manager: 状态管理器（可选）
|
||||
"""
|
||||
|
||||
STEP_NAME = "analyze"
|
||||
|
||||
async def process(self, input_data: AnalyzeInput) -> AnalyzeOutput:
|
||||
"""执行分析处理
|
||||
|
||||
策略：
|
||||
1. 检查步骤是否已完成（跳过已完成）
|
||||
2. 加载元数据和内容
|
||||
3. 调用 DeepAnalyzer
|
||||
4. 保存分析报告
|
||||
5. 记录状态
|
||||
|
||||
Args:
|
||||
input_data: Analyze 输入数据
|
||||
|
||||
Returns:
|
||||
Analyze 处理结果
|
||||
|
||||
Raises:
|
||||
Exception: 处理失败时抛出异常
|
||||
"""
|
||||
try:
|
||||
# Step 1: 检查是否已完成（幂等执行）
|
||||
if self.state_manager:
|
||||
state = self.state_manager.load()
|
||||
if state and state.is_step_completed(self.STEP_NAME):
|
||||
logger.info(f"Step {self.STEP_NAME} already completed, skipping")
|
||||
# 返回已有结果
|
||||
existing_file = Path(input_data.output_file)
|
||||
if existing_file.exists():
|
||||
return AnalyzeOutput(
|
||||
success=True,
|
||||
analysis_file=str(existing_file),
|
||||
summary="Skipped - already completed",
|
||||
main_topics=[],
|
||||
key_insights=[],
|
||||
tags=[],
|
||||
stats={"skipped": True}
|
||||
)
|
||||
|
||||
# Step 2: 加载元数据和内容
|
||||
metadata = self._load_json(input_data.metadata_file)
|
||||
content = self._load_json(input_data.content_file)
|
||||
|
||||
# 处理 Whisper 字幕格式 - 提取完整文本
|
||||
if content and 'transcription' in content:
|
||||
# Whisper JSON 格式，提取文本
|
||||
text_parts = []
|
||||
for item in content.get('transcription', []):
|
||||
text_parts.append(item.get('text', ''))
|
||||
content = {'text': ' '.join(text_parts)}
|
||||
logger.info(f"Extracted {len(content['text'])} chars from Whisper transcription")
|
||||
|
||||
# 对于文章类型，内容可能是纯文本文件
|
||||
if not content and input_data.content_type in ("article", "answer", "zhuanlan"):
|
||||
content_path = Path(input_data.content_file)
|
||||
if content_path.exists():
|
||||
text = content_path.read_text(encoding="utf-8")
|
||||
# 去除 YAML frontmatter（如果存在）
|
||||
text = self._strip_frontmatter(text)
|
||||
content = {"text": text}
|
||||
|
||||
# Step 3: 调用 DeepAnalyzer
|
||||
result = await self._run_analysis(
|
||||
metadata=metadata,
|
||||
content=content,
|
||||
content_type=input_data.content_type
|
||||
)
|
||||
|
||||
# Step 4: 生成并保存分析报告
|
||||
output_path = Path(input_data.output_file)
|
||||
output_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
self._write_analysis_report(
|
||||
output_path=output_path,
|
||||
metadata=metadata,
|
||||
result=result,
|
||||
content_type=input_data.content_type
|
||||
)
|
||||
|
||||
# Step 5: 记录成功状态
|
||||
if self.state_manager:
|
||||
state = self.state_manager.load()
|
||||
if state:
|
||||
summary = f"topics={len(result.main_topics)}, insights={len(result.key_insights)}"
|
||||
state.mark_step_completed(self.STEP_NAME, summary)
|
||||
self.state_manager.save(state)
|
||||
|
||||
return AnalyzeOutput(
|
||||
success=True,
|
||||
analysis_file=str(output_path),
|
||||
summary=result.summary,
|
||||
main_topics=result.main_topics,
|
||||
key_insights=result.key_insights,
|
||||
tags=result.recommended_tags,
|
||||
stats={
|
||||
"topics_count": len(result.main_topics),
|
||||
"insights_count": len(result.key_insights),
|
||||
"tags_count": len(result.recommended_tags)
|
||||
}
|
||||
)
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Analyze failed: {e}")
|
||||
|
||||
# 记录失败状态
|
||||
if self.state_manager:
|
||||
state = self.state_manager.load()
|
||||
if state:
|
||||
state.mark_step_failed(self.STEP_NAME, str(e))
|
||||
self.state_manager.save(state)
|
||||
|
||||
return AnalyzeOutput(
|
||||
success=False,
|
||||
error_message=str(e),
|
||||
analysis_file="",
|
||||
summary="",
|
||||
main_topics=[],
|
||||
key_insights=[],
|
||||
tags=[],
|
||||
stats={}
|
||||
)
|
||||
|
||||
def _load_json(self, file_path: str) -> Dict[str, Any]:
|
||||
"""加载 JSON 文件
|
||||
|
||||
Args:
|
||||
file_path: JSON 文件路径
|
||||
|
||||
Returns:
|
||||
解析后的字典，如果文件不存在或解析失败返回空字典
|
||||
"""
|
||||
if not file_path:
|
||||
return {}
|
||||
|
||||
path = Path(file_path)
|
||||
if not path.exists():
|
||||
logger.warning(f"File not found: {file_path}")
|
||||
return {}
|
||||
|
||||
try:
|
||||
with open(path, 'r', encoding='utf-8') as f:
|
||||
return json.load(f)
|
||||
except json.JSONDecodeError as e:
|
||||
logger.error(f"Failed to parse JSON: {file_path}, error: {e}")
|
||||
return {}
|
||||
|
||||
async def _run_analysis(self, metadata: Dict, content: Dict, content_type: str):
|
||||
"""调用 DeepAnalyzer 执行分析
|
||||
|
||||
Args:
|
||||
metadata: 元数据字典
|
||||
content: 内容字典
|
||||
content_type: 内容类型
|
||||
|
||||
Returns:
|
||||
AnalysisResult 对象
|
||||
"""
|
||||
# 初始化 AI 客户端
|
||||
ai_client = self._create_ai_client()
|
||||
|
||||
# 创建分析器实例
|
||||
analyzer = DeepAnalyzer(ai_client=ai_client)
|
||||
|
||||
# 执行分析
|
||||
return analyzer.analyze(metadata, content, content_type)
|
||||
|
||||
def _create_ai_client(self):
|
||||
"""创建 AI 客户端
|
||||
|
||||
Returns:
|
||||
SimpleAIClient 实例
|
||||
"""
|
||||
# 从环境变量获取配置
|
||||
api_key = os.environ.get("AI_API_KEY", "")
|
||||
provider = os.environ.get("AI_PROVIDER", "")
|
||||
model = os.environ.get("AI_MODEL", "glm-5")
|
||||
base_url = os.environ.get("AI_BASE_URL", "")
|
||||
|
||||
return SimpleAIClient(
|
||||
api_key=api_key,
|
||||
provider=provider,
|
||||
model=model,
|
||||
base_url=base_url
|
||||
)
|
||||
|
||||
def _normalize_metadata(self, metadata: Dict) -> Dict:
|
||||
"""标准化元数据结构（扁平化嵌套结构）
|
||||
|
||||
Args:
|
||||
metadata: 原始元数据字典，可能是嵌套结构
|
||||
|
||||
Returns:
|
||||
扁平化的元数据字典
|
||||
"""
|
||||
if not metadata:
|
||||
return {}
|
||||
|
||||
# 如果已经是扁平结构（顶层有 title）
|
||||
if 'title' in metadata:
|
||||
return metadata
|
||||
|
||||
# 处理嵌套结构 {metadata: {...}, stats: {...}}
|
||||
if 'metadata' in metadata:
|
||||
return metadata['metadata'].copy()
|
||||
|
||||
return metadata
|
||||
|
||||
def _strip_frontmatter(self, text: str) -> str:
|
||||
"""去除 Markdown 文件中的 YAML frontmatter
|
||||
|
||||
Args:
|
||||
text: 可能包含 frontmatter 的文本
|
||||
|
||||
Returns:
|
||||
去除 frontmatter 后的文本
|
||||
"""
|
||||
if not text:
|
||||
return text
|
||||
|
||||
# 检查是否以 --- 开头
|
||||
if not text.startswith('---'):
|
||||
return text
|
||||
|
||||
# 找到第二个 ---
|
||||
lines = text.split('\n')
|
||||
end_index = -1
|
||||
for i in range(1, len(lines)):
|
||||
if lines[i].strip() == '---':
|
||||
end_index = i
|
||||
break
|
||||
|
||||
if end_index == -1:
|
||||
return text
|
||||
|
||||
# 返回 frontmatter 之后的内容（跳过空行）
|
||||
remaining = '\n'.join(lines[end_index + 1:])
|
||||
return remaining.lstrip('\n')
|
||||
|
||||
def _write_analysis_report(
|
||||
self,
|
||||
output_path: Path,
|
||||
metadata: Dict,
|
||||
result,
|
||||
content_type: str
|
||||
) -> None:
|
||||
"""写入分析报告 Markdown 文件
|
||||
|
||||
Args:
|
||||
output_path: 输出文件路径
|
||||
metadata: 元数据字典
|
||||
result: AnalysisResult 对象
|
||||
content_type: 内容类型
|
||||
"""
|
||||
# 标准化元数据（扁平化嵌套结构）
|
||||
normalized_metadata = self._normalize_metadata(metadata)
|
||||
|
||||
# 生成 frontmatter
|
||||
frontmatter = generate_frontmatter(normalized_metadata, result, content_type)
|
||||
|
||||
# 获取完整分析内容
|
||||
full_analysis = result.full_analysis
|
||||
|
||||
# 文章/回答类型：追加原文
|
||||
if content_type in ("article", "answer", "zhuanlan") and result.original_text:
|
||||
full_analysis = full_analysis.rstrip() + "\n\n## 原文\n\n" + result.original_text
|
||||
|
||||
# 写入文件
|
||||
with open(output_path, 'w', encoding='utf-8') as f:
|
||||
f.write(frontmatter + "\n\n")
|
||||
f.write(full_analysis)
|
||||
|
||||
logger.info(f"Analysis report saved to: {output_path}")
|
||||
@@ -0,0 +1,170 @@
|
||||
"""ArchiveProcessor - Copy analysis reports to archive directory
|
||||
|
||||
Final step in pipeline that copies the generated analysis.md
|
||||
to the archive directory organized by platform.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
import shutil
|
||||
from datetime import datetime
|
||||
from pathlib import Path
|
||||
from typing import Optional
|
||||
|
||||
from kb_workflow.pipeline.config import get_settings
|
||||
from kb_workflow.pipeline.schemas import ArchiveInput, ArchiveOutput
|
||||
from kb_workflow.pipeline.state import WorkflowStateManager
|
||||
from kb_workflow.pipeline.processors.base import BaseProcessor
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class ArchiveProcessor(BaseProcessor):
|
||||
"""Archive 处理器 - 归档分析报告
|
||||
|
||||
处理流程：
|
||||
1. 检查步骤是否已完成（幂等执行）
|
||||
2. 检查归档目录是否存在相同 content_id 的文件
|
||||
3. 如果存在重复，询问用户是否覆盖
|
||||
4. 复制 md 文件到归档目录
|
||||
5. 记录处理状态到 workflow_state.json
|
||||
|
||||
归档目录结构：
|
||||
{KB_ARCHIVE_PATH}/{platform}/{YYYYMMDD}_{content_id}_analysis.md
|
||||
"""
|
||||
|
||||
STEP_NAME = "archive"
|
||||
|
||||
async def process(self, input_data: ArchiveInput) -> ArchiveOutput:
|
||||
"""执行归档处理
|
||||
|
||||
Args:
|
||||
input_data: Archive 输入数据
|
||||
|
||||
Returns:
|
||||
Archive 处理结果
|
||||
"""
|
||||
try:
|
||||
# Step 1: 检查是否已完成（幂等执行）
|
||||
if self.state_manager:
|
||||
state = self.state_manager.load()
|
||||
if state and state.is_step_completed(self.STEP_NAME):
|
||||
logger.info(f"Step {self.STEP_NAME} already completed, skipping")
|
||||
return ArchiveOutput(
|
||||
success=True,
|
||||
skipped=True,
|
||||
platform=input_data.platform,
|
||||
archive_file=""
|
||||
)
|
||||
|
||||
# Step 2: 准备归档路径
|
||||
archive_base = Path(self.settings.archive_path)
|
||||
platform_dir = archive_base / input_data.platform
|
||||
platform_dir.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
# 文件命名: {YYYYMMDD}_{content_id}_analysis.md
|
||||
timestamp = datetime.now().strftime("%Y%m%d")
|
||||
archive_filename = f"{timestamp}_{input_data.content_id}_analysis.md"
|
||||
archive_path = platform_dir / archive_filename
|
||||
|
||||
# Step 3: 检查重复
|
||||
existing_file = self._find_existing_archive(
|
||||
platform_dir, input_data.content_id
|
||||
)
|
||||
if existing_file:
|
||||
should_overwrite = self._ask_overwrite(existing_file)
|
||||
if not should_overwrite:
|
||||
logger.info(f"User chose not to overwrite, skipping archive")
|
||||
return ArchiveOutput(
|
||||
success=True,
|
||||
skipped=True,
|
||||
platform=input_data.platform,
|
||||
archive_file=str(existing_file)
|
||||
)
|
||||
# 删除旧文件
|
||||
existing_file.unlink()
|
||||
logger.info(f"Removed existing archive: {existing_file}")
|
||||
|
||||
# Step 4: 复制文件
|
||||
source_path = Path(input_data.analysis_file)
|
||||
if not source_path.exists():
|
||||
raise FileNotFoundError(f"Analysis file not found: {source_path}")
|
||||
|
||||
shutil.copy2(source_path, archive_path)
|
||||
logger.info(f"Archived: {source_path} -> {archive_path}")
|
||||
|
||||
# Step 5: 记录成功状态
|
||||
if self.state_manager:
|
||||
state = self.state_manager.load()
|
||||
if state:
|
||||
state.mark_step_completed(
|
||||
self.STEP_NAME,
|
||||
f"archived to {archive_path}"
|
||||
)
|
||||
self.state_manager.save(state)
|
||||
|
||||
return ArchiveOutput(
|
||||
success=True,
|
||||
skipped=False,
|
||||
platform=input_data.platform,
|
||||
archive_file=str(archive_path)
|
||||
)
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Archive failed: {e}")
|
||||
|
||||
if self.state_manager:
|
||||
state = self.state_manager.load()
|
||||
if state:
|
||||
state.mark_step_failed(self.STEP_NAME, str(e))
|
||||
self.state_manager.save(state)
|
||||
|
||||
return ArchiveOutput(
|
||||
success=False,
|
||||
error_message=str(e),
|
||||
platform=input_data.platform,
|
||||
archive_file=""
|
||||
)
|
||||
|
||||
def _find_existing_archive(self, platform_dir: Path, content_id: str) -> Optional[Path]:
|
||||
"""查找已存在的归档文件
|
||||
|
||||
使用 glob 模式 `*_{content_id}_analysis.md` 搜索
|
||||
|
||||
Args:
|
||||
platform_dir: 平台目录路径
|
||||
content_id: 内容 ID
|
||||
|
||||
Returns:
|
||||
已存在的文件路径，如果不存在返回 None
|
||||
"""
|
||||
pattern = f"*_{content_id}_analysis.md"
|
||||
matches = list(platform_dir.glob(pattern))
|
||||
return matches[0] if matches else None
|
||||
|
||||
def _ask_overwrite(self, existing_file: Path) -> bool:
|
||||
"""询问用户是否覆盖已存在的文件
|
||||
|
||||
使用 rich.prompt.Confirm 模式（与 Pipeline._ask_rerun_step 相同）
|
||||
|
||||
Args:
|
||||
existing_file: 已存在的文件路径
|
||||
|
||||
Returns:
|
||||
True 如果用户选择覆盖
|
||||
"""
|
||||
try:
|
||||
from rich.console import Console
|
||||
from rich.prompt import Confirm
|
||||
|
||||
console = Console()
|
||||
return Confirm.ask(
|
||||
f"Archive already exists: {existing_file.name}\nOverwrite?",
|
||||
default=False
|
||||
)
|
||||
except ImportError:
|
||||
response = input(
|
||||
f"Archive already exists: {existing_file.name}\nOverwrite? [y/N] "
|
||||
)
|
||||
return response.lower() in ('y', 'yes')
|
||||
@@ -0,0 +1,685 @@
|
||||
"""ASRProcessor - Audio-to-text conversion
|
||||
|
||||
Unified processor for B站官方字幕下载 and Whisper.cpp 语音识别.
|
||||
Follows the pipeline pattern established in Phase 1/2.
|
||||
|
||||
Strategy:
|
||||
1. Try official subtitle first (highest quality, no compute cost)
|
||||
2. Fall back to Whisper.cpp if no subtitle available
|
||||
3. Auto-enable speaker diarization if HF_TOKEN is configured
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
import subprocess
|
||||
from pathlib import Path
|
||||
from typing import Optional, Dict, Any, List
|
||||
|
||||
import aiohttp
|
||||
|
||||
from kb_workflow.pipeline.config import get_settings
|
||||
from kb_workflow.pipeline.schemas import ASRInput, ASROutput
|
||||
from kb_workflow.pipeline.state import WorkflowStateManager
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class WhisperNotFoundError(RuntimeError):
|
||||
"""Whisper.cpp 路径未找到"""
|
||||
pass
|
||||
|
||||
|
||||
def get_whisper_cpp_path() -> str:
|
||||
"""获取 whisper.cpp 安装路径
|
||||
|
||||
检查环境变量 WHISPER_CPP_PATH 或默认路径 ~/whisper.cpp
|
||||
|
||||
Returns:
|
||||
whisper.cpp 目录路径
|
||||
|
||||
Raises:
|
||||
WhisperNotFoundError: 如果路径不存在
|
||||
"""
|
||||
# 检查环境变量
|
||||
path = os.environ.get('WHISPER_CPP_PATH')
|
||||
if path and os.path.exists(path):
|
||||
return os.path.expanduser(path)
|
||||
|
||||
# 检查默认路径
|
||||
default = os.path.expanduser('~/whisper.cpp')
|
||||
if os.path.exists(default):
|
||||
return default
|
||||
|
||||
raise WhisperNotFoundError(
|
||||
"未找到 whisper.cpp。请设置 WHISPER_CPP_PATH 环境变量或安装到 ~/whisper.cpp"
|
||||
)
|
||||
|
||||
|
||||
def convert_to_wav(input_path: str, output_path: Optional[str] = None) -> str:
|
||||
"""使用 ffmpeg 将音频/视频转换为 16kHz WAV 格式
|
||||
|
||||
Args:
|
||||
input_path: 输入文件路径
|
||||
output_path: 输出文件路径（可选，默认为 input_path_audio.wav）
|
||||
|
||||
Returns:
|
||||
输出 WAV 文件路径
|
||||
"""
|
||||
if output_path is None:
|
||||
base = os.path.splitext(input_path)[0]
|
||||
output_path = f"{base}_audio.wav"
|
||||
|
||||
cmd = [
|
||||
'ffmpeg', '-y', '-i', input_path,
|
||||
'-ar', '16000', # 16kHz 采样率
|
||||
'-ac', '1', # 单声道
|
||||
'-c:a', 'pcm_s16le', # 16-bit PCM
|
||||
output_path
|
||||
]
|
||||
|
||||
subprocess.run(cmd, capture_output=True, check=True)
|
||||
return output_path
|
||||
|
||||
|
||||
class ASRProcessor:
|
||||
"""ASR 处理器 - 统一字幕下载和语音识别
|
||||
|
||||
处理流程：
|
||||
1. 如果是 B站视频，尝试下载官方字幕
|
||||
2. 无官方字幕则调用 Whisper.cpp
|
||||
3. HF_TOKEN 配置时自动启用说话者分离
|
||||
4. 记录处理状态到 workflow_state.json
|
||||
|
||||
Attributes:
|
||||
STEP_NAME: 步骤名称，用于状态记录
|
||||
settings: KB 配置
|
||||
state_manager: 状态管理器（可选）
|
||||
"""
|
||||
|
||||
STEP_NAME = "asr"
|
||||
|
||||
def __init__(self, state_manager: Optional[WorkflowStateManager] = None):
|
||||
"""初始化 ASR 处理器
|
||||
|
||||
Args:
|
||||
state_manager: 状态管理器（可选，用于记录处理状态）
|
||||
"""
|
||||
self.settings = get_settings()
|
||||
self.state_manager = state_manager
|
||||
|
||||
async def process(self, input_data: ASRInput) -> ASROutput:
|
||||
"""执行 ASR 处理
|
||||
|
||||
策略：
|
||||
0. 先检查是否已有字幕文件（幂等性）
|
||||
1. 先尝试官方字幕（仅 B站）
|
||||
2. 无官方字幕则调用 Whisper
|
||||
3. 记录状态
|
||||
|
||||
Args:
|
||||
input_data: ASR 输入数据
|
||||
|
||||
Returns:
|
||||
ASR 处理结果
|
||||
|
||||
Raises:
|
||||
Exception: 处理失败时抛出异常
|
||||
"""
|
||||
result = None
|
||||
|
||||
try:
|
||||
# Step 0: 检查是否已有字幕文件（幂等性）
|
||||
existing_result = self._check_existing_subtitle_files(input_data)
|
||||
if existing_result:
|
||||
logger.info(f"Found existing subtitle files, skipping processing")
|
||||
result = existing_result
|
||||
else:
|
||||
# Step 1: 尝试官方字幕（仅 B站）
|
||||
if input_data.platform == "bilibili":
|
||||
result = await self._try_official_subtitle(input_data)
|
||||
|
||||
# Step 2: 无官方字幕则调用 Whisper
|
||||
if result is None:
|
||||
result = await self._run_whisper(input_data)
|
||||
|
||||
# Step 3: 记录成功状态
|
||||
if self.state_manager:
|
||||
state = self.state_manager.load()
|
||||
if state:
|
||||
summary = f"tool={result.tool}, segments={len(result.segments)}"
|
||||
if result.speakers:
|
||||
summary += f", speakers={len(result.speakers)}"
|
||||
state.mark_step_completed(self.STEP_NAME, summary)
|
||||
self.state_manager.save(state)
|
||||
|
||||
return result
|
||||
|
||||
except Exception as e:
|
||||
# 记录失败状态（只在所有重试都失败后）
|
||||
if self.state_manager:
|
||||
state = self.state_manager.load()
|
||||
if state:
|
||||
state.mark_step_failed(self.STEP_NAME, str(e))
|
||||
self.state_manager.save(state)
|
||||
raise
|
||||
|
||||
def _check_existing_subtitle_files(self, input_data: ASRInput) -> Optional[ASROutput]:
|
||||
"""检查是否已存在字幕文件（幂等性检查）
|
||||
|
||||
Args:
|
||||
input_data: ASR 输入数据
|
||||
|
||||
Returns:
|
||||
如果找到现有字幕文件返回 ASROutput，否则返回 None
|
||||
"""
|
||||
output_dir = Path(input_data.content_dir)
|
||||
|
||||
# 尝试两种命名格式：新格式（带 _subtitle 后缀）和旧格式
|
||||
base_names = [
|
||||
f"{input_data.platform}_{input_data.content_id}_subtitle",
|
||||
f"{input_data.platform}_{input_data.content_id}",
|
||||
]
|
||||
|
||||
for base_name in base_names:
|
||||
json_path = output_dir / f"{base_name}.json"
|
||||
srt_path = output_dir / f"{base_name}.srt"
|
||||
txt_path = output_dir / f"{base_name}.txt"
|
||||
|
||||
if json_path.exists():
|
||||
# 读取现有字幕数据
|
||||
segments = []
|
||||
text = ""
|
||||
duration = 0.0
|
||||
|
||||
try:
|
||||
with open(json_path, 'r', encoding='utf-8') as f:
|
||||
data = json.load(f)
|
||||
|
||||
# 解析字幕格式
|
||||
if 'transcription' in data:
|
||||
# Whisper 格式
|
||||
for i, item in enumerate(data['transcription']):
|
||||
offsets = item.get('offsets', {})
|
||||
if isinstance(offsets, dict):
|
||||
# 新格式：{"from": 0, "to": 5320}
|
||||
start = offsets.get('from', 0) / 1000.0
|
||||
end = offsets.get('to', 0) / 1000.0
|
||||
else:
|
||||
# 旧格式：[0, 5320]
|
||||
start = offsets[0] / 1000.0 if len(offsets) > 0 else 0
|
||||
end = offsets[1] / 1000.0 if len(offsets) > 1 else 0
|
||||
segments.append({
|
||||
'id': i,
|
||||
'start': start,
|
||||
'end': end,
|
||||
'text': item.get('text', '')
|
||||
})
|
||||
elif 'body' in data:
|
||||
# B站官方字幕格式
|
||||
for i, item in enumerate(data['body']):
|
||||
segments.append({
|
||||
'id': i,
|
||||
'start': item.get('from', 0),
|
||||
'end': item.get('to', 0),
|
||||
'text': item.get('content', '')
|
||||
})
|
||||
|
||||
text = ' '.join(s['text'] for s in segments)
|
||||
if segments:
|
||||
duration = segments[-1]['end']
|
||||
|
||||
output_files = {'json': str(json_path)}
|
||||
if srt_path.exists():
|
||||
output_files['srt'] = str(srt_path)
|
||||
if txt_path.exists():
|
||||
output_files['txt'] = str(txt_path)
|
||||
|
||||
return ASROutput(
|
||||
success=True,
|
||||
tool="existing",
|
||||
has_subtitle=True,
|
||||
output_files=output_files,
|
||||
segments=segments,
|
||||
text=text,
|
||||
duration_seconds=duration
|
||||
)
|
||||
except Exception as e:
|
||||
logger.warning(f"Failed to read existing subtitle file: {e}")
|
||||
continue
|
||||
|
||||
return None
|
||||
|
||||
async def _try_official_subtitle(self, input_data: ASRInput) -> Optional[ASROutput]:
|
||||
"""尝试下载 B站官方字幕
|
||||
|
||||
Args:
|
||||
input_data: ASR 输入数据
|
||||
|
||||
Returns:
|
||||
如果找到字幕返回 ASROutput，否则返回 None
|
||||
|
||||
Note:
|
||||
无字幕返回 None（不是抛异常），让流程继续到 Whisper
|
||||
"""
|
||||
# 使用 content_dir 作为输出目录（而不是从 video_path 推断）
|
||||
output_dir = Path(input_data.content_dir)
|
||||
# 使用 content_id 作为基础文件名
|
||||
base_name = f"{input_data.platform}_{input_data.content_id}_subtitle"
|
||||
|
||||
# 输出文件路径
|
||||
txt_path = output_dir / f"{base_name}.txt"
|
||||
srt_path = output_dir / f"{base_name}.srt"
|
||||
json_path = output_dir / f"{base_name}.json"
|
||||
|
||||
async with aiohttp.ClientSession() as session:
|
||||
try:
|
||||
# Step 1: 获取 CID
|
||||
bvid = input_data.content_id
|
||||
pagelist_url = f"https://api.bilibili.com/x/player/pagelist?bvid={bvid}"
|
||||
async with session.get(pagelist_url) as resp:
|
||||
pagelist_data = await resp.json()
|
||||
|
||||
pages = pagelist_data.get('data', [])
|
||||
if not pages:
|
||||
logger.info(f"No pagelist data for {bvid}")
|
||||
return None
|
||||
|
||||
cid = pages[0].get('cid')
|
||||
if not cid:
|
||||
logger.info(f"No cid found for {bvid}")
|
||||
return None
|
||||
|
||||
# Step 2: 获取字幕列表
|
||||
player_url = f"https://api.bilibili.com/x/player/v2?bvid={bvid}&cid={cid}"
|
||||
async with session.get(player_url) as resp:
|
||||
player_data = await resp.json()
|
||||
|
||||
subtitle_info = player_data.get('data', {}).get('subtitle', {})
|
||||
subtitles = subtitle_info.get('list', [])
|
||||
|
||||
if not subtitles:
|
||||
logger.info(f"No official subtitle for {bvid}")
|
||||
return None
|
||||
|
||||
# Step 3: 下载第一个字幕（优先 AI 生成的中文字幕）
|
||||
subtitle_url = None
|
||||
for sub in subtitles:
|
||||
if sub.get('lan', '').startswith('zh'):
|
||||
subtitle_url = sub.get('subtitle_url')
|
||||
break
|
||||
|
||||
if not subtitle_url and subtitles:
|
||||
subtitle_url = subtitles[0].get('subtitle_url')
|
||||
|
||||
if not subtitle_url:
|
||||
return None
|
||||
|
||||
# 补全 URL
|
||||
if subtitle_url.startswith('//'):
|
||||
subtitle_url = f"https:{subtitle_url}"
|
||||
|
||||
async with session.get(subtitle_url) as resp:
|
||||
subtitle_content = await resp.json()
|
||||
|
||||
# Step 4: 转换为输出格式
|
||||
segments = []
|
||||
text_parts = []
|
||||
body = subtitle_content.get('body', [])
|
||||
|
||||
for i, item in enumerate(body):
|
||||
start = item.get('from', 0)
|
||||
end = item.get('to', 0)
|
||||
content = item.get('content', '')
|
||||
|
||||
segments.append({
|
||||
'id': i,
|
||||
'start': start,
|
||||
'end': end,
|
||||
'text': content
|
||||
})
|
||||
text_parts.append(content)
|
||||
|
||||
full_text = ' '.join(text_parts)
|
||||
duration = segments[-1]['end'] if segments else 0
|
||||
|
||||
# Step 5: 写入文件
|
||||
# TXT 格式
|
||||
txt_path.write_text(full_text, encoding='utf-8')
|
||||
|
||||
# SRT 格式
|
||||
srt_content = self._convert_to_srt(segments)
|
||||
srt_path.write_text(srt_content, encoding='utf-8')
|
||||
|
||||
# JSON 格式
|
||||
json_data = {
|
||||
'success': True,
|
||||
'tool': 'subtitle',
|
||||
'segments': segments,
|
||||
'text': full_text,
|
||||
'duration_seconds': duration
|
||||
}
|
||||
json_path.write_text(json.dumps(json_data, ensure_ascii=False, indent=2), encoding='utf-8')
|
||||
|
||||
return ASROutput(
|
||||
success=True,
|
||||
tool="subtitle",
|
||||
has_subtitle=True,
|
||||
output_files={
|
||||
'txt': str(txt_path),
|
||||
'srt': str(srt_path),
|
||||
'json': str(json_path)
|
||||
},
|
||||
segments=segments,
|
||||
text=full_text,
|
||||
duration_seconds=duration
|
||||
)
|
||||
|
||||
except aiohttp.ClientError as e:
|
||||
logger.warning(f"Failed to fetch subtitle: {e}")
|
||||
return None
|
||||
|
||||
def _convert_to_srt(self, segments: List[Dict]) -> str:
|
||||
"""将字幕分段转换为 SRT 格式
|
||||
|
||||
Args:
|
||||
segments: 字幕分段列表，每个包含 start, end, text
|
||||
|
||||
Returns:
|
||||
SRT 格式字符串
|
||||
"""
|
||||
lines = []
|
||||
for i, seg in enumerate(segments, 1):
|
||||
start_time = self._seconds_to_srt_time(seg['start'])
|
||||
end_time = self._seconds_to_srt_time(seg['end'])
|
||||
lines.append(f"{i}")
|
||||
lines.append(f"{start_time} --> {end_time}")
|
||||
lines.append(seg['text'])
|
||||
lines.append("") # 空行
|
||||
|
||||
return '\n'.join(lines)
|
||||
|
||||
def _seconds_to_srt_time(self, seconds: float) -> str:
|
||||
"""将秒数转换为 SRT 时间格式 (HH:MM:SS,mmm)"""
|
||||
hours = int(seconds // 3600)
|
||||
minutes = int((seconds % 3600) // 60)
|
||||
secs = int(seconds % 60)
|
||||
millis = int((seconds % 1) * 1000)
|
||||
return f"{hours:02d}:{minutes:02d}:{secs:02d},{millis:03d}"
|
||||
|
||||
async def _run_whisper(self, input_data: ASRInput) -> ASROutput:
|
||||
"""执行 Whisper 语音识别
|
||||
|
||||
根据 HF_TOKEN 配置自动选择普通模式或说话者分离模式。
|
||||
|
||||
Args:
|
||||
input_data: ASR 输入数据
|
||||
|
||||
Returns:
|
||||
ASR 处理结果
|
||||
"""
|
||||
hf_token = self.settings.hf_token.get_secret_value()
|
||||
if hf_token:
|
||||
return await self._run_whisper_with_diarization(input_data)
|
||||
else:
|
||||
return await self._run_whisper_basic(input_data)
|
||||
|
||||
async def _run_whisper_basic(self, input_data: ASRInput) -> ASROutput:
|
||||
"""执行普通 Whisper.cpp 语音识别
|
||||
|
||||
Args:
|
||||
input_data: ASR 输入数据
|
||||
|
||||
Returns:
|
||||
ASR 处理结果
|
||||
|
||||
Raises:
|
||||
RuntimeError: Whisper.cpp 路径不存在或执行失败，或 video_path 为空
|
||||
"""
|
||||
# 检查 video_path 是否有效
|
||||
if not input_data.video_path:
|
||||
raise RuntimeError(
|
||||
"video_path 为空，无法执行 Whisper 语音识别。"
|
||||
"请确保视频文件已下载（对于没有官方字幕的视频，需要先下载视频文件）。"
|
||||
)
|
||||
|
||||
video_path = Path(input_data.video_path)
|
||||
output_dir = Path(input_data.content_dir)
|
||||
base_name = f"{input_data.platform}_{input_data.content_id}_whisper"
|
||||
|
||||
# 获取 Whisper.cpp 路径
|
||||
try:
|
||||
whisper_path = get_whisper_cpp_path()
|
||||
except WhisperNotFoundError as e:
|
||||
raise RuntimeError(str(e))
|
||||
|
||||
# 模型路径
|
||||
model_path = os.path.join(whisper_path, 'models', 'ggml-base.bin')
|
||||
if not os.path.exists(model_path):
|
||||
raise RuntimeError(f"Whisper model not found: {model_path}")
|
||||
|
||||
# 转换音频格式
|
||||
audio_path = convert_to_wav(str(video_path))
|
||||
temp_audio_created = audio_path != str(video_path)
|
||||
|
||||
try:
|
||||
# 构建 Whisper 命令
|
||||
output_base = os.path.join(output_dir, base_name)
|
||||
cmd = [
|
||||
os.path.join(whisper_path, 'whisper-cli'),
|
||||
'-m', model_path,
|
||||
'-f', audio_path,
|
||||
'-of', output_base,
|
||||
'-l', 'zh',
|
||||
'--output-txt', '--output-srt', '--output-json',
|
||||
]
|
||||
|
||||
# 执行 Whisper
|
||||
result = subprocess.run(cmd, capture_output=True, text=True, timeout=3600)
|
||||
|
||||
if result.returncode != 0:
|
||||
raise RuntimeError(f"Whisper failed: {result.stderr}")
|
||||
|
||||
# 解析输出
|
||||
return self._parse_whisper_output(output_dir, base_name)
|
||||
|
||||
finally:
|
||||
# 清理临时音频文件
|
||||
if temp_audio_created and os.path.exists(audio_path):
|
||||
try:
|
||||
os.remove(audio_path)
|
||||
except Exception:
|
||||
pass # 静默失败，不阻塞主流程
|
||||
|
||||
def _parse_whisper_output(self, output_dir: Path, base_name: str) -> ASROutput:
|
||||
"""解析 Whisper 输出文件
|
||||
|
||||
Args:
|
||||
output_dir: 输出目录
|
||||
base_name: 文件基础名
|
||||
|
||||
Returns:
|
||||
ASR 处理结果
|
||||
"""
|
||||
txt_path = output_dir / f"{base_name}.txt"
|
||||
srt_path = output_dir / f"{base_name}.srt"
|
||||
json_path = output_dir / f"{base_name}.json"
|
||||
|
||||
# 读取 JSON 输出
|
||||
segments = []
|
||||
text = ""
|
||||
duration = 0.0
|
||||
|
||||
if json_path.exists():
|
||||
with open(json_path, 'r', encoding='utf-8') as f:
|
||||
data = json.load(f)
|
||||
# Whisper.cpp JSON 格式
|
||||
if 'transcription' in data:
|
||||
for i, item in enumerate(data['transcription']):
|
||||
offsets = item.get('offsets', {})
|
||||
if isinstance(offsets, dict):
|
||||
# 新格式：{"from": 0, "to": 5320}
|
||||
start = offsets.get('from', 0) / 1000.0
|
||||
end = offsets.get('to', 0) / 1000.0
|
||||
else:
|
||||
# 旧格式：[0, 5320]
|
||||
start = offsets[0] / 1000.0 if len(offsets) > 0 else 0
|
||||
end = offsets[1] / 1000.0 if len(offsets) > 1 else 0
|
||||
segments.append({
|
||||
'id': i,
|
||||
'start': start,
|
||||
'end': end,
|
||||
'text': item.get('text', '')
|
||||
})
|
||||
text = ' '.join(s['text'] for s in segments)
|
||||
if segments:
|
||||
duration = segments[-1]['end']
|
||||
|
||||
# 读取文本
|
||||
if txt_path.exists():
|
||||
text = txt_path.read_text(encoding='utf-8').strip()
|
||||
|
||||
return ASROutput(
|
||||
success=True,
|
||||
tool="whisper",
|
||||
has_subtitle=False,
|
||||
output_files={
|
||||
'txt': str(txt_path),
|
||||
'srt': str(srt_path),
|
||||
'json': str(json_path)
|
||||
},
|
||||
segments=segments,
|
||||
text=text,
|
||||
duration_seconds=duration
|
||||
)
|
||||
|
||||
async def _run_whisper_with_diarization(self, input_data: ASRInput) -> ASROutput:
|
||||
"""执行带说话者分离的 Whisper
|
||||
|
||||
Args:
|
||||
input_data: ASR 输入数据
|
||||
|
||||
Returns:
|
||||
ASR 处理结果，包含说话者信息
|
||||
"""
|
||||
# 先运行普通 Whisper
|
||||
basic_result = await self._run_whisper_basic(input_data)
|
||||
|
||||
try:
|
||||
# 运行说话者分离
|
||||
diarization_result = run_diarization(input_data.video_path)
|
||||
|
||||
# 合并结果
|
||||
merged_segments = assign_speakers_to_segments(
|
||||
basic_result.segments,
|
||||
diarization_result
|
||||
)
|
||||
|
||||
# 提取说话者列表
|
||||
speakers = list(set(
|
||||
seg.get('speaker') for seg in merged_segments if seg.get('speaker')
|
||||
))
|
||||
|
||||
return ASROutput(
|
||||
success=True,
|
||||
tool="whisper-diarization",
|
||||
has_subtitle=False,
|
||||
output_files=basic_result.output_files,
|
||||
segments=merged_segments,
|
||||
text=basic_result.text,
|
||||
duration_seconds=basic_result.duration_seconds,
|
||||
speakers=sorted(speakers)
|
||||
)
|
||||
|
||||
except ImportError:
|
||||
logger.warning("pyannote.audio not available, falling back to basic whisper")
|
||||
return basic_result
|
||||
except Exception as e:
|
||||
logger.warning(f"Diarization failed: {e}, falling back to basic whisper")
|
||||
return basic_result
|
||||
|
||||
|
||||
def run_diarization(audio_path: str, num_speakers: Optional[int] = None) -> Dict[str, List[Dict]]:
|
||||
"""运行 pyannote.audio 说话者分离
|
||||
|
||||
Args:
|
||||
audio_path: 音频文件路径
|
||||
num_speakers: 预期说话者数量（可选）
|
||||
|
||||
Returns:
|
||||
说话者时间线字典 {speaker_name: [{start, end}, ...]}
|
||||
|
||||
Raises:
|
||||
ImportError: 如果 pyannote.audio 未安装
|
||||
RuntimeError: 如果 HF_TOKEN 未配置
|
||||
"""
|
||||
hf_token = os.environ.get('HF_TOKEN')
|
||||
if not hf_token:
|
||||
raise RuntimeError("需要设置 HF_TOKEN 环境变量以使用说话者分离功能")
|
||||
|
||||
from huggingface_hub import login
|
||||
from pyannote.audio import Pipeline
|
||||
import torch
|
||||
|
||||
login(token=hf_token)
|
||||
|
||||
pipeline = Pipeline.from_pretrained("pyannote/speaker-diarization-3.1")
|
||||
|
||||
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
||||
pipeline.to(device)
|
||||
|
||||
diarization = pipeline(audio_path, num_speakers=num_speakers)
|
||||
|
||||
# 解析结果
|
||||
speakers = {}
|
||||
for turn, _, speaker in diarization.itertracks(yield_label=True):
|
||||
if speaker not in speakers:
|
||||
speakers[speaker] = []
|
||||
speakers[speaker].append({
|
||||
"start": round(turn.start, 2),
|
||||
"end": round(turn.end, 2)
|
||||
})
|
||||
|
||||
return speakers
|
||||
|
||||
|
||||
def assign_speakers_to_segments(
|
||||
segments: List[Dict],
|
||||
diarization: Dict[str, List[Dict]]
|
||||
) -> List[Dict]:
|
||||
"""为每个段落分配说话者
|
||||
|
||||
根据时间戳匹配说话者分离结果到 Whisper 段落。
|
||||
|
||||
Args:
|
||||
segments: Whisper 段落列表
|
||||
diarization: 说话者分离结果
|
||||
|
||||
Returns:
|
||||
带有 speaker 字段的段落列表
|
||||
"""
|
||||
result = []
|
||||
|
||||
for seg in segments:
|
||||
seg_start = seg.get('start', 0)
|
||||
seg_end = seg.get('end', 0)
|
||||
seg_mid = (seg_start + seg_end) / 2
|
||||
|
||||
# 找到包含中点时间的说话者
|
||||
assigned_speaker = None
|
||||
for speaker, turns in diarization.items():
|
||||
for turn in turns:
|
||||
if turn['start'] <= seg_mid <= turn['end']:
|
||||
assigned_speaker = speaker
|
||||
break
|
||||
if assigned_speaker:
|
||||
break
|
||||
|
||||
new_seg = dict(seg)
|
||||
new_seg['speaker'] = assigned_speaker
|
||||
result.append(new_seg)
|
||||
|
||||
return result
|
||||
@@ -0,0 +1,55 @@
|
||||
"""BaseProcessor - Abstract base class for all pipeline processors
|
||||
|
||||
Provides a unified interface that all processors (parse, download, asr, analyze)
|
||||
must follow, enabling consistent state management and error handling.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from abc import ABC, abstractmethod
|
||||
from typing import Optional, TYPE_CHECKING
|
||||
|
||||
from kb_workflow.pipeline.config import get_settings
|
||||
from kb_workflow.pipeline.schemas import StepInput, StepOutput
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from kb_workflow.pipeline.state import WorkflowStateManager
|
||||
|
||||
|
||||
class BaseProcessor(ABC):
|
||||
"""Abstract base class for all pipeline processors.
|
||||
|
||||
All processors (ParseProcessor, DownloadProcessor, ASRProcessor, AnalyzeProcessor)
|
||||
must inherit from this class and implement the process() method.
|
||||
|
||||
Attributes:
|
||||
STEP_NAME: Class attribute identifying the processor step name.
|
||||
settings: KBSettings instance from get_settings().
|
||||
state_manager: Optional WorkflowStateManager for state recording.
|
||||
"""
|
||||
|
||||
STEP_NAME: str
|
||||
|
||||
def __init__(self, state_manager: Optional[WorkflowStateManager] = None):
|
||||
"""Initialize the processor.
|
||||
|
||||
Args:
|
||||
state_manager: Optional WorkflowStateManager for recording step state.
|
||||
"""
|
||||
self.settings = get_settings()
|
||||
self.state_manager = state_manager
|
||||
|
||||
@abstractmethod
|
||||
async def process(self, input_data: StepInput) -> StepOutput:
|
||||
"""Process the input data and return the result.
|
||||
|
||||
Args:
|
||||
input_data: Input data for the processor step.
|
||||
|
||||
Returns:
|
||||
StepOutput with processing results.
|
||||
|
||||
Raises:
|
||||
Exception: If processing fails.
|
||||
"""
|
||||
pass
|
||||
@@ -0,0 +1,338 @@
|
||||
"""DownloadProcessor - Content download processor
|
||||
|
||||
Handles metadata extraction and content downloading for all platforms.
|
||||
Works with PlatformAdapter implementations (BilibiliAdapter, ZhihuAdapter).
|
||||
|
||||
Purpose: Unify metadata and content fetching into a single processor
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import logging
|
||||
from pathlib import Path
|
||||
from typing import Optional, TYPE_CHECKING, Dict, Any
|
||||
|
||||
from kb_workflow.pipeline.schemas import DownloadInput, DownloadOutput
|
||||
from kb_workflow.pipeline.processors.base import BaseProcessor
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from kb_workflow.pipeline.state import WorkflowStateManager
|
||||
from kb_workflow.pipeline.platforms.base import PlatformAdapter
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class DownloadProcessor(BaseProcessor):
|
||||
"""Content download processor.
|
||||
|
||||
Fetches metadata and content from platforms using PlatformAdapter.
|
||||
Supports video, article, answer, and zhuanlan content types.
|
||||
|
||||
Attributes:
|
||||
STEP_NAME: "download" - identifies this processor step
|
||||
"""
|
||||
|
||||
STEP_NAME = "download"
|
||||
|
||||
async def process(self, input_data: DownloadInput) -> DownloadOutput:
|
||||
"""Process the download request.
|
||||
|
||||
Args:
|
||||
input_data: DownloadInput with content_id, platform, content_dir, content_type
|
||||
|
||||
Returns:
|
||||
DownloadOutput with metadata_file, content_file, title, author, etc.
|
||||
"""
|
||||
content_dir = Path(input_data.content_dir)
|
||||
content_dir.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
# Check if step already completed (idempotency)
|
||||
if self.state_manager:
|
||||
state = self.state_manager.load()
|
||||
if state and state.is_step_completed(self.STEP_NAME):
|
||||
logger.info(f"Step {self.STEP_NAME} already completed, skipping")
|
||||
# Return a success output based on existing files
|
||||
return self._get_existing_output(input_data, content_dir)
|
||||
|
||||
# Get appropriate adapter
|
||||
try:
|
||||
adapter = self._get_adapter(input_data.platform)
|
||||
except ValueError as e:
|
||||
return self._record_failure(str(e))
|
||||
|
||||
try:
|
||||
# Fetch metadata (pass source_url for zhihu answers)
|
||||
if input_data.platform == "zhihu":
|
||||
metadata = await adapter.fetch_metadata(
|
||||
input_data.content_id,
|
||||
input_data.content_type,
|
||||
source_url=input_data.source_url
|
||||
)
|
||||
else:
|
||||
metadata = await adapter.fetch_metadata(
|
||||
input_data.content_id,
|
||||
input_data.content_type
|
||||
)
|
||||
|
||||
# Save metadata
|
||||
metadata_filename = adapter.get_output_filename(
|
||||
input_data.content_id,
|
||||
input_data.content_type,
|
||||
"metadata.json"
|
||||
)
|
||||
metadata_path = content_dir / metadata_filename
|
||||
self._save_metadata(metadata, metadata_path)
|
||||
|
||||
# Extract title and author from metadata
|
||||
title = self._extract_title(metadata, input_data.platform)
|
||||
author = self._extract_author(metadata, input_data.platform)
|
||||
duration = self._extract_duration(metadata, input_data.content_type)
|
||||
|
||||
# Fetch content (for text-based content)
|
||||
content_file = ""
|
||||
if input_data.content_type in ["article", "answer", "zhuanlan"]:
|
||||
if input_data.platform == "zhihu":
|
||||
content = await adapter.fetch_content(
|
||||
input_data.content_id,
|
||||
input_data.content_type,
|
||||
source_url=input_data.source_url
|
||||
)
|
||||
else:
|
||||
content = await adapter.fetch_content(
|
||||
input_data.content_id,
|
||||
input_data.content_type
|
||||
)
|
||||
content_filename = adapter.get_output_filename(
|
||||
input_data.content_id,
|
||||
input_data.content_type,
|
||||
"content.md"
|
||||
)
|
||||
content_path = content_dir / content_filename
|
||||
self._save_content(content, content_path)
|
||||
content_file = str(content_path)
|
||||
|
||||
# Download video for ASR (Whisper needs video file)
|
||||
if input_data.content_type == "video" and input_data.platform == "bilibili":
|
||||
download_result = await adapter.download_video(
|
||||
input_data.content_id,
|
||||
content_dir
|
||||
)
|
||||
if download_result["success"]:
|
||||
content_file = download_result["filepath"]
|
||||
logger.info(f"Video downloaded for ASR: {content_file}")
|
||||
else:
|
||||
logger.warning(f"Video download failed: {download_result['message']}")
|
||||
|
||||
# Record success
|
||||
result = DownloadOutput(
|
||||
success=True,
|
||||
metadata_file=str(metadata_path),
|
||||
content_file=content_file,
|
||||
title=title,
|
||||
author=author,
|
||||
duration=duration
|
||||
)
|
||||
|
||||
if self.state_manager:
|
||||
state = self.state_manager.load()
|
||||
if state:
|
||||
summary = f"title={title}, type={input_data.content_type}"
|
||||
state.mark_step_completed(self.STEP_NAME, summary)
|
||||
self.state_manager.save(state)
|
||||
|
||||
return result
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Download failed: {e}")
|
||||
return self._record_failure(str(e))
|
||||
|
||||
def _get_adapter(self, platform: str) -> PlatformAdapter:
|
||||
"""Get the appropriate PlatformAdapter for the platform.
|
||||
|
||||
Args:
|
||||
platform: Platform name ("bilibili" or "zhihu")
|
||||
|
||||
Returns:
|
||||
PlatformAdapter instance
|
||||
|
||||
Raises:
|
||||
ValueError: If platform is not supported
|
||||
"""
|
||||
if platform == "bilibili":
|
||||
from kb_workflow.pipeline.platforms.bilibili import BilibiliAdapter
|
||||
return BilibiliAdapter(state_manager=self.state_manager)
|
||||
elif platform == "zhihu":
|
||||
from kb_workflow.pipeline.platforms.zhihu import ZhihuAdapter
|
||||
return ZhihuAdapter()
|
||||
else:
|
||||
raise ValueError(f"Unsupported platform: {platform}")
|
||||
|
||||
def _save_metadata(self, metadata: Dict[str, Any], path: Path) -> None:
|
||||
"""Save metadata to JSON file.
|
||||
|
||||
Args:
|
||||
metadata: Metadata dictionary
|
||||
path: Output file path
|
||||
"""
|
||||
with open(path, 'w', encoding='utf-8') as f:
|
||||
json.dump(metadata, f, ensure_ascii=False, indent=2)
|
||||
logger.info(f"Saved metadata to {path}")
|
||||
|
||||
def _save_content(self, content: str, path: Path) -> None:
|
||||
"""Save content to Markdown file.
|
||||
|
||||
Args:
|
||||
content: Markdown content string
|
||||
path: Output file path
|
||||
"""
|
||||
with open(path, 'w', encoding='utf-8') as f:
|
||||
f.write(content)
|
||||
logger.info(f"Saved content to {path}")
|
||||
|
||||
def _extract_title(self, metadata: Dict[str, Any], platform: str) -> str:
|
||||
"""Extract title from metadata.
|
||||
|
||||
Args:
|
||||
metadata: Metadata dictionary
|
||||
platform: Platform name
|
||||
|
||||
Returns:
|
||||
Title string
|
||||
"""
|
||||
if platform == "bilibili":
|
||||
# Bilibili metadata structure: {metadata: {title: ...}, ...}
|
||||
if "metadata" in metadata:
|
||||
return metadata["metadata"].get("title", "")
|
||||
return metadata.get("title", "")
|
||||
else:
|
||||
return metadata.get("title", "")
|
||||
|
||||
def _extract_author(self, metadata: Dict[str, Any], platform: str) -> str:
|
||||
"""Extract author from metadata.
|
||||
|
||||
Args:
|
||||
metadata: Metadata dictionary
|
||||
platform: Platform name
|
||||
|
||||
Returns:
|
||||
Author string (author name, not mid)
|
||||
"""
|
||||
if platform == "bilibili":
|
||||
# Bilibili metadata structure: {metadata: {author: ..., mid: ...}, ...}
|
||||
if "metadata" in metadata:
|
||||
return str(metadata["metadata"].get("author", ""))
|
||||
return str(metadata.get("author", ""))
|
||||
else:
|
||||
return metadata.get("author", "")
|
||||
|
||||
def _extract_duration(self, metadata: Dict[str, Any], content_type: str) -> Optional[int]:
|
||||
"""Extract duration from metadata for video content.
|
||||
|
||||
Args:
|
||||
metadata: Metadata dictionary
|
||||
content_type: Content type
|
||||
|
||||
Returns:
|
||||
Duration in seconds for videos, None otherwise
|
||||
"""
|
||||
if content_type != "video":
|
||||
return None
|
||||
|
||||
# Bilibili video duration
|
||||
if "metadata" in metadata and "duration" in metadata["metadata"]:
|
||||
duration = metadata["metadata"]["duration"]
|
||||
# Handle formatted duration string (e.g., "12:34" or "01:23:45")
|
||||
if isinstance(duration, str):
|
||||
if not duration.strip():
|
||||
return None
|
||||
parts = duration.split(':')
|
||||
try:
|
||||
if len(parts) == 2:
|
||||
return int(parts[0]) * 60 + int(parts[1])
|
||||
elif len(parts) == 3:
|
||||
return int(parts[0]) * 3600 + int(parts[1]) * 60 + int(parts[2])
|
||||
except (ValueError, IndexError):
|
||||
return None
|
||||
elif isinstance(duration, (int, float)):
|
||||
return int(duration)
|
||||
|
||||
return None
|
||||
|
||||
def _record_failure(self, error_message: str) -> DownloadOutput:
|
||||
"""Record failure to state manager and return failure output.
|
||||
|
||||
Args:
|
||||
error_message: Error message
|
||||
|
||||
Returns:
|
||||
DownloadOutput with success=False
|
||||
"""
|
||||
if self.state_manager:
|
||||
state = self.state_manager.load()
|
||||
if state:
|
||||
state.mark_step_failed(self.STEP_NAME, error_message)
|
||||
self.state_manager.save(state)
|
||||
|
||||
return DownloadOutput(
|
||||
success=False,
|
||||
error_message=error_message,
|
||||
metadata_file="",
|
||||
content_file="",
|
||||
title="",
|
||||
author=""
|
||||
)
|
||||
|
||||
def _get_existing_output(
|
||||
self,
|
||||
input_data: DownloadInput,
|
||||
content_dir: Path
|
||||
) -> DownloadOutput:
|
||||
"""Get output from existing files when step was already completed.
|
||||
|
||||
Args:
|
||||
input_data: DownloadInput
|
||||
content_dir: Content directory path
|
||||
|
||||
Returns:
|
||||
DownloadOutput based on existing files
|
||||
"""
|
||||
adapter = self._get_adapter(input_data.platform)
|
||||
|
||||
metadata_filename = adapter.get_output_filename(
|
||||
input_data.content_id,
|
||||
input_data.content_type,
|
||||
"metadata.json"
|
||||
)
|
||||
metadata_path = content_dir / metadata_filename
|
||||
|
||||
# Try to read metadata for title/author
|
||||
title = ""
|
||||
author = ""
|
||||
duration = None
|
||||
if metadata_path.exists():
|
||||
try:
|
||||
with open(metadata_path, 'r', encoding='utf-8') as f:
|
||||
metadata = json.load(f)
|
||||
title = self._extract_title(metadata, input_data.platform)
|
||||
author = self._extract_author(metadata, input_data.platform)
|
||||
duration = self._extract_duration(metadata, input_data.content_type)
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
content_filename = adapter.get_output_filename(
|
||||
input_data.content_id,
|
||||
input_data.content_type,
|
||||
"content.md"
|
||||
)
|
||||
content_path = content_dir / content_filename
|
||||
content_file = str(content_path) if content_path.exists() else ""
|
||||
|
||||
return DownloadOutput(
|
||||
success=True,
|
||||
metadata_file=str(metadata_path),
|
||||
content_file=content_file,
|
||||
title=title,
|
||||
author=author,
|
||||
duration=duration
|
||||
)
|
||||
@@ -0,0 +1,83 @@
|
||||
"""ParseProcessor - URL parsing processor
|
||||
|
||||
Parses URLs from various platforms (Bilibili, Zhihu) and extracts
|
||||
platform, content type, and content ID information.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Optional, TYPE_CHECKING
|
||||
|
||||
from kb_workflow.pipeline.config import get_settings
|
||||
from kb_workflow.pipeline.link_parser import LinkParser
|
||||
from kb_workflow.pipeline.schemas import ParseInput, ParseOutput
|
||||
from kb_workflow.pipeline.processors.base import BaseProcessor
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from kb_workflow.pipeline.state import WorkflowStateManager
|
||||
|
||||
|
||||
class ParseProcessor(BaseProcessor):
|
||||
"""URL parsing processor.
|
||||
|
||||
Parses URLs from Bilibili and Zhihu platforms, extracting
|
||||
platform, content type, and content ID.
|
||||
|
||||
Attributes:
|
||||
STEP_NAME: "parse" - identifies this processor step
|
||||
"""
|
||||
|
||||
STEP_NAME = "parse"
|
||||
|
||||
async def process(self, input_data: ParseInput) -> ParseOutput:
|
||||
"""Parse the URL and extract platform information.
|
||||
|
||||
Args:
|
||||
input_data: ParseInput with URL to parse
|
||||
|
||||
Returns:
|
||||
ParseOutput with parsed information or error message
|
||||
"""
|
||||
parser = LinkParser()
|
||||
parsed = parser.parse(input_data.url)
|
||||
|
||||
if parsed is None:
|
||||
# Unknown URL format
|
||||
result = ParseOutput(
|
||||
success=False,
|
||||
error_message="Unsupported URL format",
|
||||
platform="",
|
||||
content_type="",
|
||||
content_id="",
|
||||
resolved_url="",
|
||||
original_url=input_data.url
|
||||
)
|
||||
|
||||
# Record failure state
|
||||
if self.state_manager:
|
||||
state = self.state_manager.load()
|
||||
if state:
|
||||
state.mark_step_failed(self.STEP_NAME, result.error_message)
|
||||
self.state_manager.save(state)
|
||||
|
||||
return result
|
||||
|
||||
# Successful parse
|
||||
result = ParseOutput(
|
||||
success=True,
|
||||
platform=parsed.platform,
|
||||
content_type=parsed.content_type,
|
||||
content_id=parsed.content_id,
|
||||
resolved_url=parsed.resolved_url,
|
||||
original_url=parsed.original_url
|
||||
)
|
||||
|
||||
# Record success state
|
||||
if self.state_manager:
|
||||
state = self.state_manager.load()
|
||||
if state:
|
||||
summary = f"platform={parsed.platform}, type={parsed.content_type}, id={parsed.content_id}"
|
||||
state.mark_step_completed(self.STEP_NAME, summary)
|
||||
self.state_manager.save(state)
|
||||
|
||||
return result
|
||||
Reference in New Issue
Block a user