Update streamlit_app.py

streamlit_app.py

@@ -1,151 +1,172 @@
-"""Streamlit UI for Agentic RAG System - Simplified Version"""
+import uuid
 
 import streamlit as st
-from pathlib import Path
-import sys
-import time
-
-# Add src to path
-sys.path.append(str(Path(__file__).parent))
-
-from src.config.config import Config
-from src.document_ingestion.document_processor import DocumentProcessor
-from src.vectorstore.vectorstore import VectorStore
-from src.graph_builder.graph_builder import GraphBuilder
-
-# Page configuration
-st.set_page_config(
-    page_title="🤖 RAG Search",
-    page_icon="🔍",
-    layout="centered"
+from langchain_core.messages import AIMessage, HumanMessage, ToolMessage
+
+from src.langraph_rag_backend import (
+    chatbot,
+    ingest_pdf,
+    retrieve_all_threads,
+    thread_document_metadata,
 )
 
-# Simple CSS
-st.markdown("""
-    <style>
-    .stButton > button {
-        width: 100%;
-        background-color: #4CAF50;
-        color: white;
-        font-weight: bold;
+
+# =========================== Utilities ===========================
+def generate_thread_id():
+    return uuid.uuid4()
+
+
+def reset_chat():
+    thread_id = generate_thread_id()
+    st.session_state["thread_id"] = thread_id
+    add_thread(thread_id)
+    st.session_state["message_history"] = []
+
+
+def add_thread(thread_id):
+    if thread_id not in st.session_state["chat_threads"]:
+        st.session_state["chat_threads"].append(thread_id)
+
+
+def load_conversation(thread_id):
+    state = chatbot.get_state(config={"configurable": {"thread_id": thread_id}})
+    return state.values.get("messages", [])
+
+
+# ======================= Session Initialization ===================
+if "message_history" not in st.session_state:
+    st.session_state["message_history"] = []
+
+if "thread_id" not in st.session_state:
+    st.session_state["thread_id"] = generate_thread_id()
+
+if "chat_threads" not in st.session_state:
+    st.session_state["chat_threads"] = retrieve_all_threads()
+
+if "ingested_docs" not in st.session_state:
+    st.session_state["ingested_docs"] = {}
+
+add_thread(st.session_state["thread_id"])
+
+thread_key = str(st.session_state["thread_id"])
+thread_docs = st.session_state["ingested_docs"].setdefault(thread_key, {})
+threads = st.session_state["chat_threads"][::-1]
+selected_thread = None
+
+# ============================ Sidebar ============================
+st.sidebar.title("LangGraph PDF Chatbot")
+st.sidebar.markdown(f"**Thread ID:** `{thread_key}`")
+
+if st.sidebar.button("New Chat", use_container_width=True):
+    reset_chat()
+    st.rerun()
+
+if thread_docs:
+    latest_doc = list(thread_docs.values())[-1]
+    st.sidebar.success(
+        f"Using `{latest_doc.get('filename')}` "
+        f"({latest_doc.get('chunks')} chunks from {latest_doc.get('documents')} pages)"
+    )
+else:
+    st.sidebar.info("No PDF indexed yet.")
+
+uploaded_pdfs = st.sidebar.file_uploader("Upload a PDF for this chat", type=["pdf"], accept_multiple_files=True)
+if uploaded_pdfs:
+    for uploaded_pdf in uploaded_pdfs:
+        if uploaded_pdf.name in thread_docs:
+            st.sidebar.info(f"`{uploaded_pdf.name}` already processed for this chat.")
+        else:
+            with st.sidebar.status("Indexing PDF…", expanded=True) as status_box:
+                summary = ingest_pdf(
+                    uploaded_pdf.getvalue(),
+                    thread_id=thread_key,
+                    filename=uploaded_pdf.name,
+                )
+                thread_docs[uploaded_pdf.name] = summary
+                status_box.update(label="✅ PDF indexed", state="complete", expanded=False)
+
+st.sidebar.subheader("Past conversations")
+if not threads:
+    st.sidebar.write("No past conversations yet.")
+else:
+    for thread_id in threads:
+        if st.sidebar.button(str(thread_id), key=f"side-thread-{thread_id}"):
+            selected_thread = thread_id
+
+# ============================ Main Layout ========================
+st.title("Multi Utility Chatbot")
+
+# Chat area
+for message in st.session_state["message_history"]:
+    with st.chat_message(message["role"]):
+        st.text(message["content"])
+
+user_input = st.chat_input("Ask about your document or use tools")
+
+if user_input:
+    st.session_state["message_history"].append({"role": "user", "content": user_input})
+    with st.chat_message("user"):
+        st.text(user_input)
+
+    CONFIG = {
+        "configurable": {"thread_id": thread_key},
+        "metadata": {"thread_id": thread_key},
+        "run_name": "chat_turn",
     }
-    </style>
-""", unsafe_allow_html=True)
-
-def init_session_state():
-    """Initialize session state variables"""
-    if 'rag_system' not in st.session_state:
-        st.session_state.rag_system = None
-    if 'initialized' not in st.session_state:
-        st.session_state.initialized = False
-    if 'history' not in st.session_state:
-        st.session_state.history = []
-
-@st.cache_resource
-def initialize_rag():
-    """Initialize the RAG system (cached)"""
-    try:
-        # Initialize components
-        llm = Config.get_llm()
-        doc_processor = DocumentProcessor(
-            chunk_size=Config.CHUNK_SIZE,
-            chunk_overlap=Config.CHUNK_OVERLAP
-        )
-        vector_store = VectorStore()
-        
-        # Use default URLs
-        urls = Config.DEFAULT_URLS
-        
-        # Process documents
-        documents = doc_processor.process_urls(urls)
-        
-        # Create vector store
-        vector_store.create_vectorstore(documents)
-        
-        # Build graph
-        graph_builder = GraphBuilder(
-            retriever=vector_store.get_retriever(),
-            llm=llm
-        )
-        graph_builder.build()
-        
-        return graph_builder, len(documents)
-    except Exception as e:
-        st.error(f"Failed to initialize: {str(e)}")
-        return None, 0
-
-def main():
-    """Main application"""
-    init_session_state()
-    
-    # Title
-    st.title("🔍 RAG Document Search")
-    st.markdown("Ask questions about the loaded documents")
-    
-    # Initialize system
-    if not st.session_state.initialized:
-        with st.spinner("Loading system..."):
-            rag_system, num_chunks = initialize_rag()
-            if rag_system:
-                st.session_state.rag_system = rag_system
-                st.session_state.initialized = True
-                st.success(f"✅ System ready! ({num_chunks} document chunks loaded)")
-    
-    st.markdown("---")
-    
-    # Search interface
-    with st.form("search_form"):
-        question = st.text_input(
-            "Enter your question:",
-            placeholder="What would you like to know?"
-        )
-        submit = st.form_submit_button("🔍 Search")
-    
-    # Process search
-    if submit and question:
-        if st.session_state.rag_system:
-            with st.spinner("Searching..."):
-                start_time = time.time()
-                
-                # Get answer
-                result = st.session_state.rag_system.run(question)
-                
-                elapsed_time = time.time() - start_time
-                
-                # Add to history
-                st.session_state.history.append({
-                    'question': question,
-                    'answer': result['answer'],
-                    'time': elapsed_time
-                })
-                
-                # Display answer
-                st.markdown("### 💡 Answer")
-                st.success(result['answer'])
-                
-                # Show retrieved docs in expander
-                with st.expander("📄 Source Documents"):
-                    for i, doc in enumerate(result['retrieved_docs'], 1):
-                        st.text_area(
-                            f"Document {i}",
-                            doc.page_content[:300] + "...",
-                            height=100,
-                            disabled=True
+
+    with st.chat_message("assistant"):
+        status_holder = {"box": None}
+
+        def ai_only_stream():
+            for message_chunk, _ in chatbot.stream(
+                {"messages": [HumanMessage(content=user_input)]},
+                config=CONFIG,
+                stream_mode="messages",
+            ):
+                if isinstance(message_chunk, ToolMessage):
+                    tool_name = getattr(message_chunk, "name", "tool")
+                    if status_holder["box"] is None:
+                        status_holder["box"] = st.status(
+                            f"🔧 Using `{tool_name}` …", expanded=True
                         )
-                
-                st.caption(f"⏱️ Response time: {elapsed_time:.2f} seconds")
-    
-    # Show history
-    if st.session_state.history:
-        st.markdown("---")
-        st.markdown("### 📜 Recent Searches")
-        
-        for item in reversed(st.session_state.history[-3:]):  # Show last 3
-            with st.container():
-                st.markdown(f"**Q:** {item['question']}")
-                st.markdown(f"**A:** {item['answer'][:200]}...")
-                st.caption(f"Time: {item['time']:.2f}s")
-                st.markdown("")
-
-if __name__ == "__main__":
-    main()
+                    else:
+                        status_holder["box"].update(
+                            label=f"🔧 Using `{tool_name}` …",
+                            state="running",
+                            expanded=True,
+                        )
+
+                if isinstance(message_chunk, AIMessage):
+                    yield message_chunk.content
+
+        ai_message = st.write_stream(ai_only_stream())
+
+        if status_holder["box"] is not None:
+            status_holder["box"].update(
+                label="✅ Tool finished", state="complete", expanded=False
+            )
+
+    st.session_state["message_history"].append(
+        {"role": "assistant", "content": ai_message}
+    )
+
+    doc_meta = thread_document_metadata(thread_key)
+    if doc_meta:
+        st.caption(
+            f"Document indexed: {doc_meta.get('filename')} "
+            f"(chunks: {doc_meta.get('chunks')}, pages: {doc_meta.get('documents')})"
+        )
+
+st.divider()
+
+if selected_thread:
+    st.session_state["thread_id"] = selected_thread
+    messages = load_conversation(selected_thread)
+
+    temp_messages = []
+    for msg in messages:
+        role = "user" if isinstance(msg, HumanMessage) else "assistant"
+        temp_messages.append({"role": role, "content": msg.content})
+    st.session_state["message_history"] = temp_messages
+    st.session_state["ingested_docs"].setdefault(str(selected_thread), {})
+    st.rerun()