Repository: topoteretes/cognee
File: cognee/modules/codingagents/coding_rule_associations.py (Python)

Code before:
            try:
                rule_id = getattr(rule, "id", None)
                if rule_id is not None:
                    rel_name = "rule_associated_from"
                    relationships.append(
                        (
                            rule_id,
                            origin_id,
                            rel_name,
                            {
                                "relationship_name": rel_name,
                                "source_node_id": rule_id,
                                "target_node_id": origin_id,
                                "ontology_valid": False,
                            },
                        )
                    )
            except Exception as e:
                logger.info(f"Warning: Skipping invalid rule due to error: {e}")
    else:
        logger.info("No valid origin_id or rules provided.")

    return relationships


async def rule_associations(data: str, rules_nodeset_name: str):
    llm_client = get_llm_client()
    graph_engine = await get_graph_engine()
    existing_rules = await get_existing_rules(rules_nodeset_name=rules_nodeset_name)

    user_context = {"chat": data, "rules": existing_rules}

    user_prompt = render_prompt("coding_rule_association_agent_user.txt", context=user_context)
    system_prompt = render_prompt("coding_rule_association_agent_system.txt", context={})

    rule_list = await llm_client.acreate_structured_output(
        text_input=user_prompt, system_prompt=system_prompt, response_model=RuleSet
    )

    rules_nodeset = NodeSet(
        id=uuid5(NAMESPACE_OID, name=rules_nodeset_name), name=rules_nodeset_name
    )
    for rule in rule_list.rules:
        rule.belongs_to_set = rules_nodeset

    edges_to_save = await get_origin_edges(data=data, rules=rule_list.rules)

    await add_data_points(data_points=rule_list.rules)

    if len(edges_to_save) > 0:
        await graph_engine.add_edges(edges_to_save)

Review comment (suggestion):
Since it is adding data, maybe name it more explicit, like `add_rule_associations` or something like that.

Code after:
            try:
                rule_id = getattr(rule, "id", None)
                if rule_id is not None:
                    rel_name = "rule_associated_from"
                    relationships.append(
                        (
                            rule_id,
                            origin_id,
                            rel_name,
                            {
                                "relationship_name": rel_name,
                                "source_node_id": rule_id,
                                "target_node_id": origin_id,
                                "ontology_valid": False,
                            },
                        )
                    )
            except Exception as e:
                logger.info(f"Warning: Skipping invalid rule due to error: {e}")
    else:
        logger.info("No valid origin_id or rules provided.")

    return relationships


async def add_rule_associations(data: str, rules_nodeset_name: str):
    llm_client = get_llm_client()
    graph_engine = await get_graph_engine()
    existing_rules = await get_existing_rules(rules_nodeset_name=rules_nodeset_name)

    user_context = {"chat": data, "rules": existing_rules}

    user_prompt = render_prompt("coding_rule_association_agent_user.txt", context=user_context)
    system_prompt = render_prompt("coding_rule_association_agent_system.txt", context={})

    rule_list = await llm_client.acreate_structured_output(
        text_input=user_prompt, system_prompt=system_prompt, response_model=RuleSet
    )

    rules_nodeset = NodeSet(
        id=uuid5(NAMESPACE_OID, name=rules_nodeset_name), name=rules_nodeset_name
    )
    for rule in rule_list.rules:
        rule.belongs_to_set = rules_nodeset

    edges_to_save = await get_origin_edges(data=data, rules=rule_list.rules)

    await add_data_points(data_points=rule_list.rules)

    if len(edges_to_save) > 0:
        await graph_engine.add_edges(edges_to_save)