diff --git a/package.json b/package.json index 1e96cfc..4023e29 100644 --- a/package.json +++ b/package.json @@ -39,6 +39,7 @@ "files": [ "server.js", "search_utils.js", + "search_dataset_schema.js", "browser_tools.js", "browser_session.js", "aria_snapshot_filter.js", diff --git a/search_dataset_schema.js b/search_dataset_schema.js new file mode 100644 index 0000000..305e03a --- /dev/null +++ b/search_dataset_schema.js @@ -0,0 +1,62 @@ +'use strict'; /*jslint node:true es9:true*/ +import {z} from 'zod'; + +export const DATASET_IDS = [ + 'gd_l1viktl72bvl7bjuj0', + 'gd_me5ppxjr2ge6icjuh0', + 'gd_l1vikfnt1wgvvqz95w', +]; + +export const dataset_id_schema = z.enum(DATASET_IDS); + +export const FILTER_OPERATORS = ['=', '!=', '<', '<=', '>', '>=', 'in', + 'not_in', 'includes', 'not_includes', 'array_includes', + 'not_array_includes', 'is_null', 'is_not_null']; + +const leaf_value_schema = z.union([ + z.string(), + z.number(), + z.boolean(), + z.array(z.union([z.string(), z.number(), z.boolean()])), +]); + +const leaf_schema = z.object({ + name: z.string().describe('Field name to filter on. Get valid field ' + +'names from the list_dataset_fields tool.'), + operator: z.string().describe('Filter operator, one of: ' + +FILTER_OPERATORS.join(', ')), + value: leaf_value_schema, +}); + +function build_node_schema(depth){ + if (depth<=1) + return leaf_schema; + const group_schema = z.object({ + operator: z.enum(['and', 'or']), + filters: z.array(build_node_schema(depth-1)).min(1), + }); + return z.union([group_schema, leaf_schema]); +} + +const MAX_NESTING = 3; +export const filter_schema = build_node_schema(MAX_NESTING+1); + +export function metadata_to_fields(metadata){ + const fields = metadata && typeof metadata=='object' + && metadata.fields && typeof metadata.fields=='object' + ? metadata.fields : {}; + const out = []; + for (const [name, meta] of Object.entries(fields)) + { + if (!meta || typeof meta!='object') + continue; + if (meta.active===false) + continue; + out.push({ + name, + type: meta.type, + description: meta.description, + }); + } + return out; +} diff --git a/server.js b/server.js index ff9e086..4bd8b34 100644 --- a/server.js +++ b/server.js @@ -7,6 +7,8 @@ import {tools as browser_tools} from './browser_tools.js'; import prompts from './prompts.js'; import {GROUPS} from './tool_groups.js'; import {parse_google_search_response} from './search_utils.js'; +import {dataset_id_schema, filter_schema, metadata_to_fields, FILTER_OPERATORS} + from './search_dataset_schema.js'; import {createRequire} from 'node:module'; import {remark} from 'remark'; import strip from 'strip-markdown'; @@ -603,6 +605,93 @@ addTool({ }), }); +const SEARCHABLE_DATASETS_DESC = [ + 'Supported dataset_id values:', + '- gd_l1viktl72bvl7bjuj0: LinkedIn people profiles', + '- gd_me5ppxjr2ge6icjuh0: LinkedIn people profiles (contact-enriched)', + '- gd_l1vikfnt1wgvvqz95w: LinkedIn company information', +].join('\n'); + +addTool({ + name: 'list_dataset_fields', + description: 'List the filterable fields of a searchable dataset ' + +'(field name, type, and description). Call this before ' + +'search_dataset to learn which field names and types you can ' + +'filter on.\n'+SEARCHABLE_DATASETS_DESC, + annotations: { + title: 'List Dataset Fields', + readOnlyHint: true, + openWorldHint: true, + }, + parameters: z.object({dataset_id: dataset_id_schema}), + execute: tool_fn('list_dataset_fields', async({dataset_id}, ctx)=>{ + let response = await base_request({ + url: `https://api.brightdata.com/datasets/${dataset_id}` + +`/metadata`, + method: 'GET', + headers: api_headers(ctx.clientName, 'list_dataset_fields'), + }); + return JSON.stringify(metadata_to_fields(response.data)); + }), +}); + +addTool({ + name: 'search_dataset', + description: 'Search a Bright Data dataset by a filter and get matching ' + +'records back directly (fast Elasticsearch-backed search; no ' + +'snapshot). Use this to FIND MANY records by criteria, as opposed ' + +'to the web_data_* tools which fetch ONE record by URL.\n' + +'First call list_dataset_fields to get valid field names.\n' + +'A filter is a tree: a group {operator:"and"|"or", filters:[...]} ' + +'or a leaf {name, value, operator}. Max nesting depth 3.\n' + +'Leaf operators: '+FILTER_OPERATORS.join(', ')+'.\n' + +SEARCHABLE_DATASETS_DESC, + annotations: { + title: 'Search Dataset', + readOnlyHint: true, + openWorldHint: true, + }, + parameters: z.object({ + dataset_id: dataset_id_schema, + filter: filter_schema.describe('Filter tree describing which ' + +'records to match. Required, cannot be empty.'), + size: z.number().int().positive().max(10).optional().default(10) + .describe('Number of records to return (max 10, default 10)'), + sort: z.union([ + z.enum(['default', 'random']), + z.array(z.record(z.enum(['asc', 'desc']))), + ]).optional().describe('Sorting: "default", "random", or a custom ' + +'array like [{"timestamp":"asc"}]. Use "default" or custom ' + +'sort to paginate with search_after.'), + search_after: z.array(z.any()).optional().describe('Pagination ' + +'cursor from a previous response\'s search_after value.'), + }), + execute: tool_fn('search_dataset', async({dataset_id, filter, size, sort, + search_after}, ctx)=> + { + let body = {mode: 'sync', filter, size}; + if (sort!==undefined) + body.sort = sort; + if (search_after!==undefined) + body.search_after = search_after; + let response = await base_request({ + url: `https://api.brightdata.com/datasets/search/${dataset_id}`, + method: 'POST', + data: body, + headers: { + ...api_headers(ctx.clientName, 'search_dataset'), + 'Content-Type': 'application/json', + }, + }); + let {hits, total_hits, took, search_after: next_cursor} + = response.data || {}; + let result = {hits, total_hits, took}; + if (next_cursor!==undefined) + result.search_after = next_cursor; + return JSON.stringify(result); + }), +}); + addTool({ name: 'session_stats', description: 'Tell the user about the tool usage during this session', diff --git a/test/search-dataset-schema.test.js b/test/search-dataset-schema.test.js new file mode 100644 index 0000000..304787a --- /dev/null +++ b/test/search-dataset-schema.test.js @@ -0,0 +1,109 @@ +'use strict'; /*jslint node:true es9:true*/ +import test from 'node:test'; +import assert from 'node:assert/strict'; +import {DATASET_IDS, dataset_id_schema, metadata_to_fields, FILTER_OPERATORS, + filter_schema} from '../search_dataset_schema.js'; + +test('DATASET_IDS lists the three supported datasets', ()=>{ + assert.deepEqual(DATASET_IDS, [ + 'gd_l1viktl72bvl7bjuj0', + 'gd_me5ppxjr2ge6icjuh0', + 'gd_l1vikfnt1wgvvqz95w', + ]); +}); + +test('dataset_id_schema accepts a supported id', ()=>{ + assert.equal(dataset_id_schema.parse('gd_l1viktl72bvl7bjuj0'), + 'gd_l1viktl72bvl7bjuj0'); +}); + +test('dataset_id_schema rejects an unknown id', ()=>{ + assert.throws(()=>dataset_id_schema.parse('gd_not_a_real_dataset')); +}); + +test('FILTER_OPERATORS lists the documented operators', ()=>{ + assert.deepEqual(FILTER_OPERATORS, [ + '=', '!=', '<', '<=', '>', '>=', + 'in', 'not_in', + 'includes', 'not_includes', + 'array_includes', 'not_array_includes', + 'is_null', 'is_not_null', + ]); +}); + +test('metadata_to_fields keeps active fields as name/type/description', ()=>{ + const metadata = { + id: 'gd_l1vijqt9jfj7olije', + fields: { + name: {type: 'text', active: true, + description: 'The name of the company'}, + url: {type: 'url', required: true, + description: 'The company URL'}, + cb_rank: {type: 'number', active: false, + description: 'Crunchbase rank'}, + }, + }; + assert.deepEqual(metadata_to_fields(metadata), [ + {name: 'name', type: 'text', description: 'The name of the company'}, + {name: 'url', type: 'url', description: 'The company URL'}, + ]); +}); + +test('metadata_to_fields tolerates missing fields object', ()=>{ + assert.deepEqual(metadata_to_fields({id: 'x'}), []); + assert.deepEqual(metadata_to_fields(null), []); +}); + +test('metadata_to_fields skips non-object field entries', ()=>{ + assert.deepEqual(metadata_to_fields({fields: {bad: null, + ok: {type: 'text', description: 'fine'}}}), + [{name: 'ok', type: 'text', description: 'fine'}]); +}); + +test('filter_schema accepts the documented flat example', ()=>{ + const filter = {operator: 'and', filters: [ + {name: 'name', value: 'Egor', operator: 'includes'}, + ]}; + assert.deepEqual(filter_schema.parse(filter), filter); +}); + +test('filter_schema accepts a single leaf node', ()=>{ + const filter = {name: 'cb_rank', value: 100, operator: '<'}; + assert.deepEqual(filter_schema.parse(filter), filter); +}); + +test('filter_schema accepts array and boolean leaf values', ()=>{ + const filter = {operator: 'or', filters: [ + {name: 'tags', value: ['a', 'b'], operator: 'array_includes'}, + {name: 'verified', value: true, operator: '='}, + ]}; + assert.deepEqual(filter_schema.parse(filter), filter); +}); + +test('filter_schema accepts nesting up to depth 3', ()=>{ + const filter = {operator: 'and', filters: [ + {operator: 'or', filters: [ + {operator: 'and', filters: [ + {name: 'name', value: 'x', operator: 'includes'}, + ]}, + ]}, + ]}; + assert.deepEqual(filter_schema.parse(filter), filter); +}); + +test('filter_schema rejects nesting deeper than 3', ()=>{ + const filter = {operator: 'and', filters: [ + {operator: 'or', filters: [ + {operator: 'and', filters: [ + {operator: 'or', filters: [ + {name: 'name', value: 'x', operator: 'includes'}, + ]}, + ]}, + ]}, + ]}; + assert.throws(()=>filter_schema.parse(filter)); +}); + +test('filter_schema rejects an empty object', ()=>{ + assert.throws(()=>filter_schema.parse({})); +}); diff --git a/tool_groups.js b/tool_groups.js index a9027fd..c21d75d 100644 --- a/tool_groups.js +++ b/tool_groups.js @@ -33,6 +33,8 @@ export const GROUPS = { 'web_data_linkedin_job_listings', 'web_data_linkedin_posts', 'web_data_linkedin_people_search', + 'list_dataset_fields', + 'search_dataset', 'web_data_instagram_profiles', 'web_data_instagram_posts', 'web_data_instagram_reels', @@ -95,6 +97,8 @@ export const GROUPS = { 'web_data_google_maps_reviews', 'web_data_zillow_properties_listing', 'web_data_booking_hotel_listings', + 'list_dataset_fields', + 'search_dataset', ], }, RESEARCH: {