feat: externalize chunkers [cog-1354] #547
Conversation
WalkthroughThe changes refactor how chunkers are instantiated across various modules. Instead of using string identifiers and a mapping configuration, all read methods now expect a class reference (e.g., Changes
Sequence Diagram(s)sequenceDiagram
participant TaskManager as Task Manager
participant Extractor as extract_chunks_from_documents
participant Document as Document.read
participant Chunker as TextChunker
TaskManager->>Extractor: Invoke extract_chunks_from_documents(documents, max_chunk_tokens, ..., TextChunker)
Extractor->>Document: Call read(chunk_size, chunker_cls=TextChunker, max_chunk_tokens)
Document->>Chunker: Instantiate and process chunking
Chunker-->>Document: Return text chunks
Document-->>Extractor: Yield processed chunks
Extractor-->>TaskManager: Return aggregated chunks
Poem
✨ Finishing Touches
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? 🪧 TipsChatThere are 3 ways to chat with CodeRabbit:
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)
Other keywords and placeholders
CodeRabbit Configuration File (
|
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 3
🧹 Nitpick comments (3)
cognee/modules/data/processing/document_types/Document.py (1)
12-12: Add type hint forchunker_clsparameter.The base class should define the expected type for
chunker_clsto ensure type safety across all document implementations.- def read(self, chunk_size: int, chunker_cls, max_chunk_tokens: Optional[int] = None) -> str: + def read(self, chunk_size: int, chunker_cls: type, max_chunk_tokens: Optional[int] = None) -> str:cognee/modules/data/processing/document_types/ImageDocument.py (1)
17-19: Simplify the get_text lambda function.The lambda function returning a single-item list could be simplified.
- self, chunk_size=chunk_size, get_text=lambda: [text], max_chunk_tokens=max_chunk_tokens + self, chunk_size=chunk_size, get_text=lambda: iter([text]), max_chunk_tokens=max_chunk_tokenscognee/tasks/documents/extract_chunks_from_documents.py (1)
32-38: Update docstring to reflect new chunker parameter type.The docstring still implies string-based chunker configuration. Update it to clarify that
chunkernow expects a class reference.""" Extracts chunks of data from a list of documents based on the specified chunking parameters. Notes: - The `read` method of the `Document` class must be implemented to support the chunking operation. - - The `chunker` parameter determines the chunking logic and should align with the document type. + - The `chunker` parameter expects a chunker class (e.g., TextChunker) that implements the chunking logic. """
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (8)
cognee/api/v1/cognify/cognify_v2.py(2 hunks)cognee/modules/data/processing/document_types/AudioDocument.py(1 hunks)cognee/modules/data/processing/document_types/ChunkerMapping.py(0 hunks)cognee/modules/data/processing/document_types/Document.py(1 hunks)cognee/modules/data/processing/document_types/ImageDocument.py(1 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(3 hunks)
💤 Files with no reviewable changes (1)
- cognee/modules/data/processing/document_types/ChunkerMapping.py
⏰ Context from checks skipped due to timeout of 90000ms (3)
- GitHub Check: lint (ubuntu-latest, 3.10.x)
- GitHub Check: docker-compose-test
- GitHub Check: Build Cognee Backend Docker App Image
🔇 Additional comments (5)
cognee/modules/data/processing/document_types/PdfDocument.py (1)
17-19: LGTM! Clean chunker instantiation.The direct instantiation of the chunker class is cleaner and more maintainable than the previous string-based mapping approach.
cognee/modules/data/processing/document_types/TextDocument.py (1)
18-20: LGTM! Clean chunker instantiation.The direct instantiation of the chunker class is cleaner and more maintainable than the previous string-based mapping approach.
cognee/modules/data/processing/document_types/AudioDocument.py (1)
15-22: LGTM! Clean refactoring of chunker instantiation.The changes simplify the code by removing the string-based chunker mapping and directly using class references. The implementation maintains functionality while reducing indirection.
cognee/tasks/documents/extract_chunks_from_documents.py (1)
8-8: LGTM! Consistent refactoring of chunker instantiation.The changes align with the broader refactoring effort:
- TextChunker import added
- Default chunker updated to use class reference
- Parameter renamed in document.read call for consistency
Also applies to: 26-43
cognee/api/v1/cognify/cognify_v2.py (1)
28-28: LGTM! Task configuration updated correctly.The changes properly integrate the new chunker class-based approach into the default task configuration:
- TextChunker import added
- Task instantiation updated with chunker parameter
Also applies to: 127-130
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (1)
cognee/modules/data/processing/document_types/UnstructuredDocument.py (1)
31-33: Consider adding Protocol class for chunker interface.The chunker instantiation is more flexible now, but we could improve type safety further.
Consider defining a Protocol class to enforce the chunker interface:
from typing import Protocol, Iterator class ChunkerProtocol(Protocol): def __init__(self, document: Document, chunk_size: int, get_text: callable, max_chunk_tokens: int) -> None: ... def read(self) -> Iterator[str]: ... def read(self, chunk_size: int, chunker_cls: type[ChunkerProtocol], max_chunk_tokens: int) -> str: ...
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (6)
cognee/modules/data/processing/document_types/AudioDocument.py(1 hunks)cognee/modules/data/processing/document_types/Document.py(1 hunks)cognee/modules/data/processing/document_types/ImageDocument.py(1 hunks)cognee/modules/data/processing/document_types/PdfDocument.py(1 hunks)cognee/modules/data/processing/document_types/TextDocument.py(2 hunks)cognee/modules/data/processing/document_types/UnstructuredDocument.py(2 hunks)
🚧 Files skipped from review as they are similar to previous changes (4)
- cognee/modules/data/processing/document_types/Document.py
- cognee/modules/data/processing/document_types/PdfDocument.py
- cognee/modules/data/processing/document_types/ImageDocument.py
- cognee/modules/data/processing/document_types/TextDocument.py
⏰ Context from checks skipped due to timeout of 90000ms (22)
- GitHub Check: run_notebook_test / test
- GitHub Check: Test on macos-15
- GitHub Check: run_eval_framework_test / test
- GitHub Check: Test on macos-13
- GitHub Check: run_dynamic_steps_example_test / test
- GitHub Check: Test on macos-13
- GitHub Check: run_simple_example_test / test
- GitHub Check: run_multimedia_example_test / test
- GitHub Check: run_notebook_test / test
- GitHub Check: test
- GitHub Check: test
- GitHub Check: test
- GitHub Check: test
- GitHub Check: test
- GitHub Check: Test on macos-13
- GitHub Check: run_notebook_test / test
- GitHub Check: windows-latest
- GitHub Check: test
- GitHub Check: test
- GitHub Check: Build Cognee Backend Docker App Image
- GitHub Check: docker-compose-test
- GitHub Check: profiler
🔇 Additional comments (2)
cognee/modules/data/processing/document_types/AudioDocument.py (2)
13-13: LGTM! The signature change improves type safety.The parameter change from
chunker: strtochunker_cls: typeis a good improvement that:
- Makes the code more type-safe by expecting a class reference instead of a string
- Provides better IDE support for code completion and type checking
- Makes the parameter's purpose clearer through its naming
18-20: LGTM! The chunker instantiation is now more straightforward.The direct instantiation of the chunker class simplifies the code by:
- Removing the unnecessary indirection through ChunkerConfig
- Maintaining a clean interface for passing the text content via lambda
- Preserving all necessary parameters (chunk_size, get_text, max_chunk_tokens)
There was a problem hiding this comment.
Actionable comments posted: 0
🧹 Nitpick comments (1)
cognee/tests/integration/documents/UnstructuredDocument_test.py (1)
71-73: LGTM! Consistent parameter updates across document types.The read() method calls have been consistently updated to use
chunker_cls=TextChunkeracross all document types while maintaining existing chunking parameters.Consider adding test cases for invalid chunker classes to ensure proper error handling:
def test_UnstructuredDocument_invalid_chunker(): document = UnstructuredDocument( id=uuid.uuid4(), name="example.txt", raw_data_location="test.txt", external_metadata="", mime_type="text/plain", ) # Test with invalid chunker class with pytest.raises(TypeError, match="chunker_cls must be a valid chunker class"): for _ in document.read(chunk_size=1024, chunker_cls=str, max_chunk_tokens=1024): passAlso applies to: 81-83, 91-93, 103-105
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (5)
cognee/tests/integration/documents/AudioDocument_test.py(2 hunks)cognee/tests/integration/documents/ImageDocument_test.py(2 hunks)cognee/tests/integration/documents/PdfDocument_test.py(2 hunks)cognee/tests/integration/documents/TextDocument_test.py(2 hunks)cognee/tests/integration/documents/UnstructuredDocument_test.py(5 hunks)
⏰ Context from checks skipped due to timeout of 90000ms (28)
- GitHub Check: Test on macos-15
- GitHub Check: test
- GitHub Check: Test on macos-13
- GitHub Check: Test on macos-15
- GitHub Check: Test on ubuntu-22.04
- GitHub Check: Test on macos-13
- GitHub Check: Test on macos-15
- GitHub Check: Test on ubuntu-22.04
- GitHub Check: Test on macos-13
- GitHub Check: Test on ubuntu-22.04
- GitHub Check: test
- GitHub Check: test
- GitHub Check: test
- GitHub Check: test
- GitHub Check: run_eval_framework_test / test
- GitHub Check: windows-latest
- GitHub Check: run_dynamic_steps_example_test / test
- GitHub Check: test
- GitHub Check: run_notebook_test / test
- GitHub Check: run_notebook_test / test
- GitHub Check: run_multimedia_example_test / test
- GitHub Check: run_notebook_test / test
- GitHub Check: run_simple_example_test / test
- GitHub Check: test
- GitHub Check: profiler
- GitHub Check: docker-compose-test
- GitHub Check: profiler
- GitHub Check: Build Cognee Backend Docker App Image
🔇 Additional comments (9)
cognee/tests/integration/documents/TextDocument_test.py (2)
5-5: LGTM! Clean import addition.The explicit import of
TextChunkerclass improves code clarity and enables static type checking.
41-41: LGTM! Improved type safety in chunker specification.The change from string identifier to class reference (
chunker_cls=TextChunker) is a good improvement that:
- Enables better IDE support and static type checking
- Reduces the risk of runtime errors from invalid string identifiers
- Makes the code more maintainable by making dependencies explicit
cognee/tests/integration/documents/PdfDocument_test.py (2)
3-3: LGTM! Import change aligns with the new chunker approach.The explicit import of
TextChunkeris appropriate for the new class-based chunker implementation.
28-28: LGTM! Test updated correctly for new chunker API.The test has been properly adapted to use the new class-based chunker approach while maintaining its effectiveness in validating document processing.
cognee/tests/integration/documents/AudioDocument_test.py (2)
3-3: LGTM!The import of
TextChunkeraligns with the refactoring to use class references instead of string identifiers for chunker instantiation.
38-38: LGTM!The update to use
chunker_cls=TextChunkerinstead of a string identifier is consistent with the PR's objective to externalize chunker instantiation. The test's validation logic remains robust and comprehensive.cognee/tests/integration/documents/UnstructuredDocument_test.py (1)
3-3: LGTM! Import aligns with chunker externalization.The import of
TextChunkerclass follows Python conventions and supports the transition from string-based to class-based chunker instantiation.cognee/tests/integration/documents/ImageDocument_test.py (2)
3-3: LGTM!The import of
TextChunkeraligns with the PR's objective of externalizing chunkers and is correctly used in the test.
26-27:Details
✅ Verification successful
LGTM! Verify consistency across document types.
The change from
chunkertochunker_cls=TextChunkeraligns with the PR's objective and maintains the test's functionality.Let's verify that this pattern is consistently applied across all document types:
🏁 Script executed:
#!/bin/bash # Description: Check for consistent usage of chunker_cls parameter across document types # Search for read method definitions in document types echo "Checking read method signatures in document types:" ast-grep --pattern 'def read(self, $$$, chunker_cls, $$$)' # Search for read method calls to ensure consistent usage echo "Checking read method calls:" rg -A 2 'read\(' | rg 'chunker_cls'Length of output: 2428
LGTM! Consistency verified across document types.
The updated call in
cognee/tests/integration/documents/ImageDocument_test.py—usingchunker_cls=TextChunker—is consistent with calls in the other document tests (Text, Audio, Pdf, and Unstructured) as well as with the corresponding method signatures in the document classes. One minor point: incognee/tasks/documents/extract_chunks_from_documents.py, the call passeschunker_cls=chunker. If this was intentional (i.e., if the variablechunkeralready refers to the desired class type), then no changes are needed. Otherwise, please verify that this usage aligns with the standardized pattern across the codebase.
borisarzentar
left a comment
borisarzentar
left a comment
There was a problem hiding this comment.
Added 35 lines, removed 57, created more flexible solution. Beautiful.
borisarzentar
changed the title
Description
DCO Affirmation
I affirm that all code in every commit of this pull request conforms to the terms of the Topoteretes Developer Certificate of Origin
Summary by CodeRabbit
New Features
Refactor
Chores