Coverage for haystack/hooks/tool_result_offloading/utils.py: 100%
58 statements
« prev ^ index » next coverage.py v7.16.1, created at 2026-10-09 14:34 +0000
1# SPDX-FileCopyrightText: 2022-present deepset GmbH <info@deepset.ai>
2#
3# SPDX-License-Identifier: Apache-2.0
5import base64
6import mimetypes
7from pathlib import Path
8from typing import Any
10from haystack import logging
11from haystack.dataclasses import ChatMessage, FileContent, ImageContent, TextContent, ToolCallResult
12from haystack.hooks.tool_result_offloading.types import ToolResultStore
14logger = logging.getLogger(__name__)
16# Meta key marking an already-offloaded tool-result message; its value is the list of store references written.
17# Stops a second offload from offloading the pointer text again, since the pointer is itself a tool result.
18_OFFLOADED_META_KEY = "tool_result_offloaded"
20# Extension used for a binary block whose MIME type is unknown or maps to no known extension.
21_FALLBACK_EXTENSION = ".bin"
24def _content_block_payload(content_block: TextContent | ImageContent | FileContent) -> str:
25 """
26 Return the string a content block contributes to the conversation.
28 For an image or a file this is the base64 payload, which is what actually occupies the context window.
30 :param content_block: The content block to inspect.
31 :returns: The content block's text or base64 payload.
32 """
33 if isinstance(content_block, TextContent):
34 return content_block.text
35 return content_block.base64_image if isinstance(content_block, ImageContent) else content_block.base64_data
38def _offloadable_content_blocks(
39 result: ToolCallResult, store: ToolResultStore
40) -> list[TextContent | ImageContent | FileContent] | None:
41 """
42 Return a result's content blocks, or None when the result is empty or `store` cannot hold its image or file content.
44 A warning is logged when a text-only store is given image or file content.
46 :param result: The tool result to inspect.
47 :param store: The store the result would be written to.
48 :returns: The result's content blocks, or None when the result cannot or should not be offloaded.
49 """
50 # A plain string result is handled as a single text block, so callers have one shape to work with.
51 content_blocks: list[TextContent | ImageContent | FileContent] = (
52 [TextContent(text=result.result)] if isinstance(result.result, str) else list(result.result)
53 )
55 # `all` also covers a result with no content blocks at all.
56 if all(isinstance(content_block, TextContent) and not content_block.text for content_block in content_blocks):
57 return None
59 if not getattr(store, "supports_binary_content", False) and not all(
60 isinstance(content_block, TextContent) for content_block in content_blocks
61 ):
62 logger.warning(
63 "Tool '{tool}' produced a result with image or file content, but {store} does not support binary "
64 "content; leaving the result in context.",
65 tool=result.origin.tool_name,
66 store=type(store).__name__,
67 )
68 return None
70 return content_blocks
73def _offloaded_message(
74 message: ChatMessage,
75 *,
76 content_blocks: list[TextContent | ImageContent | FileContent],
77 store: ToolResultStore,
78 key_prefix: str,
79 preview_chars: int,
80 additional_meta: dict[str, Any] | None = None,
81) -> ChatMessage:
82 """
83 Write a tool result to the store and return a message with the same origin and error flag that points to it.
85 :param message: The tool-result message being offloaded. Must carry a tool result.
86 :param content_blocks: The result's content blocks, as returned by `_offloadable_content_blocks`.
87 :param store: The store to write the content blocks to.
88 :param key_prefix: The result's store key prefix, unique within the run.
89 :param preview_chars: Number of leading characters of each offloaded text to include in the pointer.
90 :param additional_meta: Extra metadata to record on the offloaded message, such as a compaction marker.
91 :returns: A new tool-result message whose content points at the stored result.
92 """
93 result = message.tool_call_result
94 if result is None:
95 raise ValueError("Only tool-result messages can be offloaded.")
97 references, pointer = _offload_content_blocks(
98 content_blocks=content_blocks, store=store, prefix=key_prefix, preview_chars=preview_chars
99 )
100 return ChatMessage.from_tool(
101 tool_result=pointer,
102 origin=result.origin,
103 error=result.error,
104 meta={**message.meta, **(additional_meta or {}), _OFFLOADED_META_KEY: references},
105 )
108def _offload_content_blocks(
109 content_blocks: list[TextContent | ImageContent | FileContent],
110 store: ToolResultStore,
111 prefix: str,
112 preview_chars: int,
113) -> tuple[list[str], str]:
114 """
115 Write a result's content blocks to the store and build the pointer that replaces them in the conversation.
117 Every content block goes to its own store entry. A single block keeps `prefix` as its key and gets a one-line
118 pointer; several blocks get position-suffixed keys and one numbered pointer line each.
120 :param content_blocks: The result's content blocks, in order.
121 :param store: The store to write to.
122 :param prefix: The result's store key prefix, as described above.
123 :param preview_chars: Number of leading characters of each offloaded text to include in the pointer.
124 :returns: The store references written, and the pointer text for the conversation.
125 """
126 references: list[str] = []
127 descriptions: list[str] = []
128 single = len(content_blocks) == 1
129 for position, content_block in enumerate(content_blocks):
130 reference, description = _offload_content_block(
131 content_block=content_block,
132 store=store,
133 key_prefix=prefix if single else f"{prefix}_{position}",
134 preview_chars=preview_chars,
135 )
136 references.append(reference)
137 descriptions.append(description)
139 if len(descriptions) == 1:
140 return references, f"Tool result offloaded to {descriptions[0]}"
142 numbered = [f"{position}. {description}" for position, description in enumerate(descriptions, start=1)]
143 return references, "\n".join([f"Tool result offloaded to {len(descriptions)} files:", *numbered])
146def _offload_content_block(
147 content_block: TextContent | ImageContent | FileContent, store: ToolResultStore, key_prefix: str, preview_chars: int
148) -> tuple[str, str]:
149 """
150 Write a single content block to the store and describe where it went.
152 :param content_block: The content block to offload.
153 :param store: The store to write to.
154 :param key_prefix: The content block's store key without its extension.
155 :param preview_chars: Number of leading characters of an offloaded text to include in its description.
156 :returns: The store reference the content block was written to, and a one-line description for the pointer.
157 """
158 if isinstance(content_block, TextContent):
159 text = content_block.text
160 reference = store.write(key=f"{key_prefix}.txt", content=text)
161 # An ellipsis marks a preview that was cut short, so the model can tell it is not the whole text.
162 preview = f"{text[:preview_chars]}{'...' if len(text) > preview_chars else ''}"
163 return reference, f"text ({len(text)} characters) at '{reference}'. Preview: {preview}"
165 if isinstance(content_block, ImageContent):
166 data = base64.b64decode(content_block.base64_image)
167 label = content_block.mime_type or "image"
168 filename = None
169 else:
170 data = base64.b64decode(content_block.base64_data)
171 label = content_block.mime_type or "file"
172 filename = content_block.filename
174 # What the tool called the file wins over its MIME type. Only the suffix is taken.
175 mime_extension = mimetypes.guess_extension(content_block.mime_type) if content_block.mime_type else None
176 extension = (Path(filename).suffix if filename else "") or mime_extension or _FALLBACK_EXTENSION
177 reference = store.write(key=f"{key_prefix}{extension}", content=data)
179 named = f" named '{filename}'" if filename else ""
180 return reference, f"{label}{named} ({len(data)} bytes) at '{reference}'"