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:
zenyanbo
2025-08-11 17:39:18 +08:00
parent a6558b4668
commit f58ae2b340
3 changed files with 49 additions and 25 deletions
+1
View File
@@ -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]]] = []
+44 -25
View File
@@ -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)},
} }
+4
View File
@@ -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"]