6dbd8e85a1
<!-- .github/pull_request_template.md --> ## Description <!-- Provide a clear description of the changes in this PR --> - Improved list handling, removed `.index` logic from `get_graph_from_model`, transitioned to fully datapoint-oriented processing - Streamlined datapoint iteration by introducing `_datapoints_generator` with nested loops - Generalized field processing to handle mixed lists: `[DataPoint, (Edge, DataPoint), (Edge, [DataPoint])]`, allowing dynamic multiple edges generation - Small improvements and refactorings - Added tests to `test_get_graph_from_model_flexible_edges()` covering weighted edges and dynamic multiple edges - Created `dynamic_multiple_edges_example.py` demonstrating dynamic multiple edges ## DCO Affirmation I affirm that all code in every commit of this pull request conforms to the terms of the Topoteretes Developer Certificate of Origin.
115 lines
3.8 KiB
Python
115 lines
3.8 KiB
Python
import asyncio
|
|
from os import path
|
|
from typing import Any
|
|
from pydantic import SkipValidation
|
|
from cognee.api.v1.visualize.visualize import visualize_graph
|
|
from cognee.infrastructure.engine import DataPoint
|
|
from cognee.infrastructure.engine.models.Edge import Edge
|
|
from cognee.tasks.storage import add_data_points
|
|
import cognee
|
|
|
|
|
|
class Employee(DataPoint):
|
|
name: str
|
|
role: str
|
|
|
|
|
|
class Company(DataPoint):
|
|
name: str
|
|
industry: str
|
|
employs: SkipValidation[Any] # Mixed list: employees with/without weights
|
|
|
|
|
|
async def main():
|
|
# Clear the database for a clean state
|
|
await cognee.prune.prune_data()
|
|
await cognee.prune.prune_system(metadata=True)
|
|
|
|
# Create employees
|
|
michael = Employee(name="Michael", role="Regional Manager")
|
|
dwight = Employee(name="Dwight", role="Assistant to the Regional Manager")
|
|
jim = Employee(name="Jim", role="Sales Representative")
|
|
pam = Employee(name="Pam", role="Receptionist")
|
|
kevin = Employee(name="Kevin", role="Accountant")
|
|
angela = Employee(name="Angela", role="Senior Accountant")
|
|
oscar = Employee(name="Oscar", role="Accountant")
|
|
stanley = Employee(name="Stanley", role="Sales Representative")
|
|
phyllis = Employee(name="Phyllis", role="Sales Representative")
|
|
|
|
# Create Dunder Mifflin with mixed employee relationships
|
|
dunder_mifflin = Company(
|
|
name="Dunder Mifflin Paper Company",
|
|
industry="Paper Sales",
|
|
employs=[
|
|
# Manager with high authority weight
|
|
(Edge(weight=0.9, relationship_type="manager"), michael),
|
|
# Sales team with performance weights
|
|
(
|
|
Edge(weights={"sales_performance": 0.8, "loyalty": 0.9}, relationship_type="sales"),
|
|
dwight,
|
|
),
|
|
(
|
|
Edge(
|
|
weights={"sales_performance": 0.7, "creativity": 0.8}, relationship_type="sales"
|
|
),
|
|
jim,
|
|
),
|
|
(
|
|
Edge(
|
|
weights={"sales_performance": 0.6, "customer_service": 0.9},
|
|
relationship_type="sales",
|
|
),
|
|
phyllis,
|
|
),
|
|
(
|
|
Edge(
|
|
weights={"sales_performance": 0.5, "experience": 0.8}, relationship_type="sales"
|
|
),
|
|
stanley,
|
|
),
|
|
# Accounting department as a group
|
|
(
|
|
Edge(
|
|
weights={"department_efficiency": 0.8, "team_cohesion": 0.9},
|
|
relationship_type="accounting",
|
|
),
|
|
[oscar, kevin, angela],
|
|
),
|
|
# Admin staff without weights (simple relationships)
|
|
pam,
|
|
],
|
|
)
|
|
|
|
all_data_points = [
|
|
michael,
|
|
dwight,
|
|
jim,
|
|
pam,
|
|
kevin,
|
|
angela,
|
|
oscar,
|
|
stanley,
|
|
phyllis,
|
|
dunder_mifflin,
|
|
]
|
|
|
|
# Add data points to the graph
|
|
await add_data_points(all_data_points)
|
|
|
|
# Visualize the graph
|
|
graph_visualization_path = path.join(path.dirname(__file__), "dunder_mifflin_graph.html")
|
|
await visualize_graph(graph_visualization_path)
|
|
|
|
print("Dynamic multiple edges graph has been created and visualized!")
|
|
print(f"Visualization saved to: {graph_visualization_path}")
|
|
print("\nTechnical features demonstrated:")
|
|
print("- Mixed list support: weighted and unweighted relationships in single field")
|
|
print("- Single weight edges with relationship types")
|
|
print("- Multiple weight edges with custom metrics")
|
|
print("- Group relationships: single edge connecting multiple nodes")
|
|
print("- Simple relationships without edge metadata")
|
|
print("- Flexible edge extraction from heterogeneous data structures")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
asyncio.run(main())
|