Create app.py

app.py

@@ -0,0 +1,218 @@
+import gradio as gr
+import groq
+import os
+import tempfile
+import uuid
+from dotenv import load_dotenv
+from langchain.text_splitter import RecursiveCharacterTextSplitter
+from langchain.vectorstores import FAISS
+from langchain.embeddings import HuggingFaceEmbeddings
+import fitz  # PyMuPDF
+import base64
+from PIL import Image
+import io
+
+# Load environment variables
+load_dotenv()
+client = groq.Client(api_key=os.getenv("GROQ_LEGAL_API_KEY"))
+embeddings = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2")
+
+# Directory to store FAISS indexes
+FAISS_INDEX_DIR = "faiss_indexes_tech"
+if not os.path.exists(FAISS_INDEX_DIR):
+    os.makedirs(FAISS_INDEX_DIR)
+
+# Dictionary to store user-specific vectorstores
+user_vectorstores = {}
+
+# Custom CSS for Tech theme
+custom_css = """
+:root {
+    --primary-color: #008080;  /* Teal */
+    --secondary-color: #006666;  /* Dark Teal */
+    --light-background: #E0FFFF;  /* Light Cyan */
+    --dark-text: #333333;
+    --white: #FFFFFF;
+    --border-color: #E5E7EB;
+}
+body { background-color: var(--light-background); font-family: 'Inter', sans-serif; }
+.container { max-width: 1200px !important; margin: 0 auto !important; padding: 10px; }
+.header { background-color: var(--white); border-bottom: 2px solid var(--border-color); padding: 15px 0; margin-bottom: 20px; border-radius: 12px 12px 0 0; box-shadow: 0 2px 4px rgba(0,0,0,0.05); }
+.header-title { color: var(--secondary-color); font-size: 1.8rem; font-weight: 700; text-align: center; }
+.header-subtitle { color: var(--dark-text); font-size: 1rem; text-align: center; margin-top: 5px; }
+.chat-container { border-radius: 12px !important; box-shadow: 0 4px 6px rgba(0,0,0,0.1) !important; background-color: var(--white) !important; border: 1px solid var(--border-color) !important; min-height: 500px; }
+.message-user { background-color: var(--primary-color) !important; color: var(--white) !important; border-radius: 18px 18px 4px 18px !important; padding: 12px 16px !important; margin-left: auto !important; max-width: 80% !important; }
+.message-bot { background-color: #F0F0F0 !important; color: var(--dark-text) !important; border-radius: 18px 18px 18px 4px !important; padding: 12px 16px !important; margin-right: auto !important; max-width: 80% !important; }
+.input-area { background-color: var(--white) !important; border-top: 1px solid var(--border-color) !important; padding: 12px !important; border-radius: 0 0 12px 12px !important; }
+.input-box { border: 1px solid var(--border-color) !important; border-radius: 24px !important; padding: 12px 16px !important; box-shadow: 0 2px 4px rgba(0,0,0,0.05) !important; }
+.send-btn { background-color: var(--secondary-color) !important; border-radius: 24px !important; color: var(--white) !important; padding: 10px 20px !important; font-weight: 500 !important; }
+.clear-btn { background-color: #F0F0F0 !important; border: 1px solid var(--border-color) !important; border-radius: 24px !important; color: var(--dark-text) !important; padding: 8px 16px !important; font-weight: 500 !important; }
+.pdf-viewer-container { border-radius: 12px !important; box-shadow: 0 4px 6px rgba(0,0,0,0.1) !important; background-color: var(--white) !important; border: 1px solid var(--border-color) !important; padding: 20px; }
+.pdf-viewer-image { max-width: 100%; height: auto; border: 1px solid var(--border-color); border-radius: 12px; box-shadow: 0 2px 4px rgba(0,0,0,0.05); }
+.stats-box { background-color: #E0F0F0; padding: 10px; border-radius: 8px; margin-top: 10px; }
+"""
+
+# Function to process PDF files (unchanged)
+def process_pdf(pdf_file):
+    if pdf_file is None:
+        return None, "No file uploaded", {"page_images": [], "total_pages": 0, "total_words": 0}
+    try:
+        session_id = str(uuid.uuid4())
+        with tempfile.NamedTemporaryFile(suffix=".pdf", delete=False) as temp_file:
+            temp_file.write(pdf_file)
+            pdf_path = temp_file.name
+        
+        doc = fitz.open(pdf_path)
+        texts = [page.get_text() for page in doc]
+        page_images = []
+        for page in doc:
+            pix = page.get_pixmap()
+            img_bytes = pix.tobytes("png")
+            img_base64 = base64.b64encode(img_bytes).decode("utf-8")
+            page_images.append(img_base64)
+        total_pages = len(doc)
+        total_words = sum(len(text.split()) for text in texts)
+        doc.close()
+
+        text_splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200)
+        chunks = text_splitter.create_documents(texts)
+        vectorstore = FAISS.from_documents(chunks, embeddings)
+        index_path = os.path.join(FAISS_INDEX_DIR, session_id)
+        vectorstore.save_local(index_path)
+        user_vectorstores[session_id] = vectorstore
+
+        os.unlink(pdf_path)
+        pdf_state = {"page_images": page_images, "total_pages": total_pages, "total_words": total_words}
+        return session_id, f"✅ Successfully processed {len(chunks)} text chunks from your PDF", pdf_state
+    except Exception as e:
+        if "pdf_path" in locals() and os.path.exists(pdf_path):
+            os.unlink(pdf_path)
+        return None, f"Error processing PDF: {str(e)}", {"page_images": [], "total_pages": 0, "total_words": 0}
+
+# Function to generate chatbot responses with Tech theme
+def generate_response(message, session_id, model_name, history):
+    if not message:
+        return history
+    try:
+        context = ""
+        if session_id and session_id in user_vectorstores:
+            vectorstore = user_vectorstores[session_id]
+            docs = vectorstore.similarity_search(message, k=3)
+            if docs:
+                context = "\n\nRelevant information from uploaded PDF:\n" + "\n".join(f"- {doc.page_content}" for doc in docs)
+        system_prompt = "You are a technical assistant specializing in analyzing tech manuals, whitepapers, and documentation."
+        if context:
+            system_prompt += " Use the following context to answer the question if relevant: " + context
+        completion = client.chat.completions.create(
+            model=model_name,
+            messages=[
+                {"role": "system", "content": system_prompt},
+                {"role": "user", "content": message}
+            ],
+            temperature=0.7,
+            max_tokens=1024
+        )
+        response = completion.choices[0].message.content
+        history.append((message, response))
+        return history
+    except Exception as e:
+        history.append((message, f"Error generating response: {str(e)}"))
+        return history
+
+# Functions to update PDF viewer (unchanged)
+def update_pdf_viewer(pdf_state):
+    if not pdf_state["total_pages"]:
+        return 0, None, "No PDF uploaded yet"
+    try:
+        img_data = base64.b64decode(pdf_state["page_images"][0])
+        img = Image.open(io.BytesIO(img_data))
+        return pdf_state["total_pages"], img, f"**Total Pages:** {pdf_state['total_pages']}\n**Total Words:** {pdf_state['total_words']}"
+    except Exception as e:
+        print(f"Error decoding image: {e}")
+        return 0, None, "Error displaying PDF"
+
+def update_image(page_num, pdf_state):
+    if not pdf_state["total_pages"] or page_num < 1 or page_num > pdf_state["total_pages"]:
+        return None
+    try:
+        img_data = base64.b64decode(pdf_state["page_images"][page_num - 1])
+        img = Image.open(io.BytesIO(img_data))
+        return img
+    except Exception as e:
+        print(f"Error decoding image: {e}")
+        return None
+
+# Gradio interface
+with gr.Blocks(css=custom_css, theme=gr.themes.Soft()) as demo:
+    current_session_id = gr.State(None)
+    pdf_state = gr.State({"page_images": [], "total_pages": 0, "total_words": 0})
+    gr.HTML("""
+    <div class="header">
+        <div class="header-title">Tech-Vision</div>
+        <div class="header-subtitle">Analyze technical documents with Groq's LLM API.</div>
+    </div>
+    """)
+    with gr.Row(elem_classes="container"):
+        with gr.Column(scale=1, min_width=300):
+            pdf_file = gr.File(label="Upload PDF Document", file_types=[".pdf"], type="binary")
+            upload_button = gr.Button("Process PDF", variant="primary")
+            pdf_status = gr.Markdown("No PDF uploaded yet")
+            model_dropdown = gr.Dropdown(
+                choices=["llama3-70b-8192", "llama3-8b-8192", "mixtral-8x7b-32768", "gemma-7b-it"],
+                value="llama3-70b-8192",
+                label="Select Groq Model"
+            )
+        with gr.Column(scale=2, min_width=600):
+            with gr.Tabs():
+                with gr.TabItem("PDF Viewer"):
+                    with gr.Column(elem_classes="pdf-viewer-container"):
+                        page_slider = gr.Slider(minimum=1, maximum=1, step=1, label="Page Number", value=1)
+                        pdf_image = gr.Image(label="PDF Page", type="pil", elem_classes="pdf-viewer-image")
+                        stats_display = gr.Markdown("No PDF uploaded yet", elem_classes="stats-box")
+    
+    with gr.Row(elem_classes="container"):
+        with gr.Column(scale=2, min_width=600):
+            chatbot = gr.Chatbot(height=500, bubble_full_width=False, show_copy_button=True, elem_classes="chat-container")
+            with gr.Row():
+                msg = gr.Textbox(show_label=False, placeholder="Ask about your technical document...", scale=5)
+                send_btn = gr.Button("Send", scale=1)
+            clear_btn = gr.Button("Clear Conversation")
+    
+    # Event Handlers (unchanged)
+    upload_button.click(
+        process_pdf,
+        inputs=[pdf_file],
+        outputs=[current_session_id, pdf_status, pdf_state]
+    ).then(
+        update_pdf_viewer,
+        inputs=[pdf_state],
+        outputs=[page_slider, pdf_image, stats_display]
+    )
+    
+    msg.submit(
+        generate_response,
+        inputs=[msg, current_session_id, model_dropdown, chatbot],
+        outputs=[chatbot]
+    ).then(lambda: "", None, [msg])
+    
+    send_btn.click(
+        generate_response,
+        inputs=[msg, current_session_id, model_dropdown, chatbot],
+        outputs=[chatbot]
+    ).then(lambda: "", None, [msg])
+    
+    clear_btn.click(
+        lambda: ([], None, "No PDF uploaded yet", {"page_images": [], "total_pages": 0, "total_words": 0}, 0, None, "No PDF uploaded yet"),
+        None,
+        [chatbot, current_session_id, pdf_status, pdf_state, page_slider, pdf_image, stats_display]
+    )
+    
+    page_slider.change(
+        update_image,
+        inputs=[page_slider, pdf_state],
+        outputs=[pdf_image]
+    )
+
+# Launch the app
+if __name__ == "__main__":
+    demo.launch()