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,7 +741,10 @@ class PromptAgent:
assert api_key is not None, "Please set the GENAI_API_KEY environment variable" assert api_key is not None, "Please set the GENAI_API_KEY environment variable"
genai.configure(api_key=api_key) genai.configure(api_key=api_key)
logger.info("Generating content with Gemini model: %s", self.model) logger.info("Generating content with Gemini model: %s", self.model)
response = genai.GenerativeModel(self.model).generate_content( request_options = {"timeout": 120}
gemini_model = genai.GenerativeModel(self.model)
try:
response = gemini_model.generate_content(
gemini_messages, gemini_messages,
generation_config={ generation_config={
"candidate_count": 1, "candidate_count": 1,
@@ -754,13 +757,14 @@ class PromptAgent:
"hate": "block_none", "hate": "block_none",
"sex": "block_none", "sex": "block_none",
"danger": "block_none" "danger": "block_none"
} },
request_options=request_options
) )
try:
return response.text return response.text
except Exception as e: 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 "" return ""
elif self.model.startswith("qwen"): elif self.model.startswith("qwen"):
messages = payload["messages"] messages = payload["messages"]