136 lines
4.7 KiB
Python
136 lines
4.7 KiB
Python
# import os
|
|
# from dotenv import load_dotenv
|
|
# from langchain_community.chat_models import ChatOpenAI
|
|
import streamlit as st
|
|
from chats.streamlit_tools import import_file # ,clear_cache
|
|
from streamlit_extras.add_vertical_space import add_vertical_space
|
|
from langchain_tools.pdf_tools import PdfLangChain
|
|
from langchain_tools.lc_tools import LangChainTools
|
|
from chats.chat_tools import MessageManager
|
|
|
|
|
|
# App title
|
|
st.set_page_config(page_title="LLMOneClusterTeam")
|
|
|
|
# sidebar
|
|
with st.sidebar:
|
|
# Cargar el logo (asegúrate de que el archivo de imagen esté en la misma carpeta que tu script)
|
|
logo_path = "documents/Logo azulblanco.png"
|
|
logo = st.sidebar.image(logo_path, width=200)
|
|
|
|
# Ajusta el ancho según sea necesario
|
|
add_vertical_space(28)
|
|
# pdf_name = import_file()
|
|
st.markdown("Built by [OneCluster](https://www.onecluster.org/).")
|
|
|
|
|
|
col1, col2 = st.columns([1.1, 1])
|
|
with col1:
|
|
st.title(
|
|
"DocumentAssist",
|
|
)
|
|
with col2:
|
|
logo_2 = st.image("documents/pdfs/logo_1-removebg-preview.png", width=110)
|
|
|
|
pdf_name = import_file()
|
|
if pdf_name:
|
|
with st.spinner("Processing the document..."):
|
|
# Inicializamos la clase PdfLangChain
|
|
pdfLangChain = PdfLangChain(pdf_name)
|
|
pdf_name = pdfLangChain.file_name
|
|
|
|
# Cargamos el documento PDF
|
|
docs: list = pdfLangChain.load_pdf()
|
|
|
|
# Dividimos los documentos en partes mas pequenas
|
|
docs_split: list = pdfLangChain.split_docs(docs)
|
|
|
|
# Instanciamos la clase LangChainTools que contiene herramientras LangChain
|
|
langChainTools = LangChainTools()
|
|
|
|
# Cargamos el modelo de embeddings
|
|
embedding_model = langChainTools.load_embedding_opnai()
|
|
|
|
# Cargamos el modelo de embeddings
|
|
# embedding_model = langChainTools.load_embedding_hf()
|
|
|
|
# Creamos el vector store
|
|
docstorage = langChainTools.create_vector_strore(
|
|
docs_split, pdf_name, embedding_model
|
|
)
|
|
|
|
# Cargamos el modelo LLM desde LangChain
|
|
llm = langChainTools.load_llm_open_source()
|
|
|
|
# Creamos la cadena que integra Vectorstroe, el LLM para hacer consultas.
|
|
# Para este caso la cadena tene el parametro de memoria.
|
|
qa = langChainTools.define_retrieval_qa_memory(
|
|
llm, docstorage, pdf_name, embedding_model
|
|
)
|
|
|
|
# Store conversation history
|
|
if "messages" not in st.session_state.keys():
|
|
st.session_state.messages = [
|
|
{
|
|
"role": "assistant",
|
|
"content": "Hola, soy una IA con el que puedes chatear con tu PDF. Haz un pregunta al documento.",
|
|
}
|
|
]
|
|
|
|
# Display or clear chat messages
|
|
for message in st.session_state.messages:
|
|
with st.chat_message(message["role"]):
|
|
st.write(message["content"])
|
|
|
|
def clear_chat_history():
|
|
st.session_state.messages = [
|
|
{
|
|
"role": "assistant",
|
|
"content": "Hola, soy una IA con el que puedes chatear con tu PDF. Haz un pregunta al documento.",
|
|
}
|
|
]
|
|
|
|
st.sidebar.button("Clear chat history", on_click=clear_chat_history)
|
|
|
|
@st.cache_resource
|
|
def get_num_tokens(prompt):
|
|
"""Get the number of tokens in a given prompt"""
|
|
return len(prompt.split())
|
|
|
|
# Function for generating Snowflake Arctic response
|
|
|
|
# User-provided prompt
|
|
if prompt := st.chat_input():
|
|
st.session_state.messages.append({"role": "user", "content": prompt})
|
|
with st.chat_message("user"):
|
|
st.write(prompt)
|
|
|
|
# Generate a new response if last message is not from assistant
|
|
if st.session_state.messages[-1]["role"] != "assistant":
|
|
with st.spinner("Thinking..."):
|
|
# Creamos la cadena que integra Vectorstroe, el LLM para hacer consultas.
|
|
# Para este caso la cadena tene el parametro de memoria.
|
|
qa = langChainTools.define_retrieval_qa_memory(
|
|
llm, docstorage, pdf_name, embedding_model
|
|
)
|
|
|
|
input = "\n".join([msg["content"] for msg in st.session_state.messages])
|
|
|
|
query = qa.invoke({"question": f"{prompt}"}, return_only_outputs=True)
|
|
|
|
response_text = query["answer"]
|
|
documents_source = query["source_documents"]
|
|
|
|
messageManager = MessageManager()
|
|
|
|
citation: str = messageManager.generate_citations(documents_source)
|
|
# st.markdown(citation)
|
|
|
|
with st.chat_message("assistant"):
|
|
st.write(response_text)
|
|
st.session_state.messages.append(
|
|
{"role": "assistant", "content": response_text}
|
|
)
|
|
expander = st.expander("Fuentes")
|
|
expander.markdown(citation)
|