Upload 3 files (#1)

- Upload 3 files (9bdf06b946387d892b1fc36cb79f21a88e0c00cd)


Co-authored-by: Gabriel Silva Rodrigues <Restodecoca@users.noreply.huggingface.co>

chatbot_server.py

@@ -1,133 +1,140 @@
-import os
-import logging
-import sys
-
-from flask import Flask, request, jsonify, Response
-# Inicializa o Flask
-app = Flask(__name__)
-
-logging.basicConfig(stream=sys.stdout, level=logging.INFO)
-
-from llama_index.llms.openai import OpenAI
-from llama_index.embeddings.openai import OpenAIEmbedding
-from llama_index.core import (
-    Settings,
-    SimpleDirectoryReader,
-    StorageContext,
-    Document,
-)
-
-Settings.llm = OpenAI(model="gpt-3.5-turbo")
-Settings.embed_model = OpenAIEmbedding(model_name="text-embedding-3-small")
-directory_path = "documentos"
-from llama_index.readers.file import PDFReader #concatenar todo o documento já vem nativo no pdfreader
-file_extractor = {".pdf": PDFReader(return_full_document = True)}
-from drive_downloader import GoogleDriveDownloader
-
-# ID da pasta no Drive e caminho local
-folder_id = "1n34bmh9rlbOtCvE_WPZRukQilKeabWsN"
-local_path = directory_path
-
-GoogleDriveDownloader().download_from_folder(folder_id, local_path)
-
-documents = SimpleDirectoryReader(
-    input_dir=directory_path,
-    file_extractor=file_extractor,
-    filename_as_id=True,
-    recursive=True
-).load_data()
-
-from document_creator import create_single_document_with_filenames
-document = create_single_document_with_filenames(directory_path = directory_path)
-documents.append(document)
-
-#from llama_index.core.ingestion import IngestionPipeline
-#ingestion pipeline vai entrar em uso quando adicionar o extrator de metadados
-from llama_index.core.node_parser import SentenceSplitter
-splitter = SentenceSplitter(chunk_size=1024, chunk_overlap=128)
-nodes = splitter.get_nodes_from_documents(documents)
-
-from llama_index.core.storage.docstore import SimpleDocumentStore
-docstore = SimpleDocumentStore()
-docstore.add_documents(nodes)
-
-from llama_index.core import VectorStoreIndex, StorageContext
-from llama_index.vector_stores.chroma import ChromaVectorStore
-import chromadb
-
-db = chromadb.PersistentClient(path="chroma_db")
-chroma_collection = db.get_or_create_collection("dense_vectors")
-vector_store = ChromaVectorStore(chroma_collection=chroma_collection)
-storage_context = StorageContext.from_defaults(
-    docstore=docstore, vector_store=vector_store
-    )
-index = VectorStoreIndex(nodes = nodes, storage_context=storage_context, show_progress = True)
-
-storage_context.docstore.persist("./docstore.json")
-
-index_retriever = index.as_retriever(similarity_top_k=2)
-import nest_asyncio
-nest_asyncio.apply()
-from llama_index.retrievers.bm25 import BM25Retriever
-bm25_retriever = BM25Retriever.from_defaults(
-    docstore=index.docstore,
-    similarity_top_k=2,
-    language = "portuguese",
-    verbose=True,
-    )
-
-from llama_index.core.retrievers import QueryFusionRetriever
-
-retriever = QueryFusionRetriever(
-    [index_retriever, bm25_retriever],
-    num_queries=1, #desativado = 1
-    mode="reciprocal_rerank",
-    use_async=True,
-    verbose=True,
-)
-
-from llama_index.core.storage.chat_store import SimpleChatStore
-from llama_index.core.memory import ChatMemoryBuffer
-chat_store = SimpleChatStore()
-chat_memory = ChatMemoryBuffer.from_defaults(
-    token_limit=3000,
-    chat_store=chat_store,
-    chat_store_key="user1",
-)
-from llama_index.core.query_engine import RetrieverQueryEngine
-query_engine = RetrieverQueryEngine.from_args(retriever)
-from llama_index.core.chat_engine import CondensePlusContextChatEngine
-chat_engine = CondensePlusContextChatEngine.from_defaults(
-    query_engine,
-    memory=chat_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}"
-        "\nInstrução: Use o histórico da conversa anterior, ou o contexto acima, para responder."
-        "No final da resposta, depois de uma quebra de linha escreva o nome do documento que contém a informação entre dois ||, como ||Documento Nome||"
-
-    ),
-)
-
-
-
-@app.route("/chat", methods=["POST"])
-def chat():
-    user_input = request.json.get("message", "")
-    if not user_input:
-        return jsonify({"error": "Mensagem vazia"}), 400
-
-    def generate_response():
-        try:
-            response = chat_engine.stream_chat(user_input)
-            for token in response.response_gen:
-                yield token  # Envia cada token
-            chat_store.persist(persist_path="chat_store.json")
-        except Exception as e:
-            yield f"Erro: {str(e)}"
-
-    return Response(generate_response(), content_type="text/plain")
-if __name__ == "__main__":
-    app.run(port=5001, debug=False)
+import os
+import logging
+import sys
+
+from flask import Flask, request, jsonify, Response
+# Inicializa o Flask
+app = Flask(__name__)
+
+logging.basicConfig(stream=sys.stdout, level=logging.INFO)
+
+from llama_index.llms.openai import OpenAI
+from llama_index.embeddings.openai import OpenAIEmbedding
+from llama_index.core import (
+    Settings,
+    SimpleDirectoryReader,
+    StorageContext,
+    Document,
+)
+
+Settings.llm = OpenAI(model="gpt-3.5-turbo")
+Settings.embed_model = OpenAIEmbedding(model_name="text-embedding-3-small")
+directory_path = "documentos"
+from llama_index.readers.file import PDFReader #concatenar todo o documento já vem nativo no pdfreader
+file_extractor = {".pdf": PDFReader(return_full_document = True)}
+from drive_downloader import GoogleDriveDownloader
+
+# ID da pasta no Drive e caminho local
+folder_id = "1n34bmh9rlbOtCvE_WPZRukQilKeabWsN"
+local_path = directory_path
+
+GoogleDriveDownloader().download_from_folder(folder_id, local_path)
+
+documents = SimpleDirectoryReader(
+    input_dir=directory_path,
+    file_extractor=file_extractor,
+    filename_as_id=True,
+    recursive=True
+).load_data()
+
+from document_creator import create_single_document_with_filenames
+document = create_single_document_with_filenames(directory_path = directory_path)
+documents.append(document)
+
+#from llama_index.core.ingestion import IngestionPipeline
+#ingestion pipeline vai entrar em uso quando adicionar o extrator de metadados
+from llama_index.core.node_parser import SentenceSplitter
+splitter = SentenceSplitter(chunk_size=1024, chunk_overlap=128)
+nodes = splitter.get_nodes_from_documents(documents)
+
+from llama_index.core.storage.docstore import SimpleDocumentStore
+docstore = SimpleDocumentStore()
+docstore.add_documents(nodes)
+
+from llama_index.core import VectorStoreIndex, StorageContext
+from llama_index.vector_stores.chroma import ChromaVectorStore
+import chromadb
+
+db = chromadb.PersistentClient(path="chroma_db")
+chroma_collection = db.get_or_create_collection("dense_vectors")
+vector_store = ChromaVectorStore(chroma_collection=chroma_collection)
+storage_context = StorageContext.from_defaults(
+    docstore=docstore, vector_store=vector_store
+    )
+index = VectorStoreIndex(nodes = nodes, storage_context=storage_context, show_progress = True)
+
+storage_context.docstore.persist("./docstore.json")
+
+index_retriever = index.as_retriever(similarity_top_k=2)
+import nest_asyncio
+nest_asyncio.apply()
+from llama_index.retrievers.bm25 import BM25Retriever
+bm25_retriever = BM25Retriever.from_defaults(
+    docstore=index.docstore,
+    similarity_top_k=2,
+    language = "portuguese",
+    verbose=True,
+    )
+
+from llama_index.core.retrievers import QueryFusionRetriever
+
+retriever = QueryFusionRetriever(
+    [index_retriever, bm25_retriever],
+    num_queries=1, #desativado = 1
+    mode="reciprocal_rerank",
+    use_async=True,
+    verbose=True,
+)
+
+
+from llama_index.core.memory import ChatMemoryBuffer
+from mysqlchatstore import MySQLChatStore
+chat_store = MySQLChatStore.from_params(
+    host=os.getenv("MYSQL_HOST"),
+    port=os.getenv("MYSQL_PORT"),
+    user=os.getenv("MYSQL_USER"),
+    password=os.getenv("MYSQL_PASSWORD"),
+    database=os.getenv("MYSQL_DATABASE"),
+    table_name=os.getenv("MYSQL_TABLE")
+)
+chat_memory = ChatMemoryBuffer.from_defaults(
+    token_limit=3000,
+    chat_store=chat_store,
+    chat_store_key="Sicoob", #Tendo algumas dificuldades ainda pra passar o user
+)
+from llama_index.core.query_engine import RetrieverQueryEngine
+query_engine = RetrieverQueryEngine.from_args(retriever)
+from llama_index.core.chat_engine import CondensePlusContextChatEngine
+chat_engine = CondensePlusContextChatEngine.from_defaults(
+    query_engine,
+    memory=chat_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}"
+        "\nInstrução: Use o histórico da conversa anterior, ou o contexto acima, para responder."
+        "No final da resposta, depois de uma quebra de linha escreva o nome do documento que contém a informação entre dois ||, como ||Documento Nome||"
+
+    ),
+)
+
+
+
+@app.route("/chat", methods=["POST"])
+def chat():
+    user_input = request.json.get("message", "")
+    if not user_input:
+        return jsonify({"error": "Mensagem vazia"}), 400
+
+    def generate_response():
+        try:
+            response = chat_engine.stream_chat(user_input)
+            for token in response.response_gen:
+                yield token  # Envia cada token
+        except Exception as e:
+            yield f"Erro: {str(e)}"
+
+    return Response(generate_response(), content_type="text/plain")
+if __name__ == "__main__":
+    app.run(port=5001, debug=False)

mysqlchatstore.py

@@ -0,0 +1,278 @@
+from typing import Optional, Any
+from sqlalchemy import create_engine, text
+from sqlalchemy.ext.asyncio import create_async_engine, AsyncSession
+from sqlalchemy.orm import sessionmaker
+from pydantic import Field
+import pymysql
+
+from llama_index.core.storage.chat_store import BaseChatStore
+from llama_index.core.llms import ChatMessage
+from llama_index.core.memory import ChatMemoryBuffer
+
+
+
+class MySQLChatStore(BaseChatStore):
+    """
+    Implementação de um ChatStore que armazena mensagens em uma tabela MySQL,
+    unindo a pergunta do usuário e a resposta do assistente na mesma linha.
+    """
+    table_name: Optional[str] = Field(default="chatstore", description="Nome da tabela MySQL.")
+
+    _session: Optional[sessionmaker] = None
+    _async_session: Optional[sessionmaker] = None
+
+    def __init__(self, session: sessionmaker, async_session: sessionmaker, table_name: str):
+        super().__init__(table_name=table_name.lower())
+        self._session = session
+        self._async_session = async_session
+        self._initialize()
+
+    @classmethod
+    def from_params(cls, host: str, port: str, database: str, user: str, password: str, table_name: str = "chatstore") -> "MySQLChatStore":
+        """
+        Cria o sessionmaker síncrono e assíncrono, retornando a instância da classe.
+        """
+        conn_str = f"mysql+pymysql://{user}:{password}@{host}:{port}/{database}"
+        async_conn_str = f"mysql+aiomysql://{user}:{password}@{host}:{port}/{database}"
+        session, async_session = cls._connect(conn_str, async_conn_str)
+        return cls(session=session, async_session=async_session, table_name=table_name)
+
+    @classmethod
+    def _connect(cls, connection_string: str, async_connection_string: str) -> tuple[sessionmaker, sessionmaker]:
+        """
+        Cria e retorna um sessionmaker síncrono e um sessionmaker assíncrono.
+        """
+        engine = create_engine(connection_string, echo=False)
+        session = sessionmaker(bind=engine)
+
+        async_engine = create_async_engine(async_connection_string)
+        async_session = sessionmaker(bind=async_engine, class_=AsyncSession)
+
+        return session, async_session
+
+    def _initialize(self):
+        """
+        Garante que a tabela exista, com colunas para armazenar user_input e response.
+        """
+        with self._session() as session:
+            session.execute(text(f"""
+                CREATE TABLE IF NOT EXISTS {self.table_name} (
+                    id INT AUTO_INCREMENT PRIMARY KEY,
+                    chat_store_key VARCHAR(255) NOT NULL,
+                    user_input TEXT,
+                    response TEXT,
+                    timestamp TIMESTAMP DEFAULT CURRENT_TIMESTAMP
+                )
+            """))
+            session.commit()
+
+    def get_keys(self) -> list[str]:
+        """
+        Retorna todas as chaves armazenadas.
+        """
+        with self._session() as session:
+            result = session.execute(text(f"""
+                SELECT DISTINCT chat_store_key FROM {self.table_name}
+            """))
+            return [row[0] for row in result.fetchall()]
+
+    def get_messages(self, key: str) -> list[ChatMessage]:
+        """
+        Retorna a conversa inteira (perguntas e respostas), na ordem de inserção (id).
+        Cada linha pode conter o user_input, o response ou ambos (caso já respondido).
+        """
+        with self._session() as session:
+            rows = session.execute(text(f"""
+                SELECT user_input, response
+                FROM {self.table_name}
+                WHERE chat_store_key = :key
+                ORDER BY id
+            """), {"key": key}).fetchall()
+
+            messages = []
+            for user_in, resp in rows:
+                if user_in is not None:
+                    messages.append(ChatMessage(role='user', content=user_in))
+                if resp is not None:
+                    messages.append(ChatMessage(role='assistant', content=resp))
+            return messages
+
+    def set_messages(self, key: str, messages: list[ChatMessage]) -> None:
+        """
+        Sobrescreve o histórico de mensagens de uma chave (apaga tudo e insere novamente).
+        Se quiser somente acrescentar, use add_message.
+        
+        Aqui, cada pergunta do usuário gera uma nova linha.
+        Assim que encontrar uma mensagem de assistente, atualiza essa mesma linha.
+        Se houver assistentes sem usuários, insere normalmente.
+        """
+        with self._session() as session:
+            # Limpa histórico anterior
+            session.execute(text(f"""
+                DELETE FROM {self.table_name} WHERE chat_store_key = :key
+            """), {"key": key})
+
+            # Reinsere na ordem
+            current_id = None
+            for msg in messages:
+                if msg.role == 'user':
+                    # Cria nova linha com user_input
+                    result = session.execute(text(f"""
+                        INSERT INTO {self.table_name} (chat_store_key, user_input)
+                        VALUES (:key, :ui)
+                    """), {"key": key, "ui": msg.content})
+                    # Pega o id do insert
+                    current_id = result.lastrowid
+
+                else:
+                    # Tenta atualizar a última linha se existir
+                    if current_id is not None:
+                        session.execute(text(f"""
+                            UPDATE {self.table_name}
+                            SET response = :resp
+                            WHERE id = :id
+                        """), {"resp": msg.content, "id": current_id})
+                        # Depois de atualizar a linha, zera o current_id
+                        current_id = None
+                    else:
+                        # Se não houver pergunta pendente, insere como nova linha
+                        session.execute(text(f"""
+                            INSERT INTO {self.table_name} (chat_store_key, response)
+                            VALUES (:key, :resp)
+                        """), {"key": key, "resp": msg.content})
+
+            session.commit()
+
+    def add_message(self, key: str, message: ChatMessage) -> None:
+        """
+        Acrescenta uma nova mensagem no fluxo. Se for do usuário, insere nova linha;
+        se for do assistente, tenta preencher a linha pendente que não tenha resposta.
+        """
+
+        with self._session() as session:
+            if message.role == 'user':
+                # Sempre cria uma nova linha para mensagens de usuário
+                insert_stmt = text(f"""
+                    INSERT INTO {self.table_name} (chat_store_key, user_input)
+                    VALUES (:key, :ui)
+                """)
+                session.execute(insert_stmt, {
+                    "key": key,
+                    "ui": message.content
+                })
+            else:
+                # Tenta encontrar a última linha sem resposta
+                
+                row = session.execute(text(f"""
+                    SELECT id
+                    FROM {self.table_name}
+                    WHERE chat_store_key = :key
+                      AND user_input IS NOT NULL
+                      AND response IS NULL
+                    ORDER BY id DESC
+                    LIMIT 1
+               """), {"key": key}).fetchone()
+
+                if row:
+                    # Atualiza com a resposta
+                    msg_id = row[0]
+                    
+                    update_stmt = text(f"""
+                        UPDATE {self.table_name}
+                        SET response = :resp
+                        WHERE id = :id
+                    """)
+                    session.execute(update_stmt, {
+                        "resp": message.content,
+                        "id": msg_id
+                    })
+                else:
+                    # Se não achar linha pendente, insere como nova
+                   
+                    insert_stmt = text(f"""
+                        INSERT INTO {self.table_name} (chat_store_key, response)
+                        VALUES (:key, :resp)
+                    """)
+                    session.execute(insert_stmt, {
+                        "key": key,
+                        "resp": message.content
+                    })
+
+            session.commit()
+            
+
+
+    def delete_messages(self, key: str) -> None:
+        """
+        Remove todas as linhas associadas a 'key'.
+        """
+        with self._session() as session:
+            session.execute(text(f"""
+                DELETE FROM {self.table_name} WHERE chat_store_key = :key
+            """), {"key": key})
+            session.commit()
+
+    def delete_last_message(self, key: str) -> Optional[ChatMessage]:
+        """
+        Apaga a última mensagem da conversa (considerando a ordem de inserção).
+        Se a última linha tiver pergunta e resposta, remove primeiro a resposta;
+        caso não exista resposta, remove a linha inteira.
+        """
+        with self._session() as session:
+            # Localiza a última linha
+            row = session.execute(text(f"""
+                SELECT id, user_input, response
+                FROM {self.table_name}
+                WHERE chat_store_key = :key
+                ORDER BY id DESC
+                LIMIT 1
+            """), {"key": key}).fetchone()
+
+            if not row:
+                return None
+
+            row_id, user_in, resp = row
+
+            # Se a linha tiver somente pergunta, apagamos a linha inteira.
+            # Se tiver também a resposta, apagamos só a parte do assistente.
+            if user_in and resp:
+                # Remove a resposta
+                session.execute(text(f"""
+                    UPDATE {self.table_name}
+                    SET response = NULL
+                    WHERE id = :id
+                """), {"id": row_id})
+                session.commit()
+                return ChatMessage(role='assistant', content=resp)
+            else:
+                # Deleta a linha inteira
+                session.execute(text(f"""
+                    DELETE FROM {self.table_name}
+                    WHERE id = :id
+                """), {"id": row_id})
+                session.commit()
+
+                if user_in:
+                    return ChatMessage(role='user', content=user_in)
+                elif resp:
+                    return ChatMessage(role='assistant', content=resp)
+                else:
+                    return None
+
+    def delete_message(self, key: str, idx: int) -> Optional[ChatMessage]:
+        """
+        Deleta a mensagem com base na ordem total do histórico. O índice 'idx' é
+        calculado após reconstruir a lista de ChatMessages (user e assistant).
+        """
+        messages = self.get_messages(key)
+        if idx < 0 or idx >= len(messages):
+            return None
+
+        removed = messages[idx]
+
+        # Agora precisamos traduzir 'idx' para saber qual registro no banco será modificado.
+        # É mais simples recriar todos os dados com set_messages sem a mensagem em 'idx':
+        messages.pop(idx)
+        self.set_messages(key, messages)
+
+        return removed

requirements.txt

@@ -1,11 +1,13 @@
-llama-index==0.12.12
-llama-index-retrievers-bm25==0.5.2
-llama-index-vector-stores-chroma==0.4.1
-llama-index-readers-google==0.6.0
-openpyxl==3.1.5
-flask==3.1.0
-streamlit==1.41.1
-streamlit-authenticator==0.4.1
-python-levenshtein==0.26.1
-streamlit_feedback
-fuzzywuzzy
+llama-index==0.12.12
+llama-index-retrievers-bm25==0.5.2
+llama-index-vector-stores-chroma==0.4.1
+llama-index-readers-google==0.6.0
+openpyxl==3.1.5
+flask==3.1.0
+streamlit==1.41.1
+streamlit-authenticator==0.4.1
+python-levenshtein==0.26.1
+streamlit_feedback
+fuzzywuzzy
+pymysql==1.1.1
+aiomysql==0.2.0