GithubBackup/gemini-balance
1
0
Fork 0
mirror of https://github.com/snailyp/gemini-balance.git synced 2026-09-08 01:06:37 +08:00
Code Issues Packages Projects Releases Wiki Activity

refactor: 修改模型检查逻辑以提高代码可读性

Browse Source
This commit is contained in:
yinpeng
2025-01-22 14:29:07 +08:00
parent 6ee3b485a4
commit 97e601a176
3 changed files with 154 additions and 120 deletions
Download Patch File Download Diff File Expand all files Collapse all files
app/services/chat/response_handler.py
+119 -87
View File
@@ -44,51 +44,67 @@ class GeminiResponseHandler(ResponseHandler):
candidate = response["candidates"][0] candidate = response["candidates"][0]
content = candidate.get("content", {}) content = candidate.get("content", {})
parts = content.get("parts", []) parts = content.get("parts", [])
if "thinking" in model: # if "thinking" in model:
if settings.SHOW_THINKING_PROCESS: # if settings.SHOW_THINKING_PROCESS:
if len(parts) == 1: # if len(parts) == 1:
if self.thinking_first: # if self.thinking_first:
self.thinking_first = False # self.thinking_first = False
self.thinking_status = True # self.thinking_status = True
text = "> thinking\n\n" + parts[0].get("text") # text = "> thinking\n\n" + parts[0].get("text")
else: # else:
text = parts[0].get("text") # text = parts[0].get("text")
if len(parts) == 2: # if len(parts) == 2:
self.thinking_status = False # self.thinking_status = False
if self.thinking_first: # if self.thinking_first:
self.thinking_first = False # self.thinking_first = False
text = ( # text = (
"> thinking\n\n" # "> thinking\n\n"
+ parts[0].get("text") # + parts[0].get("text")
+ "\n\n---\n> output\n\n" # + "\n\n---\n> output\n\n"
+ parts[1].get("text") # + parts[1].get("text")
) # )
else: # else:
text = ( # text = (
parts[0].get("text") # parts[0].get("text")
+ "\n\n---\n> output\n\n" # + "\n\n---\n> output\n\n"
+ parts[1].get("text") # + parts[1].get("text")
) # )
else: # else:
if len(parts) == 1: # if len(parts) == 1:
if self.thinking_first: # if self.thinking_first:
self.thinking_first = False # self.thinking_first = False
self.thinking_status = True # self.thinking_status = True
text = "" # text = ""
elif self.thinking_status: # elif self.thinking_status:
text = "" # text = ""
else: # else:
text = parts[0].get("text") # text = parts[0].get("text")
if len(parts) == 2: # if len(parts) == 2:
self.thinking_status = False # self.thinking_status = False
if self.thinking_first: # if self.thinking_first:
self.thinking_first = False # self.thinking_first = False
text = parts[1].get("text") # text = parts[1].get("text")
else: # else:
text = parts[1].get("text") # text = parts[1].get("text")
else: # else:
# if "text" in parts[0]:
# text = parts[0].get("text")
# elif "executableCode" in parts[0]:
# text = _format_code_block(parts[0]["executableCode"])
# elif "codeExecution" in parts[0]:
# text = _format_code_block(parts[0]["codeExecution"])
# elif "executableCodeResult" in parts[0]:
# text = _format_execution_result(
# parts[0]["executableCodeResult"]
# )
# elif "codeExecutionResult" in parts[0]:
# text = _format_execution_result(
# parts[0]["codeExecutionResult"]
# )
# else:
# text = ""
if "text" in parts[0]: if "text" in parts[0]:
text = parts[0].get("text") text = parts[0].get("text")
elif "executableCode" in parts[0]: elif "executableCode" in parts[0]:
@@ -167,7 +183,6 @@ class OpenAIResponseHandler(ResponseHandler):
}] }]
} }
def _handle_normal_response(self, response: Dict[str, Any], model: str, finish_reason: str) -> Dict[str, Any]: def _handle_normal_response(self, response: Dict[str, Any], model: str, finish_reason: str) -> Dict[str, Any]:
text = self._extract_text(response, model, stream=False) text = self._extract_text(response, model, stream=False)
return { return {
@@ -197,51 +212,68 @@ class OpenAIResponseHandler(ResponseHandler):
candidate = response["candidates"][0] candidate = response["candidates"][0]
content = candidate.get("content", {}) content = candidate.get("content", {})
parts = content.get("parts", []) parts = content.get("parts", [])
if "thinking" in model: # if "thinking" in model:
if settings.SHOW_THINKING_PROCESS: # if settings.SHOW_THINKING_PROCESS:
if len(parts) == 1: # if len(parts) == 1:
if self.thinking_first: # if self.thinking_first:
self.thinking_first = False # self.thinking_first = False
self.thinking_status = True # self.thinking_status = True
text = "> thinking\n\n" + parts[0].get("text") # text = "> thinking\n\n" + parts[0].get("text")
else: # else:
text = parts[0].get("text") # text = parts[0].get("text")
if len(parts) == 2: # if len(parts) == 2:
self.thinking_status = False # self.thinking_status = False
if self.thinking_first: # if self.thinking_first:
self.thinking_first = False # self.thinking_first = False
text = ( # text = (
"> thinking\n\n" # "> thinking\n\n"
+ parts[0].get("text") # + parts[0].get("text")
+ "\n\n---\n> output\n\n" # + "\n\n---\n> output\n\n"
+ parts[1].get("text") # + parts[1].get("text")
) # )
else: # else:
text = ( # text = (
parts[0].get("text") # parts[0].get("text")
+ "\n\n---\n> output\n\n" # + "\n\n---\n> output\n\n"
+ parts[1].get("text") # + parts[1].get("text")
) # )
else: # else:
if len(parts) == 1: # if len(parts) == 1:
if self.thinking_first: # if self.thinking_first:
self.thinking_first = False # self.thinking_first = False
self.thinking_status = True # self.thinking_status = True
text = "" # text = ""
elif self.thinking_status: # elif self.thinking_status:
text = "" # text = ""
else: # else:
text = parts[0].get("text") # text = parts[0].get("text")
if len(parts) == 2: # if len(parts) == 2:
self.thinking_status = False # self.thinking_status = False
if self.thinking_first: # if self.thinking_first:
self.thinking_first = False # self.thinking_first = False
text = parts[1].get("text") # text = parts[1].get("text")
else: # else:
text = parts[1].get("text") # text = parts[1].get("text")
else: # else:
# if "text" in parts[0]:
# text = parts[0].get("text")
# elif "executableCode" in parts[0]:
# text = _format_code_block(parts[0]["executableCode"])
# elif "codeExecution" in parts[0]:
# text = _format_code_block(parts[0]["codeExecution"])
# elif "executableCodeResult" in parts[0]:
# text = _format_execution_result(
# parts[0]["executableCodeResult"]
# )
# elif "codeExecutionResult" in parts[0]:
# text = _format_execution_result(
# parts[0]["codeExecutionResult"]
# )
# else:
# text = ""
# text = _add_search_link_text(model, candidate, text)
if "text" in parts[0]: if "text" in parts[0]:
text = parts[0].get("text") text = parts[0].get("text")
elif "executableCode" in parts[0]: elif "executableCode" in parts[0]:
app/services/gemini_chat_service.py
+2 -1
View File
@@ -32,6 +32,7 @@ class GeminiChatService:
while retries < max_retries: while retries < max_retries:
try: try:
async for line in self.api_client.stream_generate_content(payload, model, api_key): async for line in self.api_client.stream_generate_content(payload, model, api_key):
# print(line)
if line.startswith("data:"): if line.startswith("data:"):
line = line[6:] line = line[6:]
line = json.dumps(self.response_handler.handle_response(json.loads(line), model, stream=True)) line = json.dumps(self.response_handler.handle_response(json.loads(line), model, stream=True))
@@ -80,7 +81,7 @@ class GeminiChatService:
def _get_safety_settings(self, model: str) -> List[Dict[str, str]]: def _get_safety_settings(self, model: str) -> List[Dict[str, str]]:
"""获取安全设置""" """获取安全设置"""
if "2.0" in model and model != "gemini-2.0-flash-thinking-exp": if "2.0" in model and "gemini-2.0-flash-thinking-exp" not in model:
return [ return [
{"category": "HARM_CATEGORY_HARASSMENT", "threshold": "OFF"}, {"category": "HARM_CATEGORY_HARASSMENT", "threshold": "OFF"},
{"category": "HARM_CATEGORY_HATE_SPEECH", "threshold": "OFF"}, {"category": "HARM_CATEGORY_HATE_SPEECH", "threshold": "OFF"},
app/services/openai_chat_service.py
+2 -1
View File
@@ -63,6 +63,7 @@ class OpenAIChatService:
while retries < max_retries: while retries < max_retries:
try: try:
async for line in self.api_client.stream_generate_content(payload, model, api_key): async for line in self.api_client.stream_generate_content(payload, model, api_key):
# print(line)
if line.startswith("data:"): if line.startswith("data:"):
chunk = json.loads(line[6:]) chunk = json.loads(line[6:])
openai_chunk = self.response_handler.handle_response( openai_chunk = self.response_handler.handle_response(
@@ -127,7 +128,7 @@ class OpenAIChatService:
def _get_safety_settings(self, model: str) -> List[Dict[str, str]]: def _get_safety_settings(self, model: str) -> List[Dict[str, str]]:
"""获取安全设置""" """获取安全设置"""
if "2.0" in model and model != "gemini-2.0-flash-thinking-exp": if "2.0" in model and "gemini-2.0-flash-thinking-exp" not in model:
return [ return [
{"category": "HARM_CATEGORY_HARASSMENT", "threshold": "OFF"}, {"category": "HARM_CATEGORY_HARASSMENT", "threshold": "OFF"},
{"category": "HARM_CATEGORY_HATE_SPEECH", "threshold": "OFF"}, {"category": "HARM_CATEGORY_HATE_SPEECH", "threshold": "OFF"},