shenju-enterprise-search-mcp/lib/mcp-server.js

绅聚企业资料检索MCP服务

/**
 * MCP服务器实现
 * 遵循Model Context Protocol规范
 */
const { v4: uuidv4 } = require('uuid');
const SearchClient = require('./search-client');

class MCPServer {
  constructor(apiKey, appId, baseUrl) {
    this.searchClient = new SearchClient(apiKey, appId, baseUrl);
    this.sessions = new Map();
    
    // 定义MCP工具
    this.tools = [
      {
        name: "document_search",
        description: "在企业资料库中搜索相关文档，返回最匹配的内容",
        input_schema: {
          type: "object",
          properties: {
            query: {
              type: "string",
              description: "检索查询内容"
            },
            filter: {
              type: "string",
              description: "过滤条件，如'部门=人力资源'",
              default: ""
            },
            limit: {
              type: "number",
              description: "返回结果数量上限",
              default: 10
            },
            include_content: {
              type: "boolean",
              description: "是否包含文档内容",
              default: false
            }
          },
          required: ["query"]
        }
      },
      {
        name: "document_qa",
        description: "根据企业资料库回答问题，提供精确答案",
        input_schema: {
          type: "object",
          properties: {
            question: {
              type: "string",
              description: "待回答的问题"
            },
            filter: {
              type: "string",
              description: "资料过滤条件，如'文档类型=政策'",
              default: ""
            },
            include_sources: {
              type: "boolean",
              description: "是否包含来源文档信息",
              default: true
            }
          },
          required: ["question"]
        }
      }
    ];
  }

  /**
   * 获取工具列表
   * @returns {Object} 工具列表对象
   */
  listTools() {
    return {
      tools: this.tools
    };
  }

  /**
   * 调用指定工具
   * @param {string} toolName 工具名称
   * @param {Object} parameters 参数对象
   * @returns {Promise<Object>} 工具调用结果
   */
  async callTool(toolName, parameters = {}) {
    // 兼容处理：如果第一个参数是sessionId，则使用旧版API格式
    let sessionId = null;
    if (arguments.length >= 2 && typeof arguments[0] === 'string' && typeof arguments[1] === 'string') {
      // 旧格式: callTool(sessionId, toolName, parameters)
      sessionId = arguments[0];
      toolName = arguments[1];
      parameters = arguments[2] || {};
      console.log(`使用旧格式API调用: sessionId=${sessionId}, toolName=${toolName}`);
    }
    
    // 如果没有sessionId，生成一个新的
    if (!sessionId) {
      sessionId = uuidv4();
      console.log(`为工具调用生成新的会话ID: ${sessionId}`);
    }

    // 获取或创建会话
    if (!this.sessions.has(sessionId)) {
      this.sessions.set(sessionId, {
        id: sessionId,
        createdAt: new Date(),
        lastUsedAt: new Date(),
        history: []
      });
      console.log(`创建新会话: ${sessionId}`);
    } else {
      const session = this.sessions.get(sessionId);
      session.lastUsedAt = new Date();
      console.log(`更新现有会话: ${sessionId}, 历史记录数: ${session.history.length}`);
    }

    // 查找对应工具
    const tool = this.tools.find(t => t.name === toolName);
    if (!tool) {
      console.error(`未找到工具: ${toolName}`);
      throw new Error(`未找到名为 ${toolName} 的工具`);
    }

    try {
      console.log(`开始调用工具: ${toolName}, 参数:`, JSON.stringify(parameters));
      let result;

      // 根据工具类型调用不同API
      switch (toolName) {
        case 'document_search': {
          const { query, filter = '', limit = 10, include_content = false } = parameters;
          
          // 获取当前会话
          const session = this.sessions.get(sessionId);
          
          // 添加查询到历史
          session.history.push({
            type: 'search',
            query,
            timestamp: new Date()
          });
          
          // 调用API
          console.log(`调用文档检索API, 查询: ${query}, 过滤器: ${filter}`);
          const response = await this.searchClient.documentSearch(
            query,
            { filter, limit, includeContent: include_content }
          );
          
          // 添加结果到历史
          session.history.push({
            type: 'search_result',
            resultCount: response.results?.length || 0,
            timestamp: new Date()
          });
          
          result = {
            result: {
              results: response.results || [],
              total: response.total || 0,
              session_id: sessionId
            }
          };
          break;
        }
        
        case 'document_qa': {
          const { question, filter = '', include_sources = true } = parameters;
          
          // 获取当前会话
          const session = this.sessions.get(sessionId);
          
          // 添加问题到历史
          session.history.push({
            type: 'question',
            question,
            timestamp: new Date()
          });
          
          // 调用API
          console.log(`调用问答API, 问题: ${question}, 过滤器: ${filter}`);
          const response = await this.searchClient.documentQA(
            question,
            { filter, includeSources: include_sources }
          );
          
          // 添加回答到历史
          session.history.push({
            type: 'answer',
            timestamp: new Date()
          });
          
          result = {
            result: {
              answer: response.answer,
              source_documents: response.sources || [],
              session_id: sessionId
            }
          };
          break;
        }
        
        default:
          console.error(`工具 ${toolName} 未实现`);
          throw new Error(`工具 ${toolName} 未实现`);
      }

      console.log(`工具调用成功: ${toolName}`);
      return result;
    } catch (error) {
      console.error(`调用工具 ${toolName} 失败:`, error);
      throw new Error(`工具调用失败: ${error.message}`);
    }
  }

  /**
   * 清理过期会话
   * 默认清理2小时未使用的会话
   */
  cleanSessions(maxAge = 7200000) {
    const now = new Date();
    let cleanCount = 0;
    
    for (const [sessionId, session] of this.sessions.entries()) {
      const lastUsedTime = session.lastUsedAt.getTime();
      if (now.getTime() - lastUsedTime > maxAge) {
        this.sessions.delete(sessionId);
        cleanCount++;
      }
    }
    
    if (cleanCount > 0) {
      console.log(`已清理 ${cleanCount} 个过期会话，当前会话数: ${this.sessions.size}`);
    }
  }
}

module.exports = MCPServer;