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() | |