server.js

// ----------------------------------------------------------------------------
// oAI backend proxy.
//
// Endpoints:
//   GET  /health                         -> { ok: true }
//   POST /search   { query, maxResults } -> { sources:[{title,url}], notes }
//   POST /ingest   { name, ext, base64 } -> { docId, chunks }
//   POST /retrieve { docIds, query, k }  -> { context, chunks }
//
// The on-device GGUF model still generates every answer. This server only does
// the parts a phone is bad at: reliable web search and PDF/DOCX text extraction
// + embeddings for RAG.
// ----------------------------------------------------------------------------

import express from 'express';
import cors from 'cors';
import crypto from 'node:crypto';

import { webSearch, buildSourceNotes } from './src/search.js';
import { extractText } from './src/extract.js';
import { embed, embedBatch, warmup } from './src/embed.js';
import { recursiveSplit, putDoc, hasDoc, retrieve } from './src/store.js';

const app = express();
app.use(cors());
app.use(express.json({ limit: '30mb' })); // base64 file uploads can be large

const PORT = process.env.PORT || 8787;

app.get('/health', (_req, res) => res.json({ ok: true }));

// ---- web search ------------------------------------------------------------
app.post('/search', async (req, res) => {
  const { query, maxResults = 6 } = req.body || {};
  if (!query || typeof query !== 'string') {
    return res.status(400).json({ error: 'query required' });
  }
  try {
    const results = await webSearch(query, maxResults);
    // FIX (BUG-1): buildSourceNotes now needs the query so it can score
    // page sentences against the query keywords (like the Python reference).
    const notes = await buildSourceNotes(query, results);
    res.json({
      sources: results.map((r) => ({ title: r.title, url: r.url })),
      notes,
    });
  } catch (e) {
    console.error('[/search]', e);
    res.status(502).json({ error: 'search_failed', message: String(e.message || e) });
  }
});

// ---- ingest a file: extract -> chunk -> embed -> store ---------------------
app.post('/ingest', async (req, res) => {
  const { name, ext, base64 } = req.body || {};
  if (!base64 || typeof base64 !== 'string') {
    return res.status(400).json({ error: 'base64 required' });
  }
  try {
    // dedupe identical uploads by content hash
    const docId = crypto.createHash('sha256').update(base64).digest('hex').slice(0, 24);

    if (hasDoc(docId)) {
      return res.json({ docId, cached: true });
    }

    const buffer = Buffer.from(base64, 'base64');
    const text = await extractText(ext || guessExt(name), buffer);
    if (!text || text.trim().length === 0) {
      return res.status(422).json({ error: 'no_text', message: 'Could not extract text from file.' });
    }

    const chunks = recursiveSplit(text);
    const vectors = await embedBatch(chunks);
    putDoc(docId, chunks, vectors);

    res.json({ docId, chunks: chunks.length });
  } catch (e) {
    console.error('[/ingest]', e);
    res.status(500).json({ error: 'ingest_failed', message: String(e.message || e) });
  }
});

// ---- retrieve top-k context for a query ------------------------------------
app.post('/retrieve', async (req, res) => {
  const { docIds, query, k = 4 } = req.body || {};
  if (!Array.isArray(docIds) || docIds.length === 0) {
    return res.status(400).json({ error: 'docIds required' });
  }
  if (!query || typeof query !== 'string') {
    return res.status(400).json({ error: 'query required' });
  }
  try {
    const qvec = await embed(query);
    const top = retrieve(docIds, qvec, k);
    const context = top.map((c, i) => `[chunk ${i + 1}]\n${c}`).join('\n\n');
    res.json({ context, chunks: top.length });
  } catch (e) {
    console.error('[/retrieve]', e);
    res.status(500).json({ error: 'retrieve_failed', message: String(e.message || e) });
  }
});

function guessExt(name) {
  return String(name || '').split('.').pop()?.toLowerCase() || 'txt';
}

app.listen(PORT, () => {
  console.log(`oAI backend listening on :${PORT}`);
  // warm the embedding model so the first /ingest isn't slow
  warmup()
    .then(() => console.log('embedding model ready'))
    .catch((e) => console.warn('warmup failed (will lazy-load):', e.message));
});