Repository: topoteretes/cognee
File: cognee/memify_pipelines/consolidate_entity_descriptions.py (Python)

Code before:
    """Iterate through all Entity nodes and fetch their edges and neighbor nodes."""
    graph_engine = await get_graph_engine()
    entity_nodes, _ = await graph_engine.get_filtered_graph_data([{"type": ["Entity"]}])

    async def fetch_one(node_id, props):
        edges, neighbors = await asyncio.gather(
            graph_engine.get_edges(node_id),
            graph_engine.get_neighbors(node_id),
        )
        return {
            "id": node_id,
            "properties": props,
            "edges": edges,
            "neighbors": neighbors,
        }

    return await asyncio.gather(*(fetch_one(node_id, props) for node_id, props in entity_nodes))


async def _consolidate_entity_descriptions(param) -> List[DataPoint]:
    system_prompt = render_prompt(prompt_name, {})

    enriched_data = []
    for node in param:
        props = node["properties"]
        text = json.dumps(
            {
                "current_description": props.get("description"),
                "neighbors": node["neighbors"],
                "edges": node["edges"],
            },
            default=str,
        )
        result = await LLMGateway.acreate_structured_output(
            text_input=text,
            system_prompt=system_prompt,  # no format()
            response_model=NodeDescription,
        )
        entity = Entity(
            id=UUID(props["id"]),
            name=props["name"],
            description=result.description,
            created_at=props.get("created_at"),
            updated_at=props.get("updated_at"),
            version=props.get("version", 1),
            topological_rank=props.get("topological_rank", 0),
            ontology_valid=props.get("ontology_valid", False),
            metadata=(
                json.loads(props["metadata"])
                if isinstance(props.get("metadata"), str)
                else {"index_fields": ["name"]}

Review comment (refactor):
The text that we send to the LLM for consolidation is too noisy. Could you please review the flow of what gets extracted and what gets sent, to make sure that we are only sending natural language relevant for summarization.

Code after:
# region get_entities_with_neighborhood helper functions
async def get_all_entity_nodes(graph_engine):
    """Retrieve all nodes of type Entity from the graph."""
    nodes, _ = await graph_engine.get_filtered_graph_data([{"type": ["Entity"]}])
    return nodes


async def get_entity_neighborhood(
    node_id: str, props: Dict[str, Any], graph_engine
) -> Dict[str, Any]:
    """Fetch and format data for a single entity node."""
    edges, neighbors = await asyncio.gather(
        graph_engine.get_edges(node_id),
        graph_engine.get_neighbors(node_id),
    )

    entity_type, filtered_neighbors = format_neighbors(neighbors)
    return {
        "properties": get_entity_properties(props),
        "edges": format_edges(edges),
        "neighbors": filtered_neighbors,
        "entity_type": entity_type,
    }


def get_entity_properties(
    props: Dict[str, Any], properties: Optional[Set[str]] = None
) -> Dict[str, Any]:
    """Keep only relevant entity properties."""
    if properties is None:
        properties = {"id", "description", "name"}
    return {k: v for k, v in props.items() if k in properties}


def format_edges(edges: List[Any]) -> Dict[str, str]:
    """Map target node IDs to relationship names."""
    return {edge[1]: edge[2]["relationship_name"] for edge in edges}


def format_neighbors(
    neighbors: List[Dict[str, Any]], node_fields: Optional[Set[str]] = None
) -> tuple[Optional[Dict[str, Any]], List[Dict[str, Any]]]:
    """Filter neighbor fields and exclude those with only an ID, returning EntityType separately."""
    if node_fields is None:
        node_fields = {"id", "name", "description", "text", "type"}

    entity_type = None
    filtered_neighbors: List[Dict[str, Any]] = []
    for neighbor in neighbors:
        if neighbor.get("type") == "EntityType":
            entity_type = neighbor