[PYTHON/LANGGRAPH] StateGraph 클래스 : compile 메소드에서 checkpointer 인자를 사용해 챗봇 메모리 기능 추가하기

■ StateGraph 클래스의 compile 메소드에서 checkpointer 인자를 사용해 챗봇 메모리 기능을 추가하는 방법을 보여준다.

※ OPENAI_API_KEY 환경 변수 값은 .env 파일에 정의한다.
※ TAVILY_API_KEY 환경 변수 값은 .env 파일에 정의한다.

▶ main.py


import json

from dotenv                                  import load_dotenv
from langchain_openai                        import ChatOpenAI
from langchain_community.tools.tavily_search import TavilySearchResults
from typing_extensions                       import TypedDict
from typing                                  import Annotated
from langgraph.graph.message                 import add_messages
from langgraph.graph                         import StateGraph
from langchain_core.messages                 import ToolMessage
from langgraph.graph                         import END
from langgraph.graph                         import START
from langgraph.checkpoint.memory             import MemorySaver

load_dotenv()

chatOpenAI = ChatOpenAI(model = "gpt-4o-mini")

tavilySearchResults = TavilySearchResults(max_results = 2)

toolList = [tavilySearchResults]

runnableBinding = chatOpenAI.bind_tools(toolList)

class State(TypedDict):
    messageList : Annotated[list, add_messages]

stateGraph = StateGraph(State)

def chat(state : State):
    messageList = state["messageList"]
    return {"messageList" : [runnableBinding.invoke(messageList)]}

stateGraph.add_node("chatbot_node", chat)

class BasicToolNode:
    """A node that runs the tools requested in the last AIMessage."""

    def __init__(self, toolList : list) -> None:
        self.toolNameDictionary = {tool.name : tool for tool in toolList}

    def __call__(self, inputDictionary : dict):
        if messageList := inputDictionary.get("messageList", []):
            message = messageList[-1]
        else:
            raise ValueError("No message found in input")
        toolMessageList = []
        for toolCall in message.tool_calls:
            toolResultDictionaryList = self.toolNameDictionary[toolCall["name"]].invoke(toolCall["args"])
            toolMessageList.append(ToolMessage(content = json.dumps(toolResultDictionaryList), name = toolCall["name"], tool_call_id = toolCall["id"]))
        return {"messageList" : toolMessageList}

basicToolNode = BasicToolNode(toolList = [tavilySearchResults])

stateGraph.add_node("tool_node", basicToolNode)

def routeToolNode(state : State):
    """
    Use in the conditional_edge to route to the ToolNode if the last message has tool calls. Otherwise, route to the end.
    """
    if isinstance(state, list):
        aiMessage = state[-1]
    elif messages := state.get("messageList", []):
        aiMessage = messages[-1]
    else:
        raise ValueError(f"No messages found in input state to tool_edge: {state}")
    if hasattr(aiMessage, "tool_calls") and len(aiMessage.tool_calls) > 0:
        return "tool_node"
    return END

# `tools_condition` 함수는 챗봇이 도구를 사용하라고 요청하면 "tool_node"를 반환하고, 직접 응답해도 괜찮으면 "END"를 반환합니다.
# 이 조건부 라우팅은 주요 에이전트 루프를 정의한다.
stateGraph.add_conditional_edges(
    "chatbot_node",
    routeToolNode,
    # 다음 사전은 그래프에 조건의 출력을 특정 노드로 해석하도록 지시할 수 있습니다.
    # 기본적으로 항등 함수로 설정되지만 "tool_node"가 아닌 다른 이름의 노드를 사용하려면 사전의 값을 다른 것으로 업데이트할 수 있습니다. 예 : "tool_node" : "test_node"
    {"tool_node" : "tool_node", END : END}
)

# 도구가 호출될 때마다 우리는 챗봇으로 돌아가 다음 단계를 결정한다.
stateGraph.add_edge(START      , "chatbot_node")
stateGraph.add_edge("tool_node", "chatbot_node")

memorySaver = MemorySaver()

compiledStateGraph = stateGraph.compile(checkpointer = memorySaver)

def query(userInput : str):
    for addableUpdatesDict in compiledStateGraph.stream({"messageList": [("user", userInput)]}, {"configurable" : {"thread_id" : "1"}}):
        if "chatbot_node" in addableUpdatesDict:
            for valueDictionary in addableUpdatesDict.values():
                content = valueDictionary["messageList"][-1].content
                if content:
                    print("Assistant :", content)

while True:
    userInput = input("User : ")
    if userInput.lower() in ["quit", "exit", "q"]:
        print("Goodbye!")
        break
    query(userInput)

100

101

102

103

104

import json

from dotenv import load_dotenv

from langchain_openai import ChatOpenAI

from langchain_community.tools.tavily_search import TavilySearchResults

from typing_extensions import TypedDict

from typing import Annotated

from langgraph.graph.message import add_messages

from langgraph.graph import StateGraph

from langchain_core.messages import ToolMessage

from langgraph.graph import END

from langgraph.graph import START

from langgraph.checkpoint.memory import MemorySaver

load_dotenv()

chatOpenAI = ChatOpenAI(model = "gpt-4o-mini")

tavilySearchResults = TavilySearchResults(max_results = 2)

toolList = [tavilySearchResults]

runnableBinding = chatOpenAI.bind_tools(toolList)

class State(TypedDict):

messageList : Annotated[list, add_messages]

stateGraph = StateGraph(State)

def chat(state : State):

messageList = state["messageList"]

return {"messageList" : [runnableBinding.invoke(messageList)]}

stateGraph.add_node("chatbot_node", chat)

class BasicToolNode:

"""A node that runs the tools requested in the last AIMessage."""

def __init__(self, toolList : list) -> None:

self.toolNameDictionary = {tool.name : tool for tool in toolList}

def __call__(self, inputDictionary : dict):

if messageList := inputDictionary.get("messageList", []):

message = messageList[-1]

else:

raise ValueError("No message found in input")

toolMessageList = []

for toolCall in message.tool_calls:

toolResultDictionaryList = self.toolNameDictionary[toolCall["name"]].invoke(toolCall["args"])

toolMessageList.append(ToolMessage(content = json.dumps(toolResultDictionaryList), name = toolCall["name"], tool_call_id = toolCall["id"]))

return {"messageList" : toolMessageList}

basicToolNode = BasicToolNode(toolList = [tavilySearchResults])

stateGraph.add_node("tool_node", basicToolNode)

def routeToolNode(state : State):

"""

Use in the conditional_edge to route to the ToolNode if the last message has tool calls. Otherwise, route to the end.

"""

if isinstance(state, list):

aiMessage = state[-1]

elif messages := state.get("messageList", []):

aiMessage = messages[-1]

else:

raise ValueError(f"No messages found in input state to tool_edge: {state}")

if hasattr(aiMessage, "tool_calls") and len(aiMessage.tool_calls) > 0:

return "tool_node"

return END

# `tools_condition` 함수는 챗봇이 도구를 사용하라고 요청하면 "tool_node"를 반환하고, 직접 응답해도 괜찮으면 "END"를 반환합니다.

# 이 조건부 라우팅은 주요 에이전트 루프를 정의한다.

stateGraph.add_conditional_edges(

"chatbot_node",

routeToolNode,

# 다음 사전은 그래프에 조건의 출력을 특정 노드로 해석하도록 지시할 수 있습니다.

# 기본적으로 항등 함수로 설정되지만 "tool_node"가 아닌 다른 이름의 노드를 사용하려면 사전의 값을 다른 것으로 업데이트할 수 있습니다. 예 : "tool_node" : "test_node"

{"tool_node" : "tool_node", END : END}

)

# 도구가 호출될 때마다 우리는 챗봇으로 돌아가 다음 단계를 결정한다.

stateGraph.add_edge(START , "chatbot_node")

stateGraph.add_edge("tool_node", "chatbot_node")

memorySaver = MemorySaver()

compiledStateGraph = stateGraph.compile(checkpointer = memorySaver)

def query(userInput : str):

for addableUpdatesDict in compiledStateGraph.stream({"messageList": [("user", userInput)]}, {"configurable" : {"thread_id" : "1"}}):

if "chatbot_node" in addableUpdatesDict:

for valueDictionary in addableUpdatesDict.values():

content = valueDictionary["messageList"][-1].content

if content:

print("Assistant :", content)

while True:

userInput = input("User : ")

if userInput.lower() in ["quit", "exit", "q"]:

print("Goodbye!")

break

query(userInput)

▶ requirements.txt


aiohappyeyeballs==2.4.4
aiohttp==3.11.11
aiosignal==1.3.2
annotated-types==0.7.0
anyio==4.7.0
attrs==24.3.0
certifi==2024.12.14
charset-normalizer==3.4.1
colorama==0.4.6
dataclasses-json==0.6.7
distro==1.9.0
frozenlist==1.5.0
greenlet==3.1.1
h11==0.14.0
httpcore==1.0.7
httpx==0.28.1
httpx-sse==0.4.0
idna==3.10
jiter==0.8.2
jsonpatch==1.33
jsonpointer==3.0.0
langchain==0.3.13
langchain-community==0.3.13
langchain-core==0.3.28
langchain-openai==0.2.14
langchain-text-splitters==0.3.4
langgraph==0.2.60
langgraph-checkpoint==2.0.9
langgraph-sdk==0.1.48
langsmith==0.2.6
marshmallow==3.23.2
msgpack==1.1.0
multidict==6.1.0
mypy-extensions==1.0.0
numpy==2.2.1
openai==1.58.1
orjson==3.10.12
packaging==24.2
propcache==0.2.1
pydantic==2.10.4
pydantic-settings==2.7.0
pydantic_core==2.27.2
python-dotenv==1.0.1
PyYAML==6.0.2
regex==2024.11.6
requests==2.32.3
requests-toolbelt==1.0.0
sniffio==1.3.1
SQLAlchemy==2.0.36
tavily-python==0.5.0
tenacity==9.0.0
tiktoken==0.8.0
tqdm==4.67.1
typing-inspect==0.9.0
typing_extensions==4.12.2
urllib3==2.3.0
yarl==1.18.3

aiohappyeyeballs==2.4.4

aiohttp==3.11.11

aiosignal==1.3.2

annotated-types==0.7.0

anyio==4.7.0

attrs==24.3.0

certifi==2024.12.14

charset-normalizer==3.4.1

colorama==0.4.6

dataclasses-json==0.6.7

distro==1.9.0

frozenlist==1.5.0

greenlet==3.1.1

h11==0.14.0

httpcore==1.0.7

httpx==0.28.1

httpx-sse==0.4.0

idna==3.10

jiter==0.8.2

jsonpatch==1.33

jsonpointer==3.0.0

langchain==0.3.13

langchain-community==0.3.13

langchain-core==0.3.28

langchain-openai==0.2.14

langchain-text-splitters==0.3.4

langgraph==0.2.60

langgraph-checkpoint==2.0.9

langgraph-sdk==0.1.48

langsmith==0.2.6

marshmallow==3.23.2

msgpack==1.1.0

multidict==6.1.0

mypy-extensions==1.0.0

numpy==2.2.1

openai==1.58.1

orjson==3.10.12

packaging==24.2

propcache==0.2.1

pydantic==2.10.4

pydantic-settings==2.7.0

pydantic_core==2.27.2

python-dotenv==1.0.1

PyYAML==6.0.2

regex==2024.11.6

requests==2.32.3

requests-toolbelt==1.0.0

sniffio==1.3.1

SQLAlchemy==2.0.36

tavily-python==0.5.0

tenacity==9.0.0

tiktoken==0.8.0

tqdm==4.67.1

typing-inspect==0.9.0

typing_extensions==4.12.2

urllib3==2.3.0

yarl==1.18.3

※ pip install python-dotenv langchain_community langchain_openai langgraph tavily-python 명령을 실행했다.

AI LANGGRAPH LLM OPENAI PYTHON TAVILY TOOL

icodebroker

[PYTHON/LANGGRAPH] StateGraph 클래스 : compile 메소드에서 checkpointer 인자를 사용해 챗봇 메모리 기능 추가하기

분류

가장 많이 읽힌 글

보관함