class Tool:
    name: str          # 工具的唯一标识符
    description: str   # 功能的详细描述，用于LLM理解何时使用
    func: Callable     # 实际执行功能的函数

from langchain.memory import ConversationBufferMemory
memory = ConversationBufferMemory()
memory.save_context({"input": "你好"}, {"output": "你好！我是AI助手。"})
memory.save_context({"input": "你叫什么名字？"}, {"output": "我叫minhuan。"})
# 查看内存内容
print(memory.load_memory_variables({}))
# 输出: {'history': 'Human: 你好\nAI: 你好！我是AI助手。\nHuman: 你叫什么名字？\nAI: 我叫minhuan。'}

from langchain.memory import ConversationBufferWindowMemory
# 只保留最近2轮对话
memory = ConversationBufferWindowMemory(k=2)
memory.save_context({"input": "你好"}, {"output": "你也好"})
memory.save_context({"input": "高兴遇见你"}, {"output": "也很高兴遇见你"})
memory.save_context({"input": "再见！"}, {"output": "好的，再见！"})
print(memory.load_memory_variables({}))
# 输出: {'history': 'Human: 高兴遇见你\nAI: 也很高兴遇见你\nHuman: 再见！\nAI: 好的，再见！'}

from langchain.memory import ConversationSummaryMemory
from langchain_community.chat_models.tongyi import ChatTongyi
llm = ChatTongyi(model="qwen-max", dashscope_api_key="你的API_KEY")
memory = ConversationSummaryMemory(llm=llm)
# 添加多轮对话
for i in range(5):
    memory.save_context(
        {"input": f"问题{i}: 关于天气的讨论"},
        {"output": f"回答{i}: 今天天气晴朗，温度适宜"}
    )
print(memory.load_memory_variables({}))
# 输出: {'history': '摘要: 对话主要讨论了天气情况，多次提到今天天气晴朗且温度适宜。'}

from langchain.memory import ConversationEntityMemory
from langchain.memory.prompt import ENTITY_EXTRACTION_PROMPT, ENTITY_SUMMARIZATION_PROMPT
memory = ConversationEntityMemory(llm=llm)
# 添加包含实体的对话
memory.save_context(
    {"input": "张三喜欢在北京工作"}, 
    {"output": "北京是个好城市"}
)
memory.save_context(
    {"input": "李四在上海读书"}, 
    {"output": "上海有很多好大学"}
)
# 查看记忆的实体
print(memory.load_memory_variables({"input": "张三在哪里工作？"}))
# 输出会包含实体信息，如: {'history': '...', 'entities': {'张三': '张三喜欢在北京工作'}}

from langchain.memory import VectorStoreRetrieverMemory
from langchain_community.vectorstores import FAISS
from langchain_community.embeddings import HuggingFaceEmbeddings
from langchain.schema import Document
import os
# 初始化嵌入模型
# embeddings = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2")
embeddings = HuggingFaceEmbeddings(model_name="D:/modelscope/hub/models/sentence-transformers/paraphrase-MiniLM-L6-v2")
# 创建或加载 FAISS 向量存储
persist_directory = "./faiss_db"
if os.path.exists(persist_directory):
    # 如果存在已有的 FAISS 索引，则加载
    vectorstore = FAISS.load_local(
        persist_directory, 
        embeddings, 
        allow_dangerous_deserialization=True
    )
    print("已加载现有的 FAISS 向量数据库")
else:
    # 否则创建新的向量存储
    # 初始时可以添加一些空文档或示例文档
    initial_docs = [Document(page_content="初始文档", metadata={"type": "initial"})]
    vectorstore = FAISS.from_documents(initial_docs, embeddings)
    # 保存向量存储到本地
    vectorstore.save_local(persist_directory)
    print("创建了新的 FAISS 向量数据库")
# 创建基于 FAISS 向量存储的记忆系统
memory = VectorStoreRetrieverMemory(
    retriever=vectorstore.as_retriever(
        search_type="similarity",
        search_kwargs={"k": 3}  # 检索最相关的3条记忆
    )
)
# 保存记忆
memory.save_context(
    {"input": "我最喜欢的颜色是蓝色"}, 
    {"output": "好的，已记住你喜欢的颜色"}
)
# 保存更多记忆示例
memory.save_context(
    {"input": "我来自北京"}, 
    {"output": "北京是中国的首都，很繁华的城市"}
)
memory.save_context(
    {"input": "我喜欢吃苹果和香蕉"}, 
    {"output": "水果很健康，已记住你的喜好"}
)
# 回忆记忆 - 基于语义相似性搜索
query = "我喜欢什么颜色？"
memories = memory.load_memory_variables({"input": query})
print(f"查询: '{query}'")
print("检索到的记忆:", memories)
# 测试另一个查询
query = "我来自哪里？"
memories = memory.load_memory_variables({"input": query})
print(f"\n查询: '{query}'")
print("检索到的记忆:", memories)
# 保存向量存储的更改
vectorstore.save_local(persist_directory)
print("\n记忆已保存到FAISS数据库")

def check_and_install_dependencies():
    """检查并安装必要的依赖"""
    required_packages = [
        "langchain",
        "langchain-community",
        "langchain-core",
        "dashscope",
        "wikipedia",
        "pydantic>=2.0.0",
        "typing-extensions>=4.0.0"
    ]
    # ... 检查并安装缺失的包

def initialize_llm():
    """初始化Qwen语言模型"""
    try:
        # 获取API密钥，如果环境变量不存在则提示用户
        api_key = os.environ.get("DASHSCOPE_API_KEY")
        if not api_key:
            print_formatted_log("WARNING", "环境变量 DASHSCOPE_API_KEY 未设置")
            print_formatted_log("INFO", "请输入通义千问API密钥:")
            api_key = input().strip()
            # ... 设置环境变量
        
        # 初始化模型
        llm = ChatTongyi(
            model="qwen-max",
            dashscope_api_key=api_key
        )
        return llm
    except Exception as e:
        logger.error(f"初始化语言模型失败: {str(e)}")
        raise

def initialize_tools(llm):
    """初始化所有可用工具"""
    # 初始化 Wikipedia 工具（带错误处理）
    try:
        wikipedia_api_wrapper = WikipediaAPIWrapper()
        wikipedia_tool = WikipediaQueryRun(api_wrapper=wikipedia_api_wrapper)
    except ImportError as e:
        logger.warning(f"无法加载 Wikipedia 工具: {str(e)}")
        # 创建替代工具
    
    # 自定义工具
    custom_tools = [
        Tool(
            name="SquareRootCalculator",
            func=calculate_square_root,
            description="当需要计算一个数的平方根时使用此工具..."
        ),
        Tool(
            name="StringReverser",
            func=reverse_string,
            description="当需要反转一个字符串时使用此工具。输入应该是一个字符串。"
        ),  
        Tool(
            name="TimeGetter",
            func=get_current_time,
            description="当需要获取当前日期和时间时使用此工具。输入可以是空字符串或任何内容。"
        )
        # ... 其他工具
    ]
    
    # 组合所有工具
    tools = []
    if wikipedia_tool:
        tools.append(wikipedia_tool)
    tools.extend(custom_tools)
    return tools

# 2. 自定义工具函数
def calculate_square_root(input_str):
    """计算一个数的平方根"""
    logger.info(f"调用 calculate_square_root 工具，输入: {input_str}")
    try:
        number = float(input_str)
        result = math.sqrt(number)
        logger.info(f"计算平方根成功: {number} 的平方根是 {result}")
        return result
    except ValueError:
        error_msg = f"输入 '{input_str}' 不是有效的数字"
        logger.error(error_msg)
        return error_msg
def reverse_string(input_str):
    """反转字符串"""
    logger.info(f"调用 reverse_string 工具，输入: {input_str}")
    result = input_str[::-1]
    logger.info(f"字符串反转成功: '{input_str}' 反转为 '{result}'")
    return result

def initialize_memory():
    """初始化对话记忆"""
    memory = ConversationBufferMemory(
        memory_key="chat_history", 
        output_key="output",  # 显式设置输出键
        return_messages=True
    )
    return memory

def create_agent(llm, tools, memory):
    """创建React Agent"""
    # ReAct 提示模板
    template = """
你是一个多功能AI助手，可以使用多种工具解决问题...
可用的工具:
{tools}
使用格式:
Thought: 我需要使用工具吗？...
Action: 要执行的动作...
Action Input: 工具的输入...
Observation: 工具执行的结果...
...
"""
    
    # 创建提示模板
    prompt = PromptTemplate(
        template=template,
        input_variables=["chat_history", "input", "agent_scratchpad"],
        partial_variables={
            "tool_names": ", ".join([tool.name for tool in tools]),
            "tools": "\n".join([f"{tool.name}: {tool.description}" for tool in tools])
        }
    )
    
    # 创建代理执行器
    agent_executor = AgentExecutor(
        agent=create_react_agent(llm, tools, prompt),
        tools=tools,
        memory=memory,
        verbose=True,
        handle_parsing_errors=True,
        max_iterations=5,  # 防止无限循环
        return_intermediate_steps=True  # 避免回调错误
    )
    return agent_executor

class CustomCallbackHandler(BaseCallbackHandler):
    """自定义回调处理器，用于处理LangChain的回调事件"""
    
    def on_chain_start(self, serialized, inputs, **kwargs):
        """链开始时的回调"""
        # 添加对 serialized 为 None 的检查
        if serialized is not None and hasattr(serialized, "get"):
            chain_type = serialized.get("name", "未知链")
            logger.info(f"开始执行链: {chain_type}")
    
    # ... 其他回调方法（on_tool_start, on_llm_start等）

def main():
    """主函数"""
    # 初始化所有组件
    llm = initialize_llm()
    tools = initialize_tools(llm)
    memory = initialize_memory()
    agent_executor = create_agent(llm, tools, memory)
    
    # 主循环
    while True:
        user_input = input("\nYou: ").strip()
        
        if user_input.lower() == 'exit':
            break  # 退出程序
        
        if user_input.lower() == 'history':
            # 显示对话历史
            history = memory.load_memory_variables({})
            if 'chat_history' in history and history['chat_history']:
                print_formatted_log("INFO", "对话历史:")
                for msg in history['chat_history']:
                    role = "Human" if msg.type == "human" else "AI"
                    print_formatted_log("DEBUG", f"  {role}: {msg.content}")
            continue
        
        # 处理用户请求
        try:
            response = agent_executor.invoke(
                {"input": user_input},
                {"callbacks": CallbackManager([CustomCallbackHandler()])}
            )
            
            if "output" in response:
                print_formatted_log("SUCCESS", f"Assistant: {response['output']}")
        except Exception as e:
            logger.error(f"执行Agent时出错: {str(e)}")
            print_formatted_log("ERROR", f"执行失败: {str(e)}")