Skip to content

utils

langroid/agent/special/arangodb/utils.py

trim_schema(schema)

Keep only edge connection info, remove properties and examples

Source code in langroid/agent/special/arangodb/utils.py
def trim_schema(
    schema: Dict[str, List[Dict[str, Any]]]
) -> Dict[str, List[Dict[str, Any]]]:
    """Keep only edge connection info, remove properties and examples"""
    trimmed: Dict[str, List[Dict[str, Any]]] = {
        "Graph Schema": schema["Graph Schema"],
        "Collection Schema": [],
    }
    for coll in schema["Collection Schema"]:
        col_info: Dict[str, Any] = {
            "collection_name": coll["collection_name"],
            "collection_type": coll["collection_type"],
        }
        if coll["collection_type"] == "edge":
            # preserve from/to info if present
            if f"example_{coll['collection_type']}" in coll:
                example = coll[f"example_{coll['collection_type']}"]
                if example and "_from" in example:
                    col_info["from_collection"] = example["_from"].split("/")[0]
                    col_info["to_collection"] = example["_to"].split("/")[0]
        trimmed["Collection Schema"].append(col_info)
    return trimmed