from langgraph.graph import StateGraph
from typing_extensions import TypedDict
from langgraph.graph import START, END
from dotenv import dotenv_values
from langchain_openai import ChatOpenAI


env_vars = dotenv_values('.env')
OPENAI_KEY = env_vars['OPENAI_API_KEY'] 
OPENAI_BASE_URL = env_vars['OPENAI_API_BASE'] 


# 定义输入的模式
class InputState(TypedDict):
    question: str

# 定义输出的模式
class OutputState(TypedDict):
    answer: str

# 将 InputState 和 OutputState 这两个 TypedDict 类型合并成一个更全面的字典类型。
class OverallState(InputState, OutputState):
    pass

def llm_node(state: InputState):
    messages = [
        # ("system","You are a helpful assistant"),
        ("user", state["question"])
    ]
    llm = ChatOpenAI(model="gpt-3.5-turbo", api_key=OPENAI_KEY,base_url=OPENAI_BASE_URL)
    # gpt-3.5-turbo o1-mini
    response = llm.invoke(messages) 
    print(123, "response", response)
    return {"answer": response.content}

# 明确指定它的输入和输出数据的结构或模式
builder = StateGraph(OverallState, input_schema=InputState, output_schema=OutputState)

# 添加节点
builder.add_node("llm_node", llm_node)

# 添加边
builder.add_edge(START, "llm_node")
builder.add_edge("llm_node", END)

# 编译图
graph = builder.compile()

final_answer = graph.invoke({"question":"how are you"})
print(final_answer["answer"])