"""
创建一个静态模型的智能体
"""
import json

from langchain.agents import create_agent
from langchain.chat_models import init_chat_model
from langchain.tools import tool
from langgraph.graph.state import CompiledStateGraph
from env_utils import DEEPSEEK_API_KEY, DEEPSEEK_BASE_URL

from langchain_core.messages import AnyMessage

"""
基本控制: 串行控制
基本控制: 分支控制
Successfully installed langchain-core-1.4.0 
langchain-protocol-0.0.15 langgraph-1.2.0 
langgraph-checkpoint-4.1.0 langgraph-prebuilt-1.1.0 ormsgpack-1.12.2

"""

from langchain_core.messages import AnyMessage
from typing_extensions import TypedDict

from langgraph.graph import START,StateGraph,END

from IPython.display import Image,display
from typing import Any, Annotated, Literal
import operator
class State(TypedDict):
   aggregate: Annotated[list,operator.add]

def a(state:State):
    print(f"Node A sees {state['aggregate']}")
    return {"aggregate":["A"]}

def b(state:State):
    print(f"Node B sees {state['aggregate']}")
    return {"aggregate": ["B"]}

builder=StateGraph(State)
builder.add_node(a)
builder.add_node(b)

def route(state: State) -> Literal["b",END]:
    if len(state["aggregate"])<7:
        return "b"
    else:
        return END

builder.add_edge(START,"a")
builder.add_conditional_edges("a",route)
builder.add_edge("b","a")

graph = builder.compile()

# from IPython.display import Image, display
#
# with open('graph.png', 'wb') as f:
#     f.write(graph.get_graph().draw_mermaid_png())
# display(Image('graph.png'))

from PIL import Image
import io

# 假设这是你的 PNG 字节数据
png_data = graph.get_graph().draw_mermaid_png()

if 'png_data' in locals():
    # 将字节数据解码为图片a
    img = Image.open(io.BytesIO(png_data))
    img.show()

res=graph.invoke({"value_1": "c"})

print(res)

def route(state: State) -> Literal["b",END]:
    if len(state["aggregate"])<7:
        return "b"
    else:
        return END

builder.add_edge(START,"a")
builder.add_conditional_edges("a",route)
builder.add_edge("b","a")

条件边, 如果返回b，就是a->b,b->a，循环， 直到b->END