Repository: jerryjliu/gpt_index
File: llama-index-core/llama_index/core/agent/workflow/base_agent.py (Python)

Code before:

        if num_iterations >= max_iterations:
            raise WorkflowRuntimeError(
                f"Max iterations of {max_iterations} reached! Either something went wrong, or you can "
                "increase the max iterations with `.run(.., max_iterations=...)`"
            )

        memory: BaseMemory = await ctx.store.get("memory")

        if ev.retry_messages:
            # Retry with the given messages to let the LLM fix potential errors
            history = await memory.aget()
            user_msg_str = await ctx.store.get("user_msg_str")

            return AgentInput(
                input=[
                    *history,
                    ChatMessage(role="user", content=user_msg_str),
                    *ev.retry_messages,
                ],
                current_agent_name=self.name,
            )

        if not ev.tool_calls:
            messages = await memory.aget()
            output = await self.finalize(ctx, ev, memory)
            cur_tool_calls: List[ToolCallResult] = await ctx.store.get(
                "current_tool_calls", default=[]
            )
            output.tool_calls.extend(cur_tool_calls)  # type: ignore

            # with the following code we try to address errors while generating the
            # structured response. These errors often stem from the fact that the memory
            # does not contain a 'final' answer by the assistant, which would be key
            # to format the output in the structured way.
            # In this sense, we first try generating the output directly from the
            # in-memory messages. If that does not work and the in-memory messages
            # are missing the final answer from the agent, we then proceed to add
            # that using the output object and re-try generating the structured response.

            if self.structured_output_fn is not None:
                try:
                    if inspect.iscoroutinefunction(self.structured_output_fn):
                        output.structured_response = await self.structured_output_fn(
                            messages
                        )
                    else:
                        output.structured_response = cast(
                            Dict[str, Any], self.structured_output_fn(messages)
                        )
                except Exception as e:

Review comment (bug):
the bug is right here. At least for the function calling agent, it first writes to a scratchpad, and then moves the scratchpad values to the memory only when finalize is called.

So, the fix here should just be to fetch the `messages =` assignment right after this line, rather than right before

Code after:

        if num_iterations >= max_iterations:
            raise WorkflowRuntimeError(
                f"Max iterations of {max_iterations} reached! Either something went wrong, or you can "
                "increase the max iterations with `.run(.., max_iterations=...)`"
            )

        memory: BaseMemory = await ctx.store.get("memory")

        if ev.retry_messages:
            # Retry with the given messages to let the LLM fix potential errors
            history = await memory.aget()
            user_msg_str = await ctx.store.get("user_msg_str")

            return AgentInput(
                input=[
                    *history,
                    ChatMessage(role="user", content=user_msg_str),
                    *ev.retry_messages,
                ],
                current_agent_name=self.name,
            )

        if not ev.tool_calls:
            # important: messages should always be fetched after calling finalize, otherwise they do not contain the agent's response
            output = await self.finalize(ctx, ev, memory)
            messages = await memory.aget()
            cur_tool_calls: List[ToolCallResult] = await ctx.store.get(
                "current_tool_calls", default=[]
            )
            output.tool_calls.extend(cur_tool_calls)  # type: ignore

            if self.structured_output_fn is not None:
                try:
                    if inspect.iscoroutinefunction(self.structured_output_fn):
                        output.structured_response = await self.structured_output_fn(
                            messages
                        )
                    else:
                        output.structured_response = cast(
                            Dict[str, Any], self.structured_output_fn(messages)
                        )
                except Exception as e:
                    warnings.warn(
                        f"There was a problem with the generation of the structured output: {e}"
                    )
            if self.output_cls is not None:
                try:
                    output.structured_response = await generate_structured_response(
                        messages=messages, llm=self.llm, output_cls=self.output_cls
                    )