-
Notifications
You must be signed in to change notification settings - Fork 2.1k
/
cloudflare_kv.ts
80 lines (73 loc) · 2.6 KB
/
cloudflare_kv.ts
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
import type { KVNamespace } from "@cloudflare/workers-types";
import {
BaseCache,
getCacheKey,
serializeGeneration,
deserializeStoredGeneration,
} from "@langchain/core/caches";
import { Generation } from "@langchain/core/outputs";
/**
* @deprecated Install and import from "@langchain/cloudflare" instead.
*
* Represents a specific implementation of a caching mechanism using Cloudflare KV
* as the underlying storage system. It extends the `BaseCache` class and
* overrides its methods to provide the Cloudflare KV-specific logic.
* @example
* ```typescript
* // Example of using OpenAI with Cloudflare KV as cache in a Cloudflare Worker
* const cache = new CloudflareKVCache(env.KV_NAMESPACE);
* const model = new ChatAnthropic({
* cache,
* });
* const response = await model.invoke("How are you today?");
* return new Response(JSON.stringify(response), {
* headers: { "content-type": "application/json" },
* });
*
* ```
*/
export class CloudflareKVCache extends BaseCache {
private binding: KVNamespace;
constructor(binding: KVNamespace) {
super();
this.binding = binding;
}
/**
* Retrieves data from the cache. It constructs a cache key from the given
* `prompt` and `llmKey`, and retrieves the corresponding value from the
* Cloudflare KV namespace.
* @param prompt The prompt used to construct the cache key.
* @param llmKey The LLM key used to construct the cache key.
* @returns An array of Generations if found, null otherwise.
*/
public async lookup(prompt: string, llmKey: string) {
let idx = 0;
let key = getCacheKey(prompt, llmKey, String(idx));
let value = await this.binding.get(key);
const generations: Generation[] = [];
while (value) {
generations.push(deserializeStoredGeneration(JSON.parse(value)));
idx += 1;
key = getCacheKey(prompt, llmKey, String(idx));
value = await this.binding.get(key);
}
return generations.length > 0 ? generations : null;
}
/**
* Updates the cache with new data. It constructs a cache key from the
* given `prompt` and `llmKey`, and stores the `value` in the Cloudflare KV
* namespace.
* @param prompt The prompt used to construct the cache key.
* @param llmKey The LLM key used to construct the cache key.
* @param value The value to be stored in the cache.
*/
public async update(prompt: string, llmKey: string, value: Generation[]) {
for (let i = 0; i < value.length; i += 1) {
const key = getCacheKey(prompt, llmKey, String(i));
await this.binding.put(
key,
JSON.stringify(serializeGeneration(value[i]))
);
}
}
}