chore: refactor docx_tool to reduce function size (#6273)

This commit is contained in:
Bradley Axen
2026-01-05 20:41:07 -08:00
committed by GitHub
parent f6042fa6fe
commit 41dbddade6
@@ -11,7 +11,7 @@ enum UpdateMode {
old_text: String,
},
InsertStructured {
level: Option<String>, // e.g., "Heading1", "Heading2", etc.
level: Option<String>,
style: Option<DocxStyle>,
},
AddImage {
@@ -46,18 +46,11 @@ impl DocxStyle {
alignment: obj
.get("alignment")
.and_then(|v| v.as_str())
.and_then(|a| match a {
"left" => Some(AlignmentType::Left),
"center" => Some(AlignmentType::Center),
"right" => Some(AlignmentType::Right),
"justified" => Some(AlignmentType::Both),
_ => None,
}),
.and_then(parse_alignment),
})
}
fn apply_to_run(&self, run: Run) -> Run {
let mut run = run;
fn apply_to_run(&self, mut run: Run) -> Run {
if self.bold {
run = run.bold();
}
@@ -76,8 +69,7 @@ impl DocxStyle {
run
}
fn apply_to_paragraph(&self, para: Paragraph) -> Paragraph {
let mut para = para;
fn apply_to_paragraph(&self, mut para: Paragraph) -> Paragraph {
if let Some(alignment) = self.alignment {
para = para.align(alignment);
}
@@ -85,6 +77,361 @@ impl DocxStyle {
}
}
fn parse_alignment(a: &str) -> Option<AlignmentType> {
match a {
"left" => Some(AlignmentType::Left),
"center" => Some(AlignmentType::Center),
"right" => Some(AlignmentType::Right),
"justified" => Some(AlignmentType::Both),
_ => None,
}
}
fn docx_error(message: impl Into<String>) -> ErrorData {
ErrorData {
code: ErrorCode::INTERNAL_ERROR,
message: Cow::from(message.into()),
data: None,
}
}
fn invalid_params(message: impl Into<String>) -> ErrorData {
ErrorData {
code: ErrorCode::INVALID_PARAMS,
message: Cow::from(message.into()),
data: None,
}
}
fn read_docx_file(path: &str) -> Result<Docx, ErrorData> {
let file =
fs::read(path).map_err(|e| docx_error(format!("Failed to read DOCX file: {}", e)))?;
read_docx(&file).map_err(|e| docx_error(format!("Failed to parse DOCX file: {}", e)))
}
fn read_or_create_docx(path: &str) -> Result<Docx, ErrorData> {
if std::path::Path::new(path).exists() {
read_docx_file(path)
} else {
Ok(Docx::new())
}
}
fn write_docx_file(path: &str, doc: Docx) -> Result<(), ErrorData> {
let mut buf = Vec::new();
doc.build()
.pack(&mut Cursor::new(&mut buf))
.map_err(|e| docx_error(format!("Failed to build DOCX: {}", e)))?;
fs::write(path, &buf).map_err(|e| docx_error(format!("Failed to write DOCX file: {}", e)))
}
fn extract_paragraph_text(p: &Paragraph) -> String {
p.children
.iter()
.filter_map(|child| {
if let ParagraphChild::Run(run) = child {
Some(
run.children
.iter()
.filter_map(|rc| {
if let RunChild::Text(t) = rc {
Some(t.text.clone())
} else {
None
}
})
.collect::<Vec<_>>()
.join(""),
)
} else {
None
}
})
.collect::<Vec<_>>()
.join("")
}
fn add_styled_paragraphs(mut doc: Docx, content: &str, style: &Option<DocxStyle>) -> Docx {
for para in content.split('\n').filter(|p| !p.trim().is_empty()) {
let mut run = Run::new().add_text(para);
let mut paragraph = Paragraph::new();
if let Some(s) = style {
run = s.apply_to_run(run);
paragraph = s.apply_to_paragraph(paragraph);
}
doc = doc.add_paragraph(paragraph.add_run(run));
}
doc
}
fn parse_update_mode(
params: Option<&serde_json::Value>,
) -> Result<(UpdateMode, Option<DocxStyle>), ErrorData> {
let Some(params) = params else {
return Ok((UpdateMode::Append, None));
};
let mode_str = params
.get("mode")
.and_then(|v| v.as_str())
.unwrap_or("append");
let style = params.get("style").and_then(DocxStyle::from_json);
let mode = match mode_str {
"append" => UpdateMode::Append,
"replace" => {
let old_text = params
.get("old_text")
.and_then(|v| v.as_str())
.ok_or_else(|| invalid_params("old_text parameter required for replace mode"))?;
UpdateMode::Replace {
old_text: old_text.to_string(),
}
}
"structured" => UpdateMode::InsertStructured {
level: params
.get("level")
.and_then(|v| v.as_str())
.map(String::from),
style: style.clone(),
},
"add_image" => {
let image_path = params
.get("image_path")
.and_then(|v| v.as_str())
.ok_or_else(|| {
invalid_params("image_path parameter required for add_image mode")
})?;
UpdateMode::AddImage {
image_path: image_path.to_string(),
width: params
.get("width")
.and_then(|v| v.as_u64())
.map(|w| w as u32),
height: params
.get("height")
.and_then(|v| v.as_u64())
.map(|h| h as u32),
}
}
_ => {
return Err(invalid_params(
"Invalid mode. Must be 'append', 'replace', 'structured', or 'add_image'",
))
}
};
Ok((mode, style))
}
fn extract_text_from_docx(docx: &Docx) -> String {
let mut text = String::new();
for element in docx.document.children.iter() {
if let DocumentChild::Paragraph(p) = element {
let para_text = extract_paragraph_text(p);
if !para_text.trim().is_empty() {
text.push_str(&para_text);
text.push('\n');
}
}
}
text
}
fn extract_structure_from_docx(docx: &Docx) -> Vec<String> {
let mut structure = Vec::new();
let mut current_level = None;
for element in docx.document.children.iter() {
if let DocumentChild::Paragraph(p) = element {
if let Some(style) = p.property.style.as_ref() {
if style.val.starts_with("Heading") {
current_level = Some(style.val.clone());
structure.push(format!("{}: ", style.val));
}
}
let para_text = extract_paragraph_text(p);
if !para_text.trim().is_empty() && current_level.is_some() {
if let Some(s) = structure.last_mut() {
s.push_str(&para_text);
}
current_level = None;
}
}
}
structure
}
fn do_extract_text(path: &str) -> Result<Vec<Content>, ErrorData> {
let docx = read_docx_file(path)?;
let text = extract_text_from_docx(&docx);
let structure = extract_structure_from_docx(&docx);
let result = if !structure.is_empty() {
format!(
"Document Structure:\n{}\n\nFull Text:\n{}",
structure.join("\n"),
text
)
} else {
format!("Extracted Text:\n{}", text)
};
Ok(vec![Content::text(result)])
}
fn do_append(
path: &str,
content: &str,
style: &Option<DocxStyle>,
) -> Result<Vec<Content>, ErrorData> {
let doc = read_or_create_docx(path)?;
let doc = add_styled_paragraphs(doc, content, style);
write_docx_file(path, doc)?;
Ok(vec![Content::text(format!(
"Successfully wrote content to {}",
path
))])
}
fn do_replace(
path: &str,
content: &str,
old_text: &str,
style: &Option<DocxStyle>,
) -> Result<Vec<Content>, ErrorData> {
let docx = read_docx_file(path)?;
let mut new_doc = Docx::new();
let mut found_text = false;
for element in docx.document.children.iter() {
if let DocumentChild::Paragraph(p) = element {
let para_text = extract_paragraph_text(p);
if para_text.contains(old_text) {
found_text = true;
new_doc = add_styled_paragraphs(new_doc, content, style);
} else {
let mut para = Paragraph::new();
if let Some(s) = &p.property.style {
para = para.style(&s.val);
}
for child in p.children.iter() {
if let ParagraphChild::Run(run) = child {
for rc in run.children.iter() {
if let RunChild::Text(t) = rc {
para = para.add_run(Run::new().add_text(&t.text));
}
}
}
}
new_doc = new_doc.add_paragraph(para);
}
}
}
if !found_text {
return Err(docx_error(format!(
"Could not find text to replace: {}",
old_text
)));
}
write_docx_file(path, new_doc)?;
Ok(vec![Content::text(format!(
"Successfully replaced content in {}",
path
))])
}
fn do_insert_structured(
path: &str,
content: &str,
level: &Option<String>,
style: &Option<DocxStyle>,
) -> Result<Vec<Content>, ErrorData> {
let mut doc = read_or_create_docx(path)?;
for para in content.split('\n').filter(|p| !p.trim().is_empty()) {
let mut run = Run::new().add_text(para);
let mut paragraph = Paragraph::new();
if let Some(lvl) = level {
paragraph = paragraph.style(lvl);
}
if let Some(s) = style {
run = s.apply_to_run(run);
paragraph = s.apply_to_paragraph(paragraph);
}
doc = doc.add_paragraph(paragraph.add_run(run));
}
write_docx_file(path, doc)?;
Ok(vec![Content::text(format!(
"Successfully added structured content to {}",
path
))])
}
fn load_image_as_png(image_path: &str) -> Result<Vec<u8>, ErrorData> {
let image_data = fs::read(image_path)
.map_err(|e| docx_error(format!("Failed to read image file: {}", e)))?;
let extension = std::path::Path::new(image_path)
.extension()
.and_then(|e| e.to_str())
.ok_or_else(|| docx_error("Invalid image file extension"))?
.to_lowercase();
if extension == "png" {
return Ok(image_data);
}
let img = image::load_from_memory(&image_data)
.map_err(|e| docx_error(format!("Failed to load image: {}", e)))?;
let mut png_data = Vec::new();
img.write_to(&mut Cursor::new(&mut png_data), ImageFormat::Png)
.map_err(|e| docx_error(format!("Failed to convert image to PNG: {}", e)))?;
Ok(png_data)
}
fn do_add_image(
path: &str,
content: &str,
image_path: &str,
width: Option<u32>,
height: Option<u32>,
style: &Option<DocxStyle>,
) -> Result<Vec<Content>, ErrorData> {
let mut doc = read_or_create_docx(path)?;
let image_data = load_image_as_png(image_path)?;
if !content.trim().is_empty() {
let mut caption = Paragraph::new();
if let Some(s) = style {
caption = s.apply_to_paragraph(caption);
caption = caption.add_run(s.apply_to_run(Run::new().add_text(content)));
} else {
caption = caption.add_run(Run::new().add_text(content));
}
doc = doc.add_paragraph(caption);
}
let mut paragraph = Paragraph::new();
if let Some(s) = style {
paragraph = s.apply_to_paragraph(paragraph);
}
let mut pic = Pic::new(&image_data);
if let (Some(w), Some(h)) = (width, height) {
pic = pic.size(w, h);
}
paragraph = paragraph.add_run(Run::new().add_image(pic));
doc = doc.add_paragraph(paragraph);
write_docx_file(path, doc)?;
Ok(vec![Content::text(format!(
"Successfully added image to {}",
path
))])
}
pub async fn docx_tool(
path: &str,
operation: &str,
@@ -92,513 +439,30 @@ pub async fn docx_tool(
params: Option<&serde_json::Value>,
) -> Result<Vec<Content>, ErrorData> {
match operation {
"extract_text" => {
let file = fs::read(path).map_err(|e| ErrorData {
code: ErrorCode::INTERNAL_ERROR,
message: Cow::from(format!("Failed to read DOCX file: {}", e)),
data: None,
})?;
let docx = read_docx(&file).map_err(|e| ErrorData {
code: ErrorCode::INTERNAL_ERROR,
message: Cow::from(format!("Failed to parse DOCX file: {}", e)),
data: None,
})?;
let mut text = String::new();
let mut structure = Vec::new();
let mut current_level = None;
// Extract document structure and text
for element in docx.document.children.iter() {
if let DocumentChild::Paragraph(p) = element {
// Check for heading style
if let Some(style) = p.property.style.as_ref() {
if style.val.starts_with("Heading") {
current_level = Some(style.val.clone());
structure.push(format!("{}: ", style.val));
}
}
// Extract text from runs
let para_text: String = p
.children
.iter()
.filter_map(|child| {
if let ParagraphChild::Run(run) = child {
Some(
run.children
.iter()
.filter_map(|rc| {
if let RunChild::Text(t) = rc {
Some(t.text.clone())
} else {
None
}
})
.collect::<Vec<_>>()
.join(""),
)
} else {
None
}
})
.collect::<Vec<_>>()
.join("");
if !para_text.trim().is_empty() {
if current_level.is_some() {
if let Some(s) = structure.last_mut() {
s.push_str(&para_text);
}
current_level = None;
}
text.push_str(&para_text);
text.push('\n');
}
}
}
let result = if !structure.is_empty() {
format!(
"Document Structure:\n{}\n\nFull Text:\n{}",
structure.join("\n"),
text
)
} else {
format!("Extracted Text:\n{}", text)
};
Ok(vec![Content::text(result)])
}
"extract_text" => do_extract_text(path),
"update_doc" => {
let content = content.ok_or_else(|| ErrorData {
code: ErrorCode::INVALID_PARAMS,
message: Cow::from("Content parameter required for update_doc"),
data: None,
})?;
// Parse update mode and style from params
let (mode, style) = if let Some(params) = params {
let mode = params
.get("mode")
.and_then(|v| v.as_str())
.unwrap_or("append");
let style = params.get("style").and_then(DocxStyle::from_json);
let mode = match mode {
"append" => UpdateMode::Append,
"replace" => {
let old_text = params
.get("old_text")
.and_then(|v| v.as_str())
.ok_or_else(|| ErrorData {
code: ErrorCode::INVALID_PARAMS,
message: Cow::from("old_text parameter required for replace mode"),
data: None,
})?;
UpdateMode::Replace {
old_text: old_text.to_string(),
}
}
"structured" => {
let level = params
.get("level")
.and_then(|v| v.as_str())
.map(String::from);
UpdateMode::InsertStructured {
level,
style: style.clone(),
}
}
"add_image" => {
let image_path = params
.get("image_path")
.and_then(|v| v.as_str())
.ok_or_else(|| ErrorData {
code: ErrorCode::INVALID_PARAMS,
message: Cow::from("image_path parameter required for add_image mode"),
data: None,
})?
.to_string();
let width = params
.get("width")
.and_then(|v| v.as_u64())
.map(|w| w as u32);
let height = params
.get("height")
.and_then(|v| v.as_u64())
.map(|h| h as u32);
UpdateMode::AddImage {
image_path,
width,
height,
}
}
_ => return Err(ErrorData {
code: ErrorCode::INVALID_PARAMS,
message: Cow::from("Invalid mode. Must be 'append', 'replace', 'structured', or 'add_image'"),
data: None,
}),
};
(mode, style)
} else {
(UpdateMode::Append, None)
};
let content = content
.ok_or_else(|| invalid_params("Content parameter required for update_doc"))?;
let (mode, style) = parse_update_mode(params)?;
match mode {
UpdateMode::Append => {
// Read existing document if it exists, or create new one
let mut doc = if std::path::Path::new(path).exists() {
let file = fs::read(path).map_err(|e| ErrorData {
code: ErrorCode::INTERNAL_ERROR,
message: Cow::from(format!("Failed to read DOCX file: {}", e)),
data: None,
})?;
read_docx(&file).map_err(|e| ErrorData {
code: ErrorCode::INTERNAL_ERROR,
message: Cow::from(format!("Failed to parse DOCX file: {}", e)),
data: None,
})?
} else {
Docx::new()
};
// Split content into paragraphs and add them
for para in content.split('\n') {
if !para.trim().is_empty() {
let mut run = Run::new().add_text(para);
let mut paragraph = Paragraph::new();
if let Some(style) = &style {
run = style.apply_to_run(run);
paragraph = style.apply_to_paragraph(paragraph);
}
doc = doc.add_paragraph(paragraph.add_run(run));
}
}
let mut buf = Vec::new();
{
let mut cursor = Cursor::new(&mut buf);
doc.build().pack(&mut cursor).map_err(|e| ErrorData {
code: ErrorCode::INTERNAL_ERROR,
message: Cow::from(format!("Failed to build DOCX: {}", e)),
data: None,
})?;
}
fs::write(path, &buf).map_err(|e| ErrorData {
code: ErrorCode::INTERNAL_ERROR,
message: Cow::from(format!("Failed to write DOCX file: {}", e)),
data: None,
})?;
Ok(vec![Content::text(format!(
"Successfully wrote content to {}",
path
))])
}
UpdateMode::Replace { old_text } => {
// Read existing document
let file = fs::read(path).map_err(|e| ErrorData {
code: ErrorCode::INTERNAL_ERROR,
message: Cow::from(format!("Failed to read DOCX file: {}", e)),
data: None,
})?;
let docx = read_docx(&file).map_err(|e| ErrorData {
code: ErrorCode::INTERNAL_ERROR,
message: Cow::from(format!("Failed to parse DOCX file: {}", e)),
data: None,
})?;
let mut new_doc = Docx::new();
let mut found_text = false;
// Process each paragraph
for element in docx.document.children.iter() {
if let DocumentChild::Paragraph(p) = element {
let para_text: String = p
.children
.iter()
.filter_map(|child| {
if let ParagraphChild::Run(run) = child {
Some(
run.children
.iter()
.filter_map(|rc| {
if let RunChild::Text(t) = rc {
Some(t.text.clone())
} else {
None
}
})
.collect::<Vec<_>>()
.join(""),
)
} else {
None
}
})
.collect::<Vec<_>>()
.join("");
if para_text.contains(&old_text) {
// Replace this paragraph with new content
found_text = true;
for para in content.split('\n') {
if !para.trim().is_empty() {
let mut run = Run::new().add_text(para);
let mut paragraph = Paragraph::new();
if let Some(style) = &style {
run = style.apply_to_run(run);
paragraph = style.apply_to_paragraph(paragraph);
}
new_doc = new_doc.add_paragraph(paragraph.add_run(run));
}
}
} else {
// Create a new paragraph with the same content and style
let mut para = Paragraph::new();
if let Some(style) = &p.property.style {
para = para.style(&style.val);
}
for child in p.children.iter() {
if let ParagraphChild::Run(run) = child {
for rc in run.children.iter() {
if let RunChild::Text(t) = rc {
para = para.add_run(Run::new().add_text(&t.text));
}
}
}
}
new_doc = new_doc.add_paragraph(para);
}
}
}
if !found_text {
return Err(ErrorData {
code: ErrorCode::INTERNAL_ERROR,
message: Cow::from(format!(
"Could not find text to replace: {}",
old_text
)),
data: None,
});
}
let mut buf = Vec::new();
{
let mut cursor = Cursor::new(&mut buf);
new_doc.build().pack(&mut cursor).map_err(|e| ErrorData {
code: ErrorCode::INTERNAL_ERROR,
message: Cow::from(format!("Failed to build DOCX: {}", e)),
data: None,
})?;
}
fs::write(path, &buf).map_err(|e| ErrorData {
code: ErrorCode::INTERNAL_ERROR,
message: Cow::from(format!("Failed to write DOCX file: {}", e)),
data: None,
})?;
Ok(vec![Content::text(format!(
"Successfully replaced content in {}",
path
))])
}
UpdateMode::InsertStructured { level, style } => {
let mut doc = if std::path::Path::new(path).exists() {
let file = fs::read(path).map_err(|e| ErrorData {
code: ErrorCode::INTERNAL_ERROR,
message: Cow::from(format!("Failed to read DOCX file: {}", e)),
data: None,
})?;
read_docx(&file).map_err(|e| ErrorData {
code: ErrorCode::INTERNAL_ERROR,
message: Cow::from(format!("Failed to parse DOCX file: {}", e)),
data: None,
})?
} else {
Docx::new()
};
// Create the paragraph with heading style if specified
for para in content.split('\n') {
if !para.trim().is_empty() {
let mut run = Run::new().add_text(para);
let mut paragraph = Paragraph::new();
// Apply heading style if specified
if let Some(level) = &level {
paragraph = paragraph.style(level);
}
// Apply custom style if specified
if let Some(style) = &style {
run = style.apply_to_run(run);
paragraph = style.apply_to_paragraph(paragraph);
}
doc = doc.add_paragraph(paragraph.add_run(run));
}
}
let mut buf = Vec::new();
{
let mut cursor = Cursor::new(&mut buf);
doc.build().pack(&mut cursor).map_err(|e| ErrorData {
code: ErrorCode::INTERNAL_ERROR,
message: Cow::from(format!("Failed to build DOCX: {}", e)),
data: None,
})?;
}
fs::write(path, &buf).map_err(|e| ErrorData {
code: ErrorCode::INTERNAL_ERROR,
message: Cow::from(format!("Failed to write DOCX file: {}", e)),
data: None,
})?;
Ok(vec![Content::text(format!(
"Successfully added structured content to {}",
path
))])
}
UpdateMode::Append => do_append(path, content, &style),
UpdateMode::Replace { old_text } => do_replace(path, content, &old_text, &style),
UpdateMode::InsertStructured {
level,
style: mode_style,
} => do_insert_structured(path, content, &level, &mode_style.or(style)),
UpdateMode::AddImage {
image_path,
width,
height,
} => {
let mut doc = if std::path::Path::new(path).exists() {
let file = fs::read(path).map_err(|e| ErrorData {
code: ErrorCode::INTERNAL_ERROR,
message: Cow::from(format!("Failed to read DOCX file: {}", e)),
data: None,
})?;
read_docx(&file).map_err(|e| ErrorData {
code: ErrorCode::INTERNAL_ERROR,
message: Cow::from(format!("Failed to parse DOCX file: {}", e)),
data: None,
})?
} else {
Docx::new()
};
// Read the image file
let image_data = fs::read(&image_path).map_err(|e| ErrorData {
code: ErrorCode::INTERNAL_ERROR,
message: Cow::from(format!("Failed to read image file: {}", e)),
data: None,
})?;
// Get image format and extension
let extension = std::path::Path::new(&image_path)
.extension()
.and_then(|e| e.to_str())
.ok_or_else(|| ErrorData {
code: ErrorCode::INTERNAL_ERROR,
message: Cow::from("Invalid image file extension".to_string()),
data: None,
})?
.to_lowercase();
// Convert to PNG if not already PNG
let image_data = if extension != "png" {
// Try to convert to PNG using the image crate
let img = image::load_from_memory(&image_data).map_err(|e| ErrorData {
code: ErrorCode::INTERNAL_ERROR,
message: Cow::from(format!("Failed to load image: {}", e)),
data: None,
})?;
let mut png_data = Vec::new();
img.write_to(&mut Cursor::new(&mut png_data), ImageFormat::Png)
.map_err(|e| ErrorData {
code: ErrorCode::INTERNAL_ERROR,
message: Cow::from(format!(
"Failed to convert image to PNG: {}",
e
)),
data: None,
})?;
png_data
} else {
image_data
};
// Add optional caption if provided
if !content.trim().is_empty() {
let mut caption = Paragraph::new();
if let Some(style) = &style {
caption = style.apply_to_paragraph(caption);
caption =
caption.add_run(style.apply_to_run(Run::new().add_text(content)));
} else {
caption = caption.add_run(Run::new().add_text(content));
}
doc = doc.add_paragraph(caption);
}
// Create a paragraph with the image
let mut paragraph = Paragraph::new();
if let Some(style) = &style {
paragraph = style.apply_to_paragraph(paragraph);
}
// Create and add the image
let mut pic = Pic::new(&image_data);
if let (Some(w), Some(h)) = (width, height) {
pic = pic.size(w, h);
}
paragraph = paragraph.add_run(Run::new().add_image(pic));
doc = doc.add_paragraph(paragraph);
let mut buf = Vec::new();
{
let mut cursor = Cursor::new(&mut buf);
doc.build().pack(&mut cursor).map_err(|e| ErrorData {
code: ErrorCode::INTERNAL_ERROR,
message: Cow::from(format!("Failed to build DOCX: {}", e)),
data: None,
})?;
}
fs::write(path, &buf).map_err(|e| ErrorData {
code: ErrorCode::INTERNAL_ERROR,
message: Cow::from(format!("Failed to write DOCX file: {}", e)),
data: None,
})?;
Ok(vec![Content::text(format!(
"Successfully added image to {}",
path
))])
}
} => do_add_image(path, content, &image_path, width, height, &style),
}
}
_ => Err(ErrorData {
code: ErrorCode::INVALID_PARAMS,
message: Cow::from(format!(
"Invalid operation: {}. Valid operations are: 'extract_text', 'update_doc'",
operation
)),
data: None,
}),
_ => Err(invalid_params(format!(
"Invalid operation: {}. Valid operations are: 'extract_text', 'update_doc'",
operation
))),
}
}