Skip to content

Commit 94bce75

Browse files
sergeibbbsergeibbb
committed
Uses organization AI configuration from GKDev
(#4300, #4338)
1 parent bdb8200 commit 94bce75

10 files changed

+158
-17
lines changed

‎src/constants.context.ts

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,7 @@ import type { AnnotationStatus, Keys } from './constants';
33
import type { SubscriptionState } from './constants.subscription';
44
import type { CustomEditorTypes, GroupableTreeViewTypes, WebviewTypes, WebviewViewTypes } from './constants.views';
55
import type { Features } from './features';
6+
import type { OrgAIProviders } from './plus/gk/models/organization';
67
import type { PromoKeys } from './plus/gk/models/promo';
78
import type { SubscriptionPlanIds } from './plus/gk/models/subscription';
89
import type { WalkthroughContextKeys } from './telemetry/walkthroughStateProvider';
@@ -14,6 +15,8 @@ export type ContextKeys = {
1415
'gitlens:enabled': boolean;
1516
'gitlens:gk:hasOrganizations': boolean;
1617
'gitlens:gk:organization:ai:enabled': boolean;
18+
'gitlens:gk:organization:ai:enforceProviders': boolean;
19+
'gitlens:gk:organization:ai:providers': OrgAIProviders;
1720
'gitlens:gk:organization:drafts:byob': boolean;
1821
'gitlens:gk:organization:drafts:enabled': boolean;
1922
'gitlens:hasVirtualFolders': boolean;

‎src/plus/ai/aiProviderService.ts

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -66,7 +66,7 @@ import type {
6666
PromptTemplateType,
6767
} from './models/promptTemplates';
6868
import type { AIChatMessage, AIProvider, AIRequestResult } from './models/provider';
69-
import { ensureAccess } from './utils/-webview/ai.utils';
69+
import { ensureAccess, getOrgAIConfig, isProviderEnabledByOrg } from './utils/-webview/ai.utils';
7070
import { getLocalPromptTemplate, resolvePrompt } from './utils/-webview/prompt.utils';
7171

7272
export interface AIResult {
@@ -379,9 +379,10 @@ export class AIProviderService implements Disposable {
379379
}
380380

381381
async getProvidersConfiguration(): Promise<Map<AIProviders, AIProviderDescriptorWithConfiguration>> {
382+
const orgAiConfig = getOrgAIConfig();
382383
const promises = await Promise.allSettled(
383384
map(
384-
supportedAIProviders.values(),
385+
[...supportedAIProviders.values()].filter(p => isProviderEnabledByOrg(p.id, orgAiConfig)),
385386
async p =>
386387
[
387388
p.id,

‎src/plus/ai/azureProvider.ts

Lines changed: 6 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -5,7 +5,7 @@ import { configuration } from '../../system/-webview/configuration';
55
import type { AIActionType, AIModel } from './models/model';
66
import { openAIModels } from './models/model';
77
import { OpenAICompatibleProviderBase } from './openAICompatibleProviderBase';
8-
import { isAzureUrl } from './utils/-webview/ai.utils';
8+
import { ensureOrgConfiguredUrl, getOrgAIProviderOfType, isAzureUrl } from './utils/-webview/ai.utils';
99

1010
type AzureModel = AIModel<typeof provider.id>;
1111
const models: AzureModel[] = openAIModels(provider);
@@ -24,10 +24,14 @@ export class AzureProvider extends OpenAICompatibleProviderBase<typeof provider.
2424
}
2525

2626
protected getUrl(_model?: AIModel<typeof provider.id>): string | undefined {
27-
return configuration.get('ai.azure.url') ?? undefined;
27+
return ensureOrgConfiguredUrl(this.id, configuration.get('ai.azure.url'));
2828
}
2929

3030
private async getOrPromptBaseUrl(silent: boolean, hasApiKey: boolean): Promise<string | undefined> {
31+
const orgConf = getOrgAIProviderOfType(this.id);
32+
if (!orgConf.enabled) return undefined;
33+
if (orgConf.url) return orgConf.url;
34+
3135
let url: string | undefined = this.getUrl();
3236

3337
if (silent || (url != null && hasApiKey)) return url;

‎src/plus/ai/ollamaProvider.ts

Lines changed: 16 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -5,7 +5,7 @@ import { configuration } from '../../system/-webview/configuration';
55
import type { AIActionType, AIModel } from './models/model';
66
import type { AIChatMessage, AIRequestResult } from './models/provider';
77
import { OpenAICompatibleProviderBase } from './openAICompatibleProviderBase';
8-
import { ensureAccount } from './utils/-webview/ai.utils';
8+
import { ensureAccount, ensureOrgConfiguredUrl, getOrgAIProviderOfType } from './utils/-webview/ai.utils';
99

1010
type OllamaModel = AIModel<typeof provider.id>;
1111

@@ -20,8 +20,12 @@ export class OllamaProvider extends OpenAICompatibleProviderBase<typeof provider
2020
};
2121

2222
override async configured(silent: boolean): Promise<boolean> {
23+
const url = await this.getOrPromptBaseUrl(silent);
24+
if (url === undefined) {
25+
return false;
26+
}
2327
// Ollama doesn't require an API key, but we'll check if the base URL is reachable
24-
return this.validateUrl(await this.getOrPromptBaseUrl(silent), silent);
28+
return this.validateUrl(url, silent);
2529
}
2630

2731
override async getApiKey(silent: boolean): Promise<string | undefined> {
@@ -77,7 +81,11 @@ export class OllamaProvider extends OpenAICompatibleProviderBase<typeof provider
7781
return [];
7882
}
7983

80-
private async getOrPromptBaseUrl(silent: boolean): Promise<string> {
84+
private async getOrPromptBaseUrl(silent: boolean): Promise<string | undefined> {
85+
const orgConf = getOrgAIProviderOfType(this.id);
86+
if (!orgConf.enabled) return undefined;
87+
if (orgConf.url) return orgConf.url;
88+
8189
let url = configuration.get('ai.ollama.url') ?? undefined;
8290
if (url) {
8391
if (silent) return url;
@@ -169,13 +177,14 @@ export class OllamaProvider extends OpenAICompatibleProviderBase<typeof provider
169177
}
170178
}
171179

172-
private getBaseUrl(): string {
180+
private getBaseUrl(): string | undefined {
173181
// Get base URL from configuration or use default
174-
return configuration.get('ai.ollama.url') || defaultBaseUrl;
182+
return ensureOrgConfiguredUrl(this.id, configuration.get('ai.ollama.url') || defaultBaseUrl);
175183
}
176184

177-
protected getUrl(_model: AIModel<typeof provider.id>): string {
178-
return `${this.getBaseUrl()}/api/chat`;
185+
protected getUrl(_model: AIModel<typeof provider.id>): string | undefined {
186+
const url = this.getBaseUrl();
187+
return url ? `${url}/api/chat` : undefined;
179188
}
180189

181190
protected override getHeaders<TAction extends AIActionType>(

‎src/plus/ai/openAICompatibleProvider.ts

Lines changed: 6 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -5,7 +5,7 @@ import { configuration } from '../../system/-webview/configuration';
55
import type { AIModel } from './models/model';
66
import { openAIModels } from './models/model';
77
import { OpenAICompatibleProviderBase } from './openAICompatibleProviderBase';
8-
import { isAzureUrl } from './utils/-webview/ai.utils';
8+
import { ensureOrgConfiguredUrl, getOrgAIProviderOfType, isAzureUrl } from './utils/-webview/ai.utils';
99

1010
type OpenAICompatibleModel = AIModel<typeof provider.id>;
1111
const models: OpenAICompatibleModel[] = openAIModels(provider);
@@ -24,10 +24,14 @@ export class OpenAICompatibleProvider extends OpenAICompatibleProviderBase<typeo
2424
}
2525

2626
protected getUrl(_model?: AIModel<typeof provider.id>): string | undefined {
27-
return configuration.get('ai.openaicompatible.url') ?? undefined;
27+
return ensureOrgConfiguredUrl(this.id, configuration.get('ai.openaicompatible.url'));
2828
}
2929

3030
private async getOrPromptBaseUrl(silent: boolean, hasApiKey: boolean): Promise<string | undefined> {
31+
const orgConf = getOrgAIProviderOfType(this.id);
32+
if (!orgConf.enabled) return undefined;
33+
if (orgConf.url) return orgConf.url;
34+
3135
let url: string | undefined = this.getUrl();
3236

3337
if (silent || (url != null && hasApiKey)) return url;

‎src/plus/ai/openAICompatibleProviderBase.ts

Lines changed: 10 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -10,7 +10,12 @@ import { startLogScope } from '../../system/logger.scope';
1010
import type { ServerConnection } from '../gk/serverConnection';
1111
import type { AIActionType, AIModel, AIProviderDescriptor } from './models/model';
1212
import type { AIChatMessage, AIChatMessageRole, AIProvider, AIRequestResult } from './models/provider';
13-
import { getActionName, getOrPromptApiKey, getValidatedTemperature } from './utils/-webview/ai.utils';
13+
import {
14+
getActionName,
15+
getOrgAIProviderOfType,
16+
getOrPromptApiKey,
17+
getValidatedTemperature,
18+
} from './utils/-webview/ai.utils';
1419

1520
export interface AIProviderConfig {
1621
url: string;
@@ -36,6 +41,10 @@ export abstract class OpenAICompatibleProviderBase<T extends AIProviders> implem
3641
}
3742

3843
async getApiKey(silent: boolean): Promise<string | undefined> {
44+
const orgConf = getOrgAIProviderOfType(this.id);
45+
if (!orgConf.enabled) return undefined;
46+
if (orgConf.key) return orgConf.key;
47+
3948
const { keyUrl, keyValidator } = this.config;
4049

4150
return getOrPromptApiKey(

‎src/plus/ai/openaiProvider.ts

Lines changed: 5 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,7 @@ import { configuration } from '../../system/-webview/configuration';
33
import type { AIActionType, AIModel } from './models/model';
44
import { openAIModels } from './models/model';
55
import { OpenAICompatibleProviderBase } from './openAICompatibleProviderBase';
6-
import { isAzureUrl } from './utils/-webview/ai.utils';
6+
import { ensureOrgConfiguredUrl, isAzureUrl } from './utils/-webview/ai.utils';
77

88
type OpenAIModel = AIModel<typeof provider.id>;
99
const models: OpenAIModel[] = openAIModels(provider);
@@ -22,7 +22,10 @@ export class OpenAIProvider extends OpenAICompatibleProviderBase<typeof provider
2222
}
2323

2424
protected getUrl(_model: AIModel<typeof provider.id>): string {
25-
return configuration.get('ai.openai.url') || 'https://api.openai.com/v1/chat/completions';
25+
return (
26+
ensureOrgConfiguredUrl(this.id, configuration.get('ai.openai.url')) ||
27+
'https://api.openai.com/v1/chat/completions'
28+
);
2629
}
2730

2831
protected override getHeaders<TAction extends AIActionType>(

‎src/plus/ai/utils/-webview/ai.utils.ts

Lines changed: 30 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,7 @@ import { getContext } from '../../../../system/-webview/context';
88
import { openSettingsEditor } from '../../../../system/-webview/vscode/editors';
99
import { formatNumeric } from '../../../../system/date';
1010
import { getPossessiveForm, pluralize } from '../../../../system/string';
11+
import type { OrgAIConfig, OrgAIProvider } from '../../../gk/models/organization';
1112
import { ensureAccountQuickPick } from '../../../gk/utils/-webview/acount.utils';
1213
import type { AIActionType, AIModel } from '../../models/model';
1314

@@ -170,6 +171,35 @@ export function isAzureUrl(url: string): boolean {
170171
return url.includes('.azure.com');
171172
}
172173

174+
export function getOrgAIConfig(): OrgAIConfig {
175+
return {
176+
aiEnabled: getContext('gitlens:gk:organization:ai:enabled', true),
177+
enforceAiProviders: getContext('gitlens:gk:organization:ai:enforceProviders', false),
178+
aiProviders: getContext('gitlens:gk:organization:ai:providers', {}),
179+
};
180+
}
181+
182+
export function getOrgAIProviderOfType(type: AIProviders, orgAiConfig?: OrgAIConfig): OrgAIProvider {
183+
orgAiConfig ??= getOrgAIConfig();
184+
if (!orgAiConfig.aiEnabled) return { type: type, enabled: false };
185+
if (!orgAiConfig.enforceAiProviders) return { type: type, enabled: true };
186+
return orgAiConfig.aiProviders[type] ?? { type: type, enabled: false };
187+
}
188+
189+
export function isProviderEnabledByOrg(type: AIProviders, orgAiConfig?: OrgAIConfig): boolean {
190+
return getOrgAIProviderOfType(type, orgAiConfig).enabled;
191+
}
192+
193+
/**
194+
* If the input value (userUrl) matches to the org configuration it returns it.
195+
*/
196+
export function ensureOrgConfiguredUrl(type: AIProviders, userUrl: null | undefined | string): string | undefined {
197+
const provider = getOrgAIProviderOfType(type);
198+
if (!provider.enabled) return undefined;
199+
200+
return provider.url || userUrl || undefined;
201+
}
202+
173203
export async function ensureAccess(options?: { showPicker?: boolean }): Promise<boolean> {
174204
const showPicker = options?.showPicker ?? false;
175205

‎src/plus/gk/models/organization.ts

Lines changed: 72 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,5 @@
1+
import type { AIProviders } from '../../../constants.ai';
2+
13
export interface Organization {
24
readonly id: string;
35
readonly name: string;
@@ -20,7 +22,10 @@ export interface OrganizationMember {
2022
}
2123

2224
export interface OrganizationSettings {
25+
aiEnabled: boolean;
26+
enforceAiProviders: boolean;
2327
aiSettings: OrganizationSetting;
28+
aiProviders: GkDevAIProviders;
2429
draftsSettings: OrganizationDraftsSettings;
2530
}
2631

@@ -39,3 +44,70 @@ export interface OrganizationDraftsSettings extends OrganizationSetting {
3944
}
4045
| undefined;
4146
}
47+
48+
export type GkDevAIProviders = Partial<Record<GkDevAIProviderType, GkDevAIProvider>>;
49+
50+
export interface GkDevAIProvider {
51+
enabled: boolean;
52+
url?: string;
53+
key?: string;
54+
}
55+
56+
export interface OrgAIProvider {
57+
readonly type: AIProviders;
58+
readonly enabled: boolean;
59+
readonly url?: string;
60+
readonly key?: string;
61+
}
62+
63+
export type OrgAIProviders = Partial<Record<AIProviders, OrgAIProvider | undefined>>;
64+
export interface OrgAIConfig {
65+
readonly aiEnabled: boolean;
66+
readonly enforceAiProviders: boolean;
67+
readonly aiProviders: OrgAIProviders;
68+
}
69+
70+
export type GkDevAIProviderType = 'anthropic' | 'azure' | 'gitkraken_ai' | 'openai' | 'openai_compatible';
71+
72+
export function fromGkDevAIProviderType(type: GkDevAIProviderType): AIProviders;
73+
export function fromGkDevAIProviderType(type: Exclude<unknown, GkDevAIProviderType>): never;
74+
export function fromGkDevAIProviderType(type: unknown): AIProviders | never {
75+
switch (type) {
76+
case 'anthropic':
77+
return 'anthropic';
78+
case 'azure':
79+
return 'azure';
80+
case 'gitkraken_ai':
81+
return 'gitkraken';
82+
case 'openai':
83+
return 'openai';
84+
case 'openai_compatible':
85+
return 'openaicompatible';
86+
case 'ollama':
87+
return 'ollama';
88+
default:
89+
throw new Error(`Unknown AI provider type: ${String(type)}`);
90+
}
91+
}
92+
93+
function fromGkDevAIProvider(type: GkDevAIProviderType, provider: GkDevAIProvider): OrgAIProvider {
94+
return {
95+
type: fromGkDevAIProviderType(type),
96+
enabled: provider.enabled,
97+
url: provider.url,
98+
key: provider.key,
99+
};
100+
}
101+
102+
export function fromGKDevAIProviders(providers?: GkDevAIProviders): OrgAIProviders {
103+
const result: OrgAIProviders = {};
104+
if (providers == null) return result;
105+
106+
Object.entries(providers).forEach(([type, provider]) => {
107+
result[fromGkDevAIProviderType(type as GkDevAIProviderType)] = fromGkDevAIProvider(
108+
type as GkDevAIProviderType,
109+
provider,
110+
);
111+
});
112+
return result;
113+
}

‎src/plus/gk/organizationService.ts

Lines changed: 7 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,7 @@ import type {
1111
OrganizationSettings,
1212
OrganizationsResponse,
1313
} from './models/organization';
14+
import { fromGKDevAIProviders } from './models/organization';
1415
import type { ServerConnection } from './serverConnection';
1516
import type { SubscriptionChangeEvent } from './subscriptionService';
1617

@@ -140,7 +141,12 @@ export class OrganizationService implements Disposable {
140141
private async updateOrganizationPermissions(orgId: string | undefined): Promise<void> {
141142
const settings = orgId != null ? await this.getOrganizationSettings(orgId) : undefined;
142143

143-
void setContext('gitlens:gk:organization:ai:enabled', settings?.aiSettings.enabled ?? true);
144+
void setContext(
145+
'gitlens:gk:organization:ai:enabled',
146+
settings?.aiSettings.enabled ?? settings?.aiEnabled ?? true,
147+
);
148+
void setContext('gitlens:gk:organization:ai:enforceProviders', settings?.enforceAiProviders ?? false);
149+
void setContext('gitlens:gk:organization:ai:providers', fromGKDevAIProviders(settings?.aiProviders) ?? {});
144150
void setContext('gitlens:gk:organization:drafts:byob', settings?.draftsSettings.bucket != null);
145151
void setContext('gitlens:gk:organization:drafts:enabled', settings?.draftsSettings.enabled ?? true);
146152
}

0 commit comments

Comments
 (0)