Update app.py

app.py

@@ -1,71 +1,106 @@
 import os
-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 gradio as gr
+import fitz  # PyMuPDF
 import pinecone
-from qdrant_client import QdrantClient
+import qdrant_client
+from openai import OpenAI
+import graphviz
 
-# 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")
+# =================== CONFIG ===================
+OPENAI_API_KEY = os.environ.get("OPENAI_API_KEY")
+PINECONE_API_KEY = os.environ.get("PINECONE_API_KEY")
+QDRANT_API_KEY = os.environ.get("QDRANT_API_KEY")
 
 client = OpenAI(api_key=OPENAI_API_KEY)
-
-# Choose vector DB here
 VECTOR_DB = "qdrant"  # change to "pinecone" if needed
 
-# 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 extract_text_from_pdf(pdf_path):
+    doc = fitz.open(pdf_path)
+    text = ""
+    for page in doc:
+        text += page.get_text()
+    return text
 
-# Streamlit UI
-st.title("📄 PDF AI Assistant")
+def ingest_text(text, doc_name="doc"):
+    if VECTOR_DB == "qdrant":
+        qclient = qdrant_client.QdrantClient(":memory:")
+        qclient.recreate_collection(
+            collection_name=doc_name,
+            vectors_config={"size": 1536, "distance": "Cosine"}
+        )
+        qclient.upload_points(
+            collection_name=doc_name,
+            points=[
+                {"id": 0, "vector": [0.0]*1536, "payload": {"text": text}}
+            ]
+        )
+    elif VECTOR_DB == "pinecone":
+        pinecone.init(api_key=PINECONE_API_KEY, environment="gcp-starter")
+        if doc_name not in pinecone.list_indexes():
+            pinecone.create_index(doc_name, dimension=1536, metric="cosine")
+        index = pinecone.Index(doc_name)
+        index.upsert([(str(0), [0.0]*1536, {"text": text})])
+    return f"Ingested {len(text.split())} words."
 
-uploaded_file = st.file_uploader("Upload your PDF", type="pdf")
+def summarize_text(text):
+    resp = client.chat.completions.create(
+        model="gpt-3.5-turbo",
+        messages=[{"role":"system","content":"Summarize clearly."},
+                  {"role":"user","content":text[:4000]}]
+    )
+    return resp.choices[0].message.content
 
-if uploaded_file:
-    with tempfile.NamedTemporaryFile(delete=False, suffix=".pdf") as tmp_file:
-        tmp_file.write(uploaded_file.read())
-        pdf_path = tmp_file.name
+def generate_diagram(text):
+    dot = graphviz.Digraph()
+    dot.node("A", "PDF Content")
+    dot.node("B", "Summary")
+    dot.node("C", "Key Ideas")
+    dot.edges([("A","B"),("B","C")])
+    return dot.pipe(format="png")
 
-    reader = PdfReader(pdf_path)
-    text = "".join([page.extract_text() for page in reader.pages if page.extract_text()])
+def chat_with_pdf(text, question):
+    resp = client.chat.completions.create(
+        model="gpt-3.5-turbo",
+        messages=[
+            {"role":"system","content":"You are a helpful assistant with access to the document."},
+            {"role":"user","content":f"Document:\n{text[:3000]}\n\nQuestion:{question}"}
+        ]
+    )
+    return resp.choices[0].message.content
 
-    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)
+# =================== GRADIO APP ===================
+with gr.Blocks(theme="soft") as demo:
+    gr.Markdown("# 📚 PDF Assistant — Summarize, Diagram, Chat")
+    
+    with gr.Row():
+        pdf_file = gr.File(label="Upload PDF", file_types=[".pdf"])
+        doc_name = gr.Textbox(label="Doc name", value="mydoc")
+    
+    ingest_btn = gr.Button("🚀 Ingest PDF")
+    ingest_status = gr.Markdown("")
+    
+    summary_btn = gr.Button("📖 Summarize")
+    summary_output = gr.Textbox(label="Summary", lines=8)
+    
+    diagram_btn = gr.Button("📝 Generate Diagram")
+    diagram_output = gr.Image(type="numpy", label="Diagram Preview")
+    
+    with gr.Row():
+        question = gr.Textbox(label="Ask the PDF a question")
+        answer = gr.Textbox(label="Answer")
+        ask_btn = gr.Button("💬 Ask")
+    
+    pdf_text_state = gr.State("")
 
-    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")
+    def handle_ingest(pdf_file, doc_name):
+        text = extract_text_from_pdf(pdf_file.name)
+        status = ingest_text(text, doc_name)
+        return text, status
 
-    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)
+    ingest_btn.click(handle_ingest, [pdf_file, doc_name], [pdf_text_state, ingest_status])
+    summary_btn.click(lambda t: summarize_text(t), pdf_text_state, summary_output)
+    diagram_btn.click(lambda t: generate_diagram(t), pdf_text_state, diagram_output)
+    ask_btn.click(lambda t, q: chat_with_pdf(t, q), [pdf_text_state, question], answer)
+
+demo.launch()