mirror of
https://github.com/snailyp/gemini-balance.git
synced 2026-09-05 23:56:37 +08:00
feat: add support for the n parameter in OpenAI-compatible requests. Now, when you make a request to the /v1/chat/completions endpoint with the n parameter, it will be correctly mapped to candidateCount in the Gemini API request, allowing you to receive multiple completions.
This commit is contained in:
@@ -12,6 +12,7 @@ class ChatRequest(BaseModel):
|
|||||||
max_tokens: Optional[int] = None
|
max_tokens: Optional[int] = None
|
||||||
top_p: Optional[float] = DEFAULT_TOP_P
|
top_p: Optional[float] = DEFAULT_TOP_P
|
||||||
top_k: Optional[int] = DEFAULT_TOP_K
|
top_k: Optional[int] = DEFAULT_TOP_K
|
||||||
|
n: Optional[int] = 1
|
||||||
stop: Optional[Union[List[str],str]] = None
|
stop: Optional[Union[List[str],str]] = None
|
||||||
reasoning_effort: Optional[str] = None
|
reasoning_effort: Optional[str] = None
|
||||||
tools: Optional[Union[List[Dict[str, Any]], Dict[str, Any]]] = []
|
tools: Optional[Union[List[Dict[str, Any]], Dict[str, Any]]] = []
|
||||||
|
|||||||
@@ -42,21 +42,35 @@ class GeminiResponseHandler(ResponseHandler):
|
|||||||
def _handle_openai_stream_response(
|
def _handle_openai_stream_response(
|
||||||
response: Dict[str, Any], model: str, finish_reason: str, usage_metadata: Optional[Dict[str, Any]]
|
response: Dict[str, Any], model: str, finish_reason: str, usage_metadata: Optional[Dict[str, Any]]
|
||||||
) -> Dict[str, Any]:
|
) -> Dict[str, Any]:
|
||||||
text, reasoning_content, tool_calls, _ = _extract_result(
|
choices = []
|
||||||
response, model, stream=True, gemini_format=False
|
candidates = response.get("candidates", [])
|
||||||
)
|
|
||||||
if not text and not tool_calls and not reasoning_content:
|
for candidate in candidates:
|
||||||
delta = {}
|
index = candidate.get("index", 0)
|
||||||
else:
|
text, reasoning_content, tool_calls, _ = _extract_result(
|
||||||
delta = {"content": text, "reasoning_content": reasoning_content, "role": "assistant"}
|
{"candidates": [candidate]}, model, stream=True, gemini_format=False
|
||||||
if tool_calls:
|
)
|
||||||
delta["tool_calls"] = tool_calls
|
|
||||||
|
if not text and not tool_calls and not reasoning_content:
|
||||||
|
delta = {}
|
||||||
|
else:
|
||||||
|
delta = {"content": text, "reasoning_content": reasoning_content, "role": "assistant"}
|
||||||
|
if tool_calls:
|
||||||
|
delta["tool_calls"] = tool_calls
|
||||||
|
|
||||||
|
choice = {
|
||||||
|
"index": index,
|
||||||
|
"delta": delta,
|
||||||
|
"finish_reason": finish_reason
|
||||||
|
}
|
||||||
|
choices.append(choice)
|
||||||
|
|
||||||
template_chunk = {
|
template_chunk = {
|
||||||
"id": f"chatcmpl-{uuid.uuid4()}",
|
"id": f"chatcmpl-{uuid.uuid4()}",
|
||||||
"object": "chat.completion.chunk",
|
"object": "chat.completion.chunk",
|
||||||
"created": int(time.time()),
|
"created": int(time.time()),
|
||||||
"model": model,
|
"model": model,
|
||||||
"choices": [{"index": 0, "delta": delta, "finish_reason": finish_reason}],
|
"choices": choices,
|
||||||
}
|
}
|
||||||
if usage_metadata:
|
if usage_metadata:
|
||||||
template_chunk["usage"] = {"prompt_tokens": usage_metadata.get("promptTokenCount", 0), "completion_tokens": usage_metadata.get("candidatesTokenCount",0), "total_tokens": usage_metadata.get("totalTokenCount", 0)}
|
template_chunk["usage"] = {"prompt_tokens": usage_metadata.get("promptTokenCount", 0), "completion_tokens": usage_metadata.get("candidatesTokenCount",0), "total_tokens": usage_metadata.get("totalTokenCount", 0)}
|
||||||
@@ -66,26 +80,31 @@ def _handle_openai_stream_response(
|
|||||||
def _handle_openai_normal_response(
|
def _handle_openai_normal_response(
|
||||||
response: Dict[str, Any], model: str, finish_reason: str, usage_metadata: Optional[Dict[str, Any]]
|
response: Dict[str, Any], model: str, finish_reason: str, usage_metadata: Optional[Dict[str, Any]]
|
||||||
) -> Dict[str, Any]:
|
) -> Dict[str, Any]:
|
||||||
text, reasoning_content, tool_calls, _ = _extract_result(
|
choices = []
|
||||||
response, model, stream=False, gemini_format=False
|
candidates = response.get("candidates", [])
|
||||||
)
|
|
||||||
|
for i, candidate in enumerate(candidates):
|
||||||
|
text, reasoning_content, tool_calls, _ = _extract_result(
|
||||||
|
{"candidates": [candidate]}, model, stream=False, gemini_format=False
|
||||||
|
)
|
||||||
|
choice = {
|
||||||
|
"index": i,
|
||||||
|
"message": {
|
||||||
|
"role": "assistant",
|
||||||
|
"content": text,
|
||||||
|
"reasoning_content": reasoning_content,
|
||||||
|
"tool_calls": tool_calls,
|
||||||
|
},
|
||||||
|
"finish_reason": finish_reason,
|
||||||
|
}
|
||||||
|
choices.append(choice)
|
||||||
|
|
||||||
return {
|
return {
|
||||||
"id": f"chatcmpl-{uuid.uuid4()}",
|
"id": f"chatcmpl-{uuid.uuid4()}",
|
||||||
"object": "chat.completion",
|
"object": "chat.completion",
|
||||||
"created": int(time.time()),
|
"created": int(time.time()),
|
||||||
"model": model,
|
"model": model,
|
||||||
"choices": [
|
"choices": choices,
|
||||||
{
|
|
||||||
"index": 0,
|
|
||||||
"message": {
|
|
||||||
"role": "assistant",
|
|
||||||
"content": text,
|
|
||||||
"reasoning_content": reasoning_content,
|
|
||||||
"tool_calls": tool_calls,
|
|
||||||
},
|
|
||||||
"finish_reason": finish_reason,
|
|
||||||
}
|
|
||||||
],
|
|
||||||
"usage": {"prompt_tokens": usage_metadata.get("promptTokenCount", 0), "completion_tokens": usage_metadata.get("candidatesTokenCount",0), "total_tokens": usage_metadata.get("totalTokenCount", 0)},
|
"usage": {"prompt_tokens": usage_metadata.get("promptTokenCount", 0), "completion_tokens": usage_metadata.get("candidatesTokenCount",0), "total_tokens": usage_metadata.get("totalTokenCount", 0)},
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|||||||
@@ -196,6 +196,10 @@ def _build_payload(
|
|||||||
|
|
||||||
# 处理 max_tokens 参数
|
# 处理 max_tokens 参数
|
||||||
_validate_and_set_max_tokens(payload, request.max_tokens, logger)
|
_validate_and_set_max_tokens(payload, request.max_tokens, logger)
|
||||||
|
|
||||||
|
# 处理 n 参数
|
||||||
|
if request.n is not None and request.n > 0:
|
||||||
|
payload["generationConfig"]["candidateCount"] = request.n
|
||||||
|
|
||||||
if request.model.endswith("-image") or request.model.endswith("-image-generation"):
|
if request.model.endswith("-image") or request.model.endswith("-image-generation"):
|
||||||
payload["generationConfig"]["responseModalities"] = ["Text", "Image"]
|
payload["generationConfig"]["responseModalities"] = ["Text", "Image"]
|
||||||
|
|||||||
Reference in New Issue
Block a user