-
Notifications
You must be signed in to change notification settings - Fork 0
/
gemini.py
28 lines (25 loc) · 808 Bytes
/
gemini.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
import google.generativeai as genai
from typing import Iterable
import os
from data_model import ChatMessage
generation_config = {
"temperature": 1,
"top_p": 0.95,
"top_k": 64,
"max_output_tokens": 8192,
}
def configure_gemini():
genai.configure(api_key=os.getenv("APIKEY"))
def send_prompt_flash(fileimage, prompt: str, history: list[ChatMessage]) -> Iterable[str]:
configure_gemini()
model = genai.GenerativeModel(
model_name="gemini-1.5-flash-latest",
generation_config=generation_config,
)
chat_session = model.start_chat(
history=[
{"role": message.role, "parts": [message.content]} for message in history
]
)
for chunk in chat_session.send_message([prompt, fileimage], stream=True):
yield chunk.text