diff options
Diffstat (limited to 'packages/core/src')
| -rw-r--r-- | packages/core/src/core/turn.test.ts | 7 | ||||
| -rw-r--r-- | packages/core/src/core/turn.ts | 6 |
2 files changed, 10 insertions, 3 deletions
diff --git a/packages/core/src/core/turn.test.ts b/packages/core/src/core/turn.test.ts index 2217e5da..aeb30229 100644 --- a/packages/core/src/core/turn.test.ts +++ b/packages/core/src/core/turn.test.ts @@ -239,6 +239,8 @@ describe('Turn', () => { candidates: [{ content: { parts: [{ text: 'First response' }] } }], usageMetadata: mockMetadata1, } as unknown as GenerateContentResponse; + // Add a small delay to ensure apiTimeMs is > 0 + await new Promise((resolve) => setTimeout(resolve, 10)); yield { functionCalls: [{ name: 'aTool' }], usageMetadata: mockMetadata2, @@ -262,7 +264,10 @@ describe('Turn', () => { expect(metadataEvent.type).toBe(GeminiEventType.UsageMetadata); // The value should be the *last* metadata object received. - expect(metadataEvent.value).toEqual(mockMetadata2); + expect(metadataEvent.value).toEqual( + expect.objectContaining(mockMetadata2), + ); + expect(metadataEvent.value.apiTimeMs).toBeGreaterThan(0); // Also check the public getter expect(turn.getUsageMetadata()).toEqual(mockMetadata2); diff --git a/packages/core/src/core/turn.ts b/packages/core/src/core/turn.ts index 34e4a494..71c02d83 100644 --- a/packages/core/src/core/turn.ts +++ b/packages/core/src/core/turn.ts @@ -104,7 +104,7 @@ export type ServerGeminiChatCompressedEvent = { export type ServerGeminiUsageMetadataEvent = { type: GeminiEventType.UsageMetadata; - value: GenerateContentResponseUsageMetadata; + value: GenerateContentResponseUsageMetadata & { apiTimeMs?: number }; }; // The original union type, now composed of the individual types @@ -137,6 +137,7 @@ export class Turn { req: PartListUnion, signal: AbortSignal, ): AsyncGenerator<ServerGeminiStreamEvent> { + const startTime = Date.now(); try { const responseStream = await this.chat.sendMessageStream({ message: req, @@ -174,9 +175,10 @@ export class Turn { } if (this.lastUsageMetadata) { + const durationMs = Date.now() - startTime; yield { type: GeminiEventType.UsageMetadata, - value: this.lastUsageMetadata, + value: { ...this.lastUsageMetadata, apiTimeMs: durationMs }, }; } } catch (error) { |
