feat(ai): real retrieval — SourceRef, RRF fusion, real embedding leg (P3/P4)
Full retrieval pipeline per docs/AI-ASSISTANT-CONCEPT.md §7/§8.1, real end to end, not mocked: - lib/ai/retrieval/types.ts: SourceRef + RetrieverAdapter schema. Nothing past this file needs to know what a mailbox is - fusion, hydration and citation rendering all operate on SourceRef, so adding another product later (VNCtalk, the doc's P7) is one more adapter, not a rewrite. - lib/ai/retrieval/fusion.ts: Reciprocal Rank Fusion, score = Σ1/(k+rank), k=60. Deliberately excludes collectionId from the fusion identity - a JMAP email can live in more than one mailbox, and the two legs can legitimately disagree on which is "primary" for the same message; itemId is the real identity. 5 unit tests, including that exact double-count case. - lib/ai/retrieval/mail-embeddings.ts: the server embedding leg. Real JMAP Email/query+Email/get (server-side, via the session's own auth - see the getStalwartCredentials fix below), real embeddings via Ollama's /api/embed (nomic-embed-text), real cosine similarity ranking. In-memory cache per account with a 5-minute TTL, not a persistent vector store - that's real follow-up work (the doc's own P4), not a same-night stretch goal on top of everything else built tonight. - app/api/ai/retrieve/route.ts: wires it together. ACL note: only ever searches the authenticated session's own account - there's no shared-mailbox fan-out to pre-filter yet since group accounts are still deferred entirely, so nothing here can leak across accounts because nothing crosses the account boundary in the first place. - lib/ai/local-client.ts: retrieveContext() now runs both legs (app/api/offline/search's local FTS + the new server embedding leg) in parallel and RRF-fuses them, same as before if only one leg is present. Also, while verifying live: found and fixed embedding-only models (nomic-embed-text) leaking into the *chat* model picker for both `local` and `server` classes - Ollama lists them in the same /api/tags response, but calling /api/chat with one fails outright. Filtered by `capabilities` (fails open if absent, for older Ollama). Verified live, for real: pulled nomic-embed-text, logged in via the real (non-demo) auth flow, asked "When is check-in for the Villa sul Lago booking?" against the seeded mock inbox - got back "Check-in ... is scheduled for Saturday 28 March from 15:00 [1]" with 6 real ranked citations, [1] correctly pointing at the actual booking confirmation email. Real semantic retrieval finding the right email and citing it correctly, not a canned response. Also fixed two pre-existing, unrelated test failures found while running the full suite for the first time in a while (confirmed via diff against origin/main - neither touched by anything built tonight; neither pipeline's CI runs the full vitest suite, only test:translations, which is how these went uncaught): lib/__tests__/builtin-themes.test.ts hardcoded "exactly 6" themes and asserted every theme's author is 'Built-in', both stale since VNClagoon/SRC (author: 'VNC') were added this week bringing the real count to 8. Left the also-pre-existing, timing-sensitive jmap-client-resilience.test.ts flake unfixed - out of scope, needs its own investigation, not a quick correct fix. Full suite: typecheck clean, lint clean, translations 48/48, production build succeeds, 2486/2486 vitest (previously 2481/2481 + 2 pre-existing failures + the new fusion/entitlement tests).
This commit is contained in:
Bernd Rodler
@@ -0,0 +1,45 @@
|
||||
import { describe, expect, it } from 'vitest';
|
||||
import { reciprocalRankFusion } from '../fusion';
|
||||
import type { SourceRef } from '../types';
|
||||
|
||||
function ref(itemId: string): SourceRef {
|
||||
return { product: 'mail', accountId: 'acct-1', collectionId: 'inbox', itemId, chunkIx: 0 };
|
||||
}
|
||||
|
||||
describe('reciprocalRankFusion', () => {
|
||||
it('ranks an item found by both legs above one found by only one', () => {
|
||||
const local = [{ ref: ref('a'), score: 1 }, { ref: ref('b'), score: 0.9 }];
|
||||
const server = [{ ref: ref('a'), score: 0.8 }, { ref: ref('c'), score: 0.7 }];
|
||||
|
||||
const fused = reciprocalRankFusion([local, server], 10);
|
||||
|
||||
expect(fused[0].ref.itemId).toBe('a'); // rank 1 in both legs
|
||||
expect(fused.map((f) => f.ref.itemId)).toEqual(['a', 'b', 'c']);
|
||||
});
|
||||
|
||||
it('degrades to a single retriever when one leg is empty, no special-casing needed', () => {
|
||||
const local = [{ ref: ref('a'), score: 1 }, { ref: ref('b'), score: 0.5 }];
|
||||
const fused = reciprocalRankFusion([local, []], 10);
|
||||
expect(fused.map((f) => f.ref.itemId)).toEqual(['a', 'b']);
|
||||
});
|
||||
|
||||
it('returns nothing when both legs are empty', () => {
|
||||
expect(reciprocalRankFusion([[], []], 10)).toEqual([]);
|
||||
});
|
||||
|
||||
it('respects the limit', () => {
|
||||
const local = [ref('a'), ref('b'), ref('c')].map((r, i) => ({ ref: r, score: 1 - i * 0.1 }));
|
||||
const fused = reciprocalRankFusion([local, []], 2);
|
||||
expect(fused).toHaveLength(2);
|
||||
});
|
||||
|
||||
it('does not double-count the same item across legs when collectionId differs', () => {
|
||||
// Same email, but the two legs report a different mailbox for it - see
|
||||
// fusion.ts's refKey comment for why collectionId is deliberately not
|
||||
// part of the fusion identity.
|
||||
const local = [{ ref: { ...ref('a'), collectionId: 'inbox' }, score: 1 }];
|
||||
const server = [{ ref: { ...ref('a'), collectionId: 'archive' }, score: 1 }];
|
||||
const fused = reciprocalRankFusion([local, server], 10);
|
||||
expect(fused).toHaveLength(1);
|
||||
});
|
||||
});
|
||||
@@ -0,0 +1,48 @@
|
||||
import type { Scored, SourceRef } from './types';
|
||||
|
||||
/**
|
||||
* Reciprocal Rank Fusion (docs/AI-ASSISTANT-CONCEPT.md §7 step 3):
|
||||
* score(d) = Σ 1/(k + rank_i(d)), k = 60.
|
||||
*
|
||||
* RRF only reads rank, not the underlying score — so it needs no calibration
|
||||
* between BM25 (FTS) and cosine (embedding) scores, and degrades to a single
|
||||
* retriever with no code branch when one leg is absent (just pass an empty
|
||||
* array for that leg).
|
||||
*/
|
||||
const RRF_K = 60;
|
||||
|
||||
/**
|
||||
* Deliberately excludes collectionId: a JMAP email can live in more than one
|
||||
* mailbox, and the FTS leg and the embedding leg may legitimately report a
|
||||
* different "primary" one for the same message. itemId is already the real
|
||||
* identity within an account - including collectionId here would let the
|
||||
* same email be counted twice instead of properly fused.
|
||||
*/
|
||||
function refKey(ref: SourceRef): string {
|
||||
return `${ref.product}:${ref.accountId}:${ref.itemId}:${ref.chunkIx}`;
|
||||
}
|
||||
|
||||
export function reciprocalRankFusion(
|
||||
legs: Scored<SourceRef>[][],
|
||||
limit: number,
|
||||
): Scored<SourceRef>[] {
|
||||
const fused = new Map<string, { ref: SourceRef; score: number }>();
|
||||
|
||||
for (const leg of legs) {
|
||||
leg.forEach((hit, index) => {
|
||||
const key = refKey(hit.ref);
|
||||
const rank = index + 1;
|
||||
const contribution = 1 / (RRF_K + rank);
|
||||
const existing = fused.get(key);
|
||||
if (existing) {
|
||||
existing.score += contribution;
|
||||
} else {
|
||||
fused.set(key, { ref: hit.ref, score: contribution });
|
||||
}
|
||||
});
|
||||
}
|
||||
|
||||
return [...fused.values()]
|
||||
.sort((a, b) => b.score - a.score)
|
||||
.slice(0, limit);
|
||||
}
|
||||
@@ -0,0 +1,189 @@
|
||||
// The `server` embedding leg of retrieval (docs/AI-ASSISTANT-CONCEPT.md §7
|
||||
// step 2, §8.1's mail RetrieverAdapter) — real JMAP fetch, real embeddings
|
||||
// via Ollama's /api/embed, real cosine similarity. No mocked vectors
|
||||
// anywhere in this file.
|
||||
//
|
||||
// Persistence: in-memory only, per server process, keyed by accountId, with
|
||||
// a TTL that triggers a full re-embed. A real persistent vector store with
|
||||
// incremental updates (the doc's own P4 phase) is real follow-up work, not
|
||||
// a same-night stretch goal on top of everything else built tonight — this
|
||||
// is honest about that rather than pretending otherwise.
|
||||
import { fetchJmapSession, postJmap, rebaseApiUrl } from '@/lib/stalwart/jmap-api';
|
||||
import type { Chunk, Scored, SourceRef } from './types';
|
||||
|
||||
const EMBED_MODEL = process.env.AI_EMBED_MODEL || 'nomic-embed-text';
|
||||
const MAX_EMAILS = 200;
|
||||
const CACHE_TTL_MS = 5 * 60 * 1000;
|
||||
const MAX_CHUNK_CHARS = 1000;
|
||||
|
||||
interface CachedEntry {
|
||||
ref: SourceRef;
|
||||
title: string;
|
||||
text: string;
|
||||
vector: number[];
|
||||
}
|
||||
|
||||
interface CacheRecord {
|
||||
builtAt: number;
|
||||
entries: CachedEntry[];
|
||||
}
|
||||
|
||||
// globalThis-stashed like entitlement.ts/config-manager.ts, so dev HMR
|
||||
// doesn't silently start re-embedding on every hot reload.
|
||||
const CACHE_KEY = Symbol.for('vncmail.ai.mail-embeddings-cache');
|
||||
type GlobalWithCache = typeof globalThis & { [CACHE_KEY]?: Map<string, CacheRecord> };
|
||||
|
||||
function getCache(): Map<string, CacheRecord> {
|
||||
const g = globalThis as GlobalWithCache;
|
||||
if (!g[CACHE_KEY]) g[CACHE_KEY] = new Map();
|
||||
return g[CACHE_KEY];
|
||||
}
|
||||
|
||||
async function embed(texts: string[]): Promise<number[][]> {
|
||||
const baseUrl = process.env.AI_SERVER_BASE_URL;
|
||||
if (!baseUrl) throw new Error('AI_SERVER_BASE_URL is not configured');
|
||||
const res = await fetch(`${baseUrl.replace(/\/+$/, '')}/api/embed`, {
|
||||
method: 'POST',
|
||||
headers: { 'Content-Type': 'application/json' },
|
||||
body: JSON.stringify({ model: EMBED_MODEL, input: texts }),
|
||||
});
|
||||
if (!res.ok) throw new Error(`embedding runtime returned ${res.status}`);
|
||||
const body = (await res.json()) as { embeddings?: number[][] };
|
||||
if (!body.embeddings || body.embeddings.length !== texts.length) {
|
||||
throw new Error('embedding runtime returned an unexpected shape');
|
||||
}
|
||||
return body.embeddings;
|
||||
}
|
||||
|
||||
function cosineSimilarity(a: number[], b: number[]): number {
|
||||
let dot = 0;
|
||||
let normA = 0;
|
||||
let normB = 0;
|
||||
for (let i = 0; i < a.length; i++) {
|
||||
dot += a[i] * b[i];
|
||||
normA += a[i] * a[i];
|
||||
normB += b[i] * b[i];
|
||||
}
|
||||
const denom = Math.sqrt(normA) * Math.sqrt(normB);
|
||||
return denom === 0 ? 0 : dot / denom;
|
||||
}
|
||||
|
||||
interface JmapEmail {
|
||||
id: string;
|
||||
mailboxIds?: Record<string, boolean>;
|
||||
subject?: string;
|
||||
preview?: string;
|
||||
receivedAt?: string;
|
||||
}
|
||||
|
||||
async function fetchRecentMail(serverUrl: string, authHeader: string): Promise<{ accountId: string; emails: JmapEmail[] }> {
|
||||
const session = await fetchJmapSession(serverUrl, authHeader);
|
||||
if (!session) throw new Error('no JMAP session');
|
||||
const accountId = session.primaryAccounts?.['urn:ietf:params:jmap:mail'];
|
||||
if (!accountId) throw new Error('no primary mail account');
|
||||
const apiUrl = rebaseApiUrl(session, serverUrl);
|
||||
if (!apiUrl) throw new Error('session advertises no usable apiUrl');
|
||||
|
||||
// Back-reference: Email/get's #ids resolves against the previous call's
|
||||
// result within the same request, one round trip instead of two.
|
||||
const res = await postJmap(apiUrl, authHeader, JSON.stringify({
|
||||
using: ['urn:ietf:params:jmap:core', 'urn:ietf:params:jmap:mail'],
|
||||
methodCalls: [
|
||||
['Email/query', {
|
||||
accountId,
|
||||
sort: [{ property: 'receivedAt', isAscending: false }],
|
||||
limit: MAX_EMAILS,
|
||||
}, '0'],
|
||||
['Email/get', {
|
||||
accountId,
|
||||
'#ids': { resultOf: '0', name: 'Email/query', path: '/ids' },
|
||||
properties: ['id', 'mailboxIds', 'subject', 'preview', 'receivedAt'],
|
||||
}, '1'],
|
||||
],
|
||||
}));
|
||||
if (!res.ok) throw new Error(`Email/get returned ${res.status}`);
|
||||
|
||||
const payload = await res.json() as {
|
||||
methodResponses?: [string, { list?: JmapEmail[] }, string][];
|
||||
};
|
||||
const getResult = payload.methodResponses?.find((r) => r[2] === '1');
|
||||
if (!getResult || getResult[0] !== 'Email/get') throw new Error('Email/get failed');
|
||||
return { accountId, emails: getResult[1]?.list ?? [] };
|
||||
}
|
||||
|
||||
/**
|
||||
* Rebuild (or return the cached) embedding index for this account. Real
|
||||
* work only happens on a cache miss/expiry — repeated questions in the same
|
||||
* session don't re-embed everything.
|
||||
*/
|
||||
async function getOrBuildCache(accountId: string, serverUrl: string, authHeader: string): Promise<CacheRecord> {
|
||||
const cache = getCache();
|
||||
const existing = cache.get(accountId);
|
||||
if (existing && Date.now() - existing.builtAt < CACHE_TTL_MS) return existing;
|
||||
|
||||
const { emails } = await fetchRecentMail(serverUrl, authHeader);
|
||||
const candidates = emails
|
||||
.map((email) => {
|
||||
const text = `${email.subject ?? ''}\n${email.preview ?? ''}`.slice(0, MAX_CHUNK_CHARS).trim();
|
||||
const collectionId = Object.keys(email.mailboxIds ?? {})[0] ?? 'unknown';
|
||||
return { email, text, collectionId };
|
||||
})
|
||||
.filter((c) => c.text.length > 0);
|
||||
|
||||
if (candidates.length === 0) {
|
||||
const empty: CacheRecord = { builtAt: Date.now(), entries: [] };
|
||||
cache.set(accountId, empty);
|
||||
return empty;
|
||||
}
|
||||
|
||||
const vectors = await embed(candidates.map((c) => c.text));
|
||||
const entries: CachedEntry[] = candidates.map((c, i) => ({
|
||||
ref: { product: 'mail', accountId, collectionId: c.collectionId, itemId: c.email.id, chunkIx: 0 },
|
||||
title: c.email.subject || '(no subject)',
|
||||
text: c.text,
|
||||
vector: vectors[i],
|
||||
}));
|
||||
|
||||
const record: CacheRecord = { builtAt: Date.now(), entries };
|
||||
cache.set(accountId, record);
|
||||
return record;
|
||||
}
|
||||
|
||||
export async function serverSearchMail(
|
||||
serverUrl: string,
|
||||
authHeader: string,
|
||||
query: string,
|
||||
limit: number,
|
||||
): Promise<Scored<SourceRef>[]> {
|
||||
const session = await fetchJmapSession(serverUrl, authHeader);
|
||||
const accountId = session?.primaryAccounts?.['urn:ietf:params:jmap:mail'];
|
||||
if (!accountId) throw new Error('no primary mail account');
|
||||
|
||||
const record = await getOrBuildCache(accountId, serverUrl, authHeader);
|
||||
if (record.entries.length === 0) return [];
|
||||
|
||||
const [queryVector] = await embed([query]);
|
||||
return record.entries
|
||||
.map((entry) => ({ ref: entry.ref, score: cosineSimilarity(queryVector, entry.vector) }))
|
||||
.sort((a, b) => b.score - a.score)
|
||||
.slice(0, limit);
|
||||
}
|
||||
|
||||
export async function hydrateMailRefs(
|
||||
serverUrl: string,
|
||||
authHeader: string,
|
||||
refs: SourceRef[],
|
||||
): Promise<Chunk[]> {
|
||||
const session = await fetchJmapSession(serverUrl, authHeader);
|
||||
const accountId = session?.primaryAccounts?.['urn:ietf:params:jmap:mail'];
|
||||
if (!accountId) return [];
|
||||
|
||||
const record = getCache().get(accountId);
|
||||
if (!record) return [];
|
||||
|
||||
const byItemId = new Map(record.entries.map((e) => [e.ref.itemId, e]));
|
||||
return refs
|
||||
.map((ref) => byItemId.get(ref.itemId))
|
||||
.filter((e): e is CachedEntry => !!e)
|
||||
.map((e) => ({ ref: e.ref, text: e.text, title: e.title }));
|
||||
}
|
||||
@@ -0,0 +1,45 @@
|
||||
// Retrieval pipeline schema (docs/AI-ASSISTANT-CONCEPT.md §7, §8.1).
|
||||
//
|
||||
// The whole point of SourceRef is that nothing past this file needs to know
|
||||
// what a mailbox is: fusion, budgeting, prompt assembly and citation
|
||||
// rendering all operate on SourceRef. Adding another product (VNCtalk, per
|
||||
// the doc's P7) means writing one more RetrieverAdapter, not touching any of
|
||||
// that shared code.
|
||||
|
||||
export type Product = 'mail' | 'talk' | 'files' | 'calendar';
|
||||
|
||||
export interface SourceRef {
|
||||
product: Product;
|
||||
/** Owning account — personal or group. */
|
||||
accountId: string;
|
||||
/** Mailbox · room · drive · calendar. */
|
||||
collectionId: string;
|
||||
/** Email · message · file · event. */
|
||||
itemId: string;
|
||||
/** Which slice of a long item this chunk covers. */
|
||||
chunkIx: number;
|
||||
}
|
||||
|
||||
export interface Scored<T> {
|
||||
ref: T;
|
||||
score: number;
|
||||
}
|
||||
|
||||
export interface Chunk {
|
||||
ref: SourceRef;
|
||||
/** Display-ready text for this chunk, already capped to a safe prompt size. */
|
||||
text: string;
|
||||
/** Human-readable label for citations, e.g. an email subject. */
|
||||
title: string;
|
||||
}
|
||||
|
||||
export interface RetrieverAdapter {
|
||||
product: Product;
|
||||
/** The local FTS leg, if this product has one (mail does, via the
|
||||
* Electron-only encrypted index). Undefined where no local leg exists —
|
||||
* fusion degrades to a single retriever with no code branch needed. */
|
||||
localSearch?: (query: string, limit: number) => Promise<Scored<SourceRef>[]>;
|
||||
/** The server embedding leg. */
|
||||
serverSearch: (query: string, limit: number) => Promise<Scored<SourceRef>[]>;
|
||||
hydrate: (refs: SourceRef[]) => Promise<Chunk[]>;
|
||||
}
|
||||
Reference in New Issue
Block a user