@nomyx/assistant/src/tests/promptRoutingTest.ts

A powerful assistant library and cli for your AI projects. works with Vertex AI (Claude and Gemini)

import { PromptAction, testPromptRouting } from '../prompts/PromptAction';
import { ExtendedActionContext } from '../prompts/types';
import { ActionSchema } from '../types/tool';
import { IAIProvider } from '../types/provider';
import { ProviderResponse } from '../types/chat';
import { StandardizedToolCall } from '../types/tool';

// Mock AI provider and logger
const mockAIProvider: IAIProvider = {
  chat: async (): Promise<ProviderResponse> => ({ content: 'Mock response', toolCalls: [] }),
  streamChat: async function* () { yield { content: 'Mock stream response', toolCalls: [] }; },
  getCapabilities: () => ({
    maxTokens: 1000,
    supportsFunctionCalling: true,
    supportsStreaming: true,
    supportedModels: ['mock-model'],
    maxSimultaneousCalls: 1,
    supportsSemanticCaching: false
  }),
  use: () => {},
  registerPlugin: () => {},
  convertToolSchema: () => ({}),
  convertToolCall: (): StandardizedToolCall => ({ name: 'mockTool', arguments: {} })
};

const mockLogger = {
  debug: console.log,
  info: console.log,
  warn: console.warn,
  error: console.error
};

// Create a sample PromptAction with routing
const sampleSchema: ActionSchema<any, any> = {
  input: {
    type: 'object',
    properties: {
      temperature: { type: 'number', description: 'The current temperature' }
    },
    required: ['temperature']
  },
  output: {
    type: 'object',
    properties: {
      recommendation: { type: 'string', description: 'A weather-based recommendation' }
    }
  }
};

const rootPrompt = new PromptAction({
  name: 'rootPrompt',
  schema: sampleSchema,
  systemContent: 'You are a helpful assistant.',
  userContent: 'Please provide a response.',
  routing: [
    {
      condition: 'params.temperature > 30',
      prompt: new PromptAction({
        name: 'hotWeatherPrompt',
        schema: sampleSchema,
        systemContent: 'You are a weather expert.',
        userContent: 'It\'s hot outside. What should we do?'
      })
    },
    {
      condition: 'params.temperature <= 30 && params.temperature > 10',
      prompt: new PromptAction({
        name: 'mildWeatherPrompt',
        schema: sampleSchema,
        systemContent: 'You are a weather expert.',
        userContent: 'The weather is mild. Any activity suggestions?'
      })
    },
    {
      condition: 'params.temperature <= 10',
      prompt: new PromptAction({
        name: 'coldWeatherPrompt',
        schema: sampleSchema,
        systemContent: 'You are a weather expert.',
        userContent: 'It\'s cold outside. How should we prepare?'
      })
    }
  ]
});

// Test function
function runRoutingTest() {
  const context: ExtendedActionContext = {
    aiProvider: mockAIProvider,
    logger: mockLogger,
    state: {},
    user: { id: 'test-user', name: 'Test User' },
    session: { id: 'test-session', timestamp: new Date().toISOString() },
    functionSchemas: [],
    executeAction: async () => ({})
  };

  console.log('Test 1: Hot weather');
  console.log(testPromptRouting(rootPrompt, { temperature: 35 }, context));

  console.log('\nTest 2: Mild weather');
  console.log(testPromptRouting(rootPrompt, { temperature: 20 }, context));

  console.log('\nTest 3: Cold weather');
  console.log(testPromptRouting(rootPrompt, { temperature: 5 }, context));

  console.log('\nTest 4: Edge case (exactly 30 degrees)');
  console.log(testPromptRouting(rootPrompt, { temperature: 30 }, context));
}

// Run the test
runRoutingTest();