from collections import OrderedDict from datetime import datetime, timezone from threading import Lock from typing import Tuple, List, Any, Dict, Optional from cognee.infrastructure.engine import DataPoint, Edge from cognee.modules.storage.utils import copy_model from cognee.shared.logging_utils import get_logger logger = get_logger() # Memoized simple-node pydantic classes. Without this, every call to # ``get_graph_from_model`` — one per DataPoint added to the graph — re-ran # ``copy_model`` and minted a new ``BaseModel`` subclass, each of which # attached fresh ``FieldInfo`` / ``SchemaValidator`` / ``SchemaSerializer`` # state to pydantic's global caches and never released it. Tracemalloc # attributed +~50 MB per large-text cognify cycle to pydantic internals; # this cache is keyed by ``(DataPoint subclass, sorted excluded fields)`` so # different call sites with the same shape share one class. # # Bounded LRU. An unbounded dict would itself grow without limit if # call-site exclusions vary, defeating the leak fix it was added for. _SIMPLE_MODEL_CACHE_SIZE = 256 _SIMPLE_MODEL_CACHE: "OrderedDict" = OrderedDict() _SIMPLE_MODEL_CACHE_LOCK = Lock() def _simple_model_for(data_point_type, excluded_fields): key = (data_point_type, tuple(sorted(excluded_fields))) with _SIMPLE_MODEL_CACHE_LOCK: cached = _SIMPLE_MODEL_CACHE.get(key) if cached is not None: _SIMPLE_MODEL_CACHE.move_to_end(key) return cached model = copy_model(data_point_type, exclude_fields=list(excluded_fields)) with _SIMPLE_MODEL_CACHE_LOCK: # Re-check after the heavy ``copy_model`` — another thread may # have raced us; if so, return the winner and discard our build. existing = _SIMPLE_MODEL_CACHE.get(key) if existing is not None: _SIMPLE_MODEL_CACHE.move_to_end(key) return existing _SIMPLE_MODEL_CACHE[key] = model if len(_SIMPLE_MODEL_CACHE) > _SIMPLE_MODEL_CACHE_SIZE: _SIMPLE_MODEL_CACHE.popitem(last=False) return model def _extract_field_data(field_value: Any) -> List[Tuple[Optional[Edge], List[DataPoint]]]: """Extract edge metadata and datapoints from a field value.""" # Handle single DataPoint if isinstance(field_value, DataPoint): return [(None, [field_value])] # Handle list - could contain DataPoints, edge tuples, or mixed if isinstance(field_value, list) and len(field_value) > 0: result = [] for item in field_value: # Handle tuple[Edge, DataPoint or list[DataPoint]] if isinstance(item, tuple) and len(item) == 2 and isinstance(item[0], Edge): edge, data_value = item if isinstance(data_value, DataPoint): result.append((edge, [data_value])) elif ( isinstance(data_value, list) and len(data_value) > 0 and isinstance(data_value[0], DataPoint) ): result.append((edge, data_value)) # Handle single DataPoint in list elif isinstance(item, DataPoint): result.append((None, [item])) return result # Handle tuple[Edge, DataPoint or list[DataPoint]] if ( isinstance(field_value, tuple) and len(field_value) == 2 and isinstance(field_value[0], Edge) ): edge_metadata, data_value = field_value if isinstance(data_value, DataPoint): return [(edge_metadata, [data_value])] elif ( isinstance(data_value, list) and len(data_value) > 0 and isinstance(data_value[0], DataPoint) ): return [(edge_metadata, data_value)] # Regular property or empty list return [] def _create_edge_properties( source_id: str, target_id: str, relationship_name: str, edge_metadata: Optional[Edge] ) -> Dict[str, Any]: """Create edge properties dictionary with metadata if present.""" properties = { "source_node_id": source_id, "target_node_id": target_id, "relationship_name": relationship_name, "updated_at": datetime.now(timezone.utc).strftime("%Y-%m-%d %H:%M:%S"), } if edge_metadata: # Add edge metadata edge_data = edge_metadata.model_dump(exclude_none=True) properties.update(edge_data) # Add individual weights as separate fields for easier querying if edge_metadata.weights is not None: for weight_name, weight_value in edge_metadata.weights.items(): properties[f"weight_{weight_name}"] = weight_value return properties def _get_relationship_key(field_name: str, edge_metadata: Optional[Edge]) -> str: """Extract relationship key from edge metadata or use field name as fallback.""" if ( edge_metadata and hasattr(edge_metadata, "relationship_type") and edge_metadata.relationship_type ): return edge_metadata.relationship_type return field_name def _generate_property_key(data_point_id: str, relationship_key: str, target_id: str) -> str: """Generate a unique property key for visited_properties tracking.""" return f"{data_point_id}_{relationship_key}_{target_id}" def _process_datapoint_field( data_point_id: str, field_name: str, edge_datapoint_pairs: List[Tuple[Optional[Edge], List[DataPoint]]], visited_properties: Dict[str, bool], properties_to_visit: set, excluded_properties: set, ) -> None: """Process a field containing DataPoints, always working with lists.""" if field_name != "belongs_to_set": excluded_properties.add(field_name) for edge_metadata, datapoints in edge_datapoint_pairs: relationship_key = _get_relationship_key(field_name, edge_metadata) for datapoint in datapoints: property_key = _generate_property_key( data_point_id, relationship_key, str(datapoint.id) ) if property_key in visited_properties: continue # Always use field_name since we're working with lists properties_to_visit.add(field_name) def _targets_generator( data_point: DataPoint, properties_to_visit: set, ) -> Tuple[DataPoint, str, Optional[Edge]]: """Generator that yields (target_datapoint, field_name, edge_metadata) tuples.""" for field_name in properties_to_visit: field_value = getattr(data_point, field_name) edge_datapoint_pairs = _extract_field_data(field_value) if not edge_datapoint_pairs: continue for edge_metadata, datapoints in edge_datapoint_pairs: for target_datapoint in datapoints: yield target_datapoint, field_name, edge_metadata async def get_graph_from_model( data_point: DataPoint, added_nodes: Optional[Dict[str, bool]] = None, added_edges: Optional[Dict[str, bool]] = None, visited_properties: Optional[Dict[str, bool]] = None, include_root: bool = True, ) -> Tuple[List[DataPoint], List[Tuple[str, str, str, Dict[str, Any]]]]: """ Extract graph representation from a DataPoint model. Args: data_point: The DataPoint to extract graph from added_nodes: Dictionary tracking already processed nodes added_edges: Dictionary tracking already processed edges visited_properties: Dictionary tracking visited properties to avoid cycles include_root: Whether to include the root node in results Returns: Tuple of (nodes, edges) extracted from the model """ if added_nodes is None: added_nodes = {} if added_edges is None: added_edges = {} if str(data_point.id) in added_nodes: logger.debug( "Skipping already processed DataPoint", extra={"datapoint_id": str(data_point.id)}, ) return [], [] nodes = [] edges = [] visited_properties = visited_properties or {} data_point_id = str(data_point.id) logger.debug( "Starting graph extraction for DataPoint", extra={ "datapoint_id": data_point_id, "datapoint_type": type(data_point).__name__, "processed_nodes_so_far": len(added_nodes), }, ) data_point_properties = {"id": data_point.id, "type": type(data_point).__name__} excluded_properties = set() properties_to_visit = set() # Analyze all fields to categorize them as properties or relationships for field_name, field_value in data_point: if field_name == "metadata": continue edge_datapoint_pairs = _extract_field_data(field_value) if not edge_datapoint_pairs: # Regular property data_point_properties[field_name] = field_value else: # DataPoint relationship _process_datapoint_field( data_point_id, field_name, edge_datapoint_pairs, visited_properties, properties_to_visit, excluded_properties, ) # We want to enable nodeset filtering on the vector database side if field_name == "belongs_to_set": node_set_names = [] for node_set in field_value: if isinstance(node_set, str): node_set_names.append(node_set) elif hasattr(node_set, "name"): node_set_names.append(node_set.name) data_point_properties[field_name] = node_set_names # Create node for current DataPoint if needed if include_root and data_point_id not in added_nodes: SimpleDataPointModel = _simple_model_for(type(data_point), excluded_properties) nodes.append(SimpleDataPointModel(**data_point_properties)) added_nodes[data_point_id] = True logger.debug( "Added node to graph", extra={ "datapoint_id": data_point_id, "node_type": type(data_point).__name__, }, ) # Process all relationships using generator for target_datapoint, field_name, edge_metadata in _targets_generator( data_point, properties_to_visit ): relationship_name = _get_relationship_key(field_name, edge_metadata) # Create edge if not already added edge_key = f"{data_point_id}_{target_datapoint.id}_{field_name}" if edge_key not in added_edges: edge_properties = _create_edge_properties( data_point.id, target_datapoint.id, relationship_name, edge_metadata ) edges.append((data_point.id, target_datapoint.id, relationship_name, edge_properties)) logger.debug( "Added edge to graph", extra={ "source_id": str(data_point.id), "target_id": str(target_datapoint.id), "relationship": relationship_name, }, ) added_edges[edge_key] = True # Mark property as visited - CRITICAL for preventing infinite loops property_key = _generate_property_key( data_point_id, relationship_name, str(target_datapoint.id) ) visited_properties[property_key] = True # Recursively process target node if not already processed if str(target_datapoint.id) in added_nodes: continue logger.debug( "Recursing into target DataPoint", extra={ "source_id": data_point_id, "target_id": str(target_datapoint.id), }, ) child_nodes, child_edges = await get_graph_from_model( target_datapoint, include_root=True, added_nodes=added_nodes, added_edges=added_edges, visited_properties=visited_properties, ) nodes.extend(child_nodes) edges.extend(child_edges) logger.info( "Completed graph extraction for DataPoint", extra={ "datapoint_id": data_point_id, "nodes_extracted": len(nodes), "edges_extracted": len(edges), }, ) return nodes, edges def get_own_property_nodes( property_nodes: List[DataPoint], property_edges: List[Tuple[str, str, str, Dict[str, Any]]] ) -> List[DataPoint]: """ Filter nodes to return only those that are not destinations of any edges. Args: property_nodes: List of all nodes property_edges: List of all edges Returns: List of nodes that are not edge destinations """ destination_node_ids = {str(edge[1]) for edge in property_edges} return [node for node in property_nodes if str(node.id) not in destination_node_ids]