Skip to content

Commit 5c29380

Browse files
committed
feat: Add LLM call memoization/caching for llm_query (Issue #82)
This PR implements memoization for llm_query and llm_query_batched calls to avoid redundant API requests during recursive RLM execution. Key features: - Thread-safe LRU cache with configurable max_size - Optional TTL for time-based expiration - Cache statistics for monitoring (hits, misses, evictions) - SHA-256 based cache keys for (prompt, model) pairs - Zero-configuration opt-in via cache parameter on RLM class Files added: - rlm/utils/cache.py: Core cache implementation (~400 lines) - tests/test_cache.py: 39 comprehensive tests Files modified: - rlm/core/rlm.py: Added cache parameter to __init__ - rlm/environments/local_repl.py: Integrated caching in _llm_query - rlm/utils/__init__.py: Export cache classes Usage: from rlm import RLM from rlm.utils import LLMCallCache cache = LLMCallCache(max_size=1000, ttl_seconds=3600) rlm = RLM(backend='openai', cache=cache) Addresses: #82
1 parent beb0603 commit 5c29380

5 files changed

Lines changed: 1010 additions & 8 deletions

File tree

‎rlm/core/rlm.py‎

Lines changed: 14 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,7 @@
11
import time
22
from collections.abc import Callable
33
from contextlib import contextmanager
4-
from typing import Any
4+
from typing import TYPE_CHECKING, Any
55

66
from rlm.clients import BaseLM, get_client
77
from rlm.core.lm_handler import LMHandler
@@ -38,6 +38,9 @@
3838
from rlm.utils.rlm_utils import filter_sensitive_keys
3939
from rlm.utils.token_utils import count_tokens, get_context_limit
4040

41+
if TYPE_CHECKING:
42+
from rlm.utils.cache import LLMCallCache
43+
4144

4245
class RLM:
4346
"""
@@ -70,6 +73,7 @@ def __init__(
7073
custom_sub_tools: dict[str, Any] | None = None,
7174
compaction: bool = False,
7275
compaction_threshold_pct: float = 0.85,
76+
cache: "LLMCallCache | None" = None,
7377
on_subcall_start: Callable[[int, str, str], None] | None = None,
7478
on_subcall_complete: Callable[[int, str, float, str | None], None] | None = None,
7579
on_iteration_start: Callable[[int, int], None] | None = None,
@@ -102,6 +106,9 @@ def __init__(
102106
when root context reaches compaction_threshold_pct of the model's context limit.
103107
compaction_threshold_pct: When compaction is on, trigger summarization when root
104108
message token count reaches this fraction of the model context limit (default 0.85).
109+
cache: Optional LLMCallCache for memoizing llm_query calls. When provided, identical prompts
110+
will return cached responses instead of making redundant API calls. Useful for recursive
111+
workloads with overlapping subproblems (e.g., Fibonacci-like decomposition).
105112
on_subcall_start: Callback fired when a child RLM starts. Args: (depth, model, prompt_preview).
106113
on_subcall_complete: Callback fired when a child RLM completes. Args: (depth, model, duration, error_or_none).
107114
on_iteration_start: Callback fired when an iteration starts. Args: (depth, iteration_num).
@@ -133,6 +140,9 @@ def __init__(
133140
self.compaction = compaction
134141
self.compaction_threshold_pct = compaction_threshold_pct
135142

143+
# LLM call cache for memoization
144+
self.cache = cache
145+
136146
self.depth = depth
137147
self.max_depth = max_depth
138148
self.max_iterations = max_iterations
@@ -238,6 +248,9 @@ def _spawn_completion_context(self, prompt: str | dict[str, Any]):
238248
env_kwargs["custom_sub_tools"] = self.custom_sub_tools
239249
if self.compaction and self.environment_type == "local":
240250
env_kwargs["compaction"] = True
251+
# Pass cache for LLM call memoization
252+
if self.cache is not None and self.environment_type == "local":
253+
env_kwargs["cache"] = self.cache
241254
environment: BaseEnv = get_environment(self.environment_type, env_kwargs)
242255

243256
if self.persistent:

‎rlm/environments/local_repl.py‎

Lines changed: 59 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -10,7 +10,7 @@
1010
import uuid
1111
from collections.abc import Callable
1212
from contextlib import contextmanager
13-
from typing import Any
13+
from typing import TYPE_CHECKING, Any
1414

1515
from rlm.core.comms_utils import LMRequest, send_lm_request, send_lm_request_batched
1616
from rlm.core.types import REPLResult, RLMChatCompletion
@@ -21,6 +21,9 @@
2121
validate_custom_tools,
2222
)
2323

24+
if TYPE_CHECKING:
25+
from rlm.utils.cache import LLMCallCache
26+
2427
# =============================================================================
2528
# Safe Builtins
2629
# =============================================================================
@@ -135,12 +138,14 @@ def __init__(
135138
custom_tools: dict[str, Any] | None = None,
136139
custom_sub_tools: dict[str, Any] | None = None,
137140
compaction: bool = False,
141+
cache: "LLMCallCache | None" = None,
138142
**kwargs,
139143
):
140144
super().__init__(persistent=persistent, depth=depth, **kwargs)
141145

142146
self.lm_handler_address = lm_handler_address
143147
self.subcall_fn = subcall_fn # Callback for recursive RLM calls (depth > 1 support)
148+
self.cache = cache # Optional LLM call cache for memoization
144149
self.original_cwd = os.getcwd()
145150
self.temp_dir = tempfile.mkdtemp(prefix=f"repl_env_{uuid.uuid4()}_")
146151
self._lock = threading.Lock()
@@ -242,6 +247,7 @@ def _llm_query(self, prompt: str, model: str | None = None) -> str:
242247
"""Query the LM with a single plain completion (no REPL, no recursion).
243248
244249
This always makes a direct LM call via the handler, regardless of depth.
250+
Supports optional caching to avoid redundant API calls.
245251
246252
Args:
247253
prompt: The prompt to send to the LM.
@@ -250,22 +256,32 @@ def _llm_query(self, prompt: str, model: str | None = None) -> str:
250256
if not self.lm_handler_address:
251257
return "Error: No LM handler configured"
252258

253-
try:
259+
# Define the actual API call function
260+
def make_api_call() -> str:
254261
request = LMRequest(prompt=prompt, model=model, depth=self.depth)
255262
response = send_lm_request(self.lm_handler_address, request)
256263

257264
if not response.success:
258-
return f"Error: {response.error}"
265+
raise RuntimeError(response.error)
259266

260267
self._pending_llm_calls.append(response.chat_completion)
261268
return response.chat_completion.response
269+
270+
try:
271+
# Use cache if available
272+
if self.cache is not None:
273+
result, was_cached = self.cache.get_or_call(prompt, model, make_api_call)
274+
return result
275+
else:
276+
return make_api_call()
262277
except Exception as e:
263278
return f"Error: LM query failed - {e}"
264279

265280
def _llm_query_batched(self, prompts: list[str], model: str | None = None) -> list[str]:
266281
"""Query the LM with multiple prompts concurrently (no REPL, no recursion).
267282
268283
This always makes direct LM calls via the handler, regardless of depth.
284+
Supports optional caching to avoid redundant API calls.
269285
270286
Args:
271287
prompts: List of prompts to send to the LM.
@@ -276,20 +292,56 @@ def _llm_query_batched(self, prompts: list[str], model: str | None = None) -> li
276292
"""
277293
if not self.lm_handler_address:
278294
return ["Error: No LM handler configured"] * len(prompts)
295+
279296
try:
297+
# With caching: check cache first, only send uncached prompts to API
298+
if self.cache is not None:
299+
results: list[str | None] = [None] * len(prompts)
300+
uncached_prompts: list[tuple[int, str]] = []
301+
302+
# Check cache for each prompt
303+
for i, prompt in enumerate(prompts):
304+
cached = self.cache.get(prompt, model)
305+
if cached is not None:
306+
results[i] = cached
307+
self.cache.stats.record_hit()
308+
else:
309+
uncached_prompts.append((i, prompt))
310+
self.cache.stats.record_miss()
311+
312+
# Make API calls only for uncached prompts
313+
if uncached_prompts:
314+
uncached_texts = [p for _, p in uncached_prompts]
315+
responses = send_lm_request_batched(
316+
self.lm_handler_address, uncached_texts, model=model, depth=self.depth
317+
)
318+
319+
for (idx, prompt), response in zip(uncached_prompts, responses):
320+
if not response.success:
321+
results[idx] = f"Error: {response.error}"
322+
else:
323+
self._pending_llm_calls.append(response.chat_completion)
324+
result = response.chat_completion.response
325+
results[idx] = result
326+
# Cache the successful response
327+
self.cache.set(prompt, model, result)
328+
329+
return [r if r is not None else "Error: Unexpected cache failure" for r in results]
330+
331+
# Without caching: send all prompts directly
280332
responses = send_lm_request_batched(
281333
self.lm_handler_address, prompts, model=model, depth=self.depth
282334
)
283335

284-
results = []
336+
results_list = []
285337
for response in responses:
286338
if not response.success:
287-
results.append(f"Error: {response.error}")
339+
results_list.append(f"Error: {response.error}")
288340
else:
289341
self._pending_llm_calls.append(response.chat_completion)
290-
results.append(response.chat_completion.response)
342+
results_list.append(response.chat_completion.response)
291343

292-
return results
344+
return results_list
293345
except Exception as e:
294346
return [f"Error: LM query failed - {e}"] * len(prompts)
295347

‎rlm/utils/__init__.py‎

Lines changed: 15 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,15 @@
1+
"""
2+
RLM Utility modules.
3+
4+
Contains helper utilities for parsing, prompts, exceptions, token counting,
5+
and LLM call caching.
6+
"""
7+
8+
from rlm.utils.cache import CacheEntry, CacheStats, LLMCallCache, create_cache
9+
10+
__all__ = [
11+
"CacheEntry",
12+
"CacheStats",
13+
"LLMCallCache",
14+
"create_cache",
15+
]

0 commit comments

Comments
 (0)