srenosh/rag-knowledge-base
0
1"""2ResolveAI — Enterprise RAG Knowledge Base3Hugging Face Spaces deployment.4"""5 6import os7import streamlit as st8from pathlib import Path9 10# ── Page config ─────────────────────────────────────────────────────────────11st.set_page_config(12 page_title="ResolveAI — Document Intelligence",13 page_icon="🔵",14 layout="wide",15 initial_sidebar_state="expanded",16 menu_items={17 'About': "ResolveAI — Intelligent document processing for federal and enterprise clients."18 }19)20 21# ── Custom CSS ───────────────────────────────────────────────────────────────22st.markdown("""23<style>24/* ── Base ── */25html, body, [class*="css"] {26 font-family: -apple-system, BlinkMacSystemFont, 'Inter', 'Segoe UI', sans-serif;27}28 29/* Hide Streamlit chrome */30#MainMenu, footer, header { visibility: hidden; }31.stDeployButton { display: none; }32[data-testid="stToolbar"] { display: none; }33 34/* ── App background ── */35.stApp {36 background-color: #0a0a0a;37}38 39/* ── Sidebar ── */40[data-testid="stSidebar"] {41 background-color: #0f0f0f;42 border-right: 1px solid rgba(255,255,255,0.06);43}44[data-testid="stSidebar"] * {45 color: rgba(255,255,255,0.75) !important;46}47 48/* ── Sidebar section headers ── */49.sidebar-label {50 font-size: 10px;51 font-weight: 600;52 letter-spacing: 0.18em;53 text-transform: uppercase;54 color: rgba(255,255,255,0.25) !important;55 margin-bottom: 10px;56 margin-top: 4px;57}58 59/* ── File uploader ── */60[data-testid="stFileUploader"] {61 border: 1px dashed rgba(255,255,255,0.1) !important;62 border-radius: 12px !important;63 background: rgba(255,255,255,0.02) !important;64 padding: 8px !important;65}66[data-testid="stFileUploader"]:hover {67 border-color: rgba(96,165,250,0.35) !important;68}69 70/* ── Buttons ── */71.stButton > button {72 background: rgba(255,255,255,0.04) !important;73 border: 1px solid rgba(255,255,255,0.1) !important;74 border-radius: 8px !important;75 color: rgba(255,255,255,0.7) !important;76 font-size: 13px !important;77 font-weight: 500 !important;78 transition: all 0.2s !important;79 height: 40px !important;80}81.stButton > button:hover {82 background: rgba(255,255,255,0.08) !important;83 border-color: rgba(255,255,255,0.2) !important;84 color: white !important;85}86 87/* Primary ingest button */88.stButton.primary > button,89button[kind="primary"] {90 background: #2563eb !important;91 border-color: transparent !important;92 color: white !important;93}94button[kind="primary"]:hover {95 background: #1d4ed8 !important;96}97 98/* ── Chat input ── */99[data-testid="stChatInput"] {100 border-top: 1px solid rgba(255,255,255,0.06) !important;101 background: #0a0a0a !important;102}103[data-testid="stChatInput"] textarea {104 background: rgba(255,255,255,0.04) !important;105 border: 1px solid rgba(255,255,255,0.1) !important;106 border-radius: 12px !important;107 color: white !important;108 font-size: 14px !important;109}110[data-testid="stChatInput"] textarea:focus {111 border-color: rgba(96,165,250,0.5) !important;112 box-shadow: 0 0 0 3px rgba(59,130,246,0.1) !important;113}114 115/* ── Chat messages ── */116[data-testid="stChatMessage"] {117 background: transparent !important;118 border: none !important;119 padding: 8px 0 !important;120}121 122/* ── User message bubble ── */123[data-testid="stChatMessage"][data-type="human"] .stMarkdown p {124 background: rgba(37,99,235,0.12) !important;125 border: 1px solid rgba(37,99,235,0.2) !important;126 border-radius: 12px 12px 4px 12px !important;127 padding: 12px 16px !important;128 display: inline-block !important;129 max-width: 85% !important;130 float: right !important;131 color: rgba(255,255,255,0.85) !important;132 font-size: 14px !important;133 line-height: 1.6 !important;134}135 136/* ── Assistant message ── */137[data-testid="stChatMessage"][data-type="ai"] .stMarkdown {138 color: rgba(255,255,255,0.75) !important;139 font-size: 14px !important;140 line-height: 1.75 !important;141}142 143/* ── Expander (sources) ── */144[data-testid="stExpander"] {145 background: rgba(255,255,255,0.02) !important;146 border: 1px solid rgba(255,255,255,0.06) !important;147 border-radius: 10px !important;148}149[data-testid="stExpander"] summary {150 font-size: 12px !important;151 color: rgba(255,255,255,0.35) !important;152}153 154/* ── Info / success / error ── */155[data-testid="stAlert"] {156 border-radius: 10px !important;157 font-size: 13px !important;158}159 160/* ── Divider ── */161hr {162 border-color: rgba(255,255,255,0.06) !important;163 margin: 16px 0 !important;164}165 166/* ── Scrollbar ── */167::-webkit-scrollbar { width: 4px; }168::-webkit-scrollbar-track { background: transparent; }169::-webkit-scrollbar-thumb { background: rgba(255,255,255,0.1); border-radius: 4px; }170 171/* ── Doc badge ── */172.doc-badge {173 display: inline-flex;174 align-items: center;175 gap: 6px;176 padding: 5px 12px;177 background: rgba(255,255,255,0.03);178 border: 1px solid rgba(255,255,255,0.08);179 border-radius: 100px;180 font-size: 12px;181 color: rgba(255,255,255,0.45);182 margin-bottom: 4px;183 width: 100%;184}185</style>186""", unsafe_allow_html=True)187 188# ── Lazy imports ─────────────────────────────────────────────────────────────189from llama_index.core import VectorStoreIndex, SimpleDirectoryReader, StorageContext190from llama_index.vector_stores.chroma import ChromaVectorStore191from llama_index.embeddings.openai import OpenAIEmbedding192from llama_index.llms.openai import OpenAI193from llama_index.core import Settings194import chromadb195 196# ── Paths ─────────────────────────────────────────────────────────────────────197UPLOADS_DIR = Path("data/uploads")198VECTORSTORE_DIR = Path("data/vectorstore")199UPLOADS_DIR.mkdir(parents=True, exist_ok=True)200VECTORSTORE_DIR.mkdir(parents=True, exist_ok=True)201 202OPENAI_API_KEY = os.environ.get("OPENAI_API_KEY", "")203 204# ── ChromaDB ──────────────────────────────────────────────────────────────────205@st.cache_resource206def get_chroma_client():207 return chromadb.PersistentClient(path=str(VECTORSTORE_DIR))208 209def get_collection():210 return get_chroma_client().get_or_create_collection("knowledge_base")211 212# ── Ingest ────────────────────────────────────────────────────────────────────213def ingest_file(file_path: str) -> dict:214 Settings.embed_model = OpenAIEmbedding(model="text-embedding-3-small", api_key=OPENAI_API_KEY)215 collection = get_collection()216 vector_store = ChromaVectorStore(chroma_collection=collection)217 storage_context = StorageContext.from_defaults(vector_store=vector_store)218 docs = SimpleDirectoryReader(input_files=[file_path]).load_data()219 if not docs:220 return {"status": "error", "message": "Could not parse document."}221 VectorStoreIndex.from_documents(docs, storage_context=storage_context)222 return {"status": "success", "count": len(docs)}223 224def list_documents() -> list:225 try:226 results = get_collection().get(include=["metadatas"])227 names = {m["file_name"] for m in results.get("metadatas", []) if m and "file_name" in m}228 return sorted(names)229 except Exception:230 return []231 232def clear_all():233 client = get_chroma_client()234 client.delete_collection("knowledge_base")235 client.get_or_create_collection("knowledge_base")236 for f in UPLOADS_DIR.iterdir():237 if f.is_file():238 f.unlink()239 240# ── Query ─────────────────────────────────────────────────────────────────────241def ask(question: str) -> dict:242 try:243 Settings.embed_model = OpenAIEmbedding(model="text-embedding-3-small", api_key=OPENAI_API_KEY)244 Settings.llm = OpenAI(245 model="gpt-4o",246 api_key=OPENAI_API_KEY,247 system_prompt=(248 "You are an enterprise knowledge base assistant. "249 "Answer questions strictly based on the provided documents. "250 "If the answer is not in the documents, say so clearly. "251 "Be concise, accurate, and professional. Always cite source documents."252 )253 )254 vector_store = ChromaVectorStore(chroma_collection=get_collection())255 index = VectorStoreIndex.from_vector_store(vector_store)256 engine = index.as_query_engine(similarity_top_k=5, response_mode="compact")257 response = engine.query(question)258 259 sources = []260 if hasattr(response, "source_nodes"):261 seen = set()262 for node in response.source_nodes:263 fname = node.metadata.get("file_name", "Unknown")264 if fname not in seen:265 sources.append({266 "file": fname,267 "score": round(node.score, 3) if node.score else None,268 "snippet": node.text[:220].strip().replace("\n", " ") + "…"269 })270 seen.add(fname)271 272 return {"status": "success", "answer": str(response), "sources": sources}273 except Exception as e:274 return {"status": "error", "answer": f"Error: {str(e)}", "sources": []}275 276# ── Session state ─────────────────────────────────────────────────────────────277if "messages" not in st.session_state:278 st.session_state.messages = []279 280# ═══════════════════════════════════════════════════════════281# SIDEBAR282# ═══════════════════════════════════════════════════════════283with st.sidebar:284 # Brand285 st.markdown("""286 <div style="padding: 4px 0 20px; border-bottom: 1px solid rgba(255,255,255,0.06); margin-bottom: 20px;">287 <div style="font-size: 17px; font-weight: 600; letter-spacing: -0.3px; color: white;">ResolveAI</div>288 <div style="font-size: 11px; color: rgba(255,255,255,0.25); margin-top: 2px; letter-spacing: 0.05em;">Document Intelligence</div>289 </div>290 """, unsafe_allow_html=True)291 292 if not OPENAI_API_KEY:293 st.error("OPENAI_API_KEY not configured.", icon="⚠️")294 295 # Upload296 st.markdown('<div class="sidebar-label">Upload Documents</div>', unsafe_allow_html=True)297 uploaded_files = st.file_uploader(298 "PDF, DOCX, TXT, MD",299 type=["pdf", "docx", "txt", "md"],300 accept_multiple_files=True,301 label_visibility="collapsed",302 )303 304 if uploaded_files:305 if st.button("Ingest Documents", use_container_width=True, type="primary"):306 for uf in uploaded_files:307 dest = UPLOADS_DIR / uf.name308 dest.write_bytes(uf.getvalue())309 with st.spinner(f"Processing {uf.name}…"):310 result = ingest_file(str(dest))311 if result["status"] == "success":312 st.success(f"{uf.name} added", icon="✅")313 else:314 st.error(f"{uf.name}: {result['message']}", icon="❌")315 st.rerun()316 317 st.markdown("<div style='height:16px'></div>", unsafe_allow_html=True)318 319 # Knowledge base320 docs = list_documents()321 st.markdown('<div class="sidebar-label">Knowledge Base</div>', unsafe_allow_html=True)322 323 if docs:324 for doc in docs:325 st.markdown(f'<div class="doc-badge">📄 {doc}</div>', unsafe_allow_html=True)326 else:327 st.markdown(328 '<div style="font-size:13px; color:rgba(255,255,255,0.2); padding: 8px 0;">No documents yet.</div>',329 unsafe_allow_html=True330 )331 332 st.markdown("<div style='height:16px'></div>", unsafe_allow_html=True)333 334 # Actions335 st.markdown('<div class="sidebar-label">Actions</div>', unsafe_allow_html=True)336 col1, col2 = st.columns(2)337 with col1:338 if st.button("Clear Chat", use_container_width=True):339 st.session_state.messages = []340 st.rerun()341 with col2:342 if st.button("Clear All", use_container_width=True):343 clear_all()344 st.session_state.messages = []345 st.rerun()346 347 # Footer348 st.markdown("""349 <div style="position: absolute; bottom: 24px; left: 24px; right: 24px;350 font-size: 11px; color: rgba(255,255,255,0.15); line-height: 1.6;">351 LlamaIndex · GPT-4o · ChromaDB<br>352 Built by <a href="https://linkedin.com/in/renosh-sunny-a69048222"353 style="color:rgba(96,165,250,0.5); text-decoration:none;">Renosh Sunny</a>354 </div>355 """, unsafe_allow_html=True)356 357 358# ═══════════════════════════════════════════════════════════359# MAIN PANEL360# ═══════════════════════════════════════════════════════════361 362# Header363st.markdown("""364<div style="padding: 48px 0 32px; border-bottom: 1px solid rgba(255,255,255,0.05); margin-bottom: 32px;">365 <div style="display:inline-flex; align-items:center; gap:8px; margin-bottom:14px;366 padding: 4px 14px; border-radius:100px;367 border: 1px solid rgba(255,255,255,0.08);368 font-size:11px; color:rgba(255,255,255,0.3); letter-spacing:0.18em; text-transform:uppercase;">369 Enterprise · RAG · AI370 </div>371 <h1 style="font-size: clamp(26px, 4vw, 38px); font-weight: 600; letter-spacing: -1px;372 color: white; margin: 0 0 10px; line-height: 1.1;">373 Document Intelligence374 </h1>375 <p style="font-size: 15px; color: rgba(255,255,255,0.35); font-weight: 300;376 margin: 0; max-width: 520px; line-height: 1.65;">377 Upload any document. Ask questions in plain English.378 The AI answers with cited sources — drawing only from your content.379 </p>380</div>381""", unsafe_allow_html=True)382 383# How to use — visible until the user starts chatting384if not st.session_state.messages:385 with st.expander("📖 How to use this app", expanded=True):386 col1, col2, col3, col4 = st.columns(4)387 with col1:388 st.markdown("#### 1️⃣ Upload")389 st.markdown("Click **Browse files** in the sidebar. Supports **PDF, DOCX, TXT, and Markdown**. You can upload multiple files at once.")390 with col2:391 st.markdown("#### 2️⃣ Ingest")392 st.markdown("Click **Ingest Documents**. The AI reads and indexes your content — takes a few seconds. Files appear under **Knowledge Base** when ready.")393 with col3:394 st.markdown("#### 3️⃣ Ask")395 st.markdown("Type any question in plain English in the chat box below. The AI answers using **only your documents**.")396 with col4:397 st.markdown("#### 4️⃣ Review Sources")398 st.markdown("Every answer has a **📎 Sources** dropdown. Click it to see which file and passage the answer came from.")399 400 st.divider()401 st.markdown("**💡 Example questions to try:**")402 st.markdown("""403- *What are the main topics covered in this document?*404- *Summarize the key findings*405- *What are the requirements for [topic]?*406- *What happens if [condition] occurs?*407- *List all deadlines or dates mentioned*408 """)409 st.caption("⚠️ This AI only answers from documents you upload — it will not search the internet. If the answer isn't in your documents, it will tell you.")410 411 if not docs:412 st.info("👈 Start by uploading a document in the sidebar.")413 414# Chat history415for msg in st.session_state.messages:416 with st.chat_message(msg["role"]):417 st.markdown(msg["content"])418 if msg.get("sources"):419 with st.expander(f"Sources · {len(msg['sources'])} document(s)"):420 for src in msg["sources"]:421 score_str = f" · Relevance: `{src['score']}`" if src.get("score") else ""422 st.markdown(f"**{src['file']}**{score_str}")423 st.caption(src["snippet"])424 st.divider()425 426# Chat input427if prompt := st.chat_input("Ask anything about your documents…"):428 if not docs:429 st.warning("Upload and ingest at least one document first.", icon="⚠️")430 else:431 st.session_state.messages.append({"role": "user", "content": prompt})432 with st.chat_message("user"):433 st.markdown(prompt)434 435 with st.chat_message("assistant"):436 with st.spinner("Searching knowledge base…"):437 result = ask(prompt)438 answer = result.get("answer", "No answer returned.")439 sources = result.get("sources", [])440 st.markdown(answer)441 if sources:442 with st.expander(f"Sources · {len(sources)} document(s)"):443 for src in sources:444 score_str = f" · Relevance: `{src['score']}`" if src.get("score") else ""445 st.markdown(f"**{src['file']}**{score_str}")446 st.caption(src["snippet"])447 st.divider()448 449 st.session_state.messages.append({450 "role": "assistant",451 "content": answer,452 "sources": sources,453 })454