Gemini api python chat sample.

.env

API_KEY=[REDACTED]
CHAT_MODEL_NAME=gemini-pro
INIT_PROMPT=You are a Pizza order bot. You take orders for Pizza This is all you can do.
INIT_PROMPT_AGREEMENT=I under stand I am a pizza bot I will only take orders for pizza.

app.py

from dataclasses import asdict, dataclass
from dotenv import load_dotenv
import google.generativeai as genai
import os

load_dotenv()

# The api key for accessing the api. stored in .env
genai.configure(api_key=os.getenv("API_KEY"))

# name of the ai model used in this call.
CHAT_MODEL_NAME = os.getenv("CHAT_MODEL_NAME")

# Init the chatbot with two prompts, there must be two.

# First user instructions
SYSTEM_PROMPT = "You are a pizza bot your job is to take orders for pizza. This is all you can do"
# The models agreement
MODEL_RESPONSE = "I understand I am a pizza bot I will take orders. For pizza"


@dataclass
class Message:
    role: str
    parts: list[str]
    Model = "model"
    User = "user"

    def __init__(self, role: str, content: str):
        """
        Initializes a new Message object with the specified role and content.

        Args:
            role (str): The role of the sender (e.g., "user", "model").
            content (str): The content of the message.
        """

        self.role = role
        self.parts = [content]


class GeminiService:
    generation_config: str = {
        'temperature': 0.9,
        'top_p': 1,
        'top_k': 40,
        'max_output_tokens': 2048,
        'stop_sequences': [],
    }

    safety_settings: list[str] = [{"category": "HARM_CATEGORY_HARASSMENT", "threshold": "BLOCK_NONE"},
                                  {"category": "HARM_CATEGORY_HATE_SPEECH", "threshold": "BLOCK_NONE"},
                                  {"category": "HARM_CATEGORY_SEXUALLY_EXPLICIT", "threshold": "BLOCK_NONE"},
                                  {"category": "HARM_CATEGORY_DANGEROUS_CONTENT", "threshold": "BLOCK_NONE"}]

    def __init__(self):
        self.model = genai.GenerativeModel(model_name=CHAT_MODEL_NAME,
                                           generation_config=self.generation_config,
                                           safety_settings=self.safety_settings
                                           )

    def get_completion(self,
                       messages: list[Message],
                       message: str) -> str:

        print(messages)
        # get the conversation history as a dict
        messages = [asdict(m) for m in messages]
        print(messages)
        # Start the chat with the conversation history
        convo = self.model.start_chat(history=messages)

        # send message with the new request from the user.
        convo.send_message(message)

        # return the response
        return convo.last.text


class ChattyUI:
    def __init__(
            self,
            system_prompt: str = SYSTEM_PROMPT,
            system_prompt_response: str = MODEL_RESPONSE
    ) -> None:
        self.service = GeminiService()
        self.system_prompt = system_prompt
        self.system_prompt_response = system_prompt_response

    def close(self):
        self.interface.close()

    def answer(self, message: str, chat_history: list[str]) -> str:
        # Init the chatbot with the starting prompts.
        # Note: we dont add this to the history because then it will show up on the screen
        #       for the user to see.
        messages = [Message(Message.User, self.system_prompt),
                    Message(Message.Model, self.system_prompt_response)]

        # add user and assistant message from history
        for user_content, assistant_content in chat_history:
            messages.append(Message(Message.User, user_content))
            messages.append(Message(Message.Model, assistant_content))
        return self.service.get_completion(messages=messages, message=message)


chatty = ChattyUI()
# create chat history
history = [{"One cheese pizza", "okay that's one cheese pizza anything else."}]
print(f"History: {history}")
# a new conversation item.
user_says = "Whats my order"
print(user_says)
model_says = chatty.answer("What is my order", history)
print(model_says)

# Adding a new item to the history list
# new_item = {user_says, model_says}
# history.append(new_item)
# print(f"History: {history}")
# https://www.gradio.app/guides/creating-a-custom-chatbot-with-blocks