app.py

import sys
import streamlit as st
from pinecone import Pinecone
from langchain_openai import ChatOpenAI
from langchain_ollama import ChatOllama
from langgraph.pregel import GraphRecursionError
import tempfile
import os
import time
from pathlib import Path
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
from src.vectorstore.pinecone_db import ingest_data, get_retriever, load_documents, process_chunks, save_to_parquet
from src.agents.workflow import run_adaptive_rag
# Page config
st.set_page_config(page_title="RAG Chat Assistant", layout="wide")

# Initialize session states
if "messages" not in st.session_state:
    st.session_state.messages = []
if "documents_processed" not in st.session_state:
    st.session_state.documents_processed = False
if "retriever" not in st.session_state:
    st.session_state.retriever = None
if "pinecone_client" not in st.session_state:
    st.session_state.pinecone_client = None

def initialize_pinecone(api_key):
    """Initialize Pinecone client with API key."""
    try:
        return Pinecone(api_key=api_key)
    except Exception as e:
        st.error(f"Error initializing Pinecone: {str(e)}")
        return None

def initialize_llm(llm_option, openai_api_key=None):
    """Initialize LLM based on user selection."""
    if llm_option == "OpenAI":
        if not openai_api_key:
            st.sidebar.warning("Please enter OpenAI API key.")
            return None
        return ChatOpenAI(api_key=openai_api_key, model="gpt-3.5-turbo")
    else:
        return ChatOllama(model="llama3.2", temperature=0.3, num_predict=512, top_p=0.6)

def clear_pinecone_index(pc, index_name="vector-index"):
    """Clear the Pinecone index."""
    try:
        pc.delete_index(index_name)
        st.session_state.documents_processed = False
        st.session_state.retriever = None
        st.success("Database cleared successfully!")
    except Exception as e:
        st.error(f"Error clearing database: {str(e)}")

def process_documents(uploaded_files, pc):
    """Process uploaded documents and store in Pinecone."""
    if not uploaded_files:
        st.warning("Please upload at least one document.")
        return False

    with st.spinner("Processing documents..."):
        temp_dir = tempfile.mkdtemp()
        file_paths = []
        markdown_path = Path(temp_dir) / "combined.md"
        parquet_path = Path(temp_dir) / "documents.parquet"
        
        for uploaded_file in uploaded_files:
            file_path = Path(temp_dir) / uploaded_file.name
            with open(file_path, "wb") as f:
                f.write(uploaded_file.getvalue())
            file_paths.append(str(file_path))

        try:
            markdown_path = load_documents(file_paths, output_path=markdown_path)
            chunks = process_chunks(markdown_path, chunk_size=256, threshold=0.6)
            print(f"Processed chunks: {chunks}")
            parquet_path = save_to_parquet(chunks, parquet_path)
            
            ingest_data(
                pc=pc,
                parquet_path=parquet_path,
                text_column="text",
                pinecone_client=pc
            )
            
            st.session_state.retriever = get_retriever(pc)
            st.session_state.documents_processed = True
            
            return True
            
        except Exception as e:
            st.error(f"Error processing documents: {str(e)}")
            return False
        finally:
            for file_path in file_paths:
                try:
                    os.remove(file_path)
                except:
                    pass
            try:
                os.rmdir(temp_dir)
            except:
                pass

def run_rag_with_streaming(retriever, question, llm, enable_web_search=False):
    """Run RAG workflow and yield streaming results."""
    try:
        response = run_adaptive_rag(
            retriever=retriever,
            question=question,
            llm=llm,
            top_k=5,
            enable_websearch=enable_web_search
        )
        
        for word in response.split():
            yield word + " "
            time.sleep(0.03)
            
    except GraphRecursionError:
        response = "I apologize, but I cannot find a sufficient answer to your question in the provided documents. Please try rephrasing your question or ask something else about the content of the documents."
        for word in response.split():
            yield word + " "
            time.sleep(0.03)
            
    except Exception as e:
        yield f"I encountered an error while processing your question: {str(e)}"

def main():
    st.title("🤖 RAG Chat Assistant")
    
    # Sidebar configuration
    st.sidebar.title("Configuration")
    
    # API Keys in sidebar
    pinecone_api_key = st.sidebar.text_input("Enter Pinecone API Key:", type="password")
    
    # LLM Selection
    llm_option = st.sidebar.selectbox("Select Language Model:", ["OpenAI"])
    openai_api_key = None
    if llm_option == "OpenAI":
        openai_api_key = st.sidebar.text_input("Enter OpenAI API Key:", type="password")
    
    # Web search tool in sidebar
    st.sidebar.markdown("---")
    st.sidebar.markdown("### Tools")
    use_web_search = st.sidebar.checkbox("Web search")
    
    # Initialize Pinecone
    if pinecone_api_key:
        if st.session_state.pinecone_client is None:
            st.session_state.pinecone_client = initialize_pinecone(pinecone_api_key)
    else:
        st.sidebar.warning("Please enter Pinecone API key to continue.")
        st.stop()
    
    # Initialize LLM
    llm = initialize_llm(llm_option, openai_api_key)
    if llm is None:
        st.stop()
    
    # Clear DB Button
    st.sidebar.markdown("---")
    if st.sidebar.button("Clear Database"):
        if st.session_state.pinecone_client:
            clear_pinecone_index(st.session_state.pinecone_client)
            st.session_state.messages = []  # Clear chat history
    
    # Document upload section
    if not st.session_state.documents_processed:
        st.header("📄 Document Upload")
        uploaded_files = st.file_uploader(
            "Upload your documents",
            accept_multiple_files=True,
            type=["pdf", "docx", "txt", "pptx", "md"]
        )
        
        if st.button("Process Documents"):
            if process_documents(uploaded_files, st.session_state.pinecone_client):
                st.success("Documents processed successfully!")
            
    # Chat interface
    if st.session_state.documents_processed:
        st.header("💬 Chat")
        
        # Display chat history
        for message in st.session_state.messages:
            with st.chat_message(message["role"]):
                st.markdown(message["content"])
        
        # Chat input
        if prompt := st.chat_input("Ask a question about your documents..."):
            # Display user message
            with st.chat_message("user"):
                if use_web_search:
                    st.markdown(prompt.strip() + ''' :red-background[Web Search]''')
                else:
                    st.markdown(prompt)
            st.session_state.messages.append({"role": "user", "content": prompt})
            
            # Generate and stream response
            with st.chat_message("assistant"):
                response_container = st.empty()
                full_response = ""
                
                # Show spinner while processing
                with st.spinner("Thinking..."):
                    # Stream the response
                    for chunk in run_rag_with_streaming(
                        retriever=st.session_state.retriever,
                        question=prompt,
                        llm=llm,
                        enable_web_search=use_web_search
                    ):
                        full_response += chunk
                        response_container.markdown(full_response + "▌")
                
                # Final update without cursor
                response_container.markdown(full_response)
                
                # Save to chat history
                st.session_state.messages.append(
                    {"role": "assistant", "content": full_response}
                )

if __name__ == "__main__":
    main()