Newer
Older
ai-team / system / llm / gemini.js
@楽曲作りまくりおじさん 楽曲作りまくりおじさん on 31 May 2024 2 KB first commit
import { GoogleGenerativeAI } from "@google/generative-ai";
import { searchDuckDuckGo } from './externalsearch.js';

// Access your API key (see "Set up your API key" above)
const genAI = new GoogleGenerativeAI(process.env.GEMINI_API_KEY);

export function generateModel() {
    const generationConfig = {
        temperature: process.env.GEMINI_TEMPERATURE,
        maxOutputTokens: process.env.GEMINI_MAX_OUTPUT_TOKENS,
        topK: process.env.GEMINI_TOP_K,
        topP: process.env.GEMINI_TOP_P
    };

    // For text-only input, use the gemini-pro model
    const model = genAI.getGenerativeModel({
        model:
            process.env.GEMINI_MODEL, generationConfig
    }, [
        { "category": "HARM_CATEGORY_HARASSMENT", "threshold": "BLOCK_NONE" },
        { "category": "HARM_CATEGORY_HATE_SPEECH", "threshold": "BLOCK_NONE" },
        { "category": "HARM_CATEGORY_SEXUALLY_EXPLICIT", "threshold": "BLOCK_NONE" },
        { "category": "HARM_CATEGORY_DANGEROUS_CONTENT", "threshold": "BLOCK_NONE" }
    ]
    );
    return model;
}

async function queryGemini(prompt) {
    try {
        const model = generateModel();
        let result = await model.generateContent(prompt);
        const response = result.response;
        const text = response.text();
        return text;
    } catch (e) {
        console.log(e.message);
        throw e;
    }
}

// プロンプトを作成する。TODOあとでsystemの追加ロジックを組込む。
export async function createPrompt(token) {

    let prompt = {
        "system": [],
        "data": [],
        "chart": [],
        "currentToken": {"role": "user", "content": "${token}"},
        "refinfo": ""
    };

    let externalInfo = await searchDuckDuckGo(token);
    if (externalInfo !== undefined && externalInfo.length > 0) {
        for (let item of externalInfo) {
            prompt.data.push(item);
            prompt.refinfo += `\n\n[${item.title}](${item.link}) `;
        }
    }
    return prompt;
}

export async function getAnswer(prompt) {
    // 質問の回答を取得する。
    let pastPrompt = [];
    pastPrompt = [
        ...prompt.system,
        ...prompt.data,
        ...prompt.chart
    ];
    pastPrompt.push(prompt.currentToken)

    let newPrompt = JSON.stringify(pastPrompt);
    let replyMessage = await queryGemini(newPrompt);
    if (prompt.refinfo !== '') {
        replyMessage += `\n\n**参考**\n${prompt.refinfo}`;
    }
    return replyMessage;
}