github/claude-relay-service
1
0
Fork 0
mirror of https://github.com/Wei-Shaw/claude-relay-service.git synced 2026-01-23 00:53:33 +00:00
Code Issues Packages Projects Releases Wiki Activity

feat: 增强 Gemini 服务支持并添加统一调度器

Browse Source 
- 新增 unifiedGeminiScheduler.js 统一账户调度服务
- 增强 geminiRoutes.js 支持更多 Gemini API 端点
- 优化 geminiAccountService.js 账户管理和 token 刷新机制
- 添加对 v1internal 端点的完整支持（loadCodeAssist、onboardUser、countTokens、generateContent、streamGenerateContent）
- 改进错误处理和流式响应管理

🤖 Generated with [Claude Code](https://claude.ai/code)

Co-Authored-By: Claude <noreply@anthropic.com>
This commit is contained in:
千羽
2025-08-04 14:47:03 +09:00
parent 6d27dd7c94
commit 33837c23aa
3 changed files with 1100 additions and 128 deletions
Download Patch File Download Diff File Expand all files Collapse all files

318
src/routes/geminiRoutes.js
View File

@@ -5,6 +5,9 @@ const { authenticateApiKey } = require('../middleware/auth');
const geminiAccountService = require('../services/geminiAccountService');
const { sendGeminiRequest, getAvailableModels } = require('../services/geminiRelayService');
const crypto = require('crypto');
const sessionHelper = require('../utils/sessionHelper');
const unifiedGeminiScheduler = require('../services/unifiedGeminiScheduler');
const { OAuth2Client } = require('google-auth-library');
// 生成会话哈希
function generateSessionHash(req) {
@@ -13,7 +16,7 @@ function generateSessionHash(req) {
req.ip,
req.headers['x-api-key']?.substring(0, 10)
].filter(Boolean).join(':');
return crypto.createHash('sha256').update(sessionData).digest('hex');
}
@@ -27,10 +30,10 @@ function checkPermissions(apiKeyData, requiredPermission = 'gemini') {
router.post('/messages', authenticateApiKey, async (req, res) => {
const startTime = Date.now();
let abortController = null;
try {
const apiKeyData = req.apiKey;
// 检查权限
if (!checkPermissions(apiKeyData, 'gemini')) {
return res.status(403).json({
@@ -40,7 +43,7 @@ router.post('/messages', authenticateApiKey, async (req, res) => {
}
});
}
// 提取请求参数
const {
messages,
@@ -49,7 +52,7 @@ router.post('/messages', authenticateApiKey, async (req, res) => {
max_tokens = 4096,
stream = false
} = req.body;
// 验证必需参数
if (!messages || !Array.isArray(messages) || messages.length === 0) {
return res.status(400).json({
@@ -59,16 +62,16 @@ router.post('/messages', authenticateApiKey, async (req, res) => {
}
});
}
// 生成会话哈希用于粘性会话
const sessionHash = generateSessionHash(req);
// 选择可用的 Gemini 账户
const account = await geminiAccountService.selectAvailableAccount(
apiKeyData.id,
sessionHash
);
if (!account) {
return res.status(503).json({
error: {
@@ -77,15 +80,15 @@ router.post('/messages', authenticateApiKey, async (req, res) => {
}
});
}
logger.info(`Using Gemini account: ${account.id} for API key: ${apiKeyData.id}`);
// 标记账户被使用
await geminiAccountService.markAccountUsed(account.id);
// 创建中止控制器
abortController = new AbortController();
// 处理客户端断开连接
req.on('close', () => {
if (abortController && !abortController.signal.aborted) {
@@ -93,7 +96,7 @@ router.post('/messages', authenticateApiKey, async (req, res) => {
abortController.abort();
}
});
// 发送请求到 Gemini
const geminiResponse = await sendGeminiRequest({
messages,
@@ -107,14 +110,14 @@ router.post('/messages', authenticateApiKey, async (req, res) => {
signal: abortController.signal,
projectId: account.projectId
});
if (stream) {
// 设置流式响应头
res.setHeader('Content-Type', 'text/event-stream');
res.setHeader('Cache-Control', 'no-cache');
res.setHeader('Connection', 'keep-alive');
res.setHeader('X-Accel-Buffering', 'no');
// 流式传输响应
for await (const chunk of geminiResponse) {
if (abortController.signal.aborted) {
@@ -122,26 +125,26 @@ router.post('/messages', authenticateApiKey, async (req, res) => {
}
res.write(chunk);
}
res.end();
} else {
// 非流式响应
res.json(geminiResponse);
}
const duration = Date.now() - startTime;
logger.info(`Gemini request completed in ${duration}ms`);
} catch (error) {
logger.error('Gemini request error:', error);
// 处理速率限制
if (error.status === 429) {
if (req.apiKey && req.account) {
await geminiAccountService.setAccountRateLimited(req.account.id, true);
}
}
// 返回错误响应
const status = error.status || 500;
const errorResponse = {
@@ -150,7 +153,7 @@ router.post('/messages', authenticateApiKey, async (req, res) => {
type: 'api_error'
}
};
res.status(status).json(errorResponse);
} finally {
// 清理资源
@@ -164,7 +167,7 @@ router.post('/messages', authenticateApiKey, async (req, res) => {
router.get('/models', authenticateApiKey, async (req, res) => {
try {
const apiKeyData = req.apiKey;
// 检查权限
if (!checkPermissions(apiKeyData, 'gemini')) {
return res.status(403).json({
@@ -174,10 +177,10 @@ router.get('/models', authenticateApiKey, async (req, res) => {
}
});
}
// 选择账户获取模型列表
const account = await geminiAccountService.selectAvailableAccount(apiKeyData.id);
if (!account) {
// 返回默认模型列表
return res.json({
@@ -192,15 +195,15 @@ router.get('/models', authenticateApiKey, async (req, res) => {
]
});
}
// 获取模型列表
const models = await getAvailableModels(account.accessToken, account.proxy);
res.json({
object: 'list',
data: models
});
} catch (error) {
logger.error('Failed to get Gemini models:', error);
res.status(500).json({
@@ -216,7 +219,7 @@ router.get('/models', authenticateApiKey, async (req, res) => {
router.get('/usage', authenticateApiKey, async (req, res) => {
try {
const usage = req.apiKey.usage;
res.json({
object: 'usage',
total_tokens: usage.total.tokens,
@@ -241,14 +244,14 @@ router.get('/usage', authenticateApiKey, async (req, res) => {
router.get('/key-info', authenticateApiKey, async (req, res) => {
try {
const keyData = req.apiKey;
res.json({
id: keyData.id,
name: keyData.name,
permissions: keyData.permissions || 'all',
token_limit: keyData.tokenLimit,
tokens_used: keyData.usage.total.tokens,
tokens_remaining: keyData.tokenLimit > 0
tokens_remaining: keyData.tokenLimit > 0
? Math.max(0, keyData.tokenLimit - keyData.usage.total.tokens)
: null,
rate_limit: {
@@ -272,4 +275,259 @@ router.get('/key-info', authenticateApiKey, async (req, res) => {
}
});
router.post('/v1internal\\:loadCodeAssist', authenticateApiKey, async (req, res) => {
try {
const sessionHash = sessionHelper.generateSessionHash(req.body);
// 使用统一调度选择账号（传递请求的模型）
const requestedModel = req.body.model;
const { accountId } = await unifiedGeminiScheduler.selectAccountForApiKey(req.apiKey, sessionHash, requestedModel);
const { accessToken, refreshToken } = await geminiAccountService.getAccount(accountId);
logger.info(`accessToken: ${accessToken}`);
const { metadata, cloudaicompanionProject } = req.body;
logger.info('LoadCodeAssist request', {
metadata: metadata || {},
cloudaicompanionProject: cloudaicompanionProject || null,
apiKeyId: req.apiKey?.id || 'unknown'
});
const client = await geminiAccountService.getOauthClient(accessToken, refreshToken);
const response = await geminiAccountService.loadCodeAssist(client, cloudaicompanionProject);
res.json(response);
} catch (error) {
logger.error('Error in loadCodeAssist endpoint', { error: error.message });
res.status(500).json({
error: 'Internal server error',
message: error.message
});
}
});
router.post('/v1internal\\:onboardUser', authenticateApiKey, async (req, res) => {
try {
const { tierId, cloudaicompanionProject, metadata } = req.body;
const sessionHash = sessionHelper.generateSessionHash(req.body);
// 使用统一调度选择账号（传递请求的模型）
const requestedModel = req.body.model;
const { accountId } = await unifiedGeminiScheduler.selectAccountForApiKey(req.apiKey, sessionHash, requestedModel);
const { accessToken, refreshToken } = await geminiAccountService.getAccount(accountId);
logger.info('OnboardUser request', {
tierId: tierId || 'not provided',
cloudaicompanionProject: cloudaicompanionProject || null,
metadata: metadata || {},
apiKeyId: req.apiKey?.id || 'unknown'
});
const client = await geminiAccountService.getOauthClient(accessToken, refreshToken);
// 如果提供了完整参数，直接调用onboardUser
if (tierId && metadata) {
const response = await geminiAccountService.onboardUser(client, tierId, cloudaicompanionProject, metadata);
res.json(response);
} else {
// 否则执行完整的setupUser流程
const response = await geminiAccountService.setupUser(client, cloudaicompanionProject, metadata);
res.json(response);
}
} catch (error) {
logger.error('Error in onboardUser endpoint', { error: error.message });
res.status(500).json({
error: 'Internal server error',
message: error.message
});
}
});
router.post('/v1internal\\:countTokens', authenticateApiKey, async (req, res) => {
try {
// 处理请求体结构，支持直接 contents 或 request.contents
const requestData = req.body.request || req.body;
const { contents, model = 'gemini-2.0-flash-exp' } = requestData;
const sessionHash = sessionHelper.generateSessionHash(req.body);
// 验证必需参数
if (!contents || !Array.isArray(contents)) {
return res.status(400).json({
error: {
message: 'Contents array is required',
type: 'invalid_request_error'
}
});
}
// 使用统一调度选择账号
const { accountId } = await unifiedGeminiScheduler.selectAccountForApiKey(req.apiKey, sessionHash, model);
const { accessToken, refreshToken } = await geminiAccountService.getAccount(accountId);
logger.info('CountTokens request', {
model: model,
contentsLength: contents.length,
apiKeyId: req.apiKey?.id || 'unknown'
});
const client = await geminiAccountService.getOauthClient(accessToken, refreshToken);
const response = await geminiAccountService.countTokens(client, contents, model);
res.json(response);
} catch (error) {
logger.error('Error in countTokens endpoint', { error: error.message });
res.status(500).json({
error: {
message: error.message || 'Internal server error',
type: 'api_error'
}
});
}
});
router.post('/v1internal\\:generateContent', authenticateApiKey, async (req, res) => {
try {
const { model, project, user_prompt_id, request: requestData } = req.body;
const sessionHash = sessionHelper.generateSessionHash(req.body);
// 验证必需参数
if (!requestData || !requestData.contents) {
return res.status(400).json({
error: {
message: 'Request contents are required',
type: 'invalid_request_error'
}
});
}
// 使用统一调度选择账号
const { accountId } = await unifiedGeminiScheduler.selectAccountForApiKey(req.apiKey, sessionHash, model);
const account = await geminiAccountService.getAccount(accountId);
const { accessToken, refreshToken } = account;
logger.info('GenerateContent request', {
model: model,
userPromptId: user_prompt_id,
projectId: project || account.projectId,
apiKeyId: req.apiKey?.id || 'unknown'
});
const client = await geminiAccountService.getOauthClient(accessToken, refreshToken);
const response = await geminiAccountService.generateContent(
client,
{ model, request: requestData },
user_prompt_id,
project || account.projectId,
req.apiKey?.id // 使用 API Key ID 作为 session ID
);
res.json(response);
} catch (error) {
logger.error('Error in generateContent endpoint', { error: error.message });
res.status(500).json({
error: {
message: error.message || 'Internal server error',
type: 'api_error'
}
});
}
});
router.post('/v1internal\\:streamGenerateContent', authenticateApiKey, async (req, res) => {
let abortController = null;
try {
const { model, project, user_prompt_id, request: requestData } = req.body;
const sessionHash = sessionHelper.generateSessionHash(req.body);
// 验证必需参数
if (!requestData || !requestData.contents) {
return res.status(400).json({
error: {
message: 'Request contents are required',
type: 'invalid_request_error'
}
});
}
// 使用统一调度选择账号
const { accountId } = await unifiedGeminiScheduler.selectAccountForApiKey(req.apiKey, sessionHash, model);
const account = await geminiAccountService.getAccount(accountId);
const { accessToken, refreshToken } = account;
logger.info('StreamGenerateContent request', {
model: model,
userPromptId: user_prompt_id,
projectId: project || account.projectId,
apiKeyId: req.apiKey?.id || 'unknown'
});
// 创建中止控制器
abortController = new AbortController();
// 处理客户端断开连接
req.on('close', () => {
if (abortController && !abortController.signal.aborted) {
logger.info('Client disconnected, aborting stream request');
abortController.abort();
}
});
const client = await geminiAccountService.getOauthClient(accessToken, refreshToken);
const streamResponse = await geminiAccountService.generateContentStream(
client,
{ model, request: requestData },
user_prompt_id,
project || account.projectId,
req.apiKey?.id, // 使用 API Key ID 作为 session ID
abortController.signal // 传递中止信号
);
// 设置 SSE 响应头
res.setHeader('Content-Type', 'text/event-stream');
res.setHeader('Cache-Control', 'no-cache');
res.setHeader('Connection', 'keep-alive');
res.setHeader('X-Accel-Buffering', 'no');
// 直接管道转发流式响应，不进行额外处理
streamResponse.pipe(res, { end: false });
streamResponse.on('end', () => {
logger.info('Stream completed successfully');
res.end();
});
streamResponse.on('error', (error) => {
logger.error('Stream error:', error);
if (!res.headersSent) {
res.status(500).json({
error: {
message: error.message || 'Stream error',
type: 'api_error'
}
});
} else {
res.end();
}
});
} catch (error) {
logger.error('Error in streamGenerateContent endpoint', { error: error.message });
if (!res.headersSent) {
res.status(500).json({
error: {
message: error.message || 'Internal server error',
type: 'api_error'
}
});
}
} finally {
// 清理资源
if (abortController) {
abortController = null;
}
}
});
module.exports = router;