diff --git a/bun.lock b/bun.lock index 0b36cb7..10ba56f 100644 --- a/bun.lock +++ b/bun.lock @@ -163,6 +163,7 @@ "scripts": { "name": "scripts", "dependencies": { + "@memory.build/embedding": "workspace:*", "yaml": "^2.7.0", }, }, diff --git a/queries/.gitignore b/queries/.gitignore new file mode 100644 index 0000000..30f1b39 --- /dev/null +++ b/queries/.gitignore @@ -0,0 +1,2 @@ +emb.txt +semantic.txt diff --git a/queries/01-fulltext.sql b/queries/01-fulltext.sql new file mode 100644 index 0000000..17351fc --- /dev/null +++ b/queries/01-fulltext.sql @@ -0,0 +1,31 @@ +-- Full-text/BM25 search. +-- Required: schema, user_id, fulltext. Optional: limit, query_prefix. +-- Mirrors buildBM25Query without additional filters. + +\ir _setup.sql + +\if :{?fulltext} +\else +\prompt 'fulltext: ' fulltext +\endif + +\timing on + +:query_prefix +select + id +, content +, meta +, tree::text +, temporal::text +, embedding is not null as has_embedding +, created_at +, created_by +, updated_at +, -(content <@> to_bm25query(:'fulltext', :'schema' || '.memory_content_bm25_idx')) as score +from :"schema".memory +where content <@> to_bm25query(:'fulltext', :'schema' || '.memory_content_bm25_idx') < 0 +order by score desc, created_at desc +limit :limit; + +\ir _teardown.sql diff --git a/queries/02-semantic.sql b/queries/02-semantic.sql new file mode 100644 index 0000000..f504ec1 --- /dev/null +++ b/queries/02-semantic.sql @@ -0,0 +1,28 @@ +-- Semantic/HNSW search. +-- Required: schema, user_id. Provide emb or semantic. Optional: limit, query_prefix. +-- Mirrors buildSemanticQuery without additional filters and without semanticThreshold. + +\ir _setup.sql +\ir _embedding.sql + +\timing on + +:query_prefix +select + id +, content +, meta +, tree::text +, temporal::text +, embedding is not null as has_embedding +, created_at +, created_by +, updated_at +, (1 - (embedding <=> :'emb'::halfvec)) as score +from :"schema".memory +where embedding is not null + and (embedding <=> :'emb'::halfvec) < 1.0 +order by score desc, created_at desc +limit :limit; + +\ir _teardown.sql diff --git a/queries/03-ltree.sql b/queries/03-ltree.sql new file mode 100644 index 0000000..039584f --- /dev/null +++ b/queries/03-ltree.sql @@ -0,0 +1,31 @@ +-- Filter-only ltree subtree query. +-- Required: schema, user_id, tree. Optional: limit, order_direction, query_prefix. +-- Mirrors buildFilterQuery with a plain ltree filter. + +\ir _setup.sql + +\if :{?tree} +\else +\prompt 'tree: ' tree +\endif + +\timing on + +:query_prefix +select + id +, content +, meta +, tree::text +, temporal::text +, embedding is not null as has_embedding +, created_at +, created_by +, updated_at +, 1.0 as score +from :"schema".memory +where tree <@ :'tree'::ltree +order by created_at :order_direction +limit :limit; + +\ir _teardown.sql diff --git a/queries/04-meta.sql b/queries/04-meta.sql new file mode 100644 index 0000000..220a91f --- /dev/null +++ b/queries/04-meta.sql @@ -0,0 +1,31 @@ +-- Filter-only JSONB metadata containment query. +-- Required: schema, user_id, meta. Optional: limit, order_direction, query_prefix. +-- Mirrors buildFilterQuery with meta @> filter. + +\ir _setup.sql + +\if :{?meta} +\else +\prompt 'meta json: ' meta +\endif + +\timing on + +:query_prefix +select + id +, content +, meta +, tree::text +, temporal::text +, embedding is not null as has_embedding +, created_at +, created_by +, updated_at +, 1.0 as score +from :"schema".memory +where meta @> :'meta'::jsonb +order by created_at :order_direction +limit :limit; + +\ir _teardown.sql diff --git a/queries/05-meta-ltree.sql b/queries/05-meta-ltree.sql new file mode 100644 index 0000000..3ef385b --- /dev/null +++ b/queries/05-meta-ltree.sql @@ -0,0 +1,37 @@ +-- Filter-only JSONB metadata + ltree subtree query. +-- Required: schema, user_id, meta, tree. Optional: limit, order_direction, query_prefix. +-- Mirrors buildFilterQuery with common filters in app order: meta, tree. + +\ir _setup.sql + +\if :{?meta} +\else +\prompt 'meta json: ' meta +\endif + +\if :{?tree} +\else +\prompt 'tree: ' tree +\endif + +\timing on + +:query_prefix +select + id +, content +, meta +, tree::text +, temporal::text +, embedding is not null as has_embedding +, created_at +, created_by +, updated_at +, 1.0 as score +from :"schema".memory +where meta @> :'meta'::jsonb + and tree <@ :'tree'::ltree +order by created_at :order_direction +limit :limit; + +\ir _teardown.sql diff --git a/queries/06-fulltext-ltree.sql b/queries/06-fulltext-ltree.sql new file mode 100644 index 0000000..f8ba4cb --- /dev/null +++ b/queries/06-fulltext-ltree.sql @@ -0,0 +1,37 @@ +-- Full-text/BM25 search with ltree subtree filter. +-- Required: schema, user_id, fulltext, tree. Optional: limit, query_prefix. +-- Mirrors buildBM25Query with common filters in app order: tree. + +\ir _setup.sql + +\if :{?fulltext} +\else +\prompt 'fulltext: ' fulltext +\endif + +\if :{?tree} +\else +\prompt 'tree: ' tree +\endif + +\timing on + +:query_prefix +select + id +, content +, meta +, tree::text +, temporal::text +, embedding is not null as has_embedding +, created_at +, created_by +, updated_at +, -(content <@> to_bm25query(:'fulltext', :'schema' || '.memory_content_bm25_idx')) as score +from :"schema".memory +where content <@> to_bm25query(:'fulltext', :'schema' || '.memory_content_bm25_idx') < 0 + and tree <@ :'tree'::ltree +order by score desc, created_at desc +limit :limit; + +\ir _teardown.sql diff --git a/queries/07-fulltext-meta.sql b/queries/07-fulltext-meta.sql new file mode 100644 index 0000000..73f5ebf --- /dev/null +++ b/queries/07-fulltext-meta.sql @@ -0,0 +1,37 @@ +-- Full-text/BM25 search with JSONB metadata filter. +-- Required: schema, user_id, fulltext, meta. Optional: limit, query_prefix. +-- Mirrors buildBM25Query with common filters in app order: meta. + +\ir _setup.sql + +\if :{?fulltext} +\else +\prompt 'fulltext: ' fulltext +\endif + +\if :{?meta} +\else +\prompt 'meta json: ' meta +\endif + +\timing on + +:query_prefix +select + id +, content +, meta +, tree::text +, temporal::text +, embedding is not null as has_embedding +, created_at +, created_by +, updated_at +, -(content <@> to_bm25query(:'fulltext', :'schema' || '.memory_content_bm25_idx')) as score +from :"schema".memory +where content <@> to_bm25query(:'fulltext', :'schema' || '.memory_content_bm25_idx') < 0 + and meta @> :'meta'::jsonb +order by score desc, created_at desc +limit :limit; + +\ir _teardown.sql diff --git a/queries/08-fulltext-meta-ltree.sql b/queries/08-fulltext-meta-ltree.sql new file mode 100644 index 0000000..6abce1a --- /dev/null +++ b/queries/08-fulltext-meta-ltree.sql @@ -0,0 +1,43 @@ +-- Full-text/BM25 search with JSONB metadata + ltree subtree filters. +-- Required: schema, user_id, fulltext, meta, tree. Optional: limit, query_prefix. +-- Mirrors buildBM25Query with common filters in app order: meta, tree. + +\ir _setup.sql + +\if :{?fulltext} +\else +\prompt 'fulltext: ' fulltext +\endif + +\if :{?meta} +\else +\prompt 'meta json: ' meta +\endif + +\if :{?tree} +\else +\prompt 'tree: ' tree +\endif + +\timing on + +:query_prefix +select + id +, content +, meta +, tree::text +, temporal::text +, embedding is not null as has_embedding +, created_at +, created_by +, updated_at +, -(content <@> to_bm25query(:'fulltext', :'schema' || '.memory_content_bm25_idx')) as score +from :"schema".memory +where content <@> to_bm25query(:'fulltext', :'schema' || '.memory_content_bm25_idx') < 0 + and meta @> :'meta'::jsonb + and tree <@ :'tree'::ltree +order by score desc, created_at desc +limit :limit; + +\ir _teardown.sql diff --git a/queries/09-semantic-ltree.sql b/queries/09-semantic-ltree.sql new file mode 100644 index 0000000..2b85cf4 --- /dev/null +++ b/queries/09-semantic-ltree.sql @@ -0,0 +1,34 @@ +-- Semantic/HNSW search with ltree subtree filter. +-- Required: schema, user_id, tree. Provide emb or semantic. Optional: limit, query_prefix. +-- Mirrors buildSemanticQuery with common filters in app order: tree. + +\ir _setup.sql +\ir _embedding.sql + +\if :{?tree} +\else +\prompt 'tree: ' tree +\endif + +\timing on + +:query_prefix +select + id +, content +, meta +, tree::text +, temporal::text +, embedding is not null as has_embedding +, created_at +, created_by +, updated_at +, (1 - (embedding <=> :'emb'::halfvec)) as score +from :"schema".memory +where embedding is not null + and (embedding <=> :'emb'::halfvec) < 1.0 + and tree <@ :'tree'::ltree +order by score desc, created_at desc +limit :limit; + +\ir _teardown.sql diff --git a/queries/10-semantic-meta.sql b/queries/10-semantic-meta.sql new file mode 100644 index 0000000..56984f2 --- /dev/null +++ b/queries/10-semantic-meta.sql @@ -0,0 +1,34 @@ +-- Semantic/HNSW search with JSONB metadata filter. +-- Required: schema, user_id, meta. Provide emb or semantic. Optional: limit, query_prefix. +-- Mirrors buildSemanticQuery with common filters in app order: meta. + +\ir _setup.sql +\ir _embedding.sql + +\if :{?meta} +\else +\prompt 'meta json: ' meta +\endif + +\timing on + +:query_prefix +select + id +, content +, meta +, tree::text +, temporal::text +, embedding is not null as has_embedding +, created_at +, created_by +, updated_at +, (1 - (embedding <=> :'emb'::halfvec)) as score +from :"schema".memory +where embedding is not null + and (embedding <=> :'emb'::halfvec) < 1.0 + and meta @> :'meta'::jsonb +order by score desc, created_at desc +limit :limit; + +\ir _teardown.sql diff --git a/queries/11-semantic-meta-ltree.sql b/queries/11-semantic-meta-ltree.sql new file mode 100644 index 0000000..2ab3635 --- /dev/null +++ b/queries/11-semantic-meta-ltree.sql @@ -0,0 +1,40 @@ +-- Semantic/HNSW search with JSONB metadata + ltree subtree filters. +-- Required: schema, user_id, meta, tree. Provide emb or semantic. Optional: limit, query_prefix. +-- Mirrors buildSemanticQuery with common filters in app order: meta, tree. + +\ir _setup.sql +\ir _embedding.sql + +\if :{?meta} +\else +\prompt 'meta json: ' meta +\endif + +\if :{?tree} +\else +\prompt 'tree: ' tree +\endif + +\timing on + +:query_prefix +select + id +, content +, meta +, tree::text +, temporal::text +, embedding is not null as has_embedding +, created_at +, created_by +, updated_at +, (1 - (embedding <=> :'emb'::halfvec)) as score +from :"schema".memory +where embedding is not null + and (embedding <=> :'emb'::halfvec) < 1.0 + and meta @> :'meta'::jsonb + and tree <@ :'tree'::ltree +order by score desc, created_at desc +limit :limit; + +\ir _teardown.sql diff --git a/queries/12-hybrid-candidates.sql b/queries/12-hybrid-candidates.sql new file mode 100644 index 0000000..e985859 --- /dev/null +++ b/queries/12-hybrid-candidates.sql @@ -0,0 +1,54 @@ +-- Hybrid candidate queries: BM25 candidate query followed by semantic candidate query. +-- Required: schema, user_id, fulltext. Provide emb or semantic. +-- Optional: candidate_limit, query_prefix. +-- The app runs these two queries in parallel, fuses IDs in TypeScript, then fetches by ID. +-- psql runs them sequentially but keeps each query faithful to the app SQL. + +\ir _setup.sql +\ir _embedding.sql + +\if :{?fulltext} +\else +\prompt 'fulltext: ' fulltext +\endif + +\timing on + +\echo BM25 candidates +:query_prefix +select + id +, content +, meta +, tree::text +, temporal::text +, embedding is not null as has_embedding +, created_at +, created_by +, updated_at +, -(content <@> to_bm25query(:'fulltext', :'schema' || '.memory_content_bm25_idx')) as score +from :"schema".memory +where content <@> to_bm25query(:'fulltext', :'schema' || '.memory_content_bm25_idx') < 0 +order by score desc, created_at desc +limit :candidate_limit; + +\echo Semantic candidates +:query_prefix +select + id +, content +, meta +, tree::text +, temporal::text +, embedding is not null as has_embedding +, created_at +, created_by +, updated_at +, (1 - (embedding <=> :'emb'::halfvec)) as score +from :"schema".memory +where embedding is not null + and (embedding <=> :'emb'::halfvec) < 1.0 +order by score desc, created_at desc +limit :candidate_limit; + +\ir _teardown.sql diff --git a/queries/13-tree-access.sql b/queries/13-tree-access.sql new file mode 100644 index 0000000..6d16bd8 --- /dev/null +++ b/queries/13-tree-access.sql @@ -0,0 +1,12 @@ +-- Direct RLS helper expansion benchmark. +-- Required: schema, user_id. Optional: query_prefix. + +\ir _setup.sql + +\timing on + +:query_prefix +select * +from :"schema".tree_access(:'user_id'::uuid, 'read'); + +\ir _teardown.sql diff --git a/queries/14-count-visible.sql b/queries/14-count-visible.sql new file mode 100644 index 0000000..ca937c5 --- /dev/null +++ b/queries/14-count-visible.sql @@ -0,0 +1,12 @@ +-- Minimal visible-row scan through RLS. +-- Required: schema, user_id. Optional: query_prefix. + +\ir _setup.sql + +\timing on + +:query_prefix +select count(*)::int as count +from :"schema".memory; + +\ir _teardown.sql diff --git a/queries/15-diagnose-semantic-visibility.sql b/queries/15-diagnose-semantic-visibility.sql new file mode 100644 index 0000000..9318567 --- /dev/null +++ b/queries/15-diagnose-semantic-visibility.sql @@ -0,0 +1,59 @@ +-- Diagnose why semantic search returns no rows under app-equivalent RLS. +-- Required: schema, user_id. Optional: memory_id, semantic, emb, query_prefix. +-- If emb is omitted, this includes _embedding.sql to generate/load queries/emb.txt. + +\ir _setup.sql +\ir _embedding.sql + +\if :{?memory_id} +\else +\prompt 'memory_id to inspect: ' memory_id +\endif + +\timing on + +\echo Current execution context +:query_prefix +select + current_user as current_user +, current_role as current_role +, current_setting('me.user_id', true) as me_user_id +, current_schema() as current_schema; + +\echo Visible memory counts through RLS +:query_prefix +select + count(*)::int as visible_rows +, count(embedding)::int as visible_embedded_rows +from :"schema".memory; + +\echo Known memory row through RLS +:query_prefix +select + id +, tree::text +, embedding is not null as has_embedding +, created_at +, created_by +, left(content, 500) as content_prefix +from :"schema".memory +where id = :'memory_id'::uuid; + +\echo Distance for known memory row through RLS +:query_prefix +select + id +, embedding is not null as has_embedding +, embedding <=> :'emb'::halfvec as distance +, (1 - (embedding <=> :'emb'::halfvec)) as score +from :"schema".memory +where id = :'memory_id'::uuid; + +\echo Effective read tree_access rows +:query_prefix +select tree_path::text +from :"schema".tree_access(:'user_id'::uuid, 'read') as ta(tree_path) +order by tree_path::text +limit 100; + +\ir _teardown.sql diff --git a/queries/16-semantic-index-order.sql b/queries/16-semantic-index-order.sql new file mode 100644 index 0000000..6d3d15c --- /dev/null +++ b/queries/16-semantic-index-order.sql @@ -0,0 +1,29 @@ +-- Semantic/HNSW search ordered by raw distance for pgvector index eligibility. +-- Required: schema, user_id. Provide emb or semantic. Optional: limit, query_prefix. +-- This is NOT exactly what the current app query does; it is the pgvector-recommended +-- order shape for using an ANN index: ORDER BY embedding <=> query LIMIT n. + +\ir _setup.sql +\ir _embedding.sql + +\timing on + +:query_prefix +select + id +, content +, meta +, tree::text +, temporal::text +, embedding is not null as has_embedding +, created_at +, created_by +, updated_at +, (1 - (embedding <=> :'emb'::halfvec)) as score +from :"schema".memory +where embedding is not null + and (embedding <=> :'emb'::halfvec) < 1.0 +order by embedding <=> :'emb'::halfvec, created_at desc +limit :limit; + +\ir _teardown.sql diff --git a/queries/17-diagnose-auth-owner.sql b/queries/17-diagnose-auth-owner.sql new file mode 100644 index 0000000..8e9bfcf --- /dev/null +++ b/queries/17-diagnose-auth-owner.sql @@ -0,0 +1,198 @@ +-- Owner/admin diagnostic for auth, grants, and a known memory row. +-- Required: schema, user_id. Optional: memory_id, query_prefix. +-- This intentionally does NOT set role to me_ro, so run it with the fork owner/admin +-- connection. It is for diagnosing why app-equivalent RLS sees no rows. + +\set ON_ERROR_STOP on +\pset pager off +\timing off + +\if :{?schema} +\else +\prompt 'schema: ' schema +\endif + +\if :{?user_id} +\else +\prompt 'user_id: ' user_id +\endif + +\if :{?memory_id} +\else +\prompt 'memory_id to inspect: ' memory_id +\endif + +\if :{?query_prefix} +\else +\set query_prefix -- +\endif + +begin; + +select + set_config('statement_timeout', '25s', true) as statement_timeout +, set_config('lock_timeout', '5s', true) as lock_timeout +, set_config('transaction_timeout', '30s', true) as transaction_timeout +, set_config('idle_in_transaction_session_timeout', '30s', true) as idle_timeout +\gset setup_ + +set local search_path to :"schema", public; + +\timing on + +\echo Owner/admin execution context +:query_prefix +select + current_user as current_user +, current_role as current_role +, current_schema() as current_schema +, current_setting('me.user_id', true) as me_user_id; + +\echo Memory table RLS flags and owner +:query_prefix +select + n.nspname as schema +, c.relname as table_name +, pg_get_userbyid(c.relowner) as table_owner +, c.relrowsecurity as rls_enabled +, c.relforcerowsecurity as force_rls +from pg_class c +join pg_namespace n on n.oid = c.relnamespace +where n.nspname = :'schema' + and c.relname = 'memory'; + +\echo Total memory counts as owner/admin +:query_prefix +select + count(*)::int as total_rows +, count(embedding)::int as embedded_rows +from :"schema".memory; + +\echo Known memory row as owner/admin +:query_prefix +select + id +, tree::text +, embedding is not null as has_embedding +, created_at +, created_by +, left(content, 500) as content_prefix +from :"schema".memory +where id = :'memory_id'::uuid; + +\echo Provided user by id +:query_prefix +select + id +, name +, identity_id +, can_login +, superuser +, createrole +, created_at +from :"schema"."user" +where id = :'user_id'::uuid; + +\echo Users whose identity_id equals provided user_id (common mix-up check) +:query_prefix +select + id +, name +, identity_id +, can_login +, superuser +, createrole +, created_at +from :"schema"."user" +where identity_id = :'user_id'::uuid; + +\echo Effective read tree_access for provided user_id +:query_prefix +select tree_path::text +from :"schema".tree_access(:'user_id'::uuid, 'read') as ta(tree_path) +order by tree_path::text +limit 100; + +\echo Direct grants for provided user_id +:query_prefix +select + g.id +, g.user_id +, u.name as user_name +, g.tree_path::text +, g.actions +, g.with_grant_option +, g.created_at +from :"schema".tree_grant g +join :"schema"."user" u on u.id = g.user_id +where g.user_id = :'user_id'::uuid +order by g.tree_path::text; + +\echo Role memberships involving provided user_id +:query_prefix +select + rm.role_id +, role_user.name as role_name +, rm.member_id +, member_user.name as member_name +, rm.with_admin_option +, rm.created_at +from :"schema".role_membership rm +join :"schema"."user" role_user on role_user.id = rm.role_id +join :"schema"."user" member_user on member_user.id = rm.member_id +where rm.role_id = :'user_id'::uuid + or rm.member_id = :'user_id'::uuid +order by rm.created_at; + +\echo All users summary +:query_prefix +select + id +, name +, identity_id +, can_login +, superuser +, createrole +, created_at +from :"schema"."user" +order by superuser desc, can_login desc, created_at +limit 100; + +\echo All grants summary +:query_prefix +select + g.user_id +, u.name as user_name +, g.tree_path::text +, g.actions +, g.with_grant_option +, g.created_at +from :"schema".tree_grant g +join :"schema"."user" u on u.id = g.user_id +order by u.name, g.tree_path::text +limit 200; + +\echo Users with read access to the known memory tree +:query_prefix +with target as ( + select tree + from :"schema".memory + where id = :'memory_id'::uuid +) +select + u.id +, u.name +, u.identity_id +, u.superuser +, exists ( + select 1 + from :"schema".tree_access(u.id, 'read') as ta(tree_path) + cross join target t + where t.tree <@ ta.tree_path + ) as can_read_known_memory +from :"schema"."user" u +order by can_read_known_memory desc, u.superuser desc, u.name +limit 100; + +\timing off +rollback; diff --git a/queries/README.md b/queries/README.md new file mode 100644 index 0000000..74e31ea --- /dev/null +++ b/queries/README.md @@ -0,0 +1,58 @@ +# psql Benchmark Queries + +These files are intended for manual benchmarking against a forked engine database. +They mirror the SQL generated by `packages/engine/ops/memory.ts` as closely as +possible, including transaction setup, local timeouts, `search_path`, `ROLE me_ro`, +and `me.user_id` for RLS. + +Common variables: + +- `schema`: engine schema, for example `me_abc123`. +- `user_id`: engine user UUID used by RLS. +- `limit`: result limit, default `10`. +- `candidate_limit`: hybrid candidate limit, default `30`. +- `query_prefix`: optional SQL prefix. Defaults to `--` so the next line is a normal `SELECT`. Use `explain (analyze, buffers, settings)` for plans. +- `fulltext`: BM25 query text. +- `semantic`: semantic query text used to generate `emb` via `scripts/embed-query.ts`. +- `emb`: embedding literal in pgvector/halfvec format, for example `[0.1,-0.2,...]`. If omitted, semantic files generate `queries/emb.txt` with `\!`. +- `tree`: ltree path filter. +- `meta`: JSON object for `meta @> ...` filters. + +Example: + +```bash +psql "$DATABASE_URL" \ + -v schema=me_xxx \ + -v user_id=00000000-0000-0000-0000-000000000000 \ + -v fulltext='postgres indexing' \ + -v "query_prefix=explain (analyze, buffers, settings)" \ + -f queries/01-fulltext.sql +``` + +Semantic example: + +```bash +psql "$DATABASE_URL" \ + -v schema=me_xxx \ + -v user_id=00000000-0000-0000-0000-000000000000 \ + -v semantic='history of postgres indexing' \ + -f queries/02-semantic.sql +``` + +Timing is enabled only around the query body, not around setup or embedding generation. + +The semantic helper writes `semantic` to `queries/semantic.txt` with `select +:'semantic'` and `\g`, then runs the embedding script with `\!`. It uses fixed +shell paths because psql does not expand psql variables inside `\!` shell +commands. Pass `emb` directly if you want to bypass generation. + +`15-diagnose-semantic-visibility.sql` checks whether the selected RLS user can +see any embedded rows and whether a known memory ID is visible/embedded. + +`17-diagnose-auth-owner.sql` runs without switching to `me_ro` so you can inspect +engine users, identity mappings, grants, role membership, and the known memory row +with the fork owner/admin connection. + +`02-semantic.sql` mirrors the current app query. `16-semantic-index-order.sql` +uses the pgvector-recommended `ORDER BY embedding <=> query LIMIT n` shape to +check whether the HNSW index can be used. diff --git a/queries/_embedding.sql b/queries/_embedding.sql new file mode 100644 index 0000000..c9b484b --- /dev/null +++ b/queries/_embedding.sql @@ -0,0 +1,18 @@ +-- Semantic query helper. If emb is not already supplied, generate it outside +-- the timed block using scripts/embed-query.ts and then load the file into emb. +-- Optional variable: semantic. Fixed helper files: queries/semantic.txt, queries/emb.txt. +-- Note: psql does not expand variables in \! shell commands, so keep these +-- paths literal unless you also update the shell commands below. + +\if :{?emb} +\else +\echo Generating embedding outside timed block... +\if :{?semantic} +select :'semantic'::text +\g (format=unaligned tuples_only=on) queries/semantic.txt +\! ./bun run scripts/embed-query.ts queries/emb.txt --query-file queries/semantic.txt +\else +\! ./bun run scripts/embed-query.ts queries/emb.txt +\endif +\set emb `cat queries/emb.txt` +\endif diff --git a/queries/_setup.sql b/queries/_setup.sql new file mode 100644 index 0000000..293e806 --- /dev/null +++ b/queries/_setup.sql @@ -0,0 +1,78 @@ +\set ON_ERROR_STOP on +\pset pager off +\timing off + +-- Shared app-equivalent transaction setup. +-- Required: schema, user_id. Optional: app_role, limit, candidate_limit, +-- statement_timeout, lock_timeout, transaction_timeout, idle_timeout, +-- order_direction, query_prefix. + +\if :{?schema} +\else +\prompt 'schema: ' schema +\endif + +\if :{?user_id} +\else +\prompt 'user_id: ' user_id +\endif + +\if :{?app_role} +\else +\set app_role me_ro +\endif + +\if :{?limit} +\else +\set limit 10 +\endif + +\if :{?candidate_limit} +\else +\set candidate_limit 30 +\endif + +\if :{?statement_timeout} +\else +\set statement_timeout 25s +\endif + +\if :{?lock_timeout} +\else +\set lock_timeout 5s +\endif + +\if :{?transaction_timeout} +\else +\set transaction_timeout 30s +\endif + +\if :{?idle_timeout} +\else +\set idle_timeout 30s +\endif + +\if :{?order_direction} +\else +\set order_direction DESC +\endif + +\if :{?query_prefix} +\else +\set query_prefix -- +\endif + +begin; + +select + set_config('statement_timeout', :'statement_timeout', true) as statement_timeout +, set_config('lock_timeout', :'lock_timeout', true) as lock_timeout +, set_config('transaction_timeout', :'transaction_timeout', true) as transaction_timeout +, set_config('idle_in_transaction_session_timeout', :'idle_timeout', true) as idle_timeout +\gset setup_ + +set local search_path to :"schema", public; +set local role :app_role; + +select set_config('me.user_id', :'user_id', true) as user_id +\gset setup_ diff --git a/queries/_teardown.sql b/queries/_teardown.sql new file mode 100644 index 0000000..a58e8df --- /dev/null +++ b/queries/_teardown.sql @@ -0,0 +1,3 @@ +\timing off +rollback; +reset role; diff --git a/queries/tuned-fulltext.sql b/queries/tuned-fulltext.sql new file mode 100644 index 0000000..5e4ebae --- /dev/null +++ b/queries/tuned-fulltext.sql @@ -0,0 +1,31 @@ +-- Full-text/BM25 search. +-- Required: schema, user_id, fulltext. Optional: limit, query_prefix. +-- Mirrors buildBM25Query without additional filters. + +\ir _setup.sql + +\if :{?fulltext} +\else +\prompt 'fulltext: ' fulltext +\endif + +\timing on + +:query_prefix +select + id +, content +, meta +, tree::text +, temporal::text +, embedding is not null as has_embedding +, created_at +, created_by +, updated_at +, -(content <@> to_bm25query(:'fulltext', :'schema' || '.memory_content_bm25_idx')) as score +from :"schema".memory +where content <@> to_bm25query(:'fulltext', :'schema' || '.memory_content_bm25_idx') < 0 +order by content <@> to_bm25query(:'fulltext', :'schema' || '.memory_content_bm25_idx'), created_at desc +limit :limit; + +\ir _teardown.sql diff --git a/queries/tuned-semantic.sql b/queries/tuned-semantic.sql new file mode 100644 index 0000000..65ca819 --- /dev/null +++ b/queries/tuned-semantic.sql @@ -0,0 +1,28 @@ +-- Semantic/HNSW search. +-- Required: schema, user_id. Provide emb or semantic. Optional: limit, query_prefix. +-- Mirrors buildSemanticQuery without additional filters and without semanticThreshold. + +\ir _setup.sql +\ir _embedding.sql + +\timing on + +:query_prefix +select + id +, content +, meta +, tree::text +, temporal::text +, embedding is not null as has_embedding +, created_at +, created_by +, updated_at +, (1 - (embedding <=> :'emb'::halfvec)) as score +from :"schema".memory +where embedding is not null + and (embedding <=> :'emb'::halfvec) < 1.0 +order by embedding <=> :'emb'::halfvec, created_at desc +limit :limit; + +\ir _teardown.sql diff --git a/scripts/embed-query.ts b/scripts/embed-query.ts new file mode 100644 index 0000000..3f67b04 --- /dev/null +++ b/scripts/embed-query.ts @@ -0,0 +1,81 @@ +import { generateEmbedding } from "@memory.build/embedding"; + +const DEFAULT_OUTPUT = "emb.txt"; + +async function loadDotEnv(path = ".env"): Promise { + const file = Bun.file(path); + if (!(await file.exists())) return; + + const text = await file.text(); + for (const rawLine of text.split(/\r?\n/)) { + const line = rawLine.trim(); + if (!line || line.startsWith("#")) continue; + + const match = /^(?:export\s+)?([A-Za-z_][A-Za-z0-9_]*)=(.*)$/.exec(line); + if (!match) continue; + + const [, key, rawValue] = match; + if (!key || process.env[key] !== undefined) continue; + + let value = rawValue?.trim() ?? ""; + const quote = value[0]; + if ((quote === '"' || quote === "'") && value.at(-1) === quote) { + value = value.slice(1, -1); + } + + process.env[key] = value; + } +} + +function parseIntegerEnv(name: string): number | undefined { + const value = process.env[name]; + if (!value) return undefined; + + const parsed = Number.parseInt(value, 10); + if (!Number.isFinite(parsed)) { + throw new Error(`${name} must be an integer, got ${value}`); + } + return parsed; +} + +await loadDotEnv(); + +const outputPath = process.argv[2] ?? DEFAULT_OUTPUT; +const args = process.argv.slice(3); +const queryFileIndex = args.indexOf("--query-file"); +const queryFile = args[queryFileIndex + 1]; +if (queryFileIndex >= 0 && !queryFile) { + throw new Error("--query-file requires a path"); +} +const queryArg = ( + queryFileIndex >= 0 + ? await Bun.file(queryFile as string).text() + : args.join(" ") +).trim(); +const query = queryArg || prompt("Semantic query: ")?.trim(); +if (!query) { + throw new Error("Semantic query is required"); +} + +const apiKey = process.env.EMBEDDING_API_KEY; +if (!apiKey) { + throw new Error("EMBEDDING_API_KEY is required in the environment or .env"); +} + +const embedding = await generateEmbedding(query, { + provider: "openai", + model: "text-embedding-3-small", + dimensions: 1536, + apiKey, + baseUrl: process.env.EMBEDDING_BASE_URL, + options: { + timeoutMs: parseIntegerEnv("EMBEDDING_TIMEOUT_MS"), + maxRetries: parseIntegerEnv("EMBEDDING_MAX_RETRIES"), + }, +}); + +await Bun.write(outputPath, `[${embedding.embedding.join(",")}]\n`); + +console.log( + `Wrote ${embedding.embedding.length}-dim embedding to ${outputPath}`, +); diff --git a/scripts/package.json b/scripts/package.json index 4213f92..8bd79d2 100644 --- a/scripts/package.json +++ b/scripts/package.json @@ -3,6 +3,7 @@ "private": true, "type": "module", "dependencies": { + "@memory.build/embedding": "workspace:*", "yaml": "^2.7.0" } }