@llmindset/hf-mcp 0.1.17 → 0.1.19

package/src/docs-search/doc-mappings.ts

@@ -0,0 +1,79 @@
+interface DocMapping {
+  repo_id: string;
+  doc_folder: string;
+}
+
+export const DOC_MAPPINGS: Record<string, DocMapping> = {
+  'tokenizers': {
+    repo_id: 'huggingface/tokenizers',
+    doc_folder: 'docs/source-doc-builder'
+  },
+  'diffusers': {
+    repo_id: 'huggingface/diffusers',
+    doc_folder: 'docs/source/en'
+  },
+  'accelerate': {
+    repo_id: 'huggingface/accelerate',
+    doc_folder: 'docs/source'
+  },
+  'huggingface_hub': {
+    repo_id: 'huggingface/huggingface_hub',
+    doc_folder: 'docs/source/en'
+  },
+  'transformers': {
+    repo_id: 'huggingface/transformers',
+    doc_folder: 'docs/source/en'
+  },
+  'hub': {
+    repo_id: 'huggingface/hub-docs',
+    doc_folder: 'docs/hub'
+  },
+  'huggingface.js': {
+    repo_id: 'huggingface/huggingface.js',
+    doc_folder: 'docs'
+  },
+  'transformers.js': {
+    repo_id: 'huggingface/transformers.js',
+    doc_folder: 'docs/source'
+  },
+  'smolagents': {
+    repo_id: 'huggingface/smolagents',
+    doc_folder: 'docs/source/en'
+  },
+  'peft': {
+    repo_id: 'huggingface/peft',
+    doc_folder: 'docs/source'
+  },
+  'trl': {
+    repo_id: 'huggingface/trl',
+    doc_folder: 'docs/source'
+  },
+  'bitsandbytes': {
+    repo_id: 'bitsandbytes-foundation/bitsandbytes',
+    doc_folder: 'docs/source'
+  },
+  'lerobot': {
+    repo_id: 'huggingface/lerobot',
+    doc_folder: 'docs/source'
+  },
+  'timm': {
+    repo_id: 'huggingface/pytorch-image-models',
+    doc_folder: 'hfdocs/source'
+  },
+  'inference-providers': {
+    repo_id: 'huggingface/hub-docs',
+    doc_folder: 'docs/inference-providers'
+  },
+  'safetensors': {
+    repo_id: 'huggingface/safetensors',
+    doc_folder: 'docs/source'
+  },
+  'inference-endpoints': {
+    repo_id: 'huggingface/hf-endpoints-documentation',
+    doc_folder: 'docs/source'
+  },
+  'dataset-viewer': {
+    repo_id: 'huggingface/dataset-viewer',
+    doc_folder: 'docs/source'
+  }
+};

package/src/docs-search/docs-semantic-search.test.ts

@@ -0,0 +1,365 @@
+import { describe, it, expect, vi, beforeEach } from 'vitest';
+import { DocSearchTool } from './docs-semantic-search.js';
+import { DOC_FETCH_CONFIG } from './doc-fetch.js';
+
+// Mock the fetch function
+const mockFetch = vi.fn();
+global.fetch = mockFetch as typeof fetch;
+
+describe('DocSearchTool', () => {
+	let docSearchTool: DocSearchTool;
+
+	beforeEach(() => {
+		docSearchTool = new DocSearchTool();
+		vi.clearAllMocks();
+	});
+
+	describe('search', () => {
+		it('should return error when query is too short', async () => {
+			await expect(docSearchTool.search({ query: 'ab' })).rejects.toThrow();
+		});
+
+		it('should return no results message when API returns empty array', async () => {
+			mockFetch.mockResolvedValueOnce({
+				ok: true,
+				json: () => Promise.resolve([]),
+			});
+
+			const result = await docSearchTool.search({ query: 'nonexistent' });
+			expect(result).toBe(`No documentation found for query 'nonexistent'`);
+		});
+
+		it('should return no results message with product filter', async () => {
+			mockFetch.mockResolvedValueOnce({
+				ok: true,
+				json: () => Promise.resolve([]),
+			});
+
+			const result = await docSearchTool.search({ query: 'nonexistent', product: 'hub' });
+			expect(result).toBe(`No documentation found for query 'nonexistent' in product 'hub'`);
+		});
+
+		it('should format results grouped by product and page', async () => {
+			const sampleResults = [
+				{
+					text: 'Download a comprehensive CSV file containing analytics',
+					product: 'hub',
+					heading1: 'Analytics',
+					source_page_url: 'https://huggingface.co/docs/hub/enterprise-hub-analytics#export-analytics-as-csv',
+					source_page_title: 'Enterprise-hub-analytics',
+					heading2: 'Export Analytics as CSV',
+				},
+				{
+					text: 'View analytics for your repositories',
+					product: 'hub',
+					heading1: 'Analytics',
+					source_page_url: 'https://huggingface.co/docs/hub/enterprise-hub-analytics#export-analytics-as-csv',
+					source_page_title: 'Enterprise-hub-analytics',
+					heading2: 'View Analytics',
+				},
+				{
+					text: 'In this quickstart, you will learn how to use the dataset viewer REST API',
+					product: 'dataset-viewer',
+					heading1: 'Quickstart',
+					source_page_url: 'https://huggingface.co/docs/dataset-viewer/quick_start#quickstart',
+					source_page_title: 'Quick start',
+				},
+			];
+
+			mockFetch.mockResolvedValueOnce({
+				ok: true,
+				json: () => Promise.resolve(sampleResults),
+			});
+
+			const result = await docSearchTool.search({ query: 'analytics' });
+
+			// Check header
+			expect(result).toContain('# Documentation Library Search Results for "analytics"');
+			expect(result).toContain('Found 3 results');
+
+			// Check product grouping - hub should come before dataset-viewer (hub has 2 results, dataset-viewer has 1)
+			const hubIndex = result.indexOf('## Results for Product: hub');
+			const datasetViewerIndex = result.indexOf('## Results for Product: dataset-viewer');
+			expect(hubIndex).toBeLessThan(datasetViewerIndex);
+			expect(hubIndex).toBeGreaterThan(-1);
+			expect(datasetViewerIndex).toBeGreaterThan(-1);
+
+			// Check that result counts are shown
+			expect(result).toContain('## Results for Product: hub (2 results)');
+			expect(result).toContain('## Results for Product: dataset-viewer (1 results)');
+
+			// Check page links (without anchors)
+			expect(result).toContain(
+				'### Results from [Analytics](https://huggingface.co/docs/hub/enterprise-hub-analytics)'
+			);
+			expect(result).toContain('### Results from [Quickstart](https://huggingface.co/docs/dataset-viewer/quick_start)');
+
+			// Check excerpts with heading2
+			expect(result).toContain('#### Excerpt from the "Export Analytics as CSV" section');
+			expect(result).toContain('#### Excerpt from the "View Analytics" section');
+
+			// Check excerpt content appears as plain text
+			expect(result).toContain('Download a comprehensive CSV file containing analytics');
+			expect(result).toContain('View analytics for your repositories');
+			expect(result).toContain('In this quickstart, you will learn how to use the dataset viewer REST API');
+
+			// Check footer
+			expect(result).toContain('Use the "' + DOC_FETCH_CONFIG.name + '" tool to fetch a document from the library.');
+		});
+
+		it('should handle results without heading2', async () => {
+			const sampleResults = [
+				{
+					text: 'This is a simple text without heading2',
+					product: 'transformers',
+					heading1: 'Introduction',
+					source_page_url: 'https://huggingface.co/docs/transformers/index',
+					source_page_title: 'Transformers',
+				},
+			];
+
+			mockFetch.mockResolvedValueOnce({
+				ok: true,
+				json: () => Promise.resolve(sampleResults),
+			});
+
+			const result = await docSearchTool.search({ query: 'transformers' });
+
+			// Should not contain "Excerpt from" when heading2 is missing
+			expect(result).not.toContain('#### Excerpt from');
+			expect(result).toContain('This is a simple text without heading2');
+		});
+
+		it('should properly escape markdown special characters', async () => {
+			const sampleResults = [
+				{
+					text: 'Text with [brackets] and *asterisks* and _underscores_',
+					product: 'hub',
+					heading1: 'Special * Characters',
+					source_page_url: 'https://huggingface.co/docs/hub/test',
+					source_page_title: 'Test',
+					heading2: 'Section with [brackets]',
+				},
+			];
+
+			mockFetch.mockResolvedValueOnce({
+				ok: true,
+				json: () => Promise.resolve(sampleResults),
+			});
+
+			const result = await docSearchTool.search({ query: 'special' });
+
+			// Check that special characters are escaped in headings and page titles
+			expect(result).toContain('Special \\* Characters');
+			// Note: heading2 appears in header text, but brackets don't get escaped
+			expect(result).toContain('#### Excerpt from the "Section with [brackets]" section');
+		});
+
+		it('should clean HTML tags from text', async () => {
+			const sampleResults = [
+				{
+					text: 'Text with <div class="test">HTML tags</div> and <img src="test.png" alt="image"/>',
+					product: 'hub',
+					heading1: 'HTML Test',
+					source_page_url: 'https://huggingface.co/docs/hub/html-test',
+					source_page_title: 'HTML Test',
+				},
+			];
+
+			mockFetch.mockResolvedValueOnce({
+				ok: true,
+				json: () => Promise.resolve(sampleResults),
+			});
+
+			const result = await docSearchTool.search({ query: 'html' });
+
+			// HTML tags should be removed
+			expect(result).toContain('Text with HTML tags and');
+			expect(result).not.toContain('<div');
+			expect(result).not.toContain('<img');
+		});
+
+		it('should sort multiple products and pages correctly by count', async () => {
+			const sampleResults = [
+				// Hub has 3 results (should be first)
+				{
+					text: 'Result from hub page 1 - first',
+					product: 'hub',
+					heading1: 'Page 1',
+					source_page_url: 'https://huggingface.co/docs/hub/page1',
+					source_page_title: 'Page 1',
+				},
+				{
+					text: 'Result from hub page 1 - second',
+					product: 'hub',
+					heading1: 'Page 1',
+					source_page_url: 'https://huggingface.co/docs/hub/page1',
+					source_page_title: 'Page 1',
+				},
+				{
+					text: 'Result from hub page 2',
+					product: 'hub',
+					heading1: 'Page 2',
+					source_page_url: 'https://huggingface.co/docs/hub/page2',
+					source_page_title: 'Page 2',
+				},
+				// Transformers has 1 result (should be second)
+				{
+					text: 'Result from transformers',
+					product: 'transformers',
+					heading1: 'Transformers Page',
+					source_page_url: 'https://huggingface.co/docs/transformers/page1',
+					source_page_title: 'Transformers',
+				},
+				// Datasets has 1 result (should be third)
+				{
+					text: 'Result from datasets',
+					product: 'datasets',
+					heading1: 'Datasets Page',
+					source_page_url: 'https://huggingface.co/docs/datasets/page1',
+					source_page_title: 'Datasets',
+				},
+			];
+
+			mockFetch.mockResolvedValueOnce({
+				ok: true,
+				json: () => Promise.resolve(sampleResults),
+			});
+
+			const result = await docSearchTool.search({ query: 'test' });
+
+			// Check product order by count: hub (3) > transformers (1) = datasets (1)
+			const hubIndex = result.indexOf('## Results for Product: hub');
+			const transformersIndex = result.indexOf('## Results for Product: transformers');
+			const datasetsIndex = result.indexOf('## Results for Product: datasets');
+
+			expect(hubIndex).toBeLessThan(transformersIndex);
+			expect(hubIndex).toBeLessThan(datasetsIndex);
+
+			// Check that hub shows total count
+			expect(result).toContain('## Results for Product: hub (3 results)');
+
+			// Check page order within hub product: page1 (2 results) should come before page2 (1 result)
+			const page1Index = result.indexOf('https://huggingface.co/docs/hub/page1');
+			const page2Index = result.indexOf('https://huggingface.co/docs/hub/page2');
+			expect(page1Index).toBeLessThan(page2Index);
+
+			// Check that page1 shows its multiple results count
+			expect(result).toContain('### Results from [Page 1](https://huggingface.co/docs/hub/page1) (2 results)');
+		});
+
+		it('should include product filter in API call when provided', async () => {
+			mockFetch.mockResolvedValueOnce({
+				ok: true,
+				json: () => Promise.resolve([]),
+			});
+
+			await docSearchTool.search({ query: 'test', product: 'hub' });
+
+			expect(mockFetch).toHaveBeenCalledWith(expect.stringContaining('q=test&product=hub'), expect.any(Object));
+		});
+
+		it('should group results from the same page with different anchors together', async () => {
+			const sampleResults = [
+				{
+					text: 'First result from section 1',
+					product: 'hub',
+					heading1: 'Analytics',
+					source_page_url: 'https://huggingface.co/docs/hub/analytics#section1',
+					source_page_title: 'Analytics Page',
+					heading2: 'Section 1',
+				},
+				{
+					text: 'Second result from section 2',
+					product: 'hub',
+					heading1: 'Analytics',
+					source_page_url: 'https://huggingface.co/docs/hub/analytics#section2',
+					source_page_title: 'Analytics Page',
+					heading2: 'Section 2',
+				},
+				{
+					text: 'Third result from section 3',
+					product: 'hub',
+					heading1: 'Analytics',
+					source_page_url: 'https://huggingface.co/docs/hub/analytics#section3',
+					source_page_title: 'Analytics Page',
+					heading2: 'Section 3',
+				},
+			];
+
+			mockFetch.mockResolvedValueOnce({
+				ok: true,
+				json: () => Promise.resolve(sampleResults),
+			});
+
+			const result = await docSearchTool.search({ query: 'analytics' });
+
+			// All three results should be grouped under one page heading (without anchor)
+			expect(result).toContain('### Results from [Analytics](https://huggingface.co/docs/hub/analytics) (3 results)');
+
+			// All three excerpts should appear under the same page
+			expect(result).toContain('First result from section 1');
+			expect(result).toContain('Second result from section 2');
+			expect(result).toContain('Third result from section 3');
+
+			// There should only be one "Results from" heading for this page
+			const resultsFromCount = (result.match(/### Results from/g) || []).length;
+			expect(resultsFromCount).toBe(1);
+		});
+
+		it('should handle API errors gracefully', async () => {
+			mockFetch.mockRejectedValueOnce(new Error('Network error'));
+
+			await expect(docSearchTool.search({ query: 'test' })).rejects.toThrow('Failed to search documentation:');
+		});
+	});
+
+	describe('groupResults', () => {
+		it('should group results by product and page URL', async () => {
+			const sampleResults = [
+				{
+					text: 'Result 1',
+					product: 'hub',
+					heading1: 'Page 1',
+					source_page_url: 'https://example.com/page1',
+					source_page_title: 'Page 1',
+				},
+				{
+					text: 'Result 2',
+					product: 'hub',
+					heading1: 'Page 1',
+					source_page_url: 'https://example.com/page1',
+					source_page_title: 'Page 1',
+				},
+				{
+					text: 'Result 3',
+					product: 'transformers',
+					heading1: 'Page 2',
+					source_page_url: 'https://example.com/page2',
+					source_page_title: 'Page 2',
+				},
+			];
+
+			mockFetch.mockResolvedValueOnce({
+				ok: true,
+				json: () => Promise.resolve(sampleResults),
+			});
+
+			const result = await docSearchTool.search({ query: 'test' });
+
+			// Verify grouping structure in output
+			expect(result).toContain('## Results for Product: hub');
+			expect(result).toContain('## Results for Product: transformers');
+
+			// Verify that both results from the same page are together
+			const result1Index = result.indexOf('Result 1');
+			const result2Index = result.indexOf('Result 2');
+			const result3Index = result.indexOf('Result 3');
+
+			// Results 1 and 2 should be close together (same page)
+			expect(Math.abs(result2Index - result1Index)).toBeLessThan(100);
+			// Result 3 should be further away (different product)
+			expect(Math.abs(result3Index - result1Index)).toBeGreaterThan(50);
+		});
+	});
+});

package/src/docs-search/docs-semantic-search.ts

@@ -0,0 +1,244 @@
+import { z } from 'zod';
+import { HfApiCall } from '../hf-api-call.js';
+import { escapeMarkdown } from '../utilities.js';
+import { DOC_FETCH_CONFIG } from './doc-fetch.js';
+
+export const DOCS_SEMANTIC_SEARCH_CONFIG = {
+	name: 'hf_doc_search',
+	description: 'Search the Hugging Face documentation library. Returns excerpts grouped by Product and Document.',
+	schema: z.object({
+		query: z
+			.string()
+			.min(3, 'Supply at least one search term')
+			.max(200, 'Query too long')
+			.describe('Semantic search query'),
+		product: z
+			.string()
+			.optional()
+			.describe(
+				'Filter by Product (e.g., "hub", "dataset-viewer", "transformers"). Supply when known for focused results'
+			),
+	}),
+	annotations: {
+		title: 'Hugging Face Documentation Library Search',
+		destructiveHint: false,
+		readOnlyHint: true,
+		openWorldHint: true,
+	},
+} as const;
+
+export type DocSearchParams = z.infer<typeof DOCS_SEMANTIC_SEARCH_CONFIG.schema>;
+
+interface DocSearchResult {
+	text: string;
+	product: string;
+	heading1: string;
+	source_page_url: string;
+	source_page_title: string;
+	heading2?: string;
+}
+
+interface DocSearchApiParams {
+	q: string;
+	product?: string;
+}
+
+/**
+ * Use the Hugging Face Semantic Document Search API
+ */
+export class DocSearchTool extends HfApiCall<DocSearchApiParams, DocSearchResult[]> {
+	/**
+	 * @param apiUrl The URL of the Hugging Face document search API
+	 * @param hfToken Optional Hugging Face token for API access
+	 */
+	constructor(hfToken?: string, apiUrl = 'https://hf.co/api/docs/search') {
+		super(apiUrl, hfToken);
+	}
+
+	/**
+	 * @param query Search query string (e.g. "rate limits", "analytics")
+	 * @param product Optional product filter
+	 */
+	async search(params: DocSearchParams): Promise<string> {
+		try {
+			if (!params.query) return 'No query provided';
+
+			const apiParams: DocSearchApiParams = { q: params.query.toLowerCase() };
+			if (params.product) {
+				apiParams.product = params.product;
+			}
+
+			const results = await this.callApi<DocSearchResult[]>(apiParams);
+
+			if (results.length === 0) {
+				return params.product
+					? `No documentation found for query '${params.query}' in product '${params.product}'`
+					: `No documentation found for query '${params.query}'`;
+			}
+
+			return formatSearchResults(params.query, results, params.product);
+		} catch (error) {
+			if (error instanceof Error) {
+				throw new Error(`Failed to search documentation: ${error.message}`);
+			}
+			throw error;
+		}
+	}
+}
+
+/**
+ * Group results by product and source page URL
+ */
+function groupResults(results: DocSearchResult[]): Map<string, Map<string, DocSearchResult[]>> {
+	const grouped = new Map<string, Map<string, DocSearchResult[]>>();
+
+	for (const result of results) {
+		if (!grouped.has(result.product)) {
+			grouped.set(result.product, new Map());
+		}
+
+		const productGroup = grouped.get(result.product);
+		if (!productGroup) continue;
+
+		// Strip the anchor (#section) from the URL for grouping purposes
+		const baseUrl = result.source_page_url.split('#')[0] || result.source_page_url;
+
+		if (!productGroup.has(baseUrl)) {
+			productGroup.set(baseUrl, []);
+		}
+
+		const pageResults = productGroup.get(baseUrl);
+		if (pageResults) {
+			pageResults.push(result);
+		}
+	}
+
+	return grouped;
+}
+
+/**
+ * Group page results by section (heading2)
+ */
+function groupBySection(pageResults: DocSearchResult[]): Map<string | undefined, DocSearchResult[]> {
+	const sectionGroups = new Map<string | undefined, DocSearchResult[]>();
+
+	for (const result of pageResults) {
+		const section = result.heading2;
+		if (!sectionGroups.has(section)) {
+			sectionGroups.set(section, []);
+		}
+		const sectionResults = sectionGroups.get(section);
+		if (sectionResults) {
+			sectionResults.push(result);
+		}
+	}
+
+	return sectionGroups;
+}
+
+/**
+ * Format excerpts from a section
+ */
+function formatSectionExcerpts(section: string | undefined, results: DocSearchResult[]): string {
+	const lines: string[] = [];
+
+	// Add section heading if present
+	if (section) {
+		if (results.length > 1) {
+			lines.push(`#### Excerpts from the "${escapeMarkdown(section)}" section`);
+		} else {
+			lines.push(`#### Excerpt from the "${escapeMarkdown(section)}" section`);
+		}
+		lines.push('');
+	}
+
+	// Add all excerpts from this section
+	for (const result of results) {
+		// Clean up the text - remove HTML tags if any
+		const cleanText = result.text
+			.replace(/<[^>]*>/g, '')
+			.replace(/\n\s*\n/g, '\n')
+			.trim();
+
+		lines.push(cleanText);
+		lines.push('');
+	}
+
+	return lines.join('\n');
+}
+
+/**
+ * Format search results grouped by product and page
+ */
+function formatSearchResults(query: string, results: DocSearchResult[], productFilter?: string): string {
+	const lines: string[] = [];
+
+	// Header
+	const filterText = productFilter ? ` (filtered by product: ${productFilter})` : '';
+	lines.push(`# Documentation Library Search Results for "${escapeMarkdown(query)}"${filterText}`);
+	lines.push('');
+	lines.push(`Found ${results.length} results`);
+	lines.push('');
+
+	// Group results
+	const grouped = groupResults(results);
+
+	// Sort products by count (most hits first)
+	const sortedProducts = Array.from(grouped.keys()).sort((a, b) => {
+		const productGroupA = grouped.get(a);
+		const productGroupB = grouped.get(b);
+		if (!productGroupA || !productGroupB) return 0;
+
+		const countA = Array.from(productGroupA.values()).reduce((sum, arr) => sum + arr.length, 0);
+		const countB = Array.from(productGroupB.values()).reduce((sum, arr) => sum + arr.length, 0);
+		return countB - countA; // Descending order
+	});
+
+	for (const product of sortedProducts) {
+		const productGroup = grouped.get(product);
+		if (!productGroup) continue;
+
+		const totalProductHits = Array.from(productGroup.values()).reduce((sum, arr) => sum + arr.length, 0);
+		lines.push(`## Results for Product: ${escapeMarkdown(product)} (${totalProductHits} results)`);
+		lines.push('');
+
+		// Sort URLs within each product by count (most hits first)
+		const sortedUrls = Array.from(productGroup.keys()).sort((a, b) => {
+			const pageResultsA = productGroup.get(a);
+			const pageResultsB = productGroup.get(b);
+			if (!pageResultsA || !pageResultsB) return 0;
+			return pageResultsB.length - pageResultsA.length;
+		});
+
+		for (const url of sortedUrls) {
+			const pageResults = productGroup.get(url);
+			if (!pageResults || pageResults.length === 0) continue;
+			const firstResult = pageResults[0];
+
+			// Skip if no results (shouldn't happen but TypeScript safety)
+			if (!firstResult) continue;
+
+			// Page header with link and hit count
+			const pageTitle = firstResult.heading1 || firstResult.source_page_title;
+			const hitCount = pageResults.length > 1 ? ` (${pageResults.length} results)` : '';
+			// Use the base URL (without anchor) for the page link
+			lines.push(`### Results from [${escapeMarkdown(pageTitle)}](${url})${hitCount}`);
+			lines.push('');
+
+			// Group results by section and format them
+			const sectionGroups = groupBySection(pageResults);
+
+			// Format each section's excerpts
+			for (const [section, sectionResults] of sectionGroups) {
+				lines.push(formatSectionExcerpts(section, sectionResults));
+			}
+		}
+	}
+
+	// Add suggestion to use doc fetch tool
+	lines.push('---');
+	lines.push('');
+	lines.push(`Use the "${DOC_FETCH_CONFIG.name}" tool to fetch a document from the library.`);
+
+	return lines.join('\n');
+}

package/src/index.ts

@@ -13,6 +13,8 @@ export * from './space-info.js';
 export * from './space-files.js';
 export * from './user-summary.js';
 export * from './paper-summary.js';
+export * from './docs-search/docs-semantic-search.js';
+export * from './docs-search/doc-fetch.js';
 
 // Export tool IDs for external use - these are the canonical tool identifiers
 export * from './tool-ids.js';