#!/usr/bin/env python3 """ Excel和CSV文件操作MCP服务器 支持Excel/CSV文件的读取、搜索、枚举值获取等操作 参考multi_keyword_search_server.py的实现方式 """ import json import os import sys import asyncio import re import chardet from typing import Any, Dict, List, Optional, Union import pandas as pd from mcp_common import ( get_allowed_directory, load_tools_from_json, resolve_file_path, find_file_in_project, is_regex_pattern, compile_pattern, create_error_response, create_success_response, create_initialize_response, create_ping_response, create_tools_list_response, handle_mcp_streaming ) def detect_encoding(file_path: str) -> str: """检测文件编码""" try: with open(file_path, 'rb') as f: raw_data = f.read(10000) # 读取前10KB来检测编码 result = chardet.detect(raw_data) return result['encoding'] or 'utf-8' except: return 'utf-8' class ExcelCSVOperator: """Excel和CSV文件操作核心类""" def __init__(self): self.supported_extensions = ['.xlsx', '.xls', '.csv'] self.encoding_cache = {} def _validate_file(self, file_path: str) -> str: """验证并处理文件路径""" # 解析文件路径,支持 folder/document.txt 和 document.txt 格式 resolved_path = resolve_file_path(file_path) # 验证文件扩展名 file_ext = os.path.splitext(resolved_path)[1].lower() if file_ext not in self.supported_extensions: raise ValueError(f"Unsupported file format: {file_ext}, supported formats: {self.supported_extensions}") return resolved_path def load_data(self, file_path: str, sheet_name: str = None) -> pd.DataFrame: """加载Excel或CSV文件数据""" file_path = self._validate_file(file_path) file_ext = os.path.splitext(file_path)[1].lower() try: if file_ext == '.csv': encoding = detect_encoding(file_path) df = pd.read_csv(file_path, encoding=encoding) else: # Excel文件 if sheet_name: df = pd.read_excel(file_path, sheet_name=sheet_name) else: # 读取第一个sheet df = pd.read_excel(file_path) # 处理空值 df = df.fillna('') return df except Exception as e: raise ValueError(f"File loading failed: {str(e)}") def get_sheets(self, file_path: str) -> List[str]: """获取Excel文件的所有sheet名称""" file_path = self._validate_file(file_path) file_ext = os.path.splitext(file_path)[1].lower() if file_ext == '.csv': return ['default'] # CSV文件只有一个默认sheet try: excel_file = pd.ExcelFile(file_path) return excel_file.sheet_names except Exception as e: raise ValueError(f"Failed to read Excel sheet list: {str(e)}") def get_schema(self, file_path: str, sheet_name: str = None) -> List[str]: """获取文件的schema字段列表""" try: df = self.load_data(file_path, sheet_name) return df.columns.tolist() except Exception as e: raise ValueError(f"Failed to get schema: {str(e)}") def full_text_search(self, file_path: str, keywords: List[str], top_k: int = 10, case_sensitive: bool = False) -> str: """全文搜索功能""" if not keywords: return "Error: Keyword list cannot be empty" # 预处理和验证关键词中的正则表达式 valid_keywords = [] regex_errors = [] for keyword in keywords: compiled = compile_pattern(keyword) if compiled is None: regex_errors.append(keyword) else: valid_keywords.append(keyword) if regex_errors: error_msg = f"Warning: The following regular expressions failed to compile and will be ignored: {', '.join(regex_errors)}" print(error_msg) if not valid_keywords: return "Error: No valid search keywords" try: # 验证文件路径 validated_path = self._validate_file(file_path) file_ext = os.path.splitext(validated_path)[1].lower() all_results = [] if file_ext == '.csv': # CSV文件只有一个数据集 results = self._search_in_file(validated_path, valid_keywords, case_sensitive, 'default') all_results.extend(results) else: # Excel文件,搜索所有sheet sheets = self.get_sheets(validated_path) for sheet in sheets: results = self._search_in_file(validated_path, valid_keywords, case_sensitive, sheet) all_results.extend(results) # 按匹配数量排序(降序) all_results.sort(key=lambda x: x['match_count'], reverse=True) # 限制结果数量 limited_results = all_results[:top_k] # 格式化为CSV输出 if not limited_results: return "No matching results found" # 构建CSV格式输出 csv_lines = [] headers = ["sheet", "row_index", "match_count", "matched_content", "match_details"] csv_lines.append(",".join(headers)) for result in limited_results: # 转义CSV中的特殊字符 sheet = str(result.get('sheet', '')).replace(',', ',') row_index = str(result.get('row_index', '')) match_count = str(result.get('match_count', 0)) matched_content = str(result.get('matched_content', '')).replace(',', ',').replace('\n', ' ') match_details = str(result.get('match_details', '')).replace(',', ',') csv_lines.append(f"{sheet},{row_index},{match_count},{matched_content},{match_details}") return "\n".join(csv_lines) except Exception as e: return f"Search failed: {str(e)}" def _search_in_file(self, file_path: str, keywords: List[str], case_sensitive: bool, sheet_name: str = None) -> List[Dict[str, Any]]: """在文件中搜索关键词""" results = [] try: df = self.load_data(file_path, sheet_name) # 预处理所有模式 processed_patterns = [] for keyword in keywords: compiled = compile_pattern(keyword) if compiled is not None: processed_patterns.append({ 'original': keyword, 'pattern': compiled, 'is_regex': isinstance(compiled, re.Pattern) }) # 逐行搜索 for row_index, row in df.iterrows(): # 将整行内容合并为字符串进行搜索 row_content = " ".join([str(cell) for cell in row.values if str(cell).strip()]) search_content = row_content if case_sensitive else row_content.lower() # 统计匹配的模式数量 matched_patterns = [] for pattern_info in processed_patterns: pattern = pattern_info['pattern'] is_regex = pattern_info['is_regex'] match_found = False match_details = None if is_regex: # 正则表达式匹配 if case_sensitive: match = pattern.search(row_content) else: # 对于不区分大小写的正则,需要重新编译 if isinstance(pattern, re.Pattern): flags = pattern.flags | re.IGNORECASE case_insensitive_pattern = re.compile(pattern.pattern, flags) match = case_insensitive_pattern.search(row_content) else: match = pattern.search(search_content) if match: match_found = True match_details = match.group(0) else: # 普通字符串匹配 search_keyword = pattern if case_sensitive else pattern.lower() if search_keyword in search_content: match_found = True match_details = pattern if match_found: matched_patterns.append({ 'original': pattern_info['original'], 'type': 'regex' if is_regex else 'keyword', 'match': match_details }) match_count = len(matched_patterns) if match_count > 0: # 构建匹配详情 match_details = [] for pattern in matched_patterns: if pattern['type'] == 'regex': match_details.append(f"[regex:{pattern['original']}={pattern['match']}]") else: match_details.append(f"[keyword:{pattern['match']}]") match_info = " ".join(match_details) results.append({ 'sheet': sheet_name, 'row_index': row_index, 'match_count': match_count, 'matched_content': row_content, 'match_details': match_info }) except Exception as e: print(f"Error searching file {file_path} (sheet: {sheet_name}): {str(e)}") return results def filter_search(self, file_path: str, filters: Dict, sheet_name: str = None) -> str: """字段过滤搜索功能""" if not filters: return "Error: Filter conditions cannot be empty" try: df = self.load_data(file_path, sheet_name) # 应用过滤条件 filtered_df = df.copy() for field_name, filter_condition in filters.items(): if field_name not in df.columns: return f"Error: Field '{field_name}' does not exist" operator = filter_condition.get('operator', 'eq') value = filter_condition.get('value') if operator == 'eq': # 等于 filtered_df = filtered_df[filtered_df[field_name] == value] elif operator == 'gt': # 大于 filtered_df = filtered_df[filtered_df[field_name] > value] elif operator == 'lt': # 小于 filtered_df = filtered_df[filtered_df[field_name] < value] elif operator == 'gte': # 大于等于 filtered_df = filtered_df[filtered_df[field_name] >= value] elif operator == 'lte': # 小于等于 filtered_df = filtered_df[filtered_df[field_name] <= value] elif operator == 'contains': # 包含 filtered_df = filtered_df[filtered_df[field_name].astype(str).str.contains(str(value), na=False)] elif operator == 'regex': # 正则表达式 try: pattern = re.compile(str(value)) filtered_df = filtered_df[filtered_df[field_name].astype(str).str.match(pattern, na=False)] except re.error as e: return f"Error: Regular expression '{value}' compilation failed: {str(e)}" else: return f"Error: Unsupported operator '{operator}'" # 格式化为CSV输出 if filtered_df.empty: return "No records matching conditions found" # 转换为CSV字符串 csv_result = filtered_df.to_csv(index=False, encoding='utf-8') return csv_result except Exception as e: return f"Filter search failed: {str(e)}" def get_field_enums(self, file_path: str, field_names: List[str], sheet_name: str = None, max_enum_count: int = 100, min_occurrence: int = 1) -> str: """获取指定字段的枚举值列表""" if not field_names: return "Error: Field name list cannot be empty" try: df = self.load_data(file_path, sheet_name) # 验证字段存在性 missing_fields = [field for field in field_names if field not in df.columns] if missing_fields: return f"Error: Fields do not exist: {', '.join(missing_fields)}" # 计算每个字段的枚举值 enum_results = {} for field in field_names: # 统计值出现次数 value_counts = df[field].value_counts() # 过滤出现次数过少的值 filtered_counts = value_counts[value_counts >= min_occurrence] # 限制返回数量 top_values = filtered_counts.head(max_enum_count) # 格式化结果 enum_values = [] for value, count in top_values.items(): enum_values.append(f"{value}({count})") enum_results[field] = { 'enum_values': enum_values, 'total_unique': len(value_counts), 'total_filtered': len(filtered_counts), 'total_rows': len(df) } # 格式化输出 output_lines = [] for field, data in enum_results.items(): enum_str = ", ".join(data['enum_values']) field_info = f"{field}: [{enum_str}] (总计: {data['total_unique']}个唯一值, 过滤后: {data['total_filtered']}个, 总行数: {data['total_rows']})" output_lines.append(field_info) return "\n".join(output_lines) except Exception as e: return f"Failed to get enum values: {str(e)}" # 全局操作器实例 operator = ExcelCSVOperator() async def handle_request(request: Dict[str, Any]) -> Dict[str, Any]: """Handle MCP request""" try: method = request.get("method") params = request.get("params", {}) request_id = request.get("id") if method == "initialize": return create_initialize_response(request_id, "excel-csv-operator") elif method == "ping": return create_ping_response(request_id) elif method == "tools/list": # 从 JSON 文件加载工具定义 tools = load_tools_from_json("excel_csv_operator_tools.json") return create_tools_list_response(request_id, tools) elif method == "tools/call": tool_name = params.get("name") arguments = params.get("arguments", {}) if tool_name == "get_excel_sheets": file_path = arguments.get("file_path") result = operator.get_sheets(file_path) return create_success_response(request_id, { "content": [ { "type": "text", "text": json.dumps(result, ensure_ascii=False, indent=2) } ] }) elif tool_name == "get_table_schema": file_path = arguments.get("file_path") sheet_name = arguments.get("sheet_name") result = operator.get_schema(file_path, sheet_name) return create_success_response(request_id, { "content": [ { "type": "text", "text": json.dumps(result, ensure_ascii=False, indent=2) } ] }) elif tool_name == "full_text_search": file_path = arguments.get("file_path") keywords = arguments.get("keywords", []) top_k = arguments.get("top_k", 10) case_sensitive = arguments.get("case_sensitive", False) result = operator.full_text_search(file_path, keywords, top_k, case_sensitive) return create_success_response(request_id, { "content": [ { "type": "text", "text": result } ] }) elif tool_name == "filter_search": file_path = arguments.get("file_path") sheet_name = arguments.get("sheet_name") filters = arguments.get("filters") result = operator.filter_search(file_path, filters, sheet_name) return create_success_response(request_id, { "content": [ { "type": "text", "text": result } ] }) elif tool_name == "get_field_enums": file_path = arguments.get("file_path") sheet_name = arguments.get("sheet_name") field_names = arguments.get("field_names", []) max_enum_count = arguments.get("max_enum_count", 100) min_occurrence = arguments.get("min_occurrence", 1) result = operator.get_field_enums(file_path, field_names, sheet_name, max_enum_count, min_occurrence) return create_success_response(request_id, { "content": [ { "type": "text", "text": result } ] }) else: return create_error_response(request_id, -32601, f"Unknown tool: {tool_name}") else: return create_error_response(request_id, -32601, f"Unknown method: {method}") except Exception as e: return create_error_response(request.get("id"), -32603, f"Internal error: {str(e)}") async def main(): """Main entry point.""" await handle_mcp_streaming(handle_request) if __name__ == "__main__": asyncio.run(main())