68 lines
2.5 KiB
TypeScript
68 lines
2.5 KiB
TypeScript
import { describe, it, expect } from "vitest"
|
|
import { ChatService, type ChatEvent, type MessageResponse } from "./chat"
|
|
|
|
describe("ChatService SSE Parsing", () => {
|
|
it("should correctly parse a text stream chunk", () => {
|
|
const rawChunk = `data: {"type": "on_chat_model_stream", "name": "summarizer", "data": {"chunk": "Hello"}}\n\n`
|
|
const events = ChatService.parseSSEChunk(rawChunk)
|
|
|
|
expect(events).toHaveLength(1)
|
|
expect(events[0]).toEqual({
|
|
type: "on_chat_model_stream",
|
|
name: "summarizer",
|
|
data: { chunk: "Hello" }
|
|
})
|
|
})
|
|
|
|
it("should handle multiple events in one chunk", () => {
|
|
const rawChunk =
|
|
`data: {"type": "on_chat_model_stream", "name": "summarizer", "data": {"chunk": "Hello"}}\n\n` +
|
|
`data: {"type": "on_chat_model_stream", "name": "summarizer", "data": {"chunk": " World"}}\n\n`
|
|
|
|
const events = ChatService.parseSSEChunk(rawChunk)
|
|
|
|
expect(events).toHaveLength(2)
|
|
expect(events[1].data.chunk).toBe(" World")
|
|
})
|
|
|
|
it("should parse encoded plots from executor node", () => {
|
|
const rawChunk = `data: {"type": "on_chain_end", "name": "executor", "data": {"encoded_plots": ["base64data"]}}\n\n`
|
|
const events = ChatService.parseSSEChunk(rawChunk)
|
|
|
|
expect(events[0].data.encoded_plots).toEqual(["base64data"])
|
|
})
|
|
|
|
it("should identify the done event", () => {
|
|
const rawChunk = `data: {"type": "done"}\n\n`
|
|
const events = ChatService.parseSSEChunk(rawChunk)
|
|
|
|
expect(events[0].type).toBe("done")
|
|
})
|
|
})
|
|
|
|
describe("ChatService Message State Management", () => {
|
|
it("should append text chunks to the last message content", () => {
|
|
const messages: MessageResponse[] = [{ id: "1", role: "assistant", content: "Initial", created_at: new Date().toISOString() }]
|
|
const event: ChatEvent = {
|
|
type: "on_chat_model_stream",
|
|
node: "summarizer",
|
|
data: { chunk: { content: " text" } }
|
|
}
|
|
|
|
const updatedMessages = ChatService.updateMessagesWithEvent(messages, event)
|
|
expect(updatedMessages[0].content).toBe("Initial text")
|
|
})
|
|
|
|
it("should add plots to the message state", () => {
|
|
const messages: MessageResponse[] = [{ id: "1", role: "assistant", content: "Analysis", created_at: new Date().toISOString(), plots: [] }]
|
|
const event: ChatEvent = {
|
|
type: "on_chain_end",
|
|
name: "executor",
|
|
data: { encoded_plots: ["plot1"] }
|
|
}
|
|
|
|
const updatedMessages = ChatService.updateMessagesWithEvent(messages, event)
|
|
expect(updatedMessages[0].plots).toEqual(["plot1"])
|
|
})
|
|
})
|