83 lines
2.1 KiB
TypeScript
83 lines
2.1 KiB
TypeScript
import type { LlmAdapter } from "./types.js";
|
|
|
|
interface GeminiResponse {
|
|
candidates: Array<{ content: { parts: Array<{ text: string }> } }>;
|
|
usageMetadata: {
|
|
promptTokenCount: number;
|
|
candidatesTokenCount: number;
|
|
totalTokenCount: number;
|
|
};
|
|
}
|
|
|
|
export class GeminiAdapter implements LlmAdapter {
|
|
private apiKey: string;
|
|
private model: string;
|
|
|
|
constructor(apiKey: string, model: string) {
|
|
this.apiKey = apiKey;
|
|
this.model = model;
|
|
}
|
|
|
|
async call(
|
|
words: string[],
|
|
systemPrompt: string,
|
|
): Promise<{
|
|
content: string;
|
|
promptTokens: number;
|
|
completionTokens: number;
|
|
totalTokens: number;
|
|
promptTimeMs: number;
|
|
completionTimeMs: number;
|
|
}> {
|
|
const url = `https://generativelanguage.googleapis.com/v1beta/models/${this.model}:generateContent?key=${this.apiKey}`;
|
|
|
|
const payload = {
|
|
contents: [
|
|
{
|
|
role: "user",
|
|
parts: [
|
|
{ text: systemPrompt + "\n\nWords: " + JSON.stringify(words) },
|
|
],
|
|
},
|
|
],
|
|
generationConfig: {
|
|
temperature: 0.1,
|
|
topP: 0.9,
|
|
maxOutputTokens: Math.ceil(words.length * 250 * 1.2),
|
|
},
|
|
};
|
|
|
|
const startTime = Date.now();
|
|
|
|
const response = await fetch(url, {
|
|
method: "POST",
|
|
headers: { "Content-Type": "application/json" },
|
|
body: JSON.stringify(payload),
|
|
});
|
|
|
|
const totalTimeMs = Date.now() - startTime;
|
|
|
|
if (!response.ok) {
|
|
throw new Error(`Gemini API responded with status: ${response.status}`);
|
|
}
|
|
|
|
const json = (await response.json()) as GeminiResponse;
|
|
|
|
const content = json.candidates[0]?.content?.parts[0]?.text;
|
|
if (!content) {
|
|
throw new Error("Gemini response content is empty");
|
|
}
|
|
|
|
const promptTokens = json.usageMetadata.promptTokenCount;
|
|
const completionTokens = json.usageMetadata.candidatesTokenCount;
|
|
|
|
return {
|
|
content,
|
|
promptTokens,
|
|
completionTokens,
|
|
totalTokens: json.usageMetadata.totalTokenCount,
|
|
promptTimeMs: totalTimeMs * 0.3,
|
|
completionTimeMs: totalTimeMs * 0.7,
|
|
};
|
|
}
|
|
}
|