Update app.py

app.py

@@ -14,6 +14,7 @@ from langchain_community.embeddings import HuggingFaceEmbeddings
 from langchain_community.chat_models import ChatOpenAI
 from langchain.prompts import ChatPromptTemplate, MessagesPlaceholder
 from langchain_core.messages import HumanMessage, AIMessage, SystemMessage, BaseMessage
+from langchain_community.document_loaders import PyPDFLoader
 import tempfile
 from utils.database import (
     create_connection,
@@ -33,7 +34,8 @@ from utils.database import (
     get_chat_messages,
     get_document_tags,
     add_document_tags,
-    delete_collection)
+    delete_collection,
+    insert_document)
 
 
 @dataclass
@@ -49,73 +51,54 @@ class SessionState:
     reinitialize_chat: bool = False
 
 
-def display_header():
-    """Display the application header with navigation."""
-    # Add custom CSS for header styling
-    st.markdown(
-        """
-        <style>
-        .stButton > button {
-            width: 100%;
-            margin-bottom: 0;
-        }
-        .header-button {
-            margin: 0 5px;
-        }
-        </style>
-        """,
-        unsafe_allow_html=True
-    )
-
-    # Create header layout
-    header_container = st.container()
-    with header_container:
-        # Main header row
-        col1, col2, col3, col4, col5, col6 = st.columns([1.5, 2.5, 1, 1, 1, 1])
-
-        # Logo
-        with col1:
-            if os.path.exists("img/logo.png"):
-                st.image("img/logo.png", width=150)
-            else:
-                st.info("Logo missing: img/logo.png")
-
-        # Title
-        with col2:
-            st.markdown("##### Synaptyx RFP Analyzer Agent")
-
-        # Navigation Buttons
-        with col3:
-            if st.button("🏠 Home", use_container_width=True, key="home_btn"):
-                st.session_state.chat_ready = False
-                st.session_state.messages = []
-                st.session_state.current_chat_id = None
-                st.session_state.show_explorer = False
-                st.rerun()
-
-        with col4:
-            if st.button("📚 Explorer", use_container_width=True, key="explorer_btn"):
-                st.session_state.show_explorer = True
-                st.session_state.chat_ready = False
-                st.rerun()
-
-        with col5:
-            if st.session_state.chat_ready:
-                if st.button("💭 New Chat", use_container_width=True, key="chat_btn"):
-                    st.session_state.messages = []
-                    st.session_state.current_chat_id = None
-                    st.rerun()
-
-        with col6:
-            if st.button("📁 Upload", use_container_width=True, key="upload_btn"):
-                st.session_state.show_collection_dialog = True
-                st.rerun()
-
-    # Add divider after header
-    st.divider()
+def initialize_session_state():
+    """Initialize session state with default values."""
+    defaults = SessionState()
+    if 'initialized' not in st.session_state:
+        # Setup data paths
+        data_path = Path('/data' if os.path.exists('/data') else 'data')
+        vector_store_path = data_path / 'vector_stores'
+
+        # Create necessary directories
+        data_path.mkdir(parents=True, exist_ok=True)
+        vector_store_path.mkdir(parents=True, exist_ok=True)
+
+        # Initialize session state
+        st.session_state.update({
+            'show_collection_dialog': defaults.show_collection_dialog,
+            'selected_collection': defaults.selected_collection,
+            'chat_ready': defaults.chat_ready,
+            'messages': [] if defaults.messages is None else defaults.messages,
+            'current_chat_id': defaults.current_chat_id,
+            'vector_store': defaults.vector_store,
+            'qa_system': defaults.qa_system,
+            'reinitialize_chat': defaults.reinitialize_chat,
+            'initialized': True,
+            'data_path': data_path,
+            'vector_store_path': vector_store_path,
+            'show_explorer': False
+        })
+
+
+def generate_document_tags(content: str) -> List[str]:
+    """Generate tags for a document using AI."""
+    try:
+        llm = ChatOpenAI(temperature=0.2, model="gpt-3.5-turbo")
+        
+        prompt = """Analyze the following document content and generate relevant tags/keywords. 
+        Focus on key themes, topics, and important terminology.
+        Return only the tags as a comma-separated list.
+        Content: {content}"""
+        
+        response = llm.invoke(prompt.format(content=content[:2000]))  # Use first 2000 chars
+        tags = [tag.strip() for tag in response.split(',')]
+        return tags
+    except Exception as e:
+        st.error(f"Error generating tags: {e}")
+        return []
 
 
-async def process_document(file_path: str, collection_id: Optional[int] = None) -> Tuple[List, str]:
+def process_document(file_path: str, collection_id: Optional[int] = None) -> Tuple[List, str, List[str]]:
     """Process a document with automatic tagging."""
     try:
         # Load PDF
@@ -126,7 +109,7 @@ async def process_document(file_path: str, collection_id: Optional[int] = None)
         full_content = "\n".join(doc.page_content for doc in documents)
         
         # Generate tags
-        tags = await generate_document_tags(full_content)
+        tags = generate_document_tags(full_content)
         
         # Create text splitter for chunks
         text_splitter = RecursiveCharacterTextSplitter(
@@ -153,7 +136,7 @@ async def process_document(file_path: str, collection_id: Optional[int] = None)
         return [], "", []
 
 
-async def handle_document_upload(uploaded_files: List, collection_id: Optional[int] = None) -> bool:
+def handle_document_upload(uploaded_files: List, collection_id: Optional[int] = None) -> bool:
     """Handle document upload with progress tracking and auto-tagging."""
     try:
         progress_container = st.empty()
@@ -183,7 +166,7 @@ async def handle_document_upload(uploaded_files: List, collection_id: Optional[i
                 tmp_file.flush()
                 
                 # Process document with tagging
-                chunks, content, tags = await process_document(tmp_file.name, collection_id)
+                chunks, content, tags = process_document(tmp_file.name, collection_id)
                 
                 # Store in database
                 doc_id = insert_document(st.session_state.db_conn, uploaded_file.name, content)
@@ -228,6 +211,72 @@ async def handle_document_upload(uploaded_files: List, collection_id: Optional[i
         return False
 
 
+def display_header():
+    """Display the application header with navigation."""
+    # Add custom CSS for header styling
+    st.markdown(
+        """
+        <style>
+        .stButton > button {
+            width: 100%;
+            margin-bottom: 0;
+        }
+        .header-button {
+            margin: 0 5px;
+        }
+        </style>
+        """,
+        unsafe_allow_html=True
+    )
+
+    # Create header layout
+    header_container = st.container()
+    with header_container:
+        # Main header row
+        col1, col2, col3, col4, col5, col6 = st.columns([1.5, 2.5, 1, 1, 1, 1])
+
+        # Logo
+        with col1:
+            if os.path.exists("img/logo.png"):
+                st.image("img/logo.png", width=150)
+            else:
+                st.info("Logo missing: img/logo.png")
+
+        # Title
+        with col2:
+            st.markdown("##### Synaptyx RFP Analyzer Agent")
+
+        # Navigation Buttons
+        with col3:
+            if st.button("🏠 Home", use_container_width=True, key="home_btn"):
+                st.session_state.chat_ready = False
+                st.session_state.messages = []
+                st.session_state.current_chat_id = None
+                st.session_state.show_explorer = False
+                st.rerun()
+
+        with col4:
+            if st.button("📚 Explorer", use_container_width=True, key="explorer_btn"):
+                st.session_state.show_explorer = True
+                st.session_state.chat_ready = False
+                st.rerun()
+
+        with col5:
+            if st.session_state.chat_ready:
+                if st.button("💭 New Chat", use_container_width=True, key="chat_btn"):
+                    st.session_state.messages = []
+                    st.session_state.current_chat_id = None
+                    st.rerun()
+
+        with col6:
+            if st.button("📁 Upload", use_container_width=True, key="upload_btn"):
+                st.session_state.show_collection_dialog = True
+                st.rerun()
+
+    # Add divider after header
+    st.divider()
+
+
 def display_collection_management():
     """Display collection management interface."""
     st.header("📁 Collection Management")
@@ -271,7 +320,7 @@ def display_collection_management():
             if uploaded_files:
                 collection_id = selected_collection[1]
                 with st.spinner("Processing documents..."):
-                    if await handle_document_upload(uploaded_files, collection_id=collection_id):
+                    if handle_document_upload(uploaded_files, collection_id=collection_id):
                         st.success("Documents added to collection successfully!")
                         st.rerun()
     
@@ -350,6 +399,140 @@ def display_chat_interface():
         st.rerun()
 
 
+def display_welcome_screen():
+    """Display welcome screen with quick actions."""
+    st.header("Quick Start")
+
+    col1, col2 = st.columns([3, 2])
+
+    with col1:
+        # Upload new documents
+        st.markdown("### Upload Documents")
+        collection_id = None
+        collections = get_collections(st.session_state.db_conn)
+        
+        if collections:
+            selected_collection = st.selectbox(
+                "Select Collection (Optional)",
+                options=[("None", None)] + [(c["name"], c["id"]) for c in collections],
+                format_func=lambda x: x[0]
+            )
+            collection_id = selected_collection[1] if selected_collection[0] != "None" else None
+            
+            # Add new collection button
+            if st.button("Create New Collection", use_container_width=True):
+                st.session_state.show_collection_dialog = True
+                st.rerun()
+
+        uploaded_files = st.file_uploader(
+            "Upload Documents",
+            type=['pdf'],
+            accept_multiple_files=True,
+            help="Upload PDF documents to start analyzing"
+        )
+
+        if uploaded_files:
+            with st.spinner("Processing documents..."):
+                if handle_document_upload(uploaded_files, collection_id=collection_id):
+                    initialize_chat_system(collection_id)
+                    st.rerun()
+
+    with col2:
+        # Display existing collections
+        st.header("Collections")
+        if collections:
+            for collection in collections:
+                with st.expander(f"📁 {collection['name']} ({collection['doc_count']} documents)"):
+                    st.write(collection.get('description', ''))
+                    if st.button("Start Chat", key=f"chat_{collection['id']}", use_container_width=True):
+                        st.session_state.selected_collection = collection
+                        if initialize_chat_system(collection['id']):
+                            st.rerun()
+
+        # Show recent documents
+        st.header("Recent Documents")
+        recent_docs = get_recent_documents(st.session_state.db_conn, limit=5)
+        for doc in recent_docs:
+            with st.expander(f"📄 {doc['name']}"):
+                st.caption(f"Upload date: {doc['upload_date']}")
+                if doc['collections']:
+                    st.caption(f"Collections: {', '.join(doc['collections'])}")
+                if st.button("Start Chat", key=f"doc_{doc['id']}", use_container_width=True):
+                    if initialize_chat_system():
+                        st.rerun()
+
+
+def display_document_chunks():
+    """Display document chunks with search and filtering capabilities."""
+    st.subheader("Document Chunk Explorer")
+
+    # Get all documents
+    documents = get_all_documents(st.session_state.db_conn)
+    if not documents:
+        st.info("No documents available.")
+        return
+
+    # Document selection
+    selected_doc = st.selectbox(
+        "Select Document",
+        options=documents,
+        format_func=lambda x: x['name']
+    )
+    if not selected_doc:
+        return
+
+    try:
+        # Load vector store for selected document
+        embeddings = get_embeddings_model()
+        chunks = []
+
+        # Search functionality
+        search_query = st.text_input("🔍 Search within chunks")
+        
+        if search_query and st.session_state.vector_store:
+            chunks = st.session_state.vector_store.similarity_search(search_query, k=5)
+        elif st.session_state.vector_store:
+            chunks = st.session_state.vector_store.similarity_search("", k=100)
+
+        # Display chunks with metadata
+        st.markdown("### Document Chunks")
+
+        # Filtering options
+        col1, col2 = st.columns(2)
+        with col1:
+            chunk_size = st.slider("Preview Size", 100, 1000, 500)
+        with col2:
+            sort_by = st.selectbox("Sort By", ["Relevance", "Position"])
+
+        # Display chunks in an organized way
+        for i, chunk in enumerate(chunks):
+            with st.expander(f"Chunk {i+1} | Source: {chunk.metadata.get('source', 'Unknown')}"):
+                # Content preview
+                st.markdown("**Content:**")
+                st.text(chunk.page_content[:chunk_size] + "..." if len(chunk.page_content) > chunk_size else chunk.page_content)
+
+                # Metadata
+                st.markdown("**Metadata:**")
+                for key, value in chunk.metadata.items():
+                    st.text(f"{key}: {value}")
+
+                # Actions
+                col1, col2 = st.columns(2)
+                with col1:
+                    if st.button("Copy", key=f"copy_{i}"):
+                        st.write("Content copied to clipboard!")
+                with col2:
+                    if st.button("Start Chat", key=f"chat_{i}"):
+                        initialize_chat_system()
+                        st.session_state.messages.append(
+                            HumanMessage(content=f"Tell me about: {chunk.page_content[:100]}...")
+                        )
+                        st.rerun()
+
+    except Exception as e:
+        st.error(f"Error loading document chunks: {e}")
+
+
 def main():
     """Main application function with improved state management."""
     st.set_page_config(
@@ -358,7 +541,7 @@ def main():
         initial_sidebar_state="collapsed"
     )
 
-    # Initialize session state
+    # Initialize session state with paths
     initialize_session_state()
 
     # Initialize database connection
@@ -372,13 +555,13 @@ def main():
     # Display header
     display_header()
 
-    # Show collection management if triggered
+    # Show different views based on application state
     if st.session_state.show_collection_dialog:
         display_collection_management()
-    # Display chat interface if ready
     elif st.session_state.chat_ready:
         display_chat_interface()
-    # Show welcome screen
+    elif st.session_state.show_explorer:
+        display_document_chunks()
     else:
         display_welcome_screen()