Coverage for src/dataknobs_llm/prompts/adapters/inmemory_adapter.py: 42%

1"""In-memory resource adapters for testing and demos.

2

3 This module provides simple in-memory adapters that return predefined results.

4 Useful for testing, demos, and examples where you need predictable behavior

5 without external dependencies.

6"""

7

8 from typing import Any, Dict, List, Optional

9 from .resource_adapter import ResourceAdapterBase, ResourceAdapter, AsyncResourceAdapter, BaseSearchLogic

10

11

12 class InMemoryAdapterBase(ResourceAdapterBase):

13 """Base class with shared logic for in-memory adapters.

14

15 This class contains all the core logic shared between sync and async

16 in-memory adapters, following the DRY principle.

17 """

18

19 def __init__(

20 self,

21 search_results: Optional[List[Dict[str, Any]]] = None,

22 data: Optional[Dict[str, Any]] = None,

23 name: str = "inmemory"

24 ):

25 """Initialize in-memory adapter base.

26

27 Args:

28 search_results: List of results to return from search().

29 Each result should be a dict with at least a "content" key.

30 Optional keys: "score", "metadata", etc.

31 data: Optional dictionary for key-value storage (used by get_value)

32 name: Name identifier for this adapter

33 """

34 super().__init__(name=name)

35 self._search_results = search_results or []

36 self._data = data or {}

37 self.search_count = 0 # Track number of search calls for testing

38

39 def reset(self):

40 """Reset search count. Useful for testing cache behavior."""

41 self.search_count = 0

42

43 def _get_value_impl(

44 self,

45 key: str,

46 default: Any = None,

47 context: Optional[Dict[str, Any]] = None

48 ) -> Any:

49 """Shared implementation for get_value.

50

51 Args:

52 key: Key to look up

53 default: Value to return if key is not found

54 context: Optional context (unused in in-memory adapter)

55

56 Returns:

57 Value at the key, or default if not found

58 """

59 return self._data.get(key, default)

60

61 def _search_impl(

62 self,

63 query: str,

64 k: int = 5,

65 filters: Optional[Dict[str, Any]] = None,

66 **kwargs

67 ) -> List[Dict[str, Any]]:

68 """Shared implementation for search.

69

70 Args:

71 query: Search query (unused - returns all configured results)

72 k: Maximum number of results to return

73 filters: Optional filters to apply via BaseSearchLogic

74 **kwargs: Additional options:

75 - min_score: Minimum score threshold

76

77 Returns:

78 List of search results (up to k items)

79 """

80 self.search_count += 1

81 results = self._search_results.copy()

82

83 # Ensure results have proper structure

84 results = [

85 BaseSearchLogic.format_search_result(

86 r.get("content", ""),

87 score=r.get("score", 1.0),

88 metadata=r.get("metadata", {})

89 )

90 for r in results

91 ]

92

93 # Apply filters if provided

94 if filters:

95 results = BaseSearchLogic.filter_results(results, filters=filters)

96

97 # Apply min_score filter if provided

98 min_score = kwargs.get('min_score', 0.0)

99 if min_score > 0:

100 results = BaseSearchLogic.filter_results(results, min_score=min_score)

101

102 return results[:k]

103

104

105 class InMemoryAdapter(InMemoryAdapterBase, ResourceAdapter):

106 """Synchronous in-memory adapter with predefined search results.

107

108 This adapter is designed for testing, demos, and examples. It returns

109 predefined search results and optionally stores key-value pairs.

110

111 Features:

112 - Predefined search results for predictable testing

113 - Optional key-value storage via get_value()

114 - Simple and fast (no external I/O)

115 - Search count tracking for testing

116

117 Example:

118 >>> # Simple usage with search results

119 >>> adapter = InMemoryAdapter(

120 ... search_results=[

121 ... {"content": "Python is a programming language", "score": 0.9},

122 ... {"content": "Python was created by Guido van Rossum", "score": 0.8}

123 ... ],

124 ... name="docs"

125 ... )

126 >>> results = adapter.search("python")

127 >>> len(results)

128 2

129 >>> results[0]["content"]

130 'Python is a programming language'

131 >>> adapter.search_count # Track how many times search was called

132 1

133

134 >>> # With key-value storage

135 >>> adapter = InMemoryAdapter(

136 ... data={"language": "Python", "version": "3.11"},

137 ... name="config"

138 ... )

139 >>> adapter.get_value("language")

140 'Python'

141 """

142

143 def __init__(

144 self,

145 search_results: Optional[List[Dict[str, Any]]] = None,

146 data: Optional[Dict[str, Any]] = None,

147 name: str = "inmemory"

148 ):

149 """Initialize synchronous in-memory adapter.

150

151 Args:

152 search_results: List of results to return from search()

153 data: Optional dictionary for key-value storage

154 name: Name identifier for this adapter

155 """

156 super().__init__(search_results=search_results, data=data, name=name)

157

158 def get_value(

159 self,

160 key: str,

161 default: Any = None,

162 context: Optional[Dict[str, Any]] = None

163 ) -> Any:

164 """Retrieve a value by key from the in-memory data.

165

166 Args:

167 key: Key to look up

168 default: Value to return if key is not found

169 context: Optional context (unused in in-memory adapter)

170

171 Returns:

172 Value at the key, or default if not found

173 """

174 return self._get_value_impl(key, default, context)

175

176 def search(

177 self,

178 query: str,

179 k: int = 5,

180 filters: Optional[Dict[str, Any]] = None,

181 **kwargs

182 ) -> List[Dict[str, Any]]:

183 """Return predefined search results.

184

185 Args:

186 query: Search query (unused - returns all configured results)

187 k: Maximum number of results to return

188 filters: Optional filters to apply via BaseSearchLogic

189 **kwargs: Additional options:

190 - min_score: Minimum score threshold

191

192 Returns:

193 List of search results (up to k items)

194 """

195 return self._search_impl(query, k, filters, **kwargs)

196

197

198 class InMemoryAsyncAdapter(InMemoryAdapterBase, AsyncResourceAdapter):

199 """Asynchronous in-memory adapter with predefined search results.

200

201 Async version of InMemoryAdapter. Useful for testing async code paths

202 or maintaining consistency in async codebases.

203

204 Example:

205 >>> adapter = InMemoryAsyncAdapter(

206 ... search_results=[

207 ... {"content": "Result 1", "score": 0.9},

208 ... {"content": "Result 2", "score": 0.8}

209 ... ]

210 ... )

211 >>> results = await adapter.search("test")

212 >>> len(results)

213 2

214 >>> adapter.search_count

215 1

216 """

217

218 def __init__(

219 self,

220 search_results: Optional[List[Dict[str, Any]]] = None,

221 data: Optional[Dict[str, Any]] = None,

222 name: str = "inmemory_async"

223 ):

224 """Initialize asynchronous in-memory adapter.

225

226 Args:

227 search_results: List of results to return from search()

228 data: Optional dictionary for key-value storage

229 name: Name identifier for this adapter

230 """

231 super().__init__(search_results=search_results, data=data, name=name)

232

233 async def get_value(

234 self,

235 key: str,

236 default: Any = None,

237 context: Optional[Dict[str, Any]] = None

238 ) -> Any:

239 """Retrieve a value by key from the in-memory data (async).

240

241 Args:

242 key: Key to look up

243 default: Value to return if key is not found

244 context: Optional context (unused in in-memory adapter)

245

246 Returns:

247 Value at the key, or default if not found

248 """

249 return self._get_value_impl(key, default, context)

250

251 async def search(

252 self,

253 query: str,

254 k: int = 5,

255 filters: Optional[Dict[str, Any]] = None,

256 **kwargs

257 ) -> List[Dict[str, Any]]:

258 """Return predefined search results (async).

259

260 Args:

261 query: Search query (unused - returns all configured results)

262 k: Maximum number of results to return

263 filters: Optional filters to apply via BaseSearchLogic

264 **kwargs: Additional options:

265 - min_score: Minimum score threshold

266

267 Returns:

268 List of search results (up to k items)

269 """

270 return self._search_impl(query, k, filters, **kwargs)