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
842 lines
25 KiB
TypeScript
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',
|
|
}),
|
|
},
|
|
],
|
|
});
|
|
});
|
|
});
|
|
});
|