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

Code before:
from cognee.infrastructure.databases.graph import get_graph_engine
from cognee.infrastructure.engine.models.DataPoint import DataPoint
from cognee.infrastructure.llm.LLMGateway import LLMGateway
from cognee.modules.pipelines.tasks.task import Task
from cognee.modules.engine.models.Entity import Entity
from cognee.tasks.storage import add_data_points
from cognee.infrastructure.llm.prompts import render_prompt

prompt_name = "consolidate_entity_details.txt"


class NodeDescription(BaseModel):
    description: str


async def _fetch_entity_neighbors(args) -> List[Dict[str, Any]]:
    """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,

Review comment (suggestion):
In general, let's try to make this object as light as possible, so we don't have more stuff in memory than needed. 

Code after:
from cognee.infrastructure.databases.graph import get_graph_engine
from cognee.infrastructure.engine.models.DataPoint import DataPoint
from cognee.infrastructure.llm.LLMGateway import LLMGateway
from cognee.modules.pipelines.tasks.task import Task
from cognee.modules.engine.models.Entity import Entity
from cognee.tasks.storage import add_data_points
from cognee.infrastructure.llm.prompts import render_prompt
from cognee.modules.engine.models import EntityType

prompt_name = "consolidate_entity_details.txt"


class NodeDescription(BaseModel):
    description: str


# 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]: