gmakstutis/cognee
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

format

Browse source
This commit is contained in:
vasilije 2025-07-09 16:53:34 +02:00
parent 816bc024a1
commit ace7f32112
3 changed files with 159 additions and 121 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

118
cognee/infrastructure/databases/hybrid/falkordb/FalkorDBAdapter.py
View file

@ -10,7 +10,13 @@ from typing import List, Dict, Any, Optional, Tuple, Type, Union
from falkordb import FalkorDB from falkordb import FalkorDB
from cognee.exceptions import InvalidValueError from cognee.exceptions import InvalidValueError
from cognee.infrastructure.databases.graph.graph_db_interface import GraphDBInterface, record_graph_changes, NodeData, EdgeData, Node from cognee.infrastructure.databases.graph.graph_db_interface import (
GraphDBInterface,
record_graph_changes,
NodeData,
EdgeData,
Node,
)
from cognee.infrastructure.databases.vector.embeddings import EmbeddingEngine from cognee.infrastructure.databases.vector.embeddings import EmbeddingEngine
from cognee.infrastructure.databases.vector.vector_db_interface import VectorDBInterface from cognee.infrastructure.databases.vector.vector_db_interface import VectorDBInterface
from cognee.infrastructure.engine import DataPoint from cognee.infrastructure.engine import DataPoint
@ -176,7 +182,14 @@ class FalkorDBAdapter(VectorDBInterface, GraphDBInterface):
return f"'{json.dumps(value).replace(chr(39), chr(34))}'" return f"'{json.dumps(value).replace(chr(39), chr(34))}'"
if type(value) is str: if type(value) is str:
# Escape single quotes and handle special characters # Escape single quotes and handle special characters
escaped_value = str(value).replace("'", "\\'").replace('"', '\\"').replace('\n', '\\n').replace('\r', '\\r').replace('\t', '\\t') escaped_value = (
str(value)
.replace("'", "\\'")
.replace('"', '\\"')
.replace("\n", "\\n")
.replace("\r", "\\r")
.replace("\t", "\\t")
)
return f"'{escaped_value}'" return f"'{escaped_value}'"
return f"'{str(value)}'" return f"'{str(value)}'"
@ -237,10 +250,7 @@ class FalkorDBAdapter(VectorDBInterface, GraphDBInterface):
""" """
).strip() ).strip()
params = { params = {"node_id": str(data_point.id), "properties": clean_properties}
"node_id": str(data_point.id),
"properties": clean_properties
}
return query, params return query, params
@ -258,15 +268,16 @@ class FalkorDBAdapter(VectorDBInterface, GraphDBInterface):
""" """
# Replace hyphens, spaces, and other special characters with underscores # Replace hyphens, spaces, and other special characters with underscores
import re import re
sanitized = re.sub(r'[^\w]', '_', relationship_name)
sanitized = re.sub(r"[^\w]", "_", relationship_name)
# Remove consecutive underscores # Remove consecutive underscores
sanitized = re.sub(r'_+', '_', sanitized) sanitized = re.sub(r"_+", "_", sanitized)
# Remove leading/trailing underscores # Remove leading/trailing underscores
sanitized = sanitized.strip('_') sanitized = sanitized.strip("_")
# Ensure it starts with a letter or underscore # Ensure it starts with a letter or underscore
if sanitized and not sanitized[0].isalpha() and sanitized[0] != '_': if sanitized and not sanitized[0].isalpha() and sanitized[0] != "_":
sanitized = '_' + sanitized sanitized = "_" + sanitized
return sanitized or 'RELATIONSHIP' return sanitized or "RELATIONSHIP"
async def create_edge_query(self, edge: tuple[str, str, str, dict]) -> str: async def create_edge_query(self, edge: tuple[str, str, str, dict]) -> str:
""" """
@ -462,10 +473,7 @@ class FalkorDBAdapter(VectorDBInterface, GraphDBInterface):
clean_properties[key] = value clean_properties[key] = value
query = "MERGE (node {id: $node_id}) SET node += $properties, node.updated_at = timestamp()" query = "MERGE (node {id: $node_id}) SET node += $properties, node.updated_at = timestamp()"
params = { params = {"node_id": node_id, "properties": clean_properties}
"node_id": node_id,
"properties": clean_properties
}
self.query(query, params) self.query(query, params)
@ -550,11 +558,13 @@ class FalkorDBAdapter(VectorDBInterface, GraphDBInterface):
# Node is in (node_id, properties) format # Node is in (node_id, properties) format
node_id, properties = node node_id, properties = node
await self.add_node(node_id, properties) await self.add_node(node_id, properties)
elif hasattr(node, 'id') and hasattr(node, 'model_dump'): elif hasattr(node, "id") and hasattr(node, "model_dump"):
# Node is a DataPoint object # Node is a DataPoint object
await self.add_node(str(node.id), node.model_dump()) await self.add_node(str(node.id), node.model_dump())
else: else:
raise ValueError(f"Invalid node format: {node}. Expected tuple (node_id, properties) or DataPoint object.") raise ValueError(
f"Invalid node format: {node}. Expected tuple (node_id, properties) or DataPoint object."
)
async def add_edge( async def add_edge(
self, self,
@ -599,7 +609,9 @@ class FalkorDBAdapter(VectorDBInterface, GraphDBInterface):
source_id, target_id, relationship_name, properties = edge source_id, target_id, relationship_name, properties = edge
await self.add_edge(source_id, target_id, relationship_name, properties) await self.add_edge(source_id, target_id, relationship_name, properties)
else: else:
raise ValueError(f"Invalid edge format: {edge}. Expected tuple (source_id, target_id, relationship_name, properties).") raise ValueError(
f"Invalid edge format: {edge}. Expected tuple (source_id, target_id, relationship_name, properties)."
)
async def has_edges(self, edges): async def has_edges(self, edges):
""" """
@ -1015,12 +1027,14 @@ class FalkorDBAdapter(VectorDBInterface, GraphDBInterface):
if result.result_set: if result.result_set:
for record in result.result_set: for record in result.result_set:
# FalkorDB returns values by index: source_id, target_id, relationship_name, properties # FalkorDB returns values by index: source_id, target_id, relationship_name, properties
edges.append(( edges.append(
record[0], # source_id (
record[1], # target_id record[0], # source_id
record[2], # relationship_name record[1], # target_id
record[3] # properties record[2], # relationship_name
)) record[3], # properties
)
)
return edges return edges
async def has_edge(self, source_id: str, target_id: str, relationship_name: str) -> bool: async def has_edge(self, source_id: str, target_id: str, relationship_name: str) -> bool:
@ -1044,7 +1058,11 @@ class FalkorDBAdapter(VectorDBInterface, GraphDBInterface):
AND (r.relationship_name = $relationship_name OR NOT EXISTS(r.relationship_name)) AND (r.relationship_name = $relationship_name OR NOT EXISTS(r.relationship_name))
RETURN COUNT(r) > 0 AS edge_exists RETURN COUNT(r) > 0 AS edge_exists
""", """,
{"source_id": source_id, "target_id": target_id, "relationship_name": relationship_name}, {
"source_id": source_id,
"target_id": target_id,
"relationship_name": relationship_name,
},
) )
if result.result_set and len(result.result_set) > 0: if result.result_set and len(result.result_set) > 0:
@ -1081,19 +1099,23 @@ class FalkorDBAdapter(VectorDBInterface, GraphDBInterface):
if include_optional: if include_optional:
# Add optional metrics - simplified implementation # Add optional metrics - simplified implementation
metrics.update({ metrics.update(
"num_selfloops": 0, # Simplified {
"diameter": -1, # Not implemented "num_selfloops": 0, # Simplified
"avg_shortest_path_length": -1, # Not implemented "diameter": -1, # Not implemented
"avg_clustering": -1, # Not implemented "avg_shortest_path_length": -1, # Not implemented
}) "avg_clustering": -1, # Not implemented
}
)
else: else:
metrics.update({ metrics.update(
"num_selfloops": -1, {
"diameter": -1, "num_selfloops": -1,
"avg_shortest_path_length": -1, "diameter": -1,
"avg_clustering": -1, "avg_shortest_path_length": -1,
}) "avg_clustering": -1,
}
)
return metrics return metrics
@ -1196,7 +1218,11 @@ class FalkorDBAdapter(VectorDBInterface, GraphDBInterface):
neighbor_result = self.query(neighbor_query, {"ids": primary_ids}) neighbor_result = self.query(neighbor_query, {"ids": primary_ids})
# FalkorDB returns values by index: id, properties # FalkorDB returns values by index: id, properties
neighbor_ids = [record[0] for record in neighbor_result.result_set] if neighbor_result.result_set else [] neighbor_ids = (
[record[0] for record in neighbor_result.result_set]
if neighbor_result.result_set
else []
)
all_ids = list(set(primary_ids + neighbor_ids)) all_ids = list(set(primary_ids + neighbor_ids))
@ -1226,12 +1252,14 @@ class FalkorDBAdapter(VectorDBInterface, GraphDBInterface):
if edges_result.result_set: if edges_result.result_set:
for record in edges_result.result_set: for record in edges_result.result_set:
# FalkorDB returns values by index: source_id, target_id, relationship_name, properties # FalkorDB returns values by index: source_id, target_id, relationship_name, properties
edges.append(( edges.append(
record[0], # source_id (
record[1], # target_id record[0], # source_id
record[2], # relationship_name record[1], # target_id
record[3] # properties record[2], # relationship_name
)) record[3], # properties
)
)
return nodes, edges return nodes, edges

7
cognee/tests/test_falkordb.py
View file

@ -14,6 +14,7 @@ async def check_falkordb_connection():
"""Check if FalkorDB is available at localhost:6379""" """Check if FalkorDB is available at localhost:6379"""
try: try:
from falkordb import FalkorDB from falkordb import FalkorDB
client = FalkorDB(host="localhost", port=6379) client = FalkorDB(host="localhost", port=6379)
# Try to list graphs to check connection # Try to list graphs to check connection
client.list_graphs() client.list_graphs()
@ -130,7 +131,7 @@ async def main():
# For FalkorDB vector engine, check if collections are empty # For FalkorDB vector engine, check if collections are empty
# Since FalkorDB is a hybrid adapter, we can check if the graph is empty # Since FalkorDB is a hybrid adapter, we can check if the graph is empty
# as the vector data is stored in the same graph # as the vector data is stored in the same graph
if hasattr(vector_engine, 'driver'): if hasattr(vector_engine, "driver"):
# This is FalkorDB - check if graphs exist # This is FalkorDB - check if graphs exist
collections = vector_engine.driver.list_graphs() collections = vector_engine.driver.list_graphs()
# The graph should be deleted, so either no graphs or empty graph # The graph should be deleted, so either no graphs or empty graph
@ -138,7 +139,9 @@ async def main():
# Graph exists but should be empty # Graph exists but should be empty
vector_graph_data = await vector_engine.get_graph_data() vector_graph_data = await vector_engine.get_graph_data()
vector_nodes, vector_edges = vector_graph_data vector_nodes, vector_edges = vector_graph_data
assert len(vector_nodes) == 0 and len(vector_edges) == 0, "FalkorDB vector database is not empty" assert len(vector_nodes) == 0 and len(vector_edges) == 0, (
"FalkorDB vector database is not empty"
)
else: else:
# Fallback for other vector engines like LanceDB # Fallback for other vector engines like LanceDB
connection = await vector_engine.get_connection() connection = await vector_engine.get_connection()

35
debug_falkordb.py
View file

@ -10,6 +10,7 @@ async def debug_falkordb():
# Check if FalkorDB is available # Check if FalkorDB is available
try: try:
from falkordb import FalkorDB from falkordb import FalkorDB
client = FalkorDB(host="localhost", port=6379) client = FalkorDB(host="localhost", port=6379)
client.list_graphs() client.list_graphs()
print("✅ FalkorDB connection successful") print("✅ FalkorDB connection successful")
@ -18,17 +19,21 @@ async def debug_falkordb():
return return
# Configure FalkorDB # Configure FalkorDB
cognee.config.set_graph_db_config({ cognee.config.set_graph_db_config(
"graph_database_url": "localhost", {
"graph_database_port": 6379, "graph_database_url": "localhost",
"graph_database_provider": "falkordb", "graph_database_port": 6379,
}) "graph_database_provider": "falkordb",
}
)
cognee.config.set_vector_db_config({ cognee.config.set_vector_db_config(
"vector_db_url": "localhost", {
"vector_db_port": 6379, "vector_db_url": "localhost",
"vector_db_provider": "falkordb", "vector_db_port": 6379,
}) "vector_db_provider": "falkordb",
}
)
# Set up directories # Set up directories
data_directory_path = str(pathlib.Path("./debug_data").resolve()) data_directory_path = str(pathlib.Path("./debug_data").resolve())
@ -71,19 +76,19 @@ async def debug_falkordb():
if nodes: if nodes:
print("\n🏷️ Sample nodes:") print("\n🏷️ Sample nodes:")
for i, (node_id, node_props) in enumerate(nodes[:3]): for i, (node_id, node_props) in enumerate(nodes[:3]):
print(f" Node {i+1}: ID={node_id}") print(f" Node {i + 1}: ID={node_id}")
print(f" Properties: {node_props}") print(f" Properties: {node_props}")
if edges: if edges:
print("\n🔗 Sample edges:") print("\n🔗 Sample edges:")
for i, edge in enumerate(edges[:3]): for i, edge in enumerate(edges[:3]):
print(f" Edge {i+1}: {edge}") print(f" Edge {i + 1}: {edge}")
# Try different search variations # Try different search variations
print("\n🔍 Testing different search queries...") print("\n🔍 Testing different search queries...")
# Get available graphs and collections # Get available graphs and collections
if hasattr(vector_engine, 'driver'): if hasattr(vector_engine, "driver"):
graphs = vector_engine.driver.list_graphs() graphs = vector_engine.driver.list_graphs()
print(f"Available graphs: {graphs}") print(f"Available graphs: {graphs}")
@ -109,7 +114,9 @@ async def debug_falkordb():
for query_desc, collection_name, query_text in search_queries: for query_desc, collection_name, query_text in search_queries:
try: try:
results = await vector_engine.search(collection_name=collection_name, query_text=query_text) results = await vector_engine.search(
collection_name=collection_name, query_text=query_text
)
print(f" {query_desc}: {len(results)} results") print(f" {query_desc}: {len(results)} results")
if results: if results:
print(f" First result: {results[0]}") print(f" First result: {results[0]}")