|
|
@@ -1,7 +1,30 @@
|
|
|
// npx vitest run src/api/providers/__tests__/zai.spec.ts
|
|
|
|
|
|
-import OpenAI from "openai"
|
|
|
-import { Anthropic } from "@anthropic-ai/sdk"
|
|
|
+// Use vi.hoisted to define mock functions that can be referenced in hoisted vi.mock() calls
|
|
|
+const { mockStreamText, mockGenerateText } = vi.hoisted(() => ({
|
|
|
+ mockStreamText: vi.fn(),
|
|
|
+ mockGenerateText: vi.fn(),
|
|
|
+}))
|
|
|
+
|
|
|
+vi.mock("ai", async (importOriginal) => {
|
|
|
+ const actual = await importOriginal<typeof import("ai")>()
|
|
|
+ return {
|
|
|
+ ...actual,
|
|
|
+ streamText: mockStreamText,
|
|
|
+ generateText: mockGenerateText,
|
|
|
+ }
|
|
|
+})
|
|
|
+
|
|
|
+vi.mock("zhipu-ai-provider", () => ({
|
|
|
+ createZhipu: vi.fn(() => {
|
|
|
+ return vi.fn(() => ({
|
|
|
+ modelId: "glm-4.6",
|
|
|
+ provider: "zhipu",
|
|
|
+ }))
|
|
|
+ }),
|
|
|
+}))
|
|
|
+
|
|
|
+import type { Anthropic } from "@anthropic-ai/sdk"
|
|
|
|
|
|
import {
|
|
|
type InternationalZAiModelId,
|
|
|
@@ -13,42 +36,41 @@ import {
|
|
|
ZAI_DEFAULT_TEMPERATURE,
|
|
|
} from "@roo-code/types"
|
|
|
|
|
|
-import { ZAiHandler } from "../zai"
|
|
|
+import type { ApiHandlerOptions } from "../../../shared/api"
|
|
|
|
|
|
-vitest.mock("openai", () => {
|
|
|
- const createMock = vitest.fn()
|
|
|
- return {
|
|
|
- default: vitest.fn(() => ({ chat: { completions: { create: createMock } } })),
|
|
|
- }
|
|
|
-})
|
|
|
+import { ZAiHandler } from "../zai"
|
|
|
|
|
|
describe("ZAiHandler", () => {
|
|
|
let handler: ZAiHandler
|
|
|
- let mockCreate: any
|
|
|
+ let mockOptions: ApiHandlerOptions
|
|
|
|
|
|
beforeEach(() => {
|
|
|
- vitest.clearAllMocks()
|
|
|
- mockCreate = (OpenAI as unknown as any)().chat.completions.create
|
|
|
+ mockOptions = {
|
|
|
+ zaiApiKey: "test-zai-api-key",
|
|
|
+ zaiApiLine: "international_coding",
|
|
|
+ apiModelId: "glm-4.6",
|
|
|
+ }
|
|
|
+ handler = new ZAiHandler(mockOptions)
|
|
|
+ vi.clearAllMocks()
|
|
|
})
|
|
|
|
|
|
- describe("International Z AI", () => {
|
|
|
- beforeEach(() => {
|
|
|
- handler = new ZAiHandler({ zaiApiKey: "test-zai-api-key", zaiApiLine: "international_coding" })
|
|
|
+ describe("constructor", () => {
|
|
|
+ it("should initialize with provided options", () => {
|
|
|
+ expect(handler).toBeInstanceOf(ZAiHandler)
|
|
|
+ expect(handler.getModel().id).toBe(mockOptions.apiModelId)
|
|
|
})
|
|
|
|
|
|
- it("should use the correct international Z AI base URL", () => {
|
|
|
- new ZAiHandler({ zaiApiKey: "test-zai-api-key", zaiApiLine: "international_coding" })
|
|
|
- expect(OpenAI).toHaveBeenCalledWith(
|
|
|
- expect.objectContaining({
|
|
|
- baseURL: "https://api.z.ai/api/coding/paas/v4",
|
|
|
- }),
|
|
|
- )
|
|
|
+ it("should default to international when no zaiApiLine is specified", () => {
|
|
|
+ const handlerDefault = new ZAiHandler({ zaiApiKey: "test-zai-api-key" })
|
|
|
+ const model = handlerDefault.getModel()
|
|
|
+ expect(model.id).toBe(internationalZAiDefaultModelId)
|
|
|
+ expect(model.info).toEqual(internationalZAiModels[internationalZAiDefaultModelId])
|
|
|
})
|
|
|
+ })
|
|
|
|
|
|
- it("should use the provided API key for international", () => {
|
|
|
- const zaiApiKey = "test-zai-api-key"
|
|
|
- new ZAiHandler({ zaiApiKey, zaiApiLine: "international_coding" })
|
|
|
- expect(OpenAI).toHaveBeenCalledWith(expect.objectContaining({ apiKey: zaiApiKey }))
|
|
|
+ describe("International Z AI", () => {
|
|
|
+ beforeEach(() => {
|
|
|
+ handler = new ZAiHandler({ zaiApiKey: "test-zai-api-key", zaiApiLine: "international_coding" })
|
|
|
})
|
|
|
|
|
|
it("should return international default model when no model is specified", () => {
|
|
|
@@ -119,19 +141,6 @@ describe("ZAiHandler", () => {
|
|
|
handler = new ZAiHandler({ zaiApiKey: "test-zai-api-key", zaiApiLine: "china_coding" })
|
|
|
})
|
|
|
|
|
|
- it("should use the correct China Z AI base URL", () => {
|
|
|
- new ZAiHandler({ zaiApiKey: "test-zai-api-key", zaiApiLine: "china_coding" })
|
|
|
- expect(OpenAI).toHaveBeenCalledWith(
|
|
|
- expect.objectContaining({ baseURL: "https://open.bigmodel.cn/api/coding/paas/v4" }),
|
|
|
- )
|
|
|
- })
|
|
|
-
|
|
|
- it("should use the provided API key for China", () => {
|
|
|
- const zaiApiKey = "test-zai-api-key"
|
|
|
- new ZAiHandler({ zaiApiKey, zaiApiLine: "china_coding" })
|
|
|
- expect(OpenAI).toHaveBeenCalledWith(expect.objectContaining({ apiKey: zaiApiKey }))
|
|
|
- })
|
|
|
-
|
|
|
it("should return China default model when no model is specified", () => {
|
|
|
const model = handler.getModel()
|
|
|
expect(model.id).toBe(mainlandZAiDefaultModelId)
|
|
|
@@ -200,21 +209,6 @@ describe("ZAiHandler", () => {
|
|
|
handler = new ZAiHandler({ zaiApiKey: "test-zai-api-key", zaiApiLine: "international_api" })
|
|
|
})
|
|
|
|
|
|
- it("should use the correct international API base URL", () => {
|
|
|
- new ZAiHandler({ zaiApiKey: "test-zai-api-key", zaiApiLine: "international_api" })
|
|
|
- expect(OpenAI).toHaveBeenCalledWith(
|
|
|
- expect.objectContaining({
|
|
|
- baseURL: "https://api.z.ai/api/paas/v4",
|
|
|
- }),
|
|
|
- )
|
|
|
- })
|
|
|
-
|
|
|
- it("should use the provided API key for international API", () => {
|
|
|
- const zaiApiKey = "test-zai-api-key"
|
|
|
- new ZAiHandler({ zaiApiKey, zaiApiLine: "international_api" })
|
|
|
- expect(OpenAI).toHaveBeenCalledWith(expect.objectContaining({ apiKey: zaiApiKey }))
|
|
|
- })
|
|
|
-
|
|
|
it("should return international default model when no model is specified", () => {
|
|
|
const model = handler.getModel()
|
|
|
expect(model.id).toBe(internationalZAiDefaultModelId)
|
|
|
@@ -239,21 +233,6 @@ describe("ZAiHandler", () => {
|
|
|
handler = new ZAiHandler({ zaiApiKey: "test-zai-api-key", zaiApiLine: "china_api" })
|
|
|
})
|
|
|
|
|
|
- it("should use the correct China API base URL", () => {
|
|
|
- new ZAiHandler({ zaiApiKey: "test-zai-api-key", zaiApiLine: "china_api" })
|
|
|
- expect(OpenAI).toHaveBeenCalledWith(
|
|
|
- expect.objectContaining({
|
|
|
- baseURL: "https://open.bigmodel.cn/api/paas/v4",
|
|
|
- }),
|
|
|
- )
|
|
|
- })
|
|
|
-
|
|
|
- it("should use the provided API key for China API", () => {
|
|
|
- const zaiApiKey = "test-zai-api-key"
|
|
|
- new ZAiHandler({ zaiApiKey, zaiApiLine: "china_api" })
|
|
|
- expect(OpenAI).toHaveBeenCalledWith(expect.objectContaining({ apiKey: zaiApiKey }))
|
|
|
- })
|
|
|
-
|
|
|
it("should return China default model when no model is specified", () => {
|
|
|
const model = handler.getModel()
|
|
|
expect(model.id).toBe(mainlandZAiDefaultModelId)
|
|
|
@@ -273,133 +252,98 @@ describe("ZAiHandler", () => {
|
|
|
})
|
|
|
})
|
|
|
|
|
|
- describe("Default behavior", () => {
|
|
|
- it("should default to international when no zaiApiLine is specified", () => {
|
|
|
- const handlerDefault = new ZAiHandler({ zaiApiKey: "test-zai-api-key" })
|
|
|
- expect(OpenAI).toHaveBeenCalledWith(
|
|
|
- expect.objectContaining({
|
|
|
- baseURL: "https://api.z.ai/api/coding/paas/v4",
|
|
|
- }),
|
|
|
- )
|
|
|
-
|
|
|
- const model = handlerDefault.getModel()
|
|
|
- expect(model.id).toBe(internationalZAiDefaultModelId)
|
|
|
- expect(model.info).toEqual(internationalZAiModels[internationalZAiDefaultModelId])
|
|
|
- })
|
|
|
-
|
|
|
- it("should use 'not-provided' as default API key when none is specified", () => {
|
|
|
- new ZAiHandler({ zaiApiLine: "international_coding" })
|
|
|
- expect(OpenAI).toHaveBeenCalledWith(expect.objectContaining({ apiKey: "not-provided" }))
|
|
|
+ describe("getModel", () => {
|
|
|
+ it("should include model parameters from getModelParams", () => {
|
|
|
+ const model = handler.getModel()
|
|
|
+ expect(model).toHaveProperty("temperature")
|
|
|
+ expect(model).toHaveProperty("maxTokens")
|
|
|
})
|
|
|
})
|
|
|
|
|
|
- describe("API Methods", () => {
|
|
|
- beforeEach(() => {
|
|
|
- handler = new ZAiHandler({ zaiApiKey: "test-zai-api-key", zaiApiLine: "international_coding" })
|
|
|
- })
|
|
|
-
|
|
|
- it("completePrompt method should return text from Z AI API", async () => {
|
|
|
- const expectedResponse = "This is a test response from Z AI"
|
|
|
- mockCreate.mockResolvedValueOnce({ choices: [{ message: { content: expectedResponse } }] })
|
|
|
- const result = await handler.completePrompt("test prompt")
|
|
|
- expect(result).toBe(expectedResponse)
|
|
|
- })
|
|
|
-
|
|
|
- it("should handle errors in completePrompt", async () => {
|
|
|
- const errorMessage = "Z AI API error"
|
|
|
- mockCreate.mockRejectedValueOnce(new Error(errorMessage))
|
|
|
- await expect(handler.completePrompt("test prompt")).rejects.toThrow(
|
|
|
- `Z.ai completion error: ${errorMessage}`,
|
|
|
- )
|
|
|
- })
|
|
|
+ describe("createMessage", () => {
|
|
|
+ const systemPrompt = "You are a helpful assistant."
|
|
|
+ const messages: Anthropic.Messages.MessageParam[] = [
|
|
|
+ {
|
|
|
+ role: "user",
|
|
|
+ content: [{ type: "text" as const, text: "Hello!" }],
|
|
|
+ },
|
|
|
+ ]
|
|
|
+
|
|
|
+ it("should handle streaming responses", async () => {
|
|
|
+ async function* mockFullStream() {
|
|
|
+ yield { type: "text-delta", text: "Test response from Z.ai" }
|
|
|
+ }
|
|
|
+
|
|
|
+ const mockUsage = Promise.resolve({
|
|
|
+ inputTokens: 10,
|
|
|
+ outputTokens: 5,
|
|
|
+ })
|
|
|
|
|
|
- it("createMessage should yield text content from stream", async () => {
|
|
|
- const testContent = "This is test content from Z AI stream"
|
|
|
-
|
|
|
- mockCreate.mockImplementationOnce(() => {
|
|
|
- return {
|
|
|
- [Symbol.asyncIterator]: () => ({
|
|
|
- next: vitest
|
|
|
- .fn()
|
|
|
- .mockResolvedValueOnce({
|
|
|
- done: false,
|
|
|
- value: { choices: [{ delta: { content: testContent } }] },
|
|
|
- })
|
|
|
- .mockResolvedValueOnce({ done: true }),
|
|
|
- }),
|
|
|
- }
|
|
|
+ mockStreamText.mockReturnValue({
|
|
|
+ fullStream: mockFullStream(),
|
|
|
+ usage: mockUsage,
|
|
|
})
|
|
|
|
|
|
- const stream = handler.createMessage("system prompt", [])
|
|
|
- const firstChunk = await stream.next()
|
|
|
+ const stream = handler.createMessage(systemPrompt, messages)
|
|
|
+ const chunks: any[] = []
|
|
|
+ for await (const chunk of stream) {
|
|
|
+ chunks.push(chunk)
|
|
|
+ }
|
|
|
|
|
|
- expect(firstChunk.done).toBe(false)
|
|
|
- expect(firstChunk.value).toEqual({ type: "text", text: testContent })
|
|
|
+ expect(chunks.length).toBeGreaterThan(0)
|
|
|
+ const textChunks = chunks.filter((chunk) => chunk.type === "text")
|
|
|
+ expect(textChunks).toHaveLength(1)
|
|
|
+ expect(textChunks[0].text).toBe("Test response from Z.ai")
|
|
|
})
|
|
|
|
|
|
- it("createMessage should yield usage data from stream", async () => {
|
|
|
- mockCreate.mockImplementationOnce(() => {
|
|
|
- return {
|
|
|
- [Symbol.asyncIterator]: () => ({
|
|
|
- next: vitest
|
|
|
- .fn()
|
|
|
- .mockResolvedValueOnce({
|
|
|
- done: false,
|
|
|
- value: {
|
|
|
- choices: [{ delta: {} }],
|
|
|
- usage: { prompt_tokens: 10, completion_tokens: 20 },
|
|
|
- },
|
|
|
- })
|
|
|
- .mockResolvedValueOnce({ done: true }),
|
|
|
- }),
|
|
|
- }
|
|
|
+ it("should include usage information", async () => {
|
|
|
+ async function* mockFullStream() {
|
|
|
+ yield { type: "text-delta", text: "Test response" }
|
|
|
+ }
|
|
|
+
|
|
|
+ const mockUsage = Promise.resolve({
|
|
|
+ inputTokens: 10,
|
|
|
+ outputTokens: 20,
|
|
|
})
|
|
|
|
|
|
- const stream = handler.createMessage("system prompt", [])
|
|
|
- const firstChunk = await stream.next()
|
|
|
+ mockStreamText.mockReturnValue({
|
|
|
+ fullStream: mockFullStream(),
|
|
|
+ usage: mockUsage,
|
|
|
+ })
|
|
|
|
|
|
- expect(firstChunk.done).toBe(false)
|
|
|
- expect(firstChunk.value).toMatchObject({ type: "usage", inputTokens: 10, outputTokens: 20 })
|
|
|
+ const stream = handler.createMessage(systemPrompt, messages)
|
|
|
+ const chunks: any[] = []
|
|
|
+ for await (const chunk of stream) {
|
|
|
+ chunks.push(chunk)
|
|
|
+ }
|
|
|
+
|
|
|
+ const usageChunks = chunks.filter((chunk) => chunk.type === "usage")
|
|
|
+ expect(usageChunks.length).toBeGreaterThan(0)
|
|
|
+ expect(usageChunks[0].inputTokens).toBe(10)
|
|
|
+ expect(usageChunks[0].outputTokens).toBe(20)
|
|
|
})
|
|
|
|
|
|
- it("createMessage should pass correct parameters to Z AI client", async () => {
|
|
|
- const modelId: InternationalZAiModelId = "glm-4.5"
|
|
|
- const modelInfo = internationalZAiModels[modelId]
|
|
|
- const handlerWithModel = new ZAiHandler({
|
|
|
- apiModelId: modelId,
|
|
|
- zaiApiKey: "test-zai-api-key",
|
|
|
- zaiApiLine: "international_coding",
|
|
|
- })
|
|
|
+ it("should pass correct parameters to streamText", async () => {
|
|
|
+ async function* mockFullStream() {
|
|
|
+ // empty stream
|
|
|
+ }
|
|
|
|
|
|
- mockCreate.mockImplementationOnce(() => {
|
|
|
- return {
|
|
|
- [Symbol.asyncIterator]: () => ({
|
|
|
- async next() {
|
|
|
- return { done: true }
|
|
|
- },
|
|
|
- }),
|
|
|
- }
|
|
|
+ mockStreamText.mockReturnValue({
|
|
|
+ fullStream: mockFullStream(),
|
|
|
+ usage: Promise.resolve({ inputTokens: 0, outputTokens: 0 }),
|
|
|
})
|
|
|
|
|
|
- const systemPrompt = "Test system prompt for Z AI"
|
|
|
- const messages: Anthropic.Messages.MessageParam[] = [{ role: "user", content: "Test message for Z AI" }]
|
|
|
-
|
|
|
- const messageGenerator = handlerWithModel.createMessage(systemPrompt, messages)
|
|
|
- await messageGenerator.next()
|
|
|
+ const stream = handler.createMessage(systemPrompt, messages)
|
|
|
+ // Consume the stream
|
|
|
+ for await (const _chunk of stream) {
|
|
|
+ // drain
|
|
|
+ }
|
|
|
|
|
|
- // Centralized 20% cap should apply to OpenAI-compatible providers like Z AI
|
|
|
- const expectedMaxTokens = Math.min(modelInfo.maxTokens, Math.ceil(modelInfo.contextWindow * 0.2))
|
|
|
-
|
|
|
- expect(mockCreate).toHaveBeenCalledWith(
|
|
|
+ expect(mockStreamText).toHaveBeenCalledWith(
|
|
|
expect.objectContaining({
|
|
|
- model: modelId,
|
|
|
- max_tokens: expectedMaxTokens,
|
|
|
- temperature: ZAI_DEFAULT_TEMPERATURE,
|
|
|
- messages: expect.arrayContaining([{ role: "system", content: systemPrompt }]),
|
|
|
- stream: true,
|
|
|
- stream_options: { include_usage: true },
|
|
|
+ system: systemPrompt,
|
|
|
+ temperature: expect.any(Number),
|
|
|
}),
|
|
|
- undefined,
|
|
|
)
|
|
|
})
|
|
|
})
|
|
|
@@ -410,27 +354,29 @@ describe("ZAiHandler", () => {
|
|
|
apiModelId: "glm-4.7",
|
|
|
zaiApiKey: "test-zai-api-key",
|
|
|
zaiApiLine: "international_coding",
|
|
|
- // No reasoningEffort setting - should use model default (medium)
|
|
|
})
|
|
|
|
|
|
- mockCreate.mockImplementationOnce(() => {
|
|
|
- return {
|
|
|
- [Symbol.asyncIterator]: () => ({
|
|
|
- async next() {
|
|
|
- return { done: true }
|
|
|
- },
|
|
|
- }),
|
|
|
- }
|
|
|
+ async function* mockFullStream() {
|
|
|
+ yield { type: "text-delta", text: "response" }
|
|
|
+ }
|
|
|
+
|
|
|
+ mockStreamText.mockReturnValue({
|
|
|
+ fullStream: mockFullStream(),
|
|
|
+ usage: Promise.resolve({ inputTokens: 0, outputTokens: 0 }),
|
|
|
})
|
|
|
|
|
|
- const messageGenerator = handlerWithModel.createMessage("system prompt", [])
|
|
|
- await messageGenerator.next()
|
|
|
+ const stream = handlerWithModel.createMessage("system prompt", [])
|
|
|
+ for await (const _chunk of stream) {
|
|
|
+ // drain
|
|
|
+ }
|
|
|
|
|
|
- // For GLM-4.7 with default reasoning (medium), thinking should be enabled
|
|
|
- expect(mockCreate).toHaveBeenCalledWith(
|
|
|
+ expect(mockStreamText).toHaveBeenCalledWith(
|
|
|
expect.objectContaining({
|
|
|
- model: "glm-4.7",
|
|
|
- thinking: { type: "enabled" },
|
|
|
+ providerOptions: {
|
|
|
+ zhipu: {
|
|
|
+ thinking: { type: "enabled" },
|
|
|
+ },
|
|
|
+ },
|
|
|
}),
|
|
|
)
|
|
|
})
|
|
|
@@ -444,24 +390,27 @@ describe("ZAiHandler", () => {
|
|
|
reasoningEffort: "disable",
|
|
|
})
|
|
|
|
|
|
- mockCreate.mockImplementationOnce(() => {
|
|
|
- return {
|
|
|
- [Symbol.asyncIterator]: () => ({
|
|
|
- async next() {
|
|
|
- return { done: true }
|
|
|
- },
|
|
|
- }),
|
|
|
- }
|
|
|
+ async function* mockFullStream() {
|
|
|
+ yield { type: "text-delta", text: "response" }
|
|
|
+ }
|
|
|
+
|
|
|
+ mockStreamText.mockReturnValue({
|
|
|
+ fullStream: mockFullStream(),
|
|
|
+ usage: Promise.resolve({ inputTokens: 0, outputTokens: 0 }),
|
|
|
})
|
|
|
|
|
|
- const messageGenerator = handlerWithModel.createMessage("system prompt", [])
|
|
|
- await messageGenerator.next()
|
|
|
+ const stream = handlerWithModel.createMessage("system prompt", [])
|
|
|
+ for await (const _chunk of stream) {
|
|
|
+ // drain
|
|
|
+ }
|
|
|
|
|
|
- // For GLM-4.7 with reasoning disabled, thinking should be disabled
|
|
|
- expect(mockCreate).toHaveBeenCalledWith(
|
|
|
+ expect(mockStreamText).toHaveBeenCalledWith(
|
|
|
expect.objectContaining({
|
|
|
- model: "glm-4.7",
|
|
|
- thinking: { type: "disabled" },
|
|
|
+ providerOptions: {
|
|
|
+ zhipu: {
|
|
|
+ thinking: { type: "disabled" },
|
|
|
+ },
|
|
|
+ },
|
|
|
}),
|
|
|
)
|
|
|
})
|
|
|
@@ -475,51 +424,109 @@ describe("ZAiHandler", () => {
|
|
|
reasoningEffort: "medium",
|
|
|
})
|
|
|
|
|
|
- mockCreate.mockImplementationOnce(() => {
|
|
|
- return {
|
|
|
- [Symbol.asyncIterator]: () => ({
|
|
|
- async next() {
|
|
|
- return { done: true }
|
|
|
- },
|
|
|
- }),
|
|
|
- }
|
|
|
+ async function* mockFullStream() {
|
|
|
+ yield { type: "text-delta", text: "response" }
|
|
|
+ }
|
|
|
+
|
|
|
+ mockStreamText.mockReturnValue({
|
|
|
+ fullStream: mockFullStream(),
|
|
|
+ usage: Promise.resolve({ inputTokens: 0, outputTokens: 0 }),
|
|
|
})
|
|
|
|
|
|
- const messageGenerator = handlerWithModel.createMessage("system prompt", [])
|
|
|
- await messageGenerator.next()
|
|
|
+ const stream = handlerWithModel.createMessage("system prompt", [])
|
|
|
+ for await (const _chunk of stream) {
|
|
|
+ // drain
|
|
|
+ }
|
|
|
|
|
|
- // For GLM-4.7 with reasoning set to medium, thinking should be enabled
|
|
|
- expect(mockCreate).toHaveBeenCalledWith(
|
|
|
+ expect(mockStreamText).toHaveBeenCalledWith(
|
|
|
expect.objectContaining({
|
|
|
- model: "glm-4.7",
|
|
|
- thinking: { type: "enabled" },
|
|
|
+ providerOptions: {
|
|
|
+ zhipu: {
|
|
|
+ thinking: { type: "enabled" },
|
|
|
+ },
|
|
|
+ },
|
|
|
}),
|
|
|
)
|
|
|
})
|
|
|
|
|
|
- it("should NOT add thinking parameter for non-thinking models like GLM-4.6", async () => {
|
|
|
+ it("should NOT add providerOptions for non-thinking models like GLM-4.6", async () => {
|
|
|
const handlerWithModel = new ZAiHandler({
|
|
|
apiModelId: "glm-4.6",
|
|
|
zaiApiKey: "test-zai-api-key",
|
|
|
zaiApiLine: "international_coding",
|
|
|
})
|
|
|
|
|
|
- mockCreate.mockImplementationOnce(() => {
|
|
|
- return {
|
|
|
- [Symbol.asyncIterator]: () => ({
|
|
|
- async next() {
|
|
|
- return { done: true }
|
|
|
- },
|
|
|
- }),
|
|
|
- }
|
|
|
+ async function* mockFullStream() {
|
|
|
+ yield { type: "text-delta", text: "response" }
|
|
|
+ }
|
|
|
+
|
|
|
+ mockStreamText.mockReturnValue({
|
|
|
+ fullStream: mockFullStream(),
|
|
|
+ usage: Promise.resolve({ inputTokens: 0, outputTokens: 0 }),
|
|
|
+ })
|
|
|
+
|
|
|
+ const stream = handlerWithModel.createMessage("system prompt", [])
|
|
|
+ for await (const _chunk of stream) {
|
|
|
+ // drain
|
|
|
+ }
|
|
|
+
|
|
|
+ const callArgs = mockStreamText.mock.calls[0][0]
|
|
|
+ expect(callArgs.providerOptions).toBeUndefined()
|
|
|
+ })
|
|
|
+
|
|
|
+ it("should handle reasoning content in streaming responses", async () => {
|
|
|
+ const handlerWithModel = new ZAiHandler({
|
|
|
+ apiModelId: "glm-4.7",
|
|
|
+ zaiApiKey: "test-zai-api-key",
|
|
|
+ zaiApiLine: "international_coding",
|
|
|
+ })
|
|
|
+
|
|
|
+ async function* mockFullStream() {
|
|
|
+ yield { type: "reasoning", text: "Let me think about this..." }
|
|
|
+ yield { type: "text-delta", text: "Here is my answer" }
|
|
|
+ }
|
|
|
+
|
|
|
+ mockStreamText.mockReturnValue({
|
|
|
+ fullStream: mockFullStream(),
|
|
|
+ usage: Promise.resolve({ inputTokens: 10, outputTokens: 20 }),
|
|
|
})
|
|
|
|
|
|
- const messageGenerator = handlerWithModel.createMessage("system prompt", [])
|
|
|
- await messageGenerator.next()
|
|
|
+ const stream = handlerWithModel.createMessage("system prompt", [])
|
|
|
+ const chunks: any[] = []
|
|
|
+ for await (const chunk of stream) {
|
|
|
+ chunks.push(chunk)
|
|
|
+ }
|
|
|
+
|
|
|
+ const reasoningChunks = chunks.filter((chunk) => chunk.type === "reasoning")
|
|
|
+ expect(reasoningChunks).toHaveLength(1)
|
|
|
+ expect(reasoningChunks[0].text).toBe("Let me think about this...")
|
|
|
+
|
|
|
+ const textChunks = chunks.filter((chunk) => chunk.type === "text")
|
|
|
+ expect(textChunks).toHaveLength(1)
|
|
|
+ expect(textChunks[0].text).toBe("Here is my answer")
|
|
|
+ })
|
|
|
+ })
|
|
|
+
|
|
|
+ describe("completePrompt", () => {
|
|
|
+ it("should complete a prompt using generateText", async () => {
|
|
|
+ mockGenerateText.mockResolvedValue({
|
|
|
+ text: "Test completion from Z.ai",
|
|
|
+ })
|
|
|
+
|
|
|
+ const result = await handler.completePrompt("Test prompt")
|
|
|
+
|
|
|
+ expect(result).toBe("Test completion from Z.ai")
|
|
|
+ expect(mockGenerateText).toHaveBeenCalledWith(
|
|
|
+ expect.objectContaining({
|
|
|
+ prompt: "Test prompt",
|
|
|
+ }),
|
|
|
+ )
|
|
|
+ })
|
|
|
+ })
|
|
|
|
|
|
- // For GLM-4.6 (no thinking support), thinking parameter should not be present
|
|
|
- const callArgs = mockCreate.mock.calls[0][0]
|
|
|
- expect(callArgs.thinking).toBeUndefined()
|
|
|
+ describe("isAiSdkProvider", () => {
|
|
|
+ it("should return true", () => {
|
|
|
+ expect(handler.isAiSdkProvider()).toBe(true)
|
|
|
})
|
|
|
})
|
|
|
})
|