feat: deletes on the fly embeddings and uses edge collections #436
Conversation
WalkthroughThe pull request introduces modifications to two key files in the Cognee project. In Changes
Sequence DiagramsequenceDiagram
participant VectorEngine
participant CogneeGraph
participant GraphEdges
CogneeGraph->>VectorEngine: get_distance_from_collection_elements
VectorEngine-->>CogneeGraph: Return distances
CogneeGraph->>GraphEdges: Assign distances to edges
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: 2
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (2)
cognee/modules/graph/cognee_graph/CogneeGraph.py(2 hunks)cognee/modules/retrieval/brute_force_triplet_search.py(1 hunks)
⏰ Context from checks skipped due to timeout of 90000ms (17)
- GitHub Check: run_notebook_test / test
- GitHub Check: run_simple_example_test / test
- GitHub Check: run_notebook_test / test
- GitHub Check: run_dynamic_steps_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
- GitHub Check: test
- GitHub Check: test
- GitHub Check: test
- GitHub Check: profiler
- GitHub Check: docker-compose-test
🔇 Additional comments (3)
cognee/modules/retrieval/brute_force_triplet_search.py (1)
110-110: Verify impact of removing duplicate filteringThe simplified dictionary comprehension allows duplicates to persist in the results. While this aligns with the PR objective, we should verify that:
- Downstream processing in
CogneeGraphcan handle duplicate entries- This change doesn't negatively impact the quality of search results
Let's check for potential duplicate handling in the codebase:
✅ Verification successful
Duplicate handling verified - safe to remove filtering
The removal of duplicate filtering is safe as the graph storage layer handles deduplication through the
deduplicate_nodes_and_edgesfunction, which is called during data point addition. Node uniqueness is also enforced at the graph level.🏁 Scripts executed
The following scripts were executed for the analysis:
Script:
#!/bin/bash # Search for duplicate handling patterns in vector and graph operations rg -g '!{test_*,*.test}*' -A 3 'duplicat|unique' cognee/Length of output: 8469
cognee/modules/graph/cognee_graph/CogneeGraph.py (2)
11-11: LGTM: Required import for async operationsThe addition of asyncio import is necessary for the async operations in the refactored code.
130-149: Verify edge collection usageThe method now uses a single edge collection ("edge_type_relationship_name") instead of on-the-fly embeddings. Let's verify this collection exists and contains the expected data.
✅ Verification successful
Collection usage is properly managed
The "edge_type_relationship_name" collection is dynamically populated by the
index_graph_edgestask, which collects relationship types from graph edges, transforms them intoEdgeTypeobjects, and indexes them in the vector database. The collection name is a constant used consistently between the indexing task and the graph querying code.🏁 Scripts executed
The following scripts were executed for the analysis:
Script:
#!/bin/bash # Check for edge collection references and their structure rg -g '!{test_*,*.test}*' "edge_type_relationship_name" cognee/ # Look for vector engine collection initialization rg -g '!{test_*,*.test}*' -B 3 -A 3 "create_collection|init.*collection" cognee/infrastructure/databases/vector/Length of output: 11001
Script:
#!/bin/bash # Look for edge relationship initialization and configuration rg -g '!{test_*,*.test}*' -B 3 -A 3 "relationship.*type|edge.*type" cognee/ # Check for configuration files and initialization fd -e yaml -e yml -e json . cognee/ --exec cat {}Length of output: 22780
| embedding_map = {result.payload["text"]: result.score for result in edge_distances} | ||
|
|
||
| # Step 4: Assign precomputed distances to edges | ||
| for edge in self.edges: | ||
| relationship_type = edge.attributes.get("relationship_type") | ||
| if not relationship_type or relationship_type not in embedding_map: | ||
| print(f"Edge {edge} has an unknown or missing relationship type.") | ||
| continue | ||
|
|
||
| # Assign the precomputed distance | ||
| edge.attributes["vector_distance"] = embedding_map[relationship_type] |
There was a problem hiding this comment.
🛠️ Refactor suggestion
Improve edge distance mapping robustness
The current implementation has potential issues:
- Silent failures with print statements
- No validation of embedding_map values
- No handling of edge cases (empty results, invalid scores)
Consider this improvement:
- embedding_map = {result.payload["text"]: result.score for result in edge_distances}
+ embedding_map = {}
+ for result in edge_distances:
+ if "text" not in result.payload:
+ raise InvalidValueError(f"Missing 'text' in payload: {result.payload}")
+ if not isinstance(result.score, (int, float)):
+ raise InvalidValueError(f"Invalid score type: {type(result.score)}")
+ embedding_map[result.payload["text"]] = result.score
for edge in self.edges:
relationship_type = edge.attributes.get("relationship_type")
if not relationship_type or relationship_type not in embedding_map:
- print(f"Edge {edge} has an unknown or missing relationship type.")
+ logging.warning("Edge %s has an unknown or missing relationship type", edge)
continue
edge.attributes["vector_distance"] = embedding_map[relationship_type]Committable suggestion skipped: line range outside the PR's diff.
| async def map_vector_distances_to_graph_edges(self, vector_engine, query) -> None: | ||
| try: | ||
| # Step 1: Generate the query embedding | ||
| query_vector = await vector_engine.embed_data([query]) | ||
| query_vector = query_vector[0] | ||
| if query_vector is None or len(query_vector) == 0: | ||
| raise ValueError("Failed to generate query embedding.") | ||
|
|
||
| # Step 2: Collect all unique relationship types | ||
| unique_relationship_types = set() | ||
| for edge in self.edges: | ||
| relationship_type = edge.attributes.get("relationship_type") | ||
| if relationship_type: | ||
| unique_relationship_types.add(relationship_type) | ||
|
|
||
| # Step 3: Embed all unique relationship types | ||
| unique_relationship_types = list(unique_relationship_types) | ||
| relationship_type_embeddings = await vector_engine.embed_data(unique_relationship_types) | ||
|
|
||
| # Step 4: Map relationship types to their embeddings and calculate distances | ||
| embedding_map = {} | ||
| for relationship_type, embedding in zip( | ||
| unique_relationship_types, relationship_type_embeddings | ||
| ): | ||
| edge_vector = np.array(embedding) | ||
|
|
||
| # Calculate cosine similarity | ||
| similarity = np.dot(query_vector, edge_vector) / ( | ||
| np.linalg.norm(query_vector) * np.linalg.norm(edge_vector) | ||
| ) | ||
| distance = 1 - similarity | ||
| edge_distances = await vector_engine.get_distance_from_collection_elements( | ||
| "edge_type_relationship_name", query_text=query | ||
| ) |
There was a problem hiding this comment.
🛠️ Refactor suggestion
Enhance error handling and logging
The method has several areas that could benefit from improved error handling:
- Query vector validation could raise a custom exception
- Collection name should be configurable
- Consider using proper logging instead of print statements
Consider this improvement:
async def map_vector_distances_to_graph_edges(self, vector_engine, query) -> None:
+ EDGE_TYPE_COLLECTION = "edge_type_relationship_name" # Move to config
try:
query_vector = await vector_engine.embed_data([query])
query_vector = query_vector[0]
if query_vector is None or len(query_vector) == 0:
- raise ValueError("Failed to generate query embedding.")
+ raise InvalidValueError("Failed to generate query embedding: empty or null vector")
edge_distances = await vector_engine.get_distance_from_collection_elements(
- "edge_type_relationship_name", query_text=query
+ EDGE_TYPE_COLLECTION, query_text=query
)
+ if not edge_distances:
+ raise InvalidValueError(f"No distances retrieved from {EDGE_TYPE_COLLECTION}")Committable suggestion skipped: line range outside the PR's diff.
…-from-bruteforce-search-and
hajdul88
requested review from
borisarzentar,
Vasilije1990,
dexters1,
alekszievr and
lxobr
…-from-bruteforce-search-and
hajdul88
deleted the
feature/cog-762-deleting-in-memory-embeddings-from-bruteforce-search-and
branch
Summary by CodeRabbit
Performance Improvements
Code Optimization
Refactor