84 lines
2.4 KiB
TypeScript
84 lines
2.4 KiB
TypeScript
import type { Model } from "@mariozechner/pi-ai/dist/types.js";
|
|
import { expect } from "vitest";
|
|
import { makeZeroUsageSnapshot } from "../agents/usage.js";
|
|
|
|
export const asRecord = (value: unknown): Record<string, unknown> => {
|
|
expect(value).toBeTruthy();
|
|
expect(typeof value).toBe("object");
|
|
expect(Array.isArray(value)).toBe(false);
|
|
return value as Record<string, unknown>;
|
|
};
|
|
|
|
type ConvertedTools = ReadonlyArray<{
|
|
functionDeclarations?: ReadonlyArray<{
|
|
parametersJsonSchema?: unknown;
|
|
parameters?: unknown;
|
|
}>;
|
|
}>;
|
|
|
|
export const getFirstToolParameters = (converted: ConvertedTools): Record<string, unknown> => {
|
|
const functionDeclaration = asRecord(converted?.[0]?.functionDeclarations?.[0]);
|
|
return asRecord(functionDeclaration.parametersJsonSchema ?? functionDeclaration.parameters);
|
|
};
|
|
|
|
export const makeModel = (id: string): Model<"google-generative-ai"> =>
|
|
({
|
|
id,
|
|
name: id,
|
|
api: "google-generative-ai",
|
|
provider: "google",
|
|
baseUrl: "https://example.invalid",
|
|
reasoning: false,
|
|
input: ["text"],
|
|
cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 },
|
|
contextWindow: 1,
|
|
maxTokens: 1,
|
|
}) as Model<"google-generative-ai">;
|
|
|
|
export const makeGeminiCliModel = (id: string): Model<"google-gemini-cli"> =>
|
|
({
|
|
id,
|
|
name: id,
|
|
api: "google-gemini-cli",
|
|
provider: "google-gemini-cli",
|
|
baseUrl: "https://example.invalid",
|
|
reasoning: false,
|
|
input: ["text"],
|
|
cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 },
|
|
contextWindow: 1,
|
|
maxTokens: 1,
|
|
}) as Model<"google-gemini-cli">;
|
|
|
|
export function makeGoogleAssistantMessage(model: string, content: unknown) {
|
|
return {
|
|
role: "assistant",
|
|
content,
|
|
api: "google-generative-ai",
|
|
provider: "google",
|
|
model,
|
|
usage: makeZeroUsageSnapshot(),
|
|
stopReason: "stop",
|
|
timestamp: 0,
|
|
};
|
|
}
|
|
|
|
export function makeGeminiCliAssistantMessage(model: string, content: unknown) {
|
|
return {
|
|
role: "assistant",
|
|
content,
|
|
api: "google-gemini-cli",
|
|
provider: "google-gemini-cli",
|
|
model,
|
|
usage: makeZeroUsageSnapshot(),
|
|
stopReason: "stop",
|
|
timestamp: 0,
|
|
};
|
|
}
|
|
|
|
export function expectConvertedRoles(contents: Array<{ role?: string }>, expectedRoles: string[]) {
|
|
expect(contents).toHaveLength(expectedRoles.length);
|
|
for (const [index, role] of expectedRoles.entries()) {
|
|
expect(contents[index]?.role).toBe(role);
|
|
}
|
|
}
|