Newer
Older
TelosDB / tools / gemini-rag-tool / main.mjs
import dotenv from 'dotenv';
import fs from 'fs';
import path from 'path';
import { searchWeb } from './utils/brave-client.mjs';
import { generateText } from './utils/gemini-client.mjs';
import logger from './utils/logger.mjs';
import { scrapeText } from './utils/scraper.mjs';

dotenv.config({ path: new URL('.env', import.meta.url).pathname.replace(/^\/([a-zA-Z]:)/, '$1') });

export function cleanTitle(title) {
    if (!title) return 'untitled';
    let cleaned = title.replace(/[\\/:*?"<>|]/g, ' ');
    if (cleaned.length > 32) cleaned = cleaned.substring(0, 32);
    return cleaned.trim();
}

export function generateFileName(dateStr, index, title) {
    const paddedIndex = String(index).padStart(2, '0');
    return `${dateStr}-${paddedIndex}-${title}.md`;
}

export async function generateKeywords(prompt) {
    const kgPrompt = `質問に対して最適な検索キーワードを5つ抽出してください。カンマ区切りで回答してください。\n\n質問:\n${prompt}`;
    try {
        const keywords = await generateText(kgPrompt);
        return keywords.trim();
    } catch (err) {
        logger.error('Failed: ' + err.message);
        return '';
    }
}

export async function run(prompt) {
    if (!prompt) {
        logger.error('Prompt is required.');
        return;
    }

    // 出力先をプロジェクト直下の docs/references フォルダに設定
    const outputDir = path.join(process.cwd(), 'docs', 'references');
    const maxSources = 3;
    const today = new Date().toISOString().slice(0, 10).replace(/-/g, '');

    if (!fs.existsSync(outputDir)) fs.mkdirSync(outputDir, { recursive: true });

    logger.info('Step 1: Keywords...');
    const keywords = await generateKeywords(prompt);
    if (!keywords) return;

    logger.info('Step 2: Searching...');
    const searchResults = await searchWeb(keywords, maxSources);
    if (!searchResults.length) return;

    logger.info('Step 3: Extracting...');
    const sources = [];
    for (const res of searchResults) {
        const content = await scrapeText(res.url);
        if (content) sources.push({ title: res.title, url: res.url, text: content });
    }

    logger.info('Step 4: Generating...');
    const sourceText = sources.map((s, i) => `[ソース${i + 1}] ${s.title}\nURL: ${s.url}\n内容: ${s.text}`).join('\n---\n');
    const answerPrompt = `ソースドキュメントをもとに回答してください。タイトルは「# タイトル: [内容]」形式で。\n\nプロンプト:\n${prompt}\n\nソース:\n${sourceText}`;
    const fullResponse = await generateText(answerPrompt);

    let title = 'RAG Report';
    const titleMatch = fullResponse.match(/^# タイトル:\s*(.+)$/m);
    if (titleMatch) title = titleMatch[1];

    const files = fs.readdirSync(outputDir);
    const existingNums = files.filter(f => f.startsWith(today)).map(f => parseInt(f.split('-')[1], 10)).filter(n => !isNaN(n));
    const nextIndex = existingNums.length > 0 ? Math.max(...existingNums) + 1 : 1;

    const fileName = generateFileName(today, nextIndex, cleanTitle(title));
    const filePath = path.join(outputDir, fileName);

    fs.writeFileSync(filePath, `${fullResponse}\n\n## 参考資料\n\n` + sources.map(s => `- [${s.title}](${s.url})`).join('\n'), 'utf8');
    logger.info(`✓ Saved to: ${filePath}`);
}

const args = process.argv.slice(2);
if (args.length > 0) run(args.join(' ')).catch(err => logger.error(err.stack));