guanghulab/server/ftchat/services/vllm-proxy.js

/**
 * vLLM SSE streaming proxy
 * Replaces the DashScope-based ft-dashscope.js for mother model inference.
 * SSE streaming to local vLLM via SSH tunnel (localhost:8000).
 * NO system prompt injection.
 */
'use strict';

const https = require('https');
const http = require('http');

const VLLM_ENDPOINT = process.env.VLLM_ENDPOINT || 'http://localhost:8000';
const VLLM_MODEL = process.env.VLLM_MODEL || 'qwen2.5-7b-sft';

/**
 * Parse VLLM_ENDPOINT URL into components
 */
function parseEndpoint() {
  const url = new URL(VLLM_ENDPOINT);
  return {
    hostname: url.hostname,
    port: url.port || (url.protocol === 'https:' ? 443 : 80),
    protocol: url.protocol,
    useTls: url.protocol === 'https:'
  };
}

/**
 * Stream chat completion from vLLM (SSE) to the response stream.
 * 
 * @param {Array<{role:string,content:string}>} messages - Messages array
 * @param {object} options
 * @param {number} options.maxTokens - Max new tokens
 * @param {number} options.temperature - Sampling temperature
 * @param {AbortSignal} [options.signal] - Abort signal
 * @returns {Promise<string>} - The complete response text
 */
async function streamChat(messages, options = {}) {
  const {
    maxTokens = 1024,
    temperature = 0.7,
    signal = null
  } = options;

  const endpoint = parseEndpoint();
  const body = JSON.stringify({
    model: VLLM_MODEL,
    messages: messages,
    max_tokens: maxTokens,
    temperature: temperature,
    stream: true
  });

  const transport = endpoint.useTls ? https : http;

  return new Promise((resolve, reject) => {
    const req = transport.request({
      hostname: endpoint.hostname,
      port: endpoint.port,
      path: '/v1/chat/completions',
      method: 'POST',
      headers: {
        'Content-Type': 'application/json',
        'Content-Length': Buffer.byteLength(body)
      }
    }, (res) => {
      let buffer = '';
      let fullResponse = '';

      res.on('data', (chunk) => {
        buffer += chunk.toString();
        const lines = buffer.split('\n');
        buffer = lines.pop() || '';

        for (const line of lines) {
          const trimmed = line.trim();
          if (!trimmed || !trimmed.startsWith('data: ')) continue;
          const jsonStr = trimmed.slice(6);
          if (jsonStr === '[DONE]') continue;

          try {
            const parsed = JSON.parse(jsonStr);
            const choices = parsed.choices || [];
            for (const choice of choices) {
              const delta = choice.delta || {};
              const content = delta.content || '';
              fullResponse += content;
            }
          } catch (e) {
            // Skip malformed JSON lines
          }
        }
      });

      res.on('end', () => resolve(fullResponse));
      res.on('error', reject);
    });

    req.on('error', reject);
    if (signal) {
      signal.addEventListener('abort', () => req.destroy());
    }
    req.write(body);
    req.end();
  });
}

/**
 * Pipe SSE stream from vLLM directly to the HTTP response
 * Used for real-time chat where the browser consumes the stream
 */
function pipeChat(messages, res, options = {}) {
  const {
    maxTokens = 1024,
    temperature = 0.7,
    signal = null
  } = options;

  const endpoint = parseEndpoint();
  const body = JSON.stringify({
    model: VLLM_MODEL,
    messages: messages,
    max_tokens: maxTokens,
    temperature: temperature,
    stream: true
  });

  const transport = endpoint.useTls ? https : http;

  res.writeHead(200, {
    'Content-Type': 'text/event-stream',
    'Cache-Control': 'no-cache',
    'Connection': 'keep-alive',
    'X-Accel-Buffering': 'no'
  });

  const req = transport.request({
    hostname: endpoint.hostname,
    port: endpoint.port,
    path: '/v1/chat/completions',
    method: 'POST',
    headers: {
      'Content-Type': 'application/json',
      'Content-Length': Buffer.byteLength(body)
    }
  }, (vllmRes) => {
    let buffer = '';

    vllmRes.on('data', (chunk) => {
      buffer += chunk.toString();
      const lines = buffer.split('\n');
      buffer = lines.pop() || '';

      for (const line of lines) {
        const trimmed = line.trim();
        if (!trimmed || !trimmed.startsWith('data: ')) continue;
        const jsonStr = trimmed.slice(6);
        if (jsonStr === '[DONE]') {
          res.write('data: [DONE]\n\n');
          continue;
        }

        try {
          const parsed = JSON.parse(jsonStr);
          const choices = parsed.choices || [];
          for (const choice of choices) {
            const delta = choice.delta || {};
            const content = delta.content || '';
            const finishReason = choice.finish_reason;

            res.write(JSON.stringify({
              delta: content,
              finish_reason: finishReason
            }) + '\n');
          }
        } catch (e) {
          // Skip malformed lines
        }
      }
    });

    vllmRes.on('end', () => {
      res.write('data: [DONE]\n\n');
      res.end();
    });

    vllmRes.on('error', (err) => {
      res.write(JSON.stringify({ error: err.message }) + '\n');
      res.end();
    });
  });

  req.on('error', (err) => {
    res.write(JSON.stringify({ error: 'vLLM connection failed: ' + err.message }) + '\n');
    res.end();
  });

  if (signal) {
    signal.addEventListener('abort', () => {
      req.destroy();
      if (!res.writableEnded) res.end();
    });
  }

  req.write(body);
  req.end();
}

/**
 * Non-streaming chat completion (for memory compression etc.)
 */
async function chatOnce(messages, options = {}) {
  const {
    maxTokens = 512,
    temperature = 0.7
  } = options;

  const endpoint = parseEndpoint();
  const body = JSON.stringify({
    model: VLLM_MODEL,
    messages: messages,
    max_tokens: maxTokens,
    temperature: temperature,
    stream: false
  });

  const transport = endpoint.useTls ? https : http;

  return new Promise((resolve, reject) => {
    const req = transport.request({
      hostname: endpoint.hostname,
      port: endpoint.port,
      path: '/v1/chat/completions',
      method: 'POST',
      headers: {
        'Content-Type': 'application/json',
        'Content-Length': Buffer.byteLength(body)
      }
    }, (res) => {
      let data = '';
      res.on('data', (chunk) => data += chunk.toString());
      res.on('end', () => {
        try {
          const parsed = JSON.parse(data);
          const content = parsed?.choices?.[0]?.message?.content || '';
          resolve(content);
        } catch (e) {
          reject(new Error('Failed to parse vLLM response: ' + e.message));
        }
      });
      res.on('error', reject);
    });

    req.on('error', reject);
    req.write(body);
    req.end();
  });
}

function getStatus() {
  return {
    endpoint: VLLM_ENDPOINT,
    model: VLLM_MODEL
  };
}

module.exports = {
  streamChat,
  pipeChat,
  chatOnce,
  getStatus
};