Spaces:
Sleeping
Sleeping
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() | |