Skip to content

Commit e681574

Browse files
authored
Add image reading capability to Gemini (#69)
I used upload file from sdk to upload images if the provider is gemini.
1 parent 28d7e46 commit e681574

File tree

1 file changed

+16
-1
lines changed

1 file changed

+16
-1
lines changed

tools/llm_api.py

Lines changed: 16 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -207,7 +207,22 @@ def query_llm(prompt: str, client=None, model=None, provider="openai", image_pat
207207

208208
elif provider == "gemini":
209209
model = client.GenerativeModel(model)
210-
response = model.generate_content(prompt)
210+
if image_path:
211+
file = genai.upload_file(image_path, mime_type="image/png")
212+
chat_session = model.start_chat(
213+
history=[{
214+
"role": "user",
215+
"parts": [file, prompt]
216+
}]
217+
)
218+
else:
219+
chat_session = model.start_chat(
220+
history=[{
221+
"role": "user",
222+
"parts": [prompt]
223+
}]
224+
)
225+
response = chat_session.send_message(prompt)
211226
return response.text
212227

213228
except Exception as e:

0 commit comments

Comments
 (0)