Update agent.py
This commit is contained in:
@@ -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"]
|
||||
|
||||
Reference in New Issue
Block a user