32 lines
930 B
TypeScript
32 lines
930 B
TypeScript
import { GoogleGenerativeAI } from '@google/generative-ai';
|
|
import { getConfig } from '@/lib/db';
|
|
|
|
const genAI = new GoogleGenerativeAI(process.env.GEMINI_API_KEY || '');
|
|
|
|
export async function getModel(): Promise<string> {
|
|
return (await getConfig('gemini_model')) || 'gemini-1.5-pro';
|
|
}
|
|
|
|
export async function generate(
|
|
systemPrompt: string,
|
|
userPrompt: string
|
|
): Promise<{ text: string; tokens: number }> {
|
|
const modelName = await getModel();
|
|
const model = genAI.getGenerativeModel({
|
|
model: modelName,
|
|
systemInstruction: systemPrompt,
|
|
});
|
|
|
|
const result = await model.generateContent(userPrompt);
|
|
const response = result.response;
|
|
const text = response.text();
|
|
|
|
const inputTokens = response.usageMetadata?.promptTokenCount ?? 0;
|
|
const outputTokens = response.usageMetadata?.candidatesTokenCount ?? 0;
|
|
const tokens = inputTokens + outputTokens;
|
|
|
|
return { text, tokens };
|
|
}
|
|
|
|
export default genAI;
|