Add search_dataset and list_dataset_fields tools

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",

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;
+}

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',

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({}));
+});

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: {