feat: vector collapse

Author: ThomasRooney

.changeset/vector-collapse-dedup.md

@@ -0,0 +1,7 @@
+---
+"@speakeasy-api/docs-mcp-core": minor
+"@speakeasy-api/docs-mcp-cli": minor
+"@speakeasy-api/docs-mcp-server": minor
+---
+
+Add `taxonomy` manifest field with `vector_collapse` option for collapsing content-equivalent search results across variant axes (e.g. the same API operation documented in multiple SDK languages). At search time, results sharing the same content identity are collapsed to the highest-scoring variant. On a realistic 30MB multi-language corpus this improved facet precision by 27%, MRR@5 by 10%, and NDCG@5 by 15%.

packages/cli/src/index.ts

@@ -16,12 +16,14 @@ import {
   embedChunksIncremental,
   loadCache,
   loadChunksFromPreviousIndex,
+  mergeTaxonomyConfigs,
   parseManifestJson,
   resolveFileConfig,
   saveCache,
   type Chunk,
   type EmbedProgressEvent,
   type EmbeddingMetadata,
+  type ManifestTaxonomyFieldConfig,
   type IndexBuildStep,
   type Manifest,
   type PreviousIndexReader
@@ -215,6 +217,8 @@ program
     const cacheSuffix = chunkCacheHits > 0 ? ` (${chunkCacheHits} cached)` : "";
     console.warn(`Chunked ${files.length} files into ${chunks.length.toLocaleString()} chunks${cacheSuffix}`);
 
+    const taxonomyConfig = mergeTaxonomyConfigs(manifestCache.values());
+
     const providerInput: {
       provider: "none" | "hash" | "openai";
       model?: string;
@@ -318,10 +322,18 @@ program
       files,
       options.description,
       embeddingMetadata,
-      sourceCommit
+      sourceCommit,
+      taxonomyConfig
     );
     const metadataKeys = Object.keys(metadata.taxonomy);
 
+    // Warn about taxonomy config keys that don't match any chunk metadata
+    for (const key of Object.keys(taxonomyConfig)) {
+      if (!metadata.taxonomy[key]) {
+        console.warn(`warn: taxonomy config key '${key}' does not match any chunk metadata — this configuration has no effect`);
+      }
+    }
+
     // Close previous index before writing the new one
     previousIndex?.close();
 
@@ -466,11 +478,12 @@ function buildMetadata(
   files: string[],
   corpusDescription: string,
   embedding: EmbeddingMetadata | null,
-  sourceCommit: string | null
+  sourceCommit: string | null,
+  taxonomyConfig: Record<string, ManifestTaxonomyFieldConfig>
 ): {
   metadata_version: string;
   corpus_description: string;
-  taxonomy: Record<string, { description: string; values: string[] }>;
+  taxonomy: Record<string, { description: string; values: string[]; vector_collapse?: boolean }>;
   stats: {
     total_chunks: number;
     total_files: number;
@@ -488,11 +501,13 @@ function buildMetadata(
     }
   }
 
-  const taxonomy: Record<string, { description: string; values: string[] }> = {};
+  const taxonomy: Record<string, { description: string; values: string[]; vector_collapse?: boolean }> = {};
   for (const [key, values] of taxonomyValues.entries()) {
+    const config = taxonomyConfig[key];
     taxonomy[key] = {
       description: `Filter results by ${key}.`,
-      values: [...values].sort((a, b) => a.localeCompare(b))
+      values: [...values].sort((a, b) => a.localeCompare(b)),
+      ...(config?.vector_collapse ? { vector_collapse: true } : {})
     };
   }
 

packages/core/src/lancedb.ts

@@ -9,6 +9,7 @@ import {
 import { decodeSearchCursor, encodeSearchCursor } from "./cursor.js";
 import {
   clampLimit,
+  dedupKey,
   isChunkIdFormat,
   makeSnippet
 } from "./search-common.js";
@@ -48,6 +49,7 @@ export interface OpenLanceDbSearchEngineOptions {
   dbPath: string;
   tableName?: string;
   metadataKeys: string[];
+  collapseKeys?: string[];
   proximityWeight?: number;
   phraseSlop?: number;
   queryEmbeddingProvider?: EmbeddingProvider;
@@ -127,6 +129,7 @@ export async function buildLanceDbIndex(
 export class LanceDbSearchEngine implements SearchEngine {
   private readonly table: Table;
   private readonly metadataKeys: string[];
+  private readonly collapseKeys: string[];
   private readonly proximityWeight: number;
   private readonly phraseSlop: number;
   private readonly queryEmbeddingProvider: EmbeddingProvider | undefined;
@@ -138,13 +141,15 @@ export class LanceDbSearchEngine implements SearchEngine {
     table: Table,
     metadataKeys: string[],
     options: DocsIndexOptions & {
+      collapseKeys?: string[];
       queryEmbeddingProvider?: EmbeddingProvider;
       vectorWeight?: number;
       onWarning?: (message: string) => void;
     } = {}
   ) {
     this.table = table;
     this.metadataKeys = [...metadataKeys];
+    this.collapseKeys = options.collapseKeys ?? [];
     this.proximityWeight = options.proximityWeight ?? 1.25;
     this.phraseSlop = normalizePhraseSlop(options.phraseSlop);
     this.queryEmbeddingProvider = options.queryEmbeddingProvider;
@@ -156,10 +161,14 @@ export class LanceDbSearchEngine implements SearchEngine {
     const db = await connect(options.dbPath);
     const table = await db.openTable(options.tableName ?? DEFAULT_TABLE_NAME);
     const engineOptions: DocsIndexOptions & {
+      collapseKeys?: string[];
       queryEmbeddingProvider?: EmbeddingProvider;
       vectorWeight?: number;
       onWarning?: (message: string) => void;
     } = {};
+    if (options.collapseKeys !== undefined) {
+      engineOptions.collapseKeys = options.collapseKeys;
+    }
     if (options.proximityWeight !== undefined) {
       engineOptions.proximityWeight = options.proximityWeight;
     }
@@ -211,14 +220,21 @@ export class LanceDbSearchEngine implements SearchEngine {
     const phraseWeight = request.rrf_weights?.phrase ?? this.proximityWeight;
     const vecWeight = request.rrf_weights?.vector ?? this.vectorWeight;
     const blended = blendRows(matchRows, phraseRows, vectorRows, phraseWeight, vecWeight, matchWeight);
-    const paged = blended.slice(offset, offset + limit);
+
+    // Collapse content-equivalent results across variant axes (e.g. same
+    // operation documented in multiple SDK languages). Skipped when active
+    // filters already restrict every collapse axis to a single value.
+    const activeCollapseKeys = this.collapseKeys.filter((k) => !filters[k]);
+    const deduped = deduplicateRows(blended, activeCollapseKeys);
+
+    const paged = deduped.slice(offset, offset + limit);
 
     const hits = paged.map((entry) =>
       toSearchHit(entry.row, entry.score, query, this.metadataKeys)
     );
 
     const nextOffset = offset + paged.length;
-    const nextCursor = nextOffset < blended.length
+    const nextCursor = nextOffset < deduped.length
       ? encodeSearchCursor({ offset: nextOffset, limit }, { query, filters })
       : null;
 
@@ -467,6 +483,28 @@ function blendRows(
     });
 }
 
+function deduplicateRows(
+  rows: Array<{ row: ChunkRow; score: number }>,
+  collapseKeys: string[]
+): Array<{ row: ChunkRow; score: number }> {
+  if (collapseKeys.length === 0) return rows;
+
+  const seen = new Set<string>();
+  return rows.filter((entry) => {
+    const key = dedupKey(
+      expectStringField(entry.row, "filepath"),
+      expectStringField(entry.row, "heading"),
+      expectStringField(entry.row, "chunk_id"),
+      (k) => expectStringField(entry.row, k),
+      collapseKeys
+    );
+    if (key === null) return true;
+    if (seen.has(key)) return false;
+    seen.add(key);
+    return true;
+  });
+}
+
 function buildWhereClause(filters: Record<string, string>, taxonomyKeys?: string[]): string | null {
   const clauses: string[] = [];
 

packages/core/src/manifest-schema.ts

@@ -51,6 +51,23 @@ export const ManifestOverrideSchema = z
     "Overrides the default chunking strategy and/or metadata for files matching a glob pattern. Within the overrides array, later matches take precedence."
   );
 
+export const TaxonomyFieldConfigSchema = z
+  .object({
+    vector_collapse: z
+      .boolean()
+      .default(false)
+      .describe(
+        "When true, this taxonomy dimension identifies content variants that are near-identical in vector space (e.g. the same API operation documented in multiple SDK languages). At search time, results sharing the same content identity — determined by normalizing this field's value out of the filepath — are collapsed to the highest-scoring result. Has no effect when a filter for this field is active, since the filter already restricts to a single value."
+      ),
+  })
+  .describe("Configuration for a taxonomy field's search-time behavior.");
+
+export const ManifestTaxonomyConfigSchema = z
+  .record(z.string(), TaxonomyFieldConfigSchema)
+  .describe(
+    "Per-field configuration for taxonomy dimensions. Controls search-time behavior such as cross-language result collapsing."
+  );
+
 export const ManifestSchema = z
   .object({
     version: z
@@ -67,6 +84,8 @@ export const ManifestSchema = z
         "Key-value pairs attached to every chunk produced from this directory tree. Each key becomes a filterable taxonomy dimension exposed as an enum parameter on the search tool."
       )
       .meta({ examples: [{ language: "typescript", scope: "sdk-specific" }] }),
+    taxonomy: ManifestTaxonomyConfigSchema.optional()
+      .meta({ examples: [{ language: { vector_collapse: true } }] }),
     overrides: z
       .array(ManifestOverrideSchema)
       .optional()

packages/core/src/manifest.ts

@@ -1,14 +1,37 @@
 import matter from "gray-matter";
 import picomatch from "picomatch";
+import { ManifestTaxonomyConfigSchema } from "./manifest-schema.js";
 import type {
   ChunkingStrategy,
   Manifest,
   ManifestOverride,
+  ManifestTaxonomyFieldConfig,
   ResolvedFileConfig
 } from "./types.js";
 
 const DEFAULT_STRATEGY: ChunkingStrategy = { chunk_by: "h2" };
 
+/**
+ * Union-merges taxonomy field configs from multiple manifests. If any manifest
+ * sets `vector_collapse: true` for a key, the merged result includes it.
+ */
+export function mergeTaxonomyConfigs(
+  manifests: Iterable<Manifest>
+): Record<string, ManifestTaxonomyFieldConfig> {
+  const merged: Record<string, ManifestTaxonomyFieldConfig> = {};
+
+  for (const manifest of manifests) {
+    if (!manifest.taxonomy) continue;
+    for (const [key, config] of Object.entries(manifest.taxonomy)) {
+      if (config.vector_collapse) {
+        merged[key] = { ...merged[key], vector_collapse: true };
+      }
+    }
+  }
+
+  return merged;
+}
+
 const HTML_HINT_REGEX = /<!--\s*mcp_chunking_hint:\s*(\{[^}]+\})\s*-->/;
 
 export function parseManifest(input: unknown): Manifest {
@@ -30,6 +53,13 @@ export function parseManifest(input: unknown): Manifest {
   if (manifest.metadata) {
     parsed.metadata = parseMetadata(manifest.metadata, "metadata");
   }
+  if (manifest.taxonomy) {
+    try {
+      parsed.taxonomy = ManifestTaxonomyConfigSchema.parse(manifest.taxonomy);
+    } catch (err) {
+      throw new Error(`Invalid taxonomy config: ${err instanceof Error ? err.message : String(err)}`);
+    }
+  }
   if (manifest.overrides) {
     parsed.overrides = parseOverrides(manifest.overrides);
   }

packages/core/src/metadata.ts

@@ -1,6 +1,16 @@
 import semver from "semver";
 import type { CorpusMetadata, EmbeddingMetadata, TaxonomyField } from "./types.js";
 
+/**
+ * Returns taxonomy keys that have `vector_collapse: true`, i.e. dimensions
+ * whose values should be collapsed at search time.
+ */
+export function getCollapseKeys(taxonomy: Record<string, TaxonomyField>): string[] {
+  return Object.entries(taxonomy)
+    .filter(([, field]) => field.vector_collapse === true)
+    .map(([key]) => key);
+}
+
 const LIMITS = {
   maxKeys: 64,
   maxKeyLength: 64,
@@ -81,9 +91,13 @@ function normalizeTaxonomy(value: unknown): Record<string, TaxonomyField> {
     const values = normalizeValues(field.values, key);
     const description = field.description === undefined ? undefined : asTrimmedString(field.description);
 
-    normalized[key] = description
-      ? { description, values }
-      : { values };
+    const vectorCollapse = field.vector_collapse === true ? true : undefined;
+
+    normalized[key] = {
+      ...(description ? { description } : {}),
+      values,
+      ...(vectorCollapse ? { vector_collapse: true } : {})
+    };
   }
 
   return normalized;

packages/core/src/search-common.ts

@@ -48,6 +48,42 @@ export function makeSnippet(content: string, query: string): string {
   return `${prefix}${normalized.slice(start, end).trim()}${suffix}`;
 }
 
+/**
+ * Computes a dedup key for collapsing content-equivalent results across
+ * taxonomy variant axes (e.g. the same operation documented in multiple SDK
+ * languages). Returns null when no collapsing applies.
+ */
+export function dedupKey(
+  filepath: string,
+  heading: string,
+  chunkId: string,
+  getMetadataValue: (key: string) => string,
+  collapseKeys: string[]
+): string | null {
+  if (collapseKeys.length === 0) return null;
+
+  const parts = filepath.split("/");
+  let anyNormalized = false;
+
+  for (const key of collapseKeys) {
+    const value = getMetadataValue(key);
+    if (!value) return null;
+
+    const idx = parts.indexOf(value);
+    if (idx >= 0) {
+      parts[idx] = "*";
+      anyNormalized = true;
+    }
+  }
+
+  if (!anyNormalized) return null;
+
+  const partMatch = chunkId.match(/-part-(\d+)$/);
+  const partSuffix = partMatch ? `:${partMatch[1]}` : "";
+
+  return `${parts.join("/")}:${heading}${partSuffix}`;
+}
+
 export function matchesMetadataFilters(
   metadata: Record<string, string>,
   filters: Record<string, string>,