diff --git a/graphile/graphile-llm/src/__tests__/embed-inputs.test.ts b/graphile/graphile-llm/src/__tests__/embed-inputs.test.ts new file mode 100644 index 0000000000..9cebb5bdee --- /dev/null +++ b/graphile/graphile-llm/src/__tests__/embed-inputs.test.ts @@ -0,0 +1,466 @@ +import type { GraphileConfig } from 'graphile-config'; +import { ConnectionFilterPreset } from 'graphile-connection-filter'; +import { createPgvectorAdapter } from 'graphile-search/adapters/pgvector'; +import { createTsvectorAdapter } from 'graphile-search/adapters/tsvector'; +import { TsvectorCodecPlugin } from 'graphile-search/codecs/tsvector-codec'; +import { VectorCodecPlugin } from 'graphile-search/codecs/vector-codec'; +import { createUnifiedSearchPlugin } from 'graphile-search/plugin'; +import type { GraphQLResponse } from 'graphile-test'; +import { getConnections, seed } from 'graphile-test'; +import type { GraphQLSchema } from 'graphql'; +import * as graphql from 'graphql'; +import { parse, validate } from 'graphql'; +import type { PgTestClient } from 'pgsql-test'; + +import { createEmbedIndex, embedField } from '../../src/embed-inputs'; +import { llmConfigStore } from '../../src/embedder'; +import { createLlmModulePlugin } from '../../src/plugins/llm-module-plugin'; +import { createLlmTextMutationPlugin } from '../../src/plugins/text-mutation-plugin'; +import { createLlmTextSearchPlugin } from '../../src/plugins/text-search-plugin'; + +type QueryFn = ( + query: string, + variables?: Record, + commit?: boolean, + reqOptions?: Record +) => Promise>; + +type DocNode = { name: string; embeddingVectorDistance: number | null }; +type DocsResult = { allDocs: { nodes: DocNode[] } }; + +// Deterministic 3-dim embedder so ranking is exact and no provider is needed. +const VECTORS: Record = { + rent: [1, 0, 0], + revenue: [0, 1, 0], + breach: [0, 0, 1] +}; +const vectorFor = (text: string) => VECTORS[text] ?? [1, 1, 1]; + +type EmbedderShape = 'result' | 'array' | 'null' | 'empty'; +let shape: EmbedderShape = 'result'; +const calls: Array<{ text: string; overrides: unknown }> = []; + +// `result` is the provider shape ({ embedding }), `array` is the metering +// wrapper's shape, `null` is the metering wrapper refusing (quota exceeded), +// `empty` is a provider answering with no vector. +function fakeEmbedderPlugin(configured: boolean): GraphileConfig.Plugin { + return { + name: 'FakeEmbedderPlugin', + version: '0.0.0', + after: ['LlmModulePlugin'], + schema: { + hooks: { + build(build) { + (build as any).llmEmbedder = configured + ? async (text: string) => { + calls.push({ text, overrides: llmConfigStore.getStore() }); + if (shape === 'null') return null; + if (shape === 'empty') return []; + const embedding = vectorFor(text); + return shape === 'array' ? embedding : { embedding, promptTokens: 1 }; + } + : null; + return build; + } + } + } + }; +} + +const SETUP_SQL = ` + CREATE EXTENSION IF NOT EXISTS vector; + CREATE SCHEMA llm_embed_inputs; + CREATE TABLE llm_embed_inputs.docs ( + id serial PRIMARY KEY, + name text NOT NULL, + body text NOT NULL DEFAULT '', + tsv tsvector GENERATED ALWAYS AS (to_tsvector('english', name || ' ' || body)) STORED, + embedding vector(3) NOT NULL + ); + INSERT INTO llm_embed_inputs.docs (name, body, embedding) VALUES + ('warehouse-lease.txt', 'warehouse lease and monthly rent', '[1,0,0]'), + ('q3-results.txt', 'quarterly revenue results', '[0,1,0]'), + ('breach-report.txt', 'customer data breach', '[0,0,1]'); +`; + +interface Harness { + db: PgTestClient; + pg: PgTestClient; + query: QueryFn; + schema: GraphQLSchema; + teardown: () => Promise; +} + +async function setup(options: { + onQuotaExceeded?: 'degrade' | 'throw'; + embedderConfigured?: boolean; +}): Promise { + let schema: GraphQLSchema; + const CaptureSchemaPlugin: GraphileConfig.Plugin = { + name: 'CaptureSchemaPlugin', + version: '0.0.0', + schema: { hooks: { finalize(s) { schema = s; return s; } } } + }; + + const connections = await getConnections({ + schemas: ['llm_embed_inputs'], + preset: { + extends: [ConnectionFilterPreset()], + plugins: [ + TsvectorCodecPlugin, + VectorCodecPlugin, + createUnifiedSearchPlugin({ + adapters: [createTsvectorAdapter(), createPgvectorAdapter()], + enableSearchScore: true, + enableUnifiedSearch: true + }), + createLlmModulePlugin(), + fakeEmbedderPlugin(options.embedderConfigured ?? true), + createLlmTextSearchPlugin({ onQuotaExceeded: options.onQuotaExceeded }), + createLlmTextMutationPlugin(), + CaptureSchemaPlugin + ] + }, + useRoot: true, + authRole: 'postgres' + }, [seed.fn(async (ctx) => { await ctx.pg.query(SETUP_SQL); })]); + + return { ...connections, query: connections.query as QueryFn, schema: schema! }; +} + +const BY_TEXT = ` + query($t: String!) { + allDocs(where: { vectorEmbedding: { text: $t, metric: COSINE } }, orderBy: [EMBEDDING_VECTOR_DISTANCE_ASC]) { + nodes { name embeddingVectorDistance } + } + } +`; +const BY_VECTOR = ` + query($v: [Float!]) { + allDocs(where: { vectorEmbedding: { vector: $v, metric: COSINE } }, orderBy: [EMBEDDING_VECTOR_DISTANCE_ASC]) { + nodes { name embeddingVectorDistance } + } + } +`; +const BY_VECTOR_AND_TEXT = ` + query($v: [Float!], $t: String) { + allDocs(where: { vectorEmbedding: { vector: $v, text: $t, metric: COSINE } }, orderBy: [EMBEDDING_VECTOR_DISTANCE_ASC]) { + nodes { name embeddingVectorDistance } + } + } +`; +const UNIFIED = ` + query($t: String!) { + allDocs(where: { unifiedSearch: $t }) { nodes { name embeddingVectorDistance } } + } +`; +const CREATE_FROM_TEXT = ` + mutation($name: String!, $t: String) { + createDoc(input: { doc: { name: $name, embeddingText: $t } }) { doc { rowId embedding } } + } +`; + +const names = (r: GraphQLResponse) => r.data!.allDocs.nodes.map((n) => n.name); + +describe('plan-time embedding (onQuotaExceeded: degrade)', () => { + let h: Harness; + + beforeAll(async () => { h = await setup({}); }, 120_000); + afterAll(async () => { if (h) await h.teardown(); }); + beforeEach(async () => { + shape = 'result'; + calls.length = 0; + await h.pg.beforeEach(); + await h.db.beforeEach(); + }); + afterEach(async () => { + await h.db.afterEach(); + await h.pg.afterEach(); + }); + + it('passes GraphQL validation for text-only search and a companion-only create on a NOT NULL vector column', () => { + const documents = [ + BY_TEXT, + CREATE_FROM_TEXT, + 'mutation { updateDocByRowId(input: { rowId: 1, docPatch: { embeddingText: "rent" } }) { doc { rowId } } }' + ]; + const errors = documents.map((d) => validate(h.schema, parse(d)).map((e) => e.message)); + expect(errors).toEqual([[], [], []]); + }); + + describe.each(['result', 'array'])('embedder returning the %s shape', (embedderShape) => { + beforeEach(() => { shape = embedderShape; }); + + it('ranks a text search by the embedded vector and embeds once', async () => { + const result = await h.query(BY_TEXT, { t: 'revenue' }); + + expect(result.errors).toBeUndefined(); + expect(names(result)[0]).toBe('q3-results.txt'); + expect(result.data!.allDocs.nodes[0].embeddingVectorDistance).toBeCloseTo(0, 6); + expect(calls.map((c) => c.text)).toEqual(['revenue']); + }); + + it('stores the embedding of {col}Text on create', async () => { + const result = await h.query<{ createDoc: { doc: { rowId: number; embedding: number[] } } }>( + CREATE_FROM_TEXT, { name: 'created.txt', t: 'breach' } + ); + + expect(result.errors).toBeUndefined(); + expect(result.data!.createDoc.doc.embedding).toEqual([0, 0, 1]); + expect(calls.map((c) => c.text)).toEqual(['breach']); + }); + }); + + it('stores the embedding of {col}Text on update', async () => { + const result = await h.query<{ updateDocByRowId: { doc: { embedding: number[] } } }>(` + mutation { updateDocByRowId(input: { rowId: 1, docPatch: { embeddingText: "revenue" } }) { doc { embedding } } } + `); + + expect(result.errors).toBeUndefined(); + expect(result.data!.updateDocByRowId.doc.embedding).toEqual([0, 1, 0]); + }); + + it('uses a client-sent vector without calling the embedder', async () => { + const result = await h.query(BY_VECTOR, { v: [0, 0, 1] }); + + expect(result.errors).toBeUndefined(); + expect(names(result)[0]).toBe('breach-report.txt'); + expect(calls).toEqual([]); + }); + + it('uses the vector when both vector and text are sent', async () => { + const result = await h.query(BY_VECTOR_AND_TEXT, { v: [0, 1, 0], t: 'rent' }); + + expect(result.errors).toBeUndefined(); + expect(names(result)[0]).toBe('q3-results.txt'); + expect(calls).toEqual([]); + }); + + it('embeds the text when the vector is empty', async () => { + const result = await h.query(BY_VECTOR_AND_TEXT, { v: [], t: 'rent' }); + + expect(result.errors).toBeUndefined(); + expect(names(result)[0]).toBe('warehouse-lease.txt'); + expect(calls.map((c) => c.text)).toEqual(['rent']); + }); + + it('rejects a VectorNearbyInput with neither vector nor text', async () => { + const result = await h.query( + 'query { allDocs(where: { vectorEmbedding: { metric: COSINE } }) { nodes { name } } }' + ); + + expect(result.errors?.[0]?.message).toMatch(/VectorNearbyInput: provide either `vector` or `text`/); + }); + + it('embeds text nested under and/not', async () => { + const result = await h.query(` + query { + allDocs( + where: { and: [{ vectorEmbedding: { text: "rent", metric: COSINE } }, { not: { rowId: { equalTo: 3 } } }] } + orderBy: [EMBEDDING_VECTOR_DISTANCE_ASC] + ) { nodes { name embeddingVectorDistance } } + } + `); + + expect(result.errors).toBeUndefined(); + expect(names(result)).toEqual(['warehouse-lease.txt', 'q3-results.txt']); + expect(result.data!.allDocs.nodes[0].embeddingVectorDistance).toBeCloseTo(0, 6); + expect(calls.map((c) => c.text)).toEqual(['rent']); + }); + + it('embeds unifiedSearch and VectorNearbyInput text in the same filter', async () => { + const result = await h.query(` + query { + allDocs(where: { unifiedSearch: "warehouse", vectorEmbedding: { text: "rent", metric: COSINE } }) { + nodes { name } + } + } + `); + + expect(result.errors).toBeUndefined(); + expect(names(result)).toEqual(['warehouse-lease.txt']); + expect(calls.map((c) => c.text).sort()).toEqual(['rent', 'warehouse']); + }); + + it('adds pgvector to unifiedSearch', async () => { + const result = await h.query(UNIFIED, { t: 'warehouse' }); + + expect(result.errors).toBeUndefined(); + expect(result.data!.allDocs.nodes).toEqual([ + { name: 'warehouse-lease.txt', embeddingVectorDistance: expect.any(Number) } + ]); + expect(calls.map((c) => c.text)).toEqual(['warehouse']); + }); + + it('does not embed a blank unifiedSearch', async () => { + const result = await h.query(UNIFIED, { t: ' ' }); + + expect(result.errors).toBeUndefined(); + expect(calls).toEqual([]); + }); + + it('drops the vector filter when the embedder returns null', async () => { + shape = 'null'; + const result = await h.query(BY_TEXT, { t: 'rent' }); + + expect(result.errors).toBeUndefined(); + expect(result.data!.allDocs.nodes.length).toBeGreaterThanOrEqual(3); + expect(result.data!.allDocs.nodes.every((n) => n.embeddingVectorDistance === null)).toBe(true); + }); + + it('falls back to text-only unifiedSearch when the embedder returns null', async () => { + shape = 'null'; + const result = await h.query(UNIFIED, { t: 'warehouse' }); + + expect(result.errors).toBeUndefined(); + expect(result.data!.allDocs.nodes).toEqual([ + { name: 'warehouse-lease.txt', embeddingVectorDistance: null } + ]); + }); + + it('fails a create when the embedder returns null, and writes nothing', async () => { + shape = 'null'; + const result = await h.query(CREATE_FROM_TEXT, { name: 'quota.txt', t: 'rent' }); + + expect(result.errors?.[0]?.message).toMatch(/EMBED_QUOTA_EXCEEDED: Cannot embed embeddingText/); + const rows = await h.pg.query("SELECT 1 FROM llm_embed_inputs.docs WHERE name = 'quota.txt'"); + expect(rows.rows).toHaveLength(0); + }); + + it('rejects a create that has neither the vector nor its text companion', async () => { + const result = await h.query(CREATE_FROM_TEXT, { name: 'nothing.txt' }); + + expect(result.errors?.[0]?.message).toMatch(/DocInput: provide either `embedding` or `embeddingText`/); + expect(calls).toEqual([]); + }); + + it('embeds with the per-database model for both search and create', async () => { + const reqOptions = { + constructive: { + useLlm: async () => ({ embeddingModel: 'tenant-embed', embeddingBaseUrl: 'http://tenant-ollama:11434' }) + } + }; + const search = await h.query(BY_TEXT, { t: 'rent' }, false, reqOptions); + const create = await h.query(CREATE_FROM_TEXT, { name: 'tenant.txt', t: 'rent' }, false, reqOptions); + + expect(search.errors).toBeUndefined(); + expect(create.errors).toBeUndefined(); + expect(calls).toHaveLength(2); + for (const call of calls) { + expect(call.overrides).toMatchObject({ + embeddingModel: 'tenant-embed', + embeddingBaseUrl: 'http://tenant-ollama:11434' + }); + } + }); + + it('does not look up per-database config when nothing needs embedding', async () => { + const useLlm = jest.fn(async () => ({ embeddingModel: 'tenant-embed' })); + const result = await h.query(BY_VECTOR, { v: [1, 0, 0] }, false, { constructive: { useLlm } }); + + expect(result.errors).toBeUndefined(); + expect(useLlm).not.toHaveBeenCalled(); + }); +}); + +describe('plan-time embedding (onQuotaExceeded: throw)', () => { + let h: Harness; + + beforeAll(async () => { h = await setup({ onQuotaExceeded: 'throw' }); }, 120_000); + afterAll(async () => { if (h) await h.teardown(); }); + beforeEach(async () => { + shape = 'null'; + calls.length = 0; + await h.pg.beforeEach(); + await h.db.beforeEach(); + }); + afterEach(async () => { + await h.db.afterEach(); + await h.pg.afterEach(); + }); + + it('fails a text search when the embedder returns null', async () => { + const result = await h.query(BY_TEXT, { t: 'rent' }); + + expect(result.errors?.[0]?.message).toMatch(/VectorNearbyInput: embedding failed/); + }); + + it('fails a text search when the embedder returns an empty vector', async () => { + shape = 'empty'; + const result = await h.query(BY_TEXT, { t: 'rent' }); + + expect(result.errors?.[0]?.message).toMatch(/VectorNearbyInput: embedding failed/); + }); + + it('fails unifiedSearch when the embedder returns null', async () => { + const result = await h.query(UNIFIED, { t: 'warehouse' }); + + expect(result.errors?.[0]?.message).toMatch(/unifiedSearch: embedding failed.*onQuotaExceeded is set to 'throw'/); + }); +}); + +describe('plan-time embedding without an embedder', () => { + let h: Harness; + + beforeAll(async () => { h = await setup({ embedderConfigured: false }); }, 120_000); + afterAll(async () => { if (h) await h.teardown(); }); + beforeEach(async () => { + await h.pg.beforeEach(); + await h.db.beforeEach(); + }); + afterEach(async () => { + await h.db.afterEach(); + await h.pg.afterEach(); + }); + + it('ignores search text and runs unifiedSearch on the text adapters', async () => { + const byText = await h.query(BY_TEXT, { t: 'rent' }); + const unified = await h.query(UNIFIED, { t: 'warehouse' }); + + expect(byText.errors).toBeUndefined(); + expect(byText.data!.allDocs.nodes.every((n) => n.embeddingVectorDistance === null)).toBe(true); + expect(unified.errors).toBeUndefined(); + expect(unified.data!.allDocs.nodes).toEqual([ + { name: 'warehouse-lease.txt', embeddingVectorDistance: null } + ]); + }); + + it('fails a create from text', async () => { + const result = await h.query(CREATE_FROM_TEXT, { name: 'unconfigured.txt', t: 'rent' }); + + expect(result.errors?.[0]?.message).toMatch(/EMBED_NOT_CONFIGURED: Cannot embed embeddingText/); + }); +}); + +describe('embeddable-input reachability', () => { + // Relation filters are cyclic: a filter type reaches another that reaches it + // back, and the tagged field may sit behind either one. + function cyclicFilters() { + const { GraphQLInputObjectType, GraphQLString } = graphql; + const tagged = new GraphQLInputObjectType({ + name: 'TaggedFilter', + fields: { text: embedField({ kind: 'unified', onQuotaExceeded: 'degrade' }, { type: GraphQLString }) } + }); + const parent: any = new GraphQLInputObjectType({ + name: 'ParentFilter', + fields: () => ({ child: { type: child }, tagged: { type: tagged } }) + }); + const child: any = new GraphQLInputObjectType({ + name: 'ChildFilter', + fields: () => ({ parent: { type: parent } }) + }); + return { parent, child }; + } + + it('finds the tagged field from every type in the cycle, whichever is asked first', () => { + const parentFirst = createEmbedIndex(graphql); + const a = cyclicFilters(); + expect(parentFirst.reaches(a.parent)).toBe(true); + expect(parentFirst.reaches(a.child)).toBe(true); + + const childFirst = createEmbedIndex(graphql); + const b = cyclicFilters(); + expect(childFirst.reaches(b.child)).toBe(true); + expect(childFirst.reaches(b.parent)).toBe(true); + }); +}); diff --git a/graphile/graphile-llm/src/__tests__/graphile-llm.test.ts b/graphile/graphile-llm/src/__tests__/graphile-llm.test.ts index 30f2b58439..effb8b4a0f 100644 --- a/graphile/graphile-llm/src/__tests__/graphile-llm.test.ts +++ b/graphile/graphile-llm/src/__tests__/graphile-llm.test.ts @@ -937,7 +937,7 @@ describe('GraphileLlmPreset toggles', () => { // // End-to-end: embeds article text with real Ollama nomic-embed-text (768 dims), // seeds the DB, then runs hybrid text+vector queries verifying RRF fusion, -// semantic relevance, and the embedTextInWhere pipeline. +// semantic relevance, and server-side embedding of unifiedSearch text. // // Requires: PostgreSQL with pgvector + Ollama with nomic-embed-text. // ============================================================================= @@ -1101,9 +1101,9 @@ describe('Real Ollama + unifiedSearch + pgvector RRF integration', () => { } }, 30_000); - // ─── unifiedSearch text-only (no vector injection) ───────────────────── + // ─── unifiedSearch text embedded server-side ─────────────────────────── - it('unifiedSearch text-only uses tsvector without pgvector', async () => { + it('unifiedSearch text embeds server-side so tsvector and pgvector both participate', async () => { const result = await query<{ allArticles: { nodes: Array<{ rowId: number; title: string; tsvRank: number | null; embeddingVectorDistance: number | null; searchScore: number }> }; }>(` @@ -1118,13 +1118,9 @@ describe('Real Ollama + unifiedSearch + pgvector RRF integration', () => { const nodes = result.data?.allArticles?.nodes ?? []; expect(nodes.length).toBeGreaterThan(0); - // tsvector should be active - const hasTsv = nodes.some((n) => n.tsvRank !== null); - expect(hasTsv).toBe(true); - - // pgvector should NOT participate (no vector injection in graphile-test) + expect(nodes.some((n) => n.tsvRank !== null)).toBe(true); for (const node of nodes) { - expect(node.embeddingVectorDistance).toBeNull(); + expect(typeof node.embeddingVectorDistance).toBe('number'); } }); diff --git a/graphile/graphile-llm/src/__tests__/plan-time-embedding.test.ts b/graphile/graphile-llm/src/__tests__/plan-time-embedding.test.ts new file mode 100644 index 0000000000..716c7ed4c0 --- /dev/null +++ b/graphile/graphile-llm/src/__tests__/plan-time-embedding.test.ts @@ -0,0 +1,203 @@ +import OllamaClient from '@agentic-kit/ollama'; +import { ConnectionFilterPreset } from 'graphile-connection-filter'; +import { createPgvectorAdapter } from 'graphile-search/adapters/pgvector'; +import { createTsvectorAdapter } from 'graphile-search/adapters/tsvector'; +import { TsvectorCodecPlugin } from 'graphile-search/codecs/tsvector-codec'; +import { VectorCodecPlugin } from 'graphile-search/codecs/vector-codec'; +import { createUnifiedSearchPlugin } from 'graphile-search/plugin'; +import type { GraphQLResponse } from 'graphile-test'; +import { getConnections, seed } from 'graphile-test'; +import type { PgTestClient } from 'pgsql-test'; + +import { buildEmbedder } from '../../src/embedder'; +import { createLlmModulePlugin } from '../../src/plugins/llm-module-plugin'; +import { createLlmTextMutationPlugin } from '../../src/plugins/text-mutation-plugin'; +import { createLlmTextSearchPlugin } from '../../src/plugins/text-search-plugin'; + +// Text-to-vector must happen while the operation is planned, before SQL is +// built: a v4-style resolve() wrapper runs after Grafast has already executed +// the plan, so it can neither filter nor write the vector. + +type QueryFn = ( + query: string, + variables?: Record +) => Promise>; + +type DocNode = { name: string; embeddingVectorDistance: number | null }; +type DocsResult = { allDocs: { nodes: DocNode[] } }; + +const OLLAMA_URL = 'http://localhost:11434'; + +const DOCS: Array<[string, string]> = [ + ['warehouse-lease.txt', 'Warehouse lease agreement: monthly rent of 12,000 USD for the Oakland storage facility, five year term.'], + ['q3-financial-results.txt', 'Third quarter revenue grew 14 percent while operating margin improved to 22 percent.'], + ['patent-portfolio.txt', 'The company holds 37 granted patents covering battery chemistry and charging circuits.'], + ['data-breach-incident.txt', 'An attacker accessed customer email addresses through a misconfigured backup server.'] +]; + +const ASK = 'how much rent do we pay for the warehouse'; + +const BY_VECTOR = ` + query($v: [Float!]!) { + allDocs(where: { vectorEmbedding: { vector: $v, metric: COSINE } }, orderBy: [EMBEDDING_VECTOR_DISTANCE_ASC]) { + nodes { name embeddingVectorDistance } + } + } +`; + +const BY_TEXT = ` + query($t: String!) { + allDocs(where: { vectorEmbedding: { text: $t, metric: COSINE } }, orderBy: [EMBEDDING_VECTOR_DISTANCE_ASC]) { + nodes { name embeddingVectorDistance } + } + } +`; + +describe('server-side embedding runs before SQL', () => { + let db: PgTestClient; + let pg: PgTestClient; + let teardown: () => Promise; + let query: QueryFn; + let askVector: number[]; + + beforeAll(async () => { + const ollama = new OllamaClient(OLLAMA_URL); + const models = await ollama.listModels(); + if (!models.some((m: string) => m.includes('nomic-embed-text'))) { + await ollama.pullModel('nomic-embed-text'); + } + + const embed = buildEmbedder({ provider: 'ollama', model: 'nomic-embed-text', baseUrl: OLLAMA_URL })!; + const docVectors = await Promise.all(DOCS.map(([, body]) => embed(body))); + askVector = (await embed(ASK)).embedding; + + const values = DOCS.map(([name, body], i) => + `('${name}', '${body.replace(/'/g, "''")}', '[${docVectors[i].embedding.join(',')}]')` + ).join(',\n'); + + const connections = await getConnections({ + schemas: ['llm_plan_embed'], + preset: { + extends: [ConnectionFilterPreset()], + plugins: [ + TsvectorCodecPlugin, + VectorCodecPlugin, + createUnifiedSearchPlugin({ + adapters: [createTsvectorAdapter(), createPgvectorAdapter()], + enableSearchScore: true, + enableUnifiedSearch: true + }), + createLlmModulePlugin({ + defaultEmbedder: { provider: 'ollama', model: 'nomic-embed-text', baseUrl: OLLAMA_URL } + }), + createLlmTextSearchPlugin(), + createLlmTextMutationPlugin() + ] + }, + useRoot: true, + authRole: 'postgres' + }, [ + seed.fn(async (ctx) => { + await ctx.pg.query(` + CREATE EXTENSION IF NOT EXISTS vector; + CREATE SCHEMA llm_plan_embed; + CREATE TABLE llm_plan_embed.docs ( + id serial PRIMARY KEY, + name text NOT NULL, + body text NOT NULL, + tsv tsvector GENERATED ALWAYS AS (to_tsvector('english', name || ' ' || body)) STORED, + embedding vector(768) + ); + INSERT INTO llm_plan_embed.docs (name, body, embedding) VALUES ${values}; + `); + }) + ]); + + db = connections.db; + pg = connections.pg; + teardown = connections.teardown; + query = connections.query; + }, 120_000); + + afterAll(async () => { + if (teardown) await teardown(); + }); + + beforeEach(async () => { await db.beforeEach(); }); + afterEach(async () => { await db.afterEach(); }); + + it('lets clients send text without a vector', async () => { + const result = await query<{ + __type: { inputFields: Array<{ name: string; type: { kind: string } }> }; + }>(`{ __type(name: "VectorNearbyInput") { inputFields { name type { kind } } } }`); + + expect(result.errors).toBeUndefined(); + const vectorField = result.data!.__type.inputFields.find((f) => f.name === 'vector'); + expect(vectorField!.type.kind).not.toBe('NON_NULL'); + }); + + it('ranks a text search exactly like the same text sent as a vector', async () => { + const byVector = await query(BY_VECTOR, { v: askVector }); + const byText = await query(BY_TEXT, { t: ASK }); + + expect(byVector.errors).toBeUndefined(); + expect(byText.errors).toBeUndefined(); + const expected = byVector.data!.allDocs.nodes; + const actual = byText.data!.allDocs.nodes; + expect(expected[0].name).toBe('warehouse-lease.txt'); + expect(actual.map((n) => n.name)).toEqual(expected.map((n) => n.name)); + actual.forEach((node, i) => { + expect(node.embeddingVectorDistance).toBeCloseTo(expected[i].embeddingVectorDistance!, 6); + }); + }); + + it('includes vector distance in unifiedSearch', async () => { + // unifiedSearch ORs each adapter's WHERE; pgvector adds none without a + // distance threshold, so rows are still selected by the tsvector match. + const result = await query(` + query($t: String!) { + allDocs(where: { unifiedSearch: $t }) { + nodes { name embeddingVectorDistance } + } + } + `, { t: 'warehouse rent' }); + + expect(result.errors).toBeUndefined(); + const lease = result.data!.allDocs.nodes.find((n) => n.name === 'warehouse-lease.txt'); + expect(lease).toBeDefined(); + expect(typeof lease!.embeddingVectorDistance).toBe('number'); + }); + + it('stores the embedding of {col}Text on create', async () => { + const created = await query<{ createDoc: { doc: { rowId: number } } }>(` + mutation($t: String!) { + createDoc(input: { doc: { name: "ask.txt", body: "ask", embeddingText: $t } }) { doc { rowId } } + } + `, { t: ASK }); + + expect(created.errors).toBeUndefined(); + const rowId = created.data!.createDoc.doc.rowId; + const stored = await pg.query<{ distance: number | null }>( + 'SELECT embedding <=> $1::vector AS distance FROM llm_plan_embed.docs WHERE id = $2', + [`[${askVector.join(',')}]`, rowId] + ); + expect(stored.rows[0].distance).not.toBeNull(); + expect(stored.rows[0].distance).toBeLessThan(1e-6); + }); + + it('stores the embedding of {col}Text on update', async () => { + const row = await pg.query<{ id: number }>("SELECT id FROM llm_plan_embed.docs WHERE name = 'patent-portfolio.txt'"); + const updated = await query(` + mutation($id: Int!, $t: String!) { + updateDocByRowId(input: { rowId: $id, docPatch: { embeddingText: $t } }) { doc { rowId } } + } + `, { id: row.rows[0].id, t: ASK }); + + expect(updated.errors).toBeUndefined(); + const stored = await pg.query<{ distance: number }>( + 'SELECT embedding <=> $1::vector AS distance FROM llm_plan_embed.docs WHERE id = $2', + [`[${askVector.join(',')}]`, row.rows[0].id] + ); + expect(stored.rows[0].distance).toBeLessThan(1e-6); + }); +}); diff --git a/graphile/graphile-llm/src/embed-inputs.ts b/graphile/graphile-llm/src/embed-inputs.ts new file mode 100644 index 0000000000..96fa77e0d5 --- /dev/null +++ b/graphile/graphile-llm/src/embed-inputs.ts @@ -0,0 +1,330 @@ +/** + * Plan-time text-to-vector embedding. + * + * Under Grafast a field's `resolve()` runs after its plan has executed, so the + * SQL has already been built by the time a resolver wrapper could embed text. + * Instead, input fields that carry text to embed are tagged with an + * `EmbedRule` where they are declared, and `planTimeEmbedding` (run from a + * `finalize` hook) wraps the plan of every object field whose arguments can + * reach a tagged input field. The wrapper shadows `fieldArgs.getRaw` so every + * consumer of those arguments (connection-filter's `where` applyPlan, the + * create/update `args.apply`) reads a step that embeds the text first. + */ + +import type { ExecutableStep, FieldArgs } from 'grafast'; +import { access, ConstantStep, context, lambda, object } from 'grafast'; +import type { GraphQLInputType, GraphQLSchema } from 'graphql'; + +import type { LlmConfigOverrides } from './embedder'; +import { llmConfigStore } from './embedder'; + +export type QuotaPolicy = 'degrade' | 'throw'; + +/** Stored on a GraphQL input field's `extensions[EMBED_EXT]`. */ +export type EmbedRule = + /** `VectorNearbyInput.text`: embedded into the sibling `into` field. */ + | { kind: 'nearby'; into: string; onQuotaExceeded: QuotaPolicy } + /** `unifiedSearch: String`: becomes `{ __text, __vector }` for graphile-search. */ + | { kind: 'unified'; onQuotaExceeded: QuotaPolicy } + /** `{col}Text` on a create input or patch. `required` means the column is NOT NULL. */ + | { kind: 'companion'; into: string; required: boolean }; + +export const EMBED_EXT = 'graphileLlmEmbed'; + +export type VectorEmbedder = (text: string) => Promise; + +export function embedField(rule: EmbedRule, spec: T): T { + return { ...spec, extensions: { ...(spec as { extensions?: object }).extensions, [EMBED_EXT]: rule } }; +} + +/** + * `build.llmEmbedder` resolves to `{ embedding }` from the providers and to + * `number[] | null` once the metering plugin has wrapped it. + */ +export function toVectorEmbedder(raw: unknown): VectorEmbedder | null { + if (typeof raw !== 'function') return null; + return async (text) => { + const result = await raw(text); + const vector = Array.isArray(result) ? result : result?.embedding; + return Array.isArray(vector) && vector.length > 0 ? vector : null; + }; +} + +function isNonEmptyVector(value: unknown): boolean { + return Array.isArray(value) && value.length > 0; +} + +export async function resolveConfigOverrides( + graphqlContext: any +): Promise { + const ctx = graphqlContext?.constructive; + if (!ctx?.useLlm) return null; + + try { + const llm = await ctx.useLlm(); + if (!llm?.embeddingModel) return null; + + return { + embeddingModel: llm.embeddingModel, + embeddingBaseUrl: llm.embeddingBaseUrl, + chatModel: llm.chatModel ?? undefined, + chatBaseUrl: llm.chatBaseUrl ?? undefined + }; + } catch { + return null; + } +} + +/** Binds the embedder to one request's per-database model overrides, resolved on first use. */ +function forRequest(embedder: VectorEmbedder | null, graphqlContext: unknown): VectorEmbedder | null { + if (!embedder) return null; + let overrides: Promise | undefined; + return async (text) => { + overrides ??= resolveConfigOverrides(graphqlContext); + const resolved = await overrides; + const startTime = Date.now(); + const vector = await llmConfigStore.run(resolved, () => embedder(text)); + if (vector) { + console.log(`[graphile-llm] embed: dims=${vector.length}, latency=${Date.now() - startTime}ms`); + } + return vector; + }; +} + +export function createEmbedIndex(graphql: any) { + const { getNamedType, isInputObjectType, isListType, isNonNullType } = graphql; + const reachesMemo = new WeakMap(); + + const inputObjectFields = (type: any): any[] => Object.values(type.getFields()); + const namedInputObject = (field: any): any => { + const named = getNamedType(field.type); + return isInputObjectType(named) ? named : null; + }; + + // Filter types are cyclic (AND/OR/NOT, forward and backward relations), so + // reachability is a fixed point over everything reachable from `type`; a + // depth-first walk that memoizes a provisional answer mid-cycle gets it wrong. + function reaches(type: GraphQLInputType): boolean { + const start = getNamedType(type); + if (!isInputObjectType(start)) return false; + const known = reachesMemo.get(start); + if (known !== undefined) return known; + + const pending: any[] = []; + const stack: any[] = [start]; + const seen = new Set(); + while (stack.length > 0) { + const node = stack.pop(); + if (seen.has(node) || reachesMemo.has(node)) continue; + seen.add(node); + pending.push(node); + for (const field of inputObjectFields(node)) { + const child = namedInputObject(field); + if (child) stack.push(child); + } + } + + const hit = new Set( + pending.filter((node) => inputObjectFields(node).some((f) => f.extensions?.[EMBED_EXT])) + ); + const isHit = (node: any) => hit.has(node) || reachesMemo.get(node) === true; + let changed = true; + while (changed) { + changed = false; + for (const node of pending) { + if (hit.has(node)) continue; + if (inputObjectFields(node).some((f) => { + const child = namedInputObject(f); + return child !== null && isHit(child); + })) { + hit.add(node); + changed = true; + } + } + } + + for (const node of pending) reachesMemo.set(node, hit.has(node)); + return reachesMemo.get(start); + } + + async function applyRule( + rule: EmbedRule, + fieldName: string, + typeName: string, + value: Record, + next: Record, + embed: VectorEmbedder | null + ): Promise { + const text = value[fieldName]; + + switch (rule.kind) { + case 'nearby': { + delete next[fieldName]; + if (isNonEmptyVector(value[rule.into])) return; + if (typeof text !== 'string') { + throw new Error(`${typeName}: provide either \`${rule.into}\` or \`${fieldName}\`.`); + } + if (!embed) return; + const vector = await embed(text); + if (vector) { + next[rule.into] = vector; + } else if (rule.onQuotaExceeded === 'throw') { + throw new Error( + `${typeName}: embedding failed (quota exceeded or provider unavailable). ` + + 'Upgrade your plan or adjust onQuotaExceeded to \'degrade\' for graceful fallback.' + ); + } + return; + } + + case 'unified': { + if (typeof text !== 'string' || text.trim().length === 0 || !embed) return; + const vector = await embed(text); + if (vector) { + next[fieldName] = { __text: text, __vector: vector }; + } else if (rule.onQuotaExceeded === 'throw') { + throw new Error( + `${fieldName}: embedding failed (quota exceeded or provider unavailable). ` + + 'onQuotaExceeded is set to \'throw\'. ' + + 'Upgrade your plan or adjust onQuotaExceeded to \'degrade\' for text-only fallback.' + ); + } + return; + } + + case 'companion': { + delete next[fieldName]; + if (isNonEmptyVector(value[rule.into])) return; + if (typeof text !== 'string') { + if (rule.required) { + throw new Error(`${typeName}: provide either \`${rule.into}\` or \`${fieldName}\`.`); + } + return; + } + if (!embed) { + throw new Error( + `EMBED_NOT_CONFIGURED: Cannot embed ${fieldName} — no embedding provider configured. ` + + 'Set defaultEmbedder in GraphileLlmPreset options or EMBEDDER_PROVIDER env var.' + ); + } + const vector = await embed(text); + if (!vector) { + throw new Error( + `EMBED_QUOTA_EXCEEDED: Cannot embed ${fieldName} — embedding quota exceeded. ` + + 'Upgrade your plan or wait for the next billing period.' + ); + } + next[rule.into] = vector; + return; + } + } + } + + /** Returns `value` itself when its type cannot reach a tagged field. */ + async function embedValue(type: any, value: any, embed: VectorEmbedder | null): Promise { + if (value == null) return value; + const inner = isNonNullType(type) ? type.ofType : type; + if (isListType(inner)) { + if (!Array.isArray(value) || !reaches(inner.ofType)) return value; + return Promise.all(value.map((item: any) => embedValue(inner.ofType, item, embed))); + } + if (!isInputObjectType(inner) || !reaches(inner)) return value; + + const next: Record = { ...value }; + const jobs: Promise[] = []; + for (const [name, field] of Object.entries(inner.getFields())) { + const rule = field.extensions?.[EMBED_EXT] as EmbedRule | undefined; + if (rule) { + jobs.push(applyRule(rule, name, inner.name, value, next, embed)); + } else if (value[name] != null && reaches(field.type)) { + jobs.push(embedValue(field.type, value[name], embed).then((embedded) => { next[name] = embedded; })); + } + } + await Promise.all(jobs); + return next; + } + + return { reaches, embedValue }; +} + +type EmbedIndex = ReturnType; +type FieldPlan = (parent: ExecutableStep, fieldArgs: FieldArgs, info: unknown) => unknown; + +function withEmbeddedArgs( + originalPlan: FieldPlan, + argNames: string[], + embeddable: Set, + index: EmbedIndex, + embedder: VectorEmbedder | null +): FieldPlan { + return function planWithEmbeddedArgs(parent, fieldArgs, info) { + const originalGetRaw = fieldArgs.getRaw as (path?: string | ReadonlyArray) => ExecutableStep; + const roots = new Map(); + + const rootStep = (argName: string): ExecutableStep => { + let $root = roots.get(argName); + if ($root) return $root; + const $raw = originalGetRaw.call(fieldArgs, argName); + const absent = $raw instanceof ConstantStep && $raw.data === undefined; + if (absent || !embeddable.has(argName)) { + $root = $raw; + } else { + const type = fieldArgs.typeAt(argName); + $root = $raw.withLayerPlan(() => + lambda({ raw: $raw, ctx: context() } as any, ({ raw, ctx }: any) => + index.embedValue(type, raw, forRequest(embedder, ctx)) + ) + ); + } + roots.set(argName, $root); + return $root; + }; + + // grafast reads `fieldArgs.getRaw` through the object (apply, processAfter, + // `$arg` getters, getBaked), so shadowing it here reaches every consumer. + fieldArgs.getRaw = function getRawEmbedded(path?: string | ReadonlyArray) { + if (path === undefined) { + const steps = Object.fromEntries(argNames.map((n) => [n, rootStep(n)])); + return steps[argNames[0]].withLayerPlan(() => object(steps)) as any; + } + const [argName, ...rest] = typeof path === 'string' ? [path] : path; + if (!embeddable.has(argName as string)) return originalGetRaw.call(fieldArgs, path); + const $root = rootStep(argName as string); + return (rest.length === 0 ? $root : $root.withLayerPlan(() => access($root, rest))) as any; + } as typeof fieldArgs.getRaw; + + return originalPlan(parent, fieldArgs, info); + }; +} + +const processedSchemas = new WeakSet(); + +/** + * Wraps the plan of every object field whose arguments reach a tagged input + * field. Runs in `finalize` because args are attached after the + * `GraphQLObjectType_fields_field` hook. Idempotent per schema, so every plugin + * that tags fields can call it without depending on the others being loaded. + */ +export function planTimeEmbedding(schema: GraphQLSchema, build: any): GraphQLSchema { + if (processedSchemas.has(schema)) return schema; + processedSchemas.add(schema); + + const { graphql } = build; + const index = createEmbedIndex(graphql); + const embedder = toVectorEmbedder(build.llmEmbedder); + + for (const type of Object.values(schema.getTypeMap())) { + if (!graphql.isObjectType(type) || type.name.startsWith('__')) continue; + for (const field of Object.values(type.getFields())) { + const grafastExt = field.extensions?.grafast; + if (typeof grafastExt?.plan !== 'function') continue; + const embeddable = new Set( + field.args.filter((a: any) => index.reaches(a.type)).map((a: any) => a.name) + ); + if (embeddable.size === 0) continue; + const argNames = field.args.map((a: any) => a.name); + grafastExt.plan = withEmbeddedArgs(grafastExt.plan, argNames, embeddable, index, embedder); + } + } + return schema; +} diff --git a/graphile/graphile-llm/src/embedder.ts b/graphile/graphile-llm/src/embedder.ts index a97dd2e587..13612b1a76 100644 --- a/graphile/graphile-llm/src/embedder.ts +++ b/graphile/graphile-llm/src/embedder.ts @@ -5,7 +5,7 @@ * embeddings. Currently supports Ollama via @agentic-kit/ollama. * * Per-request model/baseUrl overrides are read from AsyncLocalStorage - * (`llmConfigStore`). The text-search-plugin sets these overrides from + * (`llmConfigStore`). embed-inputs.ts sets these overrides from * the database's `llm_module` config via `ctx.useLlm()`. This means * the same embedder function — and its metering wrapper — handle every * request, regardless of which tenant's model config applies. @@ -34,7 +34,7 @@ export interface LlmConfigOverrides { /** * AsyncLocalStorage for per-request LLM config overrides. * - * Set by the text-search-plugin resolver wrapper (from ctx.useLlm()), + * Set by the plan-time embedding step in embed-inputs.ts (from ctx.useLlm()), * read by the embedder at call time. The metering wrapper sits between * them and is completely unaware — it just wraps the same function. */ diff --git a/graphile/graphile-llm/src/index.ts b/graphile/graphile-llm/src/index.ts index 47e5d36bad..efefa7aed6 100644 --- a/graphile/graphile-llm/src/index.ts +++ b/graphile/graphile-llm/src/index.ts @@ -43,6 +43,10 @@ export { createLlmRagPlugin } from './plugins/rag-plugin'; export { createLlmTextMutationPlugin } from './plugins/text-mutation-plugin'; export { createLlmTextSearchPlugin, embedTextInWhere } from './plugins/text-search-plugin'; +// Plan-time embedding of tagged input fields +export type { EmbedRule } from './embed-inputs'; +export { embedField, planTimeEmbedding } from './embed-inputs'; + // Metering plugin (opt-in billing integration) export { createLlmMeteringPlugin } from './plugins/metering-plugin'; diff --git a/graphile/graphile-llm/src/plugins/llm-module-plugin.ts b/graphile/graphile-llm/src/plugins/llm-module-plugin.ts index b50c2925c0..cbec7c1b39 100644 --- a/graphile/graphile-llm/src/plugins/llm-module-plugin.ts +++ b/graphile/graphile-llm/src/plugins/llm-module-plugin.ts @@ -17,7 +17,7 @@ * * Per-database model/baseUrl overrides (from `llm_module` via `ctx.useLlm()`) * are applied at request time via `llmConfigStore` (AsyncLocalStorage) in - * the text-search-plugin resolver wrapper. The same embedder function — and + * the plan-time embedding step (embed-inputs.ts). The same embedder function — and * its metering wrapper — handles every request; only the model name and * base URL parameters change per-tenant. * diff --git a/graphile/graphile-llm/src/plugins/text-mutation-plugin.ts b/graphile/graphile-llm/src/plugins/text-mutation-plugin.ts index 3051a20475..451a535594 100644 --- a/graphile/graphile-llm/src/plugins/text-mutation-plugin.ts +++ b/graphile/graphile-llm/src/plugins/text-mutation-plugin.ts @@ -17,9 +17,10 @@ * filter/query-side text-to-vector). Together they let clients work entirely * with text/prompts instead of raw float vectors. * - * Runtime embedding uses the v4-style resolver wrapping approach (same as - * graphile-upload-plugin and graphile-bucket-provisioner-plugin). grafserv v5 - * supports this through its backwards-compatibility layer. + * The text is embedded while the mutation is planned (see embed-inputs.ts), + * before the INSERT/UPDATE is built. On a create input for a NOT NULL vector + * column the vector field becomes optional, and one of `{col}` or `{col}Text` + * is required at execution time. If both are given, the vector is used. * * The companion fields are only added when the LLM plugin is loaded. * If no embedder is configured, the fields are still registered for schema @@ -31,6 +32,7 @@ import 'graphile-build-pg'; import type { GraphileConfig } from 'graphile-config'; +import { embedField, planTimeEmbedding } from '../embed-inputs'; // ─── TypeScript Augmentation ──────────────────────────────────────────────── @@ -49,56 +51,17 @@ function isVectorCodec(codec: any): boolean { return codec?.name === 'vector'; } -/** - * Collect the text→vector field mapping for a codec's vector columns. - * Returns a map of textFieldName → vectorFieldName. - */ -function getTextToVectorMapping( - pgCodec: any, - build: any -): Record { - const mapping: Record = {}; - if (!pgCodec?.attributes) return mapping; - - const attributeFieldNames = new Set( - Object.keys(pgCodec.attributes as Record).map( - (attributeName) => - build.inflection.attribute({ codec: pgCodec, attributeName }) - ) - ); - - for (const [attributeName, attribute] of Object.entries( - pgCodec.attributes as Record - )) { - if (isVectorCodec(attribute.codec)) { - const fieldName = build.inflection.attribute({ - codec: pgCodec, - attributeName - }); - const textFieldName = `${fieldName}Text`; - // Skip when a physical column already inflects to the companion name - // (e.g. `embedding_text` next to `embedding`) — the physical field owns it. - if (attributeFieldNames.has(textFieldName)) continue; - mapping[textFieldName] = fieldName; - } - } - return mapping; -} - /** * Creates the LlmTextMutationPlugin. * * Hooks into GraphQLInputObjectType_fields for create/update input types - * and adds `{columnName}Text: String` for each vector column. - * - * Also wraps mutation resolvers via GraphQLObjectType_fields_field to - * intercept `*Text` companion field values, embed them, and inject the - * resulting vectors before the mutation executes. + * and adds `{columnName}Text: String` for each vector column, tagged for + * plan-time embedding. */ export function createLlmTextMutationPlugin(): GraphileConfig.Plugin { return { name: 'LlmTextMutationPlugin', - version: '0.2.0', + version: '0.3.0', description: 'Adds text companion fields on mutation inputs for vector columns — ' + 'text is embedded server-side before storing', @@ -146,7 +109,7 @@ export function createLlmTextMutationPlugin(): GraphileConfig.Plugin { } const { - graphql: { GraphQLString } + graphql: { GraphQLString, getNullableType, isNonNullType } } = build; // Find vector columns on this table @@ -180,16 +143,28 @@ export function createLlmTextMutationPlugin(): GraphileConfig.Plugin { continue; } + const vectorField = newFields[fieldName]; + const required = !!vectorField && isNonNullType(vectorField.type); + if (required) { + newFields = { + ...newFields, + [fieldName]: { ...vectorField, type: getNullableType(vectorField.type) as any } + }; + } + newFields = build.extend( newFields, { - [textFieldName]: { - type: GraphQLString, - description: - `Natural language text to embed server-side into the \`${fieldName}\` vector column. ` + - `Mutually exclusive with \`${fieldName}\` — provide one or the other. ` + - 'Requires the LLM plugin to be configured with an embedding provider.' - } + [textFieldName]: embedField( + { kind: 'companion', into: fieldName, required }, + { + type: GraphQLString, + description: + `Natural language text to embed server-side into the \`${fieldName}\` vector column. ` + + `Provide \`${fieldName}\` or \`${textFieldName}\`; if both are given, \`${fieldName}\` is used. ` + + 'Requires the LLM plugin to be configured with an embedding provider.' + } + ) }, `LlmTextMutationPlugin adding ${textFieldName} companion field for vector column '${columnName}'` ); @@ -198,109 +173,8 @@ export function createLlmTextMutationPlugin(): GraphileConfig.Plugin { return newFields; }, - /** - * Wrap create/update mutation resolvers to intercept `*Text` companion - * field values, embed them using the configured embedder, and inject - * the resulting vector into the corresponding vector field. - * - * Uses the same v4-style resolver wrapping pattern as graphile-upload-plugin - * and graphile-bucket-provisioner-plugin. grafserv v5 supports this through - * its backwards-compatibility layer. - * - * If the embedder returns null (e.g. quota exceeded), throws an error. - */ - GraphQLObjectType_fields_field(field, build, context) { - const { - scope: { isRootMutation, fieldName, pgCodec } - } = context as any; - - // Only wrap root mutation fields on tables with attributes - if (!isRootMutation || !pgCodec || !pgCodec.attributes) { - return field; - } - - // Only wrap create/update mutations - const isCreate = fieldName.startsWith('create'); - const isUpdate = fieldName.startsWith('update'); - if (!isCreate && !isUpdate) { - return field; - } - - // Build the text→vector mapping for this codec - const textToVectorMap = getTextToVectorMapping(pgCodec, build); - if (Object.keys(textToVectorMap).length === 0) { - return field; - } - - const embedder = (build as any).llmEmbedder as - | ((text: string) => Promise) - | null; - - const defaultResolver = (obj: any) => obj[fieldName]; - const { resolve: oldResolve = defaultResolver, ...rest } = field; - - return { - ...rest, - async resolve(source: any, args: any, graphqlContext: any, info: any) { - // Walk through the input args and embed any *Text companion fields - async function embedTextFields(obj: any): Promise { - if (!obj || typeof obj !== 'object') return; - - const pending: Promise[] = []; - - for (const key of Object.keys(obj)) { - const value = obj[key]; - - // Check if this key is a *Text companion field - if (key in textToVectorMap && typeof value === 'string') { - const vectorFieldName = textToVectorMap[key]; - - pending.push((async () => { - if (!embedder) { - throw new Error( - `EMBED_NOT_CONFIGURED: Cannot embed ${key} — no embedding provider configured. ` + - 'Set defaultEmbedder in GraphileLlmPreset options or EMBEDDER_PROVIDER env var.' - ); - } - - const startTime = Date.now(); - const vector = await embedder(value); - const latencyMs = Date.now() - startTime; - - if (vector === null) { - throw new Error( - `EMBED_QUOTA_EXCEEDED: Cannot embed ${key} — embedding quota exceeded. ` + - 'Upgrade your plan or wait for the next billing period.' - ); - } - - console.log( - `[graphile-llm] Mutation embed: field=${key}, dims=${vector.length}, latency=${latencyMs}ms` - ); - - // Inject the vector into the corresponding field - obj[vectorFieldName] = vector; - // Remove the consumed *Text field - delete obj[key]; - })()); - continue; - } - - // Recurse into nested objects (e.g. input.article.embeddingText) - if (value !== null && typeof value === 'object' && !Array.isArray(value)) { - pending.push(embedTextFields(value)); - } - } - - if (pending.length > 0) { - await Promise.all(pending); - } - } - - await embedTextFields(args); - return oldResolve(source, args, graphqlContext, info); - } - }; + finalize(schema, build) { + return planTimeEmbedding(schema, build); } } } diff --git a/graphile/graphile-llm/src/plugins/text-search-plugin.ts b/graphile/graphile-llm/src/plugins/text-search-plugin.ts index cb6d8e34e9..e247a2e246 100644 --- a/graphile/graphile-llm/src/plugins/text-search-plugin.ts +++ b/graphile/graphile-llm/src/plugins/text-search-plugin.ts @@ -10,30 +10,26 @@ * accepts a compound input (point + distance) and the plugin handles * the conversion to SQL internally. * - * The `text` field is mutually exclusive with `vector`: clients provide - * one or the other. When `text` is provided, the plugin embeds it and - * injects the resulting vector into the normal pgvector pipeline. + * Clients provide `text` or `vector`; `vector` becomes optional. When both + * are given the vector wins. `unifiedSearch` text is embedded too, so pgvector + * participates alongside the text adapters. * - * Runtime embedding for query filters uses the v4-style resolver wrapping - * approach (same as graphile-upload-plugin). When a connection query's - * `where` argument includes a VectorNearbyInput with `text`, the resolver - * wrapper embeds the text and replaces it with the resulting vector before - * the plan executes. + * The text is embedded while the operation is planned (see embed-inputs.ts), + * before SQL is built. * - * If the embedder is not configured, the `text` field is still registered - * (so the schema is stable) but will return a clear error at execution time. + * If no embedder is configured, `text` is ignored and `unifiedSearch` runs + * on the text adapters only. * * If the embedder returns null (e.g. quota exceeded when the metering * plugin is loaded), behavior depends on `onQuotaExceeded`: - * - `'degrade'` (default): silently removes the text field and continues - * with text-only search as a graceful fallback. - * - `'throw'`: always throws an error, even if text adapters could handle it. + * - `'degrade'` (default): drops the text and continues without the vector + * filter (text-only search for `unifiedSearch`). + * - `'throw'`: throws an error. */ import type { GraphileConfig } from 'graphile-config'; -import type { LlmConfigOverrides } from '../embedder'; -import { llmConfigStore } from '../embedder'; +import { embedField, planTimeEmbedding } from '../embed-inputs'; // ─── TypeScript Augmentation ──────────────────────────────────────────────── @@ -45,23 +41,16 @@ declare global { } } -/** - * Check if a codec has any pgvector `vector` columns. - */ -function hasVectorColumns(pgCodec: any): boolean { - if (!pgCodec?.attributes) return false; - for (const attribute of Object.values(pgCodec.attributes as Record)) { - if (attribute.codec?.name === 'vector') return true; - } - return false; -} - /** * Recursively walk a `where` argument object and embed any VectorNearbyInput * values that have `text` instead of `vector`. * * If the embedder returns null (e.g. quota exceeded), the text field is * removed so the pgvector filter is skipped — graceful text-only fallback. + * + * @deprecated The plugins no longer call this: under Grafast a resolver runs + * after the plan's SQL, so embedding there is too late. They embed at plan + * time through `embedField` and `planTimeEmbedding` instead. */ export async function embedTextInWhere( obj: any, @@ -153,43 +142,12 @@ export async function embedTextInWhere( } } -/** Embedder signature expected by embedTextInWhere */ -type WhereEmbedder = (text: string) => Promise; - -/** - * Resolve per-DB LLM config overrides from express-context. - * Returns an LlmConfigOverrides object if per-DB config is available, - * or null to use defaults. These overrides are string parameters only - * (model name, base URL) — the embedder function and its metering - * wrapper stay the same. - */ -async function resolveConfigOverrides( - graphqlContext: any -): Promise { - const ctx = graphqlContext?.constructive; - if (!ctx?.useLlm) return null; - - try { - const llm = await ctx.useLlm(); - if (!llm?.embeddingModel) return null; - - return { - embeddingModel: llm.embeddingModel, - embeddingBaseUrl: llm.embeddingBaseUrl, - chatModel: llm.chatModel ?? undefined, - chatBaseUrl: llm.chatBaseUrl ?? undefined, - }; - } catch { - return null; - } -} - /** * Creates the LlmTextSearchPlugin. * * Hooks into VectorNearbyInput to add a `text` field alongside the - * existing `vector` field. When a user provides `text`, the plugin's - * resolver wrapper embeds it before passing to pgvector. + * existing `vector` field, and tags `text` and `unifiedSearch` for + * plan-time embedding. */ export function createLlmTextSearchPlugin( options: { onQuotaExceeded?: 'degrade' | 'throw' } = {} @@ -197,7 +155,7 @@ export function createLlmTextSearchPlugin( const { onQuotaExceeded = 'degrade' } = options; return { name: 'LlmTextSearchPlugin', - version: '0.2.0', + version: '0.3.0', description: 'Adds text-to-vector embedding support on VectorNearbyInput filter fields', after: [ @@ -208,117 +166,52 @@ export function createLlmTextSearchPlugin( schema: { hooks: { - /** - * Add the `text: String` field to VectorNearbyInput. - * - * We intercept VectorNearbyInput specifically and add a `text` field. - * The field is optional — clients provide either `text` or `vector`. - */ GraphQLInputObjectType_fields(fields, build, context) { - const typeName = context.Self.name; - - if (typeName !== 'VectorNearbyInput') { + if (context.Self.name !== 'VectorNearbyInput') { return fields; } const { - graphql: { GraphQLString } + graphql: { GraphQLString, getNullableType } } = build; return build.extend( - fields, { - text: { - type: GraphQLString, - description: - 'Natural language text to embed server-side for similarity search. ' + - 'Mutually exclusive with `vector` — provide one or the other. ' + - 'Requires the LLM plugin to be configured with an embedding provider.' - } + ...fields, + ...(fields.vector && { + vector: { ...fields.vector, type: getNullableType(fields.vector.type) as any } + }) + }, + { + text: embedField( + { kind: 'nearby', into: 'vector', onQuotaExceeded }, + { + type: GraphQLString, + description: + 'Natural language text to embed server-side for similarity search. ' + + 'Provide `text` or `vector`; if both are given, `vector` is used. ' + + 'Requires the LLM plugin to be configured with an embedding provider.' + } + ) }, 'LlmTextSearchPlugin adding text field to VectorNearbyInput' ); }, - /** - * Wrap connection query resolvers to intercept `where` arguments that - * contain VectorNearbyInput with `text` or `unifiedSearch` with text, - * embed the text, and inject the resulting vector before the plan executes. - * - * For tables with vector columns: embeds VectorNearbyInput.text → vector - * For ALL tables with unifiedSearch: embeds unifiedSearch text → { __text, __vector } - * - * Uses the same v4-style resolver wrapping pattern as graphile-upload-plugin - * and graphile-bucket-provisioner-plugin. - */ - GraphQLObjectType_fields_field(field, build, context) { - const { - scope: { isRootQuery, pgCodec } - } = context as any; - - if (!isRootQuery || !pgCodec) return field; - - // Wrap if the table has vector columns OR has any searchable columns - // (for unifiedSearch embedding support) - const hasVector = hasVectorColumns(pgCodec); - const hasSearchableColumns = pgCodec.attributes && Object.values( - pgCodec.attributes as Record - ).some((attr: any) => - attr.codec?.name === 'tsvector' || attr.codec?.name === 'vector' - ); - - if (!hasVector && !hasSearchableColumns) return field; - - const embedder = (build as any).llmEmbedder as - | WhereEmbedder - | null; - if (!embedder) return field; - - // Determine if this table has text-based search adapters for fallback logic - const hasTextAdapters = pgCodec.attributes && Object.values( - pgCodec.attributes as Record - ).some((attr: any) => attr.codec?.name === 'tsvector'); - - const defaultResolver = (obj: any) => obj[context.scope.fieldName]; - const { resolve: oldResolve = defaultResolver, ...rest } = field; - - return { - ...rest, - async resolve(source: any, args: any, graphqlContext: any, info: any) { - // Resolve per-DB model/baseUrl overrides from express-context. - // These are just string parameters — the embedder function - // (and its metering wrapper) stays the same. - const overrides = await resolveConfigOverrides(graphqlContext); - - // Run within llmConfigStore so the embedder picks up per-DB - // model/baseUrl at call time. Metering wraps the same function - // and is unaffected. - return llmConfigStore.run(overrides, async () => { - if (args?.where) { - await embedTextInWhere(args.where, embedder, !!hasTextAdapters, onQuotaExceeded); - } - - if (args?.filter) { - await embedTextInWhere(args.filter, embedder, !!hasTextAdapters, onQuotaExceeded); - } - - return oldResolve(source, args, graphqlContext, info); - }); - } - }; + GraphQLInputObjectType_fields_field(field, _build, context) { + const { fieldName, isPgConnectionFilterField } = context.scope as any; + if (fieldName !== 'unifiedSearch' || !isPgConnectionFilterField) return field; + return embedField({ kind: 'unified', onQuotaExceeded }, field); }, finalize(schema, build) { - const embedder = (build as any).llmEmbedder; - - if (!embedder) { + if (!(build as any).llmEmbedder) { console.log( - '[graphile-llm] No embedder available — text field on VectorNearbyInput ' + - 'will return errors if used. Configure an embedding provider to enable.' + '[graphile-llm] No embedder available — text on VectorNearbyInput is ignored ' + + 'and unifiedSearch runs without pgvector. Configure an embedding provider to enable.' ); } - - return schema; + return planTimeEmbedding(schema, build); } } } diff --git a/graphile/graphile-settings/__tests__/preset-integration.test.ts b/graphile/graphile-settings/__tests__/preset-integration.test.ts index bd2aec987a..b950ee9b3a 100644 --- a/graphile/graphile-settings/__tests__/preset-integration.test.ts +++ b/graphile/graphile-settings/__tests__/preset-integration.test.ts @@ -967,7 +967,7 @@ describe('Kitchen sink (multi-plugin queries)', () => { # 6. pgvector similarity (PgVectorPlugin) # WHERE vector_embedding <=> '[0,1,0]' < 2.0 - vectorEmbedding: { nearby: { embedding: [0, 1, 0], distance: 2.0 } } + vectorEmbedding: { vector: [0, 1, 0], metric: COSINE, distance: 2.0 } # 7. PostGIS spatial (PostGISFilterPlugin) # WHERE ST_Intersects(geom, $bbox::geometry)