feat: add chunk associations task with LLM classifier #1834
Conversation
- Implement create_chunk_associations function with vector similarity search - Add LLM prompts for association decision making - Integrate task into memify pipeline as default enrichment - Add integration test for chunk associations - Fix add_rule_associations to return data for pipeline flow
Please make sure all the checkboxes are checked:
|
WalkthroughAdds LLM-driven semantic chunk association: new prompts, an async task to detect and persist chunk-to-chunk associations via vector search + LLM classification, integrates the task into the memify pipeline, exposes it from the package, and adds an integration test and a small pipeline bugfix. Changes
Estimated code review effort🎯 4 (Complex) | ⏱️ ~45 minutes
Pre-merge checks and finishing touches✅ Passed checks (5 passed)
✨ Finishing touches
🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![github-actions[bot]](https://avatars.githubusercontent.com/in/15368?s=60&v=4){kind=small}
There was a problem hiding this comment.
Hello @kckoh, thank you for submitting a PR! We will respond as soon as possible.
![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 (4)
cognee/infrastructure/llm/prompts/chunk_association_classifier_system.txt (1)
1-1: Remove leading whitespace.Line 1 has 2 leading spaces before "You". This may be unintentional and could affect prompt rendering consistency.
- You are a semantic association classifier for knowledge graph construction. +You are a semantic association classifier for knowledge graph construction.cognee/tasks/memify/create_chunk_associations.py (3)
19-46: Consider adding error handling for LLM failures.If the LLM call fails (network issues, invalid response, etc.), the exception will propagate and may terminate the entire pipeline. Consider wrapping the LLM call in a try-except block and returning a safe default (e.g.,
False) to allow graceful degradation.+ try: decision = await LLMGateway.acreate_structured_output( text_input=user_prompt, system_prompt=system_prompt, response_model=AssociationDecision ) - - return decision.should_associate + return decision.should_associate + except Exception as e: + logger.warning(f"LLM association decision failed: {e}. Defaulting to no association.") + return False
57-66: Consider reducing log verbosity for potentially large data.Logging full
dataandchunksat INFO level could produce excessive logs for large datasets and may inadvertently log sensitive content. Consider using DEBUG level or truncating the logged data.- logger.info(f"create_chunk_associations called with data type: {type(data)}, data: {data}") + logger.debug(f"create_chunk_associations called with data type: {type(data)}, data: {data}") ... - logger.info(f"Processing chunks: {chunks}") + logger.debug(f"Processing chunks: {chunks}")
101-109: Guard againstNonepayload.Accessing
similar_chunk.payload.get('text', '')assumespayloadis notNone. IfpayloadisNone, this will raise anAttributeError.logger.info( f"Asking LLM for chunks with similarity {similar_chunk.score}: " - f"'{chunk_text[:50]}...' <-> '{similar_chunk.payload.get('text', '')[:50]}...'" + f"'{chunk_text[:50]}...' <-> '{(similar_chunk.payload or {}).get('text', '')[:50]}...'" ) # Ask LLM llm_decision = await ask_llm_for_association( - chunk_text, similar_chunk.payload.get("text", ""), similar_chunk.score + chunk_text, (similar_chunk.payload or {}).get("text", ""), similar_chunk.score )
📜 Review details
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (7)
cognee/infrastructure/llm/prompts/chunk_association_classifier_system.txt(1 hunks)cognee/infrastructure/llm/prompts/chunk_association_classifier_user.txt(1 hunks)cognee/modules/memify/memify.py(2 hunks)cognee/tasks/codingagents/coding_rule_associations.py(1 hunks)cognee/tasks/memify/__init__.py(1 hunks)cognee/tasks/memify/create_chunk_associations.py(1 hunks)cognee/tests/test_chunk_associations.py(1 hunks)
🧰 Additional context used
📓 Path-based instructions (5)
**/*.py
📄 CodeRabbit inference engine (AGENTS.md)
**/*.py: Use 4-space indentation in Python code
Use snake_case for Python module and function names
Use PascalCase for Python class names
Use ruff format before committing Python code
Use ruff check for import hygiene and style enforcement with line-length 100 configured in pyproject.toml
Prefer explicit, structured error handling in Python code
Files:
cognee/tasks/codingagents/coding_rule_associations.pycognee/tasks/memify/__init__.pycognee/modules/memify/memify.pycognee/tasks/memify/create_chunk_associations.pycognee/tests/test_chunk_associations.py
⚙️ CodeRabbit configuration file
**/*.py: When reviewing Python code for this project:
- Prioritize portability over clarity, especially when dealing with cross-Python compatibility. However, with the priority in mind, do still consider improvements to clarity when relevant.
- As a general guideline, consider the code style advocated in the PEP 8 standard (excluding the use of spaces for indentation) and evaluate suggested changes for code style compliance.
- As a style convention, consider the code style advocated in CEP-8 and evaluate suggested changes for code style compliance.
- As a general guideline, try to provide any relevant, official, and supporting documentation links to any tool's suggestions in review comments. This guideline is important for posterity.
- As a general rule, undocumented function definitions and class definitions in the project's Python code are assumed incomplete. Please consider suggesting a short summary of the code for any of these incomplete definitions as docstrings when reviewing.
Files:
cognee/tasks/codingagents/coding_rule_associations.pycognee/tasks/memify/__init__.pycognee/modules/memify/memify.pycognee/tasks/memify/create_chunk_associations.pycognee/tests/test_chunk_associations.py
cognee/**/*.py
📄 CodeRabbit inference engine (AGENTS.md)
Use shared logging utilities from cognee.shared.logging_utils in Python code
Files:
cognee/tasks/codingagents/coding_rule_associations.pycognee/tasks/memify/__init__.pycognee/modules/memify/memify.pycognee/tasks/memify/create_chunk_associations.pycognee/tests/test_chunk_associations.py
cognee/{modules,infrastructure,tasks}/**/*.py
📄 CodeRabbit inference engine (AGENTS.md)
Co-locate feature-specific helpers under their respective package (modules/, infrastructure/, or tasks/)
Files:
cognee/tasks/codingagents/coding_rule_associations.pycognee/tasks/memify/__init__.pycognee/modules/memify/memify.pycognee/tasks/memify/create_chunk_associations.py
cognee/tests/**/*.py
📄 CodeRabbit inference engine (AGENTS.md)
cognee/tests/**/*.py: Place Python tests under cognee/tests/ organized by type (unit, integration, cli_tests)
Name Python test files test_*.py and use pytest.mark.asyncio for async tests
Files:
cognee/tests/test_chunk_associations.py
cognee/tests/*
⚙️ CodeRabbit configuration file
cognee/tests/*: When reviewing test code:
- Prioritize portability over clarity, especially when dealing with cross-Python compatibility. However, with the priority in mind, do still consider improvements to clarity when relevant.
- As a general guideline, consider the code style advocated in the PEP 8 standard (excluding the use of spaces for indentation) and evaluate suggested changes for code style compliance.
- As a style convention, consider the code style advocated in CEP-8 and evaluate suggested changes for code style compliance, pointing out any violations discovered.
- As a general guideline, try to provide any relevant, official, and supporting documentation links to any tool's suggestions in review comments. This guideline is important for posterity.
- As a project rule, Python source files with names prefixed by the string "test_" and located in the project's "tests" directory are the project's unit-testing code. It is safe, albeit a heuristic, to assume these are considered part of the project's minimal acceptance testing unless a justifying exception to this assumption is documented.
- As a project rule, any files without extensions and with names prefixed by either the string "check_" or the string "test_", and located in the project's "tests" directory, are the project's non-unit test code. "Non-unit test" in this context refers to any type of testing other than unit testing, such as (but not limited to) functional testing, style linting, regression testing, etc. It can also be assumed that non-unit testing code is usually written as Bash shell scripts.
Files:
cognee/tests/test_chunk_associations.py
🧠 Learnings (1)
📚 Learning: 2024-11-13T14:55:05.912Z
Learnt from: 0xideas
Repo: topoteretes/cognee PR: 205
File: cognee/tests/unit/processing/chunks/chunk_by_paragraph_test.py:7-7
Timestamp: 2024-11-13T14:55:05.912Z
Learning: When changes are made to the chunking implementation in `cognee/tasks/chunks`, the ground truth values in the corresponding tests in `cognee/tests/unit/processing/chunks` need to be updated accordingly.
Applied to files:
cognee/modules/memify/memify.pycognee/tasks/memify/create_chunk_associations.pycognee/tests/test_chunk_associations.py
🧬 Code graph analysis (4)
cognee/tasks/memify/__init__.py (1)
cognee/tasks/memify/create_chunk_associations.py (1)
create_chunk_associations(49-137)
cognee/modules/memify/memify.py (1)
cognee/tasks/memify/create_chunk_associations.py (1)
create_chunk_associations(49-137)
cognee/tasks/memify/create_chunk_associations.py (6)
cognee/infrastructure/databases/graph/get_graph_engine.py (1)
get_graph_engine(10-24)cognee/infrastructure/databases/vector/get_vector_engine.py (1)
get_vector_engine(5-7)cognee/infrastructure/llm/LLMGateway.py (1)
LLMGateway(6-66)cognee/infrastructure/llm/prompts/render_prompt.py (1)
render_prompt(5-42)cognee/shared/logging_utils.py (1)
get_logger(212-224)cognee/tasks/storage/index_graph_edges.py (1)
index_graph_edges(42-77)
cognee/tests/test_chunk_associations.py (3)
cognee/shared/logging_utils.py (1)
get_logger(212-224)cognee/modules/memify/memify.py (1)
memify(29-124)cognee/infrastructure/databases/graph/get_graph_engine.py (1)
get_graph_engine(10-24)
🔇 Additional comments (12)
cognee/tasks/codingagents/coding_rule_associations.py (1)
128-130: LGTM! Pipeline flow continuity fix.Returning
dataenables downstream tasks to continue processing. This aligns with the pattern used in the newcreate_chunk_associationstask.cognee/modules/memify/memify.py (2)
21-21: LGTM!Import correctly placed alongside related memify task imports.
79-83: LGTM! New enrichment task properly integrated.The
create_chunk_associationstask is well-configured with a 0.7 similarity threshold. Thebatch_size=1ensures sequential LLM calls for association decisions, which is reasonable for initial correctness, though it may be a candidate for optimization with larger batch sizes in future iterations if performance becomes a concern.cognee/tasks/memify/__init__.py (1)
1-1: LGTM!Re-export follows the existing pattern and makes
create_chunk_associationsavailable at the package level.cognee/infrastructure/llm/prompts/chunk_association_classifier_system.txt (1)
5-9: LGTM!The criteria for creating associations are well-defined: related topics/concepts, selectivity for clear relationships, and focus on meaningful connections beyond keywords.
cognee/infrastructure/llm/prompts/chunk_association_classifier_user.txt (1)
1-11: LGTM!The prompt template is well-structured with clear placeholders that match the context dictionary in
ask_llm_for_association. Providing both chunk texts and the similarity score gives the LLM sufficient context for decision-making.cognee/tests/test_chunk_associations.py (2)
45-52: Kuzu-specific query may limit portability.The Cypher query uses
NodeandEDGElabels which are Kuzu-specific conventions. If the codebase supports multiple graph backends, this test will only work with Kuzu. Consider adding a comment or using a backend-agnostic query method if available.
69-76: LGTM with minor note.The assertion verifies that at least one association was created, which is appropriate for an integration test with semantically related sample texts. The cleanup at the end mirrors the setup, ensuring a clean state.
cognee/tasks/memify/create_chunk_associations.py (4)
1-10: LGTM!Imports are well-organized and follow project conventions. Logger setup uses
get_loggerfrom shared.logging_utils as per coding guidelines.
13-16: LGTM!Simple and well-documented Pydantic model for structured LLM output.
85-85: Assumption that first result is self-match may not hold.The code assumes
similar_chunks[0]is always the origin chunk itself. While this is typically true for vector similarity search, it's not guaranteed in all cases (e.g., if the chunk text was modified or if the vector engine uses approximate search). Consider adding a validation or fallback.
128-137: LGTM!Edge persistence and indexing follow the established pattern. Returning
dataensures pipeline continuity.
|
|
||
| origin_chunk_id = similar_chunks[0].id # First result is the chunk itself | ||
|
|
||
| for similar_chunk in similar_chunks[1:]: # Skip first (self) |
There was a problem hiding this comment.
I assume that this will take a lot of time when run on big documents, even if we keep track of compared pairs. Did you try running it with some bigger document and compare time to small one? It would be nice to see how it behaves.
There was a problem hiding this comment.
Great point! I haven't tested this with larger documents yet. The current test only uses few small chunks.
With N chunks and limit=15, we could make up to N×14 comparisons, and each one above the similarity threshold triggers an LLM call. This could get slow and expensive with 100+ chunks.
Let me run a benchmark with a realistic document (50-100 chunks) and measure the time and LLM call count. If it's too slow, we can look up optimizations. Let me test and report back with numbers.
| for chunk_text in chunks: | ||
| # Get this chunk's ID and find similar chunks in one search | ||
| similar_chunks = await vector_engine.search( | ||
| collection_name="DocumentChunk_text", query_text=chunk_text, limit=5 |
There was a problem hiding this comment.
Question: Why do we limit the results of the search to 5 here?
There was a problem hiding this comment.
Thank you for the question! At first I thought limit=5 was intentional to reduce the number of LLM calls (since each similar chunk above the threshold triggers an LLM classification).
However, I found that all vector adapters default to limit=15, which suggests that's a well-thought-out default value.
i will remove the hard-coded limit=5 but should we consider making it configurable if different use cases need different limits:
async def create_chunk_associations(
data,
similarity_threshold: float = 0.7,
search_limit: int = 15
):
What do you think? Should we just remove the
hardcoded 5, or make it configurable too?
@siillee
…/cognee into feature/chunk-associations
There was a problem hiding this comment.
Actionable comments posted: 0
♻️ Duplicate comments (1)
cognee/tasks/memify/create_chunk_associations.py (1)
78-80: Hardcodedlimit=5should be configurable.This was discussed in a previous review. The hardcoded limit restricts the number of candidate pairs and LLM calls, but it may be too restrictive for some use cases. Consider making it a function parameter with a sensible default.
-async def create_chunk_associations(data, similarity_threshold: float = 0.7): +async def create_chunk_associations( + data, similarity_threshold: float = 0.7, search_limit: int = 5 +) -> list:Then use
search_limitin place of the hardcoded5:similar_chunks = await vector_engine.search( - collection_name="DocumentChunk_text", query_text=chunk_text, limit=5 + collection_name="DocumentChunk_text", query_text=chunk_text, limit=search_limit )
🧹 Nitpick comments (3)
cognee/tasks/memify/create_chunk_associations.py (3)
42-46: Consider adding error handling for LLM call failures.The LLM call could fail due to network issues, rate limits, or invalid responses. Currently, exceptions will propagate and potentially fail the entire pipeline. Consider whether to wrap this in a try-except to return a safe default (e.g.,
False) on failure, or if fail-fast behavior is intentional.- decision = await LLMGateway.acreate_structured_output( - text_input=user_prompt, system_prompt=system_prompt, response_model=AssociationDecision - ) - - return decision.should_associate + try: + decision = await LLMGateway.acreate_structured_output( + text_input=user_prompt, system_prompt=system_prompt, response_model=AssociationDecision + ) + return decision.should_associate + except Exception as e: + logger.warning(f"LLM association decision failed: {e}, defaulting to False") + return False
49-56: Add return type hint for pipeline clarity.The function returns
datafor pipeline continuity but lacks a return type annotation. Adding the type hint improves readability and IDE support.-async def create_chunk_associations(data, similarity_threshold: float = 0.7): +async def create_chunk_associations(data, similarity_threshold: float = 0.7) -> list:
101-110: Guard against missingtextpayload key.If
similar_chunk.payloadisNoneor doesn't have a'text'key, the.get("text", "")will work, butsimilar_chunk.payloaditself could beNonecausing anAttributeError.logger.info( f"Asking LLM for chunks with similarity {similar_chunk.score}: " - f"'{chunk_text[:50]}...' <-> '{similar_chunk.payload.get('text', '')[:50]}...'" + f"'{chunk_text[:50]}...' <-> '{(similar_chunk.payload or {}).get('text', '')[:50]}...'" ) # Ask LLM llm_decision = await ask_llm_for_association( - chunk_text, similar_chunk.payload.get("text", ""), similar_chunk.score + chunk_text, (similar_chunk.payload or {}).get("text", ""), similar_chunk.score )
📜 Review details
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (2)
cognee/tasks/memify/create_chunk_associations.py(1 hunks)cognee/tests/test_chunk_associations.py(1 hunks)
🚧 Files skipped from review as they are similar to previous changes (1)
- cognee/tests/test_chunk_associations.py
🧰 Additional context used
📓 Path-based instructions (3)
**/*.py
📄 CodeRabbit inference engine (AGENTS.md)
**/*.py: Use 4-space indentation in Python code
Use snake_case for Python module and function names
Use PascalCase for Python class names
Use ruff format before committing Python code
Use ruff check for import hygiene and style enforcement with line-length 100 configured in pyproject.toml
Prefer explicit, structured error handling in Python code
Files:
cognee/tasks/memify/create_chunk_associations.py
⚙️ CodeRabbit configuration file
**/*.py: When reviewing Python code for this project:
- Prioritize portability over clarity, especially when dealing with cross-Python compatibility. However, with the priority in mind, do still consider improvements to clarity when relevant.
- As a general guideline, consider the code style advocated in the PEP 8 standard (excluding the use of spaces for indentation) and evaluate suggested changes for code style compliance.
- As a style convention, consider the code style advocated in CEP-8 and evaluate suggested changes for code style compliance.
- As a general guideline, try to provide any relevant, official, and supporting documentation links to any tool's suggestions in review comments. This guideline is important for posterity.
- As a general rule, undocumented function definitions and class definitions in the project's Python code are assumed incomplete. Please consider suggesting a short summary of the code for any of these incomplete definitions as docstrings when reviewing.
Files:
cognee/tasks/memify/create_chunk_associations.py
cognee/**/*.py
📄 CodeRabbit inference engine (AGENTS.md)
Use shared logging utilities from cognee.shared.logging_utils in Python code
Files:
cognee/tasks/memify/create_chunk_associations.py
cognee/{modules,infrastructure,tasks}/**/*.py
📄 CodeRabbit inference engine (AGENTS.md)
Co-locate feature-specific helpers under their respective package (modules/, infrastructure/, or tasks/)
Files:
cognee/tasks/memify/create_chunk_associations.py
🧠 Learnings (2)
📓 Common learnings
Learnt from: 0xideas
Repo: topoteretes/cognee PR: 205
File: cognee/tests/unit/processing/chunks/chunk_by_paragraph_test.py:7-7
Timestamp: 2024-11-13T14:55:05.912Z
Learning: When changes are made to the chunking implementation in `cognee/tasks/chunks`, the ground truth values in the corresponding tests in `cognee/tests/unit/processing/chunks` need to be updated accordingly.
📚 Learning: 2024-11-13T14:55:05.912Z
Learnt from: 0xideas
Repo: topoteretes/cognee PR: 205
File: cognee/tests/unit/processing/chunks/chunk_by_paragraph_test.py:7-7
Timestamp: 2024-11-13T14:55:05.912Z
Learning: When changes are made to the chunking implementation in `cognee/tasks/chunks`, the ground truth values in the corresponding tests in `cognee/tests/unit/processing/chunks` need to be updated accordingly.
Applied to files:
cognee/tasks/memify/create_chunk_associations.py
🧬 Code graph analysis (1)
cognee/tasks/memify/create_chunk_associations.py (7)
cognee/infrastructure/databases/graph/get_graph_engine.py (1)
get_graph_engine(10-24)cognee/infrastructure/databases/vector/get_vector_engine.py (1)
get_vector_engine(5-7)cognee/infrastructure/llm/LLMGateway.py (1)
LLMGateway(6-66)cognee/infrastructure/llm/prompts/render_prompt.py (1)
render_prompt(5-42)cognee/shared/logging_utils.py (1)
get_logger(212-224)cognee/tasks/storage/index_graph_edges.py (1)
index_graph_edges(42-77)cognee/modules/graph/cognee_graph/CogneeGraph.py (1)
score(248-252)
⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (1)
- GitHub Check: GitGuardian Security Checks
🔇 Additional comments (5)
cognee/tasks/memify/create_chunk_associations.py (5)
1-10: LGTM!Imports are well-organized and the logger is properly instantiated using
get_loggerfrom shared logging utilities, as per coding guidelines.
13-17: LGTM!Clean Pydantic model for structured LLM response validation. The single
should_associatefield aligns with the binary classification requirement.
128-134: Add explicit error handling for graph persistence operations.The current code lacks error handling for
graph_engine.add_edges()andindex_graph_edges()failures. Per project guidelines, prefer explicit, structured error handling in Python code. Consider wrapping these operations in a try-except block to either log failures and continue, or let exceptions propagate for fail-fast behavior. Verify which approach aligns with the pipeline's error handling strategy used in similar tasks.
85-87: Verify if first search result is guaranteed to be the source chunk.The concern that the first vector search result may not be the source chunk is conceptually valid—vector similarity engines rank results by score, not guaranteed order. However, this requires examining: (1) how
similar_chunksis populated, (2) the vector search implementation's behavior, and (3) whether chunk ID-based retrieval is available as an alternative. Manual verification of the actualquery_chunksimplementation and vector database configuration is needed to determine if this is a fragile assumption or acceptable given the search guarantees.
94-99: Verify score semantics: is lower or higher better?The code assumes
similar_chunk.score < similarity_thresholdtreats lower scores as "below threshold" (implying higher = more similar). However, vector engine score semantics depend on the metric type—similarity metrics return higher=better (cosine similarity, dot product), while distance metrics return lower=better (Euclidean, Manhattan). Verify which metric your vector engine uses and confirm this threshold comparison aligns with its score semantics.
3 participants
Chunk Associations Task
A task for creating semantic associations between document chunks and organizing them into common node sets based on their similarity.
steps:
Closes Add a task that creates common association from a group of chunks #1333
Description
Type of Change
Screenshots/Videos (if applicable)
Pre-submission Checklist
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
Bug Fixes
Tests
✏️ Tip: You can customize this high-level summary in your review settings.