Skip to content

Commit

Permalink
steps test
Browse files Browse the repository at this point in the history
  • Loading branch information
@lgrammel
lgrammel committed Feb 7, 2025
1 parent e2048ac commit e1fcd02
Show file tree
Hide file tree
Showing 4 changed files with 132 additions and 49 deletions.
2 changes: 1 addition & 1 deletion packages/ai/core/generate-text/__snapshots__/generate-text.test.ts.snap
View file
Original file line number Diff line number Diff line change
Expand Up @@ -859,7 +859,7 @@ exports[`result.steps > should add the reasoning from the model response to the
]
`;

exports[`result.toolCalls > result.steps should contain sources 1`] = `
exports[`result.steps > should contain sources 1`] = `
[
{
"experimental_providerMetadata": undefined,
Expand Down
65 changes: 65 additions & 0 deletions packages/ai/core/generate-text/__snapshots__/stream-text.test.ts.snap
View file
Original file line number Diff line number Diff line change
Expand Up @@ -3389,6 +3389,71 @@ exports[`streamText > result.steps > should add the reasoning from the model res
]
`;

exports[`streamText > result.steps > should contain sources 1`] = `
[
{
"experimental_providerMetadata": undefined,
"finishReason": "stop",
"isContinued": false,
"logprobs": undefined,
"reasoning": undefined,
"request": {},
"response": {
"headers": undefined,
"id": "id-1",
"messages": [
{
"content": [
{
"text": "Hello!",
"type": "text",
},
],
"id": "msg-0",
"role": "assistant",
},
],
"modelId": "mock-model-id",
"timestamp": 1970-01-01T00:00:00.000Z,
},
"sources": [
{
"id": "123",
"providerMetadata": {
"provider": {
"custom": "value",
},
},
"sourceType": "url",
"title": "Example",
"url": "https://example.com",
},
{
"id": "456",
"providerMetadata": {
"provider": {
"custom": "value2",
},
},
"sourceType": "url",
"title": "Example 2",
"url": "https://example.com/2",
},
],
"stepType": "initial",
"text": "Hello!",
"toolCalls": [],
"toolResults": [],
"usage": {
"completionTokens": 10,
"promptTokens": 3,
"totalTokens": 13,
},
"warnings": undefined,
},
]
`;

exports[`streamText > result.toDataStream > should create a data stream 1`] = `
[
"f:{"messageId":"msg-0"}
Expand Down
28 changes: 14 additions & 14 deletions packages/ai/core/generate-text/generate-text.test.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -124,6 +124,20 @@ describe('result.steps', () => {

expect(result.steps).toMatchSnapshot();
});

it('should contain sources', async () => {
const result = await generateText({
model: modelWithSources,
prompt: 'prompt',
experimental_generateMessageId: mockId({ prefix: 'msg' }),
_internal: {
generateId: mockId({ prefix: 'id' }),
currentDate: () => new Date(0),
},
});

expect(result.steps).toMatchSnapshot();
});
});

describe('result.toolCalls', () => {
Expand Down Expand Up @@ -210,20 +224,6 @@ describe('result.toolCalls', () => {
},
]);
});

it('result.steps should contain sources', async () => {
const result = await generateText({
model: modelWithSources,
prompt: 'prompt',
experimental_generateMessageId: mockId({ prefix: 'msg' }),
_internal: {
generateId: mockId({ prefix: 'id' }),
currentDate: () => new Date(0),
},
});

expect(result.steps).toMatchSnapshot();
});
});

describe('result.toolResults', () => {
Expand Down
86 changes: 52 additions & 34 deletions packages/ai/core/generate-text/stream-text.test.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -69,6 +69,41 @@ function createTestModel({
});
}

const modelWithSources = new MockLanguageModelV1({
doStream: async () => ({
stream: convertArrayToReadableStream([
{
type: 'source',
source: {
sourceType: 'url' as const,
id: '123',
url: 'https://example.com',
title: 'Example',
providerMetadata: { provider: { custom: 'value' } },
},
},
{ type: 'text-delta', textDelta: 'Hello!' },
{
type: 'source',
source: {
sourceType: 'url' as const,
id: '456',
url: 'https://example.com/2',
title: 'Example 2',
providerMetadata: { provider: { custom: 'value2' } },
},
},
{
type: 'finish',
finishReason: 'stop',
logprobs: undefined,
usage: { completionTokens: 10, promptTokens: 3 },
},
]),
rawCall: { rawPrompt: 'prompt', rawSettings: {} },
}),
});

describe('streamText', () => {
describe('result.textStream', () => {
it('should send text deltas', async () => {
Expand Down Expand Up @@ -287,40 +322,7 @@ describe('streamText', () => {

it('should include sources in fullStream', async () => {
const result = streamText({
model: new MockLanguageModelV1({
doStream: async () => ({
stream: convertArrayToReadableStream([
{
type: 'source',
source: {
sourceType: 'url' as const,
id: '123',
url: 'https://example.com',
title: 'Example',
providerMetadata: { provider: { custom: 'value' } },
},
},
{ type: 'text-delta', textDelta: 'Hello!' },
{
type: 'source',
source: {
sourceType: 'url' as const,
id: '456',
url: 'https://example.com/2',
title: 'Example 2',
providerMetadata: { provider: { custom: 'value2' } },
},
},
{
type: 'finish',
finishReason: 'stop',
logprobs: undefined,
usage: { completionTokens: 10, promptTokens: 3 },
},
]),
rawCall: { rawPrompt: 'prompt', rawSettings: {} },
}),
}),
model: modelWithSources,
prompt: 'test-input',
_internal: {
currentDate: mockValues(new Date(2000)),
Expand Down Expand Up @@ -1707,6 +1709,22 @@ describe('streamText', () => {

expect(await result.steps).toMatchSnapshot();
});

it('should contain sources', async () => {
const result = streamText({
model: modelWithSources,
prompt: 'prompt',
experimental_generateMessageId: mockId({ prefix: 'msg' }),
_internal: {
generateId: mockId({ prefix: 'id' }),
currentDate: () => new Date(0),
},
});

result.consumeStream();

expect(await result.steps).toMatchSnapshot();
});
});

describe('result.toolCalls', () => {
Expand Down

0 comments on commit e1fcd02

Please sign in to comment.