Update app.py

app.py

@@ -1,188 +1,147 @@
 import streamlit as st
-import os
 from groq import Groq
+import os
 from PyPDF2 import PdfReader
 import requests
+import numpy as np
+import faiss
+from sentence_transformers import SentenceTransformer
 
 # ---------------------------
 # PAGE CONFIG
 # ---------------------------
-st.set_page_config(page_title="Krish GPT Pro", layout="wide")
-
-# ---------------------------
-# THEME & CSS
-# ---------------------------
-st.markdown("""
-<style>
-body {
-    background-color: #0a0b0d;
-    color: #d1d5db;
-}
-.chat-container {
-    max-width: 800px;
-    margin: auto;
-    padding-bottom: 180px;
-    height: 70vh;
-    overflow-y: auto;
-}
-.chat-bubble {
-    padding: 12px 16px;
-    border-radius: 15px;
-    margin-bottom: 8px;
-    display: inline-block;
-    max-width: 75%;
-    word-wrap: break-word;
-    font-size: 15px;
-    line-height: 1.4;
-}
-.user {
-    background-color: #4f46e5;
-    color: white;
-    margin-left: auto;
-}
-.assistant {
-    background-color: #1f2937;
-    color: #d1d5db;
-    margin-right: auto;
-}
-.input-container {
-    display: flex;
-    gap: 10px;
-    position: fixed;
-    bottom: 20px;
-    width: 80%;
-    max-width: 800px;
-    margin-left: auto;
-    margin-right: auto;
-    background-color: #111827;
-    padding: 10px;
-    border-radius: 12px;
-    box-shadow: 0 0 10px rgba(0,0,0,0.5);
-}
-textarea {
-    border-radius: 12px;
-    padding: 10px;
-    flex: 1;
-    background-color: #1f2937;
-    color: #d1d5db;
-    border: none;
-    font-size: 16px;
-    resize: none;
-    min-height: 40px;
-    max-height: 150px;
-}
-button {
-    background-color: #4f46e5;
-    color: white;
-    padding: 10px 16px;
-    border-radius: 12px;
-    border: none;
-    cursor: pointer;
-}
-button:hover {
-    background-color: #4338ca;
-}
-input[type="file"] {
-    border-radius: 12px;
-    padding: 5px;
-    background-color: #1f2937;
-    color: #d1d5db;
-}
-</style>
-""", unsafe_allow_html=True)
-
-# ---------------------------
-# SESSION STATE
+st.set_page_config(page_title="Krish GPT Multi-Modal RAG", layout="wide")
+st.title("🤖 Krish GPT Multi-Modal RAG")
+st.caption("PDF + Image OCR + RAG using Groq LLM 🚀")
+
+# ---------------------------
+# API KEYS
+# ---------------------------
+groq_api_key = os.getenv("GROQ_API_KEY")
+ocr_api_key = os.getenv("OCR_API_KEY")
+
+if not groq_api_key:
+    groq_api_key = st.text_input("Enter GROQ API Key", type="password")
+
+if not ocr_api_key:
+    ocr_api_key = st.text_input("Enter OCR.Space API Key", type="password")
+
+if not groq_api_key or not ocr_api_key:
+    st.stop()
+
+client = Groq(api_key=groq_api_key)
+
+# ---------------------------
+# EMBEDDING MODEL
+# ---------------------------
+@st.cache_resource
+def load_embedder():
+    return SentenceTransformer("all-MiniLM-L6-v2")
+
+embedder = load_embedder()
+
+# ---------------------------
+# OCR Function
+# ---------------------------
+def ocr_space_image(file, api_key):
+    url = "https://api.ocr.space/parse/image"
+    files = {'file': file}
+    data = {'apikey': api_key, 'language': 'eng'}
+    r = requests.post(url, files=files, data=data)
+    try:
+        result = r.json()
+        text = result['ParsedResults'][0]['ParsedText']
+    except:
+        text = ""
+    return text
+
+# ---------------------------
+# FILE UPLOAD
+# ---------------------------
+uploaded_file = st.file_uploader(
+    "Upload PDF or Image", type=["pdf", "png", "jpg", "jpeg"]
+)
+file_text = ""
+
+if uploaded_file:
+    if uploaded_file.type == "application/pdf":
+        reader = PdfReader(uploaded_file)
+        for page in reader.pages:
+            t = page.extract_text()
+            if t:
+                file_text += t
+    elif "image" in uploaded_file.type:
+        file_text = ocr_space_image(uploaded_file, ocr_api_key)
+
+# ---------------------------
+# TEXT CHUNKING & FAISS
+# ---------------------------
+def chunk_text(text, chunk_size=500):
+    chunks = []
+    for i in range(0, len(text), chunk_size):
+        chunks.append(text[i:i+chunk_size])
+    return chunks
+
+def build_index(chunks):
+    embeddings = embedder.encode(chunks)
+    dim = embeddings.shape[1]
+    index = faiss.IndexFlatL2(dim)
+    index.add(np.array(embeddings))
+    return index, embeddings
+
+def search(query, chunks, index):
+    q_emb = embedder.encode([query])
+    D, I = index.search(np.array(q_emb), k=min(3, len(chunks)))
+    results = [chunks[i] for i in I[0]]
+    return "\n".join(results)
+
+# ---------------------------
+# PROCESS FILE
+# ---------------------------
+if uploaded_file and file_text:
+    chunks = chunk_text(file_text)
+    index, embeddings = build_index(chunks)
+    st.session_state.rag_data = (chunks, index)
+
+# ---------------------------
+# CHAT MEMORY
 # ---------------------------
 if "messages" not in st.session_state:
     st.session_state.messages = []
 
+for msg in st.session_state.messages:
+    with st.chat_message(msg["role"]):
+        st.markdown(msg["content"])
+
+# ---------------------------
+# USER PROMPT
 # ---------------------------
-# GROQ CLIENT
-# ---------------------------
-client = Groq(api_key=os.getenv("GROQ_API_KEY"))
-OCR_API_KEY = os.getenv("OCR_API_KEY")
-
-# ---------------------------
-# DISPLAY CHAT
-# ---------------------------
-chat_container = st.empty()
-with chat_container.container():
-    for msg in st.session_state.messages:
-        role_class = "user" if msg["role"] == "user" else "assistant"
-        st.markdown(
-            f'<div class="chat-bubble {role_class}">{msg["content"]}</div>',
-            unsafe_allow_html=True
-        )
-
-# ---------------------------
-# INPUT + FILE UPLOAD (Merged)
-# ---------------------------
-with st.form("chat_form", clear_on_submit=True):
-    st.markdown('<div class="input-container">', unsafe_allow_html=True)
-    col1, col2, col3 = st.columns([6, 3, 1])
-    with col1:
-        prompt = st.text_area(
-            "Type a message...",
-            key="input_text",
-            placeholder="Press Enter to send, Ctrl+Enter for new line",
-            height=50
-        )
-    with col2:
-        uploaded_file = st.file_uploader("", label_visibility="collapsed")
-    with col3:
-        submitted = st.form_submit_button("Send")
-    st.markdown('</div>', unsafe_allow_html=True)
-
-    if submitted and (prompt.strip() != "" or uploaded_file):
-        context = ""
-        if uploaded_file:
-            if uploaded_file.type == "application/pdf":
-                reader = PdfReader(uploaded_file)
-                for page in reader.pages:
-                    text = page.extract_text()
-                    if text:
-                        context += text
-            else:
-                try:
-                    res = requests.post(
-                        "https://api.ocr.space/parse/image",
-                        files={"file": uploaded_file},
-                        data={"apikey": OCR_API_KEY}
-                    )
-                    context = res.json()['ParsedResults'][0]['ParsedText']
-                except:
-                    context = ""
-
-        st.session_state.messages.append({"role": "user", "content": prompt})
-
-        final_prompt = (prompt or "") + "\n" + context[:2000]
-
-        with st.spinner("🤖 Thinking..."):
-            try:
-                response = client.chat.completions.create(
-                    model="llama-3.3-70b-versatile",
-                    messages=[{"role": "user", "content": final_prompt}]
-                )
-                reply = response.choices[0].message.content
-            except:
-                reply = "⚠️ Something went wrong. Try again."
-
-        st.session_state.messages.append({"role": "assistant", "content": reply})
-
-# ---------------------------
-# JS Trick: Enter = Send, Ctrl+Enter = New line
-# ---------------------------
-st.markdown("""
-<script>
-const textarea = window.parent.document.querySelector('textarea');
-if (textarea) {
-    textarea.addEventListener('keydown', function(e){
-        if(e.key === 'Enter' && !e.ctrlKey){
-            e.preventDefault();
-            document.querySelector('button[kind="primary"]').click();
-        }
-    });
-}
-</script>
-""", unsafe_allow_html=True)
+prompt = st.chat_input("Ask anything...")
+
+if prompt:
+    st.session_state.messages.append({"role": "user", "content": prompt})
+    with st.chat_message("user"):
+        st.markdown(prompt)
+
+    context = ""
+    if "rag_data" in st.session_state:
+        chunks, index = st.session_state.rag_data
+        context = search(prompt, chunks, index)
+
+    with st.chat_message("assistant"):
+        try:
+            response = client.chat.completions.create(
+                model="llama-3.3-70b-versatile",
+                messages=[
+                    {"role": "system", "content": f"Context:\n{context}"},
+                    *st.session_state.messages
+                ],
+                temperature=0.7,
+                max_tokens=1024
+            )
+            reply = response.choices[0].message.content
+        except Exception as e:
+            reply = f"❌ Error: {str(e)}"
+
+        st.markdown(reply)
+        st.session_state.messages.append({"role": "assistant", "content": reply})