feat: outsources chunking parameters to extract chunk from documents …

[hajdul88]

…task

Summary by CodeRabbit

• New Features

  • Introduced a new ChunkerConfig class for flexible chunking strategy selection.
  • Updated document processing methods to support dynamic chunker selection.

• Bug Fixes

  • Enhanced error handling for unsupported chunker names.

• Documentation

  • Updated method signatures across various document classes to include a new chunker parameter.

• Chores

  • Removed direct dependencies on TextChunker, streamlining chunking implementation.

[coderabbitai]

Walkthrough

The changes involve modifications to several document processing classes, including AudioDocument, ImageDocument, PdfDocument, TextDocument, and Document. Each class's read method has been updated to accept an additional chunker: str parameter, facilitating dynamic selection of chunking strategies. The new ChunkerMapping.py file introduces the ChunkerConfig class, which maps chunker types to their corresponding classes. The extract_chunks_from_documents function has also been updated to include the chunker parameter, allowing for more flexible document processing.

Changes

File Path	Change Summary
cognee/modules/data/processing/document_types/AudioDocument.py	Updated read method signature to include chunker: str. Removed TextChunker import, using ChunkerConfig instead.
cognee/modules/data/processing/document_types/ChunkerMapping.py	Added ChunkerConfig class with get_chunker method for mapping chunker types to classes.
cognee/modules/data/processing/document_types/Document.py	Updated read method signature to include chunker: str as an optional parameter.
cognee/modules/data/processing/document_types/ImageDocument.py	Updated read method signature to include chunker: str. Removed TextChunker import, using ChunkerConfig instead.
cognee/modules/data/processing/document_types/PdfDocument.py	Updated read method signature to include chunker: str. Removed TextChunker import, using ChunkerConfig instead.
cognee/modules/data/processing/document_types/TextDocument.py	Updated read method signature to include chunker: str. Removed TextChunker import, using ChunkerConfig instead.
cognee/tasks/documents/extract_chunks_from_documents.py	Updated extract_chunks_from_documents function signature to include chunker parameter.
cognee/tests/integration/documents/AudioDocument_test.py	Updated test to call document.read with chunker='text_chunker'.
cognee/tests/integration/documents/ImageDocument_test.py	Updated test to call document.read with chunker='text_chunker'.
cognee/tests/integration/documents/PdfDocument_test.py	Updated test to call document.read with chunker='text_chunker'.
cognee/tests/integration/documents/TextDocument_test.py	Updated test to call document.read with chunker='text_chunker'.

Possibly related PRs

• Small cleanup pull request #201: The changes in the AudioDocument class are related to the modifications in the classify_documents function, as both involve the AudioDocument class and its instantiation, reflecting a broader context of handling different document types.

Suggested reviewers

• Vasilije1990

Poem

In the land of code where documents dwell,
A new way to chunk, oh what a spell!
With ChunkerConfig guiding the way,
Flexible processing brightens the day.
From audio to text, all types in a row,
Let’s hop to the future, watch our systems grow! 🐰✨

---

📜 Recent review details

Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro

📥 Commits

Reviewing files that changed from the base of the PR and between 7768239 and 65e1c92.

📒 Files selected for processing (1)

• cognee/modules/data/processing/document_types/Document.py (1 hunks)

🚧 Files skipped from review as they are similar to previous changes (1)

• cognee/modules/data/processing/document_types/Document.py

---

Thank you for using CodeRabbit. We offer it for free to the OSS community and would appreciate your support in helping us grow. If you find it useful, would you consider giving us a shout-out on your favorite social media?

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

🪧 Tips

Chat

There are 3 ways to chat with CodeRabbit:

• Review comments: Directly reply to a review comment made by CodeRabbit. Example:
  • I pushed a fix in commit <commit_id>, please review it.
  • Generate unit testing code for this file.
  • Open a follow-up GitHub issue for this discussion.
• Files and specific lines of code (under the "Files changed" tab): Tag @coderabbitai in a new review comment at the desired location with your query. Examples:
  • @coderabbitai generate unit testing code for this file.
  • @coderabbitai modularize this function.
• PR comments: Tag @coderabbitai in a new PR comment to ask questions about the PR branch. For the best results, please provide a very specific query, as very limited context is provided in this mode. Examples:
  • @coderabbitai gather interesting stats about this repository and render them as a table. Additionally, render a pie chart showing the language distribution in the codebase.
  • @coderabbitai read src/utils.ts and generate unit testing code.
  • @coderabbitai read the files in the src/scheduler package and generate a class diagram using mermaid and a README in the markdown format.
  • @coderabbitai help me debug CodeRabbit configuration file.

Note: Be mindful of the bot's finite context window. It's strongly recommended to break down tasks such as reading entire modules into smaller chunks. For a focused discussion, use review comments to chat about specific files and their changes, instead of using the PR comments.

CodeRabbit Commands (Invoked using PR comments)

• @coderabbitai pause to pause the reviews on a PR.
• @coderabbitai resume to resume the paused reviews.
• @coderabbitai review to trigger an incremental review. This is useful when automatic reviews are disabled for the repository.
• @coderabbitai full review to do a full review from scratch and review all the files again.
• @coderabbitai summary to regenerate the summary of the PR.
• @coderabbitai generate docstrings to generate docstrings for this PR. (Beta)
• @coderabbitai resolve resolve all the CodeRabbit review comments.
• @coderabbitai configuration to show the current CodeRabbit configuration for the repository.
• @coderabbitai help to get help.

Other keywords and placeholders

• Add @coderabbitai ignore anywhere in the PR description to prevent this PR from being reviewed.
• Add @coderabbitai summary to generate the high-level summary at a specific location in the PR description.
• Add @coderabbitai anywhere in the PR title to generate the title automatically.

CodeRabbit Configuration File (.coderabbit.yaml)

• You can programmatically configure CodeRabbit by adding a .coderabbit.yaml file to the root of your repository.
• Please see the configuration documentation for more information.
• If your editor has YAML language server enabled, you can add the path at the top of this file to enable auto-completion and validation: # yaml-language-server: $schema=https://coderabbit.ai/integrations/schema.v2.json

Documentation and Community

• Visit our Documentation for detailed information on how to use CodeRabbit.
• Join our Discord Community to get help, request features, and share feedback.
• Follow us on X/Twitter for updates and announcements.

[coderabbitai]

Actionable comments posted: 6

🧹 Outside diff range and nitpick comments (4)

cognee/modules/data/processing/document_types/ChunkerMapping.py (1)

4-6: Consider making the chunker mapping immutable

The class-level dictionary could be modified at runtime. Consider using a frozen dictionary or tuple-based mapping for immutability.

-    chunker_mapping = {
-        "text_chunker": TextChunker
-    }
+    chunker_mapping = frozendict({
+        "text_chunker": TextChunker
+    })

cognee/modules/data/processing/document_types/PdfDocument.py (1)

8-8: Add type hints and update docstring

The method signature should include type hints for better code maintainability and IDE support.

-    def read(self, chunk_size: int, chunker: str):
+    def read(self, chunk_size: int, chunker: str) -> Generator[str, None, None]:
+        """Read and chunk PDF document content.
+        
+        Args:
+            chunk_size (int): Size of each chunk
+            chunker (str): Type of chunker to use (e.g., "text", "semantic")
+            
+        Yields:
+            str: Document chunks
+            
+        Raises:
+            ValueError: If chunker type is invalid
+        """

cognee/modules/data/processing/document_types/TextDocument.py (1)

Line range hint 7-15: Ensure proper file handle cleanup

The file handle should be properly closed even if an error occurs during reading.

     def read(self, chunk_size: int, chunker: str):
         def get_text():
-            with open(self.raw_data_location, mode = "r", encoding = "utf-8") as file:
-                while True:
-                    text = file.read(1024)
+            BUFFER_SIZE = 1024 * 4  # 4KB buffer for better performance
+            try:
+                with open(self.raw_data_location, mode="r", encoding="utf-8") as file:
+                    while True:
+                        text = file.read(BUFFER_SIZE)
-                    if len(text.strip()) == 0:
-                        break
+                        if not text:
+                            break
-                    yield text
+                        yield text
+            except IOError as e:
+                raise RuntimeError(f"Failed to read text file: {e}") from e

cognee/modules/data/processing/document_types/AudioDocument.py (1)

1-1: Consider architectural improvements for better maintainability

The current implementation has several areas for improvement:

1. Move common chunker initialization logic to the base Document class to reduce code duplication
2. Add a configuration validation system to verify chunker types and parameters at startup
3. Implement a progress monitoring system for long-running operations like transcription
4. Consider using the Template Method pattern for the read method

Would you like me to provide example implementations for any of these improvements?

📜 Review details

Configuration used: CodeRabbit UI
Review profile: CHILL

📥 Commits

Reviewing files that changed from the base of the PR and between ce96431 and 29a02e0.

📒 Files selected for processing (7)

• cognee/modules/data/processing/document_types/AudioDocument.py (2 hunks)
• cognee/modules/data/processing/document_types/ChunkerMapping.py (1 hunks)
• cognee/modules/data/processing/document_types/Document.py (1 hunks)
• cognee/modules/data/processing/document_types/ImageDocument.py (2 hunks)
• cognee/modules/data/processing/document_types/PdfDocument.py (1 hunks)
• cognee/modules/data/processing/document_types/TextDocument.py (2 hunks)
• cognee/tasks/documents/extract_chunks_from_documents.py (1 hunks)

🔇 Additional comments (4)

cognee/modules/data/processing/document_types/ChunkerMapping.py (1)

9-15: Add return type hint to get_chunker method

The method should specify its return type for better type safety.

-    def get_chunker(cls, chunker_name: str):
+    def get_chunker(cls, chunker_name: str) -> type[TextChunker]:

Also, let's verify if there are other chunker implementations that should be included in the mapping.

✅ Verification successful

Let's gather more information about the chunker implementations to verify the return type and available chunkers.

---

Add return type hint to get_chunker method

The return type hint should be type[TextChunker] as verified from the codebase:

• The chunker_mapping only contains TextChunker class
• TextChunker is the base class for chunking implementation
• No other chunker implementations were found in the codebase

-    def get_chunker(cls, chunker_name: str):
+    def get_chunker(cls, chunker_name: str) -> type[TextChunker]:

🏁 Scripts executed

The following scripts were executed for the analysis:

Script:

#!/bin/bash
# Search for other chunker implementations
echo "Searching for potential chunker classes..."
rg -l "class.*Chunker" --type py

Length of output: 232

---

Script:

#!/bin/bash
# Check the content of both files to understand the chunker hierarchy and implementations
echo "=== Content of TextChunker.py ==="
cat cognee/modules/chunking/TextChunker.py

echo -e "\n=== Content of ChunkerMapping.py ==="
cat cognee/modules/data/processing/document_types/ChunkerMapping.py

Length of output: 4505

cognee/tasks/documents/extract_chunks_from_documents.py (1)

4-7: 🛠️ Refactor suggestion

Add type hints and update documentation

The function needs proper type hints and documentation for the new parameter.

-async def extract_chunks_from_documents(documents: list[Document], chunk_size: int = 1024, chunker = 'text_chunker'):
+async def extract_chunks_from_documents(
+    documents: list[Document],
+    chunk_size: int = 1024,
+    chunker: str = 'text_chunker'
+) -> AsyncIterator[str]:
+    """Extract chunks from multiple documents.
+    
+    Args:
+        documents: List of documents to process
+        chunk_size: Size of each chunk
+        chunker: Type of chunker to use (default: "text_chunker")
+    
+    Yields:
+        Chunks of document content
+    """

Let's verify if all document types support the new chunker parameter.

✅ Verification successful

Let me verify the implementation details of the read method in these document types to ensure they properly handle the chunker parameter.

---

Let me correct the paths and try again.

---

Let me try one more time with ripgrep to see the actual implementations.

---

Type hints and documentation changes approved - all document types support chunker parameter

The verification shows that:

1. All document types (Image, Audio, Text, PDF) have implemented the read method with the chunker parameter
2. They all use the parameter correctly by passing it to ChunkerConfig.get_chunker
3. The base Document class declares the parameter in its interface

The suggested type hints and documentation improvements are valid and should be implemented.

🏁 Scripts executed

The following scripts were executed for the analysis:

Script:

#!/bin/bash
# Check if all document types have been updated
echo "Checking document type implementations..."
rg -l "class.*Document.*\(.*\):" --type py | xargs rg "def read.*chunker.*:"

Length of output: 741

---

Script:

#!/bin/bash
# Check the implementation of read method in document types
echo "Document.py base class implementation:"
ast-grep --pattern 'def read($$$) { $$$ }' cognee/modules/data/processing/document_types/Document.py

echo -e "\nConcrete implementations:"
for file in ImageDocument.py AudioDocument.py TextDocument.py PdfDocument.py; do
    echo -e "\n$file implementation:"
    ast-grep --pattern 'def read($$$) { $$$ }' cognee/modules/data/processing/document_types/document_types/$file
done

Length of output: 1876

---

Script:

#!/bin/bash
# Check the implementation of read method in document types
echo "Document.py base class implementation:"
ast-grep --pattern 'def read($$$) { $$$ }' cognee/modules/data/processing/document_types/Document.py

echo -e "\nConcrete implementations:"
for file in ImageDocument AudioDocument TextDocument PdfDocument; do
    echo -e "\n$file implementation:"
    ast-grep --pattern 'def read($$$) { $$$ }' cognee/modules/data/processing/document_types/${file}.py
done

Length of output: 1247

---

Script:

#!/bin/bash
# Check implementations with context
echo "Document.py base class implementation:"
rg -A 5 "def read.*chunker.*:" cognee/modules/data/processing/document_types/Document.py

echo -e "\nConcrete implementations:"
for file in ImageDocument AudioDocument TextDocument PdfDocument; do
    echo -e "\n${file}.py implementation:"
    rg -A 5 "def read.*chunker.*:" cognee/modules/data/processing/document_types/${file}.py
done

Length of output: 2195

cognee/modules/data/processing/document_types/ImageDocument.py (2)

3-3: LGTM: Import changes align with new chunking architecture

The import change supports the transition to configurable chunking strategies.

---

17-18: Verify error handling for invalid chunker types

The implementation looks clean, but we should ensure proper error handling for:

1. Invalid chunker types
2. Chunker instantiation failures

Let's check the error handling implementation:

✅ Verification successful

Let me gather more specific information about the ChunkerConfig implementation to verify error handling.

---

Let me try a different approach to locate the chunker implementation and error handling.

---

Based on the search results, I can now see that ChunkerConfig.get_chunker is implemented in ChunkerMapping.py and it does have proper error handling. Let me verify the complete implementation of ChunkerMapping.py.

---

Based on the implementation found, I can now provide the final response:

Error handling for chunker types is properly implemented

The ChunkerConfig implementation has appropriate error handling:

• Invalid chunker types raise a NotImplementedError with a descriptive message listing available options
• Chunker instantiation is handled through a clean factory pattern with type safety

The implementation is sufficient for both validation points mentioned in the review:

1. Invalid chunker types ✓ (handled via NotImplementedError)
2. Chunker instantiation failures ✓ (handled via type-safe factory pattern)

🏁 Scripts executed

The following scripts were executed for the analysis:

Script:

#!/bin/bash
# Check ChunkerConfig implementation for error handling
ast-grep --pattern 'class ChunkerConfig {
  $$$
  get_chunker($_) {
    $$$
  }
}'

# Look for any error handling tests
rg -l "test.*chunker" | xargs rg "invalid|error|exception"

Length of output: 30859

---

Script:

#!/bin/bash
# Search for ChunkerConfig implementation
ast-grep --pattern 'class ChunkerConfig {
  $$$
}'

# Look for chunker-related tests
rg -A 5 "test.*chunker" "cognee/tests"

# Search for chunker-related exceptions
rg "class.*Error.*chunker" -A 3

Length of output: 125

---

Script:

#!/bin/bash
# Search for chunker-related files
fd chunker

# Look for get_chunker method implementation
rg -A 10 "get_chunker"

# Search for chunker-related error handling
rg -A 5 "raise.*chunker"

Length of output: 2892

---

Script:

#!/bin/bash
# Get complete implementation of ChunkerMapping
cat cognee/modules/data/processing/document_types/ChunkerMapping.py

# Look for any chunker-related exceptions in the codebase
rg "class.*Error" cognee/modules/chunking/

Length of output: 620

[coderabbitai]

Actionable comments posted: 4

🧹 Outside diff range and nitpick comments (1)

cognee/tests/integration/documents/AudioDocument_test.py (1)

Line range hint 25-25: Consider implementing a shared test suite for chunking strategies

The current implementation duplicates chunking validation across different document types. Consider:

1. Creating a shared test suite for common chunking behavior
2. Defining a standard interface for chunking strategy validation
3. Adding integration tests that verify chunking consistency across document types

Example shared test structure:

class BaseChunkerTest:
    def test_basic_chunking(self, document, chunker):
        # Common chunking tests
    
    def test_invalid_chunker(self, document):
        # Common error cases

class TestPdfDocumentChunking(BaseChunkerTest):
    # PDF-specific tests

class TestTextDocumentChunking(BaseChunkerTest):
    # Text-specific tests

class TestAudioDocumentChunking(BaseChunkerTest):
    # Audio-specific tests

Also applies to: 36-36, 34-34

📜 Review details

Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro

📥 Commits

Reviewing files that changed from the base of the PR and between 29a02e0 and 7768239.

📒 Files selected for processing (5)

• cognee/modules/data/processing/document_types/Document.py (1 hunks)
• cognee/tests/integration/documents/AudioDocument_test.py (1 hunks)
• cognee/tests/integration/documents/ImageDocument_test.py (1 hunks)
• cognee/tests/integration/documents/PdfDocument_test.py (1 hunks)
• cognee/tests/integration/documents/TextDocument_test.py (1 hunks)

🚧 Files skipped from review as they are similar to previous changes (1)

• cognee/modules/data/processing/document_types/Document.py

🔇 Additional comments (1)

cognee/tests/integration/documents/AudioDocument_test.py (1)

34-34: Consider audio-specific chunking strategies and test coverage

While the test covers basic text chunking of transcribed audio, consider:

1. Audio-specific chunking strategies (e.g., time-based chunks, speaker-based chunks)
2. Testing the interaction between transcription and chunking
3. Verifying that chunking preserves important audio context (e.g., speaker transitions)

[dexters1]

What about the UnstructuredDocument class? Do these changes need to be implemented there as well?

[hajdul88]

What about the UnstructuredDocument class? Do these changes need to be implemented there as well?

If we outsource chunkers like this, then yes, but we discuss this with Boris.