Skip to content

Commit 61a7d1c

Browse files
samclark2015samclark2015
committed
Refreshing models on startup with an expired token now works
1 parent a4143fa commit 61a7d1c

File tree

2 files changed

+49
-16
lines changed

2 files changed

+49
-16
lines changed

‎extensions/positron-assistant/src/extension.ts‎

Lines changed: 19 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -45,6 +45,15 @@ class ModelDisposable implements vscode.Disposable {
4545
}
4646
}
4747

48+
/**
49+
* An error thrown by the assistant that can optionally be displayed to the user.
50+
*/
51+
export class AssistantError extends Error {
52+
constructor(message: string, public readonly display: boolean = true) {
53+
super(message);
54+
}
55+
}
56+
4857
/**
4958
* Dispose chat and/or completion models registered with Positron.
5059
* @param id If specified, only dispose models with the given ID. Otherwise, dispose all models.
@@ -134,8 +143,11 @@ export async function registerModels(context: vscode.ExtensionContext, storage:
134143
}
135144

136145
} catch (e) {
137-
const failedMessage = vscode.l10n.t('Positron Assistant: Failed to load model configurations.');
138-
vscode.window.showErrorMessage(`${failedMessage} ${e}`);
146+
if (!(e instanceof AssistantError) || e.display) {
147+
const failedMessage = vscode.l10n.t('Positron Assistant: Failed to load model configurations.');
148+
vscode.window.showErrorMessage(`${failedMessage} ${e}`);
149+
}
150+
139151
return;
140152
}
141153

@@ -145,7 +157,9 @@ export async function registerModels(context: vscode.ExtensionContext, storage:
145157
await registerModelWithAPI(config, context, storage);
146158
registeredModels.push(config);
147159
} catch (e) {
148-
vscode.window.showErrorMessage(`${e}`);
160+
if (!(e instanceof AssistantError) || e.display) {
161+
vscode.window.showErrorMessage(`${e}`);
162+
}
149163
}
150164
}
151165

@@ -172,13 +186,13 @@ export async function registerModels(context: vscode.ExtensionContext, storage:
172186
* @param modelConfig the language model's config
173187
* @param context the extension context
174188
*/
175-
async function registerModelWithAPI(modelConfig: ModelConfig, context: vscode.ExtensionContext, storage: SecretStorage) {
189+
export async function registerModelWithAPI(modelConfig: ModelConfig, context: vscode.ExtensionContext, storage: SecretStorage, instance?: positron.ai.LanguageModelChatProvider<vscode.LanguageModelChatInformation>) {
176190
// Register with Language Model API
177191
if (modelConfig.type === 'chat') {
178192
// const models = availableModels.get(modelConfig.provider);
179193
// const modelsCopy = models ? [...models] : [];
180194

181-
const languageModel = newLanguageModelChatProvider(modelConfig, context, storage);
195+
const languageModel = instance ?? newLanguageModelChatProvider(modelConfig, context, storage);
182196

183197
try {
184198
const error = await languageModel.resolveConnection(new vscode.CancellationTokenSource().token);

‎extensions/positron-assistant/src/models.ts‎

Lines changed: 30 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -6,7 +6,7 @@
66
import * as vscode from 'vscode';
77
import * as positron from 'positron';
88
import * as ai from 'ai';
9-
import { getMaxConnectionAttempts, getProviderTimeoutMs, ModelConfig, SecretStorage } from './config';
9+
import { expandConfigToSource, getMaxConnectionAttempts, getProviderTimeoutMs, ModelConfig, SecretStorage } from './config';
1010
import { AnthropicProvider, createAnthropic } from '@ai-sdk/anthropic';
1111
import { AzureOpenAIProvider, createAzure } from '@ai-sdk/azure';
1212
import { createVertex, GoogleVertexProvider } from '@ai-sdk/google-vertex';
@@ -20,11 +20,12 @@ import { AmazonBedrockProvider, createAmazonBedrock } from '@ai-sdk/amazon-bedro
2020
import { fromNodeProviderChain } from '@aws-sdk/credential-providers';
2121
import { AnthropicLanguageModel, DEFAULT_ANTHROPIC_MODEL_MATCH, DEFAULT_ANTHROPIC_MODEL_NAME } from './anthropic';
2222
import { DEFAULT_MAX_TOKEN_INPUT, DEFAULT_MAX_TOKEN_OUTPUT } from './constants.js';
23-
import { log, recordRequestTokenUsage, recordTokenUsage } from './extension.js';
23+
import { AssistantError, log, recordRequestTokenUsage, recordTokenUsage, registerModelWithAPI } from './extension.js';
2424
import { TokenUsage } from './tokens.js';
2525
import { BedrockClient, FoundationModelSummary, InferenceProfileSummary, ListFoundationModelsCommand, ListInferenceProfilesCommand } from '@aws-sdk/client-bedrock';
2626
import { PositLanguageModel } from './posit.js';
2727
import { applyModelFilters } from './modelFilters';
28+
import { PositronAssistantApi } from './api.js';
2829

2930
/**
3031
* Models used by chat participants and for vscode.lm.* API functionality.
@@ -281,7 +282,7 @@ abstract class AILanguageModel implements positron.ai.LanguageModelChatProvider
281282
constructor(
282283
protected readonly _config: ModelConfig,
283284
protected readonly _context?: vscode.ExtensionContext,
284-
private readonly _storage?: SecretStorage,
285+
protected readonly _storage?: SecretStorage,
285286
) {
286287
this.id = _config.id;
287288
this.name = _config.name;
@@ -1045,14 +1046,14 @@ export class AWSLanguageModel extends AILanguageModel implements positron.ai.Lan
10451046
bedrockClient: BedrockClient;
10461047
inferenceProfiles: InferenceProfileSummary[] = [];
10471048

1048-
constructor(_config: ModelConfig, _context?: vscode.ExtensionContext) {
1049+
constructor(_config: ModelConfig, _context?: vscode.ExtensionContext, _storage?: SecretStorage) {
10491050
// Update a stale model configuration to the latest defaults
10501051
const models = availableModels.get('amazon-bedrock')?.map(m => m.identifier) || [];
10511052
if (!(_config.model in models)) {
10521053
_config.name = AWSLanguageModel.source.defaults.name;
10531054
_config.model = AWSLanguageModel.source.defaults.model;
10541055
}
1055-
super(_config, _context);
1056+
super(_config, _context, _storage);
10561057

10571058
const environmentSettings = vscode.workspace.getConfiguration('positron.assistant.providerVariables').get<BedrockProviderVariables>('bedrock', {});
10581059
log.debug(`[BedrockLanguageModel] positron.assistant.providerVariables.bedrock settings: ${JSON.stringify(environmentSettings)}`);
@@ -1112,12 +1113,16 @@ export class AWSLanguageModel extends AILanguageModel implements positron.ai.Lan
11121113
if (message.includes('aws sso login')) {
11131114
// Give the user the option to login automatically
11141115
// Display an error message with an action the user can take
1116+
const isConnectionTest = error.stack?.includes('resolveConnection');
11151117
const action = { title: vscode.l10n.t('Run in Terminal'), id: 'aws-sso-login' };
11161118
vscode.window.showErrorMessage(`Amazon Bedrock: ${message}`, action).then(async selection => {
11171119
if (selection?.id === action.id) {
1120+
// User chose to login, so we need to refresh the credentials
1121+
11181122
// Grab the profile & region to refresh from the Bedrock client config
11191123
const profile = this.bedrockClient.config.profile;
1120-
const region = this.bedrockClient.config.region;
1124+
// Region may be an async function or a string, so handle both cases
1125+
const region = typeof this.bedrockClient.config.region === 'function' ? await this.bedrockClient.config.region() : this.bedrockClient.config.region;
11211126
// Execute the AWS SSO login command as a native task
11221127
const taskExecution = await vscode.tasks.executeTask(new vscode.Task(
11231128
{ type: 'shell' },
@@ -1130,7 +1135,8 @@ export class AWSLanguageModel extends AILanguageModel implements positron.ai.Lan
11301135
vscode.tasks.onDidEndTaskProcess(e => {
11311136
if (e.execution === taskExecution) {
11321137
// Notify the user of the result
1133-
if (e.exitCode === 0 || e.exitCode === undefined) {
1138+
const success = e.exitCode === 0 || e.exitCode === undefined;
1139+
if (success) {
11341140
// Success
11351141
vscode.window.showInformationMessage(vscode.l10n.t('AWS login completed successfully'));
11361142
} else {
@@ -1142,18 +1148,31 @@ export class AWSLanguageModel extends AILanguageModel implements positron.ai.Lan
11421148
// This is a little sneaky, but works + no other native method
11431149
const redirectUri = vscode.Uri.from({ scheme: vscode.env.uriScheme });
11441150
vscode.env.openExternal(redirectUri);
1151+
1152+
if (success && isConnectionTest) {
1153+
// If we were in a connection test, re-run it now that we've logged in
1154+
registerModelWithAPI(
1155+
this._config,
1156+
this._context,
1157+
this._storage,
1158+
this
1159+
).then(() => {
1160+
positron.ai.addLanguageModelConfig(expandConfigToSource(this._config));
1161+
PositronAssistantApi.get().notifySignIn(this._config.name);
1162+
});
1163+
}
11451164
}
11461165
});
11471166
}
11481167
});
11491168

1150-
if (error.stack?.includes('resolveConnection')) {
1151-
// We're in a connection test, so just return undefined and use this own method's error display
1152-
return undefined;
1169+
if (isConnectionTest) {
1170+
// We're in a connection test, so throw an AssistantError to avoid showing a message box
1171+
// but that stops the model provider from being registered in core
1172+
throw new AssistantError(message, false);
11531173
} else {
11541174
// We are in a chat response, so we should return an error to display in the chat pane
11551175
throw new Error(vscode.l10n.t(`AWS login required. Please run \`aws sso login --profile ${this.bedrockClient.config.profile} --region ${this.bedrockClient.config.region}\` in the terminal, and retry this request.`));
1156-
// return vscode.l10n.t(`AWS login required. Please run \`aws sso login --profile ${this.bedrockClient.config.profile}\` in the terminal to authenticate.`);
11571176
}
11581177
} else {
11591178
return vscode.l10n.t(`Invalid AWS credentials. {0}`, message);

0 commit comments

Comments
 (0)