#!/usr/bin/env python3 """ Multimodal RAG CLI — 命令行交互工具。 用法: python -m ui.cli ingest --path ./data/ python -m ui.cli query "图中有什么?" python -m ui.cli cross-modal "ocean waves" python -m ui.cli status python -m ui.cli clear """ import sys from pathlib import Path from typing import Optional # 确保项目根目录在 sys.path 中 sys.path.insert(0, str(Path(__file__).resolve().parent.parent)) import typer from rich.console import Console from rich.panel import Panel from rich.table import Table from rich.markdown import Markdown from src.pipeline.ingestion import IngestionPipeline from src.pipeline.query import QueryPipeline from src.storage.vector_store import VectorStore from src.embeddings.jina_embedder import Modality from config.settings import settings app = typer.Typer( name="multimodal-rag", help="多模态 RAG 智能问答系统 CLI", add_completion=False, ) console = Console() # 延迟初始化的组件 _ingestion_pipeline: Optional[IngestionPipeline] = None _query_pipeline: Optional[QueryPipeline] = None _vector_store: Optional[VectorStore] = None def get_ingestion_pipeline() -> IngestionPipeline: global _ingestion_pipeline if _ingestion_pipeline is None: _ingestion_pipeline = IngestionPipeline() return _ingestion_pipeline def get_query_pipeline() -> QueryPipeline: global _query_pipeline if _query_pipeline is None: _query_pipeline = QueryPipeline() return _query_pipeline def get_vector_store() -> VectorStore: global _vector_store if _vector_store is None: _vector_store = VectorStore() return _vector_store # ============================================================ # ingest — 入库数据 # ============================================================ @app.command() def ingest( path: str = typer.Argument(..., help="文件或目录路径"), recursive: bool = typer.Option(True, help="递归处理子目录"), incremental: bool = typer.Option(True, help="跳过已索引的文件"), ): """入库文件或目录到向量库。""" console.print(f"[bold blue]入库路径:[/] {path}") pipeline = get_ingestion_pipeline() pipeline.incremental = incremental p = Path(path) if p.is_file(): stats = pipeline.ingest_file(str(p)) elif p.is_dir(): stats = pipeline.ingest_directory(str(p), recursive=recursive) else: console.print(f"[bold red]路径不存在:[/] {path}") raise typer.Exit(code=1) # 输出统计 table = Table(title="入库统计") table.add_column("指标", style="cyan") table.add_column("值", style="green") table.add_row("处理文件数", str(stats.files_processed)) table.add_row("失败文件数", str(stats.files_failed)) table.add_row("生成 Chunks", str(stats.chunks_created)) for mod, count in stats.chunks_by_modality.items(): table.add_row(f" └ {mod}", str(count)) table.add_row("总耗时", f"{stats.total_latency_ms/1000:.1f}s") if stats.errors: table.add_row("错误", "\n".join(stats.errors[:5])) console.print(table) # ============================================================ # query — 问答 # ============================================================ @app.command() def query( question: str = typer.Argument(..., help="要查询的问题"), top_k: int = typer.Option(10, help="检索结果数"), modality: Optional[str] = typer.Option(None, help="按模态过滤: text/image/audio/video"), cross_modal: bool = typer.Option(False, "--cross-modal", "-x", help="启用跨模态检索"), no_generate: bool = typer.Option(False, "--no-generate", help="仅检索, 不生成回答"), ): """查询向量库并生成回答。""" pipeline = get_query_pipeline() modality_filter = None if modality: try: modality_filter = Modality(modality) except ValueError: console.print(f"[bold red]无效的模态:[/] {modality}") console.print("可用: text, image, audio, video") raise typer.Exit(code=1) with console.status("[bold green]处理中..."): if no_generate: result = pipeline.retrieve_only( query=question, top_k=top_k, modality_filter=modality_filter, ) _display_retrieve_results(result) else: response = pipeline.query( question=question, top_k=top_k, modality_filter=modality_filter, cross_modal=cross_modal, ) _display_query_response(response) def _display_query_response(response): """显示查询响应。""" # 回答 console.print() console.print(Panel( Markdown(response.answer), title="[bold green]回答[/]", border_style="green", )) # 延迟 lb = response.latency_breakdown console.print( f"[dim]⏱ 总耗时: {lb['total_ms']:.0f}ms " f"(Embed: {lb['embed_ms']:.0f}ms | " f"检索: {lb['retrieve_ms']:.0f}ms | " f"生成: {lb['generate_ms']:.0f}ms)[/]" ) # 来源 if response.sources: console.print() console.print(f"[bold]📎 来源 ({len(response.sources)} 条):[/]") table = Table(show_header=True) table.add_column("#", style="dim", width=4) table.add_column("模态", width=8) table.add_column("来源", style="cyan") table.add_column("预览", style="white") table.add_column("相关度", justify="right") for i, src in enumerate(response.sources[:10], 1): mod = src["modality"] mod_emoji = {"text": "📝", "image": "🖼", "audio": "🎵", "video": "🎬"}.get(mod, "❓") preview = src.get("content_preview", "")[:80] table.add_row( str(i), f"{mod_emoji} {mod}", src.get("source_file_name", "")[:30], preview, f"{src['score']:.3f}", ) console.print(table) def _display_retrieve_results(result): """显示仅检索结果。""" console.print() console.print(f"[bold]查询:[/] {result['query']}") console.print(f"[bold]找到 {result['total_results']} 条结果[/]") console.print(f"[dim]模态分布: {result['modality_breakdown']}[/]") table = Table(title="检索结果") table.add_column("#", style="dim", width=4) table.add_column("模态", width=8) table.add_column("来源", style="cyan") table.add_column("预览", style="white") table.add_column("相关度", justify="right") for i, src in enumerate(result["results"][:20], 1): mod = src["modality"] mod_emoji = {"text": "📝", "image": "🖼", "audio": "🎵", "video": "🎬"}.get(mod, "❓") preview = src.get("content_preview", "")[:80] table.add_row( str(i), f"{mod_emoji} {mod}", src.get("source_file_name", "")[:30], preview, f"{src['score']:.3f}", ) console.print(table) # ============================================================ # cross-modal — 跨模态搜索 # ============================================================ @app.command(name="cross-modal") def cross_modal( query_text: str = typer.Argument(..., help="搜索查询文本"), top_k: int = typer.Option(5, help="每种模态的结果数"), modalities: Optional[str] = typer.Option(None, help="模态列表 (逗号分隔): text,image,audio,video"), ): """跨模态搜索: 文本查询检索所有模态的内容。""" pipeline = get_query_pipeline() mod_list = None if modalities: mod_list = [Modality(m.strip()) for m in modalities.split(",")] with console.status("[bold green]跨模态搜索中..."): result = pipeline.retrieve_all_modalities( query=query_text, top_k_per_modality=top_k, ) # 手动扩展以支持模态过滤 if mod_list: filtered = {} for mod_str, data in result["modalities"].items(): if Modality(mod_str) in mod_list: filtered[mod_str] = data result["modalities"] = filtered console.print() console.print(f"[bold]跨模态搜索:[/] {query_text}") for mod_str in ["text", "image", "audio", "video"]: if mod_str not in result["modalities"]: continue data = result["modalities"][mod_str] mod_emoji = {"text": "📝", "image": "🖼", "audio": "🎵", "video": "🎬"}.get(mod_str, "❓") console.print() console.print(f"[bold]{mod_emoji} {mod_str.upper()} ({data['count']} 条)[/]") for i, r in enumerate(data["results"][:5], 1): preview = r.get("content_preview", "")[:100] score = r.get("score", 0) source = r.get("source_file_name", "") console.print(f" {i}. [{score:.3f}] {source}: {preview}") # ============================================================ # status — 状态查询 # ============================================================ @app.command() def status(): """查看向量库状态和统计。""" store = get_vector_store() stats = store.get_stats() console.print() console.print("[bold blue]📊 向量库状态[/]") console.print() table = Table() table.add_column("指标", style="cyan") table.add_column("值", style="green") table.add_row("Collection", stats.get("collection_name", "N/A")) table.add_row("存储路径", stats.get("persist_directory", "N/A")) table.add_row("总 Chunks", str(stats.get("total_chunks", 0))) table.add_row("嵌入维度", str(stats.get("embedding_dim", "N/A"))) table.add_row("唯一来源数", str(stats.get("unique_sources", 0))) breakdown = stats.get("modality_breakdown", {}) for mod, count in breakdown.items(): emoji = {"text": "📝", "image": "🖼", "audio": "🎵", "video": "🎬"}.get(mod, "❓") table.add_row(f" {emoji} {mod}", str(count)) console.print(table) # 来源列表 sources = store.list_sources() if sources: console.print() console.print("[bold]📁 已索引来源:[/]") for s in sources[:20]: mods = ", ".join(f"{m}:{c}" for m, c in s.get("modalities", {}).items()) console.print(f" • {s['source_file_name']} ({s['count']} chunks: {mods})") # ============================================================ # clear — 清空 / 删除 # ============================================================ @app.command() def clear( source: Optional[str] = typer.Option(None, help="按来源文件删除"), all_: bool = typer.Option(False, "--all", help="清空整个向量库"), ): """删除索引数据。""" store = get_vector_store() if all_: confirm = typer.confirm("⚠️ 确定要清空整个向量库吗? 此操作不可逆!") if confirm: count = store.clear() console.print(f"[green]✓ 已清空, 删除了 {count} 个 chunks[/]") else: console.print("已取消") elif source: count = store.delete_by_source(source) console.print(f"[green]✓ 已删除 {count} 个 chunks (来源: {source})[/]") else: console.print("[yellow]请指定 --source 或 --all[/]") # ============================================================ # list-sources — 列出来源 # ============================================================ @app.command(name="list-sources") def list_sources(): """列出所有已索引的来源文件。""" store = get_vector_store() sources = store.list_sources() if not sources: console.print("[yellow]暂无已索引的来源[/]") return console.print() console.print(f"[bold]📁 已索引来源 ({len(sources)} 个):[/]") console.print() table = Table() table.add_column("文件", style="cyan") table.add_column("Chunks", justify="right") table.add_column("模态分布", style="green") for s in sources: mods = ", ".join(f"{m}:{c}" for m, c in s.get("modalities", {}).items()) table.add_row( s.get("source_file_name", "?"), str(s.get("count", 0)), mods, ) console.print(table) # ============================================================ # main # ============================================================ def main(): """CLI 入口。""" app() if __name__ == "__main__": main()