fix: response parsing bug for gemini

This commit is contained in:
elf-mouse
2024-04-12 15:49:23 +08:00
parent 1f2b36a4f0
commit ee7306d216

View File

@@ -86,7 +86,7 @@ def _generate_response(prompt: str) -> str:
if llm_provider == "gemini": if llm_provider == "gemini":
import google.generativeai as genai import google.generativeai as genai
genai.configure(api_key=api_key) genai.configure(api_key=api_key, transport='rest')
generation_config = { generation_config = {
"temperature": 0.5, "temperature": 0.5,
@@ -118,10 +118,14 @@ def _generate_response(prompt: str) -> str:
generation_config=generation_config, generation_config=generation_config,
safety_settings=safety_settings) safety_settings=safety_settings)
convo = model.start_chat(history=[]) try:
response = model.generate_content(prompt)
candidates = response.candidates
generated_text = candidates[0].content.parts[0].text
except (AttributeError, IndexError) as e:
print("Gemini Error:", e)
convo.send_message(prompt) return generated_text
return convo.last.text
if llm_provider == "cloudflare": if llm_provider == "cloudflare":
import requests import requests