Fix: Export PDFTextWritingTool and add comprehensive tests

crewai_tools/__init__.py

@@ -53,6 +53,7 @@
     PatronusLocalEvaluatorTool,
     PatronusPredefinedCriteriaEvalTool,
     PDFSearchTool,
+    PDFTextWritingTool,
     PGSearchTool,
     QdrantVectorSearchTool,
     RagTool,

crewai_tools/tools/__init__.py

@@ -57,6 +57,7 @@
     PatronusPredefinedCriteriaEvalTool,
 )
 from .pdf_search_tool.pdf_search_tool import PDFSearchTool
+from .pdf_text_writing_tool.pdf_text_writing_tool import PDFTextWritingTool
 from .pg_search_tool.pg_search_tool import PGSearchTool
 from .qdrant_vector_search_tool.qdrant_search_tool import QdrantVectorSearchTool
 from .rag.rag_tool import RagTool

crewai_tools/tools/pdf_text_writing_tool/pdf_text_writing_tool.py

@@ -2,7 +2,8 @@
 from typing import Optional, Type
 
 from pydantic import BaseModel, Field
-from pypdf import ContentStream, Font, NameObject, PageObject, PdfReader, PdfWriter
+from pypdf import PdfReader, PdfWriter
+from pypdf.annotations import FreeText
 
 from crewai_tools.tools.rag.rag_tool import RagTool
 
@@ -23,7 +24,7 @@ class PDFTextWritingToolSchema(BaseModel):
         default="F1", description="Font name for standard fonts"
     )
     font_file: Optional[str] = Field(
-        None, description="Path to a .ttf font file for custom font usage"
+        None, description="Path to a .ttf font file for custom font usage (currently not supported)"
     )
     page_number: int = Field(default=0, description="Page number to add text to")
 
@@ -44,47 +45,50 @@ def run(
         position: tuple,
         font_size: int,
         font_color: str,
-        font_name: str = "F1",
+        font_name: str = "Arial",
         font_file: Optional[str] = None,
         page_number: int = 0,
     ) -> str:
-        reader = PdfReader(pdf_path)
-        writer = PdfWriter()
-
-        if page_number >= len(reader.pages):
-            return "Page number out of range."
-
-        page: PageObject = reader.pages[page_number]
-        content = ContentStream(page["/Contents"].data, reader)
-
-        if font_file:
-            # Check if the font file exists
-            if not Path(font_file).exists():
-                return "Font file does not exist."
-
-            # Embed the custom font
-            font_name = self.embed_font(writer, font_file)
-
-        # Prepare text operation with the custom or standard font
-        x_position, y_position = position
-        text_operation = f"BT /{font_name} {font_size} Tf {x_position} {y_position} Td ({text}) Tj ET"
-        content.operations.append([font_color])  # Set color
-        content.operations.append([text_operation])  # Add text
-
-        # Replace old content with new content
-        page[NameObject("/Contents")] = content
-        writer.add_page(page)
-
-        # Save the new PDF
-        output_pdf_path = "modified_output.pdf"
-        with open(output_pdf_path, "wb") as out_file:
-            writer.write(out_file)
-
-        return f"Text added to {output_pdf_path} successfully."
-
-    def embed_font(self, writer: PdfWriter, font_file: str) -> str:
-        """Embeds a TTF font into the PDF and returns the font name."""
-        with open(font_file, "rb") as file:
-            font = Font.true_type(file.read())
-        font_ref = writer.add_object(font)
-        return font_ref
+        try:
+            reader = PdfReader(pdf_path)
+            writer = PdfWriter()
+
+            if page_number >= len(reader.pages):
+                return "Page number out of range."
+
+            for page in reader.pages:
+                writer.add_page(page)
+
+            x_position, y_position = position
+            rect = (x_position, y_position, x_position + 200, y_position + 50)
+
+            if font_color == "0 0 0 rg":
+                color = "000000"  # black
+            elif font_color == "1 0 0 rg":
+                color = "ff0000"  # red
+            elif font_color == "0 1 0 rg":
+                color = "00ff00"  # green
+            elif font_color == "0 0 1 rg":
+                color = "0000ff"  # blue
+            else:
+                color = "000000"  # default to black
+
+            annotation = FreeText(
+                text=text,
+                rect=rect,
+                font=font_name,
+                font_size=f"{font_size}pt",
+                font_color=color,
+            )
+
+            writer.add_annotation(page_number=page_number, annotation=annotation)
+
+            # Save the new PDF
+            output_pdf_path = "modified_output.pdf"
+            with open(output_pdf_path, "wb") as out_file:
+                writer.write(out_file)
+
+            return f"Text added to {output_pdf_path} successfully."
+
+        except Exception as e:
+            return f"Error adding text to PDF: {str(e)}"

tests/tools/test_pdf_text_writing_tool.py

@@ -0,0 +1,192 @@
+import os
+import tempfile
+from pathlib import Path
+from unittest.mock import MagicMock, patch
+
+import pytest
+from pypdf import PdfReader, PdfWriter
+from pypdf.annotations import FreeText
+
+from crewai_tools.tools.pdf_text_writing_tool.pdf_text_writing_tool import (
+    PDFTextWritingTool,
+    PDFTextWritingToolSchema,
+)
+
+
+@pytest.fixture
+def sample_pdf():
+    """Create a simple PDF file for testing."""
+    with tempfile.NamedTemporaryFile(suffix=".pdf", delete=False) as temp_file:
+        writer = PdfWriter()
+        from pypdf import PageObject
+
+        page = PageObject.create_blank_page(width=612, height=792)
+        writer.add_page(page)
+
+        with open(temp_file.name, "wb") as output_file:
+            writer.write(output_file)
+
+        yield temp_file.name
+
+    os.unlink(temp_file.name)
+
+
+@pytest.fixture
+def sample_font_file():
+    """Create a mock TTF font file for testing."""
+    with tempfile.NamedTemporaryFile(suffix=".ttf", delete=False) as temp_file:
+        temp_file.write(b"mock font data")
+        yield temp_file.name
+
+    os.unlink(temp_file.name)
+
+
+def test_pdf_text_writing_tool_schema():
+    """Test the PDFTextWritingToolSchema validation."""
+    schema = PDFTextWritingToolSchema(
+        pdf_path="test.pdf",
+        text="Hello World",
+        position=(100, 200),
+        font_size=14,
+        font_color="0 0 1 rg",
+        page_number=0
+    )
+
+    assert schema.pdf_path == "test.pdf"
+    assert schema.text == "Hello World"
+    assert schema.position == (100, 200)
+    assert schema.font_size == 14
+    assert schema.font_color == "0 0 1 rg"
+    assert schema.page_number == 0
+
+
+def test_pdf_text_writing_tool_initialization():
+    """Test PDFTextWritingTool initialization."""
+    tool = PDFTextWritingTool()
+
+    assert tool.name == "PDF Text Writing Tool"
+    assert "write text to a specific position in a PDF document" in tool.description
+    assert tool.args_schema == PDFTextWritingToolSchema
+
+
+def test_pdf_text_writing_tool_basic_functionality(sample_pdf):
+    """Test basic text writing functionality."""
+    tool = PDFTextWritingTool()
+
+    result = tool.run(
+        pdf_path=sample_pdf,
+        text="Test Text",
+        position=(100, 200),
+        font_size=12,
+        font_color="0 0 0 rg",
+        page_number=0
+    )
+
+    assert "Text added to modified_output.pdf successfully" in result
+
+
+def test_pdf_text_writing_tool_page_out_of_range(sample_pdf):
+    """Test error handling for page number out of range."""
+    tool = PDFTextWritingTool()
+
+    result = tool.run(
+        pdf_path=sample_pdf,
+        text="Test Text",
+        position=(100, 200),
+        font_size=12,
+        font_color="0 0 0 rg",
+        page_number=999
+    )
+
+    assert result == "Page number out of range."
+
+
+def test_pdf_text_writing_tool_missing_font_file(sample_pdf):
+    """Test that font file parameter is ignored (not supported)."""
+    tool = PDFTextWritingTool()
+
+    result = tool.run(
+        pdf_path=sample_pdf,
+        text="Test Text",
+        position=(100, 200),
+        font_size=12,
+        font_color="0 0 0 rg",
+        font_file="nonexistent_font.ttf",
+        page_number=0
+    )
+
+    assert "Text added to modified_output.pdf successfully" in result
+
+
+def test_pdf_text_writing_tool_with_custom_font(sample_pdf, sample_font_file):
+    """Test text writing with custom font (currently not supported)."""
+    tool = PDFTextWritingTool()
+
+    result = tool.run(
+        pdf_path=sample_pdf,
+        text="Test Text",
+        position=(100, 200),
+        font_size=12,
+        font_color="0 0 0 rg",
+        font_file=sample_font_file,
+        page_number=0
+    )
+
+    assert "Text added to modified_output.pdf successfully" in result
+
+
+def test_pdf_text_writing_tool_different_positions(sample_pdf):
+    """Test text writing at different positions."""
+    tool = PDFTextWritingTool()
+
+    positions = [(50, 100), (200, 300), (400, 500)]
+
+    for position in positions:
+        result = tool.run(
+            pdf_path=sample_pdf,
+            text=f"Text at {position}",
+            position=position,
+            font_size=12,
+            font_color="0 0 0 rg",
+            page_number=0
+        )
+
+        assert "Text added to modified_output.pdf successfully" in result
+
+
+def test_pdf_text_writing_tool_different_font_sizes(sample_pdf):
+    """Test text writing with different font sizes."""
+    tool = PDFTextWritingTool()
+
+    font_sizes = [8, 12, 16, 24]
+
+    for font_size in font_sizes:
+        result = tool.run(
+            pdf_path=sample_pdf,
+            text=f"Size {font_size} text",
+            position=(100, 200),
+            font_size=font_size,
+            font_color="0 0 0 rg",
+            page_number=0
+        )
+
+        assert "Text added to modified_output.pdf successfully" in result
+
+
+def test_pdf_text_writing_tool_different_colors(sample_pdf):
+    """Test text writing with different colors."""
+    tool = PDFTextWritingTool()
+
+    colors = ["0 0 0 rg", "1 0 0 rg", "0 1 0 rg", "0 0 1 rg"]
+
+    for color in colors:
+        result = tool.run(
+            pdf_path=sample_pdf,
+            text="Colored text",
+            position=(100, 200),
+            font_size=12,
+            font_color=color,
+            page_number=0
+        )
+
+        assert "Text added to modified_output.pdf successfully" in result

tests/tools/test_search_tools.py

@@ -15,6 +15,7 @@
     JSONSearchTool,
     MDXSearchTool,
     PDFSearchTool,
+    PDFTextWritingTool,
     TXTSearchTool,
     WebsiteSearchTool,
     XMLSearchTool,