oc-chat_pdf/chats/chat_tools.py

53 lines
1.7 KiB
Python

from colorama import Fore, Back, Style
class MessageManager:
def create_chat(self, qa):
# Emoticon de robot
ia_emoticon = "\U0001f916" # Emoticon de robot Unicode
humano_emoticon = "\U0001f604" # Emoticon de carita feliz Unicode
# Imprimir el texto en amarillo y negrita con el emoticon de robot
# Definicimo el mensaje de la IA
print(
f"{ia_emoticon} "
+ Style.BRIGHT
+ Fore.YELLOW
+ "IA: "
+ Style.RESET_ALL
+ "Pregunta algo al documento"
)
while True:
input_usuario = input(
Style.BRIGHT + Fore.BLUE + f"{humano_emoticon} You: " + Style.RESET_ALL
)
if input_usuario.lower() == "salir":
break
bot_response = qa.invoke(
{"question": f"{input_usuario}"}, return_only_outputs=True
)
print(
f"{ia_emoticon} "
+ Style.BRIGHT
+ Fore.YELLOW
+ "IA:"
+ Style.RESET_ALL
+ f'{bot_response["answer"]}'
)
def generate_citations(self, documents_source: list) -> str:
text_source: str = ""
for index, document in enumerate(documents_source):
quote: str = document.page_content
source: str = document.metadata["source"].replace("documents/pdfs/", "")
page: str = document.metadata["page"] + 1
fuente: str = (
f"**Fuente #{index + 1}:** \n '{quote}'\n(*{source}, P.{page})*"
)
text_source += fuente + "\n\n\n"
return text_source