Update src/streamlit_app.py

src/streamlit_app.py

@@ -1,40 +1,110 @@
-import altair as alt
-import numpy as np
-import pandas as pd
+import re
+import os
+import shutil
 import streamlit as st
+from langchain_huggingface import HuggingFaceEndpoint, HuggingFaceEmbeddings, ChatHuggingFace
+from langchain_text_splitters import RecursiveCharacterTextSplitter
+from langchain_community.vectorstores import Chroma
+from langchain_community.document_loaders import PyPDFLoader
+from langchain_core.messages import HumanMessage, SystemMessage
 
-"""
-# Welcome to Streamlit!
-
-Edit `/streamlit_app.py` to customize this app to your heart's desire :heart:.
-If you have any questions, checkout our [documentation](https://docs.streamlit.io) and [community
-forums](https://discuss.streamlit.io).
-
-In the meantime, below is an example of what you can do with just a few lines of code:
-"""
-
-num_points = st.slider("Number of points in spiral", 1, 10000, 1100)
-num_turns = st.slider("Number of turns in spiral", 1, 300, 31)
-
-indices = np.linspace(0, 1, num_points)
-theta = 2 * np.pi * num_turns * indices
-radius = indices
-
-x = radius * np.cos(theta)
-y = radius * np.sin(theta)
-
-df = pd.DataFrame({
-    "x": x,
-    "y": y,
-    "idx": indices,
-    "rand": np.random.randn(num_points),
-})
-
-st.altair_chart(alt.Chart(df, height=700, width=700)
-    .mark_point(filled=True)
-    .encode(
-        x=alt.X("x", axis=None),
-        y=alt.Y("y", axis=None),
-        color=alt.Color("idx", legend=None, scale=alt.Scale()),
-        size=alt.Size("rand", legend=None, scale=alt.Scale(range=[1, 150])),
-    ))
+# -----------------------------
+# 1. Page Configuration
+# -----------------------------
+st.set_page_config(page_title="AI Study Assistant", layout="wide")
+st.title("🎓 AI Study Assistant (Llama 3)")
+st.markdown("---")
+
+# Get token from secrets/environment
+token = os.environ.get("HUGGINGFACEHUB_API_TOKEN2")
+
+# -----------------------------
+# 2. RAG Logic
+# -----------------------------
+def process_lecture_pdf(uploaded_file):
+    temp_path = os.path.join("/tmp", uploaded_file.name)
+    with open(temp_path, "wb") as f:
+        f.write(uploaded_file.getbuffer())
+    
+    loader = PyPDFLoader(temp_path)
+    docs = loader.load()
+    text_splitter = RecursiveCharacterTextSplitter(chunk_size=700, chunk_overlap=100)
+    chunks = text_splitter.split_documents(docs)
+    
+    embeddings = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2")
+    
+    db_path = "/tmp/chroma_db"
+    if os.path.exists(db_path):
+        shutil.rmtree(db_path)
+        
+    vectorstore = Chroma.from_documents(
+        documents=chunks,
+        embedding=embeddings,
+        persist_directory=db_path
+    )
+    return vectorstore.as_retriever(search_kwargs={"k": 3}), docs
+
+# -----------------------------
+# 3. Model Setup (Llama 3 8B Instruct)
+# -----------------------------
+llm_endpoint = HuggingFaceEndpoint(
+    # Updated Repo ID for Llama 3
+    repo_id="meta-llama/Meta-Llama-3-8B-Instruct",
+    task="conversational",
+    huggingfacehub_api_token=token,
+    max_new_tokens=1024, # Llama 3 handles longer contexts well
+    temperature=0.6      # Slightly higher for better prose
+)
+
+chat_llm = ChatHuggingFace(llm=llm_endpoint)
+
+# -----------------------------
+# 4. User Interface
+# -----------------------------
+col1, col2 = st.columns([1, 2])
+
+with col1:
+    st.header("📂 Upload Notes")
+    uploaded_file = st.file_uploader("Upload Lecture PDF", type="pdf")
+    
+    if uploaded_file:
+        if 'retriever' not in st.session_state or st.session_state.get('last_file') != uploaded_file.name:
+            with st.spinner("Analyzing PDF with Llama 3..."):
+                retriever, full_docs = process_lecture_pdf(uploaded_file)
+                st.session_state.retriever = retriever
+                st.session_state.full_text = "\n".join([d.page_content for d in full_docs])
+                st.session_state.last_file = uploaded_file.name
+            st.success("Ready to study!")
+
+    st.header("📝 Summarize")
+    if st.button("Summarize Content"):
+        if 'full_text' in st.session_state:
+            with st.spinner("Llama 3 is summarizing..."):
+                # We can now use SystemMessage to give Llama 3 a "persona"
+                messages = [
+                    SystemMessage(content="You are a helpful university teaching assistant. Summarize the following text clearly."),
+                    HumanMessage(content=f"Notes: {st.session_state.full_text[:4000]}")
+                ]
+                response = chat_llm.invoke(messages)
+                st.write(response.content)
+        else:
+            st.warning("Please upload a PDF first.")
+
+with col2:
+    st.header("💬 Ask Questions")
+    user_query = st.text_input("Ask about your lecture:")
+    
+    if user_query:
+        if 'retriever' in st.session_state:
+            with st.spinner("Searching..."):
+                context_docs = st.session_state.retriever.invoke(user_query)
+                context_text = "\n\n".join([doc.page_content for doc in context_docs])
+                
+                messages = [
+                    SystemMessage(content="Use the provided context to answer the student's question accurately."),
+                    HumanMessage(content=f"Context: {context_text}\n\nQuestion: {user_query}")
+                ]
+                response = chat_llm.invoke(messages)
+                st.info(response.content)
+        else:
+            st.warning("Upload a PDF to start.")