import chromadb

import streamlit as st

from llama_index.core import Settings

from llama_index.embeddings.huggingface import HuggingFaceEmbedding

from commons import contexto, load_llm, load_agent, \

load_engines, load_query_index

# Ruta donde se encuentra la base de datos de Chroma

chroma_path = './chromadb'

llm_string = 'Phi3 Mini' # No cambiar

# LlamaIndex gestiona el uso del modelo y el embedding a través de la clase Settings

Settings.llm = load_llm(llm_string=llm_string)

Settings.embed_model = HuggingFaceEmbedding(model_name="BAAI/bge-small-en-v1.5")

chroma_client = chromadb.PersistentClient(path=chroma_path)

as_chatbot = st.sidebar.toggle('Chatbot', help='Si no se habilita como chatbot el asistente no tendrá memoria',

on_change=st.session_state.clear())

# Si se habilita como chatbot se carga el agente para que tenga memoria. Sólo usamos un motor de búsqueda

if as_chatbot:

engines = load_engines(knowledge_bases=['All'], _chroma_client=chroma_client)

agent = load_agent(_engines=engines, llm_string=llm_string, _chroma_client=chroma_client, contexto=contexto)

else:

query_index = load_query_index(collection='All', chroma_client=chroma_client, chroma_path='./chromadb',

chunk_size=512)

# Frontend

st.header("Asistente de ciberseguridad haciendo uso de Phi3 Mini 🛺💨")

if "messages" not in st.session_state.keys():

st.session_state.messages = [

{"role": "assistant",

"content": "Hola, soy tu asistente virtual especializado en ciberseguridad. ¿En qué puedo ayudarte?"}

]

if as_chatbot:

# noinspection PyUnboundLocalVariable

st.session_state.chat_engine = agent

else:

# noinspection PyUnboundLocalVariable

st.session_state.query_engine = query_index

if prompt := st.chat_input("Your question"):

st.session_state.messages.append({"role": "user", "content": prompt})

for message in st.session_state.messages:

with st.chat_message(message["role"]):

st.write(message["content"])

if st.session_state.messages[-1]["role"] != "assistant":

with st.chat_message("assistant"):

with st.spinner("Thinking..."):

if as_chatbot:

response = st.session_state.chat_engine.chat(prompt)

else:

response = st.session_state.query_engine.query(prompt)

st.write(response.response)

message = {"role": "assistant", "content": response.response}

st.session_state.messages.append(message)