Update app.py

app.py

@@ -1,125 +1,67 @@
-'''
 import streamlit as st
-import os
-import shutil
-from retriever import find_video, build_embeddings
+from pinecone import Pinecone
+from langchain_groq import ChatGroq
+from langchain.prompts import PromptTemplate
+from langchain.chains.llm import LLMChain
 
-# Create directories
-os.makedirs("data/videos", exist_ok=True)
-os.makedirs("data/machine_catalogs", exist_ok=True)
+# ========== CONFIGURATION ========== #
+PINECONE_API_KEY = "your_pinecone_api_key"
+GROQ_API_KEY = "your_groq_api_key"
+INDEX_NAME = "rag-granite-index"
+NAMESPACE = "rag-ns"
 
-st.set_page_config(page_title="Maintenance Knowledge Portal", layout="centered")
+# ========== SETUP ========== #
+st.set_page_config(page_title="RAG Assistant", page_icon="🤖")
+st.title("💬 RAG-Powered Q&A Assistant")
 
-st.title("Maintenance Knowledge Portal")
-st.markdown("Upload videos, explore catalogs, and ask the AI assistant for help.")
+# Init Pinecone
+pc = Pinecone(api_key=PINECONE_API_KEY)
+index = pc.Index(INDEX_NAME)
 
-# --- Upload Section ---
-st.header("Upload Maintenance Video")
-video_file = st.file_uploader("Upload Video", type=["mp4", "mov", "avi"])
-if video_file:
-    with open(os.path.join("data/videos", video_file.name), "wb") as f:
-        f.write(video_file.read())
-    st.success("Video uploaded successfully!")
-    build_embeddings("data/videos")  # Update embeddings
+# Init Groq LLM
+llm = ChatGroq(
+    model_name="llama3-70b-8192",
+    api_key=GROQ_API_KEY
+)
 
+# Prompt Template
+prompt = PromptTemplate(
+    input_variables=["context", "question"],
+    template="""
+You are a smart assistant. Based on the provided context, answer the question in 1 to 2 lines only.
 
-st.header("Upload Machine Catalog (PDF)")
-pdf_file = st.file_uploader("Upload PDF", type=["pdf"])
-if pdf_file:
-    with open(os.path.join("data/machine_catalogs", pdf_file.name), "wb") as f:
-        f.write(pdf_file.read())
-    st.success("Catalog uploaded!")
+Context:
+{context}
 
+Question: {question}
 
-# --- Video Viewer Section ---
-st.header("Available Maintenance Videos")
-videos = os.listdir("data/videos")
-if videos:
-    for vid in videos:
-        st.video(os.path.join("data/videos", vid))
-else:
-    st.info("No videos uploaded yet.")
+Answer:"""
+)
 
+llm_chain = LLMChain(llm=llm, prompt=prompt)
 
-# --- Catalog Viewer Section ---
-st.header("Machine Catalogs")
-catalogs = os.listdir("data/machine_catalogs")
-if catalogs:
-    for cat in catalogs:
-        st.markdown(f"[{cat}](data/machine_catalogs/{cat})")
-else:
-    st.info("No catalogs available.")
+# ========== STREAMLIT UI ========== #
+user_query = st.text_input("Ask a question from the document 👇")
 
-
-# --- AI Assistant Section ---
-st.header("AI Assistant - Ask for Help")
-user_query = st.text_input("Describe the task or ask about a machine...")
 if user_query:
-    try:
-        video_file, desc = find_video(user_query)
-        st.subheader("Recommended Video:")
-        st.write(f"**{desc}**")
-        st.video(os.path.join("data/videos", video_file))
-    except Exception as e:
-        st.error("Sorry, couldn't find a matching video. Try uploading more videos.")
-'''
-
-
-import os
-import google.generativeai as genai
-import streamlit as st
-
-# Configure Google Gemini API Key
-genai.configure(api_key="AIzaSyBTt66oOvxpLeYn41sR-KkjSYPK2vOAqkU")
-
-# Define the model
-model = genai.GenerativeModel(model_name="models/gemini-1.5-pro")
-
-
-# Directory containing uploaded videos
-VIDEO_DIR = "videos"
-
-# --- Helper Function to Find Video ---
-def find_video(query):
-    """Search the video directory for relevant video files based on user input."""
-    matches = []
-    for filename in os.listdir(VIDEO_DIR):
-        # Simple match based on keywords in the filename
-        if query.lower() in filename.lower():
-            matches.append(filename)
-    return matches
-
-# --- Chatbot Section ---
-st.header("AI Assistant - Ask for Help")
-
-if "chat_history" not in st.session_state:
-    st.session_state.chat_history = []
-
-user_input = st.chat_input("Ask about machine, maintenance, or videos...")
-
-if user_input:
-    st.chat_message("user").write(user_input)
-
-    # Add to Gemini chat history
-    st.session_state.chat_history.append({"role": "user", "parts": [user_input]})
-
-    # Generate response from Gemini model
-    response = model.generate_content(st.session_state.chat_history)
-    reply = response.text
-
-    # Show Gemini's response
-    st.chat_message("assistant").write(reply)
-
-    # Save bot's reply to the session state
-    st.session_state.chat_history.append({"role": "model", "parts": [reply]})
-
-    # --- Video Recommendation based on query ---
-    matched_videos = find_video(user_input)
-
-    if matched_videos:
-        st.write("### Recommended Videos for You:")
-        for video in matched_videos:
-            # Show video using Streamlit video player
-            st.video(os.path.join(VIDEO_DIR, video))
-    else:
-        st.info("No video found for your query. Try uploading one or rephrasing.")
+    with st.spinner("Fetching answer..."):
+        from sentence_transformers import SentenceTransformer
+        embedder = SentenceTransformer("all-MiniLM-L6-v2")  # Match your embedding model
+        query_embedding = embedder.encode(user_query).tolist()
+
+        results = index.query(
+            namespace=NAMESPACE,
+            vector=query_embedding,
+            top_k=3,
+            include_metadata=True
+        )
+
+        context = "\n\n".join(match['metadata']['text'] for match in results['matches'])
+
+        response = llm_chain.invoke({
+            "context": context,
+            "question": user_query
+        })
+
+        st.success("Answer:")
+        st.write(response["text"])