Skip to content

Commit 4917e6d

Browse files
committed
feat(vcr): add LangChain4J model wrappers for VCR
Implements VCR wrappers for LangChain4J models: - VCREmbeddingModel: wraps EmbeddingModel for recording/replaying embeddings - VCRChatModel: wraps ChatLanguageModel for recording/replaying chat responses - VCREmbeddingInterceptor: low-level interceptor for embedding operations Features: - Support for all VCR modes (PLAYBACK, RECORD, PLAYBACK_OR_RECORD, OFF) - Redis-backed cassette storage integration - In-memory cassette cache for testing - Call counter tracking for unique cassette keys - Statistics tracking (cache hits, misses, recorded count) Unit tests verify recording and playback behavior.
1 parent 5087e3d commit 4917e6d

File tree

6 files changed

+1847
-0
lines changed

6 files changed

+1847
-0
lines changed
Lines changed: 264 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,264 @@
1+
package com.redis.vl.test.vcr;
2+
3+
import dev.langchain4j.data.message.AiMessage;
4+
import dev.langchain4j.data.message.ChatMessage;
5+
import dev.langchain4j.model.chat.ChatLanguageModel;
6+
import dev.langchain4j.model.output.Response;
7+
import edu.umd.cs.findbugs.annotations.SuppressFBWarnings;
8+
import java.util.HashMap;
9+
import java.util.List;
10+
import java.util.Map;
11+
import java.util.concurrent.atomic.AtomicInteger;
12+
13+
/**
14+
* VCR wrapper for LangChain4J ChatLanguageModel that records and replays LLM responses.
15+
*
16+
* <p>This class implements the ChatLanguageModel interface, allowing it to be used as a drop-in
17+
* replacement for any LangChain4J chat model. It provides VCR (Video Cassette Recorder)
18+
* functionality to record LLM responses during test execution and replay them in subsequent runs.
19+
*
20+
* <p>Usage:
21+
*
22+
* <pre>{@code
23+
* ChatLanguageModel openAiModel = OpenAiChatModel.builder()
24+
* .apiKey(System.getenv("OPENAI_API_KEY"))
25+
* .build();
26+
*
27+
* VCRChatModel vcrModel = new VCRChatModel(openAiModel);
28+
* vcrModel.setMode(VCRMode.PLAYBACK_OR_RECORD);
29+
* vcrModel.setTestId("MyTest.testMethod");
30+
*
31+
* // Use exactly like the original model
32+
* Response<AiMessage> response = vcrModel.generate(UserMessage.from("Hello"));
33+
* }</pre>
34+
*/
35+
@SuppressFBWarnings(
36+
value = "EI_EXPOSE_REP2",
37+
justification = "Delegate is intentionally stored and exposed for VCR functionality")
38+
public final class VCRChatModel implements ChatLanguageModel {
39+
40+
private final ChatLanguageModel delegate;
41+
private VCRCassetteStore cassetteStore;
42+
private VCRMode mode = VCRMode.PLAYBACK_OR_RECORD;
43+
private String testId = "unknown";
44+
private final AtomicInteger callCounter = new AtomicInteger(0);
45+
46+
// In-memory cassette storage for unit tests
47+
private final Map<String, String> cassettes = new HashMap<>();
48+
49+
// Statistics
50+
private int cacheHits = 0;
51+
private int cacheMisses = 0;
52+
private int recordedCount = 0;
53+
54+
/**
55+
* Creates a new VCRChatModel wrapping the given delegate.
56+
*
57+
* @param delegate The actual ChatLanguageModel to wrap
58+
*/
59+
public VCRChatModel(ChatLanguageModel delegate) {
60+
this.delegate = delegate;
61+
}
62+
63+
/**
64+
* Creates a new VCRChatModel wrapping the given delegate with Redis storage.
65+
*
66+
* @param delegate The actual ChatLanguageModel to wrap
67+
* @param cassetteStore The cassette store for persistence
68+
*/
69+
@SuppressFBWarnings(
70+
value = "EI_EXPOSE_REP2",
71+
justification = "VCRCassetteStore is intentionally shared")
72+
public VCRChatModel(ChatLanguageModel delegate, VCRCassetteStore cassetteStore) {
73+
this.delegate = delegate;
74+
this.cassetteStore = cassetteStore;
75+
}
76+
77+
/**
78+
* Sets the VCR mode.
79+
*
80+
* @param mode The VCR mode to use
81+
*/
82+
public void setMode(VCRMode mode) {
83+
this.mode = mode;
84+
}
85+
86+
/**
87+
* Gets the current VCR mode.
88+
*
89+
* @return The current VCR mode
90+
*/
91+
public VCRMode getMode() {
92+
return mode;
93+
}
94+
95+
/**
96+
* Sets the test identifier for cassette key generation.
97+
*
98+
* @param testId The test identifier (typically ClassName.methodName)
99+
*/
100+
public void setTestId(String testId) {
101+
this.testId = testId;
102+
}
103+
104+
/**
105+
* Gets the current test identifier.
106+
*
107+
* @return The current test identifier
108+
*/
109+
public String getTestId() {
110+
return testId;
111+
}
112+
113+
/** Resets the call counter. Useful when starting a new test method. */
114+
public void resetCallCounter() {
115+
callCounter.set(0);
116+
}
117+
118+
/**
119+
* Gets the underlying delegate model.
120+
*
121+
* @return The wrapped ChatLanguageModel
122+
*/
123+
@SuppressFBWarnings(
124+
value = "EI_EXPOSE_REP",
125+
justification = "Intentional exposure of delegate for advanced use cases")
126+
public ChatLanguageModel getDelegate() {
127+
return delegate;
128+
}
129+
130+
/**
131+
* Preloads a cassette for testing purposes.
132+
*
133+
* @param key The cassette key
134+
* @param response The response text to cache
135+
*/
136+
public void preloadCassette(String key, String response) {
137+
cassettes.put(key, response);
138+
}
139+
140+
/**
141+
* Gets the number of cache hits.
142+
*
143+
* @return Cache hit count
144+
*/
145+
public int getCacheHits() {
146+
return cacheHits;
147+
}
148+
149+
/**
150+
* Gets the number of cache misses.
151+
*
152+
* @return Cache miss count
153+
*/
154+
public int getCacheMisses() {
155+
return cacheMisses;
156+
}
157+
158+
/**
159+
* Gets the number of recorded responses.
160+
*
161+
* @return Recorded count
162+
*/
163+
public int getRecordedCount() {
164+
return recordedCount;
165+
}
166+
167+
/** Resets all statistics. */
168+
public void resetStatistics() {
169+
cacheHits = 0;
170+
cacheMisses = 0;
171+
recordedCount = 0;
172+
}
173+
174+
@Override
175+
public Response<AiMessage> generate(List<ChatMessage> messages) {
176+
return generateInternal(messages);
177+
}
178+
179+
@Override
180+
public String generate(String userMessage) {
181+
Response<AiMessage> response = generateInternal(userMessage);
182+
return response.content().text();
183+
}
184+
185+
private Response<AiMessage> generateInternal(Object input) {
186+
if (mode == VCRMode.OFF) {
187+
return callDelegate(input);
188+
}
189+
190+
String key = formatKey();
191+
192+
if (mode.isPlaybackMode()) {
193+
String cached = loadCassette(key);
194+
if (cached != null) {
195+
cacheHits++;
196+
return Response.from(AiMessage.from(cached));
197+
}
198+
199+
if (mode == VCRMode.PLAYBACK) {
200+
throw new VCRCassetteMissingException(key, testId);
201+
}
202+
203+
// PLAYBACK_OR_RECORD - fall through to record
204+
}
205+
206+
// Record mode or cache miss in PLAYBACK_OR_RECORD
207+
cacheMisses++;
208+
Response<AiMessage> response = callDelegate(input);
209+
String responseText = response.content().text();
210+
saveCassette(key, responseText);
211+
recordedCount++;
212+
213+
return response;
214+
}
215+
216+
private String loadCassette(String key) {
217+
// Check in-memory first
218+
String inMemory = cassettes.get(key);
219+
if (inMemory != null) {
220+
return inMemory;
221+
}
222+
223+
// Check Redis if available
224+
if (cassetteStore != null) {
225+
com.google.gson.JsonObject cassette = cassetteStore.retrieve(key);
226+
if (cassette != null && cassette.has("response")) {
227+
return cassette.get("response").getAsString();
228+
}
229+
}
230+
231+
return null;
232+
}
233+
234+
private void saveCassette(String key, String response) {
235+
// Save to in-memory
236+
cassettes.put(key, response);
237+
238+
// Save to Redis if available
239+
if (cassetteStore != null) {
240+
com.google.gson.JsonObject cassette = new com.google.gson.JsonObject();
241+
cassette.addProperty("response", response);
242+
cassette.addProperty("testId", testId);
243+
cassette.addProperty("type", "chat");
244+
cassetteStore.store(key, cassette);
245+
}
246+
}
247+
248+
@SuppressWarnings("unchecked")
249+
private Response<AiMessage> callDelegate(Object input) {
250+
if (input instanceof List) {
251+
return delegate.generate((List<ChatMessage>) input);
252+
} else if (input instanceof String) {
253+
String text = delegate.generate((String) input);
254+
return Response.from(AiMessage.from(text));
255+
} else {
256+
throw new IllegalArgumentException("Unsupported input type: " + input.getClass());
257+
}
258+
}
259+
260+
private String formatKey() {
261+
int index = callCounter.incrementAndGet();
262+
return String.format("vcr:chat:%s:%04d", testId, index);
263+
}
264+
}

0 commit comments

Comments
 (0)