Update app.py

app.py

@@ -1,43 +1,217 @@
 import streamlit as st
-import asyncio
-import websockets
+from langchain.chains import RetrievalQA
+from langchain.vectorstores import Milvus
+from langchain.embeddings import HuggingFaceEmbeddings
+from transformers import AutoTokenizer
+from langchain_groq import ChatGroq
+import os
+from docling.document_converter import DocumentConverter, PdfFormatOption
+from docling.datamodel.base_models import InputFormat
+from docling.datamodel.pipeline_options import PdfPipelineOptions
+from docling_core.transforms.chunker.hybrid_chunker import HybridChunker
+from docling_core.types.doc.document import TableItem
+from langchain_core.documents import Document
+import itertools
+from docling_core.types.doc.labels import DocItemLabel
+import google.generativeai as genai
+from PIL import Image
+import base64
+import io
 
-st.markdown('<h1 style="color: darkblue;">AI Voice Assistant</h1>', unsafe_allow_html=True)
+# Initialize components (similar to your notebook)
+@st.cache_resource
+def initialize_components():
+    # Initialize embeddings
+    embeddings_model_path = "ibm-granite/granite-embedding-30m-english"
+    embeddings_model = HuggingFaceEmbeddings(model_name=embeddings_model_path)
+    embeddings_tokenizer = AutoTokenizer.from_pretrained(embeddings_model_path)
+    
+    # Initialize language model
+    GROQ_API_KEY = "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"
+    model = ChatGroq(model_name="llama3-70b-8192", api_key=GROQ_API_KEY)
+    
+    # Initialize vision model
+    GOOGLE_API_KEY = "AIzaSyBTt66oOvxpLeYn41sR-KkjSYPK2vOAqkU"
+    genai.configure(api_key=GOOGLE_API_KEY)
+    vision_model = genai.GenerativeModel(model_name="gemini-1.5-flash")
+    
+    return embeddings_model, embeddings_tokenizer, model, vision_model
 
-# JavaScript for real-time voice streaming
-audio_recorder_js = """
-<script>
-let mediaRecorder;
-let ws;
-
-function startRecording() {
-    navigator.mediaDevices.getUserMedia({ audio: true }).then(stream => {
-        ws = new WebSocket("ws://localhost:8765"); // Replace with your server's WebSocket URL
-        mediaRecorder = new MediaRecorder(stream);
-        mediaRecorder.start();
+def process_pdf(file_path, embeddings_tokenizer, vision_model):
+    # PDF processing (similar to your notebook)
+    pdf_pipeline_options = PdfPipelineOptions(
+        do_ocr=True,
+        generate_picture_images=True
+    )
+    
+    format_options = {
+        InputFormat.PDF: PdfFormatOption(pipeline_options=pdf_pipeline_options),
+    }
+    
+    converter = DocumentConverter(format_options=format_options)
+    sources = [file_path]
+    conversions = {
+        source: converter.convert(source=source).document for source in sources
+    }
+    
+    # Process text chunks
+    doc_id = 0
+    texts = []
+    
+    for source, docling_document in conversions.items():
+        chunker = HybridChunker(tokenizer=embeddings_tokenizer)
         
-        mediaRecorder.ondataavailable = event => {
-            ws.send(event.data);
-        };
+        for chunk in chunker.chunk(docling_document):
+            items = chunk.meta.doc_items
+            
+            if len(items) == 1 and isinstance(items[0], TableItem):
+                continue
+                
+            refs = "".join(item.get_ref().cref for item in items)
+            text = chunk.text
+            
+            document = Document(
+                page_content=text,
+                metadata={
+                    "doc_id": (doc_id := doc_id + 1),
+                    "source": source,
+                    "ref": refs,
+                }
+            )
+            texts.append(document)
+    
+    # Process tables (if any)
+    tables = []
+    for source, docling_document in conversions.items():
+        for table in docling_document.tables:
+            if table.label == DocItemLabel.TABLE:
+                ref = table.get_ref().cref
+                text = table.export_to_markdown()
+                
+                document = Document(
+                    page_content=text,
+                    metadata={
+                        "doc_id": (doc_id := doc_id + 1),
+                        "source": source,
+                        "ref": ref,
+                    },
+                )
+                tables.append(document)
+    
+    # Process images (if any)
+    pictures = []
+    start_doc_id = len(texts) + len(tables) + 1
+    
+    for source, docling_document in conversions.items():
+        if hasattr(docling_document, 'pictures') and docling_document.pictures:
+            for picture in docling_document.pictures:
+                try:
+                    ref = picture.get_ref().cref
+                    image = picture.get_image(docling_document)
+                    
+                    if image:
+                        response = vision_model.generate_content([
+                            "Extract all text and describe key visual elements in this image. "
+                            "Include any numbers, labels, or important details.",
+                            image
+                        ])
+                        
+                        document = Document(
+                            page_content=response.text,
+                            metadata={
+                                "doc_id": doc_id,
+                                "source": source,
+                                "ref": ref,
+                            }
+                        )
+                        pictures.append(document)
+                        doc_id += 1
+                except Exception as e:
+                    print(f"Error processing image: {str(e)}")
+    
+    return texts + tables + pictures
 
-        ws.onmessage = function(event) {
-            document.getElementById("response").innerHTML += "<br><b>AI:</b> " + event.data;
-        };
-    });
-}
+def create_vector_store(docs, embeddings_model):
+    # Create vector store (using Milvus as in your notebook)
+    # Note: You'll need to have Milvus running
+    vector_store = Milvus.from_documents(
+        docs,
+        embeddings_model,
+        connection_args={"host": "127.0.0.1", "port": "19530"},
+        collection_name="pdf_manual"
+    )
+    return vector_store
 
-function stopRecording() {
-    mediaRecorder.stop();
-    ws.close();
-}
-</script>
-"""
+def main():
+    st.title("PDF Manual Chatbot")
+    
+    # Initialize components
+    embeddings_model, embeddings_tokenizer, model, vision_model = initialize_components()
+    
+    # File upload
+    uploaded_file = st.file_uploader("Upload a PDF manual", type="pdf")
+    
+    if uploaded_file is not None:
+        # Save the uploaded file
+        file_path = os.path.join("temp", uploaded_file.name)
+        os.makedirs("temp", exist_ok=True)
+        with open(file_path, "wb") as f:
+            f.write(uploaded_file.getbuffer())
+        
+        # Process the PDF
+        with st.spinner("Processing PDF..."):
+            docs = process_pdf(file_path, embeddings_tokenizer, vision_model)
+            vector_store = create_vector_store(docs, embeddings_model)
+        
+        st.success("PDF processed successfully!")
+        
+        # Initialize chat history
+        if "messages" not in st.session_state:
+            st.session_state.messages = []
+        
+        # Display chat messages from history on app rerun
+        for message in st.session_state.messages:
+            with st.chat_message(message["role"]):
+                st.markdown(message["content"])
+        
+        # Accept user input
+        if prompt := st.chat_input("Ask a question about the manual"):
+            # Add user message to chat history
+            st.session_state.messages.append({"role": "user", "content": prompt})
+            
+            # Display user message in chat message container
+            with st.chat_message("user"):
+                st.markdown(prompt)
+            
+            # Create QA chain
+            qa_chain = RetrievalQA.from_chain_type(
+                llm=model,
+                chain_type="stuff",
+                retriever=vector_store.as_retriever(),
+                return_source_documents=True
+            )
+            
+            # Get response
+            with st.spinner("Thinking..."):
+                result = qa_chain({"query": prompt})
+                response = result["result"]
+                source_docs = result["source_documents"]
+            
+            # Display assistant response in chat message container
+            with st.chat_message("assistant"):
+                st.markdown(response)
+                
+                # Show sources if available
+                if source_docs:
+                    with st.expander("Source Documents"):
+                        for i, doc in enumerate(source_docs):
+                            st.write(f"Source {i+1}:")
+                            st.write(doc.page_content)
+                            st.write(f"Metadata: {doc.metadata}")
+                            st.write("---")
+            
+            # Add assistant response to chat history
+            st.session_state.messages.append({"role": "assistant", "content": response})
 
-# Display buttons for real-time recording
-st.components.v1.html(
-    audio_recorder_js + """
-    <button onclick="startRecording()">🎤 Start Talking</button>
-    <button onclick="stopRecording()">🛑 Stop</button>
-    <div id="response" style="margin-top: 10px; padding: 10px; border: 1px solid #ccc;"></div>
-    """, height=200
-)
+if __name__ == "__main__":
+    main()