error

(base) liorkesos@mbp-14 symptoms % python main.py
- RELATED_TO_FILE: from Symptom to File
- RELATED_TO_CONFIG: from Symptom to Configuration
- RELATED_TO_TECHNOLOGY: from Symptom to Technology
<falkordb.graph_schema.GraphSchema object at 0x106888620>
Traceback (most recent call last):
  File "/Users/liorkesos/hoss/research/llamaindex/symptoms/main.py", line 53, in <module>
    symptoms_kg.add_edge(
  File "/Users/liorkesos/.pyenv/versions/3.12.2/lib/python3.12/site-packages/graphrag_sdk/kg.py", line 242, in add_edge
    self._validate_relation(
  File "/Users/liorkesos/.pyenv/versions/3.12.2/lib/python3.12/site-packages/graphrag_sdk/kg.py", line 276, in _validate_relation
    raise Exception(f"Relation {relation} not found in ontology")
Exception: Relation RELATED_TO_CONFIG not found in ontology

main.py

from json import dumps
from dotenv import load_dotenv

from graphrag_sdk.models.openai import OpenAiGenerativeModel
from symptoms_ontology import load_symptoms_ontology
from utils import get_technology_symptoms, list_available_technologies
from graphrag_sdk import KnowledgeGraph, KnowledgeGraphModelConfig,Ontology

load_dotenv()

# Define the model
model = OpenAiGenerativeModel("gpt-4o")

# Load the ontology
symptoms_ontology = load_symptoms_ontology()
# Create the KG from the predefined ontology.

symptoms_kg = KnowledgeGraph(
    name="symptomGraph",
    ontology=symptoms_ontology,
    model_config=KnowledgeGraphModelConfig.with_model(model),
)

# List relations
for relation in symptoms_ontology.relations:
    print(f"- {relation.label}: from {relation.source.label} to {relation.target.label}")
# Add the single node for the technology
# Get symptoms for specific technology
tech_symptoms = get_technology_symptoms("cassandra")
print(symptoms_kg.graph.schema)
symptoms_kg.add_node(
    "Technology",
    {
        "name": tech_symptoms["technology"],
    },
)

for symptom in tech_symptoms["symptoms"]:
    symptoms_kg.add_node(
        "Symptom",
        {
            "description": symptom["description"],
            "severity": symptom["severity"]  
        }
    )
    for config in symptom["configurations"]:
        symptoms_kg.add_node(
            "Configuration",
            {
                "name": config,
            },
        )
        symptoms_kg.add_edge(
            "RELATED_TO_CONFIG",
            symptom["description"],
            config
        )
    for file in symptom["files"]:
        symptoms_kg.add_node(
            "File",
            {
                "name": file,
            },
        )
        symptoms_kg.add_edge(
            "RELATED_TO_FILE",
            symptom["description"],
            file
        )
    
    
# If you still want to save it to disk
with open("ontology.json", "w", encoding="utf-8") as file:
    file.write(dumps(symptoms_ontology.to_json(), indent=2)) 

symptoms_ontology.py

from graphrag_sdk import (
    Ontology, Entity, Relation, Attribute, AttributeType
)

def load_symptoms_ontology() -> Ontology:
    symptoms_ontology  = Ontology()
    # Add Technology entity
    symptoms_ontology.add_entity(
        Entity(
            label="Technology",
            attributes=[
                Attribute(
                    name="name",
                    attr_type=AttributeType.STRING,
                    required=True,
                    unique=True,
                ),
            ],
        )
    )
    
    # Add Symptom entity
    symptoms_ontology.add_entity(
        Entity(
            label="Symptom",
            attributes=[
                Attribute(
                    name="description",
                    attr_type=AttributeType.STRING,
                    required=True,
                    unique=True,
                ),
                Attribute(
                    name="severity",
                    attr_type=AttributeType.STRING,
                    required=True,
                    unique=True,
                ),
            ],
        )
    )
    
    # Add File entity
    symptoms_ontology.add_entity(
        Entity(
            label="File",
            attributes=[
                Attribute(
                    name="name",
                    attr_type=AttributeType.STRING,
                    required=True,
                    unique=True,
                ),
            ],
        )
    )
    
    # Add Configuration entity
    symptoms_ontology.add_entity(
        Entity(
            label="Configuration",
            attributes=[
                Attribute(
                    name="name",
                    attr_type=AttributeType.STRING,
                    required=True,
                    unique=True,
                ),
            ],
        )
    )
    
    # Add relations
    symptoms_ontology.add_relation(
        Relation(
            label="RELATED_TO_FILE",
            source="Symptom",
            target="File",
        )
    )
    symptoms_ontology.add_relation(
        Relation(
            label="RELATED_TO_CONFIG",
            source="Symptom",
            target="Configuration",
        )
    )
    symptoms_ontology.add_relation(
        Relation(
            label="RELATED_TO_TECHNOLOGY",
            source="Symptom",
            target="Technology",
        )
    )
    
    return symptoms_ontology