diff --git a/newrelic/hooks/mlmodel_langgraph.py b/newrelic/hooks/mlmodel_langgraph.py index f7a3380f04..a2ab62a81a 100644 --- a/newrelic/hooks/mlmodel_langgraph.py +++ b/newrelic/hooks/mlmodel_langgraph.py @@ -82,7 +82,7 @@ def wrap_AsyncBackgroundExecutor_submit(wrapped, instance, args, kwargs): return wrapped(*args, **kwargs) context = ContextOf(trace=trace, strict=True) - func = coroutine_wrapper(wrapped, context) + func = coroutine_wrapper(func, context) return wrapped(func, *args, **kwargs) diff --git a/tests/mlmodel_langchain/_mock_external_openai_server.py b/tests/mlmodel_langchain/_mock_external_openai_server.py index 172374baf1..74740ba520 100644 --- a/tests/mlmodel_langchain/_mock_external_openai_server.py +++ b/tests/mlmodel_langchain/_mock_external_openai_server.py @@ -214,6 +214,121 @@ }, ], ], + "user: What is the capital of France? Answer in one word.": [ + { + "content-type": "text/event-stream; charset=utf-8", + "openai-organization": "user-rk8wq9voijy9sejrncvgi0iw", + "openai-processing-ms": "134", + "openai-project": "proj_0Wv6taeZjWf793P67JMswYY3", + "openai-version": "2020-10-01", + "x-ratelimit-limit-requests": "10000", + "x-ratelimit-limit-tokens": "50000000", + "x-ratelimit-remaining-requests": "9999", + "x-ratelimit-remaining-tokens": "49999985", + "x-ratelimit-reset-requests": "6ms", + "x-ratelimit-reset-tokens": "0s", + "x-request-id": "req_22204b237d22427fbfd99c665d8a9964", + }, + 200, + [ + { + "id": "", + "choices": [], + "created": 0, + "model": "", + "object": "", + "service_tier": None, + "system_fingerprint": None, + "usage": None, + "prompt_filter_results": [ + { + "prompt_index": 0, + "content_filter_results": { + "hate": {"filtered": False, "severity": "safe"}, + "self_harm": {"filtered": False, "severity": "safe"}, + "sexual": {"filtered": False, "severity": "safe"}, + "violence": {"filtered": False, "severity": "safe"}, + }, + } + ], + }, + { + "id": "chatcmpl-DelITaJCJy951hwON0psdz2H9dF7i", + "choices": [ + { + "delta": { + "content": "", + "function_call": None, + "refusal": None, + "role": "assistant", + "tool_calls": None, + }, + "finish_reason": None, + "index": 0, + "logprobs": None, + "content_filter_results": {}, + } + ], + "created": 1780079082, + "model": "gpt-3.5-turbo-0125", + "object": "chat.completion.chunk", + "service_tier": "default", + "system_fingerprint": "fp_a7294185dc", + "usage": None, + "obfuscation": "LK4Bn", + }, + { + "id": "chatcmpl-DelITaJCJy951hwON0psdz2H9dF7i", + "choices": [ + { + "delta": { + "content": "Paris", + "function_call": None, + "refusal": None, + "role": None, + "tool_calls": None, + }, + "finish_reason": None, + "index": 0, + "logprobs": None, + "content_filter_results": {}, + } + ], + "created": 1780079082, + "model": "gpt-3.5-turbo-0125", + "object": "chat.completion.chunk", + "service_tier": "default", + "system_fingerprint": "fp_a7294185dc", + "usage": None, + "obfuscation": "QP", + }, + { + "id": "chatcmpl-DelITaJCJy951hwON0psdz2H9dF7i", + "choices": [ + { + "delta": { + "content": None, + "function_call": None, + "refusal": None, + "role": None, + "tool_calls": None, + }, + "finish_reason": "stop", + "index": 0, + "logprobs": None, + "content_filter_results": {}, + } + ], + "created": 1780079082, + "model": "gpt-3.5-turbo-0125", + "object": "chat.completion.chunk", + "service_tier": "default", + "system_fingerprint": "fp_a7294185dc", + "usage": None, + "obfuscation": "s", + }, + ], + ], } RESPONSES_V1 = { 'system: You are a text manipulation algorithm. | user: Use a tool to add an exclamation to the word "Hello"': [ @@ -286,7 +401,7 @@ "x-ratelimit-remaining-tokens": "49999970", "x-ratelimit-reset-requests": "6ms", "x-ratelimit-reset-tokens": "0s", - "x-request-id": "req_e9add199e2c543f1b0f1dc5318690171", + "x-request-id": "req_619548c272db4f1ab380b83de9fdedef", }, 200, { @@ -754,6 +869,87 @@ "system_fingerprint": None, }, ], + "user: What is the capital of France? Answer in one word.": [ + { + "content-type": "application/json", + "openai-organization": "user-rk8wq9voijy9sejrncvgi0iw", + "openai-processing-ms": "238", + "openai-project": "proj_0Wv6taeZjWf793P67JMswYY3", + "openai-version": "2020-10-01", + "x-ratelimit-limit-requests": "10000", + "x-ratelimit-limit-tokens": "50000000", + "x-ratelimit-remaining-requests": "9999", + "x-ratelimit-remaining-tokens": "49999985", + "x-ratelimit-reset-requests": "6ms", + "x-ratelimit-reset-tokens": "0s", + "x-request-id": "req_22204b237d22427fbfd99c665d8a9964", + }, + 200, + { + "id": "chatcmpl-DelITaJCJy951hwON0psdz2H9dF7i", + "choices": [ + { + "finish_reason": "stop", + "index": 0, + "logprobs": None, + "message": { + "content": "Paris", + "refusal": None, + "role": "assistant", + "annotations": [], + "audio": None, + "function_call": None, + "tool_calls": None, + }, + "content_filter_results": { + "hate": {"filtered": False, "severity": "safe"}, + "self_harm": {"filtered": False, "severity": "safe"}, + "sexual": {"filtered": False, "severity": "safe"}, + "violence": {"filtered": False, "severity": "safe"}, + }, + } + ], + "created": 1780077445, + "model": "gpt-3.5-turbo-0125", + "object": "chat.completion", + "service_tier": "default", + "system_fingerprint": "fp_a7294185dc", + "usage": { + "completion_tokens": 2, + "prompt_tokens": 19, + "total_tokens": 21, + "completion_tokens_details": { + "accepted_prediction_tokens": 0, + "audio_tokens": 0, + "reasoning_tokens": 0, + "rejected_prediction_tokens": 0, + }, + "prompt_tokens_details": {"audio_tokens": 0, "cached_tokens": 0}, + "latency_checkpoint": { + "engine_tbt_ms": 11, + "engine_ttft_ms": 56, + "engine_ttlt_ms": 78, + "pre_inference_ms": 108, + "service_tbt_ms": 16, + "service_ttft_ms": 196, + "service_ttlt_ms": 219, + "total_duration_ms": 120, + "user_visible_ttft_ms": 88, + }, + }, + "prompt_filter_results": [ + { + "prompt_index": 0, + "content_filter_results": { + "hate": {"filtered": False, "severity": "safe"}, + "self_harm": {"filtered": False, "severity": "safe"}, + "sexual": {"filtered": False, "severity": "safe"}, + "violence": {"filtered": False, "severity": "safe"}, + }, + } + ], + }, + ], } diff --git a/tests/mlmodel_langchain/test_pregel_executor.py b/tests/mlmodel_langchain/test_pregel_executor.py new file mode 100644 index 0000000000..64d3934e84 --- /dev/null +++ b/tests/mlmodel_langchain/test_pregel_executor.py @@ -0,0 +1,83 @@ +# Copyright 2010 New Relic, Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import asyncio +from threading import Event + +from langgraph.pregel._executor import AsyncBackgroundExecutor, BackgroundExecutor + +from newrelic.api.background_task import background_task +from newrelic.api.time_trace import current_trace + + +@background_task() +def test_background_executor_submit_propagates_context(): + trace = current_trace() + test_ran = Event() + + def task(): + assert current_trace() is trace + test_ran.set() + + with BackgroundExecutor(config={}) as submit: + future = submit(task) + future.result() + + assert test_ran.is_set() + + +@background_task() +def test_async_background_executor_submit_propagates_context(loop): + trace = current_trace() + test_ran = Event() + + async def task(): + assert current_trace() is trace + test_ran.set() + + async def _test(): + async with AsyncBackgroundExecutor(config={}) as submit: + future = submit(task) + await asyncio.wrap_future(future) + + loop.run_until_complete(_test()) + + assert test_ran.is_set() + + +def test_background_executor_submit_pass_through_outside_transaction(): + test_ran = Event() + + def task(): + test_ran.set() + + with BackgroundExecutor(config={}) as submit: + submit(task).result() + + assert test_ran.is_set() + + +def test_async_background_executor_submit_pass_through_outside_transaction(loop): + test_ran = Event() + + async def task(): + test_ran.set() + + async def _test(): + async with AsyncBackgroundExecutor(config={}) as submit: + await asyncio.wrap_future(submit(task)) + + loop.run_until_complete(_test()) + + assert test_ran.is_set() diff --git a/tests/mlmodel_langchain/test_state_graph.py b/tests/mlmodel_langchain/test_state_graph.py index 778dc22dcd..a47ad5f3d6 100644 --- a/tests/mlmodel_langchain/test_state_graph.py +++ b/tests/mlmodel_langchain/test_state_graph.py @@ -14,14 +14,16 @@ import pytest from langchain.messages import HumanMessage +from langchain.tools import tool from testing_support.fixtures import reset_core_stats_engine from testing_support.validators.validate_custom_events import validate_custom_events from newrelic.api.background_task import background_task -PROMPT = {"messages": [HumanMessage('Use a tool to add an exclamation to the word "Hello"')]} +CLIENT_PROMPT = {"messages": [HumanMessage("What is the capital of France? Answer in one word.")]} +AGENT_PROMPT = {"messages": [HumanMessage('Use a tool to add an exclamation to the word "Hello"')]} -chat_completion_recorded_events = [ +client_recorded_events = [ [ {"type": "LlmChatCompletionSummary"}, { @@ -52,8 +54,8 @@ {"type": "LlmChatCompletionMessage"}, { "completion_id": None, - "content": 'Use a tool to add an exclamation to the word "Hello"', - "id": "chatcmpl-Dd0Na8gXEDyFIhYMsL72TYk3bSZun-0", + "content": "What is the capital of France? Answer in one word.", + "id": "chatcmpl-DelITaJCJy951hwON0psdz2H9dF7i-0", "ingest_source": "Python", "request_id": "req_22204b237d22427fbfd99c665d8a9964", "response.model": "gpt-3.5-turbo-0125", @@ -69,8 +71,8 @@ {"type": "LlmChatCompletionMessage"}, { "completion_id": None, - "content": "Hello!", - "id": "chatcmpl-Dd0Na8gXEDyFIhYMsL72TYk3bSZun-1", + "content": "Paris", + "id": "chatcmpl-DelITaJCJy951hwON0psdz2H9dF7i-1", "ingest_source": "Python", "is_response": True, "request_id": "req_22204b237d22427fbfd99c665d8a9964", @@ -84,7 +86,7 @@ ], ] -chat_completion_stream_recorded_events = [ +client_stream_recorded_events = [ [ {"type": "LlmChatCompletionSummary"}, { @@ -93,7 +95,7 @@ "ingest_source": "Python", "request.model": "gpt-3.5-turbo", "request.temperature": 0.7, - "request_id": "req_4566af5dd7224f00a2407fa1d3e32864", + "request_id": "req_22204b237d22427fbfd99c665d8a9964", "response.choices.finish_reason": "stop", "response.headers.llmVersion": "2020-10-01", "response.headers.ratelimitLimitRequests": 10000, @@ -116,10 +118,10 @@ {"type": "LlmChatCompletionMessage"}, { "completion_id": None, - "content": 'Use a tool to add an exclamation to the word "Hello"', + "content": "What is the capital of France? Answer in one word.", "id": "chatcmpl-DelITaJCJy951hwON0psdz2H9dF7i-0", "ingest_source": "Python", - "request_id": "req_4566af5dd7224f00a2407fa1d3e32864", + "request_id": "req_22204b237d22427fbfd99c665d8a9964", "response.model": "gpt-3.5-turbo-0125", "role": "user", "sequence": 0, @@ -133,11 +135,11 @@ {"type": "LlmChatCompletionMessage"}, { "completion_id": None, - "content": "Hello!", + "content": "Paris", "id": "chatcmpl-DelITaJCJy951hwON0psdz2H9dF7i-1", "ingest_source": "Python", "is_response": True, - "request_id": "req_4566af5dd7224f00a2407fa1d3e32864", + "request_id": "req_22204b237d22427fbfd99c665d8a9964", "response.model": "gpt-3.5-turbo-0125", "role": "assistant", "sequence": 1, @@ -149,6 +151,216 @@ ] +agent_recorded_events = [ + [ + {"timestamp": None, "type": "LlmChatCompletionSummary"}, + { + "duration": None, + "id": None, + "ingest_source": "Python", + "request.model": "gpt-3.5-turbo", + "request.temperature": 0.7, + "request_id": "req_619548c272db4f1ab380b83de9fdedef", + "response.choices.finish_reason": "tool_calls", + "response.headers.llmVersion": "2020-10-01", + "response.headers.ratelimitLimitRequests": 10000, + "response.headers.ratelimitLimitTokens": 50000000, + "response.headers.ratelimitRemainingRequests": 9999, + "response.headers.ratelimitRemainingTokens": 49999974, + "response.headers.ratelimitResetRequests": "6ms", + "response.headers.ratelimitResetTokens": "0s", + "response.model": "gpt-3.5-turbo-0125", + "response.number_of_messages": 2, + "response.organization": "user-rk8wq9voijy9sejrncvgi0iw", + "span_id": None, + "timestamp": None, + "trace_id": None, + "vendor": "openai", + }, + ], + [ + {"timestamp": None, "type": "LlmChatCompletionMessage"}, + { + "completion_id": None, + "content": "You are a text manipulation algorithm.", + "id": "chatcmpl-CukvsGfSQihNO9I3FTqaNKERWtUca-0", + "ingest_source": "Python", + "request_id": "req_619548c272db4f1ab380b83de9fdedef", + "response.model": "gpt-3.5-turbo-0125", + "role": "system", + "sequence": 0, + "span_id": None, + "timestamp": None, + "trace_id": None, + "vendor": "openai", + }, + ], + [ + {"timestamp": None, "type": "LlmChatCompletionMessage"}, + { + "completion_id": None, + "content": 'Use a tool to add an exclamation to the word "Hello"', + "id": "chatcmpl-CukvsGfSQihNO9I3FTqaNKERWtUca-1", + "ingest_source": "Python", + "request_id": "req_619548c272db4f1ab380b83de9fdedef", + "response.model": "gpt-3.5-turbo-0125", + "role": "user", + "sequence": 1, + "span_id": None, + "timestamp": None, + "trace_id": None, + "vendor": "openai", + }, + ], + [ + {"timestamp": None, "type": "LlmTool"}, + { + "agent_name": "my_agent", + "duration": None, + "id": None, + "ingest_source": "Python", + "input": "{'message': 'Hello'}", + "name": "add_exclamation", + "output": "Hello!", + "run_id": "call_ymnsNurMgr3atFVr7BnJ2XYK", + "span_id": None, + "trace_id": None, + "vendor": "langchain", + }, + ], + [ + {"timestamp": None, "type": "LlmChatCompletionSummary"}, + { + "duration": None, + "id": None, + "ingest_source": "Python", + "request.model": "gpt-3.5-turbo", + "request.temperature": 0.7, + "request_id": "req_619548c272db4f1ab380b83de9fdedef", + "response.choices.finish_reason": "stop", + "response.headers.llmVersion": "2020-10-01", + "response.headers.ratelimitLimitRequests": 10000, + "response.headers.ratelimitLimitTokens": 50000000, + "response.headers.ratelimitRemainingRequests": 9999, + "response.headers.ratelimitRemainingTokens": 49999970, + "response.headers.ratelimitResetRequests": "6ms", + "response.headers.ratelimitResetTokens": "0s", + "response.model": "gpt-3.5-turbo-0125", + "response.number_of_messages": 5, + "response.organization": "user-rk8wq9voijy9sejrncvgi0iw", + "span_id": None, + "timestamp": None, + "trace_id": None, + "vendor": "openai", + }, + ], + [ + {"timestamp": None, "type": "LlmChatCompletionMessage"}, + { + "completion_id": None, + "content": "You are a text manipulation algorithm.", + "id": "chatcmpl-CukvtgYHPS8HRHqCQiQgQrs7a2Tx1-0", + "ingest_source": "Python", + "request_id": "req_619548c272db4f1ab380b83de9fdedef", + "response.model": "gpt-3.5-turbo-0125", + "role": "system", + "sequence": 0, + "span_id": None, + "timestamp": None, + "trace_id": None, + "vendor": "openai", + }, + ], + [ + {"timestamp": None, "type": "LlmChatCompletionMessage"}, + { + "completion_id": None, + "content": 'Use a tool to add an exclamation to the word "Hello"', + "id": "chatcmpl-CukvtgYHPS8HRHqCQiQgQrs7a2Tx1-1", + "ingest_source": "Python", + "request_id": "req_619548c272db4f1ab380b83de9fdedef", + "response.model": "gpt-3.5-turbo-0125", + "role": "user", + "sequence": 1, + "span_id": None, + "timestamp": None, + "trace_id": None, + "vendor": "openai", + }, + ], + [ + {"timestamp": None, "type": "LlmChatCompletionMessage"}, + { + "completion_id": None, + "id": "chatcmpl-CukvtgYHPS8HRHqCQiQgQrs7a2Tx1-2", + "ingest_source": "Python", + "request_id": "req_619548c272db4f1ab380b83de9fdedef", + "response.model": "gpt-3.5-turbo-0125", + "role": "assistant", + "sequence": 2, + "span_id": None, + "timestamp": None, + "trace_id": None, + "vendor": "openai", + }, + ], + [ + {"timestamp": None, "type": "LlmChatCompletionMessage"}, + { + "completion_id": None, + "content": "Hello!", + "id": "chatcmpl-CukvtgYHPS8HRHqCQiQgQrs7a2Tx1-3", + "ingest_source": "Python", + "request_id": "req_619548c272db4f1ab380b83de9fdedef", + "response.model": "gpt-3.5-turbo-0125", + "role": "tool", + "sequence": 3, + "span_id": None, + "timestamp": None, + "trace_id": None, + "vendor": "openai", + }, + ], + [ + {"timestamp": None, "type": "LlmChatCompletionMessage"}, + { + "completion_id": None, + "content": 'The word "Hello" with an exclamation mark added is "Hello!"', + "id": "chatcmpl-CukvtgYHPS8HRHqCQiQgQrs7a2Tx1-4", + "ingest_source": "Python", + "is_response": True, + "request_id": "req_619548c272db4f1ab380b83de9fdedef", + "response.model": "gpt-3.5-turbo-0125", + "role": "assistant", + "sequence": 4, + "span_id": None, + "trace_id": None, + "vendor": "openai", + }, + ], + [ + {"timestamp": None, "type": "LlmAgent"}, + { + "duration": None, + "id": None, + "ingest_source": "Python", + "name": "my_agent", + "span_id": None, + "trace_id": None, + "vendor": "langchain", + }, + ], +] + + +@tool +def add_exclamation(message: str) -> str: + """Adds an exclamation mark to the input message.""" + if "exc" in message: + raise RuntimeError("Oops") + return f"{message}!" + + def _build_graph(node): from langgraph.graph import END, START, MessagesState, StateGraph @@ -159,51 +371,133 @@ def _build_graph(node): return builder.compile() +@pytest.fixture(scope="session") +def create_agent(chat_openai_client): + def _create_agent(model="gpt-5.1", tools=None, system_prompt=None, name="my_agent"): + from langchain.agents import create_agent + + client = chat_openai_client.with_config(model=model, timeout=30) + + return create_agent(model=client, tools=tools, system_prompt=system_prompt, name=name) + + return _create_agent + + @reset_core_stats_engine() -@validate_custom_events(chat_completion_recorded_events) +@validate_custom_events(client_recorded_events) @background_task() def test_state_graph_with_client_invoke(chat_openai_client, exercise_graph): def state_graph_invoke(state): response = chat_openai_client.invoke(state["messages"]) return {"messages": [response]} - response = exercise_graph(_build_graph(state_graph_invoke), PROMPT) + response = exercise_graph(_build_graph(state_graph_invoke), CLIENT_PROMPT) assert response @reset_core_stats_engine() -@validate_custom_events(chat_completion_recorded_events) +@validate_custom_events(client_recorded_events) @background_task() def test_state_graph_with_client_ainvoke(chat_openai_client, exercise_graph): async def state_graph_ainvoke(state): response = await chat_openai_client.ainvoke(state["messages"]) return {"messages": [response]} - response = exercise_graph(_build_graph(state_graph_ainvoke), PROMPT) + response = exercise_graph(_build_graph(state_graph_ainvoke), CLIENT_PROMPT) assert response @reset_core_stats_engine() -@validate_custom_events(chat_completion_stream_recorded_events) +@validate_custom_events(client_stream_recorded_events) @background_task() def test_state_graph_with_client_stream(chat_openai_client, exercise_graph): def state_graph_stream(state): chunks = list(chat_openai_client.stream(state["messages"])) return {"messages": ["".join(chunk.content for chunk in chunks)]} - response = exercise_graph(_build_graph(state_graph_stream), PROMPT) + response = exercise_graph(_build_graph(state_graph_stream), CLIENT_PROMPT) assert response @reset_core_stats_engine() -@validate_custom_events(chat_completion_stream_recorded_events) +@validate_custom_events(client_stream_recorded_events) @background_task() def test_state_graph_with_client_astream(chat_openai_client, exercise_graph): async def state_graph_astream(state): chunks = [chunk async for chunk in chat_openai_client.astream(state["messages"])] return {"messages": ["".join(chunk.content for chunk in chunks)]} - response = exercise_graph(_build_graph(state_graph_astream), PROMPT) + response = exercise_graph(_build_graph(state_graph_astream), CLIENT_PROMPT) + assert response + + +@reset_core_stats_engine() +@validate_custom_events(agent_recorded_events) +@background_task() +def test_state_graph_with_agent_invoke(exercise_graph, create_agent): + my_agent = create_agent(tools=[add_exclamation], system_prompt="You are a text manipulation algorithm.") + + def state_graph_invoke(state): + response = my_agent.invoke({"messages": state["messages"]}) + return {"messages": response.get("messages", [])} + + response = exercise_graph(_build_graph(state_graph_invoke), AGENT_PROMPT) + assert response + + +@reset_core_stats_engine() +@validate_custom_events(agent_recorded_events) +@background_task() +def test_state_graph_with_agent_ainvoke(exercise_graph, create_agent): + my_agent = create_agent(tools=[add_exclamation], system_prompt="You are a text manipulation algorithm.") + + async def state_graph_ainvoke(state): + response = await my_agent.ainvoke({"messages": state["messages"]}) + return {"messages": response.get("messages", [])} + + response = exercise_graph(_build_graph(state_graph_ainvoke), AGENT_PROMPT) + assert response + + +@reset_core_stats_engine() +@validate_custom_events(agent_recorded_events) +@background_task() +def test_state_graph_with_agent_stream(exercise_graph, create_agent): + my_agent = create_agent(tools=[add_exclamation], system_prompt="You are a text manipulation algorithm.") + + def state_graph_stream(state): + chunks = list(my_agent.stream({"messages": state["messages"]})) + messages = [] + for event in chunks: + if not isinstance(event, dict): + continue + for value in event.values(): + if isinstance(value, dict): + messages.extend(value.get("messages", [])) + return {"messages": messages} + + response = exercise_graph(_build_graph(state_graph_stream), AGENT_PROMPT) + assert response + + +@reset_core_stats_engine() +@validate_custom_events(agent_recorded_events) +@background_task() +def test_state_graph_with_agent_astream(exercise_graph, create_agent): + my_agent = create_agent(tools=[add_exclamation], system_prompt="You are a text manipulation algorithm.") + + async def state_graph_astream(state): + chunks = [chunk async for chunk in my_agent.astream({"messages": state["messages"]})] + messages = [] + for event in chunks: + if not isinstance(event, dict): + continue + for value in event.values(): + if isinstance(value, dict): + messages.extend(value.get("messages", [])) + return {"messages": messages} + + response = exercise_graph(_build_graph(state_graph_astream), AGENT_PROMPT) assert response