StateGraph generate step produces empty "Ai Message"

[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

I have the following StateGraph:

@tool
def ground_search(
    query: str, *, config: Annotated[RunnableConfig, InjectedToolArg]
)-> Optional[list[str]]:
    """Search for general web results.

    This function performs a search using the Google search engine, which is designed
    to provide comprehensive, accurate, and trusted results. It's particularly useful
    for answering questions about current events.
    """
    client = genai.Client(api_key=os.environ.get("VERTEX_API_KEY"))
    model_id = "gemini-2.0-flash"
    google_search_tool = Tool(
        google_search = GoogleSearch()
    )
    response = client.models.generate_content(
        model=model_id,
        contents=[{"role": "user", "parts": [{"text": query}]}], 
        config=GenerateContentConfig(
            tools=[google_search_tool],
            system_instruction="You are a helpful AI assistant named Bob.",
            response_modalities=["TEXT"],
        )
    )
    return "\n\n".join(content.text for content in response.candidates[0].content.parts)

@tool(response_format="content_and_artifact")
async def retrieve(query: str, *, config: RunnableConfig):
    """Retrieve information related to a query."""
    retrieved_docs = await vector_store.asimilarity_search(query, k=2)
    serialized = "\n\n".join(
        (f"Source: {doc.metadata}\n" f"Content: {doc.page_content}")
        for doc in retrieved_docs
    )
    return serialized, retrieved_docs

---

    # Step 3: Generate a response using the retrieved content.
    async def Generate(self, state: MessagesState, config: RunnableConfig):
        """Generate answer."""
        print(f"\n=== {self.Generate.__name__} ===")
        # Get generated ToolMessages
        recent_tool_messages = []
        for message in reversed(state["messages"]):
            if message.type == "tool":
                recent_tool_messages.append(message)
            else:
                break
        tool_messages = recent_tool_messages[::-1]
        # Format into prompt
        docs_content = "\n\n".join(doc.content for doc in tool_messages)
        system_message_content = f"""
        You are an assistant for question-answering tasks. 
        Use the following pieces of retrieved context to answer 
        the question. If you don't know the answer, say that you 
        don't know. Use three sentences maximum and keep the 
        answer concise.
        \n\n
        {docs_content}
        """
        conversation_messages = [
            message
            for message in state["messages"]
            if message.type in ("human", "system")
            or (message.type == "ai" and not message.tool_calls)
        ]
        prompt = [SystemMessage(system_message_content)] + conversation_messages
        # Run
        response = await self._llm.ainvoke(prompt, config)
        return {"messages": [response]}

---

        graph_builder = StateGraph(MessagesState)
        graph_builder.add_node("query_or_respond", self.query_or_respond)
        graph_builder.add_node("tools", ToolNode(TOOLS)) # Execute the retrieval.
        graph_builder.add_node("generate", self.generate)
        graph_builder.set_entry_point("query_or_respond")
        graph_builder.add_conditional_edges(
            "query_or_respond",
            tools_condition,
            {END: END, "tools": "tools"},
        )
        graph_builder.add_edge("tools", "generate")
        graph_builder.add_edge("generate", END)
        return graph_builder.compile(store=InMemoryStore(), checkpointer=MemorySaver(), name="Checkedpoint StateGraph RAG")

async def Chat(graph, config, messages: List[str]):
    print(f"\n=== {Chat.__name__} ===")
    messages = ["What is Task Decomposition?", "Can you look up some common ways of doing it?"]
    for message in messages:
        async for step in graph.astream(
            {"messages": [{"role": "user", "content": message}]},
            stream_mode="values",
            config = config
        ):
            step["messages"][-1].pretty_print()

Description

The generate step spits out "Ai Message" only one time, for the first question but not the follow-up question:

$ p -m src.rag_agent.CheckpointedRAG
================================ Human Message =================================

What is Task Decomposition?
================================== Ai Message ==================================
Tool Calls:
  retrieve (e1e31448-a4e7-4602-8d1c-ecf41cffa663)
 Call ID: e1e31448-a4e7-4602-8d1c-ecf41cffa663
  Args:
    query: Task Decomposition

=== asimilarity_search ===
================================= Tool Message =================================
Name: retrieve

Source: {'source': 'https://lilianweng.github.io/posts/2023-06-23-agent/'}
Content: Fig. 1. Overview of a LLM-powered autonomous agent system.
Component One: Planning#
A complicated task usually involves many steps. An agent needs to know what they are and plan ahead.
Task Decomposition#
Chain of thought (CoT; Wei et al. 2022) has become a standard prompting technique for enhancing model performance on complex tasks. The model is instructed to “think step by step” to utilize more test-time computation to decompose hard tasks into smaller and simpler steps. CoT transforms big tasks into multiple manageable tasks and shed lights into an interpretation of the model’s thinking process.

Source: {'source': 'https://lilianweng.github.io/posts/2023-06-23-agent/'}
Content: Tree of Thoughts (Yao et al. 2023) extends CoT by exploring multiple reasoning possibilities at each step. It first decomposes the problem into multiple thought steps and generates multiple thoughts per step, creating a tree structure. The search process can be BFS (breadth-first search) or DFS (depth-first search) with each state evaluated by a classifier (via a prompt) or majority vote.
Task decomposition can be done (1) by LLM with simple prompting like "Steps for XYZ.\n1.", "What are the subgoals for achieving XYZ?", (2) by using task-specific instructions; e.g. "Write a story outline." for writing a novel, or (3) with human inputs.

=== Generate ===
================================== Ai Message ==================================

Task decomposition is a method of breaking down a complex task into smaller, more manageable steps. This can be done by prompting an LLM, using task-specific instructions, or with human input. Chain of Thought (CoT) is a standard prompting technique for task decomposition, instructing the model to "think step by step". Tree of Thoughts extends CoT by exploring multiple reasoning possibilities at each step, creating a tree structure.
================================ Human Message =================================

Can you look up some common ways of doing it?
================================== Ai Message ==================================
Tool Calls:
  retrieve (4e69fd1b-ab9e-4a47-a00b-51ce34584004)
 Call ID: 4e69fd1b-ab9e-4a47-a00b-51ce34584004
  Args:
    query: common methods for task decomposition in AI

=== asimilarity_search ===
================================= Tool Message =================================
Name: retrieve

Source: {'source': 'https://lilianweng.github.io/posts/2023-06-23-agent/'}
Content: Fig. 1. Overview of a LLM-powered autonomous agent system.
Component One: Planning#
A complicated task usually involves many steps. An agent needs to know what they are and plan ahead.
Task Decomposition#
Chain of thought (CoT; Wei et al. 2022) has become a standard prompting technique for enhancing model performance on complex tasks. The model is instructed to “think step by step” to utilize more test-time computation to decompose hard tasks into smaller and simpler steps. CoT transforms big tasks into multiple manageable tasks and shed lights into an interpretation of the model’s thinking process.

Source: {'source': 'https://lilianweng.github.io/posts/2023-06-23-agent/'}
Content: Tree of Thoughts (Yao et al. 2023) extends CoT by exploring multiple reasoning possibilities at each step. It first decomposes the problem into multiple thought steps and generates multiple thoughts per step, creating a tree structure. The search process can be BFS (breadth-first search) or DFS (depth-first search) with each state evaluated by a classifier (via a prompt) or majority vote.
Task decomposition can be done (1) by LLM with simple prompting like "Steps for XYZ.\n1.", "What are the subgoals for achieving XYZ?", (2) by using task-specific instructions; e.g. "Write a story outline." for writing a novel, or (3) with human inputs.

=== Generate ===
================================== Ai Message ==================================
Tool Calls:
  retrieve (78d2ae63-8fab-4489-9422-5e5d95788f5b)
 Call ID: 78d2ae63-8fab-4489-9422-5e5d95788f5b
  Args:
    query: common ways to do task decomposition

^^^ Where is the response? ^^^

System Info

langchain==0.3.20
langchain-community==0.3.19
langchain-core==0.3.45
langchain-google-genai==2.1.0
langchain-google-vertexai==2.0.9
langchain-openai==0.3.8
langchain-text-splitters==0.3.6

[khteh]

I check the log output of the Generate step:

conversation_messages: [HumanMessage(content='Hello, who are you?', additional_kwargs={}, response_metadata={}, id='838cc7bf-9361-4777-8359-e7a7a7b141c1'), AIMessage(content='I am Gemini, a large language model built by Google.\n', additional_kwargs={}, response_metadata={'is_blocked': False, 'safety_ratings': [], 'usage_metadata': {'prompt_token_count': 16, 'candidates_token_count': 13, 'total_token_count': 29, 'prompt_tokens_details': [{'modality': 1, 'token_count': 16}], 'candidates_tokens_details': [{'modality': 1, 'token_count': 13}], 'cached_content_token_count': 0, 'cache_tokens_details': []}, 'finish_reason': 'STOP', 'avg_logprobs': -0.002166313763994437}, id='run-0393aaea-5368-4183-86f3-2dd2d7ae3cbc-0', usage_metadata={'input_tokens': 16, 'output_tokens': 13, 'total_tokens': 29}), HumanMessage(content='What is Task Decomposition?', additional_kwargs={}, response_metadata={}, id='928ca68f-0353-49e1-9523-4fe6efb17963'), AIMessage(content='Task decomposition is breaking down a complex task into smaller, simpler steps. This can be done by prompting an LLM, using task-specific instructions, or with human input. Chain of Thought and Tree of Thoughts are examples of task decomposition techniques.\n', additional_kwargs={}, response_metadata={'is_blocked': False, 'safety_ratings': [], 'usage_metadata': {'prompt_token_count': 458, 'candidates_token_count': 50, 'total_token_count': 508, 'prompt_tokens_details': [{'modality': 1, 'token_count': 458}], 'candidates_tokens_details': [{'modality': 1, 'token_count': 50}], 'cached_content_token_count': 0, 'cache_tokens_details': []}, 'finish_reason': 'STOP', 'avg_logprobs': -0.16305770874023437}, id='run-01dd3c6c-7296-47bf-8e9a-40656c7650b8-0', usage_metadata={'input_tokens': 458, 'output_tokens': 50, 'total_tokens': 508}), HumanMessage(content='Can you look up some common ways of doing it?', additional_kwargs={}, response_metadata={}, id='d9fbbb13-febe-48ff-8543-988529c725b4')]

prompt: [SystemMessage(content='\n        You are an assistant for question-answering tasks. \n        Use the following pieces of retrieved context to answer \n        the question. If you don\'t know the answer, say that you \n        don\'t know. Use three sentences maximum and keep the \n        answer concise.\n        \n\n\n        Source: {\'source\': \'https://lilianweng.github.io/posts/2023-06-23-agent/\'}\nContent: Fig. 1. Overview of a LLM-powered autonomous agent system.\nComponent One: Planning#\nA complicated task usually involves many steps. An agent needs to know what they are and plan ahead.\nTask Decomposition#\nChain of thought (CoT; Wei et al. 2022) has become a standard prompting technique for enhancing model performance on complex tasks. The model is instructed to “think step by step” to utilize more test-time computation to decompose hard tasks into smaller and simpler steps. CoT transforms big tasks into multiple manageable tasks and shed lights into an interpretation of the model’s thinking process.\n\nSource: {\'source\': \'https://lilianweng.github.io/posts/2023-06-23-agent/\'}\nContent: Tree of Thoughts (Yao et al. 2023) extends CoT by exploring multiple reasoning possibilities at each step. It first decomposes the problem into multiple thought steps and generates multiple thoughts per step, creating a tree structure. The search process can be BFS (breadth-first search) or DFS (depth-first search) with each state evaluated by a classifier (via a prompt) or majority vote.\nTask decomposition can be done (1) by LLM with simple prompting like "Steps for XYZ.\\n1.", "What are the subgoals for achieving XYZ?", (2) by using task-specific instructions; e.g. "Write a story outline." for writing a novel, or (3) with human inputs.\n        ', additional_kwargs={}, response_metadata={}), HumanMessage(content='Hello, who are you?', additional_kwargs={}, response_metadata={}, id='838cc7bf-9361-4777-8359-e7a7a7b141c1'), AIMessage(content='I am Gemini, a large language model built by Google.\n', additional_kwargs={}, response_metadata={'is_blocked': False, 'safety_ratings': [], 'usage_metadata': {'prompt_token_count': 16, 'candidates_token_count': 13, 'total_token_count': 29, 'prompt_tokens_details': [{'modality': 1, 'token_count': 16}], 'candidates_tokens_details': [{'modality': 1, 'token_count': 13}], 'cached_content_token_count': 0, 'cache_tokens_details': []}, 'finish_reason': 'STOP', 'avg_logprobs': -0.002166313763994437}, id='run-0393aaea-5368-4183-86f3-2dd2d7ae3cbc-0', usage_metadata={'input_tokens': 16, 'output_tokens': 13, 'total_tokens': 29}), HumanMessage(content='What is Task Decomposition?', additional_kwargs={}, response_metadata={}, id='928ca68f-0353-49e1-9523-4fe6efb17963'), AIMessage(content='Task decomposition is breaking down a complex task into smaller, simpler steps. This can be done by prompting an LLM, using task-specific instructions, or with human input. Chain of Thought and Tree of Thoughts are examples of task decomposition techniques.\n', additional_kwargs={}, response_metadata={'is_blocked': False, 'safety_ratings': [], 'usage_metadata': {'prompt_token_count': 458, 'candidates_token_count': 50, 'total_token_count': 508, 'prompt_tokens_details': [{'modality': 1, 'token_count': 458}], 'candidates_tokens_details': [{'modality': 1, 'token_count': 50}], 'cached_content_token_count': 0, 'cache_tokens_details': []}, 'finish_reason': 'STOP', 'avg_logprobs': -0.16305770874023437}, id='run-01dd3c6c-7296-47bf-8e9a-40656c7650b8-0', usage_metadata={'input_tokens': 458, 'output_tokens': 50, 'total_tokens': 508}), HumanMessage(content='Can you look up some common ways of doing it?', additional_kwargs={}, response_metadata={}, id='d9fbbb13-febe-48ff-8543-988529c725b4')]

Generate() response: content='' additional_kwargs={'function_call': {'name': 'retrieve', 'arguments': '{"query": "common ways to do task decomposition"}'}} response_metadata={'is_blocked': False, 'safety_ratings': [], 'usage_metadata': {'prompt_token_count': 519, 'candidates_token_count': 8, 'total_token_count': 527, 'prompt_tokens_details': [{'modality': 1, 'token_count': 519}], 'candidates_tokens_details': [{'modality': 1, 'token_count': 8}], 'cached_content_token_count': 0, 'cache_tokens_details': []}, 'finish_reason': 'STOP', 'avg_logprobs': -0.0082077132537961} id='run-9d950ea3-bcdc-4234-be8f-1e14de4f4838-0' tool_calls=[{'name': 'retrieve', 'args': {'query': 'common ways to do task decomposition'}, 'id': 'cc6e4484-a863-48c4-8381-47f1e3d73f73', 'type': 'tool_call'}] usage_metadata={'input_tokens': 519, 'output_tokens': 8, 'total_tokens': 527}

The response content=''. Is it because it runs out of tokens?