ai-agent-book 精选快照（<2MB 代码与文档，来自 github.com/bojieli/ai-agent-book）

chapter4/multimodal-agent/demo.py

@@ -0,0 +1,340 @@
+"""
+Demo script showcasing different extraction techniques
+"""
+
+import argparse
+import asyncio
+import sys
+from pathlib import Path
+from typing import List, Optional, Tuple
+
+from agent import MultimodalAgent, MultimodalContent
+from config import ExtractionMode
+
+
+class _Tee:
+    """Duplicate stdout writes to a file so --output can save the transcript."""
+
+    def __init__(self, stream, file_handle):
+        self._stream = stream
+        self._file = file_handle
+
+    def write(self, data):
+        self._stream.write(data)
+        self._file.write(data)
+
+    def flush(self):
+        self._stream.flush()
+        self._file.flush()
+
+
+async def compare_extraction_modes(file_path: str, query: str, model: str = "gemini-3.5-flash"):
+    """Compare different extraction modes for the same content"""
+    
+    print(f"\n{'='*80}")
+    print(f"COMPARING EXTRACTION MODES")
+    print(f"File: {file_path}")
+    print(f"Query: {query}")
+    print(f"{'='*80}\n")
+    
+    # Determine content type
+    path = Path(file_path)
+    suffix = path.suffix.lower()
+    
+    if suffix == '.pdf':
+        content_type = "pdf"
+    elif suffix in ['.jpg', '.jpeg', '.png', '.gif', '.bmp', '.webp']:
+        content_type = "image"
+    elif suffix in ['.mp3', '.wav', '.m4a', '.flac', '.aac', '.ogg']:
+        content_type = "audio"
+    else:
+        print(f"Unsupported file type: {suffix}")
+        return
+    
+    # Test with native mode (Gemini)
+    print("\n" + "-"*60)
+    print(f"1. NATIVE MULTIMODAL MODE ({model})")
+    print("-"*60)
+
+    agent_native = MultimodalAgent(
+        model=model,
+        mode=ExtractionMode.NATIVE,
+        enable_tools=False
+    )
+    
+    content = MultimodalContent(type=content_type, path=file_path)
+    
+    try:
+        result = await agent_native.process_multimodal_content(content, query)
+        print(result)
+    except Exception as e:
+        print(f"Error: {e}")
+    
+    # Test with extract-to-text mode
+    print("\n" + "-"*60)
+    print("2. EXTRACT TO TEXT MODE")
+    print("-"*60)
+    
+    agent_extract = MultimodalAgent(
+        model=model,
+        mode=ExtractionMode.EXTRACT_TO_TEXT,
+        enable_tools=False
+    )
+    
+    try:
+        # First extract the content
+        print("Extracting content to text...")
+        extracted = await agent_extract._extract_single_content(content)
+        print("\nExtracted text:")
+        print(extracted)
+        
+        # Then answer the query
+        print(f"\nAnswering query with extracted text...")
+        result = await agent_extract._answer_with_context(extracted, query)
+        print(result)
+    except Exception as e:
+        print(f"Error: {e}")
+    
+    # Test with extract-to-text + tools mode
+    print("\n" + "-"*60)
+    print("3. EXTRACT TO TEXT + MULTIMODAL TOOLS")
+    print("-"*60)
+    
+    agent_tools = MultimodalAgent(
+        model=model,
+        mode=ExtractionMode.EXTRACT_TO_TEXT,
+        enable_tools=True
+    )
+    
+    try:
+        print("Using extract-to-text with tools enabled for follow-up questions...")
+        
+        # Initial processing
+        extracted = await agent_tools._extract_single_content(content)
+        print(f"Extracted {len(extracted)} characters")
+        
+        # Simulate a conversation with follow-up
+        async for chunk in agent_tools.chat(query, content, stream=True):
+            print(chunk, end="", flush=True)
+        print()
+        
+        # Follow-up question that might use tools
+        if content_type == "image":
+            follow_up = f"What colors are dominant in the image at {file_path}?"
+        elif content_type == "pdf":
+            follow_up = f"What specific data or figures are mentioned in the PDF at {file_path}?"
+        else:  # audio
+            follow_up = f"What is the tone or mood of the audio at {file_path}?"
+            
+        print(f"\nFollow-up question: {follow_up}")
+        async for chunk in agent_tools.chat(follow_up, None, stream=True):
+            print(chunk, end="", flush=True)
+        print()
+        
+    except Exception as e:
+        print(f"Error: {e}")
+
+
+async def compare_models(file_path: str, query: str):
+    """Compare different models for the same task"""
+    
+    print(f"\n{'='*80}")
+    print(f"COMPARING MODELS")
+    print(f"File: {file_path}")
+    print(f"Query: {query}")
+    print(f"{'='*80}\n")
+    
+    # Determine content type
+    path = Path(file_path)
+    suffix = path.suffix.lower()
+    
+    if suffix == '.pdf':
+        content_type = "pdf"
+    elif suffix in ['.jpg', '.jpeg', '.png', '.gif', '.bmp', '.webp']:
+        content_type = "image"
+    elif suffix in ['.mp3', '.wav', '.m4a', '.flac', '.aac', '.ogg']:
+        content_type = "audio"
+    else:
+        print(f"Unsupported file type: {suffix}")
+        return
+        
+    content = MultimodalContent(type=content_type, path=file_path)
+    
+    # Test with different models
+    models = ["gemini-3.5-flash", "gpt-5.6-luna", "doubao-1.6"]
+    
+    for model in models:
+        print("\n" + "-"*60)
+        print(f"Model: {model}")
+        print("-"*60)
+        
+        try:
+            # Skip if API key not configured
+            from config import Config
+            config = Config()
+            
+            if model == "gemini-3.5-flash" and not config.gemini_api_key:
+                print("Skipping: Gemini API key not configured")
+                continue
+            elif model in ["gpt-5.6-luna", "gpt-5"] and not (config.openai_api_key or config.openrouter_api_key):
+                print("Skipping: OpenAI API key not configured")
+                continue
+            elif model == "doubao-1.6" and not config.doubao_api_key:
+                print("Skipping: Doubao API key not configured")
+                continue
+            
+            agent = MultimodalAgent(
+                model=model,
+                mode=ExtractionMode.NATIVE if content_type != "audio" or model == "gemini-3.5-flash" else ExtractionMode.EXTRACT_TO_TEXT,
+                enable_tools=False
+            )
+            
+            result = await agent.process_multimodal_content(content, query)
+            print(result)
+            
+        except Exception as e:
+            print(f"Error: {e}")
+
+
+async def demo_conversation_with_tools():
+    """Demonstrate a conversation with multimodal tools"""
+    
+    print(f"\n{'='*80}")
+    print("DEMO: CONVERSATION WITH MULTIMODAL TOOLS")
+    print(f"{'='*80}\n")
+    
+    agent = MultimodalAgent(
+        model="gemini-3.5-flash",
+        mode=ExtractionMode.EXTRACT_TO_TEXT,
+        enable_tools=True
+    )
+    
+    # Simulate a conversation
+    conversations = [
+        ("I need help analyzing some documents. I have PDFs, images, and audio files.", None),
+        ("Can you analyze the image at test_files/sample.jpg and tell me what you see?", None),
+        ("Now analyze the PDF at test_files/document.pdf and summarize its main points.", None),
+        ("What's in the audio file at test_files/recording.mp3?", None),
+        ("Based on all these files, what's the common theme?", None)
+    ]
+    
+    for message, content in conversations:
+        print(f"\nUser: {message}")
+        print("Assistant: ", end="", flush=True)
+        
+        try:
+            async for chunk in agent.chat(message, content, stream=True):
+                print(chunk, end="", flush=True)
+            print()
+        except Exception as e:
+            print(f"\nError: {e}")
+            print("(File might not exist - this is a demo)")
+
+
+def build_parser() -> argparse.ArgumentParser:
+    """构建实验 4-2 的命令行接口。"""
+    parser = argparse.ArgumentParser(
+        description=(
+            "实验 4-2：多模态信息提取的三种技术范式对比（原生多模态 / 提取为文本 / 带工具）。\n"
+            "将同一多模态文件和同一问题分别交给三种模式处理，观察表现差异。"
+        ),
+        formatter_class=argparse.RawDescriptionHelpFormatter,
+        epilog=(
+            "示例：\n"
+            "  # 先离线生成含图表的样例（无需 API Key）\n"
+            "  python demo.py --generate-sample\n"
+            "  # 用生成的图表跑三种范式对比（需要 API Key）\n"
+            "  python demo.py --file test_files/sample_chart.png \\\n"
+            '      --query \"Which quarter had the highest revenue, and what was the exact value?\"\n'
+            "  # 兼容旧写法（位置参数）\n"
+            "  python demo.py document.pdf \"总结这份文档的要点\""
+        ),
+    )
+    parser.add_argument(
+        "file", nargs="?", default=None,
+        help="要处理的多模态文件（图像 / PDF 文档 / 音频）。也可用 --file 指定",
+    )
+    parser.add_argument(
+        "query", nargs="?", default=None,
+        help="向该文件提出的问题。也可用 --query 指定",
+    )
+    parser.add_argument(
+        "--file", dest="file_opt", default=None,
+        help="要处理的多模态文件（等价于位置参数 file）",
+    )
+    parser.add_argument(
+        "--query", dest="query_opt", default=None,
+        help="向该文件提出的问题（等价于位置参数 query）",
+    )
+    parser.add_argument(
+        "--model", default="gemini-3.5-flash",
+        help="原生 / 提取模式使用的模型（默认：gemini-3.5-flash）",
+    )
+    parser.add_argument(
+        "--skip-model-comparison", action="store_true",
+        help="只跑三种范式对比，跳过跨模型对比",
+    )
+    parser.add_argument(
+        "--generate-sample", action="store_true",
+        help="离线生成含图表的样例文件到 test_files/ 后退出（无需 API Key）",
+    )
+    parser.add_argument(
+        "--output", "-o", default=None,
+        help="将完整对比结果同时写入指定文件（如 result.txt）",
+    )
+    return parser
+
+
+async def run_comparison(file_path: str, query: str, model: str, skip_model_comparison: bool):
+    """运行三种范式对比，可选跨模型对比。"""
+    print("="*80)
+    print("MULTIMODAL AGENT DEMO")
+    print("="*80)
+
+    await compare_extraction_modes(file_path, query, model=model)
+    if not skip_model_comparison:
+        await compare_models(file_path, query)
+
+
+async def main():
+    """实验入口：解析参数并运行对比。"""
+    parser = build_parser()
+    args = parser.parse_args()
+
+    # 离线样例生成：不需要 API Key，直接产出图表 + PDF 报告
+    if args.generate_sample:
+        import create_sample
+        sys.argv = ["create_sample.py"]  # 用默认输出目录 test_files/
+        create_sample.main()
+        return
+
+    file_path = args.file_opt or args.file
+    query = args.query_opt or args.query
+
+    # 缺少文件或问题时，回退到无需真实文件的对话演示
+    if not file_path or not query:
+        print("="*80)
+        print("MULTIMODAL AGENT DEMO")
+        print("="*80)
+        print("\n未提供 <file> 与 <query>，改为运行对话演示。")
+        print("用法：python demo.py --file <文件> --query <问题>")
+        print("先生成样例：python demo.py --generate-sample\n")
+        await demo_conversation_with_tools()
+        return
+
+    # 支持 --output：把整段对比结果同时落盘
+    if args.output:
+        with open(args.output, "w", encoding="utf-8") as fh:
+            original_stdout = sys.stdout
+            sys.stdout = _Tee(original_stdout, fh)
+            try:
+                await run_comparison(file_path, query, args.model, args.skip_model_comparison)
+            finally:
+                sys.stdout = original_stdout
+        print(f"\n完整对比结果已写入：{args.output}")
+    else:
+        await run_comparison(file_path, query, args.model, args.skip_model_comparison)
+
+
+if __name__ == "__main__":
+    asyncio.run(main())