Update agent.py

This commit is contained in:
Yiheng Xu
2024-03-22 12:54:22 +08:00
committed by GitHub
parent d4e81afae7
commit 5f2802292a

View File

@@ -741,26 +741,30 @@ class PromptAgent:
assert api_key is not None, "Please set the GENAI_API_KEY environment variable"
genai.configure(api_key=api_key)
logger.info("Generating content with Gemini model: %s", self.model)
response = genai.GenerativeModel(self.model).generate_content(
gemini_messages,
generation_config={
"candidate_count": 1,
"max_output_tokens": max_tokens,
"top_p": top_p,
"temperature": temperature
},
safety_settings={
"harassment": "block_none",
"hate": "block_none",
"sex": "block_none",
"danger": "block_none"
}
)
request_options = {"timeout": 120}
gemini_model = genai.GenerativeModel(self.model)
try:
response = gemini_model.generate_content(
gemini_messages,
generation_config={
"candidate_count": 1,
"max_output_tokens": max_tokens,
"top_p": top_p,
"temperature": temperature
},
safety_settings={
"harassment": "block_none",
"hate": "block_none",
"sex": "block_none",
"danger": "block_none"
},
request_options=request_options
)
return response.text
except Exception as e:
logger.error("Meet exception when calling Gemini API, " + str(e))
logger.error("Meet exception when calling Gemini API, " + str(e.__class__.__name__) + str(e))
logger.error(f"count_tokens: {gemini_model.count_tokens(gemini_messages)}")
logger.error(f"generation_config: {max_tokens}, {top_p}, {temperature}")
return ""
elif self.model.startswith("qwen"):
messages = payload["messages"]