LangGraph + PromptFlow KeyError __pf_main__ on instruction StateGraph(State)

[NinoCucc]

Checked other resources

• I added a very descriptive title to this issue.
• I searched the LangChain documentation with the integrated search.
• I used the GitHub search to find a similar question and didn't find it.
• I am sure that this is a bug in LangChain rather than my code.
• The bug is not resolved by updating to the latest stable version of LangChain (or the specific integration package).

Example Code

from langgraph.graph import StateGraph
from langgraph.graph.message import add_messages
from typing import Annotated, TypedDict

class State(TypedDict):
question: Optional[str]
messages: Annotated[list, add_messages]
check_ups: Optional[list]
check_ups_preparations: Optional[list]
triggered_actions: Optional[list]
triggered_actions_indexes: Optional[list]
triggered_actions_thought: Optional[str]
step_is_manual: Optional[bool]
step_has_preparation: Optional[bool]
manual_step_executed: Optional[dict]
interaction_attempt_executed: Optional[dict]

graph_builder = StateGraph(State)

Error Message and Stack Trace (if applicable)

KeyError pf_main

Description

Hi,
I am building RAG applications using Azure Promptflow and LangGraph, but the my app crashes when the StateGraph and more specifically its schema is initiated (on instructions: StateGraph(State)), then i get an error that says key error: pf_main.

The crash happens on instruction "self._add_schema(state_schema)" on StateGraph init of the script ...\Lib\site-packages\langgraph\graph\state.py of LangGraph.

I tried even with the latest versions of the libraries LangChain (0.2.17), LangGraph (0.2.61), PromptFlow (1.17.1)

It seems to be due to the addition of state_schema handling in LangGraph and the use of PromptFlow in Azure. Does anyone have any idea how to fix this?

Thank you in advance

System Info

C:\XX\pythonProject3\YY\Scripts\python.exe -m promptflow._cli._pf.entry flow test --flow c:\XX\YY\agent_assistant\flows\standard --user-agent "prompt-flow-extension/1.20.2 (win32; x64) VSCode/1.94.0"
Prompt flow service has started...