rag/app.py

51 lines
1.7 KiB
Python

# python -m streamlit run app.py
import streamlit as st
from rag import RAG
import re
import logging
@st.cache_resource
def init_rag():
llm_url = 'http://127.0.0.1:8080/completion'
# embed_model_name = 'intfloat/multilingual-e5-large'
embed_model_name = 'dangvantuan/sentence-camembert-large'
collection_name = 'cera'
chromadb_path = './chromadb'
rag = RAG(llm_url, embed_model_name, collection_name, chromadb_path, mulitlingual_e5=False)
return rag
rag = init_rag()
st.title("CERA Chat")
if st.sidebar.button("Nouvelle conversation"):
st.session_state.messages = []
rag.reset_history()
if "messages" not in st.session_state:
st.session_state.messages = []
for message in st.session_state.messages:
with st.chat_message(message["role"]):
st.markdown(message["content"])
if prompt := st.chat_input("Comment puis-je vous aider ?"):
st.session_state.messages.append({"role": "user", "content": prompt})
with st.chat_message("user"):
st.markdown(prompt)
with st.chat_message("assistant"):
message_placeholder = st.empty()
full_response = ""
for response in rag.chat(prompt):
full_response += response
message_placeholder.markdown(full_response + "")
message_placeholder.markdown(full_response)
# url_pattern = r"URL :\n(https?://[^\s]+)"
# urls = re.findall(url_pattern, rag.chat_history[-1]['assistant'])
# markdown_urls = "\n".join([f"- {url}" for url in urls])
# logging.info(f"URLs: \n{markdown_urls}")
# message_placeholder.markdown(markdown_urls)
st.session_state.messages.append({"role": "assistant", "content": full_response})