ChatWithGemini / app.py
AyoubChLin's picture
Create app.py
9fa982f verified
import os
import getpass
import streamlit as st
from typing import List
from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
from langchain_core.messages import HumanMessage, AIMessage
from langchain_google_genai import ChatGoogleGenerativeAI
class Chat:
def __init__(self):
self.chat = ChatGoogleGenerativeAI(model="gemini-pro", convert_system_message_to_human=True)
self.history = []
if "GOOGLE_API_KEY" not in os.environ:
os.environ["GOOGLE_API_KEY"] = st.text_input("Provide your API key: ", type="password")
self.prompt = ChatPromptTemplate.from_messages([
("system", "You are a senior software engineer. Answer all questions to the best of your ability."),
MessagesPlaceholder(variable_name="messages"),
])
def add_message(self, message: str):
self.history.append(HumanMessage(content=message))
def add_ai_message(self, message: str):
self.history.append(AIMessage(content=message))
def chat_with_user(self):
st.title("Chat with AI")
user_input = st.text_input("You: ", "")
if st.button("Send"):
if user_input:
self.add_message(user_input)
response = self.chat.invoke(self.prompt.format_prompt(messages=self.history).to_messages())
self.add_ai_message(response.content)
st.write(f"AI: {response.content}")
# Display chat history
for msg in self.history:
if isinstance(msg, HumanMessage):
st.write(f"User: {msg.content}")
else:
st.write(f"AI: {msg.content}")
if __name__ == "__main__":
chat_instance = Chat()
chat_instance.chat_with_user()