import streamlit as st import pypdf import numpy as np import os import re import ast import json from datetime import datetime # ── Page config ─────────────────────────────────────────────────────────────── st.set_page_config( page_title="Nexus AI", page_icon="⚡", layout="wide", initial_sidebar_state="expanded", ) # ── CSS ─────────────────────────────────────────────────────────────────────── st.markdown(""" """, unsafe_allow_html=True) # ── Model loading (cached) ──────────────────────────────────────────────────── @st.cache_resource(show_spinner=False) def load_embedder(): from sentence_transformers import SentenceTransformer return SentenceTransformer('all-MiniLM-L6-v2') # ── PDF helpers ─────────────────────────────────────────────────────────────── def extract_pdf(file, doc_name: str): file.seek(0) reader = pypdf.PdfReader(file) chunks = [] full_text = "" CHUNK_SIZE, OVERLAP = 280, 55 for page_idx, page in enumerate(reader.pages): page_text = page.extract_text() or "" full_text += page_text + "\n" words = page_text.split() start = 0 while start < len(words): chunk_text = " ".join(words[start: start + CHUNK_SIZE]) if chunk_text.strip(): chunks.append({ "text": chunk_text, "page": page_idx + 1, "doc": doc_name, }) start += CHUNK_SIZE - OVERLAP meta = { "pages": len(reader.pages), "words": len(full_text.split()), "chunks": len(chunks), "full_text": full_text[:5000], } return chunks, meta # ── Embeddings ──────────────────────────────────────────────────────────────── def embed_chunks(chunks): model = load_embedder() texts = [c["text"] for c in chunks] return model.encode(texts, show_progress_bar=False, batch_size=32) def semantic_search(query: str, all_chunks, all_embeddings, k=5, doc_filter=None): from sklearn.metrics.pairwise import cosine_similarity if all_embeddings is None or len(all_chunks) == 0: return [] model = load_embedder() q_emb = model.encode([query]) if doc_filter: indices = [i for i, c in enumerate(all_chunks) if c["doc"] in doc_filter] else: indices = list(range(len(all_chunks))) if not indices: return [] filtered_embs = all_embeddings[indices] scores = cosine_similarity(q_emb, filtered_embs)[0] top_local = np.argsort(scores)[::-1][:k] results = [] for local_idx in top_local: global_idx = indices[local_idx] results.append({ "text": all_chunks[global_idx]["text"], "page": all_chunks[global_idx]["page"], "doc": all_chunks[global_idx]["doc"], "score": float(scores[local_idx]), }) return results # ── LLM helpers ─────────────────────────────────────────────────────────────── def _build_messages(system, user_msg, history): msgs = [{"role": "system", "content": system}] for h in history[-4:]: msgs.append({"role": "user", "content": h["q"]}) msgs.append({"role": "assistant", "content": h["a"]}) msgs.append({"role": "user", "content": user_msg}) return msgs def call_groq(api_key: str, system: str, user_msg: str, history): try: from groq import Groq client = Groq(api_key=api_key) r = client.chat.completions.create( model="llama-3.1-8b-instant", messages=_build_messages(system, user_msg, history), max_tokens=1024, temperature=0.3, ) return r.choices[0].message.content.strip() except Exception as e: err = str(e) if "429" in err: return "⚠️ Rate limit — wait a moment and retry." if "401" in err: return "⚠️ Invalid Groq API key. Check sidebar." return f"⚠️ Groq error: {err}" def call_gemini(api_key: str, system: str, user_msg: str, history): try: import google.generativeai as genai genai.configure(api_key=api_key) model = genai.GenerativeModel( model_name="gemini-1.5-flash", system_instruction=system, ) hist = [] for h in history[-4:]: hist.append({"role": "user", "parts": [h["q"]]}) hist.append({"role": "model", "parts": [h["a"]]}) chat = model.start_chat(history=hist) return chat.send_message(user_msg).text except Exception as e: err = str(e) if "API_KEY" in err or "invalid" in err.lower(): return "⚠️ Invalid Gemini API key. Check sidebar." return f"⚠️ Gemini error: {err}" def call_llm(provider, keys, system, user_msg, history): if provider == "Groq": return call_groq(keys.get("groq", ""), system, user_msg, history) return call_gemini(keys.get("gemini", ""), system, user_msg, history) # ── Core RAG answer ─────────────────────────────────────────────────────────── def get_answer(provider, keys, results, question, history): context_parts = [] for r in results: score_pct = int(r["score"] * 100) context_parts.append( f"[Source: {r['doc']} | Page {r['page']} | Relevance: {score_pct}%]\n{r['text']}" ) context = "\n\n---\n\n".join(context_parts) system = ( "You are Nexus, a precise AI document assistant.\n" "Rules:\n" "- Answer ONLY from the provided context. Never fabricate information.\n" "- If the answer is not in the context, say: 'This information is not available in the uploaded documents.'\n" "- Always cite the source page when referencing specific facts (e.g. 'According to Page 4...').\n" "- Use markdown formatting: **bold** for key terms, bullet lists for multi-point answers.\n" "- Be concise but complete." ) user_msg = f"RETRIEVED CONTEXT:\n{context}\n\nQUESTION: {question}" return call_llm(provider, keys, system, user_msg, history) # ── Auto summary ────────────────────────────────────────────────────────────── def get_summary(provider, keys, full_text): system = "You are a document analyst. Be structured and concise." user_msg = ( "Provide a structured document summary:\n\n" "**Main Topic:** (1 sentence)\n" "**Key Points:** (4 bullet points)\n" "**Important Data/Facts:** (3 bullet points)\n" "**Conclusion:** (1-2 sentences)\n\n" f"Document:\n{full_text[:4500]}" ) return call_llm(provider, keys, system, user_msg, []) # ── Suggested questions ─────────────────────────────────────────────────────── def get_suggestions(provider, keys, full_text): system = 'Output ONLY a Python list of 4 strings. No markdown, no explanation, no preamble.' user_msg = f'4 insightful questions for this document. Format exactly: ["Q1?","Q2?","Q3?","Q4?"]\n\n{full_text[:2500]}' raw = call_llm(provider, keys, system, user_msg, []) try: match = re.search(r'\[.*?\]', raw, re.DOTALL) if match: qs = ast.literal_eval(match.group()) if isinstance(qs, list) and len(qs) >= 4: return qs[:4] except Exception: pass return [ "What is the main topic of this document?", "What are the key findings or conclusions?", "What data or evidence is presented?", "What are the recommendations or next steps?", ] # ── Chat export ─────────────────────────────────────────────────────────────── def build_export(history, doc_names): lines = [ "NEXUS AI — CHAT EXPORT", "=" * 50, f"Date: {datetime.now().strftime('%Y-%m-%d %H:%M')}", f"Documents: {', '.join(doc_names)}", "=" * 50, "", ] for i, turn in enumerate(history, 1): lines.append(f"Q{i}: {turn['q']}") lines.append(f"A{i}: {turn['a']}") if turn.get("citations"): cites = ", ".join( f"p.{c['page']} in '{c['doc']}' ({int(c['score']*100)}%)" for c in turn["citations"] ) lines.append(f"Sources: {cites}") lines.append("") return "\n".join(lines) # ── Session state ───────────────────────────────────────────────────────────── _defaults = { "history": [], "all_chunks": [], "all_embeddings": None, "docs": {}, # doc_name -> meta "summaries": {}, # doc_name -> summary text "suggestions": [], "prefill": "", "provider": "Groq", "doc_filter": None, # None = all docs } for k, v in _defaults.items(): if k not in st.session_state: st.session_state[k] = v # ── API keys (env or session) ───────────────────────────────────────────────── _env_groq = os.environ.get("GROQ_API_KEY", "") _env_gemini = os.environ.get("GEMINI_API_KEY", "") if "groq_key" not in st.session_state: st.session_state.groq_key = _env_groq if "gemini_key" not in st.session_state: st.session_state.gemini_key = _env_gemini def keys(): return {"groq": st.session_state.groq_key, "gemini": st.session_state.gemini_key} def active_key_ok(): p = st.session_state.provider k = keys() return bool(k["groq"] if p == "Groq" else k["gemini"]) # ═══════════════════════════════════════════════════════════════════════════════ # SIDEBAR # ═══════════════════════════════════════════════════════════════════════════════ with st.sidebar: # ── Logo ────────────────────────────────────────────────────────────────── st.markdown(""" """, unsafe_allow_html=True) # ── Provider ────────────────────────────────────────────────────────────── st.markdown('
AI Provider', unsafe_allow_html=True) provider = st.radio( "provider", ["Groq", "Gemini"], horizontal=True, index=0 if st.session_state.provider == "Groq" else 1, label_visibility="collapsed", ) st.session_state.provider = provider st.markdown('
', unsafe_allow_html=True) # ── API Keys ────────────────────────────────────────────────────────────── st.markdown('
API Key', unsafe_allow_html=True) if provider == "Groq": k = st.text_input("groq_key", value=st.session_state.groq_key, placeholder="gsk_...", type="password", label_visibility="collapsed") st.session_state.groq_key = k else: k = st.text_input("gemini_key", value=st.session_state.gemini_key, placeholder="AIza...", type="password", label_visibility="collapsed") st.session_state.gemini_key = k st.markdown('
', unsafe_allow_html=True) # ── Upload ──────────────────────────────────────────────────────────────── st.markdown('
', unsafe_allow_html=True) st.markdown('
Upload PDF', unsafe_allow_html=True) uploaded = st.file_uploader("pdf", type="pdf", label_visibility="collapsed") if uploaded: if not active_key_ok(): st.error("Add your API key above first.") elif uploaded.name not in st.session_state.docs: with st.spinner("Indexing…"): new_chunks, meta = extract_pdf(uploaded, uploaded.name) new_embs = embed_chunks(new_chunks) st.session_state.all_chunks.extend(new_chunks) if st.session_state.all_embeddings is None: st.session_state.all_embeddings = new_embs else: st.session_state.all_embeddings = np.vstack( [st.session_state.all_embeddings, new_embs] ) st.session_state.docs[uploaded.name] = meta # auto summary summary = get_summary(provider, keys(), meta["full_text"]) st.session_state.summaries[uploaded.name] = summary # suggestions (from latest doc) st.session_state.suggestions = get_suggestions( provider, keys(), meta["full_text"] ) st.session_state.history = [] st.success(f"✓ Indexed {meta['pages']} pages") st.markdown('
', unsafe_allow_html=True) # ── Active documents ────────────────────────────────────────────────────── if st.session_state.docs: st.markdown('
', unsafe_allow_html=True) st.markdown('
Documents', unsafe_allow_html=True) for doc_name, meta in list(st.session_state.docs.items()): words_display = ( f"{meta['words']:,}" if meta["words"] < 10000 else f"{round(meta['words']/1000,1)}k" ) short_name = doc_name if len(doc_name) <= 22 else doc_name[:19] + "…" st.markdown(f"""
📄
{short_name}
{meta['pages']} pages · {words_display} words
{meta['chunks']}c
""", unsafe_allow_html=True) total_pages = sum(m["pages"] for m in st.session_state.docs.values()) total_words = sum(m["words"] for m in st.session_state.docs.values()) total_chunks = sum(m["chunks"] for m in st.session_state.docs.values()) tw = f"{round(total_words/1000,1)}k" if total_words >= 1000 else str(total_words) st.markdown(f"""
{len(st.session_state.docs)}
docs
{total_pages}
pages
{tw}
words
""", unsafe_allow_html=True) st.markdown('
', unsafe_allow_html=True) if st.button("🗑 Clear all documents", use_container_width=True): for k_s in ["history", "all_chunks", "suggestions", "summaries"]: st.session_state[k_s] = [] if k_s != "summaries" else {} st.session_state.all_embeddings = None st.session_state.docs = {} st.rerun() st.markdown('
', unsafe_allow_html=True) # ── Model info footer ───────────────────────────────────────────────────── st.markdown("""
Groq · Llama 3.1 8B Instant
Gemini · 1.5 Flash
Embeddings · all-MiniLM-L6-v2
""", unsafe_allow_html=True) # ═══════════════════════════════════════════════════════════════════════════════ # MAIN AREA # ═══════════════════════════════════════════════════════════════════════════════ st.markdown('
', unsafe_allow_html=True) # ── No documents loaded → hero ───────────────────────────────────────────── if not st.session_state.docs: st.markdown("""
Semantic Document Intelligence
Talk to any PDF.
Instantly.
Upload documents and get precise, cited answers powered by real semantic search — not keyword guessing.
01 /
Upload
Drop one or multiple PDFs
02 /
Index
Semantic vectors built instantly
03 /
Ask
Get answers with page citations
""", unsafe_allow_html=True) else: # ── Auto summary (collapsible) ───────────────────────────────────────── for doc_name, summary in st.session_state.summaries.items(): short = doc_name if len(doc_name) <= 30 else doc_name[:27] + "…" with st.expander(f"📋 Auto-summary — {short}", expanded=False): st.markdown(f"""
⚡ Document Summary
{summary}
""", unsafe_allow_html=True) # ── Suggested questions ──────────────────────────────────────────────── if not st.session_state.history and st.session_state.suggestions: st.markdown('
Suggested questions
', unsafe_allow_html=True) cols = st.columns(2) for i, q in enumerate(st.session_state.suggestions): with cols[i % 2]: if st.button(q, key=f"sq_{i}"): st.session_state.prefill = q st.rerun() st.markdown("
", unsafe_allow_html=True) # ── Chat history ─────────────────────────────────────────────────────── if st.session_state.history: st.markdown('
', unsafe_allow_html=True) for turn in st.session_state.history: # User bubble st.markdown(f"""
{turn['q']}
YOU
""", unsafe_allow_html=True) # AI bubble st.markdown(f"""
{turn['a']}
""", unsafe_allow_html=True) # Citation cards if turn.get("citations"): cite_html = "" for c in turn["citations"]: score_pct = int(c["score"] * 100) doc_short = c["doc"][:18] + "…" if len(c["doc"]) > 20 else c["doc"] cite_html += f"""
📄 {doc_short} p.{c['page']} {score_pct}%
""" st.markdown(f'
{cite_html}
', unsafe_allow_html=True) st.markdown('
', unsafe_allow_html=True) # Export link export_text = build_export(st.session_state.history, list(st.session_state.docs.keys())) st.download_button( "↓ Export chat", data=export_text, file_name=f"nexus_chat_{datetime.now().strftime('%Y%m%d_%H%M')}.txt", mime="text/plain", ) st.markdown('
', unsafe_allow_html=True) # ── Chat input ───────────────────────────────────────────────────────────── if st.session_state.docs: placeholder = ( f"Ask across {len(st.session_state.docs)} document(s)…" if len(st.session_state.docs) > 1 else "Ask anything about your document…" ) question = st.chat_input(placeholder) # Handle suggestion prefill if st.session_state.prefill and not question: question = st.session_state.prefill st.session_state.prefill = "" if question: if not active_key_ok(): st.error("Please add your API key in the sidebar.") else: with st.spinner("Searching & reasoning…"): results = semantic_search( question, st.session_state.all_chunks, st.session_state.all_embeddings, k=5, doc_filter=st.session_state.doc_filter, ) answer = get_answer( st.session_state.provider, keys(), results, question, st.session_state.history, ) st.session_state.history.append({ "q": question, "a": answer, "citations": results, }) st.rerun() else: st.chat_input("Upload a PDF to begin…", disabled=True)