mirror of
https://github.com/infiniflow/ragflow.git
synced 2026-07-21 15:11:08 +08:00
Refactor: reformat all code for lefthook using ruff and gofmt (#16585)
This commit is contained in:
Wang Qi
@@ -39,71 +39,52 @@ class ExcelProcessorParam(ComponentParamBase):
|
||||
"""
|
||||
Define the ExcelProcessor component parameters.
|
||||
"""
|
||||
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
# Input configuration
|
||||
self.input_files = [] # Variable references to uploaded files
|
||||
self.operation = "read" # read, merge, transform, output
|
||||
|
||||
|
||||
# Processing options
|
||||
self.sheet_selection = "all" # all, first, or comma-separated sheet names
|
||||
self.merge_strategy = "concat" # concat, join
|
||||
self.join_on = "" # Column name for join operations
|
||||
|
||||
|
||||
# Transform options (for LLM-guided transformations)
|
||||
self.transform_instructions = ""
|
||||
self.transform_data = "" # Variable reference to transformation data
|
||||
|
||||
|
||||
# Output options
|
||||
self.output_format = "xlsx" # xlsx, csv
|
||||
self.output_filename = "output"
|
||||
|
||||
|
||||
# Component outputs
|
||||
self.outputs = {
|
||||
"data": {
|
||||
"type": "object",
|
||||
"value": {}
|
||||
},
|
||||
"summary": {
|
||||
"type": "str",
|
||||
"value": ""
|
||||
},
|
||||
"markdown": {
|
||||
"type": "str",
|
||||
"value": ""
|
||||
}
|
||||
}
|
||||
|
||||
self.outputs = {"data": {"type": "object", "value": {}}, "summary": {"type": "str", "value": ""}, "markdown": {"type": "str", "value": ""}}
|
||||
|
||||
def check(self):
|
||||
self.check_valid_value(
|
||||
self.operation,
|
||||
"[ExcelProcessor] Operation",
|
||||
["read", "merge", "transform", "output"]
|
||||
)
|
||||
self.check_valid_value(
|
||||
self.output_format,
|
||||
"[ExcelProcessor] Output format",
|
||||
["xlsx", "csv"]
|
||||
)
|
||||
self.check_valid_value(self.operation, "[ExcelProcessor] Operation", ["read", "merge", "transform", "output"])
|
||||
self.check_valid_value(self.output_format, "[ExcelProcessor] Output format", ["xlsx", "csv"])
|
||||
return True
|
||||
|
||||
|
||||
class ExcelProcessor(ComponentBase, ABC):
|
||||
"""
|
||||
Excel processing component for RAGFlow agents.
|
||||
|
||||
|
||||
Operations:
|
||||
- read: Parse Excel files into structured data
|
||||
- merge: Combine multiple Excel files
|
||||
- transform: Apply data transformations based on instructions
|
||||
- output: Generate Excel file output
|
||||
"""
|
||||
|
||||
component_name = "ExcelProcessor"
|
||||
|
||||
def get_input_form(self) -> dict[str, dict]:
|
||||
"""Define input form for the component."""
|
||||
res = {}
|
||||
for ref in (self._param.input_files or []):
|
||||
for ref in self._param.input_files or []:
|
||||
for k, o in self.get_input_elements_from_text(ref).items():
|
||||
res[k] = {"name": o.get("name", ""), "type": "file"}
|
||||
if self._param.transform_data:
|
||||
@@ -111,13 +92,13 @@ class ExcelProcessor(ComponentBase, ABC):
|
||||
res[k] = {"name": o.get("name", ""), "type": "object"}
|
||||
return res
|
||||
|
||||
@timeout(int(os.environ.get("COMPONENT_EXEC_TIMEOUT", 10*60)))
|
||||
@timeout(int(os.environ.get("COMPONENT_EXEC_TIMEOUT", 10 * 60)))
|
||||
def _invoke(self, **kwargs):
|
||||
if self.check_if_canceled("ExcelProcessor processing"):
|
||||
return
|
||||
|
||||
operation = self._param.operation.lower()
|
||||
|
||||
|
||||
if operation == "read":
|
||||
self._read_excels()
|
||||
elif operation == "merge":
|
||||
@@ -137,7 +118,7 @@ class ExcelProcessor(ComponentBase, ABC):
|
||||
value = self._canvas.get_variable_value(file_ref)
|
||||
if value is None:
|
||||
return None, None
|
||||
|
||||
|
||||
# Handle different value formats
|
||||
if isinstance(value, dict):
|
||||
# File reference from Begin/UserFillUp component
|
||||
@@ -154,12 +135,13 @@ class ExcelProcessor(ComponentBase, ABC):
|
||||
# Could be base64 encoded or a path
|
||||
if value.startswith("data:"):
|
||||
import base64
|
||||
|
||||
# Extract base64 content
|
||||
_, encoded = value.split(",", 1)
|
||||
return base64.b64decode(encoded), "uploaded.xlsx"
|
||||
|
||||
|
||||
return None, None
|
||||
|
||||
|
||||
def _get_file_content_from_list(self, item) -> tuple[bytes, str]:
|
||||
"""Extract file content from a list item."""
|
||||
if isinstance(item, dict):
|
||||
@@ -170,15 +152,15 @@ class ExcelProcessor(ComponentBase, ABC):
|
||||
"""Parse Excel content into a dictionary of DataFrames (one per sheet)."""
|
||||
try:
|
||||
excel_file = BytesIO(content)
|
||||
|
||||
|
||||
if filename.lower().endswith(".csv"):
|
||||
df = pd.read_csv(excel_file)
|
||||
return {"Sheet1": df}
|
||||
else:
|
||||
# Read all sheets
|
||||
xlsx = pd.ExcelFile(excel_file, engine='openpyxl')
|
||||
xlsx = pd.ExcelFile(excel_file, engine="openpyxl")
|
||||
sheet_selection = self._param.sheet_selection
|
||||
|
||||
|
||||
if sheet_selection == "all":
|
||||
sheets_to_read = xlsx.sheet_names
|
||||
elif sheet_selection == "first":
|
||||
@@ -187,12 +169,12 @@ class ExcelProcessor(ComponentBase, ABC):
|
||||
# Comma-separated sheet names
|
||||
requested = [s.strip() for s in sheet_selection.split(",")]
|
||||
sheets_to_read = [s for s in requested if s in xlsx.sheet_names]
|
||||
|
||||
|
||||
dfs = {}
|
||||
for sheet in sheets_to_read:
|
||||
dfs[sheet] = pd.read_excel(xlsx, sheet_name=sheet)
|
||||
return dfs
|
||||
|
||||
|
||||
except Exception as e:
|
||||
logging.error(f"Error parsing Excel file {filename}: {e}")
|
||||
return {}
|
||||
@@ -202,36 +184,36 @@ class ExcelProcessor(ComponentBase, ABC):
|
||||
all_data = {}
|
||||
summaries = []
|
||||
markdown_parts = []
|
||||
|
||||
for file_ref in (self._param.input_files or []):
|
||||
|
||||
for file_ref in self._param.input_files or []:
|
||||
if self.check_if_canceled("ExcelProcessor reading"):
|
||||
return
|
||||
|
||||
|
||||
# Get variable value
|
||||
value = self._canvas.get_variable_value(file_ref)
|
||||
self.set_input_value(file_ref, str(value)[:200] if value else "")
|
||||
|
||||
|
||||
if value is None:
|
||||
continue
|
||||
|
||||
|
||||
# Handle file content
|
||||
content, filename = self._get_file_content(file_ref)
|
||||
if content is None:
|
||||
continue
|
||||
|
||||
|
||||
# Parse Excel
|
||||
dfs = self._parse_excel_to_dataframes(content, filename)
|
||||
|
||||
|
||||
for sheet_name, df in dfs.items():
|
||||
key = f"{filename}_{sheet_name}" if len(dfs) > 1 else filename
|
||||
all_data[key] = df.to_dict(orient="records")
|
||||
|
||||
|
||||
# Build summary
|
||||
summaries.append(f"**{key}**: {len(df)} rows, {len(df.columns)} columns ({', '.join(df.columns.tolist()[:5])}{'...' if len(df.columns) > 5 else ''})")
|
||||
|
||||
|
||||
# Build markdown table
|
||||
markdown_parts.append(f"### {key}\n\n{df.head(10).to_markdown(index=False)}\n")
|
||||
|
||||
|
||||
# Set outputs
|
||||
self.set_output("data", all_data)
|
||||
self.set_output("summary", "\n".join(summaries) if summaries else "No Excel files found")
|
||||
@@ -240,29 +222,29 @@ class ExcelProcessor(ComponentBase, ABC):
|
||||
def _merge_excels(self):
|
||||
"""Merge multiple Excel files/sheets into one."""
|
||||
all_dfs = []
|
||||
|
||||
for file_ref in (self._param.input_files or []):
|
||||
|
||||
for file_ref in self._param.input_files or []:
|
||||
if self.check_if_canceled("ExcelProcessor merging"):
|
||||
return
|
||||
|
||||
|
||||
value = self._canvas.get_variable_value(file_ref)
|
||||
self.set_input_value(file_ref, str(value)[:200] if value else "")
|
||||
|
||||
|
||||
if value is None:
|
||||
continue
|
||||
|
||||
|
||||
content, filename = self._get_file_content(file_ref)
|
||||
if content is None:
|
||||
continue
|
||||
|
||||
|
||||
dfs = self._parse_excel_to_dataframes(content, filename)
|
||||
all_dfs.extend(dfs.values())
|
||||
|
||||
|
||||
if not all_dfs:
|
||||
self.set_output("data", {})
|
||||
self.set_output("summary", "No data to merge")
|
||||
return
|
||||
|
||||
|
||||
# Merge strategy
|
||||
if self._param.merge_strategy == "concat":
|
||||
merged_df = pd.concat(all_dfs, ignore_index=True)
|
||||
@@ -273,7 +255,7 @@ class ExcelProcessor(ComponentBase, ABC):
|
||||
merged_df = merged_df.merge(df, on=self._param.join_on, how="outer")
|
||||
else:
|
||||
merged_df = pd.concat(all_dfs, ignore_index=True)
|
||||
|
||||
|
||||
self.set_output("data", {"merged": merged_df.to_dict(orient="records")})
|
||||
self.set_output("summary", f"Merged {len(all_dfs)} sources into {len(merged_df)} rows, {len(merged_df.columns)} columns")
|
||||
self.set_output("markdown", merged_df.head(20).to_markdown(index=False))
|
||||
@@ -285,14 +267,14 @@ class ExcelProcessor(ComponentBase, ABC):
|
||||
if not transform_ref:
|
||||
self.set_output("summary", "No transform data reference provided")
|
||||
return
|
||||
|
||||
|
||||
data = self._canvas.get_variable_value(transform_ref)
|
||||
self.set_input_value(transform_ref, str(data)[:300] if data else "")
|
||||
|
||||
|
||||
if data is None:
|
||||
self.set_output("summary", "Transform data is empty")
|
||||
return
|
||||
|
||||
|
||||
# Convert to DataFrame
|
||||
if isinstance(data, dict):
|
||||
# Could be {"sheet": [rows]} format
|
||||
@@ -315,7 +297,7 @@ class ExcelProcessor(ComponentBase, ABC):
|
||||
else:
|
||||
self.set_output("data", {"raw": str(data)})
|
||||
self.set_output("markdown", str(data))
|
||||
|
||||
|
||||
self.set_output("summary", "Transformed data ready for processing")
|
||||
|
||||
def _output_excel(self):
|
||||
@@ -325,14 +307,14 @@ class ExcelProcessor(ComponentBase, ABC):
|
||||
if not transform_ref:
|
||||
self.set_output("summary", "No data reference for output")
|
||||
return
|
||||
|
||||
|
||||
data = self._canvas.get_variable_value(transform_ref)
|
||||
self.set_input_value(transform_ref, str(data)[:300] if data else "")
|
||||
|
||||
|
||||
if data is None:
|
||||
self.set_output("summary", "No data to output")
|
||||
return
|
||||
|
||||
|
||||
try:
|
||||
# Prepare DataFrames
|
||||
if isinstance(data, dict):
|
||||
@@ -346,10 +328,10 @@ class ExcelProcessor(ComponentBase, ABC):
|
||||
else:
|
||||
self.set_output("summary", "Invalid data format for Excel output")
|
||||
return
|
||||
|
||||
|
||||
# Generate output
|
||||
doc_id = get_uuid()
|
||||
|
||||
|
||||
if self._param.output_format == "csv":
|
||||
# For CSV, only output first sheet
|
||||
first_df = list(dfs.values())[0]
|
||||
@@ -358,7 +340,7 @@ class ExcelProcessor(ComponentBase, ABC):
|
||||
else:
|
||||
# Excel output
|
||||
excel_io = BytesIO()
|
||||
with pd.ExcelWriter(excel_io, engine='openpyxl') as writer:
|
||||
with pd.ExcelWriter(excel_io, engine="openpyxl") as writer:
|
||||
for sheet_name, df in dfs.items():
|
||||
# Sanitize sheet name (max 31 chars, no special chars)
|
||||
safe_name = sheet_name[:31].replace("/", "_").replace("\\", "_")
|
||||
@@ -366,23 +348,19 @@ class ExcelProcessor(ComponentBase, ABC):
|
||||
excel_io.seek(0)
|
||||
binary_content = excel_io.read()
|
||||
filename = f"{self._param.output_filename}.xlsx"
|
||||
|
||||
|
||||
# Store file
|
||||
settings.STORAGE_IMPL.put(self._canvas._tenant_id, doc_id, binary_content)
|
||||
|
||||
|
||||
# Set attachment output
|
||||
self.set_output("attachment", {
|
||||
"doc_id": doc_id,
|
||||
"format": self._param.output_format,
|
||||
"file_name": filename
|
||||
})
|
||||
|
||||
self.set_output("attachment", {"doc_id": doc_id, "format": self._param.output_format, "file_name": filename})
|
||||
|
||||
total_rows = sum(len(df) for df in dfs.values())
|
||||
self.set_output("summary", f"Generated {filename} with {len(dfs)} sheet(s), {total_rows} total rows")
|
||||
self.set_output("data", {k: v.to_dict(orient="records") for k, v in dfs.items()})
|
||||
|
||||
|
||||
logging.info(f"ExcelProcessor: Generated {filename} as {doc_id}")
|
||||
|
||||
|
||||
except Exception as e:
|
||||
logging.error(f"ExcelProcessor output error: {e}")
|
||||
self.set_output("summary", f"Error generating output: {str(e)}")
|
||||
|
||||
Reference in New Issue
Block a user