feat: control display of tool execution results to llm (#920)

Author: puzzle-solver

packages/ragbits-agents/CHANGELOG.md

@@ -17,6 +17,7 @@
 - Add supervisor post-processor (#830)
 - Add support for todo lists generated by agents with examples (#827)
 - Add long-term semantic memory tools for agents (#839)
+- Add control over what output from the tool is passed to the LLM (#920)
 - Add support for confirmation requests in agents (#853) (#914)
 - Add hooks system (pre- and post-tool) for lifecycle event interception (#914)
 

packages/ragbits-agents/src/ragbits/agents/_main.py

@@ -42,7 +42,7 @@
     StreamingPostProcessor,
     stream_with_post_processing,
 )
-from ragbits.agents.tool import Tool, ToolCallResult, ToolChoice
+from ragbits.agents.tool import Tool, ToolCallResult, ToolChoice, ToolReturn
 from ragbits.core.audit.traces import trace
 from ragbits.core.llms.base import (
     LLM,
@@ -538,7 +538,9 @@ async def _run_without_post_processing(
                 ):
                     if isinstance(result, ToolCallResult):
                         tool_calls.append(result)
-                        prompt_with_history = prompt_with_history.add_tool_use_message(**result.__dict__)
+                        prompt_with_history = prompt_with_history.add_tool_use_message(
+                            id=result.id, name=result.name, arguments=result.arguments, result=result.result
+                        )
 
                 turn_count += 1
             else:
@@ -764,7 +766,9 @@ async def _stream_internal(  # noqa: PLR0912, PLR0915
                         elif isinstance(result, ToolCallResult):
                             # Add ALL tool results to history (including pending confirmations)
                             # This allows the agent to see them in the next turn
-                            prompt_with_history = prompt_with_history.add_tool_use_message(**result.__dict__)
+                            prompt_with_history = prompt_with_history.add_tool_use_message(
+                                id=result.id, name=result.name, arguments=result.arguments, result=result.result
+                            )
                             returned_tool_call = True
 
                     # If we have pending confirmations, prepare for text-only summary generation
@@ -1004,20 +1008,23 @@ async def _execute_tool(
                     else asyncio.to_thread(tool.on_tool_call, **call_args)
                 )
 
-                if isinstance(tool_output, AgentResultStreaming):
+                if isinstance(tool_output, ToolReturn):
+                    tool_return = tool_output
+                elif isinstance(tool_output, AgentResultStreaming):
                     async for downstream_item in tool_output:
                         if context.stream_downstream_events:
                             yield DownstreamAgentResult(agent_id=tool.id, item=downstream_item)
-
-                    tool_output = {
-                        "content": tool_output.content,
+                    metadata = {
                         "metadata": tool_output.metadata,
                         "tool_calls": tool_output.tool_calls,
                         "usage": tool_output.usage,
                     }
+                    tool_return = ToolReturn(value=tool_output.content, metadata=metadata)
+                else:
+                    tool_return = ToolReturn(value=tool_output, metadata=None)
 
                 outputs.result = {
-                    "tool_output": tool_output,
+                    "tool_output": tool_return.value,
                     "tool_call_id": tool_call.id,
                 }
 
@@ -1044,7 +1051,8 @@ async def _execute_tool(
             id=tool_call.id,
             name=tool_call.name,
             arguments=tool_call.arguments,
-            result=tool_output,
+            result=tool_return.value,
+            metadata=tool_return.metadata,
         )
 
     @requires_dependencies(["a2a.types"], "a2a")

packages/ragbits-agents/src/ragbits/agents/tool.py

@@ -18,6 +18,19 @@
     from pydantic_ai import Tool as PydanticAITool
 
 
+@dataclass
+class ToolReturn:
+    """
+    Represents an object returned from the tool. If a tool wants to return a value with some content hidden
+    from LLM, it needs to return an object of this class directly.
+    """
+
+    value: Any
+    "Value passed directly to LLM as a result of the tool"
+    metadata: Any = None
+    "Metadata not passed to the LLM, but which can be used in the application later on"
+
+
 @dataclass
 class ToolCallResult:
     """
@@ -31,7 +44,9 @@ class ToolCallResult:
     arguments: dict[str, Any]
     """Dictionary containing the arguments passed to the tool"""
     result: Any
-    """The output from the tool call."""
+    """The output from the tool call passed to the LLM"""
+    metadata: Any = None
+    """Metadata returned from a tool that is not meant to be seen by the LLM"""
 
 
 @dataclass