fix: types for open ai driver (#2950)

* fix: types for open ai driver

* fix: add extra time to tests

* fix: test

* fix: together ai tests

* fix: tests
This commit is contained in:
Daniel Salazar
2026-05-07 14:48:47 -07:00
committed by GitHub
parent 25fe9a3cdf
commit 8d5e9ef9cf
18 changed files with 72 additions and 33 deletions
@@ -54,6 +54,7 @@ import {
normalize_single_message,
} from './utils/Messages.js';
import { AIChatStream } from './utils/Streaming.js';
import { EventMap } from '../../clients/event/types.js';
const MAX_FALLBACKS = 4; // includes first attempt
@@ -163,7 +164,8 @@ export class ChatCompletionDriver extends PuterDriver {
.replaceAll('-', '')
.slice(0, 25);
const validateEvent: Record<string, unknown> = {
const validateEvent: EventMap['ai.prompt.validate'] = {
username: actor.user?.username || '',
actor,
completionId,
allow: true,
@@ -308,10 +310,8 @@ export class ChatCompletionDriver extends PuterDriver {
// Credits can be exhausted mid-fallback by parallel requests;
// re-check before another upstream hit. Same bail as the
// pre-flight above.
const fallbackUsageAllowed = await metering.hasEnoughCredits(
actor,
1,
);
const fallbackUsageAllowed =
await this.services.metering.hasEnoughCredits(actor, 1);
if (!fallbackUsageAllowed) {
throw new HttpError(402, 'No usage left for request.', {
legacyCode: 'insufficient_funds',
@@ -369,13 +369,13 @@ export class ChatCompletionDriver extends PuterDriver {
this.clients.event.emit(
'ai.prompt.complete',
{
username,
username: username!,
completionId,
intended_service: intendedProvider,
parameters: args,
result: { usage: enrichedUsage, stream: true },
model_used: model.id,
service_used: model.provider,
service_used: model.provider!,
},
{},
);
@@ -425,13 +425,13 @@ export class ChatCompletionDriver extends PuterDriver {
this.clients.event.emit(
'ai.prompt.complete',
{
username,
username: username!,
completionId,
intended_service: intendedProvider,
parameters: args,
result: res,
model_used: model.id,
service_used: model.provider,
service_used: model.provider!,
},
{},
);
@@ -601,7 +601,7 @@ export class ChatCompletionDriver extends PuterDriver {
'ai.prompt.cost-calculated',
{
completionId,
username,
username: username!,
usage,
input_tokens: inputTokens,
output_tokens: outputTokens,
@@ -610,11 +610,11 @@ export class ChatCompletionDriver extends PuterDriver {
total_ucents: inputMicroCents + outputMicroCents,
costs_currency: model.costs_currency,
model_used: model.id,
service_used: model.provider,
service_used: model.provider!,
intended_service: intendedProvider,
model_details: {
id: model.id,
provider: model.provider,
provider: model.provider!,
input_cost_key: inputKey,
output_cost_key: outputKey,
costs: model.costs,
@@ -29,6 +29,7 @@
import { describe, expect, it } from 'vitest';
import {
INTEGRATION_TEST_TIMEOUT_MS,
makeMeteringStub,
optionalEnv,
skipUnlessEnv,
@@ -49,7 +50,7 @@ describe.skipIf(skipUnlessEnv(ENV_VAR))('ClaudeProvider (integration)', () => {
{ apiKey: optionalEnv(ENV_VAR)! },
);
it('returns a non-empty completion from claude-haiku-4-5', async () => {
it('returns a non-empty completion from claude-haiku-4-5', { timeout: INTEGRATION_TEST_TIMEOUT_MS }, async () => {
const provider = buildProvider();
const result = await withTestActor(() =>
provider.complete({
@@ -26,6 +26,7 @@
import { describe, expect, it } from 'vitest';
import {
INTEGRATION_TEST_TIMEOUT_MS,
makeMeteringStub,
optionalEnv,
skipUnlessEnv,
@@ -38,7 +39,7 @@ const ENV_VAR = 'PUTER_TEST_AI_DEEPSEEK_API_KEY';
describe.skipIf(skipUnlessEnv(ENV_VAR))(
'DeepSeekProvider (integration)',
() => {
it('returns a non-empty completion from deepseek-chat', async () => {
it('returns a non-empty completion from deepseek-chat', { timeout: INTEGRATION_TEST_TIMEOUT_MS }, async () => {
const provider = new DeepSeekProvider(
{ apiKey: optionalEnv(ENV_VAR)! },
makeMeteringStub(),
@@ -27,6 +27,7 @@
import { describe, expect, it } from 'vitest';
import {
INTEGRATION_TEST_TIMEOUT_MS,
makeMeteringStub,
optionalEnv,
skipUnlessEnv,
@@ -39,7 +40,7 @@ const ENV_VAR = 'PUTER_TEST_AI_GEMINI_API_KEY';
describe.skipIf(skipUnlessEnv(ENV_VAR))(
'GeminiChatProvider (integration)',
() => {
it('returns a non-empty completion from gemini-2.0-flash-lite', async () => {
it('returns a non-empty completion from gemini-2.0-flash-lite', { timeout: INTEGRATION_TEST_TIMEOUT_MS }, async () => {
const provider = new GeminiChatProvider(makeMeteringStub(), {
apiKey: optionalEnv(ENV_VAR)!,
});
@@ -27,6 +27,7 @@
import { describe, expect, it } from 'vitest';
import {
INTEGRATION_TEST_TIMEOUT_MS,
makeMeteringStub,
optionalEnv,
skipUnlessEnv,
@@ -37,7 +38,7 @@ import { GroqAIProvider } from './GroqAIProvider.js';
const ENV_VAR = 'PUTER_TEST_AI_GROQ_API_KEY';
describe.skipIf(skipUnlessEnv(ENV_VAR))('GroqAIProvider (integration)', () => {
it('returns a non-empty completion from llama-3.1-8b-instant', async () => {
it('returns a non-empty completion from llama-3.1-8b-instant', { timeout: INTEGRATION_TEST_TIMEOUT_MS }, async () => {
const provider = new GroqAIProvider(
{ apiKey: optionalEnv(ENV_VAR)! },
makeMeteringStub(),
@@ -26,6 +26,7 @@
import { describe, expect, it } from 'vitest';
import {
INTEGRATION_TEST_TIMEOUT_MS,
makeMeteringStub,
optionalEnv,
skipUnlessEnv,
@@ -38,7 +39,7 @@ const ENV_VAR = 'PUTER_TEST_AI_MISTRAL_API_KEY';
describe.skipIf(skipUnlessEnv(ENV_VAR))(
'MistralAIProvider (integration)',
() => {
it('returns a non-empty completion from mistral-small', async () => {
it('returns a non-empty completion from mistral-small', { timeout: INTEGRATION_TEST_TIMEOUT_MS }, async () => {
const provider = new MistralAIProvider(
{ apiKey: optionalEnv(ENV_VAR)! },
makeMeteringStub(),
@@ -26,6 +26,7 @@
import { describe, expect, it } from 'vitest';
import {
INTEGRATION_TEST_TIMEOUT_MS,
makeMeteringStub,
optionalEnv,
skipUnlessEnv,
@@ -38,7 +39,7 @@ const ENV_VAR = 'PUTER_TEST_AI_MOONSHOT_API_KEY';
describe.skipIf(skipUnlessEnv(ENV_VAR))(
'MoonshotProvider (integration)',
() => {
it('returns a non-empty completion from moonshot-v1-8k', async () => {
it('returns a non-empty completion from moonshot-v1-8k', { timeout: INTEGRATION_TEST_TIMEOUT_MS }, async () => {
const provider = new MoonshotProvider(
{ apiKey: optionalEnv(ENV_VAR)! },
makeMeteringStub(),
@@ -29,6 +29,7 @@
import { describe, expect, it } from 'vitest';
import {
INTEGRATION_TEST_TIMEOUT_MS,
makeMeteringStub,
optionalEnv,
skipUnlessEnv,
@@ -49,7 +50,7 @@ describe.skipIf(skipUnlessEnv(ENV_VAR))(
{ apiKey: optionalEnv(ENV_VAR)! },
);
it('returns a non-empty completion from gpt-4o-mini', async () => {
it('returns a non-empty completion from gpt-4o-mini', { timeout: INTEGRATION_TEST_TIMEOUT_MS }, async () => {
const provider = buildProvider();
const result = await withTestActor(() =>
provider.complete({
@@ -27,6 +27,7 @@
import { describe, expect, it } from 'vitest';
import {
INTEGRATION_TEST_TIMEOUT_MS,
makeMeteringStub,
optionalEnv,
skipUnlessEnv,
@@ -39,19 +40,25 @@ const ENV_VAR = 'PUTER_TEST_AI_OPENROUTER_API_KEY';
describe.skipIf(skipUnlessEnv(ENV_VAR))(
'OpenRouterProvider (integration)',
() => {
it('returns a non-empty completion via OpenRouter', async () => {
it('returns a non-empty completion via OpenRouter', { timeout: INTEGRATION_TEST_TIMEOUT_MS }, async () => {
const provider = new OpenRouterProvider(
{ apiKey: optionalEnv(ENV_VAR)! },
makeMeteringStub(),
);
const result = await withTestActor(() =>
// OpenRouter's `complete` destructure types all params
// as `any` and lists all six explicitly, so TS demands
// every field — pass undefineds for the unused ones.
provider.complete({
model: 'openrouter:google/gemini-2.0-flash-lite-001',
messages: [
{ role: 'user', content: 'Say hi in one word.' },
],
max_tokens: 16,
stream: undefined,
tools: undefined,
temperature: undefined,
}),
);
@@ -20,12 +20,17 @@
/**
* Integration test for the Together AI provider.
*
* Uses the `meta-llama/Meta-Llama-3.1-8B-Instruct-Turbo` cheap variant
* via Together. Skipped when `PUTER_TEST_AI_TOGETHER_API_KEY` is unset.
* Uses `Qwen/Qwen2.5-7B-Instruct-Turbo` non-Llama, small, cheap, and
* stays on Together's serverless tier. Llama variants on Together get
* rotated to dedicated endpoints often enough that they're not safe
* defaults. If Qwen also disappears, pick another live serverless
* model from https://api.together.ai/models?type=serverless. Skipped
* when `PUTER_TEST_AI_TOGETHER_API_KEY` is unset.
*/
import { describe, expect, it } from 'vitest';
import {
INTEGRATION_TEST_TIMEOUT_MS,
makeMeteringStub,
optionalEnv,
skipUnlessEnv,
@@ -38,7 +43,7 @@ const ENV_VAR = 'PUTER_TEST_AI_TOGETHER_API_KEY';
describe.skipIf(skipUnlessEnv(ENV_VAR))(
'TogetherAIProvider (integration)',
() => {
it('returns a non-empty completion from Llama 3.1 8B', async () => {
it('returns a non-empty completion from Qwen2.5 7B', { timeout: INTEGRATION_TEST_TIMEOUT_MS }, async () => {
const provider = new TogetherAIProvider(
{ apiKey: optionalEnv(ENV_VAR)! },
makeMeteringStub(),
@@ -46,7 +51,7 @@ describe.skipIf(skipUnlessEnv(ENV_VAR))(
const result = await withTestActor(() =>
provider.complete({
model: 'togetherai:meta-llama/Meta-Llama-3.1-8B-Instruct-Turbo',
model: 'togetherai:Qwen/Qwen2.5-7B-Instruct-Turbo',
messages: [
{ role: 'user', content: 'Say hi in one word.' },
],
@@ -26,6 +26,7 @@
import { describe, expect, it } from 'vitest';
import {
INTEGRATION_TEST_TIMEOUT_MS,
makeMeteringStub,
optionalEnv,
skipUnlessEnv,
@@ -36,7 +37,7 @@ import { XAIProvider } from './XAIProvider.js';
const ENV_VAR = 'PUTER_TEST_AI_XAI_API_KEY';
describe.skipIf(skipUnlessEnv(ENV_VAR))('XAIProvider (integration)', () => {
it('returns a non-empty completion from grok-3-mini', async () => {
it('returns a non-empty completion from grok-3-mini', { timeout: INTEGRATION_TEST_TIMEOUT_MS }, async () => {
const provider = new XAIProvider(
{ apiKey: optionalEnv(ENV_VAR)! },
makeMeteringStub(),
@@ -20,12 +20,17 @@
/**
* Integration test for the Z.AI (GLM) provider.
*
* Uses `glm-4.7-flashx` the cheapest text variant on the price
* sheet. Skipped when `PUTER_TEST_AI_ZAI_API_KEY` is unset.
* Uses `glm-4.6` with `thinking: disabled` passed through `custom`.
* GLM models default to reasoning mode and route their tokens to a
* `reasoning_content` field, leaving `content` empty under tight
* budgets. Disabling thinking forces a plain text response so the
* usual `message.content` assertion works. Skipped when
* `PUTER_TEST_AI_ZAI_API_KEY` is unset.
*/
import { describe, expect, it } from 'vitest';
import {
INTEGRATION_TEST_TIMEOUT_MS,
makeMeteringStub,
optionalEnv,
skipUnlessEnv,
@@ -36,7 +41,7 @@ import { ZAIProvider } from './ZAIProvider.js';
const ENV_VAR = 'PUTER_TEST_AI_ZAI_API_KEY';
describe.skipIf(skipUnlessEnv(ENV_VAR))('ZAIProvider (integration)', () => {
it('returns a non-empty completion from glm-4.7-flashx', async () => {
it('returns a non-empty completion from glm-4.6', { timeout: INTEGRATION_TEST_TIMEOUT_MS }, async () => {
const provider = new ZAIProvider(
{ apiKey: optionalEnv(ENV_VAR)! },
makeMeteringStub(),
@@ -44,9 +49,10 @@ describe.skipIf(skipUnlessEnv(ENV_VAR))('ZAIProvider (integration)', () => {
const result = await withTestActor(() =>
provider.complete({
model: 'glm-4.7-flashx',
model: 'glm-4.6',
messages: [{ role: 'user', content: 'Say hi in one word.' }],
max_tokens: 16,
custom: { thinking: { type: 'disabled' } },
}),
);
+1
View File
@@ -115,6 +115,7 @@ export interface IChatMessageResult {
stream?: never;
finally_fn?: never;
normalized?: boolean;
via_ai_chat_service?: boolean;
}
export type IChatCompleteResult = IChatStreamResult | IChatMessageResult;
@@ -27,6 +27,7 @@
import { describe, expect, it } from 'vitest';
import {
INTEGRATION_TEST_TIMEOUT_MS,
makeMeteringStub,
optionalEnv,
skipUnlessEnv,
@@ -39,7 +40,7 @@ const ENV_VAR = 'PUTER_TEST_AI_GEMINI_API_KEY';
describe.skipIf(skipUnlessEnv(ENV_VAR))(
'GeminiImageProvider (integration)',
() => {
it('returns image data from imagen-4.0-fast', async () => {
it('returns image data from imagen-4.0-fast', { timeout: INTEGRATION_TEST_TIMEOUT_MS }, async () => {
const provider = new GeminiImageProvider(
{ apiKey: optionalEnv(ENV_VAR)! },
makeMeteringStub(),
@@ -27,6 +27,7 @@
import { describe, expect, it } from 'vitest';
import {
INTEGRATION_TEST_TIMEOUT_MS,
makeMeteringStub,
optionalEnv,
skipUnlessEnv,
@@ -39,7 +40,7 @@ const ENV_VAR = 'PUTER_TEST_AI_OPENAI_API_KEY';
describe.skipIf(skipUnlessEnv(ENV_VAR))(
'OpenAiImageProvider (integration)',
() => {
it('returns an image url/data from dall-e-2 at 256x256', async () => {
it('returns an image url/data from dall-e-2 at 256x256', { timeout: INTEGRATION_TEST_TIMEOUT_MS }, async () => {
const provider = new OpenAiImageProvider(
{ apiKey: optionalEnv(ENV_VAR)! },
makeMeteringStub(),
@@ -27,6 +27,7 @@
import { Readable } from 'node:stream';
import { describe, expect, it } from 'vitest';
import {
INTEGRATION_TEST_TIMEOUT_MS,
makeMeteringStub,
optionalEnv,
skipUnlessEnv,
@@ -39,7 +40,7 @@ const ENV_VAR = 'PUTER_TEST_AI_ELEVENLABS_API_KEY';
describe.skipIf(skipUnlessEnv(ENV_VAR))(
'ElevenLabsTTSProvider (integration)',
() => {
it('returns an audio stream from eleven_flash_v2_5', async () => {
it('returns an audio stream from eleven_flash_v2_5', { timeout: INTEGRATION_TEST_TIMEOUT_MS }, async () => {
const provider = new ElevenLabsTTSProvider(makeMeteringStub(), {
apiKey: optionalEnv(ENV_VAR)!,
});
@@ -28,6 +28,7 @@
import { Readable } from 'node:stream';
import { describe, expect, it } from 'vitest';
import {
INTEGRATION_TEST_TIMEOUT_MS,
makeMeteringStub,
optionalEnv,
skipUnlessEnv,
@@ -40,7 +41,7 @@ const ENV_VAR = 'PUTER_TEST_AI_OPENAI_API_KEY';
describe.skipIf(skipUnlessEnv(ENV_VAR))(
'OpenAITTSProvider (integration)',
() => {
it('returns an audio stream from tts-1', async () => {
it('returns an audio stream from tts-1', { timeout: INTEGRATION_TEST_TIMEOUT_MS }, async () => {
const provider = new OpenAITTSProvider(makeMeteringStub(), {
apiKey: optionalEnv(ENV_VAR)!,
});
@@ -49,6 +49,14 @@ export const optionalEnv = (name: string): string | undefined => {
*/
export const skipUnlessEnv = (name: string): boolean => !optionalEnv(name);
/**
* Per-test timeout for provider integration tests. The default 5s
* vitest timeout is way too short for real API calls image
* generation in particular routinely takes 1530s. Pass this as the
* third argument to `it(...)`.
*/
export const INTEGRATION_TEST_TIMEOUT_MS = 90_000;
/**
* Returns a no-op MeteringService stub. Real metering would write to
* DynamoDB / Redis, which integration tests for AI providers don't