78 lines
2.5 KiB
TypeScript
78 lines
2.5 KiB
TypeScript
import { generateText, streamText } from "ai";
|
|
import { testConfig } from "../config/test-config";
|
|
|
|
interface Provider {
|
|
name: string;
|
|
activeModel: string;
|
|
apiKey: string | undefined;
|
|
}
|
|
|
|
describe.each(testConfig.providers)('TEXT/STREAM PROPERTIES: Tests with model %s', (provider: Provider) => {
|
|
const { userId } = testConfig;
|
|
let mem0: ReturnType<typeof testConfig.createTestClient>;
|
|
jest.setTimeout(50000);
|
|
|
|
beforeEach(() => {
|
|
mem0 = testConfig.createTestClient(provider);
|
|
});
|
|
|
|
it("should stream text with onChunk handler", async () => {
|
|
const chunkTexts: string[] = [];
|
|
const { textStream } = await streamText({
|
|
model: mem0(provider.activeModel, {
|
|
user_id: userId, // Use the uniform userId
|
|
}),
|
|
prompt: "Write only the name of the car I prefer and its color.",
|
|
onChunk({ chunk }) {
|
|
if (chunk.type === "text-delta") {
|
|
// Store chunk text for assertions
|
|
chunkTexts.push(chunk.textDelta);
|
|
}
|
|
},
|
|
});
|
|
|
|
// Wait for the stream to complete
|
|
for await (const _ of textStream) {
|
|
}
|
|
|
|
// Ensure chunks are collected
|
|
expect(chunkTexts.length).toBeGreaterThan(0);
|
|
expect(chunkTexts.every((text) => typeof text === "string")).toBe(true);
|
|
});
|
|
|
|
it("should call onFinish handler without throwing an error", async () => {
|
|
await streamText({
|
|
model: mem0(provider.activeModel, {
|
|
user_id: userId, // Use the uniform userId
|
|
}),
|
|
prompt: "Write only the name of the car I prefer and its color.",
|
|
onFinish({ text, finishReason, usage }) {
|
|
|
|
},
|
|
});
|
|
});
|
|
|
|
it("should generate fullStream with expected usage", async () => {
|
|
const {
|
|
text, // combined text
|
|
usage, // combined usage of all steps
|
|
} = await generateText({
|
|
model: mem0.completion(provider.activeModel), // Ensure the model name is correct
|
|
maxSteps: 5, // Enable multi-step calls
|
|
experimental_continueSteps: true,
|
|
prompt:
|
|
"Suggest me some good cars to buy. Each response MUST HAVE at least 200 words.",
|
|
});
|
|
|
|
// Ensure text is a string
|
|
expect(typeof text).toBe("string");
|
|
|
|
// Check usage
|
|
// promptTokens is a number, so we use toBeCloseTo instead of toBe and it should be in the range 155 to 165
|
|
expect(usage.promptTokens).toBeGreaterThanOrEqual(100);
|
|
expect(usage.promptTokens).toBeLessThanOrEqual(500);
|
|
expect(usage.completionTokens).toBeGreaterThanOrEqual(250); // Check completion tokens are above 250
|
|
expect(usage.totalTokens).toBeGreaterThan(400); // Check total tokens are above 400
|
|
});
|
|
});
|