Skip to content

Commit beca974

Browse files
SuperMarioYLcodelionclaude
authored
Add opt-in file persistence to the memory plugin (#111) (#315)
* Add opt-in file persistence to the memory plugin (#111) The memory plugin builds a fresh in-RAM Memory() per request, so extracted memories never survive across calls. This adds an opt-in file-backed store, gated on the OPTILLM_MEMORY_FILE env var: - Memory(persist_path=...) loads saved items on init and writes after each add() - run() reads the path from OPTILLM_MEMORY_FILE; unset => behaviour unchanged - loads degrade gracefully on missing/corrupt/non-list files (logged, no raise) - saves are atomic (temp file + os.replace) and bounded by max_size - README documents the env var; unit test covers round-trip, corrupt/missing files, max_size truncation, and the default no-I/O path * Bump version to 0.3.18 Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> * Add Frame SAST scan of PR-changed files to CI Adds a security-scan workflow that runs the Frame neuro-symbolic SAST tool (lambdasec/frame, pinned) on the Python files changed by a pull request. Scanning only the PR's added/modified files surfaces issues introduced by the change without failing on pre-existing findings elsewhere in the tree. The job fails only on high/critical severity, so lower-confidence categories (e.g. insecure_random on algorithmic sampling) do not block merges. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> --------- Co-authored-by: Asankhaya Sharma <codelion@users.noreply.github.com> Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
1 parent 6901e2d commit beca974

6 files changed

Lines changed: 207 additions & 7 deletions

File tree

Lines changed: 69 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,69 @@
1+
name: Security Scan (Frame SAST)
2+
3+
# Runs the Frame neuro-symbolic SAST tool (https://github.com/lambdasec/frame)
4+
# on the Python files changed by a pull request. Scanning only the PR's changed
5+
# files surfaces issues introduced by the change without failing on pre-existing
6+
# findings elsewhere in the tree. The job fails only on high/critical severity.
7+
8+
on:
9+
pull_request:
10+
branches: [ main ]
11+
12+
permissions:
13+
contents: read
14+
15+
jobs:
16+
frame-scan:
17+
name: Frame SAST (changed files)
18+
runs-on: ubuntu-latest
19+
steps:
20+
- name: Checkout (full history for diff)
21+
uses: actions/checkout@v4
22+
with:
23+
fetch-depth: 0
24+
25+
- name: Set up Python
26+
uses: actions/setup-python@v4
27+
with:
28+
python-version: '3.12'
29+
30+
- name: Install Frame (pinned)
31+
run: |
32+
git clone https://github.com/lambdasec/frame.git /tmp/frame
33+
git -C /tmp/frame checkout 3223ac44320b4870782d9aad03514b4d3c876e0a
34+
pip install "/tmp/frame[scan]"
35+
36+
- name: Scan Python files changed in this PR
37+
env:
38+
BASE_SHA: ${{ github.event.pull_request.base.sha }}
39+
run: |
40+
set -uo pipefail
41+
42+
# Added/copied/modified/renamed Python files in this PR (skip deletions).
43+
mapfile -t FILES < <(git diff --name-only --diff-filter=ACMR "$BASE_SHA" HEAD -- '*.py')
44+
45+
if [ "${#FILES[@]}" -eq 0 ]; then
46+
echo "No Python files changed in this PR - nothing to scan."
47+
exit 0
48+
fi
49+
50+
echo "Scanning ${#FILES[@]} changed Python file(s) (fail on high/critical):"
51+
printf ' %s\n' "${FILES[@]}"
52+
53+
FAIL=0
54+
for f in "${FILES[@]}"; do
55+
# File may have been renamed away or removed in a later commit.
56+
[ -f "$f" ] || continue
57+
echo "::group::Frame scan $f"
58+
if ! frame scan "$f" --fail-on high; then
59+
FAIL=1
60+
echo "::error file=$f::Frame flagged a high/critical severity issue in $f"
61+
fi
62+
echo "::endgroup::"
63+
done
64+
65+
if [ "$FAIL" -ne 0 ]; then
66+
echo "Frame SAST found high/critical severity issue(s) in changed files."
67+
exit 1
68+
fi
69+
echo "No high/critical severity issues in changed files."

‎README.md‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -190,7 +190,7 @@ optillm
190190
| MCP Client | `mcp` | Implements the model context protocol (MCP) client, enabling you to use any LLM with any MCP Server |
191191
| Router | `router` | Uses the [optillm-modernbert-large](https://huggingface.co/codelion/optillm-modernbert-large) model to route requests to different approaches based on the user prompt |
192192
| Chain-of-Code | `coc` | Implements a chain of code approach that combines CoT with code execution and LLM based code simulation |
193-
| Memory | `memory` | Implements a short term memory layer, enables you to use unbounded context length with any LLM |
193+
| Memory | `memory` | Implements a short term memory layer, enables you to use unbounded context length with any LLM. Set `OPTILLM_MEMORY_FILE` to opt in to file-backed persistence so memories survive across requests |
194194
| Privacy | `privacy` | Anonymize PII data in request and deanonymize it back to original value in response |
195195
| Read URLs | `readurls` | Reads all URLs found in the request, fetches the content at the URL and adds it to the context |
196196
| Execute Code | `executecode` | Enables use of code interpreter to execute python code in requests and LLM generated responses |

‎optillm/__init__.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,5 @@
11
# Version information
2-
__version__ = "0.3.17"
2+
__version__ = "0.3.18"
33

44
import os as _os
55

‎optillm/plugins/memory_plugin.py‎

Lines changed: 72 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,24 +1,90 @@
1+
import json
2+
import logging
3+
import os
14
import re
2-
from typing import Tuple, List
5+
import tempfile
6+
from typing import Optional, Tuple, List
37
import numpy as np
48
from sklearn.feature_extraction.text import TfidfVectorizer
59
from sklearn.metrics.pairwise import cosine_similarity
610

711
SLUG = "memory"
812

13+
# Environment variable that opts a run in to file-backed memory. When it is set,
14+
# the Memory store loads any previously saved items on init and persists after
15+
# every add(); when it is unset the plugin behaves exactly as before (in-RAM,
16+
# reset per request).
17+
MEMORY_FILE_ENV = "OPTILLM_MEMORY_FILE"
18+
19+
logger = logging.getLogger(__name__)
20+
921
class Memory:
10-
def __init__(self, max_size: int = 100):
22+
def __init__(self, max_size: int = 100, persist_path: Optional[str] = None):
1123
self.max_size = max_size
1224
self.items: List[str] = []
1325
self.vectorizer = TfidfVectorizer()
1426
self.vectors = None
1527
self.completion_tokens = 0
28+
self.persist_path = persist_path
29+
if self.persist_path:
30+
self._load_from_file()
1631

1732
def add(self, item: str):
1833
if len(self.items) >= self.max_size:
1934
self.items.pop(0)
2035
self.items.append(item)
2136
self.vectors = None # Reset vectors to force recalculation
37+
if self.persist_path:
38+
self._save_to_file()
39+
40+
def _load_from_file(self):
41+
"""Load persisted items from ``persist_path`` (opt-in).
42+
43+
A missing file (first run) or a corrupt/unreadable one degrades
44+
gracefully to an empty in-memory store, so persistence can never make a
45+
request fail at startup.
46+
"""
47+
try:
48+
with open(self.persist_path, "r", encoding="utf-8") as f:
49+
data = json.load(f)
50+
except FileNotFoundError:
51+
return # nothing persisted yet
52+
except (OSError, ValueError) as e:
53+
logger.warning("Could not load memory from %s: %s", self.persist_path, e)
54+
return
55+
56+
if not isinstance(data, list):
57+
logger.warning(
58+
"Ignoring memory file %s: expected a JSON list of strings",
59+
self.persist_path,
60+
)
61+
return
62+
63+
# Keep only strings and honour max_size (most recent items win).
64+
items = [x for x in data if isinstance(x, str)]
65+
self.items = items[-self.max_size:]
66+
self.vectors = None
67+
68+
def _save_to_file(self):
69+
"""Atomically persist the current items to ``persist_path`` (opt-in).
70+
71+
Writes to a temp file in the same directory and ``os.replace``s it into
72+
place so a crash mid-write cannot corrupt an existing store. Any I/O
73+
error is logged and swallowed rather than raised.
74+
"""
75+
try:
76+
directory = os.path.dirname(os.path.abspath(self.persist_path))
77+
os.makedirs(directory, exist_ok=True)
78+
fd, tmp_path = tempfile.mkstemp(dir=directory, suffix=".tmp")
79+
try:
80+
with os.fdopen(fd, "w", encoding="utf-8") as f:
81+
json.dump(self.items, f)
82+
os.replace(tmp_path, self.persist_path)
83+
finally:
84+
if os.path.exists(tmp_path):
85+
os.remove(tmp_path)
86+
except OSError as e:
87+
logger.warning("Could not save memory to %s: %s", self.persist_path, e)
2288

2389
def get_relevant(self, query: str, n: int = 10) -> List[str]:
2490
if not self.items:
@@ -91,7 +157,10 @@ def extract_key_information(system_message, text: str, query: str, client, model
91157
return margins, response.usage.completion_tokens
92158

93159
def run(system_prompt: str, initial_query: str, client, model: str) -> Tuple[str, int]:
94-
memory = Memory()
160+
# Opt-in file-backed memory: when OPTILLM_MEMORY_FILE is set, items persist
161+
# across requests; when unset, behaviour is unchanged (fresh in-RAM store).
162+
persist_path = os.environ.get(MEMORY_FILE_ENV) or None
163+
memory = Memory(persist_path=persist_path)
95164
query, context = extract_query(initial_query)
96165
completion_tokens = 0
97166

‎pyproject.toml‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
44

55
[project]
66
name = "optillm"
7-
version = "0.3.17"
7+
version = "0.3.18"
88
description = "An optimizing inference proxy for LLMs."
99
readme = "README.md"
1010
license = "Apache-2.0"

‎tests/test_plugins.py‎

Lines changed: 63 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -68,6 +68,62 @@ def test_memory_plugin_structure():
6868
assert hasattr(plugin, 'Memory') # Check for Memory class
6969

7070

71+
def test_memory_plugin_persistence():
72+
"""Test opt-in file-backed persistence for the memory plugin (issue #111).
73+
74+
Covers the round trip, graceful degradation on a missing/corrupt file, and
75+
that the default (no persist_path) never touches the filesystem.
76+
"""
77+
import json
78+
import tempfile
79+
from optillm.plugins.memory_plugin import Memory
80+
81+
with tempfile.TemporaryDirectory() as tmp:
82+
path = os.path.join(tmp, "memory.json")
83+
84+
# Round trip: items written by one instance load into the next.
85+
m1 = Memory(persist_path=path)
86+
m1.add("alpha")
87+
m1.add("beta")
88+
assert os.path.exists(path), "add() should persist when persist_path is set"
89+
90+
m2 = Memory(persist_path=path)
91+
assert m2.items == ["alpha", "beta"], "persisted items should load on init"
92+
93+
# A missing file is a valid first run, not an error.
94+
missing = os.path.join(tmp, "does_not_exist.json")
95+
m3 = Memory(persist_path=missing)
96+
assert m3.items == []
97+
98+
# A corrupt file degrades to an empty store without raising.
99+
corrupt = os.path.join(tmp, "corrupt.json")
100+
with open(corrupt, "w", encoding="utf-8") as f:
101+
f.write("{not valid json")
102+
m4 = Memory(persist_path=corrupt)
103+
assert m4.items == []
104+
105+
# A non-list JSON payload is ignored rather than trusted.
106+
wrong_shape = os.path.join(tmp, "wrong.json")
107+
with open(wrong_shape, "w", encoding="utf-8") as f:
108+
json.dump({"items": ["x"]}, f)
109+
m5 = Memory(persist_path=wrong_shape)
110+
assert m5.items == []
111+
112+
# max_size is honoured on load (most recent items win).
113+
big = os.path.join(tmp, "big.json")
114+
with open(big, "w", encoding="utf-8") as f:
115+
json.dump([str(i) for i in range(10)], f)
116+
m6 = Memory(max_size=3, persist_path=big)
117+
assert m6.items == ["7", "8", "9"]
118+
119+
# Default behaviour is unchanged: no persist_path means no file I/O.
120+
no_persist = os.path.join(tmp, "should_not_be_created.json")
121+
m7 = Memory()
122+
m7.add("gamma")
123+
assert not os.path.exists(no_persist)
124+
assert m7.persist_path is None
125+
126+
71127
def test_genselect_plugin():
72128
"""Test genselect plugin module"""
73129
import optillm.plugins.genselect_plugin as plugin
@@ -392,7 +448,13 @@ def test_no_relative_import_errors():
392448
print("✅ Memory plugin structure test passed")
393449
except Exception as e:
394450
print(f"❌ Memory plugin structure test failed: {e}")
395-
451+
452+
try:
453+
test_memory_plugin_persistence()
454+
print("✅ Memory plugin persistence test passed")
455+
except Exception as e:
456+
print(f"❌ Memory plugin persistence test failed: {e}")
457+
396458
try:
397459
test_genselect_plugin()
398460
print("✅ GenSelect plugin test passed")

0 commit comments

Comments
 (0)