348610e73c
* fix: handle rate limit error coming from llm model * fix: fixes lost edges and nodes in get_graph_from_model * fix: fixes database pruning issue in pgvector (#261) * fix: cognee_demo notebook pipeline is not saving summaries --------- Co-authored-by: hajdul88 <52442977+hajdul88@users.noreply.github.com>
63 lines
2 KiB
Python
63 lines
2 KiB
Python
import argparse
|
|
import asyncio
|
|
|
|
from .benchmark_function import benchmark_function
|
|
|
|
from cognee.modules.graph.utils import get_graph_from_model
|
|
from cognee.tests.unit.interfaces.graph.util import (
|
|
PERSON_NAMES,
|
|
create_organization_recursive,
|
|
)
|
|
|
|
# Example usage:
|
|
if __name__ == "__main__":
|
|
parser = argparse.ArgumentParser(
|
|
description="Benchmark graph model with configurable recursive depth"
|
|
)
|
|
parser.add_argument(
|
|
"--recursive-depth",
|
|
type=int,
|
|
default=3,
|
|
help="Recursive depth for graph generation (default: 3)",
|
|
)
|
|
parser.add_argument(
|
|
"--runs", type=int, default=5, help="Number of benchmark runs (default: 5)"
|
|
)
|
|
args = parser.parse_args()
|
|
|
|
society = create_organization_recursive(
|
|
"society", "Society", PERSON_NAMES, args.recursive_depth
|
|
)
|
|
added_nodes = {}
|
|
added_edges = {}
|
|
visited_properties = {}
|
|
nodes, edges = asyncio.run(get_graph_from_model(
|
|
society,
|
|
added_nodes = added_nodes,
|
|
added_edges = added_edges,
|
|
visited_properties = visited_properties,
|
|
))
|
|
|
|
def get_graph_from_model_sync(model):
|
|
added_nodes = {}
|
|
added_edges = {}
|
|
visited_properties = {}
|
|
|
|
return asyncio.run(get_graph_from_model(
|
|
model,
|
|
added_nodes = added_nodes,
|
|
added_edges = added_edges,
|
|
visited_properties = visited_properties,
|
|
))
|
|
|
|
results = benchmark_function(get_graph_from_model_sync, society, num_runs=args.runs)
|
|
print("\nBenchmark Results:")
|
|
print(
|
|
f"N nodes: {len(nodes)}, N edges: {len(edges)}, Recursion depth: {args.recursive_depth}"
|
|
)
|
|
print(f"Mean Peak Memory: {results['mean_peak_memory_mb']:.2f} MB")
|
|
print(f"Mean CPU Usage: {results['mean_cpu_percent']:.2f}%")
|
|
print(f"Mean Execution Time: {results['mean_execution_time']:.4f} seconds")
|
|
|
|
if "std_execution_time" in results:
|
|
print(f"Execution Time Std: {results['std_execution_time']:.4f} seconds")
|