CoolFace
Apppublic

doreamong/YoutubeRAG

sourceHugging Faceupdated 1y agoView on Hugging Face
0likes
main.py230 linesDownload Raw Back to root
1#!/usr/bin/env python3.122"""3YouTube RAG System - Main Entry Point4YouTube RAG系统 - 主入口程序5 6A lightweight YouTube video Q&A tool with persistent session management.7轻量级YouTube视频问答工具,支持持久化会话管理。8"""9 10import argparse11import sys12from pathlib import Path13 14# Add src to Python path15src_path = Path(__file__).parent / "src"16sys.path.insert(0, str(src_path))17 18from youtube_rag_system import YouTubeRAG19from youtube_rag_system.ui.gradio_interface import create_interface20 21 22def main():23    """Main entry point"""24    parser = argparse.ArgumentParser(25        description="YouTube RAG System - YouTube视频RAG问答工具",26        formatter_class=argparse.RawDescriptionHelpFormatter,27        epilog="""28Examples / 示例:29  python3.12 main.py --ui                                    # Launch web interface / 启动网页界面30  python3.12 main.py --url "https://youtube.com/watch?v=..."  # CLI mode / 命令行模式31  python3.12 main.py --list-sessions                        # List saved sessions / 列出保存的会话32        """33    )34    35    parser.add_argument(36        "--ui", 37        action="store_true",38        help="Launch Gradio web interface / 启动Gradio网页界面"39    )40    41    parser.add_argument(42        "--url", 43        type=str,44        help="YouTube video URL for CLI mode / 命令行模式的YouTube视频URL"45    )46    47    parser.add_argument(48        "--model", 49        type=str, 50        default="gpt-3.5-turbo",51        help="OpenAI model name / OpenAI模型名称 (default: gpt-3.5-turbo)"52    )53    54    parser.add_argument(55        "--chunk-size", 56        type=int, 57        default=1000,58        help="Text chunk size / 文本块大小 (default: 1000)"59    )60    61    parser.add_argument(62        "--chunk-overlap", 63        type=int, 64        default=20,65        help="Text chunk overlap / 文本块重叠 (default: 20)"66    )67    68    parser.add_argument(69        "--list-sessions", 70        action="store_true",71        help="List all saved sessions / 列出所有保存的会话"72    )73    74    parser.add_argument(75        "--load-session", 76        type=str,77        help="Load a saved session by name / 根据名称加载保存的会话"78    )79    80    parser.add_argument(81        "--delete-session", 82        type=str,83        help="Delete a saved session by name / 根据名称删除保存的会话"84    )85    86    args = parser.parse_args()87    88    # Default to UI mode if no arguments provided89    if len(sys.argv) == 1:90        args.ui = True91    92    try:93        if args.ui:94            # Launch web interface95            print("🚀 Launching YouTube RAG System Web Interface...")96            print("🚀 启动YouTube RAG系统网页界面...")97            interface = create_interface()98            interface.launch(99                server_name="0.0.0.0",100                server_port=7860,101                share=True,102                show_error=True103            )104            105        elif args.list_sessions:106            # List sessions107            rag = YouTubeRAG()108            sessions = rag.list_sessions()109            if not sessions:110                print("📋 No saved sessions found / 没有找到保存的会话")111            else:112                print(f"📋 Found {len(sessions)} saved sessions / 找到{len(sessions)}个保存的会话:\n")113                for i, session in enumerate(sessions, 1):114                    created_time = session['created_at'][:19].replace('T', ' ')115                    print(f"{i}. {session['name']}")116                    print(f"   🕒 Created: {created_time}")117                    print(f"   🎥 Video: {session['video_url']}")118                    print(f"   🤖 Model: {session['model_name']}")119                    print()120                    121        elif args.delete_session:122            # Delete session123            rag = YouTubeRAG()124            if rag.delete_session(args.delete_session):125                print(f"✅ Session '{args.delete_session}' deleted successfully")126            else:127                print(f"❌ Failed to delete session '{args.delete_session}'")128                129        elif args.load_session:130            # Load and interact with session131            rag = YouTubeRAG(132                chunk_size=args.chunk_size,133                chunk_overlap=args.chunk_overlap,134                model_name=args.model,135                status_callback=print136            )137            138            session_data = rag.load_session(args.load_session)139            if not session_data:140                print(f"❌ Failed to load session '{args.load_session}'")141                sys.exit(1)142            143            qa_chain = session_data["qa_chain"]144            print(f"\n✅ Session '{args.load_session}' loaded successfully!")145            print(f"\n📄 Summary:\n{'-' * 50}")146            print(session_data['summary'])147            print('-' * 50)148            149            # Interactive Q&A150            print("\n🤖 Q&A system ready! Type 'quit' to exit / 问答系统就绪!输入 'quit' 退出")151            print("-" * 50)152            153            while True:154                try:155                    question = input("\n❓ Ask a question / 请提问: ").strip()156                    if question.lower() in ['quit', 'exit', '退出']:157                        break158                    159                    if not question:160                        continue161                    162                    print("🤔 Thinking... / 思考中...")163                    answer = rag.ask_question(qa_chain, question)164                    print(f"\n💡 {answer}")165                    166                except KeyboardInterrupt:167                    break168                except Exception as e:169                    print(f"❌ Error: {e}")170            171            print("\n👋 Goodbye! / 再见!")172            173        elif args.url:174            # CLI mode with URL175            def status_print(msg):176                print(msg)177            178            rag = YouTubeRAG(179                chunk_size=args.chunk_size,180                chunk_overlap=args.chunk_overlap,181                model_name=args.model,182                status_callback=status_print183            )184            185            print(f"🎥 Processing video: {args.url}")186            result = rag.process_video(args.url, allow_transcription=True, save_summary=True, save_original=True)187            188            print(f"\n📄 Video Summary:\n{'-' * 50}")189            print(result['summary'])190            print('-' * 50)191            192            if result['session_name']:193                print(f"\n💾 Session saved as: {result['session_name']}")194            195            # Interactive Q&A196            qa_chain = result['qa_chain']197            print("\n🤖 Q&A system ready! Type 'quit' to exit / 问答系统就绪!输入 'quit' 退出")198            print("-" * 50)199            200            while True:201                try:202                    question = input("\n❓ Ask a question / 请提问: ").strip()203                    if question.lower() in ['quit', 'exit', '退出']:204                        break205                    206                    if not question:207                        continue208                    209                    print("🤔 Thinking... / 思考中...")210                    answer = rag.ask_question(qa_chain, question)211                    print(f"\n💡 {answer}")212                    213                except KeyboardInterrupt:214                    break215                except Exception as e:216                    print(f"❌ Error: {e}")217            218            print("\n👋 Goodbye! / 再见!")219            220        else:221            # No valid arguments provided222            parser.print_help()223            224    except Exception as e:225        print(f"❌ Error: {e}")226        sys.exit(1)227 228 229if __name__ == "__main__":230    main()