Coverage for haystack/hooks/tool_result_offloading/stores.py: 100%

1# SPDX-FileCopyrightText: 2022-present deepset GmbH <info@deepset.ai> 

2# 

3# SPDX-License-Identifier: Apache-2.0 

4 

5from pathlib import Path 

6from typing import Any 

7 

8from haystack.core.serialization import default_from_dict, default_to_dict 

9from haystack.hooks.tool_result_offloading.types import ToolResultStore 

10 

11 

12class FileSystemToolResultStore(ToolResultStore): 

13 """ 

14 A `ToolResultStore` that writes offloaded tool results to files under a root directory on the local file system. 

15 

16 ```python 

17 from haystack.hooks.tool_result_offloading import FileSystemToolResultStore 

18 

19 store = FileSystemToolResultStore(root="tool_results") 

20 reference = store.write(key="search_1.txt", content="...") 

21 store.read(reference) 

22 ``` 

23 

24 Binary content is supported too: `write` takes bytes (an offloaded image or file) and `read` returns them 

25 unchanged. 

26 """ 

27 

28 supports_binary_content = True 

29 

30 def __init__(self, root: str | Path) -> None: 

31 """ 

32 Initialize the store with the root directory results are written under. 

33 

34 :param root: Directory under which result files are written. Created on first write if it does not exist. 

35 """ 

36 self.root = Path(root) 

37 

38 def _resolve_in_root(self, path_like: str | Path, *, subject: str) -> Path: 

39 """ 

40 Resolve a path-like value and ensure it stays within the configured store root. 

41 

42 Relative values are interpreted relative to `self.root`; absolute values are used as-is. 

43 

44 :param path_like: Relative or absolute path-like value to resolve. 

45 :param subject: Human-readable label used in the error message. 

46 :returns: The resolved absolute path within the store root. 

47 :raises ValueError: If the resolved path escapes the store root. 

48 """ 

49 root = self.root.resolve() 

50 path = Path(path_like) 

51 candidate = path if path.is_absolute() else root / path 

52 resolved = candidate.resolve() 

53 if not resolved.is_relative_to(root): 

54 raise ValueError(f"{subject} '{path_like}' resolves outside the store root '{root}'.") 

55 return resolved 

56 

57 def write(self, *, key: str, content: str | bytes) -> str: 

58 """ 

59 Write `content` to `<root>/<key>`, creating parent directories, and return the file path. 

60 

61 Text is written UTF-8 encoded; bytes (an offloaded image or file) are written verbatim. 

62 

63 The resolved target must stay within the root directory: a `key` that escapes it (e.g. containing `../` or an 

64 absolute path) is rejected, so a tool-provided key cannot write outside the store. 

65 

66 :param key: Relative file name for the result within the store root. 

67 :param content: The tool result to persist, as text or as raw bytes. 

68 :returns: The absolute path the content was written to, as a string, for use with `read`. 

69 :raises ValueError: If `key` resolves to a location outside the store root. 

70 """ 

71 path = self._resolve_in_root(key, subject="Result key") 

72 path.parent.mkdir(parents=True, exist_ok=True) 

73 if isinstance(content, bytes): 

74 path.write_bytes(content) 

75 else: 

76 path.write_text(content, encoding="utf-8") 

77 return str(path) 

78 

79 def read(self, reference: str) -> str | bytes: 

80 """ 

81 Read back the content previously written to `reference`. 

82 

83 A file whose bytes are valid UTF-8 is returned as a string, so text results round trip unchanged; anything 

84 else (an offloaded image or file) is returned as raw bytes. 

85 

86 The resolved reference must stay within the store root: it is a store-scoped reference returned by `write`, 

87 to be passed back unchanged, not an arbitrary filesystem path callers can build themselves. 

88 

89 :param reference: A store reference returned by `write`. 

90 :returns: The stored content, as text when it decodes as UTF-8 and as bytes otherwise. 

91 :raises ValueError: If `reference` resolves to a location outside the store root. 

92 """ 

93 data = self._resolve_in_root(reference, subject="Result reference").read_bytes() 

94 try: 

95 return data.decode("utf-8") 

96 except UnicodeDecodeError: 

97 return data 

98 

99 def to_dict(self) -> dict[str, Any]: 

100 """ 

101 Serialize the store, storing its root directory as a string. 

102 

103 :returns: A dictionary representation of the store. 

104 """ 

105 return default_to_dict(self, root=str(self.root)) 

106 

107 @classmethod 

108 def from_dict(cls, data: dict[str, Any]) -> "FileSystemToolResultStore": 

109 """ 

110 Deserialize the store from a dictionary. 

111 

112 :param data: A dictionary representation produced by `to_dict`. 

113 :returns: The deserialized `FileSystemToolResultStore`. 

114 """ 

115 return default_from_dict(cls, data)