Skip to content

Commit 5087e3d

Browse files
committed
feat(vcr): add Redis cassette store for VCR persistence
Implements VCRCassetteStore for storing and retrieving recorded API responses in Redis: - Store cassettes as Redis JSON documents with structured keys - Support for single and batch embedding cassettes - Handle both array and object JSON responses from Jedis - Add VCRCassetteMissingException for strict playback mode - Unit tests for cassette store operations Cassette key format: vcr:{type}:{testId}:{callIndex}
1 parent 2519978 commit 5087e3d

File tree

3 files changed

+491
-0
lines changed

3 files changed

+491
-0
lines changed
Lines changed: 49 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,49 @@
1+
package com.redis.vl.test.vcr;
2+
3+
/**
4+
* Exception thrown when a VCR cassette is not found during playback mode.
5+
*
6+
* <p>This exception indicates that the test expected to find a recorded cassette but none was
7+
* available. To fix this, run the test in RECORD or PLAYBACK_OR_RECORD mode first.
8+
*/
9+
public class VCRCassetteMissingException extends RuntimeException {
10+
11+
private static final long serialVersionUID = 1L;
12+
13+
private final String cassetteKey;
14+
private final String testId;
15+
16+
/**
17+
* Creates a new exception.
18+
*
19+
* @param cassetteKey the key that was not found
20+
* @param testId the test identifier
21+
*/
22+
public VCRCassetteMissingException(String cassetteKey, String testId) {
23+
super(
24+
String.format(
25+
"VCR cassette not found for test '%s'%nCassette key: %s%n"
26+
+ "Run with VCRMode.RECORD or VCRMode.PLAYBACK_OR_RECORD to record this interaction",
27+
testId, cassetteKey));
28+
this.cassetteKey = cassetteKey;
29+
this.testId = testId;
30+
}
31+
32+
/**
33+
* Gets the cassette key that was not found.
34+
*
35+
* @return the cassette key
36+
*/
37+
public String getCassetteKey() {
38+
return cassetteKey;
39+
}
40+
41+
/**
42+
* Gets the test identifier.
43+
*
44+
* @return the test ID
45+
*/
46+
public String getTestId() {
47+
return testId;
48+
}
49+
}
Lines changed: 249 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,249 @@
1+
package com.redis.vl.test.vcr;
2+
3+
import com.google.gson.Gson;
4+
import com.google.gson.JsonArray;
5+
import com.google.gson.JsonObject;
6+
import edu.umd.cs.findbugs.annotations.SuppressFBWarnings;
7+
import java.util.Objects;
8+
import redis.clients.jedis.JedisPooled;
9+
import redis.clients.jedis.json.Path2;
10+
11+
/**
12+
* Stores and retrieves VCR cassettes (recorded API responses) in Redis.
13+
*
14+
* <p>Cassettes are stored as Redis JSON documents with the following key format:
15+
*
16+
* <pre>vcr:{type}:{testId}:{callIndex}</pre>
17+
*
18+
* Where:
19+
*
20+
* <ul>
21+
* <li>{type} - The type of cassette (e.g., "embedding", "llm", "chat")
22+
* <li>{testId} - The unique test identifier (e.g., "MyTest.testMethod")
23+
* <li>{callIndex} - Zero-padded call index within the test (e.g., "0001")
24+
* </ul>
25+
*/
26+
public class VCRCassetteStore {
27+
28+
private static final String KEY_PREFIX = "vcr";
29+
private static final Gson GSON = new Gson();
30+
31+
private final JedisPooled jedis;
32+
33+
/**
34+
* Creates a new cassette store.
35+
*
36+
* @param jedis the Redis client
37+
*/
38+
@SuppressFBWarnings(
39+
value = "EI_EXPOSE_REP2",
40+
justification = "JedisPooled is intentionally shared for connection pooling")
41+
public VCRCassetteStore(JedisPooled jedis) {
42+
this.jedis = jedis;
43+
}
44+
45+
/**
46+
* Formats a cassette key.
47+
*
48+
* @param type the cassette type
49+
* @param testId the test identifier
50+
* @param callIndex the call index (1-based)
51+
* @return the formatted key
52+
*/
53+
public static String formatKey(String type, String testId, int callIndex) {
54+
return String.format("%s:%s:%s:%04d", KEY_PREFIX, type, testId, callIndex);
55+
}
56+
57+
/**
58+
* Parses a cassette key into its components.
59+
*
60+
* @param key the key to parse
61+
* @return array of [prefix, type, testId, callIndex] or null if invalid
62+
*/
63+
public static String[] parseKey(String key) {
64+
if (key == null) {
65+
return null;
66+
}
67+
String[] parts = key.split(":");
68+
if (parts.length != 4 || !KEY_PREFIX.equals(parts[0])) {
69+
return null;
70+
}
71+
return parts;
72+
}
73+
74+
/**
75+
* Creates a cassette JSON object for an embedding.
76+
*
77+
* @param embedding the embedding vector
78+
* @param testId the test identifier
79+
* @param model the model name
80+
* @return the cassette JSON object
81+
*/
82+
public static JsonObject createEmbeddingCassette(float[] embedding, String testId, String model) {
83+
Objects.requireNonNull(embedding, "embedding cannot be null");
84+
Objects.requireNonNull(testId, "testId cannot be null");
85+
Objects.requireNonNull(model, "model cannot be null");
86+
87+
JsonObject cassette = new JsonObject();
88+
cassette.addProperty("type", "embedding");
89+
cassette.addProperty("testId", testId);
90+
cassette.addProperty("model", model);
91+
cassette.addProperty("timestamp", System.currentTimeMillis());
92+
93+
JsonArray embeddingArray = new JsonArray();
94+
for (float value : embedding) {
95+
embeddingArray.add(value);
96+
}
97+
cassette.add("embedding", embeddingArray);
98+
99+
return cassette;
100+
}
101+
102+
/**
103+
* Creates a cassette JSON object for batch embeddings.
104+
*
105+
* @param embeddings the embedding vectors
106+
* @param testId the test identifier
107+
* @param model the model name
108+
* @return the cassette JSON object
109+
*/
110+
public static JsonObject createBatchEmbeddingCassette(
111+
float[][] embeddings, String testId, String model) {
112+
Objects.requireNonNull(embeddings, "embeddings cannot be null");
113+
Objects.requireNonNull(testId, "testId cannot be null");
114+
Objects.requireNonNull(model, "model cannot be null");
115+
116+
JsonObject cassette = new JsonObject();
117+
cassette.addProperty("type", "batch_embedding");
118+
cassette.addProperty("testId", testId);
119+
cassette.addProperty("model", model);
120+
cassette.addProperty("timestamp", System.currentTimeMillis());
121+
122+
JsonArray embeddingsArray = new JsonArray();
123+
for (float[] embedding : embeddings) {
124+
JsonArray embeddingArray = new JsonArray();
125+
for (float value : embedding) {
126+
embeddingArray.add(value);
127+
}
128+
embeddingsArray.add(embeddingArray);
129+
}
130+
cassette.add("embeddings", embeddingsArray);
131+
132+
return cassette;
133+
}
134+
135+
/**
136+
* Extracts embedding from a cassette JSON object.
137+
*
138+
* @param cassette the cassette object
139+
* @return the embedding vector or null if not present
140+
*/
141+
public static float[] extractEmbedding(JsonObject cassette) {
142+
if (cassette == null || !cassette.has("embedding")) {
143+
return null;
144+
}
145+
146+
JsonArray embeddingArray = cassette.getAsJsonArray("embedding");
147+
float[] embedding = new float[embeddingArray.size()];
148+
for (int i = 0; i < embeddingArray.size(); i++) {
149+
embedding[i] = embeddingArray.get(i).getAsFloat();
150+
}
151+
return embedding;
152+
}
153+
154+
/**
155+
* Extracts batch embeddings from a cassette JSON object.
156+
*
157+
* @param cassette the cassette object
158+
* @return the embedding vectors or null if not present
159+
*/
160+
public static float[][] extractBatchEmbeddings(JsonObject cassette) {
161+
if (cassette == null || !cassette.has("embeddings")) {
162+
return null;
163+
}
164+
165+
JsonArray embeddingsArray = cassette.getAsJsonArray("embeddings");
166+
float[][] embeddings = new float[embeddingsArray.size()][];
167+
168+
for (int i = 0; i < embeddingsArray.size(); i++) {
169+
JsonArray embeddingArray = embeddingsArray.get(i).getAsJsonArray();
170+
embeddings[i] = new float[embeddingArray.size()];
171+
for (int j = 0; j < embeddingArray.size(); j++) {
172+
embeddings[i][j] = embeddingArray.get(j).getAsFloat();
173+
}
174+
}
175+
return embeddings;
176+
}
177+
178+
/**
179+
* Stores a cassette in Redis.
180+
*
181+
* @param key the cassette key
182+
* @param cassette the cassette data
183+
*/
184+
public void store(String key, JsonObject cassette) {
185+
if (jedis == null) {
186+
throw new IllegalStateException("Redis client not initialized");
187+
}
188+
jedis.jsonSet(key, Path2.ROOT_PATH, GSON.toJson(cassette));
189+
}
190+
191+
/**
192+
* Retrieves a cassette from Redis.
193+
*
194+
* @param key the cassette key
195+
* @return the cassette data or null if not found
196+
*/
197+
public JsonObject retrieve(String key) {
198+
if (jedis == null) {
199+
return null;
200+
}
201+
Object result = jedis.jsonGet(key, Path2.ROOT_PATH);
202+
if (result == null) {
203+
return null;
204+
}
205+
206+
// Handle both array and object responses from jsonGet
207+
// Some Jedis versions/configurations return arrays when using ROOT_PATH
208+
String jsonString = result.toString();
209+
com.google.gson.JsonElement element =
210+
GSON.fromJson(jsonString, com.google.gson.JsonElement.class);
211+
212+
if (element.isJsonArray()) {
213+
com.google.gson.JsonArray array = element.getAsJsonArray();
214+
if (array.isEmpty()) {
215+
return null;
216+
}
217+
// Return the first element if it's wrapped in an array
218+
return array.get(0).getAsJsonObject();
219+
} else if (element.isJsonObject()) {
220+
return element.getAsJsonObject();
221+
}
222+
223+
return null;
224+
}
225+
226+
/**
227+
* Checks if a cassette exists.
228+
*
229+
* @param key the cassette key
230+
* @return true if the cassette exists
231+
*/
232+
public boolean exists(String key) {
233+
if (jedis == null) {
234+
return false;
235+
}
236+
return jedis.exists(key);
237+
}
238+
239+
/**
240+
* Deletes a cassette.
241+
*
242+
* @param key the cassette key
243+
*/
244+
public void delete(String key) {
245+
if (jedis != null) {
246+
jedis.del(key);
247+
}
248+
}
249+
}

0 commit comments

Comments
 (0)