knowledge_agent/web_app.py

"""
web_app.py — FastAPI web UI + REST API for the Knowledge Agent
==============================================================
Run with:  uvicorn web_app:app --host 0.0.0.0 --port 8000 --reload

Endpoints:
  GET  /              → Web UI (chat interface)
  POST /api/ask       → JSON Q&A
  GET  /api/stream    → SSE streaming answer
  POST /api/index     → Trigger re-indexing
  POST /api/upload    → Upload and index a new document
  GET  /api/docs-list → List indexed documents
  GET  /api/stats     → Knowledge base stats
  DELETE /api/doc/{filename} → Remove a document
"""

import os
import shutil
import asyncio
from pathlib import Path
from typing import Optional

from fastapi import FastAPI, UploadFile, File, HTTPException, Query
from fastapi.responses import HTMLResponse, StreamingResponse, JSONResponse
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel

from knowledge_api import (
    query_knowledge,
    index_docs,
    index_single,
    get_knowledge_stats,
    list_indexed_docs,
    delete_doc,
    stream_knowledge_answer,
)

DOCS_DIR = os.getenv("DOCS_DIR", "./documents")

# ── App setup ─────────────────────────────────────────────────────────────────
app = FastAPI(
    title="Knowledge Agent",
    description="Personal RAG system over your documents",
    version="1.0.0",
)

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_methods=["*"],
    allow_headers=["*"],
)


# ── Request / Response models ─────────────────────────────────────────────────
class AskRequest(BaseModel):
    question: str
    top_k: int = 5

class IndexRequest(BaseModel):
    docs_dir: Optional[str] = None
    force: bool = False


# ── REST API endpoints ────────────────────────────────────────────────────────

@app.post("/api/ask")
async def ask(req: AskRequest):
    """Ask a question and get a grounded answer from your documents."""
    if not req.question.strip():
        raise HTTPException(400, "Question cannot be empty.")
    result = query_knowledge(req.question, top_k=req.top_k)
    return result


@app.get("/api/stream")
async def stream(
    question: str = Query(..., description="Your question"),
    top_k: int = Query(5, description="Number of chunks to retrieve"),
):
    """
    SSE endpoint — streams the answer token by token.
    Used by the web UI for a typing-cursor effect.
    """
    if not question.strip():
        raise HTTPException(400, "Question cannot be empty.")

    def event_generator():
        for token in stream_knowledge_answer(question, top_k=top_k):
            # SSE format: "data: <token>\n\n"
            yield f"data: {token}\n\n"
        yield "data: [DONE]\n\n"

    return StreamingResponse(
        event_generator(),
        media_type="text/event-stream",
        headers={
            "Cache-Control": "no-cache",
            "X-Accel-Buffering": "no",
        },
    )


@app.post("/api/index")
async def trigger_index(req: IndexRequest):
    """Trigger indexing of the documents directory."""
    docs_dir = req.docs_dir or DOCS_DIR
    # Run in a thread pool so we don't block the event loop
    loop = asyncio.get_event_loop()
    await loop.run_in_executor(
        None, lambda: index_docs(docs_dir=docs_dir, force=req.force)
    )
    stats = get_knowledge_stats()
    return {"message": "Indexing complete.", "stats": stats}


@app.post("/api/upload")
async def upload_document(file: UploadFile = File(...)):
    """
    Upload a document file; save it to the documents dir and index it.
    Supports: .pdf, .docx, .txt, .md
    """
    allowed = {".pdf", ".docx", ".doc", ".txt", ".md", ".markdown"}
    ext = Path(file.filename).suffix.lower()
    if ext not in allowed:
        raise HTTPException(400, f"Unsupported file type: {ext}. Allowed: {allowed}")

    os.makedirs(DOCS_DIR, exist_ok=True)
    dest = os.path.join(DOCS_DIR, file.filename)

    # Save file
    with open(dest, "wb") as f:
        shutil.copyfileobj(file.file, f)

    # Index it
    loop = asyncio.get_event_loop()
    await loop.run_in_executor(None, lambda: index_single(dest))

    return {
        "message": f"'{file.filename}' uploaded and indexed.",
        "path": dest,
    }


@app.get("/api/docs-list")
async def docs_list():
    """List all documents currently in the knowledge base."""
    return {"documents": list_indexed_docs()}


@app.get("/api/stats")
async def stats():
    """Return knowledge base statistics."""
    return get_knowledge_stats()


@app.delete("/api/doc/{filename}")
async def remove_doc(filename: str):
    """Remove a document and all its chunks from the index."""
    delete_doc(filename)
    return {"message": f"Deleted '{filename}' from the knowledge base."}


# ── Web UI ────────────────────────────────────────────────────────────────────

@app.get("/", response_class=HTMLResponse)
async def ui():
    """Serve the single-page chat interface."""
    return HTMLResponse(content=WEB_UI_HTML)


# ── HTML for the web UI (self-contained, no build step needed) ────────────────
WEB_UI_HTML = """<!DOCTYPE html>
<html lang="en">
<head>
<meta charset="UTF-8">
<meta name="viewport" content="width=device-width, initial-scale=1.0">
<title>Knowledge Agent</title>
<link href="https://fonts.googleapis.com/css2?family=IBM+Plex+Mono:wght@400;500&family=IBM+Plex+Sans:wght@300;400;600&display=swap" rel="stylesheet">
<style>
  :root {
    --bg:       #0d0f14;
    --surface:  #161921;
    --border:   #2a2f3d;
    --accent:   #5b8def;
    --accent2:  #7ecfb0;
    --text:     #d4dae8;
    --muted:    #5c6480;
    --danger:   #e05c7a;
    --mono:     'IBM Plex Mono', monospace;
    --sans:     'IBM Plex Sans', sans-serif;
  }
  * { box-sizing: border-box; margin: 0; padding: 0; }
  body {
    background: var(--bg);
    color: var(--text);
    font-family: var(--sans);
    height: 100vh;
    display: flex;
    flex-direction: column;
  }

  /* ── Header ── */
  header {
    display: flex;
    align-items: center;
    justify-content: space-between;
    padding: 14px 24px;
    border-bottom: 1px solid var(--border);
    background: var(--surface);
    flex-shrink: 0;
  }
  .logo {
    font-family: var(--mono);
    font-size: 1rem;
    color: var(--accent);
    letter-spacing: .04em;
  }
  .logo span { color: var(--accent2); }
  .header-actions { display: flex; gap: 10px; }
  .btn {
    font-family: var(--mono);
    font-size: .75rem;
    padding: 6px 14px;
    border: 1px solid var(--border);
    border-radius: 4px;
    background: transparent;
    color: var(--text);
    cursor: pointer;
    transition: all .15s;
  }
  .btn:hover { border-color: var(--accent); color: var(--accent); }
  .btn.primary { background: var(--accent); border-color: var(--accent); color: #fff; }
  .btn.primary:hover { opacity: .85; }
  .btn.danger { border-color: var(--danger); color: var(--danger); }

  /* ── Layout ── */
  .main { display: flex; flex: 1; overflow: hidden; }

  /* ── Sidebar ── */
  .sidebar {
    width: 260px;
    border-right: 1px solid var(--border);
    background: var(--surface);
    display: flex;
    flex-direction: column;
    flex-shrink: 0;
    overflow: hidden;
  }
  .sidebar-section { padding: 16px; border-bottom: 1px solid var(--border); }
  .sidebar-title {
    font-family: var(--mono);
    font-size: .65rem;
    text-transform: uppercase;
    letter-spacing: .12em;
    color: var(--muted);
    margin-bottom: 10px;
  }
  #stats-panel { font-size: .78rem; line-height: 1.9; color: var(--muted); }
  #stats-panel strong { color: var(--text); }
  .doc-list { flex: 1; overflow-y: auto; padding: 8px 0; }
  .doc-item {
    padding: 8px 16px;
    font-size: .78rem;
    color: var(--muted);
    display: flex;
    align-items: center;
    gap: 8px;
    cursor: default;
  }
  .doc-item .dot { width: 6px; height: 6px; border-radius: 50%; background: var(--accent2); flex-shrink: 0; }
  .doc-item .del {
    margin-left: auto;
    color: var(--danger);
    cursor: pointer;
    font-size: .9rem;
    opacity: 0;
    transition: opacity .15s;
  }
  .doc-item:hover .del { opacity: 1; }

  /* Upload zone */
  .upload-zone {
    margin: 0 16px 16px;
    border: 1px dashed var(--border);
    border-radius: 6px;
    padding: 12px;
    text-align: center;
    font-size: .75rem;
    color: var(--muted);
    cursor: pointer;
    transition: border-color .15s;
  }
  .upload-zone:hover { border-color: var(--accent); color: var(--accent); }
  #file-input { display: none; }

  /* ── Chat area ── */
  .chat-area {
    flex: 1;
    display: flex;
    flex-direction: column;
    overflow: hidden;
  }
  .messages {
    flex: 1;
    overflow-y: auto;
    padding: 28px 32px;
    display: flex;
    flex-direction: column;
    gap: 24px;
  }
  .msg { display: flex; gap: 14px; }
  .msg-avatar {
    width: 30px; height: 30px;
    border-radius: 4px;
    display: flex; align-items: center; justify-content: center;
    font-family: var(--mono);
    font-size: .65rem;
    font-weight: 500;
    flex-shrink: 0;
    margin-top: 2px;
  }
  .msg.user .msg-avatar  { background: #2b3354; color: var(--accent); }
  .msg.agent .msg-avatar { background: #1e3330; color: var(--accent2); }
  .msg-body { flex: 1; }
  .msg-label {
    font-size: .68rem;
    font-family: var(--mono);
    color: var(--muted);
    margin-bottom: 5px;
    text-transform: uppercase;
    letter-spacing: .08em;
  }
  .msg-text {
    font-size: .875rem;
    line-height: 1.75;
    color: var(--text);
    white-space: pre-wrap;
    word-break: break-word;
  }
  .msg.user .msg-text { color: #aebcdc; }
  .sources-bar {
    margin-top: 10px;
    font-size: .7rem;
    font-family: var(--mono);
    color: var(--muted);
  }
  .sources-bar span {
    display: inline-block;
    background: #1a2238;
    border: 1px solid var(--border);
    border-radius: 3px;
    padding: 2px 8px;
    margin: 2px 3px 0 0;
    color: var(--accent);
  }

  /* ── Input row ── */
  .input-row {
    padding: 18px 32px;
    border-top: 1px solid var(--border);
    display: flex;
    gap: 10px;
    background: var(--surface);
    flex-shrink: 0;
  }
  #question-input {
    flex: 1;
    background: var(--bg);
    border: 1px solid var(--border);
    border-radius: 6px;
    padding: 12px 16px;
    font-family: var(--sans);
    font-size: .875rem;
    color: var(--text);
    resize: none;
    outline: none;
    transition: border-color .15s;
    line-height: 1.5;
    min-height: 48px;
    max-height: 160px;
  }
  #question-input:focus { border-color: var(--accent); }
  #question-input::placeholder { color: var(--muted); }
  #send-btn {
    align-self: flex-end;
    height: 48px;
    padding: 0 22px;
  }

  /* ── Status bar ── */
  #status-bar {
    padding: 6px 32px;
    font-size: .7rem;
    font-family: var(--mono);
    color: var(--muted);
    background: var(--bg);
    border-top: 1px solid var(--border);
    flex-shrink: 0;
  }

  /* Scrollbar */
  ::-webkit-scrollbar { width: 6px; }
  ::-webkit-scrollbar-track { background: transparent; }
  ::-webkit-scrollbar-thumb { background: var(--border); border-radius: 3px; }

  /* Loading dots */
  .typing { display: flex; gap: 5px; align-items: center; padding: 4px 0; }
  .typing .dot { width: 6px; height: 6px; border-radius: 50%; background: var(--accent); animation: blink 1.2s infinite; }
  .typing .dot:nth-child(2) { animation-delay: .2s; }
  .typing .dot:nth-child(3) { animation-delay: .4s; }
  @keyframes blink { 0%,80%,100%{opacity:.2} 40%{opacity:1} }

  .empty-state {
    flex: 1;
    display: flex;
    flex-direction: column;
    align-items: center;
    justify-content: center;
    color: var(--muted);
    font-size: .875rem;
    gap: 10px;
    text-align: center;
  }
  .empty-state .icon { font-size: 2.5rem; margin-bottom: 4px; }
</style>
</head>
<body>

<header>
  <div class="logo">knowledge<span>_agent</span></div>
  <div class="header-actions">
    <button class="btn" onclick="loadStats()">↻ Refresh</button>
    <button class="btn primary" onclick="triggerIndex()">⚡ Re-index</button>
  </div>
</header>

<div class="main">
  <!-- Sidebar -->
  <div class="sidebar">
    <div class="sidebar-section">
      <div class="sidebar-title">Knowledge Base</div>
      <div id="stats-panel">Loading…</div>
    </div>
    <div class="sidebar-section" style="flex:1;overflow:hidden;display:flex;flex-direction:column;padding-bottom:0">
      <div class="sidebar-title">Indexed Documents</div>
      <div class="doc-list" id="doc-list"></div>
    </div>
    <div class="upload-zone" onclick="document.getElementById('file-input').click()">
      ＋ Upload document
    </div>
    <input type="file" id="file-input" accept=".pdf,.docx,.txt,.md,.markdown"
           onchange="uploadFile(this.files[0])">
  </div>

  <!-- Chat -->
  <div class="chat-area">
    <div class="messages" id="messages">
      <div class="empty-state" id="empty-state">
        <div class="icon">🧠</div>
        <strong>Ask anything about your documents</strong>
        <span>Add files via the sidebar, then start asking questions.</span>
      </div>
    </div>

    <div class="input-row">
      <textarea id="question-input"
                placeholder="Ask a question about your documents…"
                rows="1"
                onkeydown="handleKey(event)"></textarea>
      <button class="btn primary" id="send-btn" onclick="sendQuestion()">Ask →</button>
    </div>
    <div id="status-bar">Ready</div>
  </div>
</div>

<script>
// ── State ──────────────────────────────────────────────────────────────────
let isLoading = false;

// ── On load ────────────────────────────────────────────────────────────────
window.addEventListener('DOMContentLoaded', () => {
  loadStats();
  loadDocs();
  autoResize(document.getElementById('question-input'));
});

// ── Stats & docs ───────────────────────────────────────────────────────────
async function loadStats() {
  try {
    const r = await fetch('/api/stats');
    const s = await r.json();
    document.getElementById('stats-panel').innerHTML =
      `<strong>${s.total_chunks}</strong> chunks indexed<br>
       <strong>${s.indexed_files}</strong> documents<br>
       <span style="font-size:.7rem">${s.embed_model}</span>`;
  } catch(e) {
    document.getElementById('stats-panel').textContent = 'Error loading stats.';
  }
}

async function loadDocs() {
  try {
    const r = await fetch('/api/docs-list');
    const data = await r.json();
    const list = document.getElementById('doc-list');
    list.innerHTML = '';
    if (!data.documents.length) {
      list.innerHTML = '<div class="doc-item" style="font-style:italic">No documents yet</div>';
      return;
    }
    data.documents.forEach(doc => {
      const div = document.createElement('div');
      div.className = 'doc-item';
      div.innerHTML = `
        <span class="dot"></span>
        <span style="overflow:hidden;text-overflow:ellipsis;white-space:nowrap"
              title="${doc.path}">${doc.name}</span>
        <span class="del" onclick="removeDoc('${doc.name}', this)" title="Remove">✕</span>`;
      list.appendChild(div);
    });
  } catch(e) {}
}

async function removeDoc(name, el) {
  if (!confirm(`Remove "${name}" from the knowledge base?`)) return;
  await fetch(`/api/doc/${encodeURIComponent(name)}`, { method: 'DELETE' });
  el.closest('.doc-item').remove();
  loadStats();
}

async function triggerIndex() {
  setStatus('Indexing documents…');
  try {
    const r = await fetch('/api/index', { method: 'POST',
      headers: {'Content-Type':'application/json'},
      body: JSON.stringify({ force: false })
    });
    const d = await r.json();
    setStatus(`Indexing complete — ${d.stats.total_chunks} chunks.`);
    loadStats(); loadDocs();
  } catch(e) { setStatus('Indexing failed.'); }
}

async function uploadFile(file) {
  if (!file) return;
  setStatus(`Uploading ${file.name}…`);
  const fd = new FormData();
  fd.append('file', file);
  try {
    const r = await fetch('/api/upload', { method: 'POST', body: fd });
    const d = await r.json();
    setStatus(d.message);
    loadStats(); loadDocs();
  } catch(e) { setStatus('Upload failed.'); }
  document.getElementById('file-input').value = '';
}

// ── Chat ───────────────────────────────────────────────────────────────────
function handleKey(e) {
  if (e.key === 'Enter' && !e.shiftKey) { e.preventDefault(); sendQuestion(); }
  setTimeout(() => autoResize(e.target), 0);
}

function autoResize(el) {
  el.style.height = 'auto';
  el.style.height = Math.min(el.scrollHeight, 160) + 'px';
}

async function sendQuestion() {
  const input = document.getElementById('question-input');
  const question = input.value.trim();
  if (!question || isLoading) return;

  isLoading = true;
  input.value = '';
  input.style.height = '48px';
  document.getElementById('send-btn').disabled = true;
  document.getElementById('empty-state')?.remove();

  addMessage('user', question);
  const agentMsg = addMessage('agent', '', true);
  setStatus('Searching knowledge base…');

  // SSE streaming
  try {
    const url = `/api/stream?question=${encodeURIComponent(question)}&top_k=5`;
    const es = new EventSource(url);
    let buffer = '';

    es.onmessage = (e) => {
      if (e.data === '[DONE]') {
        es.close();
        finishMessage(agentMsg, buffer);
        isLoading = false;
        document.getElementById('send-btn').disabled = false;
        setStatus('Ready');
        return;
      }
      buffer += e.data;
      agentMsg.querySelector('.msg-text').textContent = buffer;
      scrollToBottom();
    };
    es.onerror = () => {
      es.close();
      if (!buffer) {
        agentMsg.querySelector('.msg-text').textContent =
          'Error: Could not reach the server.';
      }
      isLoading = false;
      document.getElementById('send-btn').disabled = false;
      setStatus('Error');
    };
  } catch(e) {
    agentMsg.querySelector('.msg-text').textContent = 'Unexpected error.';
    isLoading = false;
    document.getElementById('send-btn').disabled = false;
    setStatus('Error');
  }
}

function addMessage(role, text, loading = false) {
  const container = document.getElementById('messages');
  const div = document.createElement('div');
  div.className = `msg ${role}`;
  const label = role === 'user' ? 'You' : 'Agent';
  const avatarText = role === 'user' ? 'YOU' : 'KB';
  div.innerHTML = `
    <div class="msg-avatar">${avatarText}</div>
    <div class="msg-body">
      <div class="msg-label">${label}</div>
      <div class="msg-text">${loading ? '<div class="typing"><div class="dot"></div><div class="dot"></div><div class="dot"></div></div>' : escapeHtml(text)}</div>
    </div>`;
  container.appendChild(div);
  scrollToBottom();
  return div;
}

function finishMessage(div, text) {
  div.querySelector('.msg-text').textContent = text;
  scrollToBottom();
}

function scrollToBottom() {
  const c = document.getElementById('messages');
  c.scrollTop = c.scrollHeight;
}

function setStatus(msg) {
  document.getElementById('status-bar').textContent = msg;
}

function escapeHtml(s) {
  return s.replace(/&/g,'&amp;').replace(/</g,'&lt;').replace(/>/g,'&gt;');
}
</script>
</body>
</html>"""