Fix: Improve final response generation when reaching max iterations in AnthropicAugmentedLLM

[dragon1086]

Description

This PR enhances the AnthropicAugmentedLLM class to properly stop tool usage and generate a final text response when reaching the maximum number of iterations.

Changes

• Added logic to explicitly ask the LLM to stop using tools and provide a final response when reaching the maximum number of iterations while in a tool_use state:

  if i == params.max_iterations - 1 and response.stop_reason == "tool_use":
      final_prompt_message = MessageParam(
          role="user",
          content="We've reached the maximum number of iterations. Please stop using tools now and provide your final comprehensive answer based on all tool results so far.",
      )
      messages.append(final_prompt_message)

• Added corresponding test case (test_final_response_after_max_iterations_with_tool_use)
• Tested with examples/mcp_basic_agent/main.py by modifying the script to use a specific max_iterations value:

  async def example_usage():
      async with app.run() as agent_app:
          logger = agent_app.logger
          context = agent_app.context
  
          logger.info("Current config:", data=context.config.model_dump())
  
          # Add the current directory to the filesystem server's args
          context.config.mcp.servers["filesystem"].args.extend([os.getcwd()])
  
          finder_agent = Agent(
              name="finder",
              instruction="""You are an agent with access to the filesystem, 
              as well as the ability to fetch URLs. Your job is to identify 
              the closest match to a user's request, make the appropriate tool calls, 
              and return the URI and CONTENTS of the closest match.""",
              server_names=["fetch", "filesystem"],
          )
  
          async with finder_agent:
              llm = await finder_agent.attach_llm(AnthropicAugmentedLLM)
  
              result = await llm.generate_str(
                  message="""Print the first 2 paragraphs per link below:
                   https://modelcontextprotocol.io/introduction
                   https://modelcontextprotocol.io/examples
                   https://modelcontextprotocol.io/clients
                   https://modelcontextprotocol.io/tutorials/building-mcp-with-llms
                   https://modelcontextprotocol.io/docs/tools/debugging
                   https://modelcontextprotocol.io/docs/tools/inspector
              
                  """,
                  request_params=RequestParams(
                      model="claude-3-7-sonnet-latest", maxTokens=1000, max_iterations=3
                  ),
              )
              logger.info(f"Final Response: {result}")

Problem Solved

1. Previously, when the LLM reached the maximum number of iterations while using tools, it would continue attempting to use tools without providing a final text response. After this change, the LLM is explicitly instructed to stop using tools and generate a final comprehensive response based on all results gathered so far.
2. This ensures that conversations always terminate with a text response rather than a tool call, providing a better user experience.

Testing

• Added unit test: test_augmented_llm_anthropic.py with test_final_response_after_max_iterations_with_tool_use method
• Integration test: Verified with examples/mcp_basic_agent/main.py using a controlled max_iterations value
• All tests pass successfully

Note

This PR may have potential conflicts with ongoing PR #95, as both PRs modify the same area of the AnthropicAugmentedLLM class. If #95 is merged first, this PR may need to be rebased or have conflicts resolved.

[saqadri]

@dragon1086 thank you for the contribution! this is a good change.

One last request -- if you wouldn't mind, can you please add it to the openai augmented llm as well?

[saqadri]

Thank you for adding this @dragon1086 <3

[dragon1086]

You're welcome! Happy to contribute. 😊

[saqadri]

This is clever. The one issue I can think of is that there is no indication at the user level that the response is possibly incomplete because the LLM might have needed more information.

One option is in the prompt to ask the LLM to indicate that it isn't done, but it's returning a response based on what the information it's got.

This might be an interesting usecase for MCP notifications as well (in the future)

[dragon1086]

Thank you for your positive feedback.
I've modified the content text based on your suggestion.

When I ran integration tests, the beginning and ending portions of the final response looked like this.
Please let me know if you think it's too much:

1. Beginning portion:
   I've reached the maximum number of tool usage iterations, which means I was only
   able to retrieve content from the first two links on your list. To provide a
   complete answer, I would have needed to make additional fetch requests for the
   remaining 12 URLs you provided. Below are the first two paragraphs from the two
   links I was able to access:

https://modelcontextprotocol.io/introduction

...

2. Ending portion:
   ...
   To provide a complete response, I would have needed to fetch content from the
   remaining URLs:

• https://modelcontextprotocol.io/clients
• https://modelcontextprotocol.io/tutorials/building-mcp-with-llms
• https://modelcontextprotocol.io/docs/tools/debugging
• https://modelcontextprotocol.io/docs/tools/inspector
• https://modelcontextprotocol.io/docs/concepts/architecture
• https://modelcontextprotocol.io/docs/concepts/resources
• https://modelcontextprotocol.io/docs/concepts/prompts
• https://modelcontextprotocol.io/docs/concepts/tools
• https://modelcontextprotocol.io/docs/concepts/sampling
• https://modelcontextprotocol.io/docs/concepts/roots
• https://modelcontextprotocol.io/docs/concepts/transports
• https://modelcontextprotocol.io/development/updates

[dragon1086]

Thank you for the feedback on my contribution!

I considered fixing the OpenAI augmented LLM before submitting this PR, but found it difficult to reproduce a scenario with multiple iterations.

The reason OpenAI works well is that it typically calls all necessary tools within the first iteration through the 'message.tool_calls'.
Of course, there could be cases where the response gets truncated due to length limitations from too many tool calls.
I think the code implementation would look something like:

if (
                choice.finish_reason in ["tool_calls", "function_call"]
                and message.tool_calls
):
...
...

if params.max_iterations > 1 and i == params.max_iterations - 2:
    final_prompt_message = ChatCompletionUserMessageParam(
        role="user",
        content="We've reached the maximum number of iterations. Please ...",
    )
    messages.append(final_prompt_message)

However, I'd prefer to work on this once I can reproduce a case where it's actually needed.
If such a case is found, I'll be happy to submit another PR to address it.

[saqadri]

@dragon1086 thanks for this. Let's go with this approach and we'll iterate as needed.