(fix) retry gemini-pro-vision 3 times

This commit is contained in:
ishaan-jaff 2024-01-11 14:39:08 +05:30
parent c41b47dc8b
commit fccd81a444

View file

@ -6,29 +6,34 @@ sys.path.insert(
import litellm import litellm
from dotenv import load_dotenv from dotenv import load_dotenv
def generate_text(): def generate_text():
try: try:
litellm.set_verbose = True
messages = [ messages = [
{ {
"role": "user", "role": "user",
"content": [ "content": [
{ {"type": "text", "text": "What is this image?"},
"type": "text",
"text": "What is this image?"
},
{ {
"type": "image_url", "type": "image_url",
"image_url": { "image_url": {
"url": "https://avatars.githubusercontent.com/u/17561003?v=4" "url": "https://avatars.githubusercontent.com/u/17561003?v=4"
} },
} },
] ],
} }
] ]
response = litellm.completion(model="gemini/gemini-pro-vision", messages=messages, stop="Hello world") response = litellm.completion(
model="gemini/gemini-pro-vision",
messages=messages,
stop="Hello world",
num_retries=3,
)
print(response) print(response)
assert isinstance(response.choices[0].message.content, str) == True assert isinstance(response.choices[0].message.content, str) == True
except Exception as exception: except Exception as exception:
raise Exception("An error occurred during text generation:", exception) raise Exception("An error occurred during text generation:", exception)
generate_text()
# generate_text()