init

app.py

@@ -1,11 +1,13 @@
 import os
 import streamlit as st
-from src.indexing.document_processor import DocumentProcessor
+from src.indexing.document_processing import DocumentProcessor
 from src.indexing.vectore_store import VectorStoreManager
-from src.retriever.retriever import RetrieverManager
+from src.tools_retrieval.retriever import RetrieverManager
+from src.workflow import RAGWorkflow
+
 
 UPLOAD_FOLDER = "uploads/"
-PERSIST_DIRECTORY = "chroma_db/"
+PERSIST_DIRECTORY = "./chroma_db"
 os.makedirs(UPLOAD_FOLDER, exist_ok=True)
 os.makedirs(PERSIST_DIRECTORY, exist_ok=True)
 
@@ -15,6 +17,9 @@ if "retriever" not in st.session_state:
     st.session_state.retriever = None
 if "vector_store" not in st.session_state:
     st.session_state.vector_store = None
+if "workflow" not in st.session_state:
+    st.session_state.workflow = None
+
 
 st.set_page_config(
     page_title="RAG Chatbot",
@@ -23,36 +28,71 @@ st.set_page_config(
 )
 st.title("Agentic RAG Chatbot")
 
-
 with st.sidebar:
     st.header("PDF Upload")
     uploaded_file = st.file_uploader("Upload your PDF", type=["pdf"])
     st.info("Supported file type: PDF")
+    process_button = st.button("Process PDF")
 
-    if uploaded_file:
-        with st.spinner("Processing PDF..."):
-            
+    if uploaded_file and process_button:    
+        with st.spinner("Processing PDF..."):  
             file_path = os.path.join(UPLOAD_FOLDER, uploaded_file.name)
             with open(file_path, "wb") as f:
                 f.write(uploaded_file.getbuffer())
 
-            
             doc_processor = DocumentProcessor()
             chunks = doc_processor.load_and_split_pdf(file_path)
 
-            # Buat vector store
             vector_store_manager = VectorStoreManager()
-            vector_store = vector_store_manager.index_documents(
-                documents=chunks,
-                collection_name=uploaded_file.name,
-                persist_directory=PERSIST_DIRECTORY
-            )
+            vector_store = vector_store_manager.index_documents(chunks, uploaded_file.name, PERSIST_DIRECTORY)
+                    
             st.session_state.vector_store = vector_store
-
-            # Setup retriever
+            st.success("PDF processed and indexed successfully!")
+                    
             retriever_manager = RetrieverManager(vector_store)
-            base_retriever = retriever_manager.create_base_retriever()
-            compression_retriever = retriever_manager.create_compression_retriever(base_retriever)
-            st.session_state.retriever = compression_retriever
+            retriever_tool = retriever_manager.create_retriever(chunks)
+            st.session_state.retriever = retriever_tool
+            st.success("Retriever tool created successfully!")
+            rag_workflow = RAGWorkflow(retriever_tool)
+            workflow = rag_workflow.compile()
+            st.session_state.workflow = workflow
+            
+            
+
+# Display chat messages
+for message in st.session_state.messages:
+    with st.chat_message(message["role"]):
+        st.markdown(message["content"])
+
+if prompt := st.chat_input("Ask a question about your document"):
+    # Add user message to chat history
+    st.session_state.messages.append({"role": "user", "content": prompt})
+    with st.chat_message("user"):
+        st.markdown(prompt)    
+
+    # Generate response
+    with st.chat_message("assistant"):
+        if st.session_state.retriever is None:
+            final_response = "Please upload a PDF document first."
+        else:
+            with st.spinner("Thinking..."):
+                # Retrieve relevant documents
+                inputs = {
+                    "messages": [
+                        ("user", prompt),
+                    ]
+                }
+
+                # Generate response using workflow
+                if st.session_state.workflow is not None:
+                    response = st.session_state.workflow.invoke(inputs)
+                    final_response = response["messages"][-1].content
+                else:
+                    final_response = "Please upload a PDF document first."
+
+        st.markdown(final_response)
+        st.session_state.messages.append({"role": "assistant", "content": final_response})
 
-            st.success("File processed successfully!")
+# Add clear chat button
+if st.sidebar.button("Clear Chat"):
+    st.session_state.messages = []    

requirements.txt

@@ -1,11 +1,12 @@
-langchain 
+langchain
 langgraph
-langchain-huggingface 
-langchain-google-genai 
-google-ai-generativelanguage==0.6.15
-langchain-community 
+langchain-huggingface
+langchain-groq
+langchain-community
+scikit-learn
 langchain-chroma
-pypdf 
-tiktoken 
-rank_bm25 
-flashrank
+pypdf==5.1.0
+tiktoken  
+rank_bm25
+fastembed
+flashrank

src/indexing/vectore_store.py

@@ -1,23 +1,27 @@
 from langchain_huggingface import HuggingFaceEmbeddings
+from langchain_community.vectorstores import SKLearnVectorStore
 from langchain_chroma import Chroma
 
 class VectorStoreManager:
     def __init__(self, embedding_model="intfloat/multilingual-e5-small"):
         self.embeddings = HuggingFaceEmbeddings(model_name=embedding_model)
         
-    def create_vector_store(self, collection_name="my_collection", persist_directory=None):
+    def create_vector_store(self, collection_name, presist_directory):
         """Create a new vector store"""
-        store_params = {
-            "collection_name": collection_name,
-            "embedding_function": self.embeddings,
-        }
-        if persist_directory:
-            store_params["persist_directory"] = persist_directory
-            
-        return Chroma(**store_params)
+#        vector_store = SKLearnVectorStore.from_documents(
+#            documents=documents,
+#            embedding=self.embeddings,
+#        )
+        vector_store = Chroma(
+            collection_name=collection_name,
+            embedding_function=self.embeddings,
+            persist_directory=presist_directory,  # Where to save data locally, remove if not necessary
+        )
+        
+        return vector_store
     
-    def index_documents(self, documents, collection_name="my_collection", persist_directory=None):
+    def index_documents(self, documents, collection_name, presist_directory):
         """Index documents into vector store"""
-        vector_store = self.create_vector_store(collection_name, persist_directory)
+        vector_store = self.create_vector_store(collection_name, presist_directory)
         vector_store.add_documents(documents=documents)
-        return vector_store
+        return vector_store 

src/llm/llm_interface.py

@@ -0,0 +1,10 @@
+import os
+from langchain_groq import ChatGroq
+     
+llm_groq = ChatGroq(
+    model="llama3-8b-8192",
+    temperature=0.1,
+    api_key=os.getenv("GROQ_API_KEY"),
+    max_retries=3,
+    streaming=True,
+)

src/state.py

@@ -0,0 +1,7 @@
+from typing import Annotated, Sequence
+from typing_extensions import TypedDict
+from langchain_core.messages import BaseMessage
+from langgraph.graph.message import add_messages
+
+class AgentState(TypedDict):
+    messages: Annotated[Sequence[BaseMessage], add_messages]

src/{retriever → tools_retrieval}/retriever.py

@@ -1,7 +1,10 @@
-from langchain.retrievers import BM25Retriever, EnsembleRetriever
+from langchain_community.retrievers import BM25Retriever
+from langchain.retrievers import EnsembleRetriever
 from langchain.retrievers import ContextualCompressionRetriever
 from langchain.retrievers.document_compressors import FlashrankRerank
 
+from langchain.tools.retriever import create_retriever_tool
+
 
 class RetrieverManager:
     def __init__(self, vector_store):
@@ -31,4 +34,13 @@ class RetrieverManager:
         return ContextualCompressionRetriever(
             base_compressor=compressor,
             base_retriever=base_retriever
+        )
+
+    def create_retriever(self, documents):
+        base_retriever = self.create_ensemble_retriever(documents)
+        compression_retriever = self.create_compression_retriever(base_retriever=base_retriever)
+        return create_retriever_tool(
+            compression_retriever,
+            "retrieve_docs",
+            "use tools for search through the user's provided documents and return relevant information about user query.",
         )

src/workflow.py

@@ -0,0 +1,128 @@
+from langchain import hub
+from langchain_core.output_parsers import StrOutputParser
+from langchain_core.messages import HumanMessage
+from typing import Literal
+from pydantic import BaseModel, Field
+from langchain_core.prompts import PromptTemplate
+
+from langgraph.graph import END, StateGraph, START
+from langgraph.prebuilt import ToolNode, tools_condition
+from .state import AgentState
+from src.llm.llm_interface import llm_groq
+
+
+class GradeDocs(BaseModel):
+    binary_score: str = Field(description="Relevance score 'yes' or 'no'")
+
+
+
+class RAGWorkflow:
+    def __init__(self, retriever_tool):
+        self.workflow = StateGraph(AgentState)
+        self.tools = [retriever_tool]
+        self.retrieve = ToolNode([retriever_tool])
+        self._setup_nodes()
+        self._setup_edges()
+
+    def _setup_nodes(self):
+        self.workflow.add_node("agent", self._agent_node)
+        self.workflow.add_node("retrieve", self.retrieve)
+        self.workflow.add_node("generate", self._generator_node)
+        
+        self.workflow.add_node("rewrite", self._rewrite_node)
+
+    def _setup_edges(self):
+        self.workflow.add_edge(START, "agent")
+        self.workflow.add_conditional_edges(
+            "agent",
+            tools_condition,
+            {
+                "tools": "retrieve",
+                END: END
+            }
+        )
+        
+        self.workflow.add_conditional_edges(
+            "retrieve",
+            self._grade_docs,
+        )
+        self.workflow.add_edge("generate", END)
+        self.workflow.add_edge("rewrite", "agent")
+        
+    def compile(self):
+        return self.workflow.compile()
+    
+    
+    def _agent_node(self, state):
+        print("---CALL AGENT---")
+        messages = state["messages"]
+            
+        model = llm_groq.bind_tools(self.tools)
+        response = model.invoke(messages[0].content)
+        return {"messages": [response]}
+
+
+    def _generator_node(self, state):
+        print("---GENERATE---")
+        messages = state["messages"]
+        question = messages[0].content
+        docs = messages[-1].content
+
+        prompt = hub.pull("rlm/rag-prompt")
+        rag_chain = prompt | llm_groq | StrOutputParser()
+            
+        response = rag_chain.invoke({"context": docs, "question": question})
+        return {"messages": [response]}
+
+
+
+    def _rewrite_node(self, state):
+        print("---REWRITE---")
+        messages = state["messages"]
+        question = messages[0].content
+
+        msg = [
+            HumanMessage(
+                content=f""" \n 
+                    Look at the input and try to reason about the underlying semantic intent / meaning. \n 
+                    Here is the initial question:
+                    \n ------- \n
+                    {question} 
+                    \n ------- \n
+                    Formulate an improved question: """,
+            )
+        ]
+
+        response = llm_groq.invoke(msg)
+        return {"messages": [response]}
+
+
+
+    def _grade_docs(self, state):
+        print("---CHECK RELEVANCE---")
+            
+        llm_with_tool = llm_groq.with_structured_output(GradeDocs)
+            
+        prompt = PromptTemplate(
+        template="""You are a grader assessing relevance of a retrieved document to a user question. \n 
+            Here is the retrieved document: \n\n {context} \n\n
+            Here is the user question: {question} \n
+            If the document contains keyword(s) or semantic meaning related to the user question, grade it as relevant. \n
+            Give a binary score 'yes' or 'no' score to indicate whether the document is relevant to the question.""",
+            input_variables=["context", "question"],
+        )
+
+        chain = prompt | llm_with_tool
+
+        messages = state["messages"]
+        question = messages[0].content
+        docs = messages[-1].content
+
+        scored_result = chain.invoke({"question": question, "context": docs})
+            
+        if scored_result.binary_score == "yes":
+            print("---DECISION: DOCS RELEVANT---")
+            return "generate"
+        print("---DECISION: DOCS NOT RELEVANT---")
+        return "rewrite"
+