Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Send server contextualization to Copilot extension #24230

Merged
Merged
Show file tree
Hide file tree
Changes from 8 commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 2 additions & 2 deletions extensions/mssql/src/contracts.ts
Original file line number Diff line number Diff line change
Expand Up @@ -1566,8 +1566,8 @@ export interface ServerContextualizationParams {
ownerUri: string;
}

export namespace GenerateServerContextualizationNotification {
export const type = new NotificationType<ServerContextualizationParams, void>('metadata/generateServerContext');
export namespace GenerateServerContextualizationRequest {
export const type = new RequestType<ServerContextualizationParams, azdata.contextualization.GenerateServerContextualizationResult, void, void>('metadata/generateServerContext');
}

export namespace GetServerContextualizationRequest {
Expand Down
12 changes: 9 additions & 3 deletions extensions/mssql/src/features.ts
Original file line number Diff line number Diff line change
Expand Up @@ -1310,7 +1310,7 @@ export class ExecutionPlanServiceFeature extends SqlOpsFeature<undefined> {
*/
export class ServerContextualizationServiceFeature extends SqlOpsFeature<undefined> {
private static readonly messagesTypes: RPCMessageType[] = [
contracts.GenerateServerContextualizationNotification.type
contracts.GenerateServerContextualizationRequest.type
];

constructor(client: SqlOpsDataClient) {
Expand All @@ -1330,12 +1330,18 @@ export class ServerContextualizationServiceFeature extends SqlOpsFeature<undefin
protected registerProvider(options: undefined): Disposable {
const client = this._client;

const generateServerContextualization = (ownerUri: string): void => {
const generateServerContextualization = (ownerUri: string): Thenable<azdata.contextualization.GenerateServerContextualizationResult> => {
const params: contracts.ServerContextualizationParams = {
ownerUri: ownerUri
};

return client.sendNotification(contracts.GenerateServerContextualizationNotification.type, params);
return client.sendRequest(contracts.GenerateServerContextualizationRequest.type, params).then(
r => r,
e => {
client.logFailedRequest(contracts.GenerateServerContextualizationRequest.type, e);
return Promise.reject(e);
}
);
};

const getServerContextualization = (ownerUri: string): Thenable<azdata.contextualization.GetServerContextualizationResult> => {
Expand Down
15 changes: 11 additions & 4 deletions src/sql/azdata.proposed.d.ts
Original file line number Diff line number Diff line change
Expand Up @@ -901,7 +901,7 @@ declare module 'azdata' {
* Copilot for improved suggestions.
* @param provider The provider to register
*/
export function registerServerContextualizationProvider(provider: contextualization.ServerContextualizationProvider): vscode.Disposable
export function registerServerContextualizationProvider(provider: contextualization.ServerContextualizationProvider): vscode.Disposable;
}

export namespace designers {
Expand Down Expand Up @@ -1783,19 +1783,26 @@ declare module 'azdata' {
}

export namespace contextualization {
export interface GenerateServerContextualizationResult {
/**
* The generated server context.
*/
context: string | undefined;
}

export interface GetServerContextualizationResult {
/**
* An array containing the generated server context.
* The retrieved server context.
*/
context: string[];
context: string | undefined;
}

export interface ServerContextualizationProvider extends DataProvider {
/**
* Generates server context.
* @param ownerUri The URI of the connection to generate context for.
*/
generateServerContextualization(ownerUri: string): void;
generateServerContextualization(ownerUri: string): Thenable<GenerateServerContextualizationResult>;

/**
* Gets server context, which can be in the form of create scripts but is left up each provider.
Expand Down
4 changes: 2 additions & 2 deletions src/sql/workbench/api/common/extHostDataProtocol.ts
Original file line number Diff line number Diff line change
Expand Up @@ -972,8 +972,8 @@ export class ExtHostDataProtocol extends ExtHostDataProtocolShape {

// Database Server Contextualization API

public override $generateServerContextualization(handle: number, ownerUri: string): void {
this._resolveProvider<azdata.contextualization.ServerContextualizationProvider>(handle).generateServerContextualization(ownerUri);
public override $generateServerContextualization(handle: number, ownerUri: string): Thenable<azdata.contextualization.GenerateServerContextualizationResult> {
return this._resolveProvider<azdata.contextualization.ServerContextualizationProvider>(handle).generateServerContextualization(ownerUri);
}

public override $getServerContextualization(handle: number, ownerUri: string): Thenable<azdata.contextualization.GetServerContextualizationResult> {
Expand Down
2 changes: 1 addition & 1 deletion src/sql/workbench/api/common/sqlExtHost.protocol.ts
Original file line number Diff line number Diff line change
Expand Up @@ -598,7 +598,7 @@ export abstract class ExtHostDataProtocolShape {
/**
* Generates server context.
*/
$generateServerContextualization(handle: number, ownerUri: string): void { throw ni(); }
$generateServerContextualization(handle: number, ownerUri: string): Thenable<azdata.contextualization.GenerateServerContextualizationResult> { throw ni(); }
/**
* Gets server context.
*/
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -19,7 +19,6 @@ import { IInstantiationService } from 'vs/platform/instantiation/common/instanti
import { EditorInput } from 'vs/workbench/common/editor/editorInput';
import { IResourceEditorInput } from 'vs/platform/editor/common/editor';
import { IServerContextualizationService } from 'sql/workbench/services/contextualization/common/interfaces';
import { IExtensionService } from 'vs/workbench/services/extensions/common/extensions';

export class FileQueryEditorInput extends QueryEditorInput {

Expand All @@ -33,10 +32,9 @@ export class FileQueryEditorInput extends QueryEditorInput {
@IQueryModelService queryModelService: IQueryModelService,
@IConfigurationService configurationService: IConfigurationService,
@IInstantiationService instantiationService: IInstantiationService,
@IServerContextualizationService serverContextualizationService: IServerContextualizationService,
@IExtensionService extensionService: IExtensionService
@IServerContextualizationService serverContextualizationService: IServerContextualizationService
) {
super(description, text, results, connectionManagementService, queryModelService, configurationService, instantiationService, serverContextualizationService, extensionService);
super(description, text, results, connectionManagementService, queryModelService, configurationService, instantiationService, serverContextualizationService);
}

public override resolve(): Promise<ITextFileEditorModel | BinaryEditorModel> {
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -23,7 +23,6 @@ import { IEditorResolverService } from 'vs/workbench/services/editor/common/edit
import { Uri } from 'vscode';
import { ILogService } from 'vs/platform/log/common/log';
import { IServerContextualizationService } from 'sql/workbench/services/contextualization/common/interfaces';
import { IExtensionService } from 'vs/workbench/services/extensions/common/extensions';

export class UntitledQueryEditorInput extends QueryEditorInput implements IUntitledQueryEditorInput {

Expand All @@ -39,10 +38,9 @@ export class UntitledQueryEditorInput extends QueryEditorInput implements IUntit
@IInstantiationService instantiationService: IInstantiationService,
@ILogService private readonly logService: ILogService,
@IEditorResolverService private readonly editorResolverService: IEditorResolverService,
@IServerContextualizationService serverContextualizationService: IServerContextualizationService,
@IExtensionService extensionService: IExtensionService
@IServerContextualizationService serverContextualizationService: IServerContextualizationService
) {
super(description, text, results, connectionManagementService, queryModelService, configurationService, instantiationService, serverContextualizationService, extensionService);
super(description, text, results, connectionManagementService, queryModelService, configurationService, instantiationService, serverContextualizationService);
// Set the mode explicitely to stop the auto language detection service from changing the mode unexpectedly.
// the auto language detection service won't do the language change only if the mode is explicitely set.
// if the mode (e.g. kusto, sql) do not exist for whatever reason, we will default it to sql.
Expand Down
49 changes: 23 additions & 26 deletions src/sql/workbench/common/editor/query/queryEditorInput.ts
Original file line number Diff line number Diff line change
Expand Up @@ -21,7 +21,6 @@ import { IQueryEditorConfiguration } from 'sql/platform/query/common/query';
import { EditorInput } from 'vs/workbench/common/editor/editorInput';
import { IInstantiationService } from 'vs/platform/instantiation/common/instantiation';
import { IServerContextualizationService } from 'sql/workbench/services/contextualization/common/interfaces';
import { IExtensionService } from 'vs/workbench/services/extensions/common/extensions';

const MAX_SIZE = 13;

Expand Down Expand Up @@ -144,8 +143,6 @@ export abstract class QueryEditorInput extends EditorInput implements IConnectab
private _state = this._register(new QueryEditorState());
public get state(): QueryEditorState { return this._state; }

private _serverContext: string[];

constructor(
private _description: string | undefined,
protected _text: AbstractTextResourceEditorInput,
Expand All @@ -154,8 +151,7 @@ export abstract class QueryEditorInput extends EditorInput implements IConnectab
@IQueryModelService private readonly queryModelService: IQueryModelService,
@IConfigurationService private readonly configurationService: IConfigurationService,
@IInstantiationService protected readonly instantiationService: IInstantiationService,
@IServerContextualizationService private readonly serverContextualizationService: IServerContextualizationService,
@IExtensionService private readonly extensionService: IExtensionService
@IServerContextualizationService private readonly serverContextualizationService: IServerContextualizationService
) {
super();

Expand Down Expand Up @@ -241,27 +237,6 @@ export abstract class QueryEditorInput extends EditorInput implements IConnectab
public override isDirty(): boolean { return this._text.isDirty(); }
public get resource(): URI { return this._text.resource; }

public async getServerContext(): Promise<string[]> {
const copilotExt = await this.extensionService.getExtension('github.copilot');

if (copilotExt && this.configurationService.getValue<IQueryEditorConfiguration>('queryEditor').githubCopilotContextualizationEnabled) {
if (!this._serverContext) {
const result = await this.serverContextualizationService.getServerContextualization(this.uri);
// TODO lewissanchez - Remove this from here once Copilot starts pulling context. That isn't implemented yet, so
// getting scripts this way for now.
this._serverContext = result.context;

return this._serverContext;
}
else {
return this._serverContext;
}
}
else {
return Promise.resolve([]);
}
}

public override getName(longForm?: boolean): string {
if (this.configurationService.getValue<IQueryEditorConfiguration>('queryEditor').showConnectionInfoInTitle) {
let profile = this.connectionManagementService.getConnectionProfile(this.uri);
Expand Down Expand Up @@ -346,6 +321,28 @@ export abstract class QueryEditorInput extends EditorInput implements IConnectab
}
}
this._onDidChangeLabel.fire();

this.contextualizeEditorForCopilot();
lewis-sanchez marked this conversation as resolved.
Show resolved Hide resolved
lewis-sanchez marked this conversation as resolved.
Show resolved Hide resolved
}

private contextualizeEditorForCopilot(): void {
// Don't need to take any actions if contextualization is not enabled and can return
if (!this.configurationService.getValue<IQueryEditorConfiguration>('queryEditor').githubCopilotContextualizationEnabled) {
lewis-sanchez marked this conversation as resolved.
Show resolved Hide resolved
return;
}

this.serverContextualizationService.getServerContextualization(this.uri)
.then(async getServerContextualizationResult => {
lewis-sanchez marked this conversation as resolved.
Show resolved Hide resolved
if (getServerContextualizationResult.context) {
await this.serverContextualizationService.sendServerContextualizationToCopilot(getServerContextualizationResult.context);
}
else {
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

So now you're changing it to only generate the context when an editor is connected?

  1. Is there any concern about the duration? As I found in my testing - getting the scripts can take a long time so it might be quite a while after an editor is opened before the context is sent
  2. If you're going with this then would it be better to just get rid of the "generate" call completely and instead just have get be the only call (and if it doesn't have context it'll go generate that before returning)?

Copy link
Contributor Author

@lewis-sanchez lewis-sanchez Aug 31, 2023

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think it would be better to do both actions in one request, but I wasn't sure if generating and reading context was considered doing too much in a single request endpoint, which is why I kept them separate.

I can consolidate the two requests into a single one, so that this new request can read the context or generate context if there isn't any available for it to send back to ADS.

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Sounds good, I think consolidating into one makes sense to keep things simpler. Then you can easily add something like a "forceRefresh" param as well later on to force regenerating the context if that's something you want (probably useful for users to be able to control like how we have a refresh intellisense command)

const generateServerContextualizationResult = await this.serverContextualizationService.generateServerContextualization(this.uri);
if (generateServerContextualizationResult.context) {
await this.serverContextualizationService.sendServerContextualizationToCopilot(generateServerContextualizationResult.context);
}
}
});
}

public onDisconnect(): void {
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -467,7 +467,7 @@ suite('commandLineService tests', () => {
let uri = URI.file(args._[0]);
const workbenchinstantiationService = workbenchInstantiationService();
const editorInput = workbenchinstantiationService.createInstance(FileEditorInput, uri, undefined, undefined, undefined, undefined, undefined, undefined);
const queryInput = new FileQueryEditorInput(undefined, editorInput, undefined, connectionManagementService.object, querymodelService.object, configurationService.object, workbenchinstantiationService, undefined, undefined);
const queryInput = new FileQueryEditorInput(undefined, editorInput, undefined, connectionManagementService.object, querymodelService.object, configurationService.object, workbenchinstantiationService, undefined);
queryInput.state.connected = true;
const editorService: TypeMoq.Mock<IEditorService> = TypeMoq.Mock.ofType<IEditorService>(TestEditorService, TypeMoq.MockBehavior.Strict);
editorService.setup(e => e.editors).returns(() => [queryInput]);
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -317,7 +317,6 @@ suite('SQL QueryEditor Tests', () => {
testinstantiationService,
undefined,
undefined,
undefined,
undefined
);
});
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -33,11 +33,17 @@ export interface IServerContextualizationService {
* Generates server context
* @param ownerUri The URI of the connection to generate context for.
*/
generateServerContextualization(ownerUri: string): void;
generateServerContextualization(ownerUri: string): Promise<azdata.contextualization.GenerateServerContextualizationResult>;

/**
* Gets all database context.
* @param ownerUri The URI of the connection to get context for.
*/
getServerContextualization(ownerUri: string): Promise<azdata.contextualization.GetServerContextualizationResult>;

/**
* Sends server context to the Copilot extension, so it can be used to generate improved suggestions.
* @param serverContext The server context to be sent to Copilot.
*/
sendServerContextualizationToCopilot(serverContext: string | undefined): Promise<void>
}
Original file line number Diff line number Diff line change
Expand Up @@ -5,10 +5,11 @@

import * as azdata from 'azdata';
import { invalidProvider } from 'sql/base/common/errors';
import { IConnectionManagementService, IConnectionParams } from 'sql/platform/connection/common/connectionManagement';
import { IConnectionManagementService } from 'sql/platform/connection/common/connectionManagement';
import { IQueryEditorConfiguration } from 'sql/platform/query/common/query';
import { IServerContextualizationService } from 'sql/workbench/services/contextualization/common/interfaces';
import { Disposable } from 'vs/base/common/lifecycle';
import { ICommandService } from 'vs/platform/commands/common/commands';
import { IConfigurationService } from 'vs/platform/configuration/common/configuration';
import { IExtensionService } from 'vs/workbench/services/extensions/common/extensions';

Expand All @@ -19,18 +20,10 @@ export class ServerContextualizationService extends Disposable implements IServe
constructor(
@IConnectionManagementService private readonly _connectionManagementService: IConnectionManagementService,
@IConfigurationService private readonly _configurationService: IConfigurationService,
@IExtensionService private readonly _extensionService: IExtensionService
@IExtensionService private readonly _extensionService: IExtensionService,
@ICommandService private readonly _commandService: ICommandService
) {
super();

this._register(this._connectionManagementService.onConnect(async (e: IConnectionParams) => {
const copilotExt = await this._extensionService.getExtension('github.copilot');

if (copilotExt && this._configurationService.getValue<IQueryEditorConfiguration>('queryEditor').githubCopilotContextualizationEnabled) {
const ownerUri = e.connectionUri;
await this.generateServerContextualization(ownerUri);
}
}));
}

/**
Expand Down Expand Up @@ -67,11 +60,19 @@ export class ServerContextualizationService extends Disposable implements IServe
* Generates server context
* @param ownerUri The URI of the connection to generate context for.
*/
public generateServerContextualization(ownerUri: string): void {
public async generateServerContextualization(ownerUri: string): Promise<azdata.contextualization.GenerateServerContextualizationResult> {
const copilotExt = await this._extensionService.getExtension('github.copilot');
const isContextualizationEnabled = this._configurationService.getValue<IQueryEditorConfiguration>('queryEditor').githubCopilotContextualizationEnabled;

const providerName = this._connectionManagementService.getProviderIdFromUri(ownerUri);
const handler = this.getProvider(providerName);
if (handler) {
handler.generateServerContextualization(ownerUri);
if (copilotExt && isContextualizationEnabled && handler) {
return await handler.generateServerContextualization(ownerUri);
}
else {
return Promise.resolve({
context: undefined
});
}
}

Expand All @@ -80,14 +81,27 @@ export class ServerContextualizationService extends Disposable implements IServe
* @param ownerUri The URI of the connection to get context for.
*/
public async getServerContextualization(ownerUri: string): Promise<azdata.contextualization.GetServerContextualizationResult> {
const copilotExt = await this._extensionService.getExtension('github.copilot');
lewis-sanchez marked this conversation as resolved.
Show resolved Hide resolved
const isContextualizationEnabled = this._configurationService.getValue<IQueryEditorConfiguration>('queryEditor').githubCopilotContextualizationEnabled;

const providerName = this._connectionManagementService.getProviderIdFromUri(ownerUri);
const handler = this.getProvider(providerName);
if (handler) {
if (copilotExt && isContextualizationEnabled && handler) {
return await handler.getServerContextualization(ownerUri);
}
else {
return Promise.resolve({
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

would it be helpful to also log an error here that the provider wasn't able to be retrieved for debugging?

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I'm going to have a follow up PR with all the logging.

context: []
context: undefined
});
}
}

public async sendServerContextualizationToCopilot(serverContext: string | undefined): Promise<void> {
const copilotExt = await this._extensionService.getExtension('github.copilot');
if (copilotExt && serverContext && this._configurationService.getValue<IQueryEditorConfiguration>('queryEditor').githubCopilotContextualizationEnabled) {
// LEWISSANCHEZ TODO: Find way to set context on untitled query editor files. Need to save first for Copilot status to say "Has Context"

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Hey @lewis-sanchez, let's chat about this --- I can change the fork slightly to better work w/ untitled files. I also want to make sure that, before the call to github.copilot.provideContext we hook in some amount of "schema compression" (using a fixed budget, as we discussed offline).

I can play around with the compression and get back to you on that. (I'd like to start w/ similar heuristics to what we use elsewhere.)

await this._commandService.executeCommand('github.copilot.provideContext', '**/*.sql', {
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

  1. Is there any way to specify what editor the context is for?
  2. Does providing context clear the other context?

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I can comment here (since I'm working on the copilot fork side): context is associated to files via a glob-style pattern. One context object is allowed per unique pattern (so, in this case, calling provideContext again with the exact same **/*.sql pattern will replace the existing context).

You could make this call a bit more granular, so that it matches only the given file (but we still need to work through what happens w/ untitled editors).

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Cool, thanks for the info. How does copilot know what editor to associate what context with?

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

On the copilot side, as you type in an editor, we'll see if the uri for that editor matches any pattern in the { pattern: context } map. For each pattern that matches, we add the context as a candidate for inclusion into the prompt.

This seemingly works well (except for untitled editors, not sure what the uri for those looks like on the Copilot side, need to do some tests there).

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Great! So yeah @lewis-sanchez, sounds like we should scope the context to be for each file.

@jjhenkel Is there any concern about context not being cleaned up if the editor is closed (or does copilot handle that already)?

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Cleanup could be done by setting the context to empty when an editor is closed, right?

Anyways, I agree it's not something that is critical to address now it sounds like - but something to keep in mind as we're designing how this will all fit together.

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Set to empty (assuming we're using the more specific per-file patterns) would be a way to implement some level of cleanup for now, yes!

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@Charles-Gagnon, I thought this was already scoped to the file, but guess not. Do you know where I would need to make changes to scope context to a file?

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@jjhenkel Should be able to help. From what he was saying earlier you have to modify the pattern you use when setting it (right now hardcoded to **/*.sql). So the change I believe it would be would be something more like <Path to file>/FileName.sql instead (although what the pattern can accept isn't something I know)

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@lewis-sanchez I'm looking at your branch today/tomorrow and can do some testing w/ stricter path constraints and what happens w/ untitled files (and modify the copilot fork if needed).

value: serverContext
});
}
}
Expand Down