@llmindset/hf-mcp 0.1.20 → 0.1.22

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

@@ -1,11 +1,16 @@
 import { z } from 'zod';
 import { HfApiCall } from '../hf-api-call.js';
-import { escapeMarkdown } from '../utilities.js';
+import { escapeMarkdown, estimateTokens } from '../utilities.js';
 import { DOC_FETCH_CONFIG } from './doc-fetch.js';
 
+/** token estimation. initial results for "how to load a image to image model in transformers" returned
+ * 121973 characters (36711 anthropic tokens) */
+
 export const DOCS_SEMANTIC_SEARCH_CONFIG = {
 	name: 'hf_doc_search',
-	description: 'Search the Hugging Face documentation library. Returns excerpts grouped by Product and Document.',
+	description:
+		'Search the Hugging Face documentation library. Use this for the most up-to-date information ' +
+		'Returns excerpts grouped by Product and Document.',
 	schema: z.object({
 		query: z
 			.string()
@@ -43,16 +48,24 @@ interface DocSearchApiParams {
 	product?: string;
 }
 
+// Token budget defaults
+const DEFAULT_TOKEN_BUDGET = 12500;
+const TRUNCATE_EXCERPT_LENGTH = 400; // chars for truncated excerpts
+
 /**
  * Use the Hugging Face Semantic Document Search API
  */
 export class DocSearchTool extends HfApiCall<DocSearchApiParams, DocSearchResult[]> {
+	private tokenBudget: number;
+
 	/**
-	 * @param apiUrl The URL of the Hugging Face document search API
 	 * @param hfToken Optional Hugging Face token for API access
+	 * @param apiUrl The URL of the Hugging Face document search API
+	 * @param tokenBudget Maximum number of tokens to return
 	 */
-	constructor(hfToken?: string, apiUrl = 'https://hf.co/api/docs/search') {
+	constructor(hfToken?: string, apiUrl = 'https://hf.co/api/docs/search', tokenBudget = DEFAULT_TOKEN_BUDGET) {
 		super(apiUrl, hfToken);
+		this.tokenBudget = tokenBudget;
 	}
 
 	/**
@@ -76,7 +89,7 @@ export class DocSearchTool extends HfApiCall<DocSearchApiParams, DocSearchResult
 					: `No documentation found for query '${params.query}'`;
 			}
 
-			return formatSearchResults(params.query, results, params.product);
+			return formatSearchResults(params.query, results, params.product, this.tokenBudget);
 		} catch (error) {
 			if (error instanceof Error) {
 				throw new Error(`Failed to search documentation: ${error.message}`);
@@ -139,70 +152,110 @@ function groupBySection(pageResults: DocSearchResult[]): Map<string | undefined,
 /**
  * Format excerpts from a section
  */
-function formatSectionExcerpts(section: string | undefined, results: DocSearchResult[]): string {
+function formatSectionExcerpts(
+	section: string | undefined,
+	results: DocSearchResult[],
+	useTruncatedMode: boolean,
+	hasAlreadyShownTruncation: boolean
+): { text: string; tokensUsed: number; wasContentTruncated: boolean } {
 	const lines: string[] = [];
+	let tokensUsed = 0;
+	let wasContentTruncated = false;
 
-	// Add section heading if present
+	// Add section heading if we have one
 	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('');
+		const heading =
+			results.length > 1
+				? `#### Excerpts from the "${escapeMarkdown(section)}" section`
+				: `#### Excerpt from the "${escapeMarkdown(section)}" section`;
+
+		lines.push(heading, '');
+		tokensUsed += estimateTokens(heading + '\n\n');
 	}
 
-	// Add all excerpts from this section
 	for (const result of results) {
-		// Clean up the text - remove HTML tags if any
-		const cleanText = result.text
+		let cleanText = result.text
 			.replace(/<[^>]*>/g, '')
 			.replace(/\n\s*\n/g, '\n')
 			.trim();
 
-		lines.push(cleanText);
-		lines.push('');
+		// Truncate if in truncated mode and we haven't shown the message yet
+		if (useTruncatedMode && cleanText.length > TRUNCATE_EXCERPT_LENGTH && !hasAlreadyShownTruncation) {
+			cleanText =
+				cleanText.substring(0, TRUNCATE_EXCERPT_LENGTH) +
+				`...\n\n*[Content truncated - use ${DOC_FETCH_CONFIG.name} for full text or narrow search terms]*`;
+			wasContentTruncated = true;
+		}
+
+		lines.push(cleanText, '');
+		tokensUsed += estimateTokens(cleanText + '\n\n');
 	}
 
-	return lines.join('\n');
+	// Remove trailing empty line
+	if (lines.length > 0 && lines[lines.length - 1] === '') {
+		lines.pop();
+	}
+
+	return { text: lines.join('\n'), tokensUsed, wasContentTruncated };
 }
 
 /**
- * Format search results grouped by product and page
+ * Format search results with simple token budget management
  */
-function formatSearchResults(query: string, results: DocSearchResult[], productFilter?: string): string {
+function formatSearchResults(
+	query: string,
+	results: DocSearchResult[],
+	productFilter?: string,
+	tokenBudget = DEFAULT_TOKEN_BUDGET
+): string {
 	const lines: string[] = [];
+	let hasShownTruncationMessage = false;
 
 	// 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('');
+	const header = `# Documentation Library Search Results for "${escapeMarkdown(query)}"${filterText}\n\nFound ${results.length} results\n`;
+	lines.push(header);
 
-	// Group results
+	// Group and sort 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
+		return countB - countA;
 	});
 
+	const linkOnlyResults: Array<{ product: string; url: string; title: string; count: number }> = [];
+
 	for (const product of sortedProducts) {
 		const productGroup = grouped.get(product);
 		if (!productGroup) continue;
 
+		// Check current size before adding anything
+		const currentText = lines.join('\n');
+		if (estimateTokens(currentText) > tokenBudget) {
+			// Over budget - add remaining products to links
+			for (const url of productGroup.keys()) {
+				const pageResults = productGroup.get(url);
+				if (!pageResults?.[0]) continue;
+				linkOnlyResults.push({
+					product,
+					url,
+					title: pageResults[0].heading1 || pageResults[0].source_page_title,
+					count: pageResults.length,
+				});
+			}
+			continue;
+		}
+
+		// Add product header
 		const totalProductHits = Array.from(productGroup.values()).reduce((sum, arr) => sum + arr.length, 0);
-		lines.push(`## Results for Product: ${escapeMarkdown(product)} (${totalProductHits} results)`);
-		lines.push('');
+		const productHeader = `## Results for Product: ${escapeMarkdown(product)} (${totalProductHits} results)\n`;
+		lines.push(productHeader);
 
-		// Sort URLs within each product by count (most hits first)
+		// Sort pages by hit count
 		const sortedUrls = Array.from(productGroup.keys()).sort((a, b) => {
 			const pageResultsA = productGroup.get(a);
 			const pageResultsB = productGroup.get(b);
@@ -212,32 +265,50 @@ function formatSearchResults(query: string, results: DocSearchResult[], productF
 
 		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;
+			if (!pageResults?.[0]) continue;
+
+			const pageTitle = pageResults[0].heading1 || pageResults[0].source_page_title;
+			
+			// Check if we're over budget - if so, add remaining pages to links
+			const currentText = lines.join('\n');
+			if (estimateTokens(currentText) > tokenBudget) {
+				linkOnlyResults.push({ product, url, title: pageTitle, count: pageResults.length });
+				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('');
+			const pageHeader = `### Results from [${escapeMarkdown(pageTitle)}](${url})${hitCount}\n`;
+			lines.push(pageHeader);
 
-			// Group results by section and format them
+			// Add all sections for this page
 			const sectionGroups = groupBySection(pageResults);
-
-			// Format each section's excerpts
 			for (const [section, sectionResults] of sectionGroups) {
-				lines.push(formatSectionExcerpts(section, sectionResults));
+				const currentTokens = estimateTokens(lines.join('\n'));
+				const useTruncatedMode = currentTokens > tokenBudget * 0.7;
+
+				const result = formatSectionExcerpts(section, sectionResults, useTruncatedMode, hasShownTruncationMessage);
+				
+				if (result.text.trim()) {
+					lines.push(result.text);
+					if (result.wasContentTruncated) {
+						hasShownTruncationMessage = true;
+					}
+				}
 			}
 		}
 	}
 
-	// Add suggestion to use doc fetch tool
-	lines.push('---');
-	lines.push('');
+	// Add link-only results
+	if (linkOnlyResults.length > 0) {
+		lines.push(`\n## Further results were found in:\n`);
+		for (const linkResult of linkOnlyResults) {
+			const hitText = linkResult.count > 1 ? ` (${linkResult.count} results)` : '';
+			lines.push(`- [${escapeMarkdown(linkResult.title)}](${linkResult.url})${hitText} *(${linkResult.product})*`);
+		}
+		lines.push('');
+	}
+
+	lines.push('---\n');
 	lines.push(`Use the "${DOC_FETCH_CONFIG.name}" tool to fetch a document from the library.`);
 
 	return lines.join('\n');

package/src/utilities.ts

@@ -62,3 +62,16 @@ export function escapeMarkdown(text: string): string {
 		.replace(/>/g, '\\>')
 		.replace(/#/g, '\\#');
 }
+
+// Token estimation constants
+const CHARS_PER_TOKEN = 3.3; // based on anthropic tokenizer for "how to load a image to image model in transformers"
+//  data: 121973 chars = 36711 tokens
+
+/**
+ * Simple token estimation based on character count
+ * @param text The text to estimate tokens for
+ * @returns Estimated number of tokens
+ */
+export function estimateTokens(text: string): number {
+	return Math.ceil(text.length / CHARS_PER_TOKEN);
+}

package/dist/docs-search/doc-mappings.d.ts

@@ -1,7 +0,0 @@
-interface DocMapping {
-    repo_id: string;
-    doc_folder: string;
-}
-export declare const DOC_MAPPINGS: Record<string, DocMapping>;
-export {};
-//# sourceMappingURL=doc-mappings.d.ts.map

package/dist/docs-search/doc-mappings.d.ts.map

@@ -1 +0,0 @@
-{"version":3,"file":"doc-mappings.d.ts","sourceRoot":"","sources":["../../src/docs-search/doc-mappings.ts"],"names":[],"mappings":"AAAA,UAAU,UAAU;IAClB,OAAO,EAAE,MAAM,CAAC;IAChB,UAAU,EAAE,MAAM,CAAC;CACpB;AAED,eAAO,MAAM,YAAY,EAAE,MAAM,CAAC,MAAM,EAAE,UAAU,CAyEnD,CAAC"}

package/dist/docs-search/doc-mappings.js

@@ -1,75 +0,0 @@
-export const DOC_MAPPINGS = {
-    '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'
-    }
-};
-//# sourceMappingURL=doc-mappings.js.map

package/dist/docs-search/doc-mappings.js.map

@@ -1 +0,0 @@
-{"version":3,"file":"doc-mappings.js","sourceRoot":"","sources":["../../src/docs-search/doc-mappings.ts"],"names":[],"mappings":"AAKA,MAAM,CAAC,MAAM,YAAY,GAA+B;IACtD,YAAY,EAAE;QACZ,OAAO,EAAE,wBAAwB;QACjC,UAAU,EAAE,yBAAyB;KACtC;IACD,WAAW,EAAE;QACX,OAAO,EAAE,uBAAuB;QAChC,UAAU,EAAE,gBAAgB;KAC7B;IACD,YAAY,EAAE;QACZ,OAAO,EAAE,wBAAwB;QACjC,UAAU,EAAE,aAAa;KAC1B;IACD,iBAAiB,EAAE;QACjB,OAAO,EAAE,6BAA6B;QACtC,UAAU,EAAE,gBAAgB;KAC7B;IACD,cAAc,EAAE;QACd,OAAO,EAAE,0BAA0B;QACnC,UAAU,EAAE,gBAAgB;KAC7B;IACD,KAAK,EAAE;QACL,OAAO,EAAE,sBAAsB;QAC/B,UAAU,EAAE,UAAU;KACvB;IACD,gBAAgB,EAAE;QAChB,OAAO,EAAE,4BAA4B;QACrC,UAAU,EAAE,MAAM;KACnB;IACD,iBAAiB,EAAE;QACjB,OAAO,EAAE,6BAA6B;QACtC,UAAU,EAAE,aAAa;KAC1B;IACD,YAAY,EAAE;QACZ,OAAO,EAAE,wBAAwB;QACjC,UAAU,EAAE,gBAAgB;KAC7B;IACD,MAAM,EAAE;QACN,OAAO,EAAE,kBAAkB;QAC3B,UAAU,EAAE,aAAa;KAC1B;IACD,KAAK,EAAE;QACL,OAAO,EAAE,iBAAiB;QAC1B,UAAU,EAAE,aAAa;KAC1B;IACD,cAAc,EAAE;QACd,OAAO,EAAE,sCAAsC;QAC/C,UAAU,EAAE,aAAa;KAC1B;IACD,SAAS,EAAE;QACT,OAAO,EAAE,qBAAqB;QAC9B,UAAU,EAAE,aAAa;KAC1B;IACD,MAAM,EAAE;QACN,OAAO,EAAE,kCAAkC;QAC3C,UAAU,EAAE,eAAe;KAC5B;IACD,qBAAqB,EAAE;QACrB,OAAO,EAAE,sBAAsB;QAC/B,UAAU,EAAE,0BAA0B;KACvC;IACD,aAAa,EAAE;QACb,OAAO,EAAE,yBAAyB;QAClC,UAAU,EAAE,aAAa;KAC1B;IACD,qBAAqB,EAAE;QACrB,OAAO,EAAE,wCAAwC;QACjD,UAAU,EAAE,aAAa;KAC1B;IACD,gBAAgB,EAAE;QAChB,OAAO,EAAE,4BAA4B;QACrC,UAAU,EAAE,aAAa;KAC1B;CACF,CAAC"}

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

@@ -1,79 +0,0 @@
-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'
-  }
-};