Files
wehub-resource-sync 0d3cb498a3
CI / Shell Format Check (push) Has been cancelled
CI / Check Ruby (3.4) (push) Has been cancelled
CI / CI Config (push) Has been cancelled
CI / Test on Node ${{ matrix.node }} and ${{ matrix.os }}${{ matrix.shard && format(' (shard {0}/3)', matrix.shard) || '' }} (push) Has been cancelled
CI / Build on Node ${{ matrix.node }} (push) Has been cancelled
CI / Style Check (push) Has been cancelled
CI / Generate Assets (push) Has been cancelled
CI / Check Python (3.14) (push) Has been cancelled
CI / Check Python (3.9) (push) Has been cancelled
CI / Build Docs (push) Has been cancelled
CI / Code Scan Action (push) Has been cancelled
CI / Site tests (push) Has been cancelled
CI / webui tests (push) Has been cancelled
CI / Run Integration Tests (push) Has been cancelled
CI / Run Smoke Tests (push) Has been cancelled
CI / Go Tests (push) Has been cancelled
CI / Share Test (push) Has been cancelled
CI / Redteam (Production API) (push) Has been cancelled
CI / Redteam (Staging API) (push) Has been cancelled
CI / GitHub Actions Lint (push) Has been cancelled
CI / Check Ruby (3.0) (push) Has been cancelled
release-please / release-please (push) Has been cancelled
release-please / build (push) Has been cancelled
release-please / publish-npm (push) Has been cancelled
release-please / publish-npm-backfill (push) Has been cancelled
release-please / docker (push) Has been cancelled
release-please / publish-code-scan-action (push) Has been cancelled
release-please / attest-code-scan-action (push) Has been cancelled
Deploy local.promptfoo.app / Deploy to Cloudflare Pages (push) Has been cancelled
Test and Publish Multi-arch Docker Image / test (push) Has been cancelled
Test and Publish Multi-arch Docker Image / build-docker-and-push-digests (map[digest-suffix:linux-amd64 platform:linux/amd64 runner:ubuntu-latest]) (push) Has been cancelled
Test and Publish Multi-arch Docker Image / build-docker-and-push-digests (map[digest-suffix:linux-arm64 platform:linux/arm64 runner:ubuntu-24.04-arm]) (push) Has been cancelled
Test and Publish Multi-arch Docker Image / merge-docker-digests (push) Has been cancelled
Test and Publish Multi-arch Docker Image / Attest Multi-arch Image (push) Has been cancelled
Validate Renovate Config / Validate Renovate Configuration (push) Has been cancelled
chore: import upstream snapshot with attribution
2026-07-13 13:24:08 +08:00

842 lines
25 KiB
TypeScript

import { trace } from '@opentelemetry/api';
import OpenAI from 'openai';
import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest';
import { disableCache, enableCache } from '../../../src/cache';
import { OpenAiAssistantProvider } from '../../../src/providers/openai/assistant';
import { mockProcessEnv } from '../../util/utils';
import { getOpenAiMissingApiKeyMessage } from './shared';
import type { CallbackContext } from '../../../src/providers/openai/types';
vi.mock('openai');
interface RecordedSpan {
name: string;
attributes: Record<string, any>;
status?: { code: number; message?: string };
ended: boolean;
}
function installTracerSpy(): RecordedSpan[] {
const spans: RecordedSpan[] = [];
const make = (name: string, attributes: Record<string, any> = {}) => {
const entry: RecordedSpan = { name, attributes: { ...attributes }, ended: false };
spans.push(entry);
return {
setAttribute: (key: string, value: unknown) => {
entry.attributes[key] = value;
},
setAttributes: (attrs: Record<string, unknown>) => Object.assign(entry.attributes, attrs),
setStatus: (status: { code: number; message?: string }) => {
entry.status = status;
},
end: () => {
entry.ended = true;
},
recordException: () => undefined,
addEvent: () => undefined,
spanContext: () => ({ traceId: 'x', spanId: 'y' }),
isRecording: () => true,
updateName: () => undefined,
};
};
vi.spyOn(trace, 'getTracer').mockReturnValue({
startSpan: (name: string, options?: { attributes?: Record<string, unknown> }) =>
make(name, options?.attributes),
startActiveSpan: (...args: any[]) => {
const name = args[0];
const options = typeof args[1] === 'object' ? args[1] : undefined;
const callback = args[args.length - 1];
return callback(make(name, options?.attributes));
},
} as any);
return spans;
}
describe('OpenAI Provider', () => {
beforeEach(() => {
vi.resetAllMocks();
disableCache();
});
afterEach(() => {
enableCache();
});
describe('OpenAiAssistantProvider', () => {
let mockClient: any;
beforeEach(() => {
vi.clearAllMocks();
mockClient = {
beta: {
threads: {
createAndRun: vi.fn(),
runs: {
retrieve: vi.fn(),
submitToolOutputs: vi.fn(),
steps: {
list: vi.fn(),
},
},
messages: {
retrieve: vi.fn(),
},
},
},
};
vi.mocked(OpenAI).mockImplementation(function (this: any) {
Object.assign(this, mockClient);
return this;
});
});
const provider = new OpenAiAssistantProvider('test-assistant-id', {
config: {
apiKey: 'test-key',
organization: 'test-org',
functionToolCallbacks: {
test_function: async (_args: string) => 'Function result',
},
},
});
it('should handle successful assistant completion', async () => {
const mockRun = {
id: 'run_123',
thread_id: 'thread_123',
status: 'completed',
};
const mockSteps = {
data: [
{
id: 'step_1',
step_details: {
type: 'message_creation',
message_creation: {
message_id: 'msg_1',
},
},
},
],
};
const mockMessage = {
role: 'assistant',
content: [
{
type: 'text',
text: {
value: 'Test response',
},
},
],
};
mockClient.beta.threads.createAndRun.mockResolvedValue(mockRun);
mockClient.beta.threads.runs.retrieve.mockResolvedValue(mockRun);
mockClient.beta.threads.runs.steps.list.mockResolvedValue(mockSteps);
mockClient.beta.threads.messages.retrieve.mockResolvedValue(mockMessage);
const result = await provider.callApi('Test prompt');
expect(OpenAI).toHaveBeenCalledWith(
expect.objectContaining({
defaultHeaders: expect.objectContaining({
'X-OpenAI-Originator': 'promptfoo',
}),
}),
);
expect(result.output).toBe('[Assistant] Test response');
expect(mockClient.beta.threads.createAndRun).toHaveBeenCalledTimes(1);
expect(mockClient.beta.threads.runs.retrieve).toHaveBeenCalledTimes(1);
expect(mockClient.beta.threads.runs.steps.list).toHaveBeenCalledTimes(1);
expect(mockClient.beta.threads.messages.retrieve).toHaveBeenCalledTimes(1);
});
it('drops the SDK organization option when a case-variant org header overrides it', async () => {
const mockRun = { id: 'run_123', thread_id: 'thread_123', status: 'completed' };
const mockSteps = {
data: [
{
id: 'step_1',
step_details: { type: 'message_creation', message_creation: { message_id: 'msg_1' } },
},
],
};
const mockMessage = {
role: 'assistant',
content: [{ type: 'text', text: { value: 'Test response' } }],
};
mockClient.beta.threads.createAndRun.mockResolvedValue(mockRun);
mockClient.beta.threads.runs.retrieve.mockResolvedValue(mockRun);
mockClient.beta.threads.runs.steps.list.mockResolvedValue(mockSteps);
mockClient.beta.threads.messages.retrieve.mockResolvedValue(mockMessage);
const overrideProvider = new OpenAiAssistantProvider('test-assistant-id', {
config: {
apiKey: 'test-key',
organization: 'test-org',
headers: { 'openai-organization': 'custom-org' },
},
});
await overrideProvider.callApi('Test prompt');
// The SDK would otherwise inject its own canonical OpenAI-Organization header
// from `organization`, duplicating (and beating) the custom override.
expect(OpenAI).toHaveBeenCalledWith(
expect.objectContaining({
organization: undefined,
defaultHeaders: expect.objectContaining({ 'openai-organization': 'custom-org' }),
}),
);
const sdkArgs = vi.mocked(OpenAI).mock.calls.at(-1)?.[0] as { defaultHeaders: object };
expect(sdkArgs.defaultHeaders).not.toHaveProperty('OpenAI-Organization');
});
it('emits a chat <model> span around the assistant run', async () => {
const spans = installTracerSpy();
const mockRun = { id: 'run_123', thread_id: 'thread_123', status: 'completed' };
const mockSteps = {
data: [
{
id: 'step_1',
step_details: { type: 'message_creation', message_creation: { message_id: 'msg_1' } },
},
],
};
const mockMessage = {
role: 'assistant',
content: [{ type: 'text', text: { value: 'Test response' } }],
};
mockClient.beta.threads.createAndRun.mockResolvedValue(mockRun);
mockClient.beta.threads.runs.retrieve.mockResolvedValue(mockRun);
mockClient.beta.threads.runs.steps.list.mockResolvedValue(mockSteps);
mockClient.beta.threads.messages.retrieve.mockResolvedValue(mockMessage);
const localProvider = new OpenAiAssistantProvider('test-assistant-id', {
config: { apiKey: 'test-key' },
});
await localProvider.callApi('Test prompt');
const chatSpan = spans.find((span) => span.name === 'chat test-assistant-id');
expect(chatSpan).toBeDefined();
expect(chatSpan?.attributes).toMatchObject({
'gen_ai.system': 'openai',
'gen_ai.operation.name': 'chat',
'gen_ai.request.model': 'test-assistant-id',
});
expect(chatSpan?.ended).toBe(true);
// SpanStatusCode.OK === 1
expect(chatSpan?.status?.code).toBe(1);
});
it('marks the chat span ERROR when the assistant run fails', async () => {
const spans = installTracerSpy();
mockClient.beta.threads.createAndRun.mockRejectedValue(new Error('assistant boom'));
const localProvider = new OpenAiAssistantProvider('test-assistant-id', {
config: { apiKey: 'test-key' },
});
const result = await localProvider.callApi('Test prompt');
expect(result.error).toBeDefined();
const chatSpan = spans.find((span) => span.name === 'chat test-assistant-id');
expect(chatSpan).toBeDefined();
// SpanStatusCode.ERROR === 2
expect(chatSpan?.status?.code).toBe(2);
expect(chatSpan?.ended).toBe(true);
});
it('should preserve an explicit temperature of 0', async () => {
const mockRun = {
id: 'run_123',
thread_id: 'thread_123',
status: 'completed',
};
const mockSteps = {
data: [
{
id: 'step_1',
step_details: {
type: 'message_creation',
message_creation: {
message_id: 'msg_1',
},
},
},
],
};
const mockMessage = {
role: 'assistant',
content: [
{
type: 'text',
text: {
value: 'Test response',
},
},
],
};
mockClient.beta.threads.createAndRun.mockResolvedValue(mockRun);
mockClient.beta.threads.runs.retrieve.mockResolvedValue(mockRun);
mockClient.beta.threads.runs.steps.list.mockResolvedValue(mockSteps);
mockClient.beta.threads.messages.retrieve.mockResolvedValue(mockMessage);
const provider = new OpenAiAssistantProvider('test-assistant-id', {
config: {
apiKey: 'test-key',
temperature: 0,
},
});
await provider.callApi('Test prompt');
expect(mockClient.beta.threads.createAndRun).toHaveBeenCalledWith(
expect.objectContaining({
temperature: 0,
}),
);
});
it('should handle function calling', async () => {
const mockRun = {
id: 'run_123',
thread_id: 'thread_123',
status: 'requires_action',
required_action: {
type: 'submit_tool_outputs',
submit_tool_outputs: {
tool_calls: [
{
id: 'call_123',
type: 'function',
function: {
name: 'test_function',
arguments: '{"arg": "value"}',
},
},
],
},
},
};
const mockCompletedRun = {
...mockRun,
status: 'completed',
required_action: null,
};
const mockSteps = {
data: [
{
id: 'step_1',
step_details: {
type: 'tool_calls',
tool_calls: [
{
type: 'function',
function: {
name: 'test_function',
arguments: '{"arg": "value"}',
output: 'Function result',
},
},
],
},
},
],
};
mockClient.beta.threads.createAndRun.mockResolvedValue(mockRun);
mockClient.beta.threads.runs.retrieve
.mockResolvedValueOnce(mockRun)
.mockResolvedValueOnce(mockCompletedRun);
mockClient.beta.threads.runs.submitToolOutputs.mockResolvedValue(mockCompletedRun);
mockClient.beta.threads.runs.steps.list.mockResolvedValue(mockSteps);
const result = await provider.callApi('Test prompt');
expect(result.output).toBe(
'[Call function test_function with arguments {"arg": "value"}]\n\n[Function output: Function result]',
);
expect(mockClient.beta.threads.createAndRun).toHaveBeenCalledTimes(1);
expect(mockClient.beta.threads.runs.retrieve).toHaveBeenCalledTimes(2);
expect(mockClient.beta.threads.runs.submitToolOutputs).toHaveBeenCalledTimes(1);
expect(mockClient.beta.threads.runs.steps.list).toHaveBeenCalledTimes(1);
});
it('should handle run failures', async () => {
const mockRun = {
id: 'run_123',
thread_id: 'thread_123',
status: 'failed',
last_error: {
message: 'Test error message',
},
};
mockClient.beta.threads.createAndRun.mockResolvedValue(mockRun);
mockClient.beta.threads.runs.retrieve.mockResolvedValue(mockRun);
const result = await provider.callApi('Test prompt');
expect(result.error).toBe('Thread run failed: Test error message');
expect(mockClient.beta.threads.createAndRun).toHaveBeenCalledTimes(1);
expect(mockClient.beta.threads.runs.retrieve).toHaveBeenCalledTimes(1);
});
it('should handle API errors', async () => {
const error = new OpenAI.APIError(500, {}, 'API Error', new Headers());
Object.defineProperty(error, 'type', {
value: 'API Error',
writable: true,
configurable: true,
});
Object.defineProperty(error, 'message', {
value: 'API Error',
writable: true,
configurable: true,
});
mockClient.beta.threads.createAndRun.mockRejectedValueOnce(error);
const provider = new OpenAiAssistantProvider('test-assistant-id', {
config: {
apiKey: 'test-key',
},
});
const result = await provider.callApi('Test prompt');
expect(result.error).toBe('API error: API Error API Error');
expect(mockClient.beta.threads.createAndRun).toHaveBeenCalledTimes(1);
});
it('should handle missing API key', async () => {
const restoreEnv = mockProcessEnv({ OPENAI_API_KEY: undefined });
try {
const providerNoKey = new OpenAiAssistantProvider('test-assistant-id', {
env: {
OPENAI_API_KEY: undefined,
},
});
await expect(providerNoKey.callApi('Test prompt')).rejects.toThrow(
getOpenAiMissingApiKeyMessage(),
);
} finally {
restoreEnv();
}
});
it('should use custom apiKeyEnvar in missing API key errors', async () => {
const restoreEnv = mockProcessEnv({
OPENAI_API_KEY: undefined,
CUSTOM_ASSISTANT_API_KEY: undefined,
});
try {
const providerNoKey = new OpenAiAssistantProvider('test-assistant-id', {
config: {
apiKeyEnvar: 'CUSTOM_ASSISTANT_API_KEY',
},
env: {
OPENAI_API_KEY: undefined,
CUSTOM_ASSISTANT_API_KEY: undefined,
},
});
await expect(providerNoKey.callApi('Test prompt')).rejects.toThrow(
getOpenAiMissingApiKeyMessage('CUSTOM_ASSISTANT_API_KEY'),
);
} finally {
restoreEnv();
}
});
});
describe('Function Callbacks with Context', () => {
let mockClient: any;
beforeEach(() => {
vi.clearAllMocks();
disableCache();
mockClient = {
beta: {
threads: {
createAndRun: vi.fn(),
runs: {
retrieve: vi.fn(),
submitToolOutputs: vi.fn(),
steps: {
list: vi.fn(),
},
},
messages: {
retrieve: vi.fn(),
},
},
},
};
vi.mocked(OpenAI).mockImplementation(function (this: any) {
Object.assign(this, mockClient);
return this;
});
});
it('should pass context to function callbacks', async () => {
const mockCallback = vi.fn().mockResolvedValue('test result');
const provider = new OpenAiAssistantProvider('asst_test', {
config: {
apiKey: 'test-key',
functionToolCallbacks: {
test_function: mockCallback,
},
},
});
const mockRun = {
id: 'run_test',
thread_id: 'thread_test',
status: 'requires_action',
required_action: {
type: 'submit_tool_outputs',
submit_tool_outputs: {
tool_calls: [
{
id: 'call_test',
type: 'function',
function: {
name: 'test_function',
arguments: '{"param": "value"}',
},
},
],
},
},
};
const mockCompletedRun = {
...mockRun,
status: 'completed',
};
const mockSteps = {
data: [
{
id: 'step_test',
step_details: {
type: 'message_creation',
message_creation: {
message_id: 'msg_test',
},
},
},
],
};
const mockMessage = {
id: 'msg_test',
role: 'assistant',
content: [
{
type: 'text',
text: {
value: 'Test response',
},
},
],
};
mockClient.beta.threads.createAndRun.mockResolvedValue(mockRun);
mockClient.beta.threads.runs.retrieve
.mockResolvedValueOnce(mockRun)
.mockResolvedValueOnce(mockCompletedRun);
mockClient.beta.threads.runs.submitToolOutputs.mockResolvedValue(mockCompletedRun);
mockClient.beta.threads.runs.steps.list.mockResolvedValue(mockSteps);
mockClient.beta.threads.messages.retrieve.mockResolvedValue(mockMessage);
await provider.callApi('test prompt');
// Verify that the callback was called with the correct context
expect(mockCallback).toHaveBeenCalledWith(
{ param: 'value' },
{
threadId: 'thread_test',
runId: 'run_test',
assistantId: 'asst_test',
provider: 'openai',
},
);
});
it('should work with callbacks that do not use context', async () => {
const oldStyleCallback = vi.fn().mockResolvedValue('old style result');
const provider = new OpenAiAssistantProvider('asst_test', {
config: {
apiKey: 'test-key',
functionToolCallbacks: {
old_function: oldStyleCallback,
},
},
});
const mockRun = {
id: 'run_test',
thread_id: 'thread_test',
status: 'requires_action',
required_action: {
type: 'submit_tool_outputs',
submit_tool_outputs: {
tool_calls: [
{
id: 'call_test',
type: 'function',
function: {
name: 'old_function',
arguments: '{"param": "value"}',
},
},
],
},
},
};
const mockCompletedRun = { ...mockRun, status: 'completed' };
const mockSteps = { data: [] };
mockClient.beta.threads.createAndRun.mockResolvedValue(mockRun);
mockClient.beta.threads.runs.retrieve
.mockResolvedValueOnce(mockRun)
.mockResolvedValueOnce(mockCompletedRun);
mockClient.beta.threads.runs.submitToolOutputs.mockResolvedValue(mockCompletedRun);
mockClient.beta.threads.runs.steps.list.mockResolvedValue(mockSteps);
await provider.callApi('test prompt');
// Callback should be called with args and context, but context is optional
expect(oldStyleCallback).toHaveBeenCalledWith(
{ param: 'value' },
expect.objectContaining({
threadId: 'thread_test',
runId: 'run_test',
assistantId: 'asst_test',
provider: 'openai',
}),
);
});
it('should handle string-based function callbacks', async () => {
const provider = new OpenAiAssistantProvider('asst_test', {
config: {
apiKey: 'test-key',
functionToolCallbacks: {
string_function:
'(args, context) => { return `received: ${JSON.stringify(args)} with context: ${JSON.stringify(context)}`; }',
},
},
});
const mockRun = {
id: 'run_test',
thread_id: 'thread_test',
status: 'requires_action',
required_action: {
type: 'submit_tool_outputs',
submit_tool_outputs: {
tool_calls: [
{
id: 'call_test',
type: 'function',
function: {
name: 'string_function',
arguments: '{"test": "data"}',
},
},
],
},
},
};
const mockCompletedRun = { ...mockRun, status: 'completed' };
const mockSteps = { data: [] };
mockClient.beta.threads.createAndRun.mockResolvedValue(mockRun);
mockClient.beta.threads.runs.retrieve
.mockResolvedValueOnce(mockRun)
.mockResolvedValueOnce(mockCompletedRun);
mockClient.beta.threads.runs.submitToolOutputs.mockResolvedValue(mockCompletedRun);
mockClient.beta.threads.runs.steps.list.mockResolvedValue(mockSteps);
await provider.callApi('test prompt');
// Check that the tool output was submitted correctly
expect(mockClient.beta.threads.runs.submitToolOutputs).toHaveBeenCalledWith('run_test', {
thread_id: 'thread_test',
tool_outputs: [
{
tool_call_id: 'call_test',
output: expect.stringContaining('received: {"test":"data"}'),
},
],
});
});
it('should handle callbacks that access context properties', async () => {
const contextAwareCallback = vi.fn().mockImplementation(function (
args: any,
context?: CallbackContext,
) {
const result = {
originalArgs: args,
contextInfo: {
threadId: context?.threadId,
provider: context?.provider,
},
};
return Promise.resolve(JSON.stringify(result));
});
const provider = new OpenAiAssistantProvider('asst_test', {
config: {
apiKey: 'test-key',
functionToolCallbacks: {
context_function: contextAwareCallback,
},
},
});
const mockRun = {
id: 'run_test',
thread_id: 'thread_test',
status: 'requires_action',
required_action: {
type: 'submit_tool_outputs',
submit_tool_outputs: {
tool_calls: [
{
id: 'call_test',
type: 'function',
function: {
name: 'context_function',
arguments: '{"user_id": "123"}',
},
},
],
},
},
};
const mockCompletedRun = { ...mockRun, status: 'completed' };
const mockSteps = { data: [] };
mockClient.beta.threads.createAndRun.mockResolvedValue(mockRun);
mockClient.beta.threads.runs.retrieve
.mockResolvedValueOnce(mockRun)
.mockResolvedValueOnce(mockCompletedRun);
mockClient.beta.threads.runs.submitToolOutputs.mockResolvedValue(mockCompletedRun);
mockClient.beta.threads.runs.steps.list.mockResolvedValue(mockSteps);
await provider.callApi('test prompt');
// Verify the callback was called with the correct parameters
expect(contextAwareCallback).toHaveBeenCalledWith(
{ user_id: '123' },
expect.objectContaining({
threadId: 'thread_test',
runId: 'run_test',
assistantId: 'asst_test',
provider: 'openai',
}),
);
// Verify the tool output contains the context information
expect(mockClient.beta.threads.runs.submitToolOutputs).toHaveBeenCalledWith('run_test', {
thread_id: 'thread_test',
tool_outputs: [
{
tool_call_id: 'call_test',
output: JSON.stringify({
originalArgs: { user_id: '123' },
contextInfo: {
threadId: 'thread_test',
provider: 'openai',
},
}),
},
],
});
});
it('should handle function callback errors gracefully', async () => {
const errorCallback = vi.fn().mockRejectedValue(new Error('Callback error'));
const provider = new OpenAiAssistantProvider('asst_test', {
config: {
apiKey: 'test-key',
functionToolCallbacks: {
error_function: errorCallback,
},
},
});
const mockRun = {
id: 'run_test',
thread_id: 'thread_test',
status: 'requires_action',
required_action: {
type: 'submit_tool_outputs',
submit_tool_outputs: {
tool_calls: [
{
id: 'call_test',
type: 'function',
function: {
name: 'error_function',
arguments: '{"param": "value"}',
},
},
],
},
},
};
const mockCompletedRun = { ...mockRun, status: 'completed' };
const mockSteps = { data: [] };
mockClient.beta.threads.createAndRun.mockResolvedValue(mockRun);
mockClient.beta.threads.runs.retrieve
.mockResolvedValueOnce(mockRun)
.mockResolvedValueOnce(mockCompletedRun);
mockClient.beta.threads.runs.submitToolOutputs.mockResolvedValue(mockCompletedRun);
mockClient.beta.threads.runs.steps.list.mockResolvedValue(mockSteps);
await provider.callApi('test prompt');
// Verify error was handled and submitted as tool output
expect(mockClient.beta.threads.runs.submitToolOutputs).toHaveBeenCalledWith('run_test', {
thread_id: 'thread_test',
tool_outputs: [
{
tool_call_id: 'call_test',
output: JSON.stringify({
error: 'Error in error_function: Callback error',
}),
},
],
});
});
});
});