mirror of
https://github.com/getcompanion-ai/co-mono.git
synced 2026-04-16 18:03:50 +00:00
fix(ai): use model.api instead of hardcoding api type in streaming functions
- anthropic.ts: use model.api instead of hardcoding 'anthropic-messages' - openai-responses.ts: use model.api instead of hardcoding 'openai-responses' - gitlab-duo: simplify to use actual model IDs, export MODELS array
This commit is contained in:
Mario Zechner
parent
7019fdb92b
commit
0d24ddbb03
4 changed files with 162 additions and 515 deletions
|
|
@ -5,16 +5,8 @@
|
|||
* Delegates to pi-ai's built-in Anthropic and OpenAI streaming implementations.
|
||||
*
|
||||
* Usage:
|
||||
* # First install dependencies
|
||||
* cd packages/coding-agent/examples/extensions/gitlab-duo && npm install
|
||||
*
|
||||
* # With OAuth (run /login gitlab-duo first)
|
||||
* pi -e ./packages/coding-agent/examples/extensions/gitlab-duo
|
||||
*
|
||||
* # With PAT
|
||||
* GITLAB_TOKEN=glpat-... pi -e ./packages/coding-agent/examples/extensions/gitlab-duo
|
||||
*
|
||||
* Then use /model to select gitlab-duo/duo-chat-sonnet-4-5
|
||||
* pi -e ./packages/coding-agent/examples/extensions/custom-provider-gitlab-duo
|
||||
* # Then /login gitlab-duo, or set GITLAB_TOKEN=glpat-...
|
||||
*/
|
||||
|
||||
import {
|
||||
|
|
@ -26,7 +18,8 @@ import {
|
|||
type OAuthCredentials,
|
||||
type OAuthLoginCallbacks,
|
||||
type SimpleStreamOptions,
|
||||
streamSimple,
|
||||
streamSimpleAnthropic,
|
||||
streamSimpleOpenAIResponses,
|
||||
} from "@mariozechner/pi-ai";
|
||||
import type { ExtensionAPI } from "@mariozechner/pi-coding-agent";
|
||||
|
||||
|
|
@ -39,42 +32,101 @@ const AI_GATEWAY_URL = "https://cloud.gitlab.com";
|
|||
const ANTHROPIC_PROXY_URL = `${AI_GATEWAY_URL}/ai/v1/proxy/anthropic/`;
|
||||
const OPENAI_PROXY_URL = `${AI_GATEWAY_URL}/ai/v1/proxy/openai/v1`;
|
||||
|
||||
// Bundled OAuth client ID for gitlab.com (from opencode-gitlab-auth, registered with localhost redirect)
|
||||
const BUNDLED_CLIENT_ID = "1d89f9fdb23ee96d4e603201f6861dab6e143c5c3c00469a018a2d94bdc03d4e";
|
||||
const OAUTH_SCOPES = ["api"];
|
||||
const REDIRECT_URI = "http://127.0.0.1:8080/callback";
|
||||
|
||||
// Direct access token cache (25 min, tokens expire after 30 min)
|
||||
const DIRECT_ACCESS_TTL = 25 * 60 * 1000;
|
||||
|
||||
// Model mappings: duo model ID -> backend config
|
||||
type OpenAIApi = "openai-completions" | "openai-responses";
|
||||
const MODEL_MAPPINGS: Record<string, { api: "anthropic-messages" | OpenAIApi; backendModel: string; baseUrl: string }> =
|
||||
// =============================================================================
|
||||
// Models - exported for use by tests
|
||||
// =============================================================================
|
||||
|
||||
type Backend = "anthropic" | "openai";
|
||||
|
||||
interface GitLabModel {
|
||||
id: string;
|
||||
name: string;
|
||||
backend: Backend;
|
||||
baseUrl: string;
|
||||
reasoning: boolean;
|
||||
input: ("text" | "image")[];
|
||||
cost: { input: number; output: number; cacheRead: number; cacheWrite: number };
|
||||
contextWindow: number;
|
||||
maxTokens: number;
|
||||
}
|
||||
|
||||
export const MODELS: GitLabModel[] = [
|
||||
// Anthropic
|
||||
{
|
||||
"duo-chat-opus-4-5": {
|
||||
api: "anthropic-messages",
|
||||
backendModel: "claude-opus-4-5-20251101",
|
||||
baseUrl: ANTHROPIC_PROXY_URL,
|
||||
},
|
||||
"duo-chat-sonnet-4-5": {
|
||||
api: "anthropic-messages",
|
||||
backendModel: "claude-sonnet-4-5-20250929",
|
||||
baseUrl: ANTHROPIC_PROXY_URL,
|
||||
},
|
||||
"duo-chat-haiku-4-5": {
|
||||
api: "anthropic-messages",
|
||||
backendModel: "claude-haiku-4-5-20251001",
|
||||
baseUrl: ANTHROPIC_PROXY_URL,
|
||||
},
|
||||
// All GPT models use Responses API for consistent tool call ID format across model switches
|
||||
"duo-chat-gpt-5-1": { api: "openai-responses", backendModel: "gpt-5.1-2025-11-13", baseUrl: OPENAI_PROXY_URL },
|
||||
"duo-chat-gpt-5-mini": {
|
||||
api: "openai-responses",
|
||||
backendModel: "gpt-5-mini-2025-08-07",
|
||||
baseUrl: OPENAI_PROXY_URL,
|
||||
},
|
||||
"duo-chat-gpt-5-codex": { api: "openai-responses", backendModel: "gpt-5-codex", baseUrl: OPENAI_PROXY_URL },
|
||||
};
|
||||
id: "claude-opus-4-5-20251101",
|
||||
name: "Claude Opus 4.5",
|
||||
backend: "anthropic",
|
||||
baseUrl: ANTHROPIC_PROXY_URL,
|
||||
reasoning: true,
|
||||
input: ["text", "image"],
|
||||
cost: { input: 15, output: 75, cacheRead: 1.5, cacheWrite: 18.75 },
|
||||
contextWindow: 200000,
|
||||
maxTokens: 32000,
|
||||
},
|
||||
{
|
||||
id: "claude-sonnet-4-5-20250929",
|
||||
name: "Claude Sonnet 4.5",
|
||||
backend: "anthropic",
|
||||
baseUrl: ANTHROPIC_PROXY_URL,
|
||||
reasoning: true,
|
||||
input: ["text", "image"],
|
||||
cost: { input: 3, output: 15, cacheRead: 0.3, cacheWrite: 3.75 },
|
||||
contextWindow: 200000,
|
||||
maxTokens: 16384,
|
||||
},
|
||||
{
|
||||
id: "claude-haiku-4-5-20251001",
|
||||
name: "Claude Haiku 4.5",
|
||||
backend: "anthropic",
|
||||
baseUrl: ANTHROPIC_PROXY_URL,
|
||||
reasoning: true,
|
||||
input: ["text", "image"],
|
||||
cost: { input: 1, output: 5, cacheRead: 0.1, cacheWrite: 1.25 },
|
||||
contextWindow: 200000,
|
||||
maxTokens: 8192,
|
||||
},
|
||||
// OpenAI (all use Responses API)
|
||||
{
|
||||
id: "gpt-5.1-2025-11-13",
|
||||
name: "GPT-5.1",
|
||||
backend: "openai",
|
||||
baseUrl: OPENAI_PROXY_URL,
|
||||
reasoning: true,
|
||||
input: ["text", "image"],
|
||||
cost: { input: 2.5, output: 10, cacheRead: 0, cacheWrite: 0 },
|
||||
contextWindow: 128000,
|
||||
maxTokens: 16384,
|
||||
},
|
||||
{
|
||||
id: "gpt-5-mini-2025-08-07",
|
||||
name: "GPT-5 Mini",
|
||||
backend: "openai",
|
||||
baseUrl: OPENAI_PROXY_URL,
|
||||
reasoning: true,
|
||||
input: ["text", "image"],
|
||||
cost: { input: 0.15, output: 0.6, cacheRead: 0, cacheWrite: 0 },
|
||||
contextWindow: 128000,
|
||||
maxTokens: 16384,
|
||||
},
|
||||
{
|
||||
id: "gpt-5-codex",
|
||||
name: "GPT-5 Codex",
|
||||
backend: "openai",
|
||||
baseUrl: OPENAI_PROXY_URL,
|
||||
reasoning: true,
|
||||
input: ["text", "image"],
|
||||
cost: { input: 2.5, output: 10, cacheRead: 0, cacheWrite: 0 },
|
||||
contextWindow: 128000,
|
||||
maxTokens: 16384,
|
||||
},
|
||||
];
|
||||
|
||||
const MODEL_MAP = new Map(MODELS.map((m) => [m.id, m]));
|
||||
|
||||
// =============================================================================
|
||||
// Direct Access Token Cache
|
||||
|
|
@ -94,13 +146,9 @@ async function getDirectAccessToken(gitlabAccessToken: string): Promise<DirectAc
|
|||
return cachedDirectAccess;
|
||||
}
|
||||
|
||||
const url = `${GITLAB_COM_URL}/api/v4/ai/third_party_agents/direct_access`;
|
||||
const response = await fetch(url, {
|
||||
const response = await fetch(`${GITLAB_COM_URL}/api/v4/ai/third_party_agents/direct_access`, {
|
||||
method: "POST",
|
||||
headers: {
|
||||
Authorization: `Bearer ${gitlabAccessToken}`,
|
||||
"Content-Type": "application/json",
|
||||
},
|
||||
headers: { Authorization: `Bearer ${gitlabAccessToken}`, "Content-Type": "application/json" },
|
||||
body: JSON.stringify({ feature_flags: { DuoAgentPlatformNext: true } }),
|
||||
});
|
||||
|
||||
|
|
@ -115,11 +163,7 @@ async function getDirectAccessToken(gitlabAccessToken: string): Promise<DirectAc
|
|||
}
|
||||
|
||||
const data = (await response.json()) as { token: string; headers: Record<string, string> };
|
||||
cachedDirectAccess = {
|
||||
token: data.token,
|
||||
headers: data.headers,
|
||||
expiresAt: now + DIRECT_ACCESS_TTL,
|
||||
};
|
||||
cachedDirectAccess = { token: data.token, headers: data.headers, expiresAt: now + DIRECT_ACCESS_TTL };
|
||||
return cachedDirectAccess;
|
||||
}
|
||||
|
||||
|
|
@ -128,7 +172,7 @@ function invalidateDirectAccessToken() {
|
|||
}
|
||||
|
||||
// =============================================================================
|
||||
// OAuth Implementation
|
||||
// OAuth
|
||||
// =============================================================================
|
||||
|
||||
async function generatePKCE(): Promise<{ verifier: string; challenge: string }> {
|
||||
|
|
@ -138,21 +182,16 @@ async function generatePKCE(): Promise<{ verifier: string; challenge: string }>
|
|||
.replace(/\+/g, "-")
|
||||
.replace(/\//g, "_")
|
||||
.replace(/=+$/, "");
|
||||
|
||||
const encoder = new TextEncoder();
|
||||
const data = encoder.encode(verifier);
|
||||
const hash = await crypto.subtle.digest("SHA-256", data);
|
||||
const hash = await crypto.subtle.digest("SHA-256", new TextEncoder().encode(verifier));
|
||||
const challenge = btoa(String.fromCharCode(...new Uint8Array(hash)))
|
||||
.replace(/\+/g, "-")
|
||||
.replace(/\//g, "_")
|
||||
.replace(/=+$/, "");
|
||||
|
||||
return { verifier, challenge };
|
||||
}
|
||||
|
||||
async function loginGitLab(callbacks: OAuthLoginCallbacks): Promise<OAuthCredentials> {
|
||||
const { verifier, challenge } = await generatePKCE();
|
||||
|
||||
const authParams = new URLSearchParams({
|
||||
client_id: BUNDLED_CLIENT_ID,
|
||||
redirect_uri: REDIRECT_URI,
|
||||
|
|
@ -165,9 +204,7 @@ async function loginGitLab(callbacks: OAuthLoginCallbacks): Promise<OAuthCredent
|
|||
|
||||
callbacks.onAuth({ url: `${GITLAB_COM_URL}/oauth/authorize?${authParams.toString()}` });
|
||||
const callbackUrl = await callbacks.onPrompt({ message: "Paste the callback URL:" });
|
||||
|
||||
const urlObj = new URL(callbackUrl);
|
||||
const code = urlObj.searchParams.get("code");
|
||||
const code = new URL(callbackUrl).searchParams.get("code");
|
||||
if (!code) throw new Error("No authorization code found in callback URL");
|
||||
|
||||
const tokenResponse = await fetch(`${GITLAB_COM_URL}/oauth/token`, {
|
||||
|
|
@ -183,14 +220,12 @@ async function loginGitLab(callbacks: OAuthLoginCallbacks): Promise<OAuthCredent
|
|||
});
|
||||
|
||||
if (!tokenResponse.ok) throw new Error(`Token exchange failed: ${await tokenResponse.text()}`);
|
||||
|
||||
const data = (await tokenResponse.json()) as {
|
||||
access_token: string;
|
||||
refresh_token: string;
|
||||
expires_in: number;
|
||||
created_at: number;
|
||||
};
|
||||
|
||||
invalidateDirectAccessToken();
|
||||
return {
|
||||
refresh: data.refresh_token,
|
||||
|
|
@ -209,16 +244,13 @@ async function refreshGitLabToken(credentials: OAuthCredentials): Promise<OAuthC
|
|||
refresh_token: credentials.refresh,
|
||||
}).toString(),
|
||||
});
|
||||
|
||||
if (!response.ok) throw new Error(`Token refresh failed: ${await response.text()}`);
|
||||
|
||||
const data = (await response.json()) as {
|
||||
access_token: string;
|
||||
refresh_token: string;
|
||||
expires_in: number;
|
||||
created_at: number;
|
||||
};
|
||||
|
||||
invalidateDirectAccessToken();
|
||||
return {
|
||||
refresh: data.refresh_token,
|
||||
|
|
@ -228,10 +260,10 @@ async function refreshGitLabToken(credentials: OAuthCredentials): Promise<OAuthC
|
|||
}
|
||||
|
||||
// =============================================================================
|
||||
// Main Stream Function - Delegates to pi-ai's built-in implementations
|
||||
// Stream Function
|
||||
// =============================================================================
|
||||
|
||||
function streamGitLabDuo(
|
||||
export function streamGitLabDuo(
|
||||
model: Model<Api>,
|
||||
context: Context,
|
||||
options?: SimpleStreamOptions,
|
||||
|
|
@ -241,57 +273,22 @@ function streamGitLabDuo(
|
|||
(async () => {
|
||||
try {
|
||||
const gitlabAccessToken = options?.apiKey;
|
||||
if (!gitlabAccessToken) {
|
||||
throw new Error("No GitLab access token. Run /login gitlab-duo or set GITLAB_TOKEN");
|
||||
}
|
||||
if (!gitlabAccessToken) throw new Error("No GitLab access token. Run /login gitlab-duo or set GITLAB_TOKEN");
|
||||
|
||||
const mapping = MODEL_MAPPINGS[model.id];
|
||||
if (!mapping) throw new Error(`Unknown model: ${model.id}`);
|
||||
const cfg = MODEL_MAP.get(model.id);
|
||||
if (!cfg) throw new Error(`Unknown model: ${model.id}`);
|
||||
|
||||
// Get direct access token (cached)
|
||||
const directAccess = await getDirectAccessToken(gitlabAccessToken);
|
||||
const modelWithBaseUrl = { ...model, baseUrl: cfg.baseUrl };
|
||||
const headers = { ...directAccess.headers, Authorization: `Bearer ${directAccess.token}` };
|
||||
const streamOptions = { ...options, apiKey: "gitlab-duo", headers };
|
||||
|
||||
// Create a proxy model that uses the backend API
|
||||
const proxyModel: Model<typeof mapping.api> = {
|
||||
...model,
|
||||
id: mapping.backendModel,
|
||||
api: mapping.api,
|
||||
baseUrl: mapping.baseUrl,
|
||||
};
|
||||
const innerStream =
|
||||
cfg.backend === "anthropic"
|
||||
? streamSimpleAnthropic(modelWithBaseUrl as Model<"anthropic-messages">, context, streamOptions)
|
||||
: streamSimpleOpenAIResponses(modelWithBaseUrl as Model<"openai-responses">, context, streamOptions);
|
||||
|
||||
// Merge GitLab headers with Authorization bearer token
|
||||
const headers = {
|
||||
...directAccess.headers,
|
||||
Authorization: `Bearer ${directAccess.token}`,
|
||||
};
|
||||
|
||||
// Delegate to pi-ai's built-in streaming
|
||||
const innerStream = streamSimple(proxyModel, context, {
|
||||
...options,
|
||||
apiKey: "gitlab-duo", // Dummy value to pass validation
|
||||
headers,
|
||||
});
|
||||
|
||||
// Forward all events
|
||||
for await (const event of innerStream) {
|
||||
// Patch the model info back to gitlab-duo
|
||||
if ("partial" in event && event.partial) {
|
||||
event.partial.api = model.api;
|
||||
event.partial.provider = model.provider;
|
||||
event.partial.model = model.id;
|
||||
}
|
||||
if ("message" in event && event.message) {
|
||||
event.message.api = model.api;
|
||||
event.message.provider = model.provider;
|
||||
event.message.model = model.id;
|
||||
}
|
||||
if ("error" in event && event.error) {
|
||||
event.error.api = model.api;
|
||||
event.error.provider = model.provider;
|
||||
event.error.model = model.id;
|
||||
}
|
||||
stream.push(event);
|
||||
}
|
||||
for await (const event of innerStream) stream.push(event);
|
||||
stream.end();
|
||||
} catch (error) {
|
||||
stream.push({
|
||||
|
|
@ -332,73 +329,21 @@ export default function (pi: ExtensionAPI) {
|
|||
baseUrl: AI_GATEWAY_URL,
|
||||
apiKey: "GITLAB_TOKEN",
|
||||
api: "gitlab-duo-api",
|
||||
|
||||
models: [
|
||||
// Anthropic models
|
||||
{
|
||||
id: "duo-chat-opus-4-5",
|
||||
name: "GitLab Duo Claude Opus 4.5",
|
||||
reasoning: false,
|
||||
input: ["text"],
|
||||
cost: { input: 15, output: 75, cacheRead: 1.5, cacheWrite: 18.75 },
|
||||
contextWindow: 200000,
|
||||
maxTokens: 32000,
|
||||
},
|
||||
{
|
||||
id: "duo-chat-sonnet-4-5",
|
||||
name: "GitLab Duo Claude Sonnet 4.5",
|
||||
reasoning: false,
|
||||
input: ["text"],
|
||||
cost: { input: 3, output: 15, cacheRead: 0.3, cacheWrite: 3.75 },
|
||||
contextWindow: 200000,
|
||||
maxTokens: 16384,
|
||||
},
|
||||
{
|
||||
id: "duo-chat-haiku-4-5",
|
||||
name: "GitLab Duo Claude Haiku 4.5",
|
||||
reasoning: false,
|
||||
input: ["text"],
|
||||
cost: { input: 1, output: 5, cacheRead: 0.1, cacheWrite: 1.25 },
|
||||
contextWindow: 200000,
|
||||
maxTokens: 8192,
|
||||
},
|
||||
// OpenAI models
|
||||
{
|
||||
id: "duo-chat-gpt-5-1",
|
||||
name: "GitLab Duo GPT-5.1",
|
||||
reasoning: false,
|
||||
input: ["text"],
|
||||
cost: { input: 2.5, output: 10, cacheRead: 0, cacheWrite: 0 },
|
||||
contextWindow: 128000,
|
||||
maxTokens: 16384,
|
||||
},
|
||||
{
|
||||
id: "duo-chat-gpt-5-mini",
|
||||
name: "GitLab Duo GPT-5 Mini",
|
||||
reasoning: false,
|
||||
input: ["text"],
|
||||
cost: { input: 0.15, output: 0.6, cacheRead: 0, cacheWrite: 0 },
|
||||
contextWindow: 128000,
|
||||
maxTokens: 16384,
|
||||
},
|
||||
{
|
||||
id: "duo-chat-gpt-5-codex",
|
||||
name: "GitLab Duo GPT-5 Codex",
|
||||
reasoning: false,
|
||||
input: ["text"],
|
||||
cost: { input: 2.5, output: 10, cacheRead: 0, cacheWrite: 0 },
|
||||
contextWindow: 128000,
|
||||
maxTokens: 16384,
|
||||
},
|
||||
],
|
||||
|
||||
models: MODELS.map(({ id, name, reasoning, input, cost, contextWindow, maxTokens }) => ({
|
||||
id,
|
||||
name,
|
||||
reasoning,
|
||||
input,
|
||||
cost,
|
||||
contextWindow,
|
||||
maxTokens,
|
||||
})),
|
||||
oauth: {
|
||||
name: "GitLab Duo",
|
||||
login: loginGitLab,
|
||||
refreshToken: refreshGitLabToken,
|
||||
getApiKey: (cred) => cred.access,
|
||||
},
|
||||
|
||||
streamSimple: streamGitLabDuo,
|
||||
});
|
||||
}
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue