(LangChain) LangGraph Tutorial - Customizing State

Customizing State

LangGraph는 모델과 툴의 상태를 State로 정의하고, 이를 통해 상태의 변화를 추적합니다.
State의 변화 관리를 통해서 프로젝트의 요구 사항에 맞춰 상태를 효율적으로 관리하고 추적할 수 있게 해줍니다. 이를 통해 다양한 노드와 툴(Tool)이 상호작용하는 과정에서 발생할 수 있는 문제를 예방하고, 더 나은 사용자 경험을 제공할 수 있습니다.

Implementation

여기서는 이전 Agent의 내용을 바탕으로 챗봇이 검색 도구를 사용하여 특정 정보를 찾고 검토를 위해 사람에게 전달하는 시나리오를 구현합니다.

1. State에 키 추가

특정 정보를 찾아 State의 수정하기 위해서, 필요한 정보를 용이하게 접근할 수 있도록 state에 key 값을 추가합니다. 여기서는 name 과 birthday 키 값을 추가합니다.

# state.py


from typing import TypedDict, Annotated, Sequence
from langgraph.graph.message import add_messages
from langchain_core.messages import BaseMessage

class AgentState(TypedDict):
    messages: Annotated[Sequence[BaseMessage], add_messages]
    name: str
    birthday: str

2. State 업데이트

LangGraph에서 Human-in-the-loop 상호작용을 위한 human_assistance 도구를 다음과 같이 개선합니다.

✅ Call ID 기반의 State 업데이트

• Call ID를 활용하여 업데이트할 정확한 State를 찾도록 수정
• 각 실행 인스턴스를 구분하고 올바른 State 업데이트 보장

✅ 응답 정확도 판별 및 반환 값 변경

• y로 시작하는 응답("y")이면 결과를 반환, 아닌 경우에는 수정된 답변으로 반환
• 응답에서 "correct" 키가 없을 경우 기본값 빈 문자열 ("")을 반환

✅ InjectedToolCallId 사용하여 tool_call_id 관리

• LangChain 내부에서 tool_call_id를 자동 관리하도록 변경
• 각 툴의 실행 결과가 섞이지 않도록 분리하여 후속 작업의 혼선 방지

✅ Command를 활용한 State 업데이트

• State 변경 사항을 명확하게 추적 가능하도록 Command 사용

# tools.py

# 도구 정의
@tool
def human_assistance(name: str, birthday: str, tool_call_id: Annotated[str, InjectedToolCallId]) -> str:
    """Request assistance from a human."""
    human_response = interrupt({
        "question": "Is this correct?",
        "name": name,
        "birthday": birthday,
    })

    if human_response.get("correct", "").lower().startswith("y"):
        verified_name = name
        verified_birthday = birthday
        response = "Correct"
    else:
        verified_name = human_response.get("name", name)
        verified_birthday = human_response.get("birthday", birthday)
        response = f"Made a correction: {human_response}"

    state_update = {
        "name": verified_name,
        "birthday": verified_birthday,
        "messages": [ToolMessage(response, tool_call_id=tool_call_id)]  ,
    }

    return Command(update=state_update) # 상태 업데이트 명령어 반환

3. 그래프(Graph) 실행

상태(State) 업데이트 결과를 실행하여 확인합니다.
최근에 화두가 되었던, DeepSeek에 대한 출시일정을 확인하는 프롬프트(Prompt)를 구성해서 확인 해보겠습니다.

# agent.py

# 그래프 실행
if __name__ == "__main__":

    config = {"configurable": {"thread_id": "USER_03"}}
    user_input = {"type": "user", "content": "Can you look up when DeepSeek was released? \
When you have the answer, use the human_assistance tool for review."}

    for event in graph.stream({"messages": [user_input]}, config, stream_mode="values"):
        event["messages"][-1].pretty_print()

    human_command = Command(
        resume={
            "name": "DeepSeek-R1",
            "birthday": "Jan 20, 2025"
        }
    )

    for event in graph.stream(human_command, config, stream_mode="values"):
        event["messages"][-1].pretty_print()

(결과)

================================ Human Message =================================

Can you look up when DeepSeek was released? When you have the answer, use the human_assistance tool for review.
================================== Ai Message ==================================
Tool Calls:
  tavily_search_results_json (call_JK559rFs2TUTnOBXzq397glS)
 Call ID: call_JK559rFs2TUTnOBXzq397glS
  Args:
    query: DeepSeek release date
================================= Tool Message =================================
Name: tavily_search_results_json

[{"url": "https://en.wikipedia.org/wiki/DeepSeek", "content": "It would not be used for stock trading and would be
separate from High-Flyer's financial business.[4] In May 2023, the company was launched as DeepSeek.[2] DeepSeek's 
development is funded by High-Flyer.[3] After releasing DeepSeek-V2 in May 2024 which offered strong performance for 
a low price, DeepSeek became known as the catalyst for China's AI model price war. On 29 November 2023, DeepSeek 
launched DeepSeek LLM (large language model) which scaled up to 67B parameters. In December 2024, DeepSeek-V3 was 
launched. Benchmark tests showed it outperformed Llama 3.1 and Qwen 2.5 while matching GPT-4o and Claude 3.5 Sonnet.
[10][11] DeepSeek's optimization on limited resources highlighted potential limits of US sanctions on China's AI 
development.[12] \"Meet DeepSeek Chat, China's latest ChatGPT rival with a 67B model\". \"Chinese start-up DeepSeek's
new AI model outperforms Meta, OpenAI products\"."}, {"url": "https://www.techtarget.com/whatis/feature/DeepSeek-
explained-Everything-you-need-to-know", "content": "DeepSeek, a Chinese AI firm, is disrupting the industry with its 
low-cost, open source large language models, challenging U.S. tech giants. What is DeepSeek? DeepSeek focuses on 
developing open source LLMs. The company's first model was released in November 2023. DeepSeek Since the company was 
created in 2023, DeepSeek has released a series of generative AI models. The low-cost development threatens 
the business model of U.S. tech companies that have invested billions in AI. In contrast with OpenAI, which is 
proprietary technology, DeepSeek is open source and free, challenging the revenue model of U.S. companies charging 
monthly fees for AI services. Being based in China, DeepSeek challenges U.S. technological dominance in AI."}]
================================== Ai Message ==================================
Tool Calls:
  human_assistance (call_IBFdno2pWzb8gSFgFKfO7ySw)
 Call ID: call_IBFdno2pWzb8gSFgFKfO7ySw
  Args:
    name: DeepSeek
    birthday: November 2023
================================== Ai Message ==================================
Tool Calls:
  human_assistance (call_IBFdno2pWzb8gSFgFKfO7ySw)
 Call ID: call_IBFdno2pWzb8gSFgFKfO7ySw
  Args:
    name: DeepSeek
    birthday: November 2023
================================= Tool Message =================================
Name: human_assistance

Made a correction: {'name': 'DeepSeek-R1', 'birthday': 'Jan 20, 2025'}
================================== Ai Message ==================================

DeepSeek was launched in May 2023, with its first model released in November 2023. However, there was a correction 
indicating that the specific model "DeepSeek-R1" will be released on January 20, 2025.

결과를 확인하면, 검색 툴(tavily)을 통해 얻은 birthday는 November 2023 으로 human_assistance 툴에 interrupt 가 발생하였고, 리뷰를 통해 정정된 메시지로 전달 된 것을 확인 할 수 있다.
출시일을 birthday에 매핑되어 출력되어 의아하게 생각할 수 있는데, 이는 흔히 birthday 를 생일 뜻으로만 인식할 수 있는데 사전을 검색해보면 다른 뜻도 있음을 확인 할 수 있습니다.(부족한 영어실력이..발목을..😭😭)

3. 수동으로 State 업데이트 하기

graph.update_state를 활용하여 수동으로 thread의 키값으로 기반으로 state를 업데이트 할 수 있습니다.
기존 name 필드에 적용되어 있는 값을 DeepSeek-R1 Model로 이름을 변경하겠습니다.

# agent.py

# 그래프 실행
if __name__ == "__main__":

    print("================================= Graph State =================================")
    snapshot = graph.get_state(config)
    result = {k: v for k, v in snapshot.values.items() if k in ("name", "birthday")}
    print(result)

    print("============================ Manually Updating State ==========================")
    update_state = graph.update_state(config, {"name": "DeepSeek-R1 Model"})
    print(update_state)    
    snapshot = graph.get_state(config)

    result = {k: v for k, v in snapshot.values.items() if k in ("name", "birthday")}
    print(result)

(결과)

================================= Graph State =================================
{'name': 'DeepSeek-R1', 'birthday': 'Jan 20, 2025'}
============================ Manually Updating State ==========================
{'configurable': {'thread_id': 'USER_03', 'checkpoint_ns': '', 'checkpoint_id': '1efe8419-d950-6098-8006-3d30fc5fee6b'}}
{'name': 'DeepSeek-R1 Model', 'birthday': 'Jan 20, 2025'}

일반적으로 interrupt 기능을 사용하는 것이 권장 되며, 이를 통해 Human-in-the-loop 에서 필요한 정보만 전달하고, 불필요한 상태 변경을 방지할 수 있습니다.

4. LangGraph Studio에서 확인하기

LangGraph Studio에서 보다 편리하게 테스트 해볼 수 있습니다.

동일 질문의 질의하면, 검색 툴을 통해서 결과 반환된 후 Interrupt가 발생하고, Name과 Birthday를 입력할 수 있도록 커서가 이동 하였습니다.

업데이트 할 메시지를 입력하고, Resume을 클릭하며 내용의 변경사항과 최종 chatbot이 반환하는 메시지를 확인할 수 있습니다.

Reference

• Customizing State