Repository: run-llama/llama_index
Pull request: Fix TaskStepOutput sources bug
# Description

Fixing an issue where the stepwise sources returned in function calling agents' `TaskStepOutput` may include `ToolOutput` from other steps. I updated the handling of these sources to ensure that references to `task.extra_state["sources"]` are no longer inserted into `TaskStepOutput`, which should ensure that each step's sources actually represent the tool calls executed by said step. I also added a few tests to ensure that this works for sync and async.

Fixes #14884 

## New Package?

Did I fill in the `tool.llamahub` section in the `pyproject.toml` and provide a detailed README.md for my new integration or package?

- [ ] Yes
- [x] No

## Version Bump?

Did I bump the version in the `pyproject.toml` file of the package I am updating? (Except for the `llama-index-core` package)

- [ ] Yes
- [x] No

## Type of Change

Please delete options that are not relevant.

- [x] Bug fix (non-breaking change which fixes an issue)

## How Has This Been Tested?

Please describe the tests that you ran to verify your changes. Provide instructions so we can reproduce. Please also list any relevant details for your test configuration

- [x] Added new unit/integration tests
- [ ] Added new notebook (that tests end-to-end)
- [x] I stared at the code and made sure it makes sense

## Suggested Checklist:

- [x] I have performed a self-review of my own code
- [x] I have commented my code, particularly in hard-to-understand areas
- [ ] I have made corresponding changes to the documentation
- [ ] I have added Google Colab support for the newly added notebooks.
- [x] My changes generate no new warnings
- [x] I have added tests that prove my fix is effective or that my feature works
- [x] New and existing unit tests pass locally with my changes
- [x] I ran `make format; make lint` to appease the lint gods


Problem statement:
[Bug]: TaskStepOutput sources may include ToolOutput from other steps
### Bug Description

After an agent executes and completes a task (i.e., via `chat`), the `TaskState` associated with this task populates its `completed_steps: List[TaskStepOutput]` field with one `TaskStepOutput` for each step executed. Nested within each of these completed steps is a field `TaskStepOutput.output.sources: List[ToolOutput]`, a list of outputs resulting from tool calls made.

Presumably, a given step's sources should only contain outputs for the tool calls actually executed _in that step_. However, due to an apparent bug, this is not the case. Currently, all sources observed across all task steps are retained in a task-level state variable `task.extra_state["sources"]`; at the end of each step, they are "inserted" into a new completed step. This insertion does not make a copy of the sources list, instead it passes a reference, thus the sources of the now-completed steps continue to be updated as subsequent steps are executed. Upon completion, all completed step sources will be pointing to `task.extra_state["sources"]`, therefore all will be identical.

### Version

0.10.56

### Steps to Reproduce

1. Create a function calling agent
2. Run the agent on a suitable task
3. Inspect `agent.state` and nested fields to observe the noted issue

### Relevant Logs/Tracbacks

_No response_

Hints from issue discussion:

Commit message: Update finalize_response to attach all task sources

Patch to review:
diff --git a/llama-index-core/llama_index/core/agent/function_calling/step.py b/llama-index-core/llama_index/core/agent/function_calling/step.py
index 0680b30f09..5c2e3692ad 100644
--- a/llama-index-core/llama_index/core/agent/function_calling/step.py
+++ b/llama-index-core/llama_index/core/agent/function_calling/step.py
@@ -309,6 +309,7 @@ def run_step(self, step: TaskStep, task: Task, **kwargs: Any) -> TaskStepOutput:
         tool_calls = self._llm.get_tool_calls_from_response(
             response, error_on_no_tool_call=False
         )
+        tool_outputs: List[ToolOutput] = []
 
         if self._verbose and response.message.content:
             print("=== LLM Response ===")
@@ -336,10 +337,10 @@ def run_step(self, step: TaskStep, task: Task, **kwargs: Any) -> TaskStepOutput:
                     tools,
                     tool_call,
                     task.extra_state["new_memory"],
-                    task.extra_state["sources"],
+                    tool_outputs,
                     verbose=self._verbose,
                 )
-
+                task.extra_state["sources"].append(tool_outputs[-1])
                 task.extra_state["n_function_calls"] += 1
 
                 # check if any of the tools return directly -- only works if there is one tool call
@@ -368,9 +369,7 @@ def run_step(self, step: TaskStep, task: Task, **kwargs: Any) -> TaskStepOutput:
         except AttributeError:
             response_str = str(response)
 
-        agent_response = AgentChatResponse(
-            response=response_str, sources=task.extra_state["sources"]
-        )
+        agent_response = AgentChatResponse(response=response_str, sources=tool_outputs)
 
         return TaskStepOutput(
             output=agent_response,
@@ -402,6 +401,7 @@ async def arun_step(
         tool_calls = self._llm.get_tool_calls_from_response(
             response, error_on_no_tool_call=False
         )
+        tool_outputs: List[ToolOutput] = []
 
         if self._verbose and response.message.content:
             print("=== LLM Response ===")
@@ -428,17 +428,18 @@ async def arun_step(
                     tools,
                     tool_call,
                     task.extra_state["new_memory"],
-                    task.extra_state["sources"],
+                    tool_outputs,
                     verbose=self._verbose,
                 )
                 for tool_call in tool_calls
             ]
             return_directs = await asyncio.gather(*tasks)
+            task.extra_state["sources"].extend(tool_outputs)
 
             # check if any of the tools return directly -- only works if there is one tool call
             if len(return_directs) == 1 and return_directs[0]:
                 is_done = True
-                response = task.extra_state["sources"][-1].content
+                response = tool_outputs[-1].content
 
             task.extra_state["n_function_calls"] += len(tool_calls)
             # put tool output in sources and memory
@@ -461,9 +462,7 @@ async def arun_step(
         except AttributeError:
             response_str = str(response)
 
-        agent_response = AgentChatResponse(
-            response=response_str, sources=task.extra_state["sources"]
-        )
+        agent_response = AgentChatResponse(response=response_str, sources=tool_outputs)
 
         return TaskStepOutput(
             output=agent_response,
diff --git a/llama-index-core/llama_index/core/agent/runner/base.py b/llama-index-core/llama_index/core/agent/runner/base.py
index bd9301bbde..57b7e056a4 100644
--- a/llama-index-core/llama_index/core/agent/runner/base.py
+++ b/llama-index-core/llama_index/core/agent/runner/base.py
@@ -551,6 +551,11 @@ def finalize_response(
         if self.delete_task_on_finish:
             self.delete_task(task_id)
 
+        # Attach all sources generated across all steps
+        step_output.output.sources = self.get_task(task_id).extra_state.get(
+            "sources", []
+        )
+
         return cast(AGENT_CHAT_RESPONSE_TYPE, step_output.output)
 
     @dispatcher.span
diff --git a/llama-index-core/tests/agent/function_calling/test_step.py b/llama-index-core/tests/agent/function_calling/test_step.py
index bbf242a0cf..196e5b674a 100644
--- a/llama-index-core/tests/agent/function_calling/test_step.py
+++ b/llama-index-core/tests/agent/function_calling/test_step.py
@@ -1,3 +1,4 @@
+import uuid
 import pytest
 from typing import Any, AsyncGenerator, Coroutine, List, Optional, Sequence, Union
 from llama_index.core.agent.function_calling.step import (
@@ -32,8 +33,21 @@
 )
 NONEXISTENT_TOOL_OUTPUT = build_missing_tool_output(NONEXISTENT_TOOL_SELECTION)
 
+TOOL_1_NAME = "Tool 1"
+TOOL_1 = FunctionTool(lambda: None, ToolMetadata("", TOOL_1_NAME))
+TOOL_1_SELECTION = ToolSelection(
+    tool_id=TOOL_1_NAME,
+    tool_name=TOOL_1_NAME,
+    tool_kwargs={},
+)
+
+TOOL_2_NAME = "Tool 2"
+TOOL_2 = FunctionTool(lambda x: None, ToolMetadata("", TOOL_2_NAME))
+
+
+class MockFunctionCallingLLM(FunctionCallingLLM):
+    use_nonexistent_tool: bool = False
 
-class MockBadFunctionCallingLLM(FunctionCallingLLM):
     def achat(
         self, messages: Sequence[ChatMessage], **kwargs: Any
     ) -> Coroutine[Any, Any, ChatResponse]:
@@ -104,27 +118,35 @@ def get_tool_calls_from_response(
         error_on_no_tool_call: bool = True,
         **kwargs: Any
     ) -> List[ToolSelection]:
-        return [NONEXISTENT_TOOL_SELECTION]
+        return [
+            (
+                NONEXISTENT_TOOL_SELECTION
+                if self.use_nonexistent_tool
+                else TOOL_1_SELECTION
+            )
+        ]
+
+
+@pytest.fixture()
+def agent_worker() -> FunctionCallingAgentWorker:
+    llm = MockFunctionCallingLLM(use_nonexistent_tool=False)
+    return FunctionCallingAgentWorker([TOOL_1], llm, [])
 
 
 @pytest.fixture()
 def missing_function_agent_worker() -> FunctionCallingAgentWorker:
-    llm = MockBadFunctionCallingLLM()
-    tool = FunctionTool(lambda x: x, ToolMetadata("", "Tool 1"))
-    return FunctionCallingAgentWorker([tool], llm, [])
+    llm = MockFunctionCallingLLM(use_nonexistent_tool=True)
+    return FunctionCallingAgentWorker([TOOL_1], llm, [])
 
 
 def test_get_function_by_name_finds_existing_tool() -> None:
-    first_tool = FunctionTool(lambda x: x, ToolMetadata("", "Tool 1"))
-    second_tool = FunctionTool(lambda x: x, ToolMetadata("", "Tool 2"))
-    tools = [first_tool, second_tool]
-    assert get_function_by_name(tools, first_tool.metadata.name) == first_tool
-    assert get_function_by_name(tools, first_tool.metadata.name) == first_tool
+    tools = [TOOL_1, TOOL_2]
+    assert get_function_by_name(tools, TOOL_1.metadata.name) == TOOL_1
+    assert get_function_by_name(tools, TOOL_1.metadata.name) == TOOL_1
 
 
 def test_get_function_by_name_returns_none_for_nonexistent_tool() -> None:
-    tool = FunctionTool(lambda x: x, ToolMetadata("", "Tool 1"))
-    assert get_function_by_name([tool], "Name of a Nonexistent Tool") is None
+    assert get_function_by_name([TOOL_1], "Name of a Nonexistent Tool") is None
     assert get_function_by_name([], "Name of a Nonexistent Tool") is None
 
 
@@ -155,3 +177,36 @@ async def test_arun_step_returns_message_if_function_not_found(
     assert len(output.next_steps) == 1
     assert len(output_chat_response.sources) == 1
     assert output_chat_response.sources[0] == NONEXISTENT_TOOL_OUTPUT
+
+
+@pytest.mark.asyncio()
+@pytest.mark.parametrize("method", ["run_step", "arun_step"])
+async def test_run_step_returns_correct_sources_history(
+    method: str,
+    agent_worker: FunctionCallingAgentWorker,
+) -> None:
+    num_steps = 4
+    task = Task(input="", memory=ChatMemoryBuffer.from_defaults(), extra_state={})
+    step_outputs: List[TaskStepOutput] = []
+
+    # Create steps
+    steps = [agent_worker.initialize_step(task)]
+    for step_idx in range(num_steps - 1):
+        steps.append(
+            steps[-1].get_next_step(
+                step_id=str(uuid.uuid4()),
+                input=None,
+            )
+        )
+
+    # Run each step, invoking a single tool call each time
+    for step_idx in range(num_steps):
+        step_outputs.append(
+            agent_worker.run_step(steps[step_idx], task)
+            if method == "run_step"
+            else await agent_worker.arun_step(steps[step_idx], task)
+        )
+
+    # Ensure that each step only has one source for its one tool call
+    for step_idx in range(num_steps):
+        assert len(step_outputs[step_idx].output.sources) == 1
diff --git a/llama-index-core/tests/agent/runner/test_base.py b/llama-index-core/tests/agent/runner/test_base.py
index 9fb2b09843..2650a2874a 100644
--- a/llama-index-core/tests/agent/runner/test_base.py
+++ b/llama-index-core/tests/agent/runner/test_base.py
@@ -21,6 +21,7 @@
     AgentRunStepStartEvent,
 )
 from llama_index.core.instrumentation.events.base import BaseEvent
+from llama_index.core.tools.types import ToolOutput
 
 dispatcher = instrument.get_dispatcher()
 
@@ -38,7 +39,7 @@ def handle(self, e: BaseEvent):
 
 # define mock agent worker
 class MockAgentWorker(BaseAgentWorker):
-    """Mock agent agent worker."""
+    """Mock agent worker."""
 
     def __init__(self, limit: int = 2):
         """Initialize."""
@@ -202,6 +203,28 @@ def finalize_task(self, task: Task, **kwargs: Any) -> None:
         """Finalize task, after all the steps are completed."""
 
 
+class MockFunctionCallingAgentWorker(MockAgentWorker):
+    """Mock function calling agent worker."""
+
+    counter = 0
+
+    def initialize_step(self, task: Task, **kwargs: Any) -> TaskStep:
+        task.extra_state["sources"] = []
+        return super().initialize_step(task, **kwargs)
+
+    def run_step(self, step: TaskStep, task: Task, **kwargs: Any) -> TaskStepOutput:
+        self.counter += 1
+        task.extra_state["sources"].append(
+            ToolOutput(
+                tool_name=f"My Tool",
+                content=f"This is the output of tool call {self.counter}",
+                raw_input={},
+                raw_output=f"This is the raw output of tool call {self.counter}",
+            )
+        )
+        return super().run_step(step, task, **kwargs)
+
+
 def test_agent() -> None:
     """Test executor."""
     agent_runner = AgentRunner(agent_worker=MockAgentWorker(limit=2))
@@ -343,3 +366,28 @@ def test_agent_dispatches_events() -> None:
     # Check AgentChatWithStepEndEvent
     assert isinstance(event_handler.events[-1], AgentChatWithStepEndEvent)
     assert event_handler.events[-1].response == response
+
+
+def test_agent_response_contains_full_source_history() -> None:
+    # Test with a mock agent worker that records one source in each step
+    num_steps = 4
+    agent_runner = AgentRunner(
+        agent_worker=MockFunctionCallingAgentWorker(limit=num_steps)
+    )
+
+    response = agent_runner.chat("hello world")
+    assert len(response.sources) == num_steps
+
+    # Expect that agent will collect all sources generated across all steps
+    for step_idx in range(num_steps):
+        assert (
+            response.sources[step_idx].content
+            == f"This is the output of tool call {step_idx + 1}"
+        )
+
+
+def test_agent_handles_nonexistent_source_history() -> None:
+    num_steps = 4
+    agent_runner = AgentRunner(agent_worker=MockAgentWorker(limit=num_steps))
+    response = agent_runner.chat("hello world")
+    assert len(response.sources) == 0
diff --git a/llama-index-integrations/agent/llama-index-agent-openai/llama_index/agent/openai/step.py b/llama-index-integrations/agent/llama-index-agent-openai/llama_index/agent/openai/step.py
index 7e4e913451..69b0a6d201 100644
--- a/llama-index-integrations/agent/llama-index-agent-openai/llama_index/agent/openai/step.py
+++ b/llama-index-integrations/agent/llama-index-agent-openai/llama_index/agent/openai/step.py
@@ -335,9 +335,11 @@ def _call_function(
         dispatcher.event(
             AgentToolCallEvent(
                 arguments=function_args_str,
-                tool=tool.metadata
-                if tool
-                else ToolMetadata(description="unknown", name=function_name),
+                tool=(
+                    tool.metadata
+                    if tool
+                    else ToolMetadata(description="unknown", name=function_name)
+                ),
             )
         )
         with self.callback_manager.event(
@@ -439,9 +441,11 @@ async def _acall_function(
         dispatcher.event(
             AgentToolCallEvent(
                 arguments=function_args_str,
-                tool=tool.metadata
-                if tool
-                else ToolMetadata(description="unknown", name=function_name),
+                tool=(
+                    tool.metadata
+                    if tool
+                    else ToolMetadata(description="unknown", name=function_name)
+                ),
             )
         )
         with self.callback_manager.event(
@@ -569,12 +573,14 @@ def _run_step(
         openai_tools = [tool.metadata.to_openai_tool() for tool in tools]
 
         llm_chat_kwargs = self._get_llm_chat_kwargs(task, openai_tools, tool_choice)
-        agent_chat_response = self._get_agent_response(
+        response: str = self._get_agent_response(
             task, mode=mode, **llm_chat_kwargs
-        )
+        ).response
 
         # TODO: implement _should_continue
         latest_tool_calls = self.get_latest_tool_calls(task) or []
+        latest_tool_outputs: List[ToolOutput] = []
+
         if not self._should_continue(
             latest_tool_calls, task.extra_state["n_function_calls"]
         ):
@@ -590,13 +596,16 @@ def _run_step(
 
                 if tool_call.type != "function":
                     raise ValueError("Invalid tool type. Unsupported by OpenAI")
+
                 # TODO: maybe execute this with multi-threading
                 return_direct = self._call_function(
                     tools,
                     tool_call,
                     task.extra_state["new_memory"],
-                    task.extra_state["sources"],
+                    latest_tool_outputs,
                 )
+                task.extra_state["sources"].append(latest_tool_outputs[-1])
+
                 # change function call to the default value, if a custom function was given
                 # as an argument (none and auto are predefined by OpenAI)
                 if tool_choice not in ("auto", "none"):
@@ -605,16 +614,7 @@ def _run_step(
 
                 if return_direct and len(latest_tool_calls) == 1:
                     is_done = True
-                    response_str = task.extra_state["sources"][-1].content
-                    chat_response = ChatResponse(
-                        message=ChatMessage(
-                            role=MessageRole.ASSISTANT, content=response_str
-                        )
-                    )
-                    agent_chat_response = self._process_message(task, chat_response)
-                    agent_chat_response.is_dummy_stream = (
-                        mode == ChatResponseMode.STREAM
-                    )
+                    response = latest_tool_outputs[-1].content
                     break
 
             new_steps = (
@@ -629,10 +629,14 @@ def _run_step(
                 else []
             )
 
-        # attach next step to task
+        agent_response = AgentChatResponse(
+            response=response,
+            sources=latest_tool_outputs,
+            is_dummy_stream=(mode == ChatResponseMode.STREAM),
+        )
 
         return TaskStepOutput(
-            output=agent_chat_response,
+            output=agent_response,
             task_step=step,
             is_last=is_done,
             next_steps=new_steps,
@@ -656,12 +660,14 @@ async def _arun_step(
         openai_tools = [tool.metadata.to_openai_tool() for tool in tools]
 
         llm_chat_kwargs = self._get_llm_chat_kwargs(task, openai_tools, tool_choice)
-        agent_chat_response = await self._get_async_agent_response(
-            task, mode=mode, **llm_chat_kwargs
-        )
+        response: str = (
+            await self._get_async_agent_response(task, mode=mode, **llm_chat_kwargs)
+        ).response
 
         # TODO: implement _should_continue
         latest_tool_calls = self.get_latest_tool_calls(task) or []
+        latest_tool_outputs: List[ToolOutput] = []
+
         if not self._should_continue(
             latest_tool_calls, task.extra_state["n_function_calls"]
         ):
@@ -676,13 +682,16 @@ async def _arun_step(
 
                 if tool_call.type != "function":
                     raise ValueError("Invalid tool type. Unsupported by OpenAI")
+
                 # TODO: maybe execute this with multi-threading
                 return_direct = await self._acall_function(
                     tools,
                     tool_call,
                     task.extra_state["new_memory"],
-                    task.extra_state["sources"],
+                    latest_tool_outputs,
                 )
+                task.extra_state["sources"].append(latest_tool_outputs[-1])
+
                 # change function call to the default value, if a custom function was given
                 # as an argument (none and auto are predefined by OpenAI)
                 if tool_choice not in ("auto", "none"):
@@ -691,16 +700,7 @@ async def _arun_step(
 
                 if return_direct and len(latest_tool_calls) == 1:
                     is_done = True
-                    response_str = task.extra_state["sources"][-1].content
-                    chat_response = ChatResponse(
-                        message=ChatMessage(
-                            role=MessageRole.ASSISTANT, content=response_str
-                        )
-                    )
-                    agent_chat_response = self._process_message(task, chat_response)
-                    agent_chat_response.is_dummy_stream = (
-                        mode == ChatResponseMode.STREAM
-                    )
+                    response = latest_tool_outputs[-1].content
                     break
 
         # generate next step, append to task queue
@@ -716,8 +716,14 @@ async def _arun_step(
             else []
         )
 
+        agent_response = AgentChatResponse(
+            response=response,
+            sources=latest_tool_outputs,
+            is_dummy_stream=(mode == ChatResponseMode.STREAM),
+        )
+
         return TaskStepOutput(
-            output=agent_chat_response,
+            output=agent_response,
             task_step=step,
             is_last=is_done,
             next_steps=new_steps,
diff --git a/llama-index-integrations/agent/llama-index-agent-openai/tests/test_openai_agent.py b/llama-index-integrations/agent/llama-index-agent-openai/tests/test_openai_agent.py
index 33f00e2ffd..18e1e076c7 100644
--- a/llama-index-integrations/agent/llama-index-agent-openai/tests/test_openai_agent.py
+++ b/llama-index-integrations/agent/llama-index-agent-openai/tests/test_openai_agent.py
@@ -1,17 +1,28 @@
-from typing import Any, AsyncGenerator, Generator, List, Sequence
+import json
+import uuid
+from typing import (
+    Any,
+    AsyncGenerator,
+    Generator,
+    List,
+    Sequence,
+)
 from unittest.mock import MagicMock, patch
-
 from llama_index.core.agent.function_calling.step import (
     build_error_tool_output,
     build_missing_tool_message,
 )
+from llama_index.core.base.agent.types import TaskStepOutput
 import pytest
 from llama_index.agent.openai.base import OpenAIAgent
 from llama_index.agent.openai.step import (
     call_tool_with_error_handling,
     advanced_tool_call_parser,
 )
-from llama_index.core.base.llms.types import ChatMessage, ChatResponse
+from llama_index.core.base.llms.types import (
+    ChatMessage,
+    ChatResponse,
+)
 from llama_index.core.chat_engine.types import (
     AgentChatResponse,
     StreamingAgentChatResponse,
@@ -19,7 +30,6 @@
 from llama_index.core.llms.mock import MockLLM
 from llama_index.core.tools.function_tool import FunctionTool
 from llama_index.llms.openai import OpenAI
-
 from openai.types.chat.chat_completion import ChatCompletion, Choice
 from openai.types.chat.chat_completion_chunk import Choice as ChunkChoice
 from openai.types.chat.chat_completion_chunk import ChatCompletionChunk, ChoiceDelta
@@ -166,11 +176,17 @@ def echo(query: str) -> str:
 
 
 @pytest.fixture()
-def echo_function() -> Function:
+def malformed_echo_function() -> Function:
     test_result: str = "This is a test"
     return Function(name="echo", arguments=f'query = "{test_result}"')
 
 
+@pytest.fixture()
+def echo_function() -> Function:
+    test_result: str = "This is a test"
+    return Function(name="echo", arguments=json.dumps({"query": test_result}))
+
+
 class MockChatLLM(MockLLM):
     def __init__(self, responses: List[ChatMessage]) -> None:
         self._i = 0  # call counter, determines which response to return
@@ -322,12 +338,12 @@ def _add(a: int, b: int) -> int:
 def test_call_tool_with_malformed_function_call(
     MockSyncOpenAI: MagicMock,
     echo_tool: FunctionTool,
-    echo_function: Function,
+    malformed_echo_function: Function,
 ) -> None:
     """Test add step."""
     mock_instance = MockSyncOpenAI.return_value
     mock_instance.chat.completions.create.return_value = mock_chat_completion_tool_call(
-        function=echo_function
+        function=malformed_echo_function
     )
 
     llm = OpenAI(model="gpt-3.5-turbo")
@@ -339,7 +355,7 @@ def test_call_tool_with_malformed_function_call(
     ## NOTE: can only take a single step before finishing,
     # since mocked chat output does not call any tools
     task = agent.create_task(
-        f"This happens if tool call is malformed like:\n{echo_function.arguments}"
+        f"This happens if tool call is malformed like:\n{malformed_echo_function.arguments}"
     )
     step_output = agent.run_step(task.task_id)
     assert (
@@ -352,13 +368,13 @@ def test_call_tool_with_malformed_function_call(
 def test_call_tool_with_malformed_function_call_and_parser(
     MockSyncOpenAI: MagicMock,
     echo_tool: FunctionTool,
-    echo_function: Function,
+    malformed_echo_function: Function,
 ) -> None:
     """Test add step."""
     mock_instance = MockSyncOpenAI.return_value
     test_result: str = "This is a test"
     mock_instance.chat.completions.create.return_value = mock_chat_completion_tool_call(
-        function=echo_function
+        function=malformed_echo_function
     )
 
     llm = OpenAI(model="gpt-3.5-turbo")
@@ -371,7 +387,7 @@ def test_call_tool_with_malformed_function_call_and_parser(
     ## NOTE: can only take a single step before finishing,
     # since mocked chat output does not call any tools
     task = agent.create_task(
-        f"This happens if tool call is malformed like:\n{echo_function.arguments}"
+        f"This happens if tool call is malformed like:\n{malformed_echo_function.arguments}"
     )
     step_output = agent.run_step(task.task_id)
     assert str(step_output.output.sources[0]) == test_result
@@ -468,11 +484,11 @@ async def test_async_add_step(
 @patch("llama_index.llms.openai.base.SyncOpenAI")
 def test_run_step_returns_message_if_tool_not_found(
     MockSyncOpenAI: MagicMock,
-    echo_function: Function,
+    malformed_echo_function: Function,
 ) -> None:
     mock_instance = MockSyncOpenAI.return_value
     mock_instance.chat.completions.create.return_value = mock_chat_completion_tool_call(
-        function=echo_function
+        function=malformed_echo_function
     )
     llm = OpenAI(model="gpt-3.5-turbo")
     agent = OpenAIAgent.from_tools(tools=[], llm=llm)
@@ -485,9 +501,9 @@ def test_run_step_returns_message_if_tool_not_found(
     assert len(step_output.next_steps) == 1
     assert len(output_chat_response.sources) == 1
     assert output_chat_response.sources[0] == build_error_tool_output(
-        echo_function.name,
-        echo_function.arguments,
-        build_missing_tool_message(echo_function.name),
+        malformed_echo_function.name,
+        malformed_echo_function.arguments,
+        build_missing_tool_message(malformed_echo_function.name),
     )
 
 
@@ -517,3 +533,57 @@ async def test_arun_step_returns_message_if_tool_not_found(
         echo_function.arguments,
         build_missing_tool_message(echo_function.name),
     )
+
+
+@pytest.mark.asyncio()
+@pytest.mark.parametrize("method", ["run_step", "arun_step"])
+@patch("llama_index.llms.openai.base.SyncOpenAI")
+@patch("llama_index.llms.openai.base.AsyncOpenAI")
+async def test_run_step_returns_correct_sources_history(
+    MockAsyncOpenAI: MagicMock,
+    MockSyncOpenAI: MagicMock,
+    method: str,
+    echo_tool: FunctionTool,
+    echo_function: Function,
+) -> None:
+    num_steps = 4
+    llm = OpenAI(model="gpt-3.5-turbo")
+    agent = OpenAIAgent.from_tools(
+        tools=[echo_tool],
+        llm=llm,
+    )
+    task = agent.create_task("")
+    step_outputs: List[TaskStepOutput] = []
+
+    if method == "run_step":
+        mock_instance = MockSyncOpenAI.return_value
+        mock_instance.chat.completions.create.return_value = (
+            mock_chat_completion_tool_call(echo_function)
+        )
+    else:
+        mock_instance = MockAsyncOpenAI.return_value
+        mock_instance.chat.completions.create.side_effect = [
+            mock_achat_completion_tool_call(echo_function) for _ in range(num_steps)
+        ]
+
+    # Create steps
+    steps = [agent.agent_worker.initialize_step(task)]
+    for step_idx in range(num_steps - 1):
+        steps.append(
+            steps[-1].get_next_step(
+                step_id=str(uuid.uuid4()),
+                input=None,
+            )
+        )
+
+    # Run each step, invoking a single tool call each time
+    for step_idx in range(num_steps):
+        step_outputs.append(
+            agent.agent_worker.run_step(steps[step_idx], task)
+            if method == "run_step"
+            else await agent.agent_worker.arun_step(steps[step_idx], task)
+        )
+
+    # Ensure that each step only has one source for its one tool call
+    for step_idx in range(num_steps):
+        assert len(step_outputs[step_idx].output.sources) == 1


Code review comments:
diff_hunk: @@ -716,8 +716,14 @@ async def _arun_step(
             else []
         )
 
+        agent_response = AgentChatResponse(
+            response=response,
+            sources=latest_tool_outputs,
+            is_dummy_stream=(mode == ChatResponseMode.STREAM),
path: llama-index-integrations/agent/llama-index-agent-openai/llama_index/agent/openai/step.py
text: @user1:
I think this should only be set when its a return direct tool, not just when its streaming 👀 

@author:
Yes good catch! Turns out we don't need to modify much of the `agent_chat_response` stuff like I did, so I've reverted it and bypassed this issue.

@user1:
Yay!