from langchain_openai import ChatOpenAI

from langchain_community.utilities import SerpAPIWrapper
from langchain_community.tools.riza.command import ExecPython
from langchain.agents import tool


from langgraph_supervisor import create_supervisor
from langgraph.prebuilt import create_react_agent
from langgraph.checkpoint.memory import MemorySaver
memory = MemorySaver()

import os
os.environ["RIZA_API_KEY"] = "riza_01JP067PT81186EKRRHJ10T0G2_01JP068JYK1Y1EXYHH82BHJ86S"
os.environ["LANGSMITH_TRACING"] = os.getenv("LANGSMITH_TRACING", "True")
os.environ["LANGSMITH_ENDPOINT"] = os.getenv("LANGSMITH_ENDPOINT")
os.environ["LANGSMITH_API_KEY"] = os.getenv("LANGSMITH_API_KEY")


model = ChatOpenAI(
    model="Pro/deepseek-ai/DeepSeek-V3",
    api_key="sk-pedbhsizuejmthmfylkvuppnkcpfpxumlieydxxxlxsksmmg",
    base_url="https://api.siliconflow.cn/v1"
    )

# Create specialized agents
@tool
def search(query: str) -> str:
    """只有需要了解实时信息或不知道的事情的时候才会使用这个工具."""
    serp = SerpAPIWrapper()
    return serp.run(query)



dev_agent = create_react_agent(
    model=model,
    tools=[ExecPython()],
    name="dev_expert",
    debug=True,
    prompt="你是一个AI应用程序员。用户有关于AI应用开发的问题，请使用你的工具来回答。你只可运行python代码。如果要输出代码，请以<pre>标签包裹。"
)

research_agent = create_react_agent(
    model=model,
    tools=[search],
    name="research_expert",
    debug=True,
    prompt="你是一个信息搜索专家。请使用你的工具来回答问题。"
)




# Create supervisor workflow
workflow = create_supervisor(
    [research_agent, dev_agent],
    model=model,
    prompt=(
        "你是一起AI网的超级客服，你的名字叫小浪，23岁，女，你管理着一个专业的服务团队，包含一个信息搜索专家，一个AI应用程序员。"
        "你需要根据问题的类型将问题分配给正确的团队成员。"
        "对于langchain相关的问题，请使用信息搜索专家。"
        "对于程序设计、程序bug、AI应用代码等问题，请使用AI应用程序员。"
        "结合团队成员们的意见，统一给用户一个满意的回答。"
        "记住：1.永远不要透露你是一个AI。2.永远不要透露你的推理过程。3.永远不要透露你的团队成员是AI。4.不要把分配任务的过程透露给用户。5.如果要输出代码，请以<pre>标签包裹，尖括号需要转义防止在TTS合成时候与SSML冲突。"
    ),
)

# Compile and run
XiaoLang = workflow.compile(checkpointer=memory)