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

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: List[Dict[str, Any]] | None = None,

22 data: Dict[str, Any] | None = 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: Dict[str, Any] | None = 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: Dict[str, Any] | None = 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 >>> first_result = results[0]

130 >>> first_result.get('content')

131 'Python is a programming language'

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

133 1

134

135 >>> # With key-value storage

136 >>> adapter = InMemoryAdapter(

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

138 ... name="config"

139 ... )

140 >>> adapter.get_value('language')

141 'Python'

142 """

143

144 def __init__(

145 self,

146 search_results: List[Dict[str, Any]] | None = None,

147 data: Dict[str, Any] | None = None,

148 name: str = "inmemory"

149 ):

150 """Initialize synchronous in-memory adapter.

151

152 Args:

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

154 data: Optional dictionary for key-value storage

155 name: Name identifier for this adapter

156 """

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

158

159 def get_value(

160 self,

161 key: str,

162 default: Any = None,

163 context: Dict[str, Any] | None = None

164 ) -> Any:

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

166

167 Args:

168 key: Key to look up

169 default: Value to return if key is not found

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

171

172 Returns:

173 Value at the key, or default if not found

174 """

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

176

177 def search(

178 self,

179 query: str,

180 k: int = 5,

181 filters: Dict[str, Any] | None = None,

182 **kwargs: Any

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

184 """Return predefined search results.

185

186 Args:

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

188 k: Maximum number of results to return

189 filters: Optional filters to apply via BaseSearchLogic

190 **kwargs: Additional options:

191 - min_score: Minimum score threshold

192

193 Returns:

194 List of search results (up to k items)

195 """

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

197

198

199 class InMemoryAsyncAdapter(InMemoryAdapterBase, AsyncResourceAdapter):

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

201

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

203 or maintaining consistency in async codebases.

204

205 Example:

206 >>> adapter = InMemoryAsyncAdapter(

207 ... search_results=[

208 ... {'content': "Result 1", 'score': 0.9},

209 ... {'content': "Result 2", 'score': 0.8}

210 ... ]

211 ... )

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

213 >>> len(results)

214 2

215 >>> adapter.search_count

216 1

217 """

218

219 def __init__(

220 self,

221 search_results: List[Dict[str, Any]] | None = None,

222 data: Dict[str, Any] | None = None,

223 name: str = "inmemory_async"

224 ):

225 """Initialize asynchronous in-memory adapter.

226

227 Args:

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

229 data: Optional dictionary for key-value storage

230 name: Name identifier for this adapter

231 """

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

233

234 async def get_value(

235 self,

236 key: str,

237 default: Any = None,

238 context: Dict[str, Any] | None = None

239 ) -> Any:

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

241

242 Args:

243 key: Key to look up

244 default: Value to return if key is not found

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

246

247 Returns:

248 Value at the key, or default if not found

249 """

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

251

252 async def search(

253 self,

254 query: str,

255 k: int = 5,

256 filters: Dict[str, Any] | None = None,

257 **kwargs

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

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

260

261 Args:

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

263 k: Maximum number of results to return

264 filters: Optional filters to apply via BaseSearchLogic

265 **kwargs: Additional options:

266 - min_score: Minimum score threshold

267

268 Returns:

269 List of search results (up to k items)

270 """

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