Upload 3 files

app.py

@@ -0,0 +1,103 @@
+import streamlit as st
+from langchain_community.document_loaders import PyPDFLoader
+from langchain_core.messages import HumanMessage, AIMessageChunk, AIMessage
+from langchain_huggingface import HuggingFaceEmbeddings
+from langchain.text_splitter import RecursiveCharacterTextSplitter
+from langchain_core.vectorstores import InMemoryVectorStore
+import os
+from langchain_core.chat_history import InMemoryChatMessageHistory, BaseChatMessageHistory
+import time
+from graph import get_graph
+
+if 'read_file' not in st.session_state:
+    st.session_state.read_file = False
+    st.session_state.retriever = None
+
+if 'chat_history' not in st.session_state:
+    st.session_state.chat_history = {}
+    st.session_state.first_msg = True
+
+def get_session_by_id(session_id: str) -> BaseChatMessageHistory:
+    if session_id not in st.session_state.chat_history:
+        st.session_state.chat_history[session_id] = InMemoryChatMessageHistory()
+        return st.session_state.chat_history[session_id]
+    return st.session_state.chat_history[session_id]
+
+if not st.session_state.read_file:
+    st.title('🤓 Upload your PDF to talk with it', anchor=False)
+    file = st.file_uploader('Upload a PDF file', type='pdf')
+    if file:
+        with st.status('🤗 Booting up the things!', expanded=True):
+            with st.spinner('📁 Uploading the PDF...', show_time=True):
+                with open('file.pdf', 'wb') as f:
+                    f.write(file.read())
+                    loader = PyPDFLoader('file.pdf')
+                    documents = loader.load_and_split(RecursiveCharacterTextSplitter(chunk_size=2000, chunk_overlap=200))
+                st.success('📁 File uploaded successfully!!!')
+            with st.spinner('🧐 Reading the file...', show_time=True):
+                vstore = InMemoryVectorStore.from_documents(documents, HuggingFaceEmbeddings(model_name='all-MiniLM-L6-v2'))
+                st.session_state.retriever = vstore.as_retriever()
+                st.success('🧐 File read successfully!!!')
+            os.remove('file.pdf')
+            with st.spinner('😴 Waking up the LLM...', show_time=True):
+                st.session_state.graph = get_graph(st.session_state.retriever)
+                st.success('😁 LLM awakened!!!')
+            st.balloons()
+        placeholder = st.empty()
+        for _ in range(5, -1, -1):
+            placeholder.write(f'⏳ Chat starting in 0{_} sec.')
+            time.sleep(1)
+        st.session_state.read_file = True
+        st.rerun()
+
+if st.session_state.read_file:
+
+    st.title('🤗 DocAI', anchor=False)
+    st.subheader('Chat with your document!', anchor=False)
+
+    if st.session_state.first_msg:
+        st.session_state.first_msg = False
+        get_session_by_id('chat42').add_message(AIMessage(content='Hello, how are you? How about we talk about the '
+                                                                  'document you sent me to read?'))
+
+    for msg in get_session_by_id('chat42').messages:
+        with st.chat_message(name='user' if isinstance(msg, HumanMessage) else 'ai'):
+            st.write(msg.content)
+
+    prompt = st.chat_input('Try to ask something about your file!')
+    if prompt:
+        with st.chat_message(name='user'):
+            st.write(prompt)
+
+        response = st.session_state.graph.stream(
+            {
+                'question': prompt,
+                'scratchpad': None,
+                'answer': None,
+                'next_node': None,
+                'history': get_session_by_id('chat42').messages,
+            },
+            stream_mode='messages'
+        )
+
+        get_session_by_id('chat42').add_message(HumanMessage(content=prompt))
+
+        def get_message():
+            for chunk, _ in response:
+                if chunk.content and isinstance(chunk, AIMessageChunk):
+                    yield chunk.content
+
+        with st.chat_message(name='ai'):
+            full_response = ''
+            placeholder = st.empty()
+
+            for msg in get_message():
+                full_response += msg
+                if '</tool>' in full_response:
+                    full_response = ''
+                    continue
+                if '<tool>' in full_response:
+                    continue
+                placeholder.write(full_response)
+        print(full_response)
+        get_session_by_id('chat42').add_message(AIMessage(content=full_response))

graph.py

@@ -0,0 +1,81 @@
+from utils import MainState, generate_uuid, llm
+from langchain_core.messages import AIMessage, ToolMessage, HumanMessage
+from langchain_core.prompts import ChatPromptTemplate
+from langgraph.graph import StateGraph, START, END
+import re
+
+def get_graph(retriever):
+    def retriever_node(state: MainState):
+        return {
+            'question': state['question'],
+            'scratchpad': state['scratchpad'] + [ToolMessage(content=retriever.invoke(state['question'].content),
+                                                             tool_call_id=state['scratchpad'][-1].tool_call_id)],
+            'answer': state['answer'],
+            'next_node': 'model_node',
+            'history': state['history']
+        }
+
+    import re
+
+    def model_node(state: MainState):
+        prompt = ChatPromptTemplate.from_template(
+            """
+            Você é um assistente de IA. Responda à pergunta abaixo da forma mais precisa possível.
+
+            Caso não tenha informações para responder à pergunte **retorne apenas** uma resposta no seguinte formato:
+            <tool>retriever</tool>,
+            ao fazer isso a task será repassada para um agente que irá complementar as informações.
+            Se a pergunta puder ser respondida sem acessar documentos enviados, forneça uma resposta **concisa e objetiva**, com no máximo três sentenças.
+
+            ### Contexto:
+            - Bloco de Notas: {scratchpad}
+            - Histórico de Conversas: {chat_history}
+
+            **Pergunta:** {question}
+            """
+        )
+
+        if isinstance(state['question'], str):
+            state['question'] = HumanMessage(content=state['question'])
+
+        qa_chain = prompt | llm
+
+        response = qa_chain.invoke({'question': state['question'].content,
+                                    'scratchpad': state['scratchpad'],
+                                    'chat_history': [
+                                        f'AI: {msg.content}' if isinstance(msg, AIMessage) else f'Human: {msg.content}'
+                                        for msg in state['history']],
+                                    })
+
+        if '<tool>' in response.content:
+            return {
+                'question': state['question'],
+                'scratchpad': state['scratchpad'] + [AIMessage(content='', tool_call_id=generate_uuid())] if state[
+                    'scratchpad'] else [AIMessage(content='', tool_call_id=generate_uuid())],
+                'answer': state['answer'],
+                'next_node': 'retriever',
+                'history': state['history']
+            }
+
+        # print(state['scratchpad'])
+        return {
+            'question': state['question'],
+            'scratchpad': state['scratchpad'],
+            'answer': response,
+            'next_node': END,
+            'history': state['history'] + [HumanMessage(content=state['question'].content), response]
+        }
+
+    def next_node(state: MainState):
+        return state['next_node']
+
+    graph = StateGraph(MainState)
+    graph.add_node('model', model_node)
+    graph.add_node('retriever', retriever_node)
+    graph.add_edge(START, 'model')
+    graph.add_edge('retriever', 'model')
+    graph.add_conditional_edges('model', next_node)
+
+    chain = graph.compile()
+
+    return chain

utils.py

@@ -0,0 +1,33 @@
+import uuid
+from typing import TypedDict
+from langchain_core.messages import AIMessage, HumanMessage, ToolMessage
+from langchain_openai import ChatOpenAI
+import os
+import re
+from dotenv import load_dotenv
+
+load_dotenv()
+
+class MainState(TypedDict):
+    question: HumanMessage | str| None
+    scratchpad: list[AIMessage | ToolMessage] | None
+    answer: AIMessage | None
+    next_node: str | None
+    history: list[HumanMessage | AIMessage]
+
+llm = ChatOpenAI(
+    model="gpt-4o-mini",
+    temperature=0,
+    api_key=os.environ.get('OPENAI_API_KEY'), # Insira sua chave aqui
+)
+
+def generate_uuid():
+    return str(uuid.uuid4())
+
+def post_process(message: AIMessage) -> AIMessage:
+    matches = re.findall(r"\[SOT\](.*)\[EOT\]", message.content, re.DOTALL)
+    matches = matches[0] if matches else None
+
+    if matches:
+        return AIMessage(content='', additional_kwargs={'custom_tool_call': matches} ,tool_call_id=generate_uuid())
+    return AIMessage(content=message.content)