AI_chatbot/app.py

from datetime import datetime

import streamlit as st
from langchain import LLMChain
from langchain.callbacks.base import BaseCallbackHandler
from langchain.callbacks.tracers.langchain import wait_for_all_tracers
from langchain.callbacks.tracers.run_collector import RunCollectorCallbackHandler
from langchain.chat_models import ChatOpenAI
from langchain.memory import StreamlitChatMessageHistory, ConversationBufferMemory
from langchain.prompts import ChatPromptTemplate, MessagesPlaceholder
from langchain.schema.runnable import RunnableConfig
from langsmith import Client
from streamlit_feedback import streamlit_feedback

st.set_page_config(
 page_title="Chat LangSmith",
 page_icon="🦜",
)


def get_llm_chain(system_prompt: str, memory: ConversationBufferMemory) -> LLMChain:
 """Return a basic LLMChain with memory."""
 prompt = ChatPromptTemplate.from_messages(
 [
 (
 "system",
 system_prompt + "\nIt's currently {time}.",
 ),
 MessagesPlaceholder(variable_name="chat_history"),
 ("human", "{input}"),
 ],
 ).partial(time=lambda: str(datetime.now()))
 llm = ChatOpenAI(temperature=0.7, streaming=True)
 return LLMChain(prompt=prompt, llm=llm, memory=memory)


client = Client()


# "# Chat🦜🛠️"
# Initialize State
if "trace_link" not in st.session_state:
 st.session_state.trace_link = None
if "run_id" not in st.session_state:
 st.session_state.run_id = None
st.sidebar.markdown(
 """
# Menu
""",
)

_DEFAULT_SYSTEM_PROMPT = "You are a helpful chatbot."

system_prompt = st.sidebar.text_area(
 "Custom Instructions",
 _DEFAULT_SYSTEM_PROMPT,
 help="Custom instructions to provide the language model to determine style, personality, etc.",
)
system_prompt = system_prompt.strip().replace("{", "{{").replace("}", "}}")
memory = ConversationBufferMemory(
 chat_memory=StreamlitChatMessageHistory(key="langchain_messages"),
 return_messages=True,
 memory_key="chat_history",
)

chain = get_llm_chain(system_prompt, memory)

if st.sidebar.button("Clear message history"):
 print("Clearing message history")
 memory.clear()
 st.session_state.trace_link = None
 st.session_state.run_id = None


# Display chat messages from history on app rerun
# NOTE: This won't be necessary for Streamlit 1.26+, you can just pass the type directly
# https://github.com/streamlit/streamlit/pull/7094
def _get_openai_type(msg):
 if msg.type == "human":
 return "user"
 if msg.type == "ai":
 return "assistant"
 return msg.role if msg.type == "chat" else msg.type


for msg in st.session_state.langchain_messages:
 streamlit_type = _get_openai_type(msg)
 avatar = "🦜" if streamlit_type == "assistant" else None
 with st.chat_message(streamlit_type, avatar=avatar):
 st.markdown(msg.content)

if st.session_state.trace_link:
 st.sidebar.markdown(
 f'<a href="{st.session_state.trace_link}" target="_blank"><button>Latest Trace: 🛠️</button></a>',
 unsafe_allow_html=True,
 )


class StreamHandler(BaseCallbackHandler):
 def __init__(self, container, initial_text=""):
 self.container = container
 self.text = initial_text

 def on_llm_new_token(self, token: str, **kwargs) -> None:
 self.text += token
 self.container.markdown(self.text)


run_collector = RunCollectorCallbackHandler()


def _reset_feedback():
 st.session_state.feedback_update = None
 st.session_state.feedback = None


if prompt := st.chat_input(placeholder="Ask me a question!"):
 st.chat_message("user").write(prompt)
 _reset_feedback()

 with st.chat_message("assistant", avatar="🦜"):
 message_placeholder = st.empty()
 stream_handler = StreamHandler(message_placeholder)
 runnable_config = RunnableConfig(
 callbacks=[run_collector, stream_handler],
 tags=["Streamlit Chat"],
 )
 full_response = chain.invoke({"input": prompt}, config=runnable_config)["text"]
 message_placeholder.markdown(full_response)

 run = run_collector.traced_runs[0]
 run_collector.traced_runs = []
 st.session_state.run_id = run.id
 wait_for_all_tracers()
 url = client.read_run(run.id).url
 st.session_state.trace_link = url

# Simple feedback section
# Optionally add a thumbs up/down button for feedback
if st.session_state.get("run_id"):
 feedback = streamlit_feedback(
 feedback_type="thumbs",
 key=f"feedback_{st.session_state.run_id}",
 )
 scores = {"👍": 1, "👎": 0}
 if feedback:
 score = scores[feedback["score"]]
 feedback = client.create_feedback(
 st.session_state.run_id,
 "user_score",
 score=score,
 )
 st.session_state.feedback = {"feedback_id": str(feedback.id), "score": score}

# Prompt for more information, if feedback was submitted
if st.session_state.get("feedback"):
 feedback = st.session_state.get("feedback")
 feedback_id = feedback["feedback_id"]
 score = feedback["score"]
 if score == 0:
 if correction := st.text_input(
 label="What would the correct or preferred response have been?",
 key=f"correction_{feedback_id}",
 ):
 st.session_state.feedback_update = {
 "correction": {"desired": correction},
 "feedback_id": feedback_id,
 }
 elif score == 1:
 if comment := st.text_input(
 label="Anything else you'd like to add about this response?",
 key=f"comment_{feedback_id}",
 ):
 st.session_state.feedback_update = {
 "comment": comment,
 "feedback_id": feedback_id,
 }
# Update the feedback if additional information was provided
if st.session_state.get("feedback_update"):
 feedback_update = st.session_state.get("feedback_update")
 feedback_id = feedback_update.pop("feedback_id")
 client.update_feedback(feedback_id, **feedback_update)
 # Clear the comment or correction box
 _reset_feedback()

# # Feedback section
# if st.session_state.get("last_run"):
# run_url = client.read_run(st.session_state.last_run).url
# st.sidebar.markdown(f"[Latest Trace: 🛠️]({run_url})")
# feedback = streamlit_feedback(
# feedback_type="faces",
# optional_text_label="[Optional] Please provide an explanation",
# key=f"feedback_{st.session_state.last_run}",
# )
# if feedback:
# scores = {"😀": 1, "🙂": 0.75, "😐": 0.5, "🙁": 0.25, "😞": 0}
# client.create_feedback(
# st.session_state.last_run,
# feedback["type"],
# score=scores[feedback["score"]],
# comment=feedback.get("text", None),
# )
# st.toast("Feedback recorded!", icon="📝")