from typing import TypedDict, Annotated, Sequencefrom langchain.chat_models import init_chat_modelfrom langchain_core.messages import BaseMessage, HumanMessage, SystemMessagefrom langgraph.graph import StateGraph, START, ENDfrom langgraph.graph.message import add_messagesfrom langgraph.prebuilt import ToolNode, tools_condition# 상태 정의class AgentState(TypedDict): messages: Annotated[Sequence[BaseMessage], add_messages]# LLM에 도구 바인딩llm = init_chat_model("gpt-4o-mini", temperature=0)llm_with_tools = llm.bind_tools(tools)SYSTEM_PROMPT = """당신은 질문에 답변하는 AI 어시스턴트입니다.주어진 도구를 활용하여 정확한 정보를 수집한 뒤 답변하세요.- search_documents: 내부 문서에서 정보를 찾을 때- web_search: 최신 정보나 내부 문서에 없는 정보가 필요할 때- calculator: 수학 계산이 필요할 때충분한 정보를 수집한 후에 답변하세요."""# 에이전트 노드def agent(state: AgentState) -> AgentState: messages = [SystemMessage(content=SYSTEM_PROMPT)] + state["messages"] response = llm_with_tools.invoke(messages) return {"messages": [response]}# 그래프 구성workflow = StateGraph(AgentState)workflow.add_node("agent", agent)workflow.add_node("tools", ToolNode(tools))workflow.add_edge(START, "agent")workflow.add_conditional_edges("agent", tools_condition)workflow.add_edge("tools", "agent")app = workflow.compile()
# 문서 검색 활용result = app.invoke({ "messages": [HumanMessage(content="Self-RAG의 Reflection Token에 대해 설명해줘")]})print(result["messages"][-1].content)# 웹 검색 활용result = app.invoke({ "messages": [HumanMessage(content="2024년 주목받은 RAG 논문은?")]})print(result["messages"][-1].content)# 복합 도구 활용 (검색 + 계산)result = app.invoke({ "messages": [HumanMessage(content="RAG 논문이 2023년 대비 2024년에 몇 배 증가했나?")]})print(result["messages"][-1].content)