feat: 优化 Gemini 模型思考过程的展示

This commit is contained in:
yinpeng
2024-12-24 23:27:20 +08:00
parent 98ba46f779
commit 5a1c3bdbe7
+50 -7
View File
@@ -71,6 +71,7 @@ class ChatService:
def __init__(self, base_url: str, key_manager=None):
self.base_url = base_url
self.key_manager = key_manager
self.thinking_first = True
def convert_gemini_response_to_openai(
self,
@@ -82,11 +83,39 @@ class ChatService:
"""Convert Gemini response to OpenAI format"""
if stream:
try:
text = ""
if response.get("candidates"):
candidate = response["candidates"][0]
content = candidate.get("content", {})
parts = content.get("parts", [])
if "thinking" in model:
if len(parts) == 1:
if self.thinking_first:
self.thinking_first = False
text = "\n🤔 **思考过程** 🤔\n---\n```\n" + parts[
0
].get("text")
else:
text = parts[0].get("text")
elif len(parts) == 2:
if self.thinking_first:
self.thinking_first = False
text = (
"\n🤔 **思考过程** 🤔\n---\n```\n"
+ parts[0].get("text")
+ "\n```\n---\n"
+ parts[1].get("text")
)
else:
text = (
parts[0].get("text")
+ "\n```\n---\n"
+ parts[1].get("text")
)
else:
text = ""
else:
if "text" in parts[0]:
text = parts[0].get("text")
elif "executableCode" in parts[0]:
@@ -105,8 +134,6 @@ class ChatService:
text = ""
text = self.add_search_link_text(model, candidate, text)
else:
text = ""
return {
"id": f"chatcmpl-{uuid.uuid4()}",
@@ -136,7 +163,9 @@ class ChatService:
"index": 0,
"message": {
"role": "assistant",
"content": response["candidates"][0]["content"]["parts"][0]["text"],
"content": response["candidates"][0]["content"]["parts"][0][
"text"
],
},
"finish_reason": finish_reason,
}
@@ -149,8 +178,17 @@ class ChatService:
}
try:
if response.get("candidates"):
text = response["candidates"][0]["content"]["parts"][0]["text"]
candidate = response["candidates"][0]
if "thinking" in model:
text = (
"\n🤔 **思考过程** 🤔\n---\n```\n"
+ candidate["content"]["parts"][0]["text"]
+ "\n```\n---\n"
+ candidate["content"]["parts"][1]["text"]
)
else:
text = candidate["content"]["parts"][0]["text"]
text = self.add_search_link_text(model, candidate, text)
res["choices"][0]["message"]["content"] = text
return res
@@ -160,7 +198,9 @@ class ChatService:
except Exception as e:
logger.error(f"Error converting Gemini response: {str(e)}")
logger.debug(f"Raw response: {response}")
res["choices"][0]["message"]["content"] = f"Error converting Gemini response: {str(e)}"
res["choices"][0]["message"][
"content"
] = f"Error converting Gemini response: {str(e)}"
return res
def add_search_link_text(self, model, candidate, text):
@@ -172,7 +212,7 @@ class ChatService:
):
grounding_chunks = candidate["groundingMetadata"]["groundingChunks"]
text += "\n\n---\n\n"
text += f"**【引用来源】**\n\n"
text += "**【引用来源】**\n\n"
for _, grounding_chunk in enumerate(grounding_chunks, 1):
if "web" in grounding_chunk:
text += create_search_link(grounding_chunk["web"])
@@ -252,6 +292,7 @@ class ChatService:
}
if stream:
async def generate():
retries = 0
max_retries = 3
@@ -344,7 +385,9 @@ class ChatService:
return generate()
else:
try:
timeout = httpx.Timeout(300.0, read=300.0) # 连接超时300秒,读取超时300秒
timeout = httpx.Timeout(
300.0, read=300.0
) # 连接超时300秒,读取超时300秒
async with httpx.AsyncClient(timeout=timeout) as client:
url = f"https://generativelanguage.googleapis.com/v1beta/models/{gemini_model}:generateContent?key={api_key}"
response = await client.post(url, json=payload)