@llmindset/hf-mcp 0.2.6 → 0.2.7
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- package/dist/dataset-detail.d.ts +2 -1
- package/dist/dataset-detail.d.ts.map +1 -1
- package/dist/dataset-detail.js +5 -1
- package/dist/dataset-detail.js.map +1 -1
- package/dist/dataset-search.d.ts +3 -2
- package/dist/dataset-search.d.ts.map +1 -1
- package/dist/dataset-search.js +15 -3
- package/dist/dataset-search.js.map +1 -1
- package/dist/docs-search/docs-semantic-search.d.ts +2 -1
- package/dist/docs-search/docs-semantic-search.d.ts.map +1 -1
- package/dist/docs-search/docs-semantic-search.js +17 -5
- package/dist/docs-search/docs-semantic-search.js.map +1 -1
- package/dist/docs-search/docs-semantic-search.test.js +71 -51
- package/dist/docs-search/docs-semantic-search.test.js.map +1 -1
- package/dist/index.d.ts +1 -0
- package/dist/index.d.ts.map +1 -1
- package/dist/index.js +1 -0
- package/dist/index.js.map +1 -1
- package/dist/model-detail.d.ts +2 -1
- package/dist/model-detail.d.ts.map +1 -1
- package/dist/model-detail.js +5 -1
- package/dist/model-detail.js.map +1 -1
- package/dist/model-search.d.ts +3 -2
- package/dist/model-search.d.ts.map +1 -1
- package/dist/model-search.js +15 -3
- package/dist/model-search.js.map +1 -1
- package/dist/paper-search.d.ts +2 -1
- package/dist/paper-search.d.ts.map +1 -1
- package/dist/paper-search.js +15 -3
- package/dist/paper-search.js.map +1 -1
- package/dist/paper-summary.js +6 -6
- package/dist/paper-summary.js.map +1 -1
- package/dist/space-search.d.ts +3 -2
- package/dist/space-search.d.ts.map +1 -1
- package/dist/space-search.js +15 -3
- package/dist/space-search.js.map +1 -1
- package/dist/types/tool-result.d.ts +6 -0
- package/dist/types/tool-result.d.ts.map +1 -0
- package/dist/types/tool-result.js +2 -0
- package/dist/types/tool-result.js.map +1 -0
- package/dist/user-summary.js +5 -5
- package/dist/user-summary.js.map +1 -1
- package/package.json +1 -1
- package/src/dataset-detail.ts +9 -4
- package/src/dataset-search.ts +19 -6
- package/src/docs-search/docs-semantic-search.test.ts +71 -51
- package/src/docs-search/docs-semantic-search.ts +20 -7
- package/src/index.ts +3 -0
- package/src/model-detail.ts +9 -4
- package/src/model-search.ts +19 -6
- package/src/paper-search.ts +19 -6
- package/src/paper-summary.ts +6 -6
- package/src/space-search.ts +19 -6
- package/src/types/tool-result.ts +24 -0
- package/src/user-summary.ts +5 -5
package/src/dataset-search.ts
CHANGED
|
@@ -1,6 +1,7 @@
|
|
|
1
1
|
import { z } from 'zod';
|
|
2
2
|
import { HfApiCall } from './hf-api-call.js';
|
|
3
3
|
import { formatDate, formatNumber } from './utilities.js';
|
|
4
|
+
import type { ToolResult } from './types/tool-result.js';
|
|
4
5
|
const TAGS_TO_RETURN = 20;
|
|
5
6
|
// Dataset Search Tool Configuration
|
|
6
7
|
export const DATASET_SEARCH_TOOL_CONFIG = {
|
|
@@ -83,7 +84,7 @@ export class DatasetSearchTool extends HfApiCall<DatasetApiParams, DatasetApiRes
|
|
|
83
84
|
/**
|
|
84
85
|
* Search for datasets with detailed parameters
|
|
85
86
|
*/
|
|
86
|
-
async searchWithParams(params: Partial<DatasetSearchParams>): Promise<
|
|
87
|
+
async searchWithParams(params: Partial<DatasetSearchParams>): Promise<ToolResult> {
|
|
87
88
|
try {
|
|
88
89
|
// Convert our params to the HF API format
|
|
89
90
|
const apiParams: DatasetApiParams = {};
|
|
@@ -118,7 +119,11 @@ export class DatasetSearchTool extends HfApiCall<DatasetApiParams, DatasetApiRes
|
|
|
118
119
|
const datasets = await this.callApi<DatasetApiResult[]>(apiParams);
|
|
119
120
|
|
|
120
121
|
if (datasets.length === 0) {
|
|
121
|
-
return
|
|
122
|
+
return {
|
|
123
|
+
formatted: `No datasets found for the given criteria.`,
|
|
124
|
+
totalResults: 0,
|
|
125
|
+
resultsShared: 0
|
|
126
|
+
};
|
|
122
127
|
}
|
|
123
128
|
|
|
124
129
|
return formatSearchResults(datasets, params);
|
|
@@ -133,7 +138,7 @@ export class DatasetSearchTool extends HfApiCall<DatasetApiParams, DatasetApiRes
|
|
|
133
138
|
/**
|
|
134
139
|
* Search for datasets with a specific filter (e.g., arxiv:XXXX.XXXXX)
|
|
135
140
|
*/
|
|
136
|
-
async searchWithFilter(filter: string, limit: number = 10): Promise<
|
|
141
|
+
async searchWithFilter(filter: string, limit: number = 10): Promise<ToolResult> {
|
|
137
142
|
try {
|
|
138
143
|
const apiParams: DatasetApiParams = {
|
|
139
144
|
filter: filter,
|
|
@@ -146,7 +151,11 @@ export class DatasetSearchTool extends HfApiCall<DatasetApiParams, DatasetApiRes
|
|
|
146
151
|
const datasets = await this.callApi<DatasetApiResult[]>(apiParams);
|
|
147
152
|
|
|
148
153
|
if (datasets.length === 0) {
|
|
149
|
-
return
|
|
154
|
+
return {
|
|
155
|
+
formatted: `No datasets found referencing ${filter}.`,
|
|
156
|
+
totalResults: 0,
|
|
157
|
+
resultsShared: 0
|
|
158
|
+
};
|
|
150
159
|
}
|
|
151
160
|
|
|
152
161
|
return formatSearchResults(datasets, { limit });
|
|
@@ -160,7 +169,7 @@ export class DatasetSearchTool extends HfApiCall<DatasetApiParams, DatasetApiRes
|
|
|
160
169
|
}
|
|
161
170
|
|
|
162
171
|
// Formatting Function
|
|
163
|
-
function formatSearchResults(datasets: DatasetApiResult[], params: Partial<DatasetSearchParams>):
|
|
172
|
+
function formatSearchResults(datasets: DatasetApiResult[], params: Partial<DatasetSearchParams>): ToolResult {
|
|
164
173
|
const r: string[] = [];
|
|
165
174
|
|
|
166
175
|
// Build search description
|
|
@@ -233,5 +242,9 @@ function formatSearchResults(datasets: DatasetApiResult[], params: Partial<Datas
|
|
|
233
242
|
r.push('');
|
|
234
243
|
}
|
|
235
244
|
|
|
236
|
-
return
|
|
245
|
+
return {
|
|
246
|
+
formatted: r.join('\n'),
|
|
247
|
+
totalResults: datasets.length,
|
|
248
|
+
resultsShared: datasets.length
|
|
249
|
+
};
|
|
237
250
|
}
|
|
@@ -26,7 +26,9 @@ describe('DocSearchTool', () => {
|
|
|
26
26
|
});
|
|
27
27
|
|
|
28
28
|
const result = await docSearchTool.search({ query: 'nonexistent' });
|
|
29
|
-
expect(result).toBe(`No documentation found for query 'nonexistent'`);
|
|
29
|
+
expect(result.formatted).toBe(`No documentation found for query 'nonexistent'`);
|
|
30
|
+
expect(result.totalResults).toBe(0);
|
|
31
|
+
expect(result.resultsShared).toBe(0);
|
|
30
32
|
});
|
|
31
33
|
|
|
32
34
|
it('should return no results message with product filter', async () => {
|
|
@@ -36,7 +38,9 @@ describe('DocSearchTool', () => {
|
|
|
36
38
|
});
|
|
37
39
|
|
|
38
40
|
const result = await docSearchTool.search({ query: 'nonexistent', product: 'hub' });
|
|
39
|
-
expect(result).toBe(`No documentation found for query 'nonexistent' in product 'hub'`);
|
|
41
|
+
expect(result.formatted).toBe(`No documentation found for query 'nonexistent' in product 'hub'`);
|
|
42
|
+
expect(result.totalResults).toBe(0);
|
|
43
|
+
expect(result.resultsShared).toBe(0);
|
|
40
44
|
});
|
|
41
45
|
|
|
42
46
|
it('should format results grouped by product and page', async () => {
|
|
@@ -74,37 +78,39 @@ describe('DocSearchTool', () => {
|
|
|
74
78
|
const result = await docSearchTool.search({ query: 'analytics' });
|
|
75
79
|
|
|
76
80
|
// Check header
|
|
77
|
-
expect(result).toContain('# Documentation Library Search Results for "analytics"');
|
|
78
|
-
expect(result).toContain('Found 3 results');
|
|
81
|
+
expect(result.formatted).toContain('# Documentation Library Search Results for "analytics"');
|
|
82
|
+
expect(result.formatted).toContain('Found 3 results');
|
|
83
|
+
expect(result.totalResults).toBe(3);
|
|
84
|
+
expect(result.resultsShared).toBe(3);
|
|
79
85
|
|
|
80
86
|
// Check product grouping - hub should come before dataset-viewer (hub has 2 results, dataset-viewer has 1)
|
|
81
|
-
const hubIndex = result.indexOf('## Results for Product: hub');
|
|
82
|
-
const datasetViewerIndex = result.indexOf('## Results for Product: dataset-viewer');
|
|
87
|
+
const hubIndex = result.formatted.indexOf('## Results for Product: hub');
|
|
88
|
+
const datasetViewerIndex = result.formatted.indexOf('## Results for Product: dataset-viewer');
|
|
83
89
|
expect(hubIndex).toBeLessThan(datasetViewerIndex);
|
|
84
90
|
expect(hubIndex).toBeGreaterThan(-1);
|
|
85
91
|
expect(datasetViewerIndex).toBeGreaterThan(-1);
|
|
86
92
|
|
|
87
93
|
// Check that result counts are shown
|
|
88
|
-
expect(result).toContain('## Results for Product: hub (2 results)');
|
|
89
|
-
expect(result).toContain('## Results for Product: dataset-viewer (1 results)');
|
|
94
|
+
expect(result.formatted).toContain('## Results for Product: hub (2 results)');
|
|
95
|
+
expect(result.formatted).toContain('## Results for Product: dataset-viewer (1 results)');
|
|
90
96
|
|
|
91
97
|
// Check page links (without anchors)
|
|
92
|
-
expect(result).toContain(
|
|
98
|
+
expect(result.formatted).toContain(
|
|
93
99
|
'### Results from [Analytics](https://huggingface.co/docs/hub/enterprise-hub-analytics)'
|
|
94
100
|
);
|
|
95
|
-
expect(result).toContain('### Results from [Quickstart](https://huggingface.co/docs/dataset-viewer/quick_start)');
|
|
101
|
+
expect(result.formatted).toContain('### Results from [Quickstart](https://huggingface.co/docs/dataset-viewer/quick_start)');
|
|
96
102
|
|
|
97
103
|
// Check excerpts with heading2
|
|
98
|
-
expect(result).toContain('#### Excerpt from the "Export Analytics as CSV" section');
|
|
99
|
-
expect(result).toContain('#### Excerpt from the "View Analytics" section');
|
|
104
|
+
expect(result.formatted).toContain('#### Excerpt from the "Export Analytics as CSV" section');
|
|
105
|
+
expect(result.formatted).toContain('#### Excerpt from the "View Analytics" section');
|
|
100
106
|
|
|
101
107
|
// Check excerpt content appears as plain text
|
|
102
|
-
expect(result).toContain('Download a comprehensive CSV file containing analytics');
|
|
103
|
-
expect(result).toContain('View analytics for your repositories');
|
|
104
|
-
expect(result).toContain('In this quickstart, you will learn how to use the dataset viewer REST API');
|
|
108
|
+
expect(result.formatted).toContain('Download a comprehensive CSV file containing analytics');
|
|
109
|
+
expect(result.formatted).toContain('View analytics for your repositories');
|
|
110
|
+
expect(result.formatted).toContain('In this quickstart, you will learn how to use the dataset viewer REST API');
|
|
105
111
|
|
|
106
112
|
// Check footer
|
|
107
|
-
expect(result).toContain('Use the "' + DOC_FETCH_CONFIG.name + '" tool to fetch a document from the library.');
|
|
113
|
+
expect(result.formatted).toContain('Use the "' + DOC_FETCH_CONFIG.name + '" tool to fetch a document from the library.');
|
|
108
114
|
});
|
|
109
115
|
|
|
110
116
|
it('should handle results without heading2', async () => {
|
|
@@ -126,8 +132,10 @@ describe('DocSearchTool', () => {
|
|
|
126
132
|
const result = await docSearchTool.search({ query: 'transformers' });
|
|
127
133
|
|
|
128
134
|
// Should not contain "Excerpt from" when heading2 is missing
|
|
129
|
-
expect(result).not.toContain('#### Excerpt from');
|
|
130
|
-
expect(result).toContain('This is a simple text without heading2');
|
|
135
|
+
expect(result.formatted).not.toContain('#### Excerpt from');
|
|
136
|
+
expect(result.formatted).toContain('This is a simple text without heading2');
|
|
137
|
+
expect(result.totalResults).toBe(1);
|
|
138
|
+
expect(result.resultsShared).toBe(1);
|
|
131
139
|
});
|
|
132
140
|
|
|
133
141
|
it('should properly escape markdown special characters', async () => {
|
|
@@ -150,9 +158,11 @@ describe('DocSearchTool', () => {
|
|
|
150
158
|
const result = await docSearchTool.search({ query: 'special' });
|
|
151
159
|
|
|
152
160
|
// Check that special characters are escaped in headings and page titles
|
|
153
|
-
expect(result).toContain('Special \\* Characters');
|
|
161
|
+
expect(result.formatted).toContain('Special \\* Characters');
|
|
154
162
|
// Note: heading2 appears in header text, but brackets don't get escaped
|
|
155
|
-
expect(result).toContain('#### Excerpt from the "Section with [brackets]" section');
|
|
163
|
+
expect(result.formatted).toContain('#### Excerpt from the "Section with [brackets]" section');
|
|
164
|
+
expect(result.totalResults).toBe(1);
|
|
165
|
+
expect(result.resultsShared).toBe(1);
|
|
156
166
|
});
|
|
157
167
|
|
|
158
168
|
it('should clean HTML tags from text', async () => {
|
|
@@ -174,9 +184,11 @@ describe('DocSearchTool', () => {
|
|
|
174
184
|
const result = await docSearchTool.search({ query: 'html' });
|
|
175
185
|
|
|
176
186
|
// HTML tags should be removed
|
|
177
|
-
expect(result).toContain('Text with HTML tags and');
|
|
178
|
-
expect(result).not.toContain('<div');
|
|
179
|
-
expect(result).not.toContain('<img');
|
|
187
|
+
expect(result.formatted).toContain('Text with HTML tags and');
|
|
188
|
+
expect(result.formatted).not.toContain('<div');
|
|
189
|
+
expect(result.formatted).not.toContain('<img');
|
|
190
|
+
expect(result.totalResults).toBe(1);
|
|
191
|
+
expect(result.resultsShared).toBe(1);
|
|
180
192
|
});
|
|
181
193
|
|
|
182
194
|
it('should sort multiple products and pages correctly by count', async () => {
|
|
@@ -229,23 +241,25 @@ describe('DocSearchTool', () => {
|
|
|
229
241
|
const result = await docSearchTool.search({ query: 'test' });
|
|
230
242
|
|
|
231
243
|
// Check product order by count: hub (3) > transformers (1) = datasets (1)
|
|
232
|
-
const hubIndex = result.indexOf('## Results for Product: hub');
|
|
233
|
-
const transformersIndex = result.indexOf('## Results for Product: transformers');
|
|
234
|
-
const datasetsIndex = result.indexOf('## Results for Product: datasets');
|
|
244
|
+
const hubIndex = result.formatted.indexOf('## Results for Product: hub');
|
|
245
|
+
const transformersIndex = result.formatted.indexOf('## Results for Product: transformers');
|
|
246
|
+
const datasetsIndex = result.formatted.indexOf('## Results for Product: datasets');
|
|
235
247
|
|
|
236
248
|
expect(hubIndex).toBeLessThan(transformersIndex);
|
|
237
249
|
expect(hubIndex).toBeLessThan(datasetsIndex);
|
|
238
250
|
|
|
239
251
|
// Check that hub shows total count
|
|
240
|
-
expect(result).toContain('## Results for Product: hub (3 results)');
|
|
252
|
+
expect(result.formatted).toContain('## Results for Product: hub (3 results)');
|
|
241
253
|
|
|
242
254
|
// Check page order within hub product: page1 (2 results) should come before page2 (1 result)
|
|
243
|
-
const page1Index = result.indexOf('https://huggingface.co/docs/hub/page1');
|
|
244
|
-
const page2Index = result.indexOf('https://huggingface.co/docs/hub/page2');
|
|
255
|
+
const page1Index = result.formatted.indexOf('https://huggingface.co/docs/hub/page1');
|
|
256
|
+
const page2Index = result.formatted.indexOf('https://huggingface.co/docs/hub/page2');
|
|
245
257
|
expect(page1Index).toBeLessThan(page2Index);
|
|
246
258
|
|
|
247
259
|
// Check that page1 shows its multiple results count
|
|
248
|
-
expect(result).toContain('### Results from [Page 1](https://huggingface.co/docs/hub/page1) (2 results)');
|
|
260
|
+
expect(result.formatted).toContain('### Results from [Page 1](https://huggingface.co/docs/hub/page1) (2 results)');
|
|
261
|
+
expect(result.totalResults).toBe(5);
|
|
262
|
+
expect(result.resultsShared).toBe(5);
|
|
249
263
|
});
|
|
250
264
|
|
|
251
265
|
it('should include product filter in API call when provided', async () => {
|
|
@@ -295,16 +309,18 @@ describe('DocSearchTool', () => {
|
|
|
295
309
|
const result = await docSearchTool.search({ query: 'analytics' });
|
|
296
310
|
|
|
297
311
|
// All three results should be grouped under one page heading (without anchor)
|
|
298
|
-
expect(result).toContain('### Results from [Analytics](https://huggingface.co/docs/hub/analytics) (3 results)');
|
|
312
|
+
expect(result.formatted).toContain('### Results from [Analytics](https://huggingface.co/docs/hub/analytics) (3 results)');
|
|
299
313
|
|
|
300
314
|
// All three excerpts should appear under the same page
|
|
301
|
-
expect(result).toContain('First result from section 1');
|
|
302
|
-
expect(result).toContain('Second result from section 2');
|
|
303
|
-
expect(result).toContain('Third result from section 3');
|
|
315
|
+
expect(result.formatted).toContain('First result from section 1');
|
|
316
|
+
expect(result.formatted).toContain('Second result from section 2');
|
|
317
|
+
expect(result.formatted).toContain('Third result from section 3');
|
|
304
318
|
|
|
305
319
|
// There should only be one "Results from" heading for this page
|
|
306
|
-
const resultsFromCount = (result.match(/### Results from/g) || []).length;
|
|
320
|
+
const resultsFromCount = (result.formatted.match(/### Results from/g) || []).length;
|
|
307
321
|
expect(resultsFromCount).toBe(1);
|
|
322
|
+
expect(result.totalResults).toBe(3);
|
|
323
|
+
expect(result.resultsShared).toBe(3);
|
|
308
324
|
});
|
|
309
325
|
|
|
310
326
|
it('should handle API errors gracefully', async () => {
|
|
@@ -349,33 +365,35 @@ describe('DocSearchTool', () => {
|
|
|
349
365
|
const result = await smallBudgetTool.search({ query: 'test' });
|
|
350
366
|
|
|
351
367
|
// Should contain the header
|
|
352
|
-
expect(result).toContain('# Documentation Library Search Results for "test"');
|
|
353
|
-
expect(result).toContain('Found 8 results');
|
|
368
|
+
expect(result.formatted).toContain('# Documentation Library Search Results for "test"');
|
|
369
|
+
expect(result.formatted).toContain('Found 8 results');
|
|
354
370
|
|
|
355
371
|
// Early results should have full content
|
|
356
|
-
expect(result).toContain('This is a very long text that repeats');
|
|
372
|
+
expect(result.formatted).toContain('This is a very long text that repeats');
|
|
357
373
|
|
|
358
374
|
// Debug: let's see what the result looks like
|
|
359
|
-
console.log('Result length:', result.length);
|
|
360
|
-
console.log('Estimated tokens:', Math.ceil(result.length / 3.3));
|
|
375
|
+
console.log('Result length:', result.formatted.length);
|
|
376
|
+
console.log('Estimated tokens:', Math.ceil(result.formatted.length / 3.3));
|
|
361
377
|
|
|
362
378
|
// Check that early pages have content
|
|
363
|
-
expect(result).toContain('#### Excerpt from the "Section 0" section');
|
|
364
|
-
expect(result).toContain('This is a very long text that repeats');
|
|
379
|
+
expect(result.formatted).toContain('#### Excerpt from the "Section 0" section');
|
|
380
|
+
expect(result.formatted).toContain('This is a very long text that repeats');
|
|
365
381
|
|
|
366
382
|
// With token budget, we should see either truncation or link-only section
|
|
367
|
-
const hasTruncation = result.includes(`*[Content truncated - use ${DOC_FETCH_CONFIG.name} for full text or narrow search terms]*`);
|
|
368
|
-
const hasAdditionalResults = result.includes('## Further results were found in:');
|
|
383
|
+
const hasTruncation = result.formatted.includes(`*[Content truncated - use ${DOC_FETCH_CONFIG.name} for full text or narrow search terms]*`);
|
|
384
|
+
const hasAdditionalResults = result.formatted.includes('## Further results were found in:');
|
|
369
385
|
|
|
370
386
|
|
|
371
387
|
// At least one of these should indicate budget management
|
|
372
388
|
expect(hasTruncation || hasAdditionalResults).toBeTruthy();
|
|
373
389
|
|
|
374
390
|
// Should still contain the footer
|
|
375
|
-
expect(result).toContain(`Use the "${DOC_FETCH_CONFIG.name}" tool to fetch a document from the library.`);
|
|
391
|
+
expect(result.formatted).toContain(`Use the "${DOC_FETCH_CONFIG.name}" tool to fetch a document from the library.`);
|
|
376
392
|
|
|
377
393
|
// Result should be around our smaller token budget (5k tokens = ~16.5k chars)
|
|
378
|
-
expect(result.length).toBeLessThan(25000); // Should be controlled by token budget
|
|
394
|
+
expect(result.formatted.length).toBeLessThan(25000); // Should be controlled by token budget
|
|
395
|
+
expect(result.totalResults).toBe(8);
|
|
396
|
+
expect(result.resultsShared).toBe(8);
|
|
379
397
|
});
|
|
380
398
|
});
|
|
381
399
|
|
|
@@ -413,18 +431,20 @@ describe('DocSearchTool', () => {
|
|
|
413
431
|
const result = await docSearchTool.search({ query: 'test' });
|
|
414
432
|
|
|
415
433
|
// Verify grouping structure in output
|
|
416
|
-
expect(result).toContain('## Results for Product: hub');
|
|
417
|
-
expect(result).toContain('## Results for Product: transformers');
|
|
434
|
+
expect(result.formatted).toContain('## Results for Product: hub');
|
|
435
|
+
expect(result.formatted).toContain('## Results for Product: transformers');
|
|
418
436
|
|
|
419
437
|
// Verify that both results from the same page are together
|
|
420
|
-
const result1Index = result.indexOf('Result 1');
|
|
421
|
-
const result2Index = result.indexOf('Result 2');
|
|
422
|
-
const result3Index = result.indexOf('Result 3');
|
|
438
|
+
const result1Index = result.formatted.indexOf('Result 1');
|
|
439
|
+
const result2Index = result.formatted.indexOf('Result 2');
|
|
440
|
+
const result3Index = result.formatted.indexOf('Result 3');
|
|
423
441
|
|
|
424
442
|
// Results 1 and 2 should be close together (same page)
|
|
425
443
|
expect(Math.abs(result2Index - result1Index)).toBeLessThan(100);
|
|
426
444
|
// Result 3 should be further away (different product)
|
|
427
445
|
expect(Math.abs(result3Index - result1Index)).toBeGreaterThan(50);
|
|
446
|
+
expect(result.totalResults).toBe(3);
|
|
447
|
+
expect(result.resultsShared).toBe(3);
|
|
428
448
|
});
|
|
429
449
|
});
|
|
430
450
|
});
|
|
@@ -2,6 +2,7 @@ import { z } from 'zod';
|
|
|
2
2
|
import { HfApiCall } from '../hf-api-call.js';
|
|
3
3
|
import { escapeMarkdown, estimateTokens } from '../utilities.js';
|
|
4
4
|
import { DOC_FETCH_CONFIG } from './doc-fetch.js';
|
|
5
|
+
import type { ToolResult } from '../types/tool-result.js';
|
|
5
6
|
|
|
6
7
|
/** token estimation. initial results for "how to load a image to image model in transformers" returned
|
|
7
8
|
* 121973 characters (36711 anthropic tokens) */
|
|
@@ -72,9 +73,13 @@ export class DocSearchTool extends HfApiCall<DocSearchApiParams, DocSearchResult
|
|
|
72
73
|
* @param query Search query string (e.g. "rate limits", "analytics")
|
|
73
74
|
* @param product Optional product filter
|
|
74
75
|
*/
|
|
75
|
-
async search(params: DocSearchParams): Promise<
|
|
76
|
+
async search(params: DocSearchParams): Promise<ToolResult> {
|
|
76
77
|
try {
|
|
77
|
-
if (!params.query) return
|
|
78
|
+
if (!params.query) return {
|
|
79
|
+
formatted: 'No query provided',
|
|
80
|
+
totalResults: 0,
|
|
81
|
+
resultsShared: 0
|
|
82
|
+
};
|
|
78
83
|
|
|
79
84
|
const apiParams: DocSearchApiParams = { q: params.query.toLowerCase() };
|
|
80
85
|
if (params.product) {
|
|
@@ -84,9 +89,13 @@ export class DocSearchTool extends HfApiCall<DocSearchApiParams, DocSearchResult
|
|
|
84
89
|
const results = await this.callApi<DocSearchResult[]>(apiParams);
|
|
85
90
|
|
|
86
91
|
if (results.length === 0) {
|
|
87
|
-
return
|
|
88
|
-
|
|
89
|
-
|
|
92
|
+
return {
|
|
93
|
+
formatted: params.product
|
|
94
|
+
? `No documentation found for query '${params.query}' in product '${params.product}'`
|
|
95
|
+
: `No documentation found for query '${params.query}'`,
|
|
96
|
+
totalResults: 0,
|
|
97
|
+
resultsShared: 0
|
|
98
|
+
};
|
|
90
99
|
}
|
|
91
100
|
|
|
92
101
|
return formatSearchResults(params.query, results, params.product, this.tokenBudget);
|
|
@@ -207,7 +216,7 @@ function formatSearchResults(
|
|
|
207
216
|
results: DocSearchResult[],
|
|
208
217
|
productFilter?: string,
|
|
209
218
|
tokenBudget = DEFAULT_TOKEN_BUDGET
|
|
210
|
-
):
|
|
219
|
+
): ToolResult {
|
|
211
220
|
const lines: string[] = [];
|
|
212
221
|
let hasShownTruncationMessage = false;
|
|
213
222
|
|
|
@@ -311,5 +320,9 @@ function formatSearchResults(
|
|
|
311
320
|
lines.push('---\n');
|
|
312
321
|
lines.push(`Use the "${DOC_FETCH_CONFIG.name}" tool to fetch a document from the library.`);
|
|
313
322
|
|
|
314
|
-
return
|
|
323
|
+
return {
|
|
324
|
+
formatted: lines.join('\n'),
|
|
325
|
+
totalResults: results.length,
|
|
326
|
+
resultsShared: results.length
|
|
327
|
+
};
|
|
315
328
|
}
|
package/src/index.ts
CHANGED
|
@@ -16,5 +16,8 @@ export * from './paper-summary.js';
|
|
|
16
16
|
export * from './docs-search/docs-semantic-search.js';
|
|
17
17
|
export * from './docs-search/doc-fetch.js';
|
|
18
18
|
|
|
19
|
+
// Export shared types
|
|
20
|
+
export * from './types/tool-result.js';
|
|
21
|
+
|
|
19
22
|
// Export tool IDs for external use - these are the canonical tool identifiers
|
|
20
23
|
export * from './tool-ids.js';
|
package/src/model-detail.ts
CHANGED
|
@@ -1,6 +1,7 @@
|
|
|
1
1
|
import { modelInfo } from '@huggingface/hub';
|
|
2
2
|
import { z } from 'zod';
|
|
3
3
|
import { formatDate, formatNumber } from './utilities.js';
|
|
4
|
+
import type { ToolResult } from './types/tool-result.js';
|
|
4
5
|
|
|
5
6
|
const SPACES_TO_INCLUDE = 12;
|
|
6
7
|
// Model Detail Tool Configuration
|
|
@@ -102,9 +103,9 @@ export class ModelDetailTool {
|
|
|
102
103
|
* Get detailed information about a specific model
|
|
103
104
|
*
|
|
104
105
|
* @param modelId The model ID to get details for (e.g., microsoft/DialoGPT-large)
|
|
105
|
-
* @returns
|
|
106
|
+
* @returns ToolResult with formatted model details
|
|
106
107
|
*/
|
|
107
|
-
async getDetails(modelId: string): Promise<
|
|
108
|
+
async getDetails(modelId: string): Promise<ToolResult> {
|
|
108
109
|
try {
|
|
109
110
|
// Define additional fields we want to retrieve (only those available in the hub library)
|
|
110
111
|
const additionalFields = [
|
|
@@ -252,7 +253,7 @@ export class ModelDetailTool {
|
|
|
252
253
|
}
|
|
253
254
|
|
|
254
255
|
// Formatting Function
|
|
255
|
-
function formatModelDetails(model: ModelInformation):
|
|
256
|
+
function formatModelDetails(model: ModelInformation): ToolResult {
|
|
256
257
|
const r: string[] = [];
|
|
257
258
|
const [authorFromName] = model.name.includes('/') ? model.name.split('/') : ['', model.name];
|
|
258
259
|
|
|
@@ -404,5 +405,9 @@ function formatModelDetails(model: ModelInformation): string {
|
|
|
404
405
|
// Link is reliable - based on model name which is required
|
|
405
406
|
r.push(`**Link:** [https://hf.co/${model.name}](https://hf.co/${model.name})`);
|
|
406
407
|
|
|
407
|
-
return
|
|
408
|
+
return {
|
|
409
|
+
formatted: r.join('\n'),
|
|
410
|
+
totalResults: 1, // Model was found
|
|
411
|
+
resultsShared: 1 // All details shared
|
|
412
|
+
};
|
|
408
413
|
}
|
package/src/model-search.ts
CHANGED
|
@@ -1,6 +1,7 @@
|
|
|
1
1
|
import { z } from 'zod';
|
|
2
2
|
import { HfApiCall } from './hf-api-call.js';
|
|
3
3
|
import { formatDate, formatNumber } from './utilities.js';
|
|
4
|
+
import type { ToolResult } from './types/tool-result.js';
|
|
4
5
|
|
|
5
6
|
export const TAGS_TO_RETURN = 20;
|
|
6
7
|
// Model Search Tool Configuration
|
|
@@ -82,7 +83,7 @@ export class ModelSearchTool extends HfApiCall<ModelApiParams, ModelApiResult[]>
|
|
|
82
83
|
/**
|
|
83
84
|
* Search for models with detailed parameters
|
|
84
85
|
*/
|
|
85
|
-
async searchWithParams(params: Partial<ModelSearchParams>): Promise<
|
|
86
|
+
async searchWithParams(params: Partial<ModelSearchParams>): Promise<ToolResult> {
|
|
86
87
|
try {
|
|
87
88
|
// Convert our params to the HF API format
|
|
88
89
|
const apiParams: ModelApiParams = {};
|
|
@@ -120,7 +121,11 @@ export class ModelSearchTool extends HfApiCall<ModelApiParams, ModelApiResult[]>
|
|
|
120
121
|
const models = await this.callApi<ModelApiResult[]>(apiParams);
|
|
121
122
|
|
|
122
123
|
if (models.length === 0) {
|
|
123
|
-
return
|
|
124
|
+
return {
|
|
125
|
+
formatted: `No models found for the given criteria.`,
|
|
126
|
+
totalResults: 0,
|
|
127
|
+
resultsShared: 0
|
|
128
|
+
};
|
|
124
129
|
}
|
|
125
130
|
|
|
126
131
|
return formatSearchResults(models, params);
|
|
@@ -135,7 +140,7 @@ export class ModelSearchTool extends HfApiCall<ModelApiParams, ModelApiResult[]>
|
|
|
135
140
|
/**
|
|
136
141
|
* Search for models with a specific filter (e.g., arxiv:XXXX.XXXXX)
|
|
137
142
|
*/
|
|
138
|
-
async searchWithFilter(filter: string, limit: number = 10): Promise<
|
|
143
|
+
async searchWithFilter(filter: string, limit: number = 10): Promise<ToolResult> {
|
|
139
144
|
try {
|
|
140
145
|
const apiParams: ModelApiParams = {
|
|
141
146
|
filter: filter,
|
|
@@ -148,7 +153,11 @@ export class ModelSearchTool extends HfApiCall<ModelApiParams, ModelApiResult[]>
|
|
|
148
153
|
const models = await this.callApi<ModelApiResult[]>(apiParams);
|
|
149
154
|
|
|
150
155
|
if (models.length === 0) {
|
|
151
|
-
return
|
|
156
|
+
return {
|
|
157
|
+
formatted: `No models found referencing ${filter}.`,
|
|
158
|
+
totalResults: 0,
|
|
159
|
+
resultsShared: 0
|
|
160
|
+
};
|
|
152
161
|
}
|
|
153
162
|
|
|
154
163
|
return formatSearchResults(models, { limit });
|
|
@@ -162,7 +171,7 @@ export class ModelSearchTool extends HfApiCall<ModelApiParams, ModelApiResult[]>
|
|
|
162
171
|
}
|
|
163
172
|
|
|
164
173
|
// Formatting Function
|
|
165
|
-
function formatSearchResults(models: ModelApiResult[], params: Partial<ModelSearchParams>):
|
|
174
|
+
function formatSearchResults(models: ModelApiResult[], params: Partial<ModelSearchParams>): ToolResult {
|
|
166
175
|
const r: string[] = [];
|
|
167
176
|
|
|
168
177
|
// Build search description
|
|
@@ -227,5 +236,9 @@ function formatSearchResults(models: ModelApiResult[], params: Partial<ModelSear
|
|
|
227
236
|
r.push('');
|
|
228
237
|
}
|
|
229
238
|
|
|
230
|
-
return
|
|
239
|
+
return {
|
|
240
|
+
formatted: r.join('\n'),
|
|
241
|
+
totalResults: models.length,
|
|
242
|
+
resultsShared: models.length
|
|
243
|
+
};
|
|
231
244
|
}
|
package/src/paper-search.ts
CHANGED
|
@@ -1,6 +1,7 @@
|
|
|
1
1
|
import { z } from 'zod';
|
|
2
2
|
import { HfApiCall } from './hf-api-call.js';
|
|
3
3
|
import { formatUnknownDate } from './utilities.js';
|
|
4
|
+
import type { ToolResult } from './types/tool-result.js';
|
|
4
5
|
|
|
5
6
|
// https://github.com/huggingface/huggingface_hub/blob/a26b93e8ba0b51ce76ce5c2044896587c47c6b60/src/huggingface_hub/hf_api.py#L1481-L1542
|
|
6
7
|
// Raw JSON response for https://hf.co/api/papers/search?q=llama%203%20herd Llama Herd is ~50,000 tokens
|
|
@@ -93,15 +94,23 @@ export class PaperSearchTool extends HfApiCall<PaperSearchParams, PaperSearchRes
|
|
|
93
94
|
* Searches for papers on the Hugging Face Hub
|
|
94
95
|
* @param query Search query string (e.g. "llama", "attention")
|
|
95
96
|
* @param limit Maximum number of results to return
|
|
96
|
-
* @returns
|
|
97
|
+
* @returns ToolResult with formatted paper information and metrics
|
|
97
98
|
*/
|
|
98
|
-
async search(query: string, limit: number = RESULTS_TO_RETURN, conciseOnly: boolean = false): Promise<
|
|
99
|
+
async search(query: string, limit: number = RESULTS_TO_RETURN, conciseOnly: boolean = false): Promise<ToolResult> {
|
|
99
100
|
try {
|
|
100
|
-
if (!query) return
|
|
101
|
+
if (!query) return {
|
|
102
|
+
formatted: 'No query',
|
|
103
|
+
totalResults: 0,
|
|
104
|
+
resultsShared: 0
|
|
105
|
+
};
|
|
101
106
|
|
|
102
107
|
const papers = await this.callApi<PaperSearchResult[]>({ q: query });
|
|
103
108
|
|
|
104
|
-
if (papers.length === 0) return
|
|
109
|
+
if (papers.length === 0) return {
|
|
110
|
+
formatted: `No papers found for query '${query}'`,
|
|
111
|
+
totalResults: 0,
|
|
112
|
+
resultsShared: 0
|
|
113
|
+
};
|
|
105
114
|
return formatSearchResults(query, papers.slice(0, limit), papers.length, conciseOnly);
|
|
106
115
|
} catch (error) {
|
|
107
116
|
if (error instanceof Error) {
|
|
@@ -122,7 +131,7 @@ function formatSearchResults(
|
|
|
122
131
|
papers: PaperSearchResult[],
|
|
123
132
|
totalCount: number,
|
|
124
133
|
conciseOnly: boolean = false
|
|
125
|
-
):
|
|
134
|
+
): ToolResult {
|
|
126
135
|
const r: string[] = [];
|
|
127
136
|
const showingText =
|
|
128
137
|
papers.length < totalCount
|
|
@@ -169,7 +178,11 @@ function formatSearchResults(
|
|
|
169
178
|
}
|
|
170
179
|
r.push('');
|
|
171
180
|
r.push('---');
|
|
172
|
-
return
|
|
181
|
+
return {
|
|
182
|
+
formatted: r.join('\n'),
|
|
183
|
+
totalResults: totalCount,
|
|
184
|
+
resultsShared: papers.length
|
|
185
|
+
};
|
|
173
186
|
}
|
|
174
187
|
|
|
175
188
|
export function authors(authors: Author[] | undefined, authorsToShow: number = DEFAULT_AUTHORS_TO_SHOW): string {
|
package/src/paper-summary.ts
CHANGED
|
@@ -267,8 +267,8 @@ export class PaperSummaryPrompt extends HfApiCall<Record<string, string>, PaperD
|
|
|
267
267
|
const modelSearch = new ModelSearchTool(this.hfToken);
|
|
268
268
|
// Use the filter parameter to search for models referencing this paper
|
|
269
269
|
const modelResults = await modelSearch.searchWithFilter(`arxiv:${arxivId}`, 25);
|
|
270
|
-
if (modelResults && !modelResults.includes('No models found')) {
|
|
271
|
-
results.models = `## Related Models\n\n${modelResults}`;
|
|
270
|
+
if (modelResults && !modelResults.formatted.includes('No models found')) {
|
|
271
|
+
results.models = `## Related Models\n\n${modelResults.formatted}`;
|
|
272
272
|
}
|
|
273
273
|
} catch (error) {
|
|
274
274
|
console.warn(`Failed to fetch related models for paper ${arxivId}:`, error);
|
|
@@ -279,8 +279,8 @@ export class PaperSummaryPrompt extends HfApiCall<Record<string, string>, PaperD
|
|
|
279
279
|
const datasetSearch = new DatasetSearchTool(this.hfToken);
|
|
280
280
|
// Use the filter parameter to search for datasets referencing this paper
|
|
281
281
|
const datasetResults = await datasetSearch.searchWithFilter(`arxiv:${arxivId}`, 25);
|
|
282
|
-
if (datasetResults && !datasetResults.includes('No datasets found')) {
|
|
283
|
-
results.datasets = `## Related Datasets\n\n${datasetResults}`;
|
|
282
|
+
if (datasetResults && !datasetResults.formatted.includes('No datasets found')) {
|
|
283
|
+
results.datasets = `## Related Datasets\n\n${datasetResults.formatted}`;
|
|
284
284
|
}
|
|
285
285
|
} catch (error) {
|
|
286
286
|
console.warn(`Failed to fetch related datasets for paper ${arxivId}:`, error);
|
|
@@ -291,8 +291,8 @@ export class PaperSummaryPrompt extends HfApiCall<Record<string, string>, PaperD
|
|
|
291
291
|
const spaceSearch = new SpaceSearchTool(this.hfToken);
|
|
292
292
|
// Use the filter parameter to search for spaces referencing this paper
|
|
293
293
|
const spaceResults = await spaceSearch.searchWithFilter(`arxiv:${arxivId}`, 25, 2);
|
|
294
|
-
if (spaceResults && !spaceResults.includes('No matching Hugging Face Spaces found')) {
|
|
295
|
-
results.spaces = `## Related Spaces\n\n${spaceResults}`;
|
|
294
|
+
if (spaceResults && !spaceResults.formatted.includes('No matching Hugging Face Spaces found')) {
|
|
295
|
+
results.spaces = `## Related Spaces\n\n${spaceResults.formatted}`;
|
|
296
296
|
}
|
|
297
297
|
} catch (error) {
|
|
298
298
|
console.warn(`Failed to fetch related spaces for paper ${arxivId}:`, error);
|
package/src/space-search.ts
CHANGED
|
@@ -1,6 +1,7 @@
|
|
|
1
1
|
import { z } from 'zod';
|
|
2
2
|
import { HfApiCall } from './hf-api-call.js';
|
|
3
3
|
import { escapeMarkdown } from './utilities.js';
|
|
4
|
+
import type { ToolResult } from './types/tool-result.js';
|
|
4
5
|
|
|
5
6
|
// Define the SearchResult interface
|
|
6
7
|
export interface SpaceSearchResult {
|
|
@@ -99,7 +100,7 @@ export class SpaceSearchTool extends HfApiCall<SpaceSearchParams, SpaceSearchRes
|
|
|
99
100
|
* Search for spaces with a specific filter (e.g., arxiv:XXXX.XXXXX)
|
|
100
101
|
* Note: For spaces, we need to use the regular API endpoint with filter parameter
|
|
101
102
|
*/
|
|
102
|
-
async searchWithFilter(filter: string, limit: number = 10, headerLevel: number = 1): Promise<
|
|
103
|
+
async searchWithFilter(filter: string, limit: number = 10, headerLevel: number = 1): Promise<ToolResult> {
|
|
103
104
|
try {
|
|
104
105
|
// For spaces, we need to use the regular spaces API endpoint with filter
|
|
105
106
|
const url = new URL('https://huggingface.co/api/spaces');
|
|
@@ -111,7 +112,11 @@ export class SpaceSearchTool extends HfApiCall<SpaceSearchParams, SpaceSearchRes
|
|
|
111
112
|
const results = await this.fetchFromApi<SpaceSearchResult[]>(url);
|
|
112
113
|
|
|
113
114
|
if (results.length === 0) {
|
|
114
|
-
return
|
|
115
|
+
return {
|
|
116
|
+
formatted: `No matching Hugging Face Spaces found referencing ${filter}.`,
|
|
117
|
+
totalResults: 0,
|
|
118
|
+
resultsShared: 0
|
|
119
|
+
};
|
|
115
120
|
}
|
|
116
121
|
|
|
117
122
|
// Format results using the existing formatter
|
|
@@ -133,16 +138,20 @@ export type SearchParams = z.infer<typeof SEMANTIC_SEARCH_TOOL_CONFIG.schema>;
|
|
|
133
138
|
/**
|
|
134
139
|
* Formats search results as a markdown table for MCP friendly output
|
|
135
140
|
* @param results The search results to format
|
|
136
|
-
* @returns A
|
|
141
|
+
* @returns A ToolResult with formatted string and metrics
|
|
137
142
|
*/
|
|
138
143
|
export const formatSearchResults = (
|
|
139
144
|
query: string,
|
|
140
145
|
results: SpaceSearchResult[],
|
|
141
146
|
totalCount: number,
|
|
142
147
|
headerLevel: number = 1
|
|
143
|
-
):
|
|
148
|
+
): ToolResult => {
|
|
144
149
|
if (results.length === 0) {
|
|
145
|
-
return
|
|
150
|
+
return {
|
|
151
|
+
formatted: `No matching Hugging Face Spaces found for the query '${query}'. Try a different query.`,
|
|
152
|
+
totalResults: 0,
|
|
153
|
+
resultsShared: 0
|
|
154
|
+
};
|
|
146
155
|
}
|
|
147
156
|
|
|
148
157
|
const showingText =
|
|
@@ -173,5 +182,9 @@ export const formatSearchResults = (
|
|
|
173
182
|
`| ${relevance} |\n`;
|
|
174
183
|
}
|
|
175
184
|
|
|
176
|
-
return
|
|
185
|
+
return {
|
|
186
|
+
formatted: markdown,
|
|
187
|
+
totalResults: totalCount,
|
|
188
|
+
resultsShared: results.length
|
|
189
|
+
};
|
|
177
190
|
};
|