Melhorias no carregamento dos documentos

app.py

@@ -1,86 +1,43 @@
 import logging
 import sys
 import os
+
 import re
 import base64
 import nest_asyncio
+nest_asyncio.apply()
 import pandas as pd
 from pathlib import Path
 from typing import Any, Dict, List, Optional
 from PIL import Image
 import streamlit as st
 import torch
-from llama_index.core import Settings, SimpleDirectoryReader, StorageContext, Document
+
+# Imports do LlamaIndex
+from llama_index.core import (
+    Settings, 
+    SimpleDirectoryReader, 
+    StorageContext, 
+    Document
+)
+
 from llama_index.core.storage.docstore import SimpleDocumentStore
-# from llama_index.llms.ollama import Ollama
-# from llama_index.embeddings.ollama import OllamaEmbedding
 from llama_index.core.node_parser import LangchainNodeParser
 from langchain.text_splitter import RecursiveCharacterTextSplitter
 from llama_index.core.storage.chat_store import SimpleChatStore
 from llama_index.core.memory import ChatMemoryBuffer
 from llama_index.core.query_engine import RetrieverQueryEngine
 from llama_index.core.chat_engine import CondensePlusContextChatEngine
-#from llama_index.retrievers.bm25 import BM25Retriever
 from llama_index.core.retrievers import QueryFusionRetriever
 from llama_index.vector_stores.chroma import ChromaVectorStore
 from llama_index.core import VectorStoreIndex
-# from llama_index.llms.huggingface import HuggingFaceLLM
-# from llama_index.llms.huggingface_api import HuggingFaceInferenceAPI
-# from llama_index.embeddings.huggingface import HuggingFaceEmbedding
+
 import chromadb
 
 ###############################################################################
 #                           MONKEY PATCH EM bm25s                             #
 ###############################################################################
 import bm25s
-
-# Guardamos a referência da função original
-orig_find_newline_positions = bm25s.utils.corpus.find_newline_positions
-
-def patched_find_newline_positions(path, show_progress=True, leave_progress=True):
-    """
-    Versão 'gambiarra' da função original, forçando uso de encoding='utf-8'
-    e ignorando erros de decodificação. Assim, evitamos UnicodeDecodeError
-    mesmo que o arquivo contenha caracteres fora da faixa UTF-8.
-    
-    (Esta referência é real, baseada em ajustes de leitura de arquivos do Python.)
-    """
-    path = str(path)
-    indexes = []
-
-    with open(path, "r", encoding="utf-8", errors="ignore") as f:
-        indexes.append(f.tell())
-        file_size = os.path.getsize(path)
-
-        try:
-            from tqdm.auto import tqdm
-            pbar = tqdm(
-                total=file_size,
-                desc="Finding newlines for mmindex",
-                unit="B",
-                unit_scale=True,
-                leave=leave_progress,
-                disable=not show_progress,
-            )
-        except ImportError:
-            pbar = None
-
-        while True:
-            line = f.readline()
-            if not line:
-                break
-            t = f.tell()
-            indexes.append(t)
-            if pbar is not None:
-                pbar.update(t - indexes[-2])
-
-        if pbar is not None:
-            pbar.close()
-
-    return indexes[:-1]
-
-# Aplicamos nosso patch
-bm25s.utils.corpus.find_newline_positions = patched_find_newline_positions
 ###############################################################################
 #                   CLASSE BM25Retriever (AJUSTADA PARA ENCODING)             #
 ###############################################################################
@@ -261,28 +218,21 @@ class BM25Retriever(BaseRetriever):
 
         return nodes
 
-#Configuração da imagem da aba
-im = Image.open("pngegg.png")
-st.set_page_config(page_title = "Chatbot Carômetro", page_icon=im, layout = "wide")
 
-#Removido loop e adicionado os.makedirs
-os.makedirs("bm25_retriever", exist_ok=True)
-os.makedirs("chat_store", exist_ok=True)
-os.makedirs("chroma_db", exist_ok=True)
-os.makedirs("documentos", exist_ok=True)
+###############################################################################
+#                   CONFIGURAÇÃO STREAMLIT E AJUSTES DA PIPELINE             #
+###############################################################################
+# Evite reindexar ou baixar dados repetidamente armazenando o estado na sessão.
+im = Image.open("pngegg.png")
+st.set_page_config(page_title="Chatbot Carômetro", page_icon=im, layout="wide")
 
-# Configuração do Streamlit
+# Seções laterais (sidebar)
 st.sidebar.title("Configuração de LLM")
 sidebar_option = st.sidebar.radio("Selecione o LLM", ["gpt-3.5-turbo"])
 
-# logo_url = 'app\logos\logo-sicoob.jpg'
-# st.sidebar.image(logo_url)
 import base64
-
-#Configuração da imagem da sidebar
 with open("sicoob-logo.png", "rb") as f:
     data = base64.b64encode(f.read()).decode("utf-8")
-
     st.sidebar.markdown(
         f"""
         <div style="display:table;margin-top:-80%;margin-left:0%;">
@@ -292,67 +242,24 @@ with open("sicoob-logo.png", "rb") as f:
         unsafe_allow_html=True,
     )
 
-#if sidebar_option == "Ollama":
-   # Settings.llm = Ollama(model="llama3.2:latest", request_timeout=500.0, num_gpu=1)
-   # Settings.embed_model = OllamaEmbedding(model_name="nomic-embed-text:latest")
 if sidebar_option == "gpt-3.5-turbo":
     from llama_index.llms.openai import OpenAI
     from llama_index.embeddings.openai import OpenAIEmbedding
     Settings.llm = OpenAI(model="gpt-3.5-turbo")
     Settings.embed_model = OpenAIEmbedding(model_name="text-embedding-ada-002")
-# elif sidebar_option == 'NuExtract-1.5':
-#     #Embedding do huggingface
-#     Settings.embed_model = HuggingFaceEmbedding(
-#         model_name="BAAI/bge-small-en-v1.5"
-#     )
-#     #Carregamento do modelo local, descomentar o modelo desejado
-
-#     llm = HuggingFaceLLM(
-#         context_window=2048,
-#         max_new_tokens=2048,
-#         generate_kwargs={"do_sample": False},
-#         #query_wrapper_prompt=query_wrapper_prompt,
-#         #model_name="Qwen/Qwen2.5-Coder-32B-Instruct",
-#         #model_name="Qwen/Qwen2.5-14B-Instruct",
-#         # model_name="meta-llama/Llama-3.2-3B",
-#         #model_name="HuggingFaceH4/zephyr-7b-beta",
-#         # model_name="meta-llama/Meta-Llama-3-8B",
-#         model_name="numind/NuExtract-1.5",
-#         #model_name="meta-llama/Llama-3.2-3B",
-#         tokenizer_name="numind/NuExtract-1.5",
-#         device_map="auto",
-#         tokenizer_kwargs={"max_length": 512},
-#         # uncomment this if using CUDA to reduce memory usage
-#         model_kwargs={"torch_dtype": torch.bfloat16},
-#     )
-#     chat = [
-#         {"role": "user", "content": "Hello, how are you?"},
-#         {"role": "assistant", "content": "I'm doing great. How can I help you today?"},
-#         {"role": "user", "content": "I'd like to show off how chat templating works!"},
-#     ]
-
-#     from transformers import AutoTokenizer
-
-#     tokenizer = AutoTokenizer.from_pretrained("numind/NuExtract-1.5")
-#     tokenizer.apply_chat_template(chat, tokenize=False)
-
-#     Settings.chunk_size = 512
-#     Settings.llm = llm
-
 else:
     raise Exception("Opção de LLM inválida!")
 
 logging.basicConfig(stream=sys.stdout, level=logging.INFO)
 logging.getLogger().addHandler(logging.StreamHandler(stream=sys.stdout))
 
-
-# Diretórios configurados pelo usuário
+# Caminhos principais
 chat_store_path = os.path.join("chat_store", "chat_store.json")
-documents_path = os.path.join("documentos")
-chroma_storage_path = os.path.join("chroma_db")  # Diretório para persistência do Chroma
-bm25_persist_path = os.path.join("bm25_retriever")
+documents_path = "documentos"
+chroma_storage_path = "chroma_db"
+bm25_persist_path = "bm25_retriever"
 
-# Classe CSV Customizada (novo código)
+# Classe CSV customizada
 class CustomPandasCSVReader:
     """PandasCSVReader modificado para incluir cabeçalhos nos documentos."""
     def __init__(
@@ -394,15 +301,20 @@ class CustomPandasCSVReader:
                 for text in text_list
             ]
 
-def clean_documents(documents):
-    """Remove caracteres não desejados diretamente nos textos dos documentos."""
-    cleaned_documents = []
+def clean_documents(documents: List[Document]) -> List[Document]:
+    """Remove caracteres indesejados diretamente nos textos."""
+    cleaned_docs = []
     for doc in documents:
         cleaned_text = re.sub(r"[^0-9A-Za-zÀ-ÿ ]", "", doc.get_content())
         doc.text = cleaned_text
-        cleaned_documents.append(doc)
-    return cleaned_documents
+        cleaned_docs.append(doc)
+    return cleaned_docs
 
+def are_docs_downloaded(directory_path: str) -> bool:
+    """Verifica se o diretório tem algum arquivo."""
+    return os.path.isdir(directory_path) and any(os.scandir(directory_path))
+
+# Simula a leitura de arquivos do Google Drive
 from llama_index.readers.google import GoogleDriveReader
 import json
 
@@ -424,10 +336,12 @@ with open(token_path, 'w') as credentials_file:
 google_drive_reader = GoogleDriveReader(credentials_path=credentials_path)
 google_drive_reader._creds = google_drive_reader._get_credentials()
 
-def are_docs_downloaded(directory_path: str) -> bool:
-    return os.path.isdir(directory_path) and any(os.scandir(directory_path))
-
-def download_original_files_from_folder(greader: GoogleDriveReader, pasta_documentos_drive: str, local_path: str):
+def download_original_files_from_folder(
+    greader: GoogleDriveReader, 
+    pasta_documentos_drive: str, 
+    local_path: str
+):
+    """Faz download dos arquivos apenas se não existirem localmente."""
     os.makedirs(local_path, exist_ok=True)
     files_meta = greader._get_fileids_meta(folder_id=pasta_documentos_drive)
     if not files_meta:
@@ -448,114 +362,161 @@ def download_original_files_from_folder(greader: GoogleDriveReader, pasta_docume
         else:
             logging.warning(f"Não foi possível baixar '{file_name}'")
 
-#DADOS/QA_database/Documentos CSV/documentos
+# Pasta do Drive
 pasta_documentos_drive = "1s0UUANcU1B0D2eyRweb1W5idUn1V5JEh"
 
-# Verifica e baixa arquivos se necessário (novo código)
+###############################################################################
+#           CRIAÇÃO/CARREGAMENTO DE RECURSOS (evita repetição de etapas)      #
+###############################################################################
+# 1. Garantir que não baixamos dados novamente se eles já existem.
 if not are_docs_downloaded(documents_path):
     logging.info("Baixando arquivos originais do Drive para 'documentos'...")
-    download_original_files_from_folder(google_drive_reader, pasta_documentos_drive, documents_path)
+    download_original_files_from_folder(
+        google_drive_reader, 
+        pasta_documentos_drive, 
+        documents_path
+    )
 else:
     logging.info("'documentos' já contém arquivos, ignorando download.")
 
-# Configuração de leitura de documentos
-file_extractor = {".csv": CustomPandasCSVReader()}
-documents = SimpleDirectoryReader(
-    input_dir=documents_path,
-    file_extractor=file_extractor,
-    filename_as_id=True,
-    recursive=True
-    #Recursive caso tenha varias pastas no drive
-).load_data()
-
-documents = clean_documents(documents)
-
-# Configuração do Chroma e BM25 com persistência
-docstore = SimpleDocumentStore()
-docstore.add_documents(documents)
-
-db = chromadb.PersistentClient(path=chroma_storage_path)
-chroma_collection = db.get_or_create_collection("dense_vectors")
-vector_store = ChromaVectorStore(chroma_collection=chroma_collection)
+# 2. Se ainda não existir docstore e index no estado da sessão, criamos.
+#    Caso contrário, apenas reutilizamos o que já existe.
+if "docstore" not in st.session_state:
+    # Carregar documentos do diretório local
+    file_extractor = {".csv": CustomPandasCSVReader()}
+    documents = SimpleDirectoryReader(
+        input_dir=documents_path,
+        file_extractor=file_extractor,
+        filename_as_id=True,
+        recursive=True
+    ).load_data()
+
+    documents = clean_documents(documents)
+
+    # Cria docstore
+    docstore = SimpleDocumentStore()
+    docstore.add_documents(documents)
+
+    st.session_state["docstore"] = docstore
+else:
+    docstore = st.session_state["docstore"]
+
+# 3. Configuramos o VectorStore + Chroma sem recriar se já estiver pronto.
+if "vector_store" not in st.session_state:
+    db = chromadb.PersistentClient(path=chroma_storage_path)
+    chroma_collection = db.get_or_create_collection("dense_vectors")
+    vector_store = ChromaVectorStore(chroma_collection=chroma_collection)
+    st.session_state["vector_store"] = vector_store
+else:
+    vector_store = st.session_state["vector_store"]
 
-# Configuração do StorageContext
 storage_context = StorageContext.from_defaults(
-    docstore=docstore, vector_store=vector_store
+    docstore=docstore, 
+    vector_store=vector_store
 )
 
-# Criação/Recarregamento do índice com embeddings
-if os.path.exists(chroma_storage_path):
-    index = VectorStoreIndex.from_vector_store(vector_store)
+# 4. Carregamos ou criamos o índice. Se já existe a base do Chroma, supõe-se
+#    que o índice foi persistido. Caso contrário, cria-se.
+if "index" not in st.session_state:
+    if os.path.exists(chroma_storage_path) and os.listdir(chroma_storage_path):
+        # Há dados salvos, então criamos índice a partir do vector_store
+        index = VectorStoreIndex.from_vector_store(vector_store)
+    else:
+        # Cria índice (chunk_size pode ser configurado conforme necessidade)
+        splitter = LangchainNodeParser(
+            RecursiveCharacterTextSplitter(chunk_size=1024, chunk_overlap=128)
+        )
+        index = VectorStoreIndex.from_documents(
+            list(docstore.docs.values()),
+            storage_context=storage_context,
+            transformations=[splitter]
+        )
+        vector_store.persist()
+    st.session_state["index"] = index
 else:
-    splitter = LangchainNodeParser(
-        RecursiveCharacterTextSplitter(chunk_size=1024, chunk_overlap=128)
-    )
-    index = VectorStoreIndex.from_documents(
-        documents,
-        storage_context=storage_context,
-        transformations=[splitter]
-    )
-    vector_store.persist()
-
-# Criação/Recarregamento do BM25 Retriever
-if os.path.exists(os.path.join(bm25_persist_path, "params.index.json")):
-    bm25_retriever = BM25Retriever.from_persist_dir(bm25_persist_path)
+    index = st.session_state["index"]
+
+# 5. Criação ou carregamento do BM25Retriever customizado
+if "bm25_retriever" not in st.session_state:
+    if (
+        os.path.exists(bm25_persist_path) 
+        and os.path.exists(os.path.join(bm25_persist_path, "bm25.index.json"))
+    ):
+        bm25_retriever = BM25Retriever.from_persist_dir(bm25_persist_path)
+    else:
+        bm25_retriever = BM25Retriever.from_defaults(
+            docstore=docstore,
+            similarity_top_k=2,
+            language="portuguese",
+            verbose=True
+        )
+        os.makedirs(bm25_persist_path, exist_ok=True)
+        bm25_retriever.persist(bm25_persist_path)
+    st.session_state["bm25_retriever"] = bm25_retriever
 else:
-    bm25_retriever = BM25Retriever.from_defaults(
-        docstore=docstore,
+    bm25_retriever = st.session_state["bm25_retriever"]
+
+# 6. Criamos ou recuperamos o retriever que fará Query Fusion (BM25 + eventual vetor)
+if "fusion_retriever" not in st.session_state:
+    vector_retriever = index.as_retriever(similarity_top_k=2)
+    fusion_retriever = QueryFusionRetriever(
+        [bm25_retriever, vector_retriever],
         similarity_top_k=2,
-        language="portuguese",  # Idioma ajustado para seu caso
+        num_queries=0,
+        mode="reciprocal_rerank",
+        use_async=True,
+        verbose=True,
+        query_gen_prompt=(
+            "Gere {num_queries} perguntas de busca relacionadas à seguinte pergunta. "
+            "Priorize o significado da pergunta sobre qualquer histórico de conversa. "
+            "Se o histórico não for relevante, ignore-o. "
+            "Não adicione explicações ou introduções. Apenas escreva as perguntas. "
+            "Pergunta: {query}\n\nPerguntas:\n"
+        ),
     )
-    os.makedirs(bm25_persist_path, exist_ok=True)
-    bm25_retriever.persist(bm25_persist_path)
-
-# Combinação de Retrievers (Embeddings + BM25)
-vector_retriever = index.as_retriever(similarity_top_k=2)
-retriever = QueryFusionRetriever(
-    [vector_retriever, bm25_retriever],
-    similarity_top_k=3,
-    num_queries=0,
-    mode="reciprocal_rerank",
-    use_async=True,
-    verbose=True,
-    query_gen_prompt=(
-        "Gere {num_queries} perguntas de busca relacionadas à seguinte pergunta. "
-        "Priorize o significado da pergunta sobre qualquer histórico de conversa. "
-        "Se o histórico não for relevante para a pergunta, ignore-o. "
-        "Não adicione explicações, notas ou introduções. Apenas escreva as perguntas. "
-        "Pergunta: {query}\n\n"
-        "Perguntas:\n"
-    ),
-)
-
-# Configuração do chat engine
-nest_asyncio.apply()
-memory = ChatMemoryBuffer.from_defaults(token_limit=3900)
-query_engine = RetrieverQueryEngine.from_args(retriever)
-chat_engine = CondensePlusContextChatEngine.from_defaults(
-    query_engine,
-    memory=memory,
-    chat_mode="context",
-    context_prompt=(
-        "Você é um assistente virtual capaz de interagir normalmente, além de"
-        " fornecer informações sobre organogramas e listar funcionários."
-        " Aqui estão os documentos relevantes para o contexto:\n"
-        "{context_str}"
-        "\nInstrução: Use o histórico da conversa anterior, ou o contexto acima, para responder."
-    ),
-    verbose=True,
-)
-
-# Armazenamento do chat
-chat_store = SimpleChatStore()
-if os.path.exists(chat_store_path):
-    chat_store = SimpleChatStore.from_persist_path(persist_path=chat_store_path)
+    st.session_state["fusion_retriever"] = fusion_retriever
+else:
+    fusion_retriever = st.session_state["fusion_retriever"]
+
+# 7. Configura o Chat Engine caso ainda não esteja na sessão
+if "chat_engine" not in st.session_state:
+    nest_asyncio.apply()
+    memory = ChatMemoryBuffer.from_defaults(token_limit=3900)
+    query_engine = RetrieverQueryEngine.from_args(fusion_retriever)
+
+    chat_engine = CondensePlusContextChatEngine.from_defaults(
+        query_engine,
+        memory=memory,
+        context_prompt=(
+            "Você é um assistente virtual capaz de interagir normalmente, além de "
+            "fornecer informações sobre organogramas e listar funcionários. "
+            "Aqui estão os documentos relevantes para o contexto:\n"
+            "{context_str}\n"
+            "Use o histórico anterior ou o contexto acima para responder."
+        ),
+        verbose=True,
+    )
+    st.session_state["chat_engine"] = chat_engine
+else:
+    chat_engine = st.session_state["chat_engine"]
+
+# 8. Armazenamento do chat
+if "chat_store" not in st.session_state:
+    if os.path.exists(chat_store_path):
+        chat_store = SimpleChatStore.from_persist_path(persist_path=chat_store_path)
+    else:
+        chat_store = SimpleChatStore()
+        chat_store.persist(persist_path=chat_store_path)
+    st.session_state["chat_store"] = chat_store
 else:
-    chat_store.persist(persist_path=chat_store_path)
+    chat_store = st.session_state["chat_store"]
 
-# Interface do Chatbot
+
+###############################################################################
+#                          INTERFACE DO CHAT EM STREAMLIT                     #
+###############################################################################
 st.title("Chatbot Carômetro")
-st.write("Este chatbot pode te ajudar a conseguir informações relevantes sobre os carômetros da Sicoob.")
+st.write("Este assistente virtual pode te ajudar a encontrar informações relevantes sobre os carômetros da Sicoob.")
 
 if 'chat_history' not in st.session_state:
     st.session_state.chat_history = []
@@ -567,23 +528,18 @@ for message in st.session_state.chat_history:
 
 user_input = st.chat_input("Digite sua pergunta")
 if user_input:
-    # Exibir a mensagem do usuário e adicionar ao histórico
     with st.chat_message('user'):
         st.write(user_input)
     st.session_state.chat_history.append(f"user: {user_input}")
 
-    # Placeholder para a mensagem do assistente
     with st.chat_message('assistant'):
         message_placeholder = st.empty()
         assistant_message = ''
 
-    # Obter a resposta em streaming do chat_engine
     response = chat_engine.stream_chat(user_input)
     for token in response.response_gen:
         assistant_message += token
-        # Atualizar o placeholder da mensagem
         message_placeholder.markdown(assistant_message + "▌")
 
-    # Remover o cursor após a conclusão
     message_placeholder.markdown(assistant_message)
-    st.session_state.chat_history.append(f"assistant: {assistant_message}")
+    st.session_state.chat_history.append(f"assistant: {assistant_message}")