How to access chat model and StateGraph from @tool and manage the resource context properly?

[khteh]

Checked other resources

• I added a very descriptive title to this question.
• I searched the LangChain documentation with the integrated search.
• I used the GitHub search to find a similar question and didn't find it.

Commit to Help

• I commit to help with one of those options 👆

Example Code

Initialize the LLM chat model:

self._llm = init_chat_model(appconfig.LLM_RAG_MODEL, model_provider="ollama", base_url=appconfig.OLLAMA_URI, configurable_fields=("user_id", "graph", "email_state"), streaming=True, temperature=0).bind_tools([email_processing_tool])

The tool which the model binds to:

@tool
async def email_processing_tool(
    email: str, escalation_criteria: str,
    config: Annotated[RunnableConfig, InjectedToolArg]
) -> EmailModel:
    logging.info(f"\n=== email_processing_tool ===")
    """Extract the user's state from the conversation and update the memory."""
    graph = EmailConfiguration.from_runnable_config(config).graph
    emailState = EmailConfiguration.from_runnable_config(config).email_state
    emailState["email"] = email
    emailState["escalation_text_criteria"] = escalation_criteria
    logging.debug(f"email: {email}, escalation_criteria: {escalation_criteria}, emailState:: {emailState}")
    results = await graph.with_config(config).ainvoke(emailState)
    logging.debug(f"result: {results}")
    return results["extract"]

The StateGraph used by the @tool:

    graph_builder = StateGraph(EmailRAGState)
    graph_builder.add_node("ParseEmail", self.ParseEmail)
    graph_builder.add_node("NeedsEscalation", self.NeedsEscalation)
    graph_builder.add_edge(START, "ParseEmail")
    graph_builder.add_edge("ParseEmail", "NeedsEscalation")
    graph_builder.add_edge("NeedsEscalation", END)
    self._graph = graph_builder.compile(store=self._in_memory_store, name=self._graphName)

The agent which manages / orchestrate the workflow:

    graph_builder = StateGraph(EmailAgentState)
    graph_builder.add_node("EmailAgent", self.call_agent_model_node)
    graph_builder.add_node("EmailTools", ToolNode([email_processing_tool]))
    graph_builder.add_edge(START, "EmailAgent")
    graph_builder.add_conditional_edges(
    # if the EmailAgent node returns a tool message, your graph moves to the EmailTools node to call the respective tool.
        "EmailAgent", self.route_agent_graph_edge, ["EmailTools", END]
    )
    graph_builder.add_edge("EmailTools", "EmailAgent")
    self._agent = graph_builder.compile(store=self._in_memory_store, name=self._agentName)

Invoke the agent:

    async for step in self._agent.with_config({"graph": self._graph, "email_state": email_state, "thread_id": uuid7str()}).astream(
        {"messages": [{"role": "user", "content": message_with_criteria}]},
        stream_mode="values",
    ):
        step["messages"][-1].pretty_print()

Description

init_chat_model would bind_tools to a @tool which requires access to the model and/or another StateGraph which are all constructed in the constructor of an application. Here are my questions:

(1) How to pass the graph and/or model to the @tool? Would using configurable_fields in init_chat_model be the best option considering the performance impact it may have on every super-step call that it has to pass these large data structure around?

(2) How to manage the resource context so that they are shutdown / closed properly thoughout the lifetime of the hosting application?

System Info

langchain==0.1.20
langchain-community==0.0.38
langchain-core==0.1.53
langchain-text-splitters==0.0.2