1import { assert, describe, it } from "@effect/vitest"
2import { Effect, Layer, Predicate, Ref, Schema } from "effect"
3import { Chat, IdGenerator, Prompt, Tool, Toolkit } from "effect/ai"
4import { Persistence } from "effect/persistence"
5import { TestClock } from "effect/testing"
6import * as TestUtils from "./utils.ts"
7
8const withConstantIdGenerator = (id: string) =>
9 Effect.provideService(IdGenerator.IdGenerator, {
10 generateId: () => Effect.succeed(id)
11 })
12
13const PersistenceLayer = Layer.provideMerge(
14 Chat.layerPersisted({ storeId: "chat" }),
15 Persistence.layerBackingMemory
16)
17
18describe("Chat", () => {
19 it.effect("does not replay a completed approved tool", () => {
20 let executions = 0
21 const toolkit = Toolkit.make(Tool.make("Counter", {
22 parameters: Schema.Struct({}),
23 success: Schema.Number,
24 needsApproval: true
25 }))
26
27 return Effect.gen(function*() {
28 const chat = yield* Chat.empty
29 const initial = yield* chat.generateText({ prompt: "Run the counter", toolkit }).pipe(
30 TestUtils.withLanguageModel({
31 generateText: [{ type: "tool-call", id: "counter-call", name: "Counter", params: {} }]
32 })
33 )
34 const approval = initial.content.find((part) => part.type === "tool-approval-request")
35 assert.isDefined(approval)
36
37 yield* chat.generateText({
38 toolkit,
39 prompt: [Prompt.toolMessage({
40 content: [Prompt.toolApprovalResponsePart({ approvalId: approval.approvalId, approved: true })]
41 })]
42 }).pipe(
43 TestUtils.withLanguageModel({ generateText: [{ type: "text", text: "Complete" }] })
44 )
45
46 yield* chat.generateText({ prompt: "Continue", toolkit }).pipe(
47 TestUtils.withLanguageModel({ generateText: [{ type: "text", text: "Continued" }] })
48 )
49
50 assert.strictEqual(executions, 1)
51 }).pipe(
52 Effect.provide(toolkit.toLayer({ Counter: () => Effect.sync(() => ++executions) }))
53 )
54 })
55
56 it.effect("should persist chat history to the backing persistence store", () =>
57 Effect.gen(function*() {
58 const storeId = "chat"
59 const chatId = "1"
60
61 const backing = yield* Persistence.BackingPersistence
62 const persistence = yield* Chat.Persistence
63
64 const store = yield* backing.make(storeId)
65 const chat = yield* persistence.getOrCreate(chatId)
66
67 yield* chat.generateText({ prompt: "test user message" }).pipe(
68 TestUtils.withLanguageModel({
69 generateText: [{
70 type: "text",
71 text: "test assistant message"
72 }]
73 })
74 )
75
76 const chatHistory = yield* Ref.get(chat.history)
77 const encodedHistory = yield* store.get(chatId)
78 const storedHistory = Predicate.isNotUndefined(encodedHistory)
79 ? yield* Schema.decodeUnknownEffect(Prompt.Prompt)(encodedHistory)
80 : undefined
81
82 const options = { [Chat.Persistence.key]: { messageId: "msg_abc123" } }
83 const expectedHistory = Prompt.make([
84 { role: "user", content: [{ type: "text", text: "test user message" }], options },
85 { role: "assistant", content: [{ type: "text", text: "test assistant message" }], options }
86 ])
87
88 assert.deepStrictEqual(chatHistory, expectedHistory)
89 assert.deepStrictEqual(chatHistory, storedHistory)
90 }).pipe(withConstantIdGenerator("msg_abc123"), Effect.provide(PersistenceLayer)))
91
92 it.effect("should respect the specified time to live", () =>
93 Effect.gen(function*() {
94 const storeId = "chat"
95 const chatId = "1"
96
97 const backing = yield* Persistence.BackingPersistence
98 const persistence = yield* Chat.Persistence
99
100 const store = yield* backing.make(storeId)
101 const chat = yield* persistence.getOrCreate(chatId, {
102 timeToLive: "30 days"
103 })
104
105 yield* chat.generateText({ prompt: "test user message" }).pipe(
106 TestUtils.withLanguageModel({
107 generateText: [{
108 type: "text",
109 text: "test assistant message"
110 }]
111 })
112 )
113
114 const encodedHistory = yield* store.get(chatId)
115 const storedHistory = Predicate.isNotUndefined(encodedHistory)
116 ? yield* Schema.decodeUnknownEffect(Prompt.Prompt)(encodedHistory)
117 : undefined
118
119 const options = { [Chat.Persistence.key]: { messageId: "msg_abc123" } }
120 const expectedHistory = Prompt.make([
121 { role: "user", content: [{ type: "text", text: "test user message" }], options },
122 { role: "assistant", content: [{ type: "text", text: "test assistant message" }], options }
123 ])
124
125 assert.deepStrictEqual(storedHistory, expectedHistory)
126
127 // Simulate chat expiration
128 yield* TestClock.adjust("30 days")
129
130 const afterExpiration = yield* store.get(chatId)
131
132 assert.isUndefined(afterExpiration)
133 }).pipe(withConstantIdGenerator("msg_abc123"), Effect.provide(PersistenceLayer)))
134
135 it.effect("should prefer the message identifier of the most recent assistant message", () =>
136 Effect.gen(function*() {
137 const storeId = "chat"
138 const chatId = "2"
139
140 const backing = yield* Persistence.BackingPersistence
141 const persistence = yield* Chat.Persistence
142
143 const store = yield* backing.make(storeId)
144 const chat = yield* persistence.getOrCreate(chatId)
145
146 const options = { [Chat.Persistence.key]: { messageId: "msg_123abc" } }
147 const history = Prompt.make([
148 { role: "user", content: "first user message", options },
149 { role: "assistant", content: "first assistant message", options }
150 ])
151 yield* Ref.set(chat.history, history)
152 yield* chat.save
153
154 yield* chat.generateText({ prompt: "second user message" }).pipe(
155 TestUtils.withLanguageModel({
156 generateText: [{
157 type: "text",
158 text: "second assistant message"
159 }]
160 })
161 )
162
163 const encodedHistory = yield* store.get(chatId)
164 const storedHistory = Predicate.isNotUndefined(encodedHistory)
165 ? yield* Schema.decodeUnknownEffect(Prompt.Prompt)(encodedHistory)
166 : undefined
167 const expectedHistory = Prompt.concat(history, [
168 { role: "user", content: "second user message", options },
169 { role: "assistant", content: "second assistant message", options }
170 ])
171
172 assert.deepStrictEqual(storedHistory, expectedHistory)
173 }).pipe(withConstantIdGenerator("msg_abc123"), Effect.provide(PersistenceLayer)))
174
175 it.effect("should raise an error when retrieving a chat that does not exist", () =>
176 Effect.gen(function*() {
177 const persistence = yield* Chat.Persistence
178
179 const result = yield* Effect.flip(persistence.get("chat-321"))
180
181 assert.instanceOf(result, Chat.ChatNotFoundError)
182 assert.strictEqual(result.chatId, "chat-321")
183 }).pipe(Effect.provide(PersistenceLayer)))
184})