1
0

main.py 5.0 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139
  1. import os
  2. import time
  3. import sys
  4. # Add the current directory to sys.path to ensure imports work correctly
  5. sys.path.append(os.getcwd())
  6. sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), "agents")))
  7. sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), "src")))
  8. from agents.outline_agent import OutlineAgent
  9. from agents.chapter_generate_agent import ChapterGenerateAgent
  10. from hello_agents import HelloAgentsLLM
  11. def print_step(step_name):
  12. print("\n" + "="*60)
  13. print(f"正在执行步骤: {step_name}")
  14. print("="*60 + "\n")
  15. def main():
  16. # Configuration
  17. novel_id = f"test_novel_{int(time.time())}"
  18. title = "测试Agent功能小说"
  19. user_idea = "一个关于AI程序员意外穿越到自己编写的代码世界中的故事,他需要修复这个世界的BUG才能回到现实。"
  20. print(f"测试配置:\n小说ID: {novel_id}\n标题: {title}\n创意: {user_idea}\n")
  21. # Initialize LLM
  22. # Assuming environment variables are set correctly for the default provider
  23. try:
  24. llm = HelloAgentsLLM()
  25. print("LLM 初始化成功。")
  26. except Exception as e:
  27. print(f"LLM 初始化失败: {e}")
  28. return
  29. # ---------------------------------------------------------
  30. # Test Outline Agent
  31. # ---------------------------------------------------------
  32. print_step("1. 初始化 OutlineAgent (大纲生成Agent)")
  33. try:
  34. outline_agent = OutlineAgent(name="TestOutlineAgent", llm=llm)
  35. print("OutlineAgent 初始化完成。")
  36. except Exception as e:
  37. print(f"OutlineAgent 初始化失败: {e}")
  38. return
  39. print_step("2. 生成大纲 (Generate Outline)")
  40. print(f"调用 outline_agent.run,输入创意: {user_idea}")
  41. start_time = time.time()
  42. try:
  43. outline_content, outline_note_id = outline_agent.run(
  44. user_input=user_idea,
  45. novel_id=novel_id,
  46. title=title,
  47. tags=["科幻", "穿越", "程序员"],
  48. target_length=1000 # Keep it short for testing
  49. )
  50. except Exception as e:
  51. print(f"大纲生成失败: {e}")
  52. # import traceback
  53. # traceback.print_exc()
  54. return
  55. end_time = time.time()
  56. print(f"大纲生成耗时: {end_time - start_time:.2f} 秒。")
  57. print(f"生成的大纲 Note ID: {outline_note_id}")
  58. print("大纲内容预览 (前500字符):")
  59. print("-" * 30)
  60. print(outline_content[:500] + "...")
  61. print("-" * 30)
  62. # ---------------------------------------------------------
  63. # Test Chapter Generate Agent
  64. # ---------------------------------------------------------
  65. print_step("3. 初始化 ChapterGenerateAgent (章节生成Agent)")
  66. try:
  67. chapter_agent = ChapterGenerateAgent(
  68. name="TestChapterAgent",
  69. llm=llm,
  70. max_steps=3, # Limit steps for testing
  71. chapter_length=1000 # Keep it short
  72. )
  73. print("ChapterGenerateAgent 初始化完成。")
  74. except Exception as e:
  75. print(f"ChapterGenerateAgent 初始化失败: {e}")
  76. return
  77. print_step("4. 生成第一章 (Generate Chapter 1)")
  78. print("调用 chapter_agent.run 生成第一章...")
  79. start_time = time.time()
  80. try:
  81. # The first run doesn't have previous chapters, so it should start fresh based on outline
  82. chapter_data, chapter_note_id = chapter_agent.run(
  83. user_input="第一章:主角醒来发现自己在代码构成的森林里。",
  84. novel_id=novel_id,
  85. novel_title=title
  86. )
  87. except Exception as e:
  88. print(f"章节生成失败: {e}")
  89. # import traceback
  90. # traceback.print_exc()
  91. return
  92. end_time = time.time()
  93. print(f"第一章生成耗时: {end_time - start_time:.2f} 秒。")
  94. print(f"生成的章节 Note ID: {chapter_note_id}")
  95. print(f"章节标题: {chapter_data.get('title')}")
  96. print(f"章节摘要: {chapter_data.get('summary')}")
  97. print("章节内容预览 (前500字符):")
  98. print("-" * 30)
  99. print(chapter_data.get('content', '')[:500] + "...")
  100. print("-" * 30)
  101. # ---------------------------------------------------------
  102. # Verification
  103. # ---------------------------------------------------------
  104. print_step("5. 验证输出文件 (Verify Output Files)")
  105. outline_path = os.path.join("outputs", f"{title}-{novel_id}", "outline")
  106. chapter_path = os.path.join("outputs", f"{title}-{novel_id}", "chapters")
  107. print(f"检查大纲目录: {outline_path}")
  108. if os.path.exists(outline_path) and os.listdir(outline_path):
  109. print("PASS: 大纲目录存在且不为空。")
  110. else:
  111. print("FAIL: 大纲目录缺失或为空。")
  112. print(f"检查章节目录: {chapter_path}")
  113. if os.path.exists(chapter_path) and os.listdir(chapter_path):
  114. print("PASS: 章节目录存在且不为空。")
  115. else:
  116. print("FAIL: 章节目录缺失或为空。")
  117. print("\n" + "="*60)
  118. print("测试流程结束")
  119. print("="*60)
  120. if __name__ == "__main__":
  121. main()