streamlit_app.py

import streamlit as st
import os
import pdfplumber
from io import BytesIO
from PIL import Image
from docx import Document
import pandas as pd
import numpy as np
import faiss
from sentence_transformers import SentenceTransformer
from huggingface_hub import InferenceClient

# ============== CONFIG ==============
CHUNK_SIZE = 500
CHUNK_OVERLAP = 50

# ============== TEXT PROCESSING ==============
def chunk_text(text: str) -> list[dict]:
    if not text or not text.strip():
        return []
    
    text = " ".join(text.strip().split())
    chunks = []
    start = 0
    chunk_index = 0
    
    while start < len(text):
        end = start + CHUNK_SIZE
        chunk_content = text[start:end]
        
        if end < len(text):
            last_period = chunk_content.rfind(". ")
            if last_period > CHUNK_SIZE * 0.5:
                chunk_content = chunk_content[:last_period + 1]
                end = start + last_period + 1
        
        chunks.append({"content": chunk_content.strip(), "chunk_index": chunk_index})
        chunk_index += 1
        start = end - CHUNK_OVERLAP
        
        if start >= len(text) - CHUNK_OVERLAP:
            break
    
    return chunks

# ============== DOCUMENT PARSERS ==============
def parse_pdf(file_bytes) -> str:
    text_parts = []
    with pdfplumber.open(BytesIO(file_bytes)) as pdf:
        for i, page in enumerate(pdf.pages):
            page_text = page.extract_text() or ""
            if page_text.strip():
                text_parts.append(f"[Page {i + 1}]\n{page_text}")
    return "\n\n".join(text_parts)

def parse_docx(file_bytes) -> str:
    doc = Document(BytesIO(file_bytes))
    paragraphs = [para.text for para in doc.paragraphs if para.text.strip()]
    return "\n\n".join(paragraphs)

def parse_txt(file_bytes) -> str:
    return file_bytes.decode("utf-8")

def parse_image(file_bytes) -> str:
    return "[Image uploaded - OCR not available in cloud version]"

def parse_csv(file_bytes) -> str:
    df = pd.read_csv(BytesIO(file_bytes))
    lines = [f"Columns: {', '.join(df.columns.tolist())}", f"Total rows: {len(df)}", "\nData:"]
    for idx, row in df.head(50).iterrows():
        row_text = " | ".join([f"{col}: {val}" for col, val in row.items()])
        lines.append(row_text)
    return "\n".join(lines)

def parse_document(file_bytes, filename) -> dict:
    ext = filename.split(".")[-1].lower()
    
    if ext == "pdf":
        text = parse_pdf(file_bytes)
    elif ext == "docx":
        text = parse_docx(file_bytes)
    elif ext == "txt":
        text = parse_txt(file_bytes)
    elif ext in ["jpg", "jpeg", "png"]:
        text = parse_image(file_bytes)
    elif ext == "csv":
        text = parse_csv(file_bytes)
    else:
        text = ""
    
    chunks = chunk_text(text)
    for chunk in chunks:
        chunk["source"] = filename
        chunk["file_type"] = ext
    
    return {"text": text, "chunks": chunks}

# ============== EMBEDDING SERVICE ==============
@st.cache_resource
def load_embedding_model():
    return SentenceTransformer("all-MiniLM-L6-v2")

def embed_texts(texts: list[str]) -> np.ndarray:
    model = load_embedding_model()
    return model.encode(texts)

# ============== VECTOR STORE ==============
class SimpleVectorStore:
    def __init__(self):
        self.index = None
        self.documents = []
        self.dimension = 384
    
    def add_documents(self, chunks: list[dict]):
        if not chunks:
            return 0
        
        texts = [c["content"] for c in chunks]
        embeddings = embed_texts(texts).astype("float32")
        
        if self.index is None:
            self.index = faiss.IndexFlatL2(self.dimension)
        
        self.index.add(embeddings)
        self.documents.extend(chunks)
        return len(chunks)
    
    def search(self, query: str, top_k: int = 5) -> list[dict]:
        if self.index is None or self.index.ntotal == 0:
            return []
        
        query_embedding = embed_texts([query]).astype("float32")
        distances, indices = self.index.search(query_embedding, top_k)
        
        results = []
        for i, idx in enumerate(indices[0]):
            if 0 <= idx < len(self.documents):
                doc = self.documents[idx].copy()
                doc["score"] = float(distances[0][i])
                results.append(doc)
        return results
    
    def clear(self):
        self.index = None
        self.documents = []

# ============== LLM SERVICE ==============
@st.cache_resource
def get_llm_client():
    token = os.getenv("HUGGINGFACE_API_KEY", "")
    if not token:
        try:
            token = st.secrets["HUGGINGFACE_API_KEY"]
        except:
            token = ""
    return InferenceClient(model="HuggingFaceH4/zephyr-7b-beta", token=token)

def generate_answer(question: str, context: str) -> str:
    prompt = f"""You are a helpful assistant. Answer based on the context below.

CONTEXT:
{context}

QUESTION: {question}

ANSWER:"""
    
    try:
        client = get_llm_client()
        response = client.chat_completion(
            messages=[{"role": "user", "content": prompt}],
            max_tokens=512,
            temperature=0.7
        )
        return response.choices[0].message.content
    except Exception as e:
        return f"Error: {str(e)}"

# ============== STREAMLIT APP ==============
st.set_page_config(page_title="Smart RAG API", page_icon="🔍", layout="wide")

st.title("🔍 Smart RAG API")
st.markdown("Upload documents and ask questions - Powered by HuggingFace")

if "vector_store" not in st.session_state:
    st.session_state.vector_store = SimpleVectorStore()

# Sidebar
with st.sidebar:
    st.header("📊 Status")
    st.success("✅ Running")
    st.metric("Documents", len(st.session_state.vector_store.documents))
    
    if st.button("🗑️ Clear All"):
        st.session_state.vector_store.clear()
        st.rerun()
    
    st.divider()
    st.markdown("**Supported:** PDF, DOCX, TXT, CSV")

# Main columns
col1, col2 = st.columns(2)

with col1:
    st.header("📁 Upload")
    uploaded_file = st.file_uploader("Choose file", type=["pdf", "docx", "txt", "csv"])
    
    if uploaded_file and st.button("📤 Process", type="primary"):
        with st.spinner("Processing..."):
            try:
                parsed = parse_document(uploaded_file.getvalue(), uploaded_file.name)
                added = st.session_state.vector_store.add_documents(parsed["chunks"])
                st.success(f"✅ Added {added} chunks")
            except Exception as e:
                st.error(f"Error: {e}")

with col2:
    st.header("💬 Ask")
    question = st.text_area("Question:", placeholder="What is this about?")
    top_k = st.slider("Sources", 1, 5, 3)
    
    if st.button("🔍 Answer", type="primary"):
        if not question:
            st.warning("Enter a question")
        elif not st.session_state.vector_store.documents:
            st.warning("Upload documents first")
        else:
            with st.spinner("Thinking..."):
                results = st.session_state.vector_store.search(question, top_k)
                if results:
                    context = "\n\n".join([f"[{r['source']}]: {r['content']}" for r in results])
                    answer = generate_answer(question, context)
                    
                    st.subheader("📝 Answer")
                    st.write(answer)
                    
                    st.subheader("📚 Sources")
                    for r in results:
                        with st.expander(r["source"]):
                            st.write(r["content"][:300])

st.divider()
st.caption("Smart RAG API - FAISS + HuggingFace")