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 

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

26 """ 

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

28 

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

30 """ 

31 self.root = Path(root) 

32 

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

34 """ 

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

36 

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

38 

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

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

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

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

43 """ 

44 root = self.root.resolve() 

45 path = Path(path_like) 

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

47 resolved = candidate.resolve() 

48 if not resolved.is_relative_to(root): 

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

50 return resolved 

51 

52 def write(self, *, key: str, content: str) -> str: 

53 """ 

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

55 

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

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

58 

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

60 :param content: The tool result to persist. 

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

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

63 """ 

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

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

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

67 return str(path) 

68 

69 def read(self, reference: str) -> str: 

70 """ 

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

72 

73 The resolved reference must stay within the store root: callers must treat it as an opaque 

74 store-scoped reference, not as an arbitrary filesystem path. 

75 

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

77 :returns: The stored content. 

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

79 """ 

80 return self._resolve_in_root(reference, subject="Result reference").read_text(encoding="utf-8") 

81 

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

83 """ 

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

85 

86 :returns: A dictionary representation of the store. 

87 """ 

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

89 

90 @classmethod 

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

92 """ 

93 Deserialize the store from a dictionary. 

94 

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

96 :returns: The deserialized `FileSystemToolResultStore`. 

97 """ 

98 return default_from_dict(cls, data)