mirror of
https://github.com/snailyp/gemini-balance.git
synced 2026-07-06 07:11:30 +08:00
Compare commits
7 Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
5a44a76c48 | ||
|
|
7b5b6c7d4c | ||
|
|
68ed4da789 | ||
|
|
cdbca7ec62 | ||
|
|
48d58ef2e8 | ||
|
|
88d483c1ef | ||
|
|
8d48db026c |
@@ -3,10 +3,10 @@ FROM python:3.10-slim
|
|||||||
WORKDIR /app
|
WORKDIR /app
|
||||||
|
|
||||||
# 复制所需文件到容器中
|
# 复制所需文件到容器中
|
||||||
COPY ./app /app/app
|
|
||||||
COPY ./requirements.txt /app
|
COPY ./requirements.txt /app
|
||||||
|
|
||||||
RUN pip install --no-cache-dir -r requirements.txt
|
RUN pip install --no-cache-dir -r requirements.txt
|
||||||
|
COPY ./app /app/app
|
||||||
ENV API_KEYS='["your_api_key_1"]'
|
ENV API_KEYS='["your_api_key_1"]'
|
||||||
ENV ALLOWED_TOKENS='["your_token_1"]'
|
ENV ALLOWED_TOKENS='["your_token_1"]'
|
||||||
ENV BASE_URL=https://generativelanguage.googleapis.com/v1beta
|
ENV BASE_URL=https://generativelanguage.googleapis.com/v1beta
|
||||||
|
|||||||
@@ -63,7 +63,7 @@ async def generate_content(
|
|||||||
logger.info(f"Using API key: {api_key}")
|
logger.info(f"Using API key: {api_key}")
|
||||||
|
|
||||||
try:
|
try:
|
||||||
response = chat_service.generate_content(
|
response = await chat_service.generate_content(
|
||||||
model=model_name,
|
model=model_name,
|
||||||
request=request,
|
request=request,
|
||||||
api_key=api_key
|
api_key=api_key
|
||||||
@@ -122,7 +122,7 @@ async def verify_key(api_key: str):
|
|||||||
)
|
)
|
||||||
]
|
]
|
||||||
)
|
)
|
||||||
response = chat_service.generate_content(settings.TEST_MODEL,gemini_requset, api_key)
|
response = await chat_service.generate_content(settings.TEST_MODEL,gemini_requset, api_key)
|
||||||
if response:
|
if response:
|
||||||
return JSONResponse({"status": "valid"})
|
return JSONResponse({"status": "valid"})
|
||||||
return JSONResponse({"status": "invalid"})
|
return JSONResponse({"status": "invalid"})
|
||||||
|
|||||||
@@ -24,13 +24,13 @@ class GeminiApiClient(ApiClient):
|
|||||||
self.base_url = base_url
|
self.base_url = base_url
|
||||||
self.timeout = timeout
|
self.timeout = timeout
|
||||||
|
|
||||||
def generate_content(self, payload: Dict[str, Any], model: str, api_key: str) -> Dict[str, Any]:
|
async def generate_content(self, payload: Dict[str, Any], model: str, api_key: str) -> Dict[str, Any]:
|
||||||
timeout = httpx.Timeout(self.timeout, read=self.timeout)
|
timeout = httpx.Timeout(self.timeout, read=self.timeout)
|
||||||
if model.endswith("-search"):
|
if model.endswith("-search"):
|
||||||
model = model[:-7]
|
model = model[:-7]
|
||||||
with httpx.Client(timeout=timeout) as client:
|
async with httpx.AsyncClient(timeout=timeout) as client:
|
||||||
url = f"{self.base_url}/models/{model}:generateContent?key={api_key}"
|
url = f"{self.base_url}/models/{model}:generateContent?key={api_key}"
|
||||||
response = client.post(url, json=payload)
|
response = await client.post(url, json=payload)
|
||||||
if response.status_code != 200:
|
if response.status_code != 200:
|
||||||
error_content = response.text
|
error_content = response.text
|
||||||
raise Exception(f"API call failed with status code {response.status_code}, {error_content}")
|
raise Exception(f"API call failed with status code {response.status_code}, {error_content}")
|
||||||
|
|||||||
@@ -1,14 +1,16 @@
|
|||||||
# app/services/chat/message_converter.py
|
# app/services/chat/message_converter.py
|
||||||
|
|
||||||
from abc import ABC, abstractmethod
|
from abc import ABC, abstractmethod
|
||||||
from typing import List, Dict, Any
|
from typing import Any, Dict, List, Optional
|
||||||
|
|
||||||
|
SUPPORTED_ROLES = ["user", "model", "system"]
|
||||||
|
|
||||||
|
|
||||||
class MessageConverter(ABC):
|
class MessageConverter(ABC):
|
||||||
"""消息转换器基类"""
|
"""消息转换器基类"""
|
||||||
|
|
||||||
@abstractmethod
|
@abstractmethod
|
||||||
def convert(self, messages: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
|
def convert(self, messages: List[Dict[str, Any]]) -> tuple[List[Dict[str, Any]], Optional[Dict[str, Any]]]:
|
||||||
pass
|
pass
|
||||||
|
|
||||||
|
|
||||||
@@ -30,16 +32,26 @@ def _convert_image(image_url: str) -> Dict[str, Any]:
|
|||||||
class OpenAIMessageConverter(MessageConverter):
|
class OpenAIMessageConverter(MessageConverter):
|
||||||
"""OpenAI消息格式转换器"""
|
"""OpenAI消息格式转换器"""
|
||||||
|
|
||||||
def convert(self, messages: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
|
def convert(self, messages: List[Dict[str, Any]]) -> tuple[List[Dict[str, Any]], Optional[Dict[str, Any]]]:
|
||||||
converted_messages = []
|
converted_messages = []
|
||||||
for msg in messages:
|
system_instruction = None
|
||||||
role = "user" if msg["role"] == "user" else "model"
|
|
||||||
parts = []
|
|
||||||
|
|
||||||
if isinstance(msg["content"], str):
|
for idx, msg in enumerate(messages):
|
||||||
# 请求 gemini 接口时如果包含 content 字段但内容为空时会返回 400 错误,所以需要判断是否为空
|
role = msg.get("role", "")
|
||||||
if msg["content"]:
|
if role not in SUPPORTED_ROLES:
|
||||||
parts.append({"text": msg["content"]})
|
if role == "tool":
|
||||||
|
role = "user"
|
||||||
|
else:
|
||||||
|
# 如果是最后一条消息,则认为是用户消息
|
||||||
|
if idx == len(messages) - 1:
|
||||||
|
role = "user"
|
||||||
|
else:
|
||||||
|
role = "model"
|
||||||
|
|
||||||
|
parts = []
|
||||||
|
if isinstance(msg["content"], str) and msg["content"]:
|
||||||
|
# 请求 gemini 接口时如果包含 content 字段但内容为空时会返回 400 错误,所以需要判断是否为空并移除
|
||||||
|
parts.append({"text": msg["content"]})
|
||||||
elif isinstance(msg["content"], list):
|
elif isinstance(msg["content"], list):
|
||||||
for content in msg["content"]:
|
for content in msg["content"]:
|
||||||
if isinstance(content, str) and content:
|
if isinstance(content, str) and content:
|
||||||
@@ -50,6 +62,10 @@ class OpenAIMessageConverter(MessageConverter):
|
|||||||
elif content["type"] == "image_url":
|
elif content["type"] == "image_url":
|
||||||
parts.append(_convert_image(content["image_url"]["url"]))
|
parts.append(_convert_image(content["image_url"]["url"]))
|
||||||
|
|
||||||
converted_messages.append({"role": role, "parts": parts})
|
if parts:
|
||||||
|
if role == "system":
|
||||||
|
system_instruction = {"role": "system", "parts": parts}
|
||||||
|
else:
|
||||||
|
converted_messages.append({"role": role, "parts": parts})
|
||||||
|
|
||||||
return converted_messages
|
return converted_messages, system_instruction
|
||||||
|
|||||||
@@ -79,10 +79,10 @@ class GeminiChatService:
|
|||||||
self.key_manager = key_manager
|
self.key_manager = key_manager
|
||||||
self.response_handler = GeminiResponseHandler()
|
self.response_handler = GeminiResponseHandler()
|
||||||
|
|
||||||
def generate_content(self, model: str, request: GeminiRequest, api_key: str) -> Dict[str, Any]:
|
async def generate_content(self, model: str, request: GeminiRequest, api_key: str) -> Dict[str, Any]:
|
||||||
"""生成内容"""
|
"""生成内容"""
|
||||||
payload = _build_payload(model, request)
|
payload = _build_payload(model, request)
|
||||||
response = self.api_client.generate_content(payload, model, api_key)
|
response = await self.api_client.generate_content(payload, model, api_key)
|
||||||
return self.response_handler.handle_response(response, model, stream=False)
|
return self.response_handler.handle_response(response, model, stream=False)
|
||||||
|
|
||||||
async def stream_generate_content(self, model: str, request: GeminiRequest, api_key: str) -> AsyncGenerator[str, None]:
|
async def stream_generate_content(self, model: str, request: GeminiRequest, api_key: str) -> AsyncGenerator[str, None]:
|
||||||
|
|||||||
@@ -2,7 +2,7 @@
|
|||||||
|
|
||||||
from copy import deepcopy
|
from copy import deepcopy
|
||||||
import json
|
import json
|
||||||
from typing import Dict, Any, AsyncGenerator, List, Union
|
from typing import Dict, Any, AsyncGenerator, List, Optional, Union
|
||||||
from app.core.logger import get_openai_logger
|
from app.core.logger import get_openai_logger
|
||||||
from app.services.chat.message_converter import OpenAIMessageConverter
|
from app.services.chat.message_converter import OpenAIMessageConverter
|
||||||
from app.services.chat.response_handler import OpenAIResponseHandler
|
from app.services.chat.response_handler import OpenAIResponseHandler
|
||||||
@@ -57,7 +57,14 @@ def _build_tools(
|
|||||||
function_declarations.append(function)
|
function_declarations.append(function)
|
||||||
|
|
||||||
if function_declarations:
|
if function_declarations:
|
||||||
tools.append({"functionDeclarations": function_declarations})
|
# 按照 function 的 name 去重
|
||||||
|
names, functions = set(), []
|
||||||
|
for item in function_declarations:
|
||||||
|
if item.get("name") not in names:
|
||||||
|
names.add(item.get("name"))
|
||||||
|
functions.append(item)
|
||||||
|
|
||||||
|
tools.append({"functionDeclarations": functions})
|
||||||
|
|
||||||
return tools
|
return tools
|
||||||
|
|
||||||
@@ -87,10 +94,10 @@ def _get_safety_settings(model: str) -> List[Dict[str, str]]:
|
|||||||
|
|
||||||
|
|
||||||
def _build_payload(
|
def _build_payload(
|
||||||
request: ChatRequest, messages: List[Dict[str, Any]]
|
request: ChatRequest, messages: List[Dict[str, Any]], instruction: Optional[Dict[str, Any]] = None
|
||||||
) -> Dict[str, Any]:
|
) -> Dict[str, Any]:
|
||||||
"""构建请求payload"""
|
"""构建请求payload"""
|
||||||
return {
|
payload = {
|
||||||
"contents": messages,
|
"contents": messages,
|
||||||
"generationConfig": {
|
"generationConfig": {
|
||||||
"temperature": request.temperature,
|
"temperature": request.temperature,
|
||||||
@@ -103,6 +110,16 @@ def _build_payload(
|
|||||||
"safetySettings": _get_safety_settings(request.model),
|
"safetySettings": _get_safety_settings(request.model),
|
||||||
}
|
}
|
||||||
|
|
||||||
|
if (
|
||||||
|
instruction
|
||||||
|
and isinstance(instruction, dict)
|
||||||
|
and instruction.get("role") == "system"
|
||||||
|
and instruction.get("parts")
|
||||||
|
):
|
||||||
|
payload["systemInstruction"] = instruction
|
||||||
|
|
||||||
|
return payload
|
||||||
|
|
||||||
|
|
||||||
class OpenAIChatService:
|
class OpenAIChatService:
|
||||||
"""聊天服务"""
|
"""聊天服务"""
|
||||||
@@ -120,20 +137,20 @@ class OpenAIChatService:
|
|||||||
) -> Union[Dict[str, Any], AsyncGenerator[str, None]]:
|
) -> Union[Dict[str, Any], AsyncGenerator[str, None]]:
|
||||||
"""创建聊天完成"""
|
"""创建聊天完成"""
|
||||||
# 转换消息格式
|
# 转换消息格式
|
||||||
messages = self.message_converter.convert(request.messages)
|
messages, instruction = self.message_converter.convert(request.messages)
|
||||||
|
|
||||||
# 构建请求payload
|
# 构建请求payload
|
||||||
payload = _build_payload(request, messages)
|
payload = _build_payload(request, messages, instruction)
|
||||||
|
|
||||||
if request.stream:
|
if request.stream:
|
||||||
return self._handle_stream_completion(request.model, payload, api_key)
|
return self._handle_stream_completion(request.model, payload, api_key)
|
||||||
return self._handle_normal_completion(request.model, payload, api_key)
|
return await self._handle_normal_completion(request.model, payload, api_key)
|
||||||
|
|
||||||
def _handle_normal_completion(
|
async def _handle_normal_completion(
|
||||||
self, model: str, payload: Dict[str, Any], api_key: str
|
self, model: str, payload: Dict[str, Any], api_key: str
|
||||||
) -> Dict[str, Any]:
|
) -> Dict[str, Any]:
|
||||||
"""处理普通聊天完成"""
|
"""处理普通聊天完成"""
|
||||||
response = self.api_client.generate_content(payload, model, api_key)
|
response = await self.api_client.generate_content(payload, model, api_key)
|
||||||
return self.response_handler.handle_response(
|
return self.response_handler.handle_response(
|
||||||
response, model, stream=False, finish_reason="stop"
|
response, model, stream=False, finish_reason="stop"
|
||||||
)
|
)
|
||||||
|
|||||||
9
docker-compose.yml
Normal file
9
docker-compose.yml
Normal file
@@ -0,0 +1,9 @@
|
|||||||
|
version: '3'
|
||||||
|
|
||||||
|
services:
|
||||||
|
gemini-balance:
|
||||||
|
build: .
|
||||||
|
ports:
|
||||||
|
- "8000:8000"
|
||||||
|
env_file:
|
||||||
|
- .env
|
||||||
Reference in New Issue
Block a user