@launchdarkly/server-sdk-ai 0.12.3 → 0.14.0

package/__tests__/LDAIClientImpl.test.ts

@@ -1,382 +1,589 @@
 import { LDContext } from '@launchdarkly/js-server-sdk-common';
 
-import { LDAIAgentDefaults } from '../src/api/agents';
-import { LDAIDefaults } from '../src/api/config';
+import {
+  LDAIAgentConfigDefault,
+  LDAICompletionConfigDefault,
+  LDAIJudgeConfigDefault,
+} from '../src/api/config/types';
+import { Judge } from '../src/api/judge/Judge';
+import { AIProviderFactory } from '../src/api/providers/AIProviderFactory';
 import { LDAIClientImpl } from '../src/LDAIClientImpl';
 import { LDClientMin } from '../src/LDClientMin';
 
+// Mock Judge and AIProviderFactory
+jest.mock('../src/api/judge/Judge');
+jest.mock('../src/api/providers/AIProviderFactory');
+
 const mockLdClient: jest.Mocked<LDClientMin> = {
   variation: jest.fn(),
   track: jest.fn(),
 };
 
+// Reset mocks before each test
+beforeEach(() => {
+  jest.clearAllMocks();
+});
+
 const testContext: LDContext = { kind: 'user', key: 'test-user' };
 
-it('returns config with interpolated messages', async () => {
-  const client = new LDAIClientImpl(mockLdClient);
-  const key = 'test-flag';
-  const defaultValue: LDAIDefaults = {
-    model: { name: 'test', parameters: { name: 'test-model' } },
-    messages: [],
-    enabled: true,
-  };
-
-  const mockVariation = {
-    model: {
-      name: 'example-model',
-      parameters: { name: 'imagination', temperature: 0.7, maxTokens: 4096 },
-    },
-    provider: {
-      name: 'example-provider',
-    },
-    messages: [
-      { role: 'system', content: 'Hello {{name}}' },
-      { role: 'user', content: 'Score: {{score}}' },
-    ],
-    _ldMeta: {
-      variationKey: 'v1',
-      enabled: true,
-    },
-  };
-
-  mockLdClient.variation.mockResolvedValue(mockVariation);
-
-  const variables = { name: 'John', score: 42 };
-  const result = await client.config(key, testContext, defaultValue, variables);
-
-  expect(result).toEqual({
-    model: {
-      name: 'example-model',
-      parameters: { name: 'imagination', temperature: 0.7, maxTokens: 4096 },
-    },
-    provider: {
-      name: 'example-provider',
-    },
-    messages: [
+describe('config evaluation', () => {
+  it('evaluates completion config successfully with variable interpolation', async () => {
+    const client = new LDAIClientImpl(mockLdClient);
+    const key = 'test-flag';
+    const defaultValue: LDAICompletionConfigDefault = {
+      enabled: false,
+    };
+
+    const mockVariation = {
+      model: {
+        name: 'example-model',
+        parameters: { name: 'imagination', temperature: 0.7, maxTokens: 4096 },
+      },
+      provider: {
+        name: 'example-provider',
+      },
+      messages: [
+        { role: 'system', content: 'Hello {{name}}' },
+        { role: 'user', content: 'Score: {{score}}' },
+      ],
+      _ldMeta: {
+        variationKey: 'v1',
+        enabled: true,
+        mode: 'completion',
+      },
+    };
+
+    mockLdClient.variation.mockResolvedValue(mockVariation);
+
+    const variables = { name: 'John', score: 42 };
+    const evaluateSpy = jest.spyOn(client as any, '_evaluate');
+    const result = await client.completionConfig(key, testContext, defaultValue, variables);
+
+    expect(evaluateSpy).toHaveBeenCalledWith(
+      key,
+      testContext,
+      defaultValue,
+      'completion',
+      variables,
+    );
+    expect(result.messages).toEqual([
       { role: 'system', content: 'Hello John' },
       { role: 'user', content: 'Score: 42' },
-    ],
-    tracker: expect.any(Object),
-    enabled: true,
-    toVercelAISDK: expect.any(Function),
+    ]);
+    expect(result.tracker).toBeDefined();
+    expect(result.enabled).toBe(true);
+    evaluateSpy.mockRestore();
   });
 
-  // Verify tracking was called
-  expect(mockLdClient.track).toHaveBeenCalledWith(
-    '$ld:ai:config:function:single',
-    testContext,
-    key,
-    1,
-  );
-});
+  it('includes context (ldctx) in variables for message interpolation', async () => {
+    const client = new LDAIClientImpl(mockLdClient);
+    const key = 'test-flag';
+    const defaultValue: LDAICompletionConfigDefault = {
+      enabled: false,
+    };
+
+    const mockVariation = {
+      messages: [{ role: 'system', content: 'User key: {{ldctx.key}}' }],
+      _ldMeta: { variationKey: 'v1', enabled: true, mode: 'completion' },
+    };
 
-it('includes context in variables for messages interpolation', async () => {
-  const client = new LDAIClientImpl(mockLdClient);
-  const key = 'test-flag';
-  const defaultValue: LDAIDefaults = {
-    model: { name: 'test', parameters: { name: 'test-model' } },
-    messages: [],
-  };
+    mockLdClient.variation.mockResolvedValue(mockVariation);
 
-  const mockVariation = {
-    messages: [{ role: 'system', content: 'User key: {{ldctx.key}}' }],
-    _ldMeta: { variationKey: 'v1', enabled: true },
-  };
+    const result = await client.completionConfig(key, testContext, defaultValue);
 
-  mockLdClient.variation.mockResolvedValue(mockVariation);
+    expect(result.messages?.[0].content).toBe('User key: test-user');
+  });
 
-  const result = await client.config(key, testContext, defaultValue);
+  it('evaluates agent config successfully with instruction interpolation', async () => {
+    const client = new LDAIClientImpl(mockLdClient);
+    const key = 'test-agent';
+    const defaultValue: LDAIAgentConfigDefault = {
+      enabled: false,
+    };
 
-  expect(result.messages?.[0].content).toBe('User key: test-user');
-  expect(result.toVercelAISDK).toEqual(expect.any(Function));
-});
+    const mockVariation = {
+      model: {
+        name: 'example-model',
+        parameters: { temperature: 0.7, maxTokens: 4096 },
+      },
+      provider: {
+        name: 'example-provider',
+      },
+      instructions:
+        'You are a helpful assistant. Your name is {{name}} and your score is {{score}}',
+      _ldMeta: {
+        variationKey: 'v1',
+        enabled: true,
+        mode: 'agent',
+      },
+    };
 
-it('handles missing metadata in variation', async () => {
-  const client = new LDAIClientImpl(mockLdClient);
-  const key = 'test-flag';
-  const defaultValue: LDAIDefaults = {
-    model: { name: 'test', parameters: { name: 'test-model' } },
-    messages: [],
-  };
-
-  const mockVariation = {
-    model: { name: 'example-provider', parameters: { name: 'imagination' } },
-    messages: [{ role: 'system', content: 'Hello' }],
-  };
-
-  mockLdClient.variation.mockResolvedValue(mockVariation);
-
-  const result = await client.config(key, testContext, defaultValue);
-
-  expect(result).toEqual({
-    model: { name: 'example-provider', parameters: { name: 'imagination' } },
-    messages: [{ role: 'system', content: 'Hello' }],
-    tracker: expect.any(Object),
-    enabled: false,
-    toVercelAISDK: expect.any(Function),
-  });
-});
+    mockLdClient.variation.mockResolvedValue(mockVariation);
 
-it('passes the default value to the underlying client', async () => {
-  const client = new LDAIClientImpl(mockLdClient);
-  const key = 'non-existent-flag';
-  const defaultValue: LDAIDefaults = {
-    model: { name: 'default-model', parameters: { name: 'default' } },
-    provider: { name: 'default-provider' },
-    messages: [{ role: 'system', content: 'Default messages' }],
-    enabled: true,
-  };
-
-  const expectedLDFlagValue = {
-    _ldMeta: { enabled: true },
-    model: defaultValue.model,
-    messages: defaultValue.messages,
-    provider: defaultValue.provider,
-  };
-
-  mockLdClient.variation.mockResolvedValue(expectedLDFlagValue);
-
-  const result = await client.config(key, testContext, defaultValue);
-
-  expect(result).toEqual({
-    model: defaultValue.model,
-    messages: defaultValue.messages,
-    provider: defaultValue.provider,
-    tracker: expect.any(Object),
-    enabled: defaultValue.enabled,
-    toVercelAISDK: expect.any(Function),
+    const variables = { name: 'John', score: 42 };
+    const evaluateSpy = jest.spyOn(client as any, '_evaluate');
+    const result = await client.agentConfig(key, testContext, defaultValue, variables);
+
+    expect(evaluateSpy).toHaveBeenCalledWith(key, testContext, defaultValue, 'agent', variables);
+    expect(result.instructions).toBe(
+      'You are a helpful assistant. Your name is John and your score is 42',
+    );
+    expect(result.tracker).toBeDefined();
+    expect(result.enabled).toBe(true);
+    evaluateSpy.mockRestore();
   });
 
-  expect(mockLdClient.variation).toHaveBeenCalledWith(key, testContext, expectedLDFlagValue);
-});
+  it('evaluates judge config successfully', async () => {
+    const client = new LDAIClientImpl(mockLdClient);
+    const key = 'test-judge';
+    const defaultValue: LDAIJudgeConfigDefault = {
+      enabled: false,
+    };
 
-// New agent-related tests
-it('returns single agent config with interpolated instructions', async () => {
-  const client = new LDAIClientImpl(mockLdClient);
-  const key = 'test-agent';
-  const defaultValue: LDAIAgentDefaults = {
-    model: { name: 'test', parameters: { name: 'test-model' } },
-    instructions: 'You are a helpful assistant.',
-    enabled: true,
-  };
-
-  const mockVariation = {
-    model: {
-      name: 'example-model',
-      parameters: { name: 'imagination', temperature: 0.7, maxTokens: 4096 },
-    },
-    provider: {
-      name: 'example-provider',
-    },
-    instructions: 'You are a helpful assistant. Your name is {{name}} and your score is {{score}}',
-    _ldMeta: {
-      variationKey: 'v1',
+    const mockVariation = {
       enabled: true,
-      mode: 'agent',
-    },
-  };
-
-  mockLdClient.variation.mockResolvedValue(mockVariation);
-
-  const variables = { name: 'John', score: 42 };
-  const result = await client.agent(key, testContext, defaultValue, variables);
-
-  expect(result).toEqual({
-    model: {
-      name: 'example-model',
-      parameters: { name: 'imagination', temperature: 0.7, maxTokens: 4096 },
-    },
-    provider: {
-      name: 'example-provider',
-    },
-    instructions: 'You are a helpful assistant. Your name is John and your score is 42',
-    tracker: expect.any(Object),
-    enabled: true,
+      model: { name: 'gpt-4' },
+      provider: { name: 'openai' },
+      evaluationMetricKeys: ['relevance', 'accuracy'],
+      messages: [{ role: 'system', content: 'You are a judge.' }],
+      _ldMeta: {
+        variationKey: 'v1',
+        enabled: true,
+        mode: 'judge',
+      },
+    };
+
+    mockLdClient.variation.mockResolvedValue(mockVariation);
+
+    const evaluateSpy = jest.spyOn(client as any, '_evaluate');
+    const result = await client.judgeConfig(key, testContext, defaultValue);
+
+    expect(evaluateSpy).toHaveBeenCalledWith(key, testContext, defaultValue, 'judge', undefined);
+    expect(result.evaluationMetricKeys).toEqual(['relevance', 'accuracy']);
+    expect(result.tracker).toBeDefined();
+    expect(result.enabled).toBe(true);
+    evaluateSpy.mockRestore();
   });
 
-  // Verify tracking was called
-  expect(mockLdClient.track).toHaveBeenCalledWith(
-    '$ld:ai:agent:function:single',
-    testContext,
-    key,
-    1,
-  );
-});
+  it('handles mode mismatch by returning disabled config', async () => {
+    const client = new LDAIClientImpl(mockLdClient);
+    const key = 'test-flag';
+    const defaultValue: LDAICompletionConfigDefault = {
+      enabled: false,
+    };
 
-it('includes context in variables for agent instructions interpolation', async () => {
-  const client = new LDAIClientImpl(mockLdClient);
-  const key = 'test-agent';
-  const defaultValue: LDAIAgentDefaults = {
-    model: { name: 'test', parameters: { name: 'test-model' } },
-    instructions: 'You are a helpful assistant.',
-    enabled: true,
-  };
+    const mockVariation = {
+      model: { name: 'example-provider', parameters: { name: 'imagination' } },
+      messages: [{ role: 'system', content: 'Hello' }],
+      _ldMeta: { variationKey: 'v1', enabled: true, mode: 'agent' }, // Wrong mode
+    };
 
-  const mockVariation = {
-    instructions: 'You are a helpful assistant. Your user key is {{ldctx.key}}',
-    _ldMeta: { variationKey: 'v1', enabled: true, mode: 'agent' },
-  };
+    mockLdClient.variation.mockResolvedValue(mockVariation);
 
-  mockLdClient.variation.mockResolvedValue(mockVariation);
+    const result = await client.completionConfig(key, testContext, defaultValue);
 
-  const result = await client.agent(key, testContext, defaultValue);
+    expect(result.enabled).toBe(false);
+    expect(result.tracker).toBeUndefined();
+  });
 
-  expect(result.instructions).toBe('You are a helpful assistant. Your user key is test-user');
-});
+  it('handles missing metadata mode by defaulting to completion mode', async () => {
+    const client = new LDAIClientImpl(mockLdClient);
+    const key = 'test-flag';
+    const defaultValue: LDAICompletionConfigDefault = {
+      enabled: false,
+    };
+
+    const mockVariation = {
+      model: { name: 'example-provider', parameters: { name: 'imagination' } },
+      messages: [{ role: 'system', content: 'Hello' }],
+      // No _ldMeta - mode defaults to completion
+    };
 
-it('handles missing metadata in agent variation', async () => {
-  const client = new LDAIClientImpl(mockLdClient);
-  const key = 'test-agent';
-  const defaultValue: LDAIAgentDefaults = {
-    model: { name: 'test', parameters: { name: 'test-model' } },
-    instructions: 'You are a helpful assistant.',
-    enabled: true,
-  };
-
-  const mockVariation = {
-    model: { name: 'example-provider', parameters: { name: 'imagination' } },
-    instructions: 'Hello.',
-  };
-
-  mockLdClient.variation.mockResolvedValue(mockVariation);
-
-  const result = await client.agent(key, testContext, defaultValue);
-
-  expect(result).toEqual({
-    model: { name: 'example-provider', parameters: { name: 'imagination' } },
-    instructions: 'Hello.',
-    tracker: expect.any(Object),
-    enabled: false,
+    mockLdClient.variation.mockResolvedValue(mockVariation);
+
+    const result = await client.completionConfig(key, testContext, defaultValue);
+
+    expect(result.enabled).toBe(false);
+    expect(result.tracker).toBeDefined();
+    expect(result.messages).toEqual([{ role: 'system', content: 'Hello' }]);
+    expect(result.model).toEqual({ name: 'example-provider', parameters: { name: 'imagination' } });
   });
-});
 
-it('passes the default value to the underlying client for single agent', async () => {
-  const client = new LDAIClientImpl(mockLdClient);
-  const key = 'non-existent-agent';
-  const defaultValue: LDAIAgentDefaults = {
-    model: { name: 'default-model', parameters: { name: 'default' } },
-    provider: { name: 'default-provider' },
-    instructions: 'Default instructions',
-    enabled: true,
-  };
-
-  const expectedLDFlagValue = {
-    _ldMeta: { enabled: defaultValue.enabled },
-    model: defaultValue.model,
-    provider: defaultValue.provider,
-    instructions: defaultValue.instructions,
-  };
-
-  mockLdClient.variation.mockResolvedValue(expectedLDFlagValue);
-
-  const result = await client.agent(key, testContext, defaultValue);
-
-  expect(result).toEqual({
-    model: defaultValue.model,
-    instructions: defaultValue.instructions,
-    provider: defaultValue.provider,
-    tracker: expect.any(Object),
-    enabled: defaultValue.enabled,
+  it('uses default value when flag does not exist', async () => {
+    const client = new LDAIClientImpl(mockLdClient);
+    const key = 'non-existent-flag';
+    const defaultValue: LDAICompletionConfigDefault = {
+      model: { name: 'default-model', parameters: { name: 'default' } },
+      provider: { name: 'default-provider' },
+      messages: [{ role: 'system', content: 'Default messages' }],
+      enabled: true,
+    };
+
+    const expectedLDFlagValue = {
+      _ldMeta: { enabled: true, mode: 'completion', variationKey: '' },
+      model: defaultValue.model,
+      messages: defaultValue.messages,
+      provider: defaultValue.provider,
+    };
+
+    mockLdClient.variation.mockResolvedValue(expectedLDFlagValue);
+
+    const result = await client.completionConfig(key, testContext, defaultValue);
+
+    expect(result.model).toEqual(defaultValue.model);
+    expect(result.messages).toEqual(defaultValue.messages);
+    expect(result.provider).toEqual(defaultValue.provider);
+    expect(result.tracker).toBeDefined();
+    expect(result.enabled).toBe(defaultValue.enabled);
+    expect(mockLdClient.variation).toHaveBeenCalledWith(
+      key,
+      testContext,
+      expect.objectContaining({
+        model: defaultValue.model,
+        provider: defaultValue.provider,
+      }),
+    );
   });
+});
+
+describe('completionConfig method', () => {
+  it('calls _evaluate with correct parameters and tracks usage', async () => {
+    const client = new LDAIClientImpl(mockLdClient);
+    const key = 'test-flag';
+    const defaultValue: LDAICompletionConfigDefault = {
+      model: { name: 'test', parameters: { name: 'test-model' } },
+      messages: [],
+      enabled: true,
+    };
+    const variables = { var1: 'value1' };
 
-  expect(mockLdClient.variation).toHaveBeenCalledWith(key, testContext, expectedLDFlagValue);
+    const mockConfig = {
+      model: { name: 'test-model' },
+      messages: [],
+      tracker: {} as any,
+      enabled: true,
+    };
+
+    const evaluateSpy = jest.spyOn(client as any, '_evaluate');
+    evaluateSpy.mockResolvedValue(mockConfig);
+
+    const result = await client.completionConfig(key, testContext, defaultValue, variables);
+
+    expect(mockLdClient.track).toHaveBeenCalledWith(
+      '$ld:ai:config:function:single',
+      testContext,
+      key,
+      1,
+    );
+    expect(evaluateSpy).toHaveBeenCalledWith(
+      key,
+      testContext,
+      defaultValue,
+      'completion',
+      variables,
+    );
+    expect(result).toBeDefined();
+    evaluateSpy.mockRestore();
+  });
 });
 
-it('returns multiple agents config with interpolated instructions', async () => {
-  const client = new LDAIClientImpl(mockLdClient);
+describe('agentConfig method', () => {
+  it('calls _evaluate with correct parameters and tracks usage', async () => {
+    const client = new LDAIClientImpl(mockLdClient);
+    const key = 'test-agent';
+    const defaultValue: LDAIAgentConfigDefault = {
+      model: { name: 'test', parameters: { name: 'test-model' } },
+      instructions: 'You are a helpful assistant.',
+      enabled: true,
+    };
+    const variables = { var1: 'value1' };
+
+    const mockConfig = {
+      model: { name: 'test-model' },
+      instructions: 'You are a helpful assistant.',
+      tracker: {} as any,
+      enabled: true,
+    };
+
+    const evaluateSpy = jest.spyOn(client as any, '_evaluate');
+    evaluateSpy.mockResolvedValue(mockConfig);
+
+    const result = await client.agentConfig(key, testContext, defaultValue, variables);
+
+    expect(mockLdClient.track).toHaveBeenCalledWith(
+      '$ld:ai:agent:function:single',
+      testContext,
+      key,
+      1,
+    );
+    expect(evaluateSpy).toHaveBeenCalledWith(key, testContext, defaultValue, 'agent', variables);
+    expect(result).toBe(mockConfig);
+    evaluateSpy.mockRestore();
+  });
+});
 
-  const agentConfigs = [
-    {
-      key: 'research-agent',
-      defaultValue: {
-        model: { name: 'test', parameters: { name: 'test-model' } },
-        instructions: 'You are a research assistant.',
+describe('agents method', () => {
+  it('retrieves multiple agent configs with interpolated instructions', async () => {
+    const client = new LDAIClientImpl(mockLdClient);
+
+    const agentConfigs = [
+      {
+        key: 'research-agent',
+        defaultValue: {
+          enabled: false,
+        },
+        variables: { topic: 'climate change' },
+      },
+      {
+        key: 'writing-agent',
+        defaultValue: {
+          enabled: false,
+        },
+        variables: { style: 'academic' },
+      },
+    ] as const;
+
+    const mockVariations = {
+      'research-agent': {
+        model: {
+          name: 'research-model',
+          parameters: { temperature: 0.3, maxTokens: 2048 },
+        },
+        provider: { name: 'openai' },
+        instructions: 'You are a research assistant specializing in {{topic}}.',
+        _ldMeta: { variationKey: 'v1', enabled: true, mode: 'agent' },
+      },
+      'writing-agent': {
+        model: {
+          name: 'writing-model',
+          parameters: { temperature: 0.7, maxTokens: 1024 },
+        },
+        provider: { name: 'anthropic' },
+        instructions: 'You are a writing assistant with {{style}} style.',
+        _ldMeta: { variationKey: 'v2', enabled: true, mode: 'agent' },
+      },
+    };
+
+    mockLdClient.variation.mockImplementation((key) =>
+      Promise.resolve(mockVariations[key as keyof typeof mockVariations]),
+    );
+
+    const result = await client.agentConfigs(agentConfigs, testContext);
+
+    expect(result).toEqual({
+      'research-agent': {
+        model: {
+          name: 'research-model',
+          parameters: { temperature: 0.3, maxTokens: 2048 },
+        },
+        provider: { name: 'openai' },
+        instructions: 'You are a research assistant specializing in climate change.',
+        tracker: expect.any(Object),
         enabled: true,
       },
-      variables: { topic: 'climate change' },
-    },
-    {
-      key: 'writing-agent',
-      defaultValue: {
-        model: { name: 'test', parameters: { name: 'test-model' } },
-        instructions: 'You are a writing assistant.',
+      'writing-agent': {
+        model: {
+          name: 'writing-model',
+          parameters: { temperature: 0.7, maxTokens: 1024 },
+        },
+        provider: { name: 'anthropic' },
+        instructions: 'You are a writing assistant with academic style.',
+        tracker: expect.any(Object),
         enabled: true,
       },
-      variables: { style: 'academic' },
-    },
-  ] as const;
+    });
+
+    expect(mockLdClient.track).toHaveBeenCalledWith(
+      '$ld:ai:agent:function:multiple',
+      testContext,
+      agentConfigs.length,
+      agentConfigs.length,
+    );
+  });
 
-  const mockVariations = {
-    'research-agent': {
-      model: {
-        name: 'research-model',
-        parameters: { temperature: 0.3, maxTokens: 2048 },
-      },
-      provider: { name: 'openai' },
-      instructions: 'You are a research assistant specializing in {{topic}}.',
-      _ldMeta: { variationKey: 'v1', enabled: true, mode: 'agent' },
-    },
-    'writing-agent': {
-      model: {
-        name: 'writing-model',
-        parameters: { temperature: 0.7, maxTokens: 1024 },
-      },
-      provider: { name: 'anthropic' },
-      instructions: 'You are a writing assistant with {{style}} style.',
-      _ldMeta: { variationKey: 'v2', enabled: true, mode: 'agent' },
-    },
-  };
+  it('handles empty agent configs array', async () => {
+    const client = new LDAIClientImpl(mockLdClient);
 
-  mockLdClient.variation.mockImplementation((key) =>
-    Promise.resolve(mockVariations[key as keyof typeof mockVariations]),
-  );
+    const result = await client.agentConfigs([], testContext);
 
-  const result = await client.agents(agentConfigs, testContext);
+    expect(result).toEqual({});
 
-  expect(result).toEqual({
-    'research-agent': {
-      model: {
-        name: 'research-model',
-        parameters: { temperature: 0.3, maxTokens: 2048 },
-      },
-      provider: { name: 'openai' },
-      instructions: 'You are a research assistant specializing in climate change.',
-      tracker: expect.any(Object),
+    expect(mockLdClient.track).toHaveBeenCalledWith(
+      '$ld:ai:agent:function:multiple',
+      testContext,
+      0,
+      0,
+    );
+  });
+});
+
+describe('judgeConfig method', () => {
+  it('calls _evaluate with correct parameters and tracks usage', async () => {
+    const client = new LDAIClientImpl(mockLdClient);
+    const key = 'test-judge';
+    const defaultValue: LDAIJudgeConfigDefault = {
+      enabled: false,
+    };
+    const variables = { metric: 'relevance' };
+
+    const mockJudgeConfig = {
       enabled: true,
-    },
-    'writing-agent': {
-      model: {
-        name: 'writing-model',
-        parameters: { temperature: 0.7, maxTokens: 1024 },
-      },
-      provider: { name: 'anthropic' },
-      instructions: 'You are a writing assistant with academic style.',
-      tracker: expect.any(Object),
+      model: { name: 'gpt-4' },
+      provider: { name: 'openai' },
+      evaluationMetricKeys: ['relevance'],
+      messages: [{ role: 'system' as const, content: 'You are a judge for {{metric}}.' }],
+      tracker: {} as any,
+      toVercelAISDK: jest.fn(),
+    };
+
+    const evaluateSpy = jest.spyOn(client as any, '_evaluate');
+    evaluateSpy.mockResolvedValue(mockJudgeConfig);
+
+    const result = await client.judgeConfig(key, testContext, defaultValue, variables);
+
+    expect(mockLdClient.track).toHaveBeenCalledWith(
+      '$ld:ai:judge:function:single',
+      testContext,
+      key,
+      1,
+    );
+    expect(evaluateSpy).toHaveBeenCalledWith(key, testContext, defaultValue, 'judge', variables);
+    expect(result).toBe(mockJudgeConfig);
+    evaluateSpy.mockRestore();
+  });
+});
+
+describe('createJudge method', () => {
+  let mockProvider: jest.Mocked<any>;
+  let mockJudge: jest.Mocked<Judge>;
+
+  beforeEach(() => {
+    mockProvider = {
+      invokeStructuredModel: jest.fn(),
+    };
+
+    mockJudge = {
+      evaluate: jest.fn(),
+      evaluateMessages: jest.fn(),
+    } as any;
+
+    // Mock AIProviderFactory.create
+    (AIProviderFactory.create as jest.Mock).mockResolvedValue(mockProvider);
+
+    // Mock Judge constructor
+    (Judge as jest.MockedClass<typeof Judge>).mockImplementation(() => mockJudge);
+  });
+
+  it('initializes judge successfully', async () => {
+    const client = new LDAIClientImpl(mockLdClient);
+    const key = 'test-judge';
+    const defaultValue: LDAIJudgeConfigDefault = {
+      enabled: false,
+    };
+
+    const mockJudgeConfig = {
       enabled: true,
-    },
+      model: { name: 'gpt-4' },
+      provider: { name: 'openai' },
+      evaluationMetricKeys: ['relevance', 'accuracy'],
+      messages: [{ role: 'system' as const, content: 'You are a judge.' }],
+      tracker: {} as any,
+      toVercelAISDK: jest.fn(),
+    };
+
+    const judgeConfigSpy = jest.spyOn(client, 'judgeConfig');
+    judgeConfigSpy.mockResolvedValue(mockJudgeConfig);
+
+    const result = await client.createJudge(key, testContext, defaultValue);
+
+    expect(mockLdClient.track).toHaveBeenCalledWith(
+      '$ld:ai:judge:function:createJudge',
+      testContext,
+      key,
+      1,
+    );
+    expect(judgeConfigSpy).toHaveBeenCalledWith(key, testContext, defaultValue, {
+      message_history: '{{message_history}}',
+      response_to_evaluate: '{{response_to_evaluate}}',
+    });
+    expect(AIProviderFactory.create).toHaveBeenCalledWith(mockJudgeConfig, undefined, undefined);
+    expect(Judge).toHaveBeenCalledWith(
+      mockJudgeConfig,
+      mockJudgeConfig.tracker,
+      mockProvider,
+      undefined,
+    );
+    expect(result).toBe(mockJudge);
+    judgeConfigSpy.mockRestore();
   });
 
-  // Verify tracking was called
-  expect(mockLdClient.track).toHaveBeenCalledWith(
-    '$ld:ai:agent:function:multiple',
-    testContext,
-    agentConfigs.length,
-    agentConfigs.length,
-  );
-});
+  it('returns undefined when judge configuration is disabled', async () => {
+    const client = new LDAIClientImpl(mockLdClient);
+    const key = 'test-judge';
+    const defaultValue: LDAIJudgeConfigDefault = {
+      enabled: false,
+    };
+
+    const mockJudgeConfig = {
+      enabled: false,
+      evaluationMetricKeys: [],
+    };
 
-it('handles empty agent configs array', async () => {
-  const client = new LDAIClientImpl(mockLdClient);
+    const judgeConfigSpy = jest.spyOn(client, 'judgeConfig');
+    judgeConfigSpy.mockResolvedValue(mockJudgeConfig);
 
-  const result = await client.agents([], testContext);
+    const result = await client.createJudge(key, testContext, defaultValue);
 
-  expect(result).toEqual({});
+    expect(result).toBeUndefined();
+    expect(AIProviderFactory.create).not.toHaveBeenCalled();
+    expect(Judge).not.toHaveBeenCalled();
+    judgeConfigSpy.mockRestore();
+  });
+
+  it('returns undefined when AIProviderFactory.create fails', async () => {
+    const client = new LDAIClientImpl(mockLdClient);
+    const key = 'test-judge';
+    const defaultValue: LDAIJudgeConfigDefault = {
+      enabled: false,
+    };
 
-  // Verify tracking was called with 0 agents
-  expect(mockLdClient.track).toHaveBeenCalledWith(
-    '$ld:ai:agent:function:multiple',
-    testContext,
-    0,
-    0,
-  );
+    const mockJudgeConfig = {
+      enabled: true,
+      model: { name: 'gpt-4' },
+      provider: { name: 'openai' },
+      evaluationMetricKeys: ['relevance'],
+      messages: [{ role: 'system' as const, content: 'You are a judge.' }],
+      tracker: {} as any,
+      toVercelAISDK: jest.fn(),
+    };
+
+    const judgeConfigSpy = jest.spyOn(client, 'judgeConfig');
+    judgeConfigSpy.mockResolvedValue(mockJudgeConfig);
+
+    (AIProviderFactory.create as jest.Mock).mockResolvedValue(undefined);
+
+    const result = await client.createJudge(key, testContext, defaultValue);
+
+    expect(result).toBeUndefined();
+    expect(AIProviderFactory.create).toHaveBeenCalledWith(mockJudgeConfig, undefined, undefined);
+    expect(Judge).not.toHaveBeenCalled();
+    judgeConfigSpy.mockRestore();
+  });
+
+  it('handles errors gracefully', async () => {
+    const client = new LDAIClientImpl(mockLdClient);
+    const key = 'test-judge';
+    const defaultValue: LDAIJudgeConfigDefault = {
+      enabled: false,
+    };
+
+    const error = new Error('Judge configuration error');
+    const judgeConfigSpy = jest.spyOn(client, 'judgeConfig');
+    judgeConfigSpy.mockRejectedValue(error);
+
+    const result = await client.createJudge(key, testContext, defaultValue);
+
+    expect(result).toBeUndefined();
+    judgeConfigSpy.mockRestore();
+  });
 });