Ungroup tools based on embedding rankings by vritant24 · Pull Request #678 · microsoft/vscode-copilot-chat

vritant24 · 2025-08-19T22:05:56Z

Add embeddings-based tool ranking and caching for virtual tool grouper to predict which tools are most relevant to user queries, improving the tool selection experience.

The feature is currently behind an experiment flag (chat.virtualTools.embeddingRanking)

… into dev/vrbhardw/embeddings

vritant24 · 2025-09-11T17:40:51Z

cc @connor4312

Copilot

Pull Request Overview

This PR adds embeddings-based tool ranking and caching to the virtual tool grouper to predict which tools are most relevant to user queries, improving the tool selection experience. The feature is currently behind an experiment flag (chat.advanced.virtualTools.embeddingRanking) and includes infrastructure for computing tool embeddings and expanding virtual tool groups based on similarity rankings.

Key changes include:

Implementing an embeddings-based prediction system for virtual tools
Adding caching for tool embeddings with both pre-computed and runtime computation
Integrating the ranking system into the tool grouping workflow

Reviewed Changes

Copilot reviewed 13 out of 13 changed files in this pull request and generated 2 comments.

Show a summary per file

File	Description
`src/platform/embeddings/common/embeddingsIndex.ts`	Adds new `Tools` remote cache type for storing tool embeddings
`src/platform/configuration/common/configurationService.ts`	Adds experiment flag for virtual tool embedding ranking feature
`src/extension/tools/test/node/virtualTools/virtualToolGrouping.spec.ts`	Updates test calls to include query parameter in grouping methods
`src/extension/tools/test/node/virtualTools/virtualToolGrouper.spec.ts`	Updates test calls and adds comprehensive tests for the new embedding-based expansion feature
`src/extension/tools/common/virtualTools/virtualToolTypes.ts`	Updates interface methods to accept query string parameter for context
`src/extension/tools/common/virtualTools/virtualToolGrouper.ts`	Core implementation of embedding-based tool prediction and virtual tool expansion logic
`src/extension/tools/common/virtualTools/toolGrouping.ts`	Updates method signatures to pass through query string for embedding computation
`src/extension/tools/common/virtualTools/toolEmbeddingsCache.ts`	New file implementing tool embeddings computation and caching infrastructure
`src/extension/tools/common/test/toolEmbeddingsCache.spec.ts`	Comprehensive unit tests for the tool embeddings computation system
`src/extension/test/node/services.ts`	Registers embeddings computer service for testing
`src/extension/prompt/node/defaultIntentRequestHandler.ts`	Updates method calls to pass query text for embedding-based tool ranking
`src/extension/intents/node/toolCallingLoop.ts`	Updates abstract method signature to include query parameter
`src/extension/intents/node/agentIntent.ts`	Updates method call to include query parameter

src/extension/tools/test/node/virtualTools/virtualToolGrouping.spec.ts

src/extension/tools/common/virtualTools/virtualToolTypes.ts

connor4312 · 2025-09-11T18:29:30Z

I get an error when we try to compare here

https://github.com/vritant24/vscode-copilot-chat/blob/cbf04b59ee7dfb9e94b5d32d4b534ad780c8c2e1/src/platform/embeddings/common/embeddingsComputer.ts#L162

embedding is simply an array of the embedding values, not the

export interface Embedding {
	readonly type: EmbeddingType;
	readonly value: EmbeddingVector;
}

as expected

src/extension/prompt/node/defaultIntentRequestHandler.ts

src/extension/intents/node/toolCallingLoop.ts

src/extension/tools/common/test/toolEmbeddingsCache.spec.ts

src/extension/tools/common/virtualTools/toolEmbeddingsCache.ts

src/extension/tools/test/node/virtualTools/virtualToolGrouper.spec.ts

src/extension/tools/common/virtualTools/virtualToolGrouper.ts

connor4312 · 2025-09-11T18:21:36Z

src/extension/tools/common/virtualTools/virtualToolGrouper.ts

+		if (virtualToolEmbeddingRankingEnabled) {
+			const predictedTools = await this._getPredictedTools(query, tools, token);
+
+			this._expandGroupsWithPredictedTools(root, predictedTools);
+		}
+
 		this._reExpandToolsToHitBudget(root);


How about we refactor this so _reExpandToolsToHitBudget is generic and takes a ranker or comparator argument. By default it would be something like this._reExpandToolsToHitBudget(root, g => g.tools.length) but for embeddings it would be like this._reExpandToolsToHitBudget(root, g => this._getGroupPredictedRelevancy(g, predictedTools))

Thanks, that's a pretty great way structure it. I've made that change but allowing for more tools with the embedding based prediction as opposed to what's the limit today. If that feeld incorrect do let me know and I can keep the limits the same

…4/vscode-copilot-chat into dev/vrbhardw/embeddings

vritant24 · 2025-09-12T18:55:16Z

I get an error when we try to compare here

https://github.com/vritant24/vscode-copilot-chat/blob/cbf04b59ee7dfb9e94b5d32d4b534ad780c8c2e1/src/platform/embeddings/common/embeddingsComputer.ts#L162

embedding is simply an array of the embedding values, not the
export interface Embedding {
	readonly type: EmbeddingType;
	readonly value: EmbeddingVector;
}
as expected

thanks, somehow I didn't run into this when testing but was able to repro it. It's now fixed

Vritant Bhardwaj added 18 commits August 19, 2025 10:25

wip

4306ba9

fix impl

82fdf6a

fix return

2f7e33a

add a basic algo to unfold grouped tools

edc5fda

only rank based on available tools

747f355

use json file

fc61505

add todo

56d065b

only compute similarity for non built in tools

69aeabc

compute tool embeddings on the fly and refactor

1d3e633

add experiment config

46969f4

add testsfor ToolEmbeddignsComputer

8b37363

add tests for ToolEmbeddingsComputer and some cleanup

d71896d

add an integration test

b04df54

use remote cacche to get embeddings

52dc3f3

add logging and error catching

f9a45a8

Merge branch 'main' of https://github.com/vritant24/vscode-copilot-chat…

ba125a3

… into dev/vrbhardw/embeddings

consume new change

f42c083

use telemetry correlation id

453e48a

vritant24 changed the title ~~Dev/vrbhardw/embeddings~~ Ungroup tools based on embedding rankings Sep 10, 2025

Vritant Bhardwaj and others added 2 commits September 10, 2025 13:50

undo change

33a048f

Update VirtualToolEmbeddingRanking setting path

354cbf5

vritant24 marked this pull request as ready for review September 11, 2025 17:40

Copilot AI review requested due to automatic review settings September 11, 2025 17:40

Merge branch 'main' into dev/vrbhardw/embeddings

cbf04b5

Copilot AI reviewed Sep 11, 2025

View reviewed changes

src/extension/tools/test/node/virtualTools/virtualToolGrouping.spec.ts Show resolved Hide resolved

src/extension/tools/common/virtualTools/virtualToolTypes.ts Show resolved Hide resolved

vs-code-engineering bot assigned chrmarti Sep 11, 2025

vs-code-engineering bot added the triage-needed label Sep 11, 2025

connor4312 assigned connor4312 and unassigned chrmarti Sep 11, 2025

connor4312 self-requested a review September 11, 2025 17:59

connor4312 requested changes Sep 11, 2025

View reviewed changes

Vritant Bhardwaj and others added 9 commits September 11, 2025 11:42

refactor query extraction

d27f53d

tiny pr refactor

65b03e7

Merge branch 'dev/vrbhardw/embeddings' of https://github.com/vritant2…

4f51bee

…4/vscode-copilot-chat into dev/vrbhardw/embeddings

refactor instantiation

1c58ccd

refactor test

e318ea4

Merge branch 'main' into dev/vrbhardw/embeddings

6f932f1

address comment

b989687

fix typing

03f6c2c

Merge branch 'dev/vrbhardw/embeddings' of https://github.com/vritant2…

5e94e88

…4/vscode-copilot-chat into dev/vrbhardw/embeddings

Merge branch 'main' into dev/vrbhardw/embeddings

1476ec2

connor4312 self-requested a review September 15, 2025 19:57

connor4312 approved these changes Sep 15, 2025

View reviewed changes

vs-code-engineering bot added this to the September 2025 milestone Sep 15, 2025

lszomoru approved these changes Sep 15, 2025

View reviewed changes

connor4312 added this pull request to the merge queue Sep 15, 2025

Merged via the queue into microsoft:main with commit 6621102 Sep 15, 2025
6 checks passed

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Ungroup tools based on embedding rankings#678

Ungroup tools based on embedding rankings#678
connor4312 merged 31 commits intomicrosoft:mainfrom
vritant24:dev/vrbhardw/embeddings

vritant24 commented Aug 19, 2025 •

edited

Loading

Uh oh!

vritant24 commented Sep 11, 2025

Uh oh!

Copilot AI left a comment

Uh oh!

Uh oh!

Uh oh!

connor4312 commented Sep 11, 2025

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

connor4312 Sep 11, 2025

Uh oh!

vritant24 Sep 12, 2025

Uh oh!

vritant24 commented Sep 12, 2025

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

5 participants

Conversation

vritant24 commented Aug 19, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

vritant24 commented Sep 11, 2025

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull Request Overview

Reviewed Changes

Uh oh!

Uh oh!

Uh oh!

connor4312 commented Sep 11, 2025

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

connor4312 Sep 11, 2025

Choose a reason for hiding this comment

Uh oh!

vritant24 Sep 12, 2025

Choose a reason for hiding this comment

Uh oh!

vritant24 commented Sep 12, 2025

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

5 participants

vritant24 commented Aug 19, 2025 •

edited

Loading