tool-calling.js

/**
 * Tool Calling Example - AetherGuard SDK
 *
 * Demonstrates the full tool-calling loop:
 *   1. Send a request with tool definitions
 *   2. Model responds with tool_calls
 *   3. Execute the tools locally
 *   4. Submit results back to get the final answer
 */

const {
  AetherGuardClient,
  createTool,
  hasToolCalls,
  extractToolCalls
} = require('@aetherguard/sdk');

// ─── Simulated tool implementations ──────────────────────────────────────────

function getWeather(location, unit = 'celsius') {
  // In a real app, this would call a weather API
  const data = {
    'London': { temp: 15, condition: 'cloudy' },
    'Tokyo': { temp: 28, condition: 'sunny' },
    'New York': { temp: 22, condition: 'partly cloudy' },
  };
  const weather = data[location] || { temp: 20, condition: 'unknown' };
  const temp = unit === 'fahrenheit' ? (weather.temp * 9/5) + 32 : weather.temp;
  return { location, temperature: temp, unit, condition: weather.condition };
}

function searchDatabase(query) {
  // Simulated database search
  return {
    results: [
      { id: 1, title: `Result for: ${query}`, relevance: 0.95 },
      { id: 2, title: `Related: ${query}`, relevance: 0.82 }
    ],
    total: 2
  };
}

// ─── Tool dispatcher ─────────────────────────────────────────────────────────

function executeTool(name, args) {
  switch (name) {
    case 'get_weather':
      return getWeather(args.location, args.unit);
    case 'search_database':
      return searchDatabase(args.query);
    default:
      return { error: `Unknown tool: ${name}` };
  }
}

// ─── Main example ────────────────────────────────────────────────────────────

async function toolCallingExample() {
  const client = new AetherGuardClient({
    apiKey: process.env.AETHERGUARD_API_KEY || 'your-api-key-here',
    baseUrl: process.env.AETHERGUARD_BASE_URL || 'http://localhost:8080'
  });

  // Define available tools
  const tools = [
    createTool('get_weather', 'Get current weather for a city', {
      type: 'object',
      properties: {
        location: { type: 'string', description: 'City name' },
        unit: { type: 'string', enum: ['celsius', 'fahrenheit'] }
      },
      required: ['location']
    }),
    createTool('search_database', 'Search the internal knowledge base', {
      type: 'object',
      properties: {
        query: { type: 'string', description: 'Search query' }
      },
      required: ['query']
    })
  ];

  const messages = [
    { role: 'system', content: 'You are a helpful assistant with access to weather data and a knowledge base.' },
    { role: 'user', content: 'What is the weather like in Tokyo and London? Also search for "AI safety" in our database.' }
  ];

  console.log('User:', messages[1].content);
  console.log('\nSending request with', tools.length, 'tools...\n');

  try {
    // Step 1: Initial request with tools
    const response = await client.createToolCompletion({
      model: 'gpt-4',
      messages,
      tools,
      tool_choice: 'auto'
    });

    if (!hasToolCalls(response)) {
      console.log('Model responded directly:', response.choices[0].message.content);
      return;
    }

    // Step 2: Extract and execute tool calls
    const calls = extractToolCalls(response);
    console.log(`Model requested ${calls.length} tool call(s):\n`);

    const toolOutputs = calls.map(tc => {
      const args = JSON.parse(tc.arguments);
      console.log(`  📞 ${tc.name}(${JSON.stringify(args)})`);

      const result = executeTool(tc.name, args);
      console.log(`  ✅ Result: ${JSON.stringify(result)}\n`);

      return { tool_call_id: tc.id, content: JSON.stringify(result) };
    });

    // Step 3: Submit tool results back
    console.log('Submitting tool results back to model...\n');

    const finalResponse = await client.submitToolOutputs(
      'gpt-4',
      [...messages, response.choices[0].message],
      toolOutputs
    );

    console.log('--- Final Answer ---');
    console.log(finalResponse.choices[0].message.content);
    console.log('\nTokens used:', finalResponse.usage.total_tokens);

  } catch (error) {
    console.error('Error:', error);
  }
}

toolCallingExample();