Update app.py

app.py

@@ -1,83 +1,71 @@
-import gradio as gr
-from transformers import pipeline
-import torch
-import tempfile
 import os
-import graphviz
+import tempfile
+import streamlit as st
+from dotenv import load_dotenv
+from PyPDF2 import PdfReader
+from openai import OpenAI
+
+# Vector DB imports (Qdrant + Pinecone)
+import pinecone
+from qdrant_client import QdrantClient
+
+# Load secrets
+load_dotenv()
+OPENAI_API_KEY = os.getenv("OPENAI_API_KEY")
+PINECONE_API_KEY = os.getenv("PINECONE_API_KEY")
+QDRANT_API_KEY = os.getenv("QDRANT_API_KEY")
+
+client = OpenAI(api_key=OPENAI_API_KEY)
+
+# Choose vector DB here
+VECTOR_DB = "qdrant"  # change to "pinecone" if needed
 
-# -------------------------
-# Load lightweight models
-# -------------------------
-summarizer = pipeline("summarization", model="sshleifer/distilbart-cnn-12-6")
-tts = pipeline("text-to-speech", model="facebook/mms-tts-eng")
+# Initialize vector DB
+if VECTOR_DB == "pinecone":
+    pinecone.init(api_key=PINECONE_API_KEY, environment="gcp-starter")
+    index_name = "pdf-index"
+    if index_name not in pinecone.list_indexes():
+        pinecone.create_index(index_name, dimension=1536)
+    vector_db = pinecone.Index(index_name)
+else:
+    vector_db = QdrantClient(
+        url="https://your-qdrant-url", api_key=QDRANT_API_KEY
+    )
 
-# -------------------------
-# Helpers
-# -------------------------
-def summarize_pdf(pdf_file):
-    try:
-        import pypdf
-        reader = pypdf.PdfReader(pdf_file.name)
-        text = ""
-        for page in reader.pages:
-            text += page.extract_text() or ""
-        if not text.strip():
-            return "❌ No text extracted from PDF."
-        # keep only first 2000 chars (model limit)
-        chunk = text[:2000]
-        summary = summarizer(chunk, max_length=120, min_length=40, do_sample=False)[0]['summary_text']
-        return summary
-    except Exception as e:
-        return f"❌ Error in summarization: {e}"
+# Streamlit UI
+st.title("📄 PDF AI Assistant")
 
-def summary_audio(summary_text):
-    try:
-        speech = tts(summary_text)
-        tmp = tempfile.NamedTemporaryFile(delete=False, suffix=".wav")
-        import soundfile as sf
-        sf.write(tmp.name, speech["audio"], speech["sampling_rate"])
-        return tmp.name
-    except Exception as e:
-        return f"❌ Error in audio generation: {e}"
+uploaded_file = st.file_uploader("Upload your PDF", type="pdf")
 
-def summary_diagram(summary_text):
-    try:
-        dot = graphviz.Digraph()
-        dot.node("Summary", "📄 Summary")
-        for i, sentence in enumerate(summary_text.split(".")[:5]):
-            s = sentence.strip()
-            if not s:
-                continue
-            dot.node(f"S{i}", s[:40] + ("..." if len(s) > 40 else ""))
-            dot.edge("Summary", f"S{i}")
-        out_path = tempfile.NamedTemporaryFile(delete=False, suffix=".png").name
-        dot.render(out_path, format="png", cleanup=True)
-        return out_path + ".png"
-    except Exception as e:
-        return f"❌ Error in diagram generation: {e}"
+if uploaded_file:
+    with tempfile.NamedTemporaryFile(delete=False, suffix=".pdf") as tmp_file:
+        tmp_file.write(uploaded_file.read())
+        pdf_path = tmp_file.name
 
-# -------------------------
-# Gradio UI
-# -------------------------
-with gr.Blocks(css=".gradio-container {background-color: #f5f5f5}") as demo:
-    gr.Markdown("<h1 style='text-align:center;color:#4CAF50;'>📚 PDF Assistant</h1>")
-    with gr.Row():
-        with gr.Column():
-            pdf_input = gr.File(label="Upload PDF", file_types=[".pdf"])
-            summarize_btn = gr.Button("Summarize 🚀", variant="primary")
-            summary_output = gr.Textbox(label="Summary")
-            audio_output = gr.Audio(label="Summary Audio")
-            diagram_output = gr.Image(label="Summary Diagram")
+    reader = PdfReader(pdf_path)
+    text = "".join([page.extract_text() for page in reader.pages if page.extract_text()])
 
-    def full_pipeline(pdf_file):
-        summary = summarize_pdf(pdf_file)
-        if summary.startswith("❌"):
-            return summary, None, None
-        audio = summary_audio(summary)
-        diagram = summary_diagram(summary)
-        return summary, audio, diagram
+    if st.button("Summarize"):
+        response = client.chat.completions.create(
+            model="gpt-3.5-turbo",
+            messages=[{"role": "user", "content": f"Summarize this: {text[:4000]}"}],
+        )
+        st.subheader("Summary")
+        st.write(response.choices[0].message.content)
 
-    summarize_btn.click(full_pipeline, inputs=pdf_input, outputs=[summary_output, audio_output, diagram_output])
+    if st.button("Generate Diagram"):
+        response = client.chat.completions.create(
+            model="gpt-3.5-turbo",
+            messages=[{"role": "user", "content": f"Make a mermaid diagram for: {text[:2000]}"}],
+        )
+        st.subheader("Diagram")
+        st.code(response.choices[0].message.content, language="mermaid")
 
-if __name__ == "__main__":
-    demo.launch()
+    st.subheader("💬 Chat with PDF")
+    query = st.text_input("Ask a question about your PDF:")
+    if query:
+        response = client.chat.completions.create(
+            model="gpt-3.5-turbo",
+            messages=[{"role": "user", "content": f"Answer based on PDF: {query}\n\n{text[:4000]}"}],
+        )
+        st.write(response.choices[0].message.content)