mirror of
https://github.com/getcompanion-ai/co-mono.git
synced 2026-04-15 13:03:42 +00:00
380 lines
11 KiB
TypeScript
380 lines
11 KiB
TypeScript
/**
|
|
* Test script for GitLab Duo - full streaming flow
|
|
* Run: npx tsx test.ts [model-id]
|
|
*
|
|
* Examples:
|
|
* npx tsx test.ts # Test default (duo-chat-sonnet-4-5)
|
|
* npx tsx test.ts duo-chat-gpt-5-codex # Test GPT-5 Codex (Responses API)
|
|
* npx tsx test.ts duo-chat-gpt-5-1 # Test GPT-5.1 (Chat Completions API)
|
|
*/
|
|
|
|
import {
|
|
type Api,
|
|
type AssistantMessageEventStream,
|
|
type Context,
|
|
createAssistantMessageEventStream,
|
|
type Model,
|
|
registerApiProvider,
|
|
type SimpleStreamOptions,
|
|
streamSimple,
|
|
} from "@mariozechner/pi-ai";
|
|
import { readFileSync } from "fs";
|
|
import { homedir } from "os";
|
|
import { join } from "path";
|
|
|
|
// =============================================================================
|
|
// Constants (copied from index.ts)
|
|
// =============================================================================
|
|
|
|
const GITLAB_COM_URL = "https://gitlab.com";
|
|
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`;
|
|
const DIRECT_ACCESS_TTL = 25 * 60 * 1000;
|
|
|
|
type OpenAIApi = "openai-completions" | "openai-responses";
|
|
const MODEL_MAPPINGS: Record<string, { api: "anthropic-messages" | OpenAIApi; backendModel: string; baseUrl: string }> =
|
|
{
|
|
"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 },
|
|
};
|
|
|
|
// Model definitions for cost tracking
|
|
const MODEL_DEFS: Record<
|
|
string,
|
|
{
|
|
name: string;
|
|
cost: { input: number; output: number; cacheRead: number; cacheWrite: number };
|
|
contextWindow: number;
|
|
maxTokens: number;
|
|
}
|
|
> = {
|
|
"duo-chat-opus-4-5": {
|
|
name: "GitLab Duo Claude Opus 4.5",
|
|
cost: { input: 15, output: 75, cacheRead: 1.5, cacheWrite: 18.75 },
|
|
contextWindow: 200000,
|
|
maxTokens: 32000,
|
|
},
|
|
"duo-chat-sonnet-4-5": {
|
|
name: "GitLab Duo Claude Sonnet 4.5",
|
|
cost: { input: 3, output: 15, cacheRead: 0.3, cacheWrite: 3.75 },
|
|
contextWindow: 200000,
|
|
maxTokens: 16384,
|
|
},
|
|
"duo-chat-haiku-4-5": {
|
|
name: "GitLab Duo Claude Haiku 4.5",
|
|
cost: { input: 1, output: 5, cacheRead: 0.1, cacheWrite: 1.25 },
|
|
contextWindow: 200000,
|
|
maxTokens: 8192,
|
|
},
|
|
"duo-chat-gpt-5-1": {
|
|
name: "GitLab Duo GPT-5.1",
|
|
cost: { input: 2.5, output: 10, cacheRead: 0, cacheWrite: 0 },
|
|
contextWindow: 128000,
|
|
maxTokens: 16384,
|
|
},
|
|
"duo-chat-gpt-5-mini": {
|
|
name: "GitLab Duo GPT-5 Mini",
|
|
cost: { input: 0.15, output: 0.6, cacheRead: 0, cacheWrite: 0 },
|
|
contextWindow: 128000,
|
|
maxTokens: 16384,
|
|
},
|
|
"duo-chat-gpt-5-codex": {
|
|
name: "GitLab Duo GPT-5 Codex",
|
|
cost: { input: 2.5, output: 10, cacheRead: 0, cacheWrite: 0 },
|
|
contextWindow: 128000,
|
|
maxTokens: 16384,
|
|
},
|
|
};
|
|
|
|
// =============================================================================
|
|
// Direct Access Token
|
|
// =============================================================================
|
|
|
|
interface DirectAccessToken {
|
|
token: string;
|
|
headers: Record<string, string>;
|
|
expiresAt: number;
|
|
}
|
|
|
|
let cachedDirectAccess: DirectAccessToken | null = null;
|
|
|
|
async function getDirectAccessToken(gitlabAccessToken: string): Promise<DirectAccessToken> {
|
|
const now = Date.now();
|
|
if (cachedDirectAccess && cachedDirectAccess.expiresAt > now) {
|
|
return cachedDirectAccess;
|
|
}
|
|
|
|
const url = `${GITLAB_COM_URL}/api/v4/ai/third_party_agents/direct_access`;
|
|
console.log("Fetching direct access token from:", url);
|
|
|
|
const response = await fetch(url, {
|
|
method: "POST",
|
|
headers: {
|
|
Authorization: `Bearer ${gitlabAccessToken}`,
|
|
"Content-Type": "application/json",
|
|
},
|
|
body: JSON.stringify({ feature_flags: { DuoAgentPlatformNext: true } }),
|
|
});
|
|
|
|
if (!response.ok) {
|
|
const errorText = await response.text();
|
|
throw new Error(`Failed to get direct access token: ${response.status} ${errorText}`);
|
|
}
|
|
|
|
const data = (await response.json()) as { token: string; headers: Record<string, string> };
|
|
console.log("Got direct access token");
|
|
cachedDirectAccess = {
|
|
token: data.token,
|
|
headers: data.headers,
|
|
expiresAt: now + DIRECT_ACCESS_TTL,
|
|
};
|
|
return cachedDirectAccess;
|
|
}
|
|
|
|
// =============================================================================
|
|
// Stream Function (copied from index.ts)
|
|
// =============================================================================
|
|
|
|
function streamGitLabDuo(
|
|
model: Model<Api>,
|
|
context: Context,
|
|
options?: SimpleStreamOptions,
|
|
): AssistantMessageEventStream {
|
|
const stream = createAssistantMessageEventStream();
|
|
|
|
(async () => {
|
|
try {
|
|
const gitlabAccessToken = options?.apiKey;
|
|
console.log("streamGitLabDuo called");
|
|
console.log(" model.id:", model.id);
|
|
console.log(" options.apiKey present:", !!gitlabAccessToken);
|
|
|
|
if (!gitlabAccessToken) {
|
|
throw new Error("No GitLab access token provided in options.apiKey");
|
|
}
|
|
|
|
const mapping = MODEL_MAPPINGS[model.id];
|
|
if (!mapping) throw new Error(`Unknown model: ${model.id}`);
|
|
|
|
// Get direct access token (cached)
|
|
const directAccess = await getDirectAccessToken(gitlabAccessToken);
|
|
|
|
// 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,
|
|
};
|
|
|
|
// Merge GitLab headers with Authorization bearer token
|
|
const headers = {
|
|
...directAccess.headers,
|
|
Authorization: `Bearer ${directAccess.token}`,
|
|
};
|
|
|
|
console.log("Calling streamSimple with proxy model:");
|
|
console.log(" proxyModel.id:", proxyModel.id);
|
|
console.log(" proxyModel.api:", proxyModel.api);
|
|
console.log(" proxyModel.baseUrl:", proxyModel.baseUrl);
|
|
|
|
// 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);
|
|
}
|
|
stream.end();
|
|
} catch (error) {
|
|
console.error("Stream error:", error);
|
|
stream.push({
|
|
type: "error",
|
|
reason: "error",
|
|
error: {
|
|
role: "assistant",
|
|
content: [],
|
|
api: model.api,
|
|
provider: model.provider,
|
|
model: model.id,
|
|
usage: {
|
|
input: 0,
|
|
output: 0,
|
|
cacheRead: 0,
|
|
cacheWrite: 0,
|
|
totalTokens: 0,
|
|
cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 },
|
|
},
|
|
stopReason: "error",
|
|
errorMessage: error instanceof Error ? error.message : String(error),
|
|
timestamp: Date.now(),
|
|
},
|
|
});
|
|
stream.end();
|
|
}
|
|
})();
|
|
|
|
return stream;
|
|
}
|
|
|
|
// =============================================================================
|
|
// Main Test
|
|
// =============================================================================
|
|
|
|
interface AuthData {
|
|
[provider: string]: {
|
|
type: "oauth" | "api_key";
|
|
refresh?: string;
|
|
access?: string;
|
|
expires?: number;
|
|
key?: string;
|
|
};
|
|
}
|
|
|
|
async function main() {
|
|
const modelId = process.argv[2] || "duo-chat-sonnet-4-5";
|
|
|
|
if (!MODEL_MAPPINGS[modelId]) {
|
|
console.error(`Unknown model: ${modelId}`);
|
|
console.error("Available models:", Object.keys(MODEL_MAPPINGS).join(", "));
|
|
process.exit(1);
|
|
}
|
|
|
|
// Read auth.json
|
|
const authPath = join(homedir(), ".pi", "agent", "auth.json");
|
|
console.log("Reading auth from:", authPath);
|
|
|
|
let authData: AuthData;
|
|
try {
|
|
authData = JSON.parse(readFileSync(authPath, "utf-8"));
|
|
} catch (e) {
|
|
console.error("Failed to read auth.json:", e);
|
|
process.exit(1);
|
|
}
|
|
|
|
const gitlabCred = authData["gitlab-duo"];
|
|
if (!gitlabCred || gitlabCred.type !== "oauth" || !gitlabCred.access) {
|
|
console.error("No gitlab-duo OAuth credentials found. Run /login gitlab-duo first.");
|
|
process.exit(1);
|
|
}
|
|
|
|
console.log("Found gitlab-duo OAuth credentials");
|
|
const gitlabAccessToken = gitlabCred.access;
|
|
|
|
// Register our custom API provider
|
|
console.log("\nRegistering gitlab-duo-api provider...");
|
|
registerApiProvider({
|
|
api: "gitlab-duo-api" as Api,
|
|
stream: (model, context, options) => streamGitLabDuo(model, context, options as SimpleStreamOptions),
|
|
streamSimple: streamGitLabDuo,
|
|
});
|
|
|
|
// Create a test model
|
|
const modelDef = MODEL_DEFS[modelId];
|
|
const mapping = MODEL_MAPPINGS[modelId];
|
|
// Enable reasoning for Anthropic models
|
|
const supportsReasoning = mapping.api === "anthropic-messages";
|
|
const testModel: Model<Api> = {
|
|
id: modelId,
|
|
name: modelDef.name,
|
|
api: "gitlab-duo-api" as Api,
|
|
provider: "gitlab-duo",
|
|
baseUrl: AI_GATEWAY_URL,
|
|
reasoning: supportsReasoning,
|
|
input: ["text"],
|
|
cost: modelDef.cost,
|
|
contextWindow: modelDef.contextWindow,
|
|
maxTokens: modelDef.maxTokens,
|
|
};
|
|
|
|
// Create test context
|
|
const context: Context = {
|
|
messages: [{ role: "user", content: "Say hello in exactly 3 words.", timestamp: Date.now() }],
|
|
};
|
|
|
|
// Check for --thinking flag
|
|
const useThinking = process.argv.includes("--thinking");
|
|
|
|
console.log("\nStarting stream test...");
|
|
console.log("Model:", testModel.id);
|
|
console.log("Backend:", MODEL_MAPPINGS[modelId].backendModel);
|
|
console.log("API:", MODEL_MAPPINGS[modelId].api);
|
|
console.log(
|
|
"Reasoning:",
|
|
supportsReasoning ? (useThinking ? "enabled" : "supported but not enabled") : "not supported",
|
|
);
|
|
console.log("Prompt:", context.messages[0].content);
|
|
console.log("");
|
|
|
|
// Call streamSimple
|
|
const stream = streamSimple(testModel, context, {
|
|
apiKey: gitlabAccessToken,
|
|
maxTokens: 100,
|
|
reasoning: useThinking && supportsReasoning ? "low" : undefined,
|
|
});
|
|
|
|
// Consume the stream
|
|
let inThinking = false;
|
|
for await (const event of stream) {
|
|
if (event.type === "thinking_start") {
|
|
inThinking = true;
|
|
console.log("[Thinking]");
|
|
} else if (event.type === "thinking_delta") {
|
|
process.stdout.write(event.delta);
|
|
} else if (event.type === "thinking_end") {
|
|
inThinking = false;
|
|
console.log("\n[/Thinking]\n");
|
|
} else if (event.type === "text_delta") {
|
|
process.stdout.write(event.delta);
|
|
} else if (event.type === "error") {
|
|
console.error("\nError:", event.error.errorMessage);
|
|
} else if (event.type === "done") {
|
|
console.log("\n\nDone! Stop reason:", event.reason);
|
|
console.log("Usage:", event.message.usage);
|
|
}
|
|
}
|
|
}
|
|
|
|
main().catch(console.error);
|