不当人了

utils/tools/CustomSearchTool.js

@@ -1,99 +1,137 @@
 import { AbstractTool } from './AbstractTool.js';
 import fetch from 'node-fetch';
 import { Config } from '../config.js';
-// import { UrlExtractionTool } from './UrlExtractionTool.js'; // 不再需要引入 UrlExtractionTool
 
 /**
- * 自定义搜索工具类
+ * 自定义搜索工具类 - 使用 Gemini API
  * @class CustomSearchTool
  * @extends {AbstractTool}
  */
 export class CustomSearchTool extends AbstractTool {
-  // 工具名称
   name = 'CustomSearchTool';
 
-  // 工具参数
   parameters = {
     properties: {
       query: {
         type: 'string',
-        description: 'Search keyword', // 修改描述，不再支持 URL
+        description: '搜索关键词',
       },
       length: {
         type: 'integer',
-        description: 'The desired length of the summary in sentences. Defaults to 3.',
+        description: '期望的摘要长度（句子数），默认为3',
       },
     },
     required: ['query'],
   };
 
+  description = '使用 Gemini API 进行智能搜索，提供全面的搜索结果和摘要。支持自定义摘要长度。';
+
   /**
    * 工具执行函数
    * @param {Object} opt - 工具参数
-   * @param {Object} ai - AI对象（未使用）
-   * @returns {Promise<string>} - 搜索结果或摘要
+   * @param {string} opt.query - 搜索关键词
+   * @param {number} [opt.length=3] - 摘要长度
+   * @returns {Promise<Object>} - 包含答案和来源的对象
    */
-  func = async function (opt, ai) {
-    let { query, length = 3 } = opt;
-    if (!query) {
-      return 'The query parameter is required.';
+  func = async function (opt) {
+    const { query, length = 3 } = opt;
+
+    if (!query?.trim()) {
+      throw new Error('搜索关键词不能为空');
     }
 
     try {
-      // 直接使用OpenAI API进行搜索或摘要
-      const result = await searchOrSummarize(query, length);
-
-      console.log(`Search or summarization result: ${result}`);
-
-      // 返回搜索结果或摘要给AI
+      const result = await this.searchWithGemini(query, length);
+      console.log(`搜索结果: ${JSON.stringify(result)}`);
       return result;
     } catch (error) {
-      console.error('Search or summarization failed:', error);
-      return `Search or summarization failed, please check the logs. ${error.message}`;
+      console.error('搜索失败:', error);
+      throw new Error(`搜索失败: ${error.message}`);
     }
   };
 
-  // 工具描述
-  description = 'Use OpenAI API for custom search, providing comprehensive search results or summaries. This tool does not handle URLs.'; // 修改描述，明确不再处理 URL
-}
+  /**
+   * 使用 Gemini API 进行搜索
+   * @param {string} query - 搜索关键词
+   * @param {number} length - 摘要长度
+   * @returns {Promise<Object>} - 包含答案和来源的对象
+   * @private
+   */
+  async searchWithGemini(query, length) {
+    const apiKey = Config.geminiKey;
+    const apiBaseUrl = Config.geminiBaseUrl;
+    const apiUrl = `${apiBaseUrl}/v1beta/models/gemini-2.0-flash-exp:generateContent?key=${apiKey}`;
 
-/**
- * 使用OpenAI API进行搜索或摘要
- * @param {string} query - 搜索关键词
- * @param {number} length - 期望的摘要长度（以句子为单位）
- * @returns {Promise<string>} - 搜索结果或摘要
- */
-async function searchOrSummarize(query, length) {
-  const apiKey = Config.apiKey;
-  const apiBaseUrl = Config.openAiBaseUrl;
-  const apiUrl = `${apiBaseUrl}/chat/completions`;
-  const model = Config.model;
+    if (!apiKey || !apiBaseUrl) {
+      throw new Error('Gemini API 配置缺失');
+    }
 
-  const response = await fetch(apiUrl, {
-    method: 'POST',
-    headers: {
-      'Content-Type': 'application/json',
-      'Authorization': `Bearer ${apiKey}`,
-    },
-    body: JSON.stringify({
-      model: model,
-      messages: [
-        {
-          role: 'system',
-          content: `You are a search and summarization assistant. Please use English to search based on the following keywords and return a summary of ${length} sentences.`, // 简化提示词
-        },
-        {
-          role: 'user',
-          content: `Search: ${query}`, // 简化提示词
-        },
-      ],
-      max_tokens: 1000 * length,
-    }),
-  });
+    const response = await fetch(apiUrl, {
+      method: 'POST',
+      headers: {
+        'Content-Type': 'application/json',
+      },
+      body: JSON.stringify({
+        contents: [{
+          parts: [{
+            text: this.constructPrompt(query, length)
+          }]
+        }],
+        tools: [{
+          googleSearch: {}
+        }]
+      })
+    });
 
-  const data = await response.json();
-  if (data.error) {
-    throw new Error(`OpenAI API Error: ${data.error.message}`);
+    const data = await response.json();
+
+    if (!response.ok) {
+      throw new Error(`API 请求失败: ${data.error?.message || '未知错误'}`);
+    }
+
+    return this.processGeminiResponse(data);
+  }
+
+  /**
+   * 构建提示词
+   * @param {string} query - 搜索关键词
+   * @param {number} length - 摘要长度
+   * @returns {string} - 格式化的提示词
+   * @private
+   */
+  constructPrompt(query, length) {
+    return `Please provide a comprehensive ${length} sentence summary for the following query. 
+    Include relevant facts and information.
+    Query: ${query}`;
+  }
+
+  /**
+   * 处理 Gemini API 响应
+   * @param {Object} data - API 响应数据
+   * @returns {Object} - 处理后的结果对象
+   * @private
+   */
+  processGeminiResponse(data) {
+    if (!data?.candidates?.[0]?.content?.parts?.[0]?.text) {
+      throw new Error('无效的 API 响应');
+    }
+
+    const answer = data.candidates[0].content.parts[0].text;
+    
+    // 提取来源信息
+    const sources = data.candidates?.[0]?.groundingMetadata?.groundingChunks
+      ?.filter(chunk => chunk.web)
+      ?.map(chunk => ({
+        title: chunk.web.title,
+        url: chunk.web.uri
+      }))
+      ?.filter((v, i, a) => 
+        a.findIndex(t => (t.title === v.title && t.url === v.url)) === i
+      ) || [];
+
+    return {
+      answer,
+      sources,
+    };
   }
-  return data.choices[0].message.content;
 }