Spaces:
Running
Running
import os | |
import google.generativeai as genai | |
import mesop as me | |
from typing_extensions import TypedDict | |
from data_model import ChatMessage, State | |
class Pad(TypedDict): | |
title: str | |
content: str | |
class Output(TypedDict): | |
intro: str | |
pad: Pad | |
conclusion: str | |
# Create the model | |
# See https://ai.google.dev/api/python/google/generativeai/GenerativeModel | |
generation_config = { | |
"temperature": 1, | |
"top_p": 0.95, | |
"top_k": 64, | |
"max_output_tokens": 8192, | |
"response_mime_type": "text/plain", | |
# "response_schema": Output, | |
} | |
def configure_gemini(): | |
state = me.state(State) | |
genai.configure(api_key=state.gemini_api_key) | |
def send_prompt_pro(prompt: str, history: list[ChatMessage]): | |
configure_gemini() | |
model = genai.GenerativeModel( | |
model_name="gemini-1.5-pro-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, stream=True): | |
yield chunk.text | |
def send_prompt_flash(prompt: str, history: list[ChatMessage]): | |
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, stream=True): | |
yield chunk.text | |