mirror of
https://github.com/Wei-Shaw/claude-relay-service.git
synced 2026-01-23 00:53:33 +00:00
feat: 完善 Gemini 功能与 Claude 保持一致
- 添加 Gemini 账户的 schedulable 字段和调度开关 API - 实现 Gemini 调度器的模型过滤功能 - 完善 Gemini 数据统计,记录 token 使用量 - 修复 Gemini 流式响应的 SSE 解析和 AbortController 支持 - 在教程页面和 README 中添加 Gemini CLI 环境变量说明 - 修复前端 Gemini 账户调度开关限制 🤖 Generated with [Claude Code](https://claude.ai/code) Co-Authored-By: Claude <noreply@anthropic.com>
This commit is contained in:
@@ -3,7 +3,7 @@ const { HttpsProxyAgent } = require('https-proxy-agent');
|
||||
const { SocksProxyAgent } = require('socks-proxy-agent');
|
||||
const logger = require('../utils/logger');
|
||||
const config = require('../../config/config');
|
||||
const { recordUsageMetrics } = require('./apiKeyService');
|
||||
const apiKeyService = require('./apiKeyService');
|
||||
|
||||
// Gemini API 配置
|
||||
const GEMINI_API_BASE = 'https://cloudcode.googleapis.com/v1';
|
||||
@@ -123,7 +123,7 @@ function convertGeminiResponse(geminiResponse, model, stream = false) {
|
||||
}
|
||||
|
||||
// 处理流式响应
|
||||
async function* handleStreamResponse(response, model, apiKeyId) {
|
||||
async function* handleStreamResponse(response, model, apiKeyId, accountId = null) {
|
||||
let buffer = '';
|
||||
let totalUsage = {
|
||||
promptTokenCount: 0,
|
||||
@@ -168,10 +168,16 @@ async function* handleStreamResponse(response, model, apiKeyId) {
|
||||
if (data.candidates?.[0]?.finishReason === 'STOP') {
|
||||
// 记录使用量
|
||||
if (apiKeyId && totalUsage.totalTokenCount > 0) {
|
||||
await recordUsageMetrics(apiKeyId, {
|
||||
inputTokens: totalUsage.promptTokenCount,
|
||||
outputTokens: totalUsage.candidatesTokenCount,
|
||||
model: model
|
||||
await apiKeyService.recordUsage(
|
||||
apiKeyId,
|
||||
totalUsage.promptTokenCount || 0, // inputTokens
|
||||
totalUsage.candidatesTokenCount || 0, // outputTokens
|
||||
0, // cacheCreateTokens (Gemini 没有这个概念)
|
||||
0, // cacheReadTokens (Gemini 没有这个概念)
|
||||
model,
|
||||
accountId
|
||||
).catch(error => {
|
||||
logger.error('❌ Failed to record Gemini usage:', error);
|
||||
});
|
||||
}
|
||||
|
||||
@@ -206,13 +212,18 @@ async function* handleStreamResponse(response, model, apiKeyId) {
|
||||
|
||||
yield 'data: [DONE]\n\n';
|
||||
} catch (error) {
|
||||
logger.error('Stream processing error:', error);
|
||||
yield `data: ${JSON.stringify({
|
||||
error: {
|
||||
message: error.message,
|
||||
type: 'stream_error'
|
||||
}
|
||||
})}\n\n`;
|
||||
// 检查是否是请求被中止
|
||||
if (error.name === 'CanceledError' || error.code === 'ECONNABORTED') {
|
||||
logger.info('Stream request was aborted by client');
|
||||
} else {
|
||||
logger.error('Stream processing error:', error);
|
||||
yield `data: ${JSON.stringify({
|
||||
error: {
|
||||
message: error.message,
|
||||
type: 'stream_error'
|
||||
}
|
||||
})}\n\n`;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -226,8 +237,10 @@ async function sendGeminiRequest({
|
||||
accessToken,
|
||||
proxy,
|
||||
apiKeyId,
|
||||
signal,
|
||||
projectId,
|
||||
location = 'us-central1'
|
||||
location = 'us-central1',
|
||||
accountId = null
|
||||
}) {
|
||||
// 确保模型名称格式正确
|
||||
if (!model.startsWith('models/')) {
|
||||
@@ -281,6 +294,12 @@ async function sendGeminiRequest({
|
||||
logger.debug('Using proxy for Gemini request');
|
||||
}
|
||||
|
||||
// 添加 AbortController 信号支持
|
||||
if (signal) {
|
||||
axiosConfig.signal = signal;
|
||||
logger.debug('AbortController signal attached to request');
|
||||
}
|
||||
|
||||
if (stream) {
|
||||
axiosConfig.responseType = 'stream';
|
||||
}
|
||||
@@ -290,23 +309,42 @@ async function sendGeminiRequest({
|
||||
const response = await axios(axiosConfig);
|
||||
|
||||
if (stream) {
|
||||
return handleStreamResponse(response, model, apiKeyId);
|
||||
return handleStreamResponse(response, model, apiKeyId, accountId);
|
||||
} else {
|
||||
// 非流式响应
|
||||
const openaiResponse = convertGeminiResponse(response.data, model, false);
|
||||
|
||||
// 记录使用量
|
||||
if (apiKeyId && openaiResponse.usage) {
|
||||
await recordUsageMetrics(apiKeyId, {
|
||||
inputTokens: openaiResponse.usage.prompt_tokens,
|
||||
outputTokens: openaiResponse.usage.completion_tokens,
|
||||
model: model
|
||||
await apiKeyService.recordUsage(
|
||||
apiKeyId,
|
||||
openaiResponse.usage.prompt_tokens || 0,
|
||||
openaiResponse.usage.completion_tokens || 0,
|
||||
0, // cacheCreateTokens
|
||||
0, // cacheReadTokens
|
||||
model,
|
||||
accountId
|
||||
).catch(error => {
|
||||
logger.error('❌ Failed to record Gemini usage:', error);
|
||||
});
|
||||
}
|
||||
|
||||
return openaiResponse;
|
||||
}
|
||||
} catch (error) {
|
||||
// 检查是否是请求被中止
|
||||
if (error.name === 'CanceledError' || error.code === 'ECONNABORTED') {
|
||||
logger.info('Gemini request was aborted by client');
|
||||
throw {
|
||||
status: 499,
|
||||
error: {
|
||||
message: 'Request canceled by client',
|
||||
type: 'canceled',
|
||||
code: 'request_canceled'
|
||||
}
|
||||
};
|
||||
}
|
||||
|
||||
logger.error('Gemini API request failed:', error.response?.data || error.message);
|
||||
|
||||
// 转换错误格式
|
||||
|
||||
Reference in New Issue
Block a user