"""
生成LangGraph Mermaid图表的脚本
使用LangGraph内置的draw_mermaid()方法
"""

import os
import sys

# 添加项目根目录到Python路径
# 从当前文件位置向上找到项目根目录 (Retriver/)
project_root = os.path.dirname(os.path.dirname(os.path.dirname(__file__)))
sys.path.insert(0, project_root)

def generate_mermaid_graph():
    """使用LangGraph内置方法生成Mermaid图表"""
    
    try:
        # 导入必要的模块
        from retriver.langgraph.iterative_retriever import IterativeRetriever
        
        print("初始化检索器...")
        # 创建一个最小化配置的迭代检索器实例
        retriever = IterativeRetriever(
            keyword="temp",
            top_k=2,
            max_iterations=3,
            max_parallel_retrievals=2
        )
        
        print("生成LangGraph Mermaid图表...")
        # 使用LangGraph内置的draw_mermaid()方法
        mermaid_graph = retriever.workflow.get_graph().draw_mermaid()
        
        # 保存到Graph.txt文件
        output_file = os.path.join(os.path.dirname(__file__), "Graph.txt")
        with open(output_file, 'w', encoding='utf-8') as f:
            f.write(mermaid_graph)
        
        print(f"LangGraph Mermaid图表已保存到: {output_file}")
        print("\n图表内容:")
        print(mermaid_graph)
        
        return mermaid_graph
        
    except ImportError as e:
        print(f"模块导入失败: {e}")
        print("请确保在正确的Python环境中运行此脚本")
        return None
    except Exception as e:
        print(f"生成Mermaid图表失败: {e}")
        print("可能是因为缺少必要的依赖或配置")
        return None


if __name__ == "__main__":
    generate_mermaid_graph()