Update app.py

app.py

@@ -6,19 +6,37 @@ from langchain_community.document_loaders import PyPDFLoader
 from langchain_groq import ChatGroq
 from langchain.text_splitter import RecursiveCharacterTextSplitter
 from langchain.chains import RetrievalQA
+from langchain.docstore.document import Document
 from tempfile import NamedTemporaryFile
+import nbformat
 
 # Load Groq API Key securely
 os.environ["GROQ_API_KEY"] = os.getenv("GROQ_API_KEY")
 
-# Function to process PDFs
-def process_pdfs(files):
+# Helper: Read .ipynb file and extract text
+def load_ipynb(file):
+    with NamedTemporaryFile(delete=False, suffix=".ipynb") as temp_file:
+        temp_file.write(file.read())
+        temp_file.flush()
+        with open(temp_file.name, "r", encoding="utf-8") as f:
+            nb = nbformat.read(f, as_version=nbformat.NO_CONVERT)
+            text = ""
+            for cell in nb.cells:
+                if cell.cell_type in ["markdown", "code"]:
+                    text += cell.source + "\n\n"
+        return [Document(page_content=text)]
+
+# Helper: Read PDF or IPYNB and build retriever chain
+def process_files(files):
     all_docs = []
     for file in files:
-        with NamedTemporaryFile(delete=False, suffix=".pdf") as temp_file:
-            temp_file.write(file.read())
-            loader = PyPDFLoader(temp_file.name)
-            all_docs.extend(loader.load())
+        if file.name.endswith(".pdf"):
+            with NamedTemporaryFile(delete=False, suffix=".pdf") as temp_file:
+                temp_file.write(file.read())
+                loader = PyPDFLoader(temp_file.name)
+                all_docs.extend(loader.load())
+        elif file.name.endswith(".ipynb"):
+            all_docs.extend(load_ipynb(file))
 
     splitter = RecursiveCharacterTextSplitter(chunk_size=500, chunk_overlap=50)
     chunks = splitter.split_documents(all_docs)
@@ -31,31 +49,29 @@ def process_pdfs(files):
     qa_chain = RetrievalQA.from_chain_type(llm=llm, retriever=retriever)
     return qa_chain
 
-# Global variable
+# Global chain
 qa_chain = None
 
-# Upload handler
-def upload_pdfs(files):
+def upload_docs(files):
     global qa_chain
-    qa_chain = process_pdfs(files)
-    return "✅ PDFs uploaded and processed. Now ask your questions."
+    qa_chain = process_files(files)
+    return "✅ PDFs or Notebooks uploaded and processed. Now ask your questions."
 
-# Question handler
 def ask_question(query):
     if qa_chain is None:
-        return "❌ Please upload Kaggle notebooks/competition PDFs first."
+        return "❌ Please upload PDFs or Kaggle Notebooks first."
     return qa_chain.run(query)
 
-# ✅ Gradio UI (fixed)
+# Gradio UI
 with gr.Blocks() as app:
-    gr.Markdown("## 🤖 Kaggle Study Assistant\nUpload PDFs from Kaggle and ask intelligent questions.")
+    gr.Markdown("## 🤖 Kaggle Study Assistant\nUpload Kaggle `.pdf` or `.ipynb` files and ask intelligent questions.")
 
     with gr.Row():
-        upload = gr.File(file_types=[".pdf"], file_count="multiple", label="Upload Kaggle PDFs")
-        btn_upload = gr.Button("📥 Process PDFs")
+        upload = gr.File(file_types=[".pdf", ".ipynb"], file_count="multiple", label="Upload Kaggle Files")
+        btn_upload = gr.Button("📥 Process Files")
 
     upload_output = gr.Textbox(label="Upload Status")
-    btn_upload.click(fn=upload_pdfs, inputs=upload, outputs=upload_output)
+    btn_upload.click(fn=upload_docs, inputs=upload, outputs=upload_output)
 
     question = gr.Textbox(label="Ask a question about uploaded notebooks")
     answer = gr.Textbox(label="Assistant Answer", interactive=False)