#大模型可用性测试

// 调用脚本的命令
// node test-nora-agent.js \
//   --api-key 'xxx' \
//   --base-url 'http://25.42.245.59/xlm-gateway-yuuvsm/sfm-api-gateway/gateway/compatible-mode/v1' \
//   --model 'rsv-fnjhvprfssmhnjjgdzoky-pozyhc'

const { ChatOpenAI } = require('@langchain/openai');
const { tool } = require('langchain');
const { z } = require('zod');

const DEFAULT_SYSTEM_PROMPT = [
  '你叫 Nora，你是一个专业的智能数据分析助手。',
  '如果问题属于数据查询或数据分析，你可以调用工具。',
  '在调用工具前不要编造任何数字。',
].join('\n');

const DEFAULT_MESSAGE = '请简单自我介绍，并说明你在需要时可以调用工具。';

function parseArgs(argv) {
  const args = new Map();

  for (let i = 0; i < argv.length; i += 1) {
    const current = argv[i];
    if (!current.startsWith('--')) {
      continue;
    }
    const key = current.slice(2);
    const next = argv[i + 1];
    if (!next || next.startsWith('--')) {
      args.set(key, 'true');
      continue;
    }
    args.set(key, next);
    i += 1;
  }

  if (args.has('help')) {
    printHelp();
    process.exit(0);
  }

  const apiKey =
    args.get('api-key') ||
    process.env.NORA_TEST_API_KEY ||
    process.env.OPENAI_API_KEY;
  const baseURL =
    args.get('base-url') ||
    process.env.NORA_TEST_BASE_URL ||
    process.env.OPENAI_BASE_URL;
  const model =
    args.get('model') ||
    process.env.NORA_TEST_MODEL ||
    process.env.OPENAI_MODEL;

  if (!apiKey || !baseURL || !model) {
    printHelp();
    throw new Error('Missing required args: --api-key, --base-url, --model');
  }

  const only = args.get('only');
  if (
    only &&
    !['basic-invoke', 'invoke-with-tools', 'stream-with-tools'].includes(only)
  ) {
    throw new Error(`Invalid --only value: ${only}`);
  }

  return {
    apiKey,
    baseURL,
    model,
    message:
      args.get('message') || process.env.NORA_TEST_MESSAGE || DEFAULT_MESSAGE,
    system:
      args.get('system') ||
      process.env.NORA_TEST_SYSTEM ||
      DEFAULT_SYSTEM_PROMPT,
    timeout: Number(
      args.get('timeout') || process.env.NORA_TEST_TIMEOUT || 60000,
    ),
    maxRetries: Number(
      args.get('max-retries') || process.env.NORA_TEST_MAX_RETRIES || 1,
    ),
    temperature: Number(
      args.get('temperature') || process.env.NORA_TEST_TEMPERATURE || 0.1,
    ),
    topP: Number(args.get('top-p') || process.env.NORA_TEST_TOP_P || 1),
    presencePenalty: Number(
      args.get('presence-penalty') ||
        process.env.NORA_TEST_PRESENCE_PENALTY ||
        0,
    ),
    frequencyPenalty: Number(
      args.get('frequency-penalty') ||
        process.env.NORA_TEST_FREQUENCY_PENALTY ||
        0,
    ),
    responseFormat:
      args.get('response-format') ||
      process.env.NORA_TEST_RESPONSE_FORMAT ||
      'text',
    only,
  };
}

function printHelp() {
  const lines = [
    'Usage:',
    '  node test-nora-agent.js --api-key <key> --base-url <url> --model <model> [options]',
    '',
    'Options:',
    '  --message <text>             User message for the test',
    '  --system <text>              System prompt for the test',
    '  --timeout <ms>               Request timeout, default 60000',
    '  --max-retries <n>            Client retries, default 1',
    '  --temperature <n>            Default 0.1',
    '  --top-p <n>                  Default 1',
    '  --presence-penalty <n>       Default 0',
    '  --frequency-penalty <n>      Default 0',
    '  --response-format <type>     text | json_object | json_schema',
    '  --only <test>                basic-invoke | invoke-with-tools | stream-with-tools',
    '',
    'Env fallbacks:',
    '  NORA_TEST_API_KEY / OPENAI_API_KEY',
    '  NORA_TEST_BASE_URL / OPENAI_BASE_URL',
    '  NORA_TEST_MODEL / OPENAI_MODEL',
  ];
  console.log(lines.join('\n'));
}

function createModel(config) {
  return new ChatOpenAI({
    apiKey: config.apiKey,
    model: config.model,
    timeout: config.timeout,
    maxRetries: config.maxRetries,
    temperature: config.temperature,
    topP: config.topP,
    presencePenalty: config.presencePenalty,
    frequencyPenalty: config.frequencyPenalty,
    modelKwargs: {
      response_format: {
        type: config.responseFormat,
      },
    },
    configuration: {
      baseURL: config.baseURL,
    },
  });
}

function createNoraLikeTools() {
  return [
    tool(
      async ({ query }) =>
        JSON.stringify({
          status: 'success',
          message: `Stub yiask-askAlisa executed with query: ${query}`,
        }),
      {
        name: 'yiask-askAlisa',
        description:
          'Nora uses this tool to delegate data query tasks to Alisa.',
        schema: z.object({
          query: z.string().describe('The user query to send to Alisa'),
        }),
      },
    ),
    tool(
      async ({ metric, dimension }) =>
        JSON.stringify({
          status: 'success',
          message: `Stub yiask-attribution executed for metric=${metric}, dimension=${dimension || ''}`,
        }),
      {
        name: 'yiask-attribution',
        description: 'Nora uses this tool for attribution analysis.',
        schema: z.object({
          metric: z.string().describe('Metric name'),
          dimension: z.string().optional().describe('Optional dimension name'),
        }),
      },
    ),
    tool(
      async ({ metric }) =>
        JSON.stringify({
          status: 'success',
          views: [`Stub metric view for ${metric}`],
        }),
      {
        name: 'yiask_sage-getMetricViews',
        description: 'Fetch metric views for a metric.',
        schema: z.object({
          metric: z.string().describe('Metric name'),
        }),
      },
    ),
    tool(
      async ({ metric }) =>
        JSON.stringify({
          status: 'success',
          dimensions: [`Stub dimension for ${metric}`],
        }),
      {
        name: 'yiask_sage-getMetricDimensions',
        description: 'Fetch metric dimensions for a metric.',
        schema: z.object({
          metric: z.string().describe('Metric name'),
        }),
      },
    ),
    tool(
      async ({ cron, task }) =>
        JSON.stringify({
          status: 'success',
          message: `Stub yiask-createSchedule executed with cron=${cron}, task=${task}`,
        }),
      {
        name: 'yiask-createSchedule',
        description: 'Create a schedule task.',
        schema: z.object({
          cron: z.string().describe('Cron expression'),
          task: z.string().describe('Task description'),
        }),
      },
    ),
  ];
}

function buildMessages(config) {
  return [
    { role: 'system', content: config.system },
    { role: 'user', content: config.message },
  ];
}

function formatError(error) {
  const lines = [
    `message=${error && error.message ? error.message : 'unknown error'}`,
    `status=${error && error.status !== undefined ? error.status : 'n/a'}`,
    `request_id=${error && (error.requestID || error.request_id) ? error.requestID || error.request_id : 'n/a'}`,
  ];

  if (error && error.headers) {
    lines.push(`headers=${JSON.stringify(error.headers)}`);
  }

  if (error && error.error) {
    lines.push(`body=${JSON.stringify(error.error)}`);
  } else if (error && error.response && error.response.data) {
    lines.push(`body=${JSON.stringify(error.response.data)}`);
  }

  return lines.join('\n');
}

function extractText(content) {
  if (typeof content === 'string') {
    return content;
  }

  if (Array.isArray(content)) {
    return content
      .map((item) => {
        if (typeof item === 'string') {
          return item;
        }
        if (item && typeof item === 'object' && 'text' in item) {
          return String(item.text || '');
        }
        return JSON.stringify(item);
      })
      .join('');
  }

  if (content == null) {
    return '';
  }

  return JSON.stringify(content);
}

async function runBasicInvoke(config) {
  const model = createModel(config);
  const messages = buildMessages(config);

  try {
    const response = await model.invoke(messages);
    return {
      name: 'basic-invoke',
      ok: true,
      summary: 'basic invoke succeeded',
      details: extractText(response.content).slice(0, 500),
    };
  } catch (error) {
    return {
      name: 'basic-invoke',
      ok: false,
      summary: 'basic invoke failed',
      details: formatError(error),
    };
  }
}

async function runInvokeWithTools(config) {
  const model = createModel(config);
  const messages = buildMessages(config);
  const tools = createNoraLikeTools();

  try {
    const chain = model.bindTools(tools);
    const response = await chain.invoke(messages);
    const toolCalls = Array.isArray(response.tool_calls)
      ? response.tool_calls.length
      : 0;
    return {
      name: 'invoke-with-tools',
      ok: true,
      summary: `invoke with tools succeeded, tool_calls=${toolCalls}`,
      details: extractText(response.content).slice(0, 500),
    };
  } catch (error) {
    return {
      name: 'invoke-with-tools',
      ok: false,
      summary: 'invoke with tools failed',
      details: formatError(error),
    };
  }
}

async function runStreamWithTools(config) {
  const model = createModel(config);
  const messages = buildMessages(config);
  const tools = createNoraLikeTools();

  try {
    const chain = model.bindTools(tools);
    const stream = await chain.streamEvents(messages, {
      version: 'v2',
      signal: new AbortController().signal,
    });

    let chunks = 0;
    let content = '';

    for await (const event of stream) {
      if (event.event !== 'on_chat_model_stream') {
        continue;
      }
      chunks += 1;
      const chunkContent = extractText(
        event.data && event.data.chunk ? event.data.chunk.content : '',
      );
      content += chunkContent;
    }

    return {
      name: 'stream-with-tools',
      ok: true,
      summary: `stream with tools succeeded, chunks=${chunks}`,
      details: content.slice(0, 500),
    };
  } catch (error) {
    return {
      name: 'stream-with-tools',
      ok: false,
      summary: 'stream with tools failed',
      details: formatError(error),
    };
  }
}

async function main() {
  const config = parseArgs(process.argv.slice(2));
  const tests = [];

  if (!config.only || config.only === 'basic-invoke') {
    tests.push(() => runBasicInvoke(config));
  }
  if (!config.only || config.only === 'invoke-with-tools') {
    tests.push(() => runInvokeWithTools(config));
  }
  if (!config.only || config.only === 'stream-with-tools') {
    tests.push(() => runStreamWithTools(config));
  }

  console.log('Nora agent connectivity test');
  console.log(`baseURL=${config.baseURL}`);
  console.log(`model=${config.model}`);
  console.log(`timeout=${config.timeout}`);
  console.log(`maxRetries=${config.maxRetries}`);
  console.log(`responseFormat=${config.responseFormat}`);
  console.log('');

  const results = [];
  for (const test of tests) {
    const result = await test();
    results.push(result);
    console.log(
      `[${result.ok ? 'PASS' : 'FAIL'}] ${result.name}: ${result.summary}`,
    );
    if (result.details) {
      console.log(result.details);
    }
    console.log('');
  }

  if (results.some((result) => !result.ok)) {
    process.exitCode = 1;
  }
}

main().catch((error) => {
  console.error('Fatal error');
  console.error(formatError(error));
  process.exit(1);
});