Modified document loading logic

app.py

@@ -3,7 +3,6 @@ from graph import build_graph
 from utils import get_retriever, load_vectorstore_from_text
 from pypdf import PdfReader
 import hashlib
-from transformers import pipeline, BartForConditionalGeneration, BartTokenizer
 
 # --- Page Config ---
 st.set_page_config(page_title="LangGraph RAG Chatbot", layout="wide")
@@ -74,30 +73,20 @@ with st.sidebar:
         else:
             st.warning("Uploaded file is empty or could not be read.")
 
-    # Show current mode
-    if "retriever" in st.session_state and st.session_state.retriever:
+    # Show current mode with proper guard
+    if ("retriever" in st.session_state and 
+        st.session_state.retriever is not None and 
+        "raw_text" in st.session_state):
         st.info("📄 **RAG Mode**: Answering from uploaded document")
     else:
         st.info("💬 **General Chat Mode**: No document loaded")
 
-# --- Initialize Summarizer ---
-if "summarizer" not in st.session_state:
-    tokenizer = BartTokenizer.from_pretrained("facebook/bart-large-cnn")
-    model = BartForConditionalGeneration.from_pretrained("facebook/bart-large-cnn")
-    st.session_state.summarizer = pipeline(
-        "summarization",
-        model=model,
-        tokenizer=tokenizer,
-        device=-1
-    )
-
 # --- Build Graph ---
 if "graph" not in st.session_state or st.session_state.get("graph_model") != model_type:
     try:
         st.session_state.graph = build_graph(
             model_type=model_type,
-            retriever=st.session_state.get("retriever"),
-            summarizer=st.session_state.get("summarizer")
+            retriever=st.session_state.get("retriever")
         )
         st.session_state.graph_model = model_type
     except Exception as e:
@@ -108,49 +97,55 @@ if "graph" not in st.session_state or st.session_state.get("graph_model") != mod
 if "history" not in st.session_state:
     st.session_state.history = []
 
+# --- Initialize current query to handle input clearing ---
+if "current_query" not in st.session_state:
+    st.session_state.current_query = ""
+
 # --- Query Input ---
-query_input = st.text_input("💬 Ask a question:")
+query_input = st.text_input("💬 Ask a question:", value=st.session_state.current_query, key="current_query")
 
 send_triggered = st.button("Send")
 
-if send_triggered:
-    if query_input.strip():
-        formatted_history = [(q, r) for q, r, _ in st.session_state.history]
-
-        with st.spinner("Generating response..."):
-            try:
-                result = st.session_state.graph(
-                    query=query_input,
-                    temperature=temperature,
-                    raw_text=st.session_state.get("raw_text"),
-                    history=formatted_history,
-                    retriever_override=st.session_state.get("retriever")
-                )
-
-                response = result.get("response", "No response generated.")
-                retrieved_docs = result.get("retrieved_docs", [])
-
-                st.markdown("### 🤖 Response")
-                st.markdown(response)
-
-                # Save to history
-                st.session_state.history.append((query_input, response, retrieved_docs))
-
-                # Show retrieved docs if available
-                if retrieved_docs:
-                    with st.expander("📄 Retrieved Chunks"):
-                        for j, doc in enumerate(retrieved_docs):
-                            content = getattr(doc, "text", str(doc))
-                            st.markdown(f"**Chunk {j+1}:**")
-                            st.code(content.strip(), language="markdown")
-
-                # Clear the input field by rerunning widget with empty value
-                st.rerun()
-
-            except Exception as e:
-                st.error(f"Query failed: {e}")
-    else:
-        st.warning("Please enter a question.")
+if send_triggered and query_input.strip():
+    formatted_history = [(q, r) for q, r, _ in st.session_state.history]
+
+    with st.spinner("Generating response..."):
+        try:
+            result = st.session_state.graph(
+                query=query_input,
+                temperature=temperature,
+                raw_text=st.session_state.get("raw_text"),
+                history=formatted_history,
+                retriever_override=st.session_state.get("retriever")
+            )
+
+            response = result.get("response", "No response generated.")
+            retrieved_docs = result.get("retrieved_docs", [])
+
+            # Display response immediately
+            st.markdown("### 🤖 Response")
+            st.markdown(response)
+
+            # Show retrieved docs if available
+            if retrieved_docs:
+                with st.expander("📄 Retrieved Chunks"):
+                    for j, doc in enumerate(retrieved_docs):
+                        content = getattr(doc, "text", str(doc))
+                        st.markdown(f"**Chunk {j+1}:**")
+                        st.code(content.strip(), language="markdown")
+
+            # Save to history after displaying
+            st.session_state.history.append((query_input, response, retrieved_docs))
+
+            # Clear the input field
+            st.session_state.current_query = ""
+            st.rerun()
+
+        except Exception as e:
+            st.error(f"Query failed: {e}")
+
+elif send_triggered and not query_input.strip():
+    st.warning("Please enter a question.")
 
 # --- Chat History Display ---
 if st.session_state.history:
@@ -168,5 +163,4 @@ if st.session_state.history:
 # --- Clear Chat ---
 if st.sidebar.button("🗑️ Clear Chat History"):
     st.session_state.history = []
-    st.session_state.current_query = ""
-    st.rerun()
+    st.session_state.current_query = ""

graph.py

@@ -6,13 +6,9 @@ from typing import TypedDict, Optional, List
 from llama_index.core.schema import Document
 from langchain_google_genai import ChatGoogleGenerativeAI
 from langchain_openai import ChatOpenAI
-from llama_index.core import VectorStoreIndex
 from llama_index.core.retrievers import BaseRetriever
-from langchain_groq import ChatGroq
-from transformers import pipeline, BartForConditionalGeneration, BartTokenizer
 from langgraph.graph import StateGraph, END
 
-# --- 1. Define the State for the Graph ---
 class GraphState(TypedDict):
     query: str
     response: Optional[str]
@@ -21,9 +17,77 @@ class GraphState(TypedDict):
     history: list
     retriever: Optional[BaseRetriever]
     llm: any
-    summarizer: Optional[any]
 
-# --- 2. Define Graph Nodes ---
+def _select_chunks_by_relevance_and_position(docs: List[Document], query: str, target_count: int) -> List[Document]:
+    """Select chunks based on semantic relevance + positional importance without summarization."""
+    if len(docs) <= target_count:
+        return docs
+    
+    selected = []
+    if len(docs) >= 2:
+        selected = [docs[0], docs[-1]]  
+        remaining_slots = target_count - 2
+        middle_docs = docs[1:-1]
+    else:
+        remaining_slots = target_count
+        middle_docs = docs[:]
+    
+    if remaining_slots > 0 and middle_docs:
+        query_words = set(query.lower().split())
+        scored_docs = []
+        
+        for doc in middle_docs:
+            content = doc.text.lower()
+            relevance_score = sum(content.count(word) for word in query_words)
+            length_bonus = len(doc.text) / 1000  
+            total_score = relevance_score + length_bonus
+            scored_docs.append((total_score, doc))
+        
+        scored_docs.sort(key=lambda x: x[0], reverse=True)
+        selected.extend([doc for _, doc in scored_docs[:remaining_slots]])
+    
+    return selected[:target_count]
+
+def _create_context_with_priorities(docs: List[Document], max_chars: int) -> str:
+    """Create context by prioritizing important chunks without summarization."""
+    if not docs:
+        return ""
+    
+    contexts = []
+    total_chars = 0
+    
+    priority_docs = []
+    
+    if len(docs) >= 2:
+        priority_docs.append(("BOUNDARY", docs[0]))
+        priority_docs.append(("BOUNDARY", docs[-1]))
+        
+        middle_docs = docs[1:-1] if len(docs) > 2 else []
+        middle_docs.sort(key=lambda d: len(d.text), reverse=True)
+        priority_docs.extend([("CONTENT", doc) for doc in middle_docs])
+    else:
+        priority_docs = [("CONTENT", doc) for doc in docs]
+    
+    for priority_type, doc in priority_docs:
+        content = doc.text
+        
+        if total_chars + len(content) > max_chars:
+            remaining_chars = max_chars - total_chars
+            if remaining_chars > 200:  
+                truncated = content[:remaining_chars]
+                last_period = truncated.rfind('.')
+                if last_period > remaining_chars * 0.8:  
+                    truncated = truncated[:last_period + 1]
+                truncated += "...[truncated]"
+                contexts.append(f"[{priority_type}] {truncated}")
+            break
+        
+        contexts.append(f"[{priority_type}] {content}")
+        total_chars += len(content)
+    
+    return "\n\n---\n\n".join(contexts)
+
+
 def router_node(state: GraphState) -> GraphState:
     """
     Router that determines the next step based on available retriever.
@@ -66,13 +130,11 @@ AI:"""
 
 def retrieve_node(state: GraphState) -> GraphState:
     """
-    Retrieves relevant documents from the vector store based on the query.
-    Summarizes context if too long, or truncates at sentence boundaries.
+    Retrieves relevant documents and creates intelligent context without summarization.
     """
     print("---NODE: RETRIEVE---")
     query = state["query"]
     retriever = state["retriever"]
-    summarizer = state["summarizer"]
     history = state["history"]
 
     context = ""
@@ -80,43 +142,43 @@ def retrieve_node(state: GraphState) -> GraphState:
     
     try:
         q_len = len(query.split())
-        top_k = 3 if q_len < 5 else (5 if q_len < 15 else 8)
+        if q_len < 5:
+            top_k = 5  
+        elif q_len < 15:
+            top_k = 8
+        else:
+            top_k = 12
 
         retrieved_docs = retriever.retrieve(query)
+        print(f"Retrieved {len(retrieved_docs)} documents")
 
         if retrieved_docs:
-            context = "\n\n---\n\n".join([doc.text for doc in retrieved_docs])
-        
-        if history:
-            history_context = "\n\n".join([f"Human: {q}\nAI: {a}" for q, a in history])
-            context = f"{context}\n\n--- Chat History ---\n{history_context}"
-
-        MAX_CONTEXT_CHARS = 4000
-        if len(context) > MAX_CONTEXT_CHARS:
-            try:
-                print("---CONTEXT TOO LONG, SUMMARIZING---")
-                summary_result = summarizer(
-                    context,
-                    max_length=500,
-                    min_length=150,
-                    do_sample=False
-                )
-                context = summary_result[0].get("summary_text", context[:MAX_CONTEXT_CHARS])
-            except Exception as e:
-                print(f"Summarizer failed: {e}")
-                sentences = re.split(r'(?<=[.!?]) +', context)
-                truncated = []
-                total_len = 0
-                for sent in sentences:
-                    if total_len + len(sent) > MAX_CONTEXT_CHARS:
-                        break
-                    truncated.append(sent)
-                    total_len += len(sent)
-                context = " ".join(truncated)
+            max_chunks = min(len(retrieved_docs), top_k)
+            selected_docs = _select_chunks_by_relevance_and_position(
+                retrieved_docs, query, max_chunks
+            )
+            
+            print(f"Selected {len(selected_docs)} chunks for context")
+            
+            MAX_CONTEXT_CHARS = 6000  
+            doc_context = _create_context_with_priorities(selected_docs, MAX_CONTEXT_CHARS)
+            
+            if history and len(doc_context) < MAX_CONTEXT_CHARS * 0.8:
+                history_context = "\n\n".join([f"Human: {q}\nAI: {a}" for q, a in history[-3:]])  # Last 3 exchanges
+                remaining_chars = MAX_CONTEXT_CHARS - len(doc_context)
+                if len(history_context) <= remaining_chars:
+                    context = f"{doc_context}\n\n--- Recent Chat History ---\n{history_context}"
+                else:
+                    context = doc_context
+            else:
+                context = doc_context
+        else:
+            context = "No relevant content found in the document."
             
     except Exception as e:
         print(f"Error in retrieve_node: {e}")
         context = f"Retriever failed: {str(e)}"
+        retrieved_docs = []
     
     return {"retrieved_docs": retrieved_docs, "context": context}
 
@@ -157,7 +219,7 @@ Answer:"""
 
     return {"response": response_text.strip()}
 
-# --- 3. Define the Router Logic ---
+
 def route_query(state: GraphState) -> str:
     """
     Checks if a retriever is available in the state to decide the next step.
@@ -169,11 +231,10 @@ def route_query(state: GraphState) -> str:
         print("---ROUTING: No PDF, routing to general_chat_node---")
         return "general_chat"
 
-# --- 4. Build the Graph ---
-def build_graph(model_type: str = "groq", retriever=None, summarizer=None):
+
+def build_graph(model_type: str = "groq", retriever=None):
     """
-    Builds the workflow graph with LLM, retriever, and optional summarizer.
-    If summarizer not provided, initializes a default HuggingFace summarizer.
+    Builds the workflow graph with LLM and retriever.
     """
     
     def make_llm(temp: float):
@@ -187,6 +248,7 @@ def build_graph(model_type: str = "groq", retriever=None, summarizer=None):
                 api_key=api_key,
                 temperature=temp,
             )
+            
         elif model_type == "gemini":
             api_key = os.getenv("GEMINI_API_KEY", "").strip()
             if not api_key:
@@ -196,16 +258,9 @@ def build_graph(model_type: str = "groq", retriever=None, summarizer=None):
                 api_key=api_key,
                 temperature=temp,
             )
+        else:
+            raise ValueError("Invalid model_type. Choose 'groq' or 'gemini'.")
     
-    def get_default_summarizer():
-        tokenizer = BartTokenizer.from_pretrained("facebook/bart-large-cnn")
-        model = BartForConditionalGeneration.from_pretrained("facebook/bart-large-cnn")
-        return pipeline("summarization", model=model, tokenizer=tokenizer, device=-1)
-
-    if summarizer is None:
-        print("---NO SUMMARIZER PROVIDED, USING DEFAULT (facebook/bart-large-cnn)---")
-        summarizer = get_default_summarizer()
-
     workflow = StateGraph(GraphState)
 
     workflow.add_node("router", router_node)
@@ -238,7 +293,6 @@ def build_graph(model_type: str = "groq", retriever=None, summarizer=None):
             "retriever": active_retriever.as_retriever() if active_retriever else None,
             "history": history or [],
             "llm": llm,
-            "summarizer": summarizer,
             "response": None,
             "retrieved_docs": None,
             "context": None,

requirements.txt

@@ -1,7 +1,5 @@
 # Core LLM + Transformers
-transformers[sentencepiece]>=4.40.0
 sentence-transformers>=2.6.0
-accelerate>=0.30.0
 
 # LangChain + LangGraph
 langchain>=0.2.1