Custom Entity Types

Graphiti supports custom ontologies, allowing you to define specific types with custom attributes for the nodes in your knowledge graph. This feature enables more precise and domain-specific knowledge representation. This guide explains how to create user-defined entities to model your domain knowledge.

Defining Custom Ontologies

You can define custom entity types as Pydantic models when adding episodes to your knowledge graph.

Creating Entity Types

The code below is an example of defining new entity types, Customer and Product, by extending the BaseModel class from Pydantic:

1
from pydantic import BaseModel, Field
2
3
class Customer(BaseModel):
4
    """A customer of the service"""
5
    name: str | None = Field(..., description="The name of the customer")
6
    email: str | None = Field(..., description="The email address of the customer")
7
    subscription_tier: str | None  = Field(..., description="The customer's subscription level")
8
9
10
class Product(BaseModel):
11
    """A product or service offering"""
12
    price: float | None  = Field(..., description="The price of the product or service")
13
    category: str | None  = Field(..., description="The category of the product")

Adding Data with Entities

Now when you add episodes to your graph, you can pass in a dictionary of the Pydantic models you created above - this will ensure that new nodes are classified into one of the provided types (or none of the provided types) and the attributes of that type are automatically populated:

1
entity_types = {"Customer": Customer, "Product": Product}
2
3
await client.add_episode(
4
            name='Message',
5
            episode_body="New customer John ([email protected]) signed up for premium tier and purchased our Analytics Pro product ($199.99) from the Software category." ,
6
            reference_time=datetime.now(),
7
            source_description='Support Ticket Log',
8
            group_id=group_id,
9
            entity_types=entity_types,
10
        )

Results

When you provide custom entity types, Graphiti will:

• Extract entities and classify them according to your defined types
• Identify and populate the provided attributes of each type

This will affect the node.labels and node.attributes fields of the extracted nodes:

1
# Example Customer Node Attributes
2
{
3
    ...
4
    "labels": ["Entity","Customer"],
5
    "attributes": {
6
        "name": "John",
7
        "email": "[email protected]",
8
        "subscription_tier": "premium",
9
    }
10
}
11
12
# Example Product Node Attributes
13
{
14
    ...
15
    "labels": ["Entity", "Product"],
16
    "attributes": {
17
        "price": 199.99,
18
        "category": "Software",
19
    }
20
}

Schema Evolution

Your knowledge graph’s schema can evolve over time as your needs change. You can update entity types by adding new attributes to existing types without breaking existing nodes. When you add new attributes, existing nodes will preserve their original attributes while supporting the new ones for future updates. This flexible approach allows your knowledge graph to grow and adapt while maintaining backward compatibility with historical data.

For example, if you initially defined a “Customer” type with basic attributes like name and email, you could later add attributes like “loyalty_tier” or “acquisition_channel” without needing to modify or migrate existing customer nodes in your graph.

Best Practices

When extracting attributes, maintain consistent naming conventions across related entity types and include a clear and thorough description of each attribute.

Additionally, attributes should be broken down into their smallest meaningful units rather than storing compound information.

Instead of:

1
class Employment(BaseModel):
2
    """User's current employment"""
3
    employment: str | None = Field(..., description="Employment details")

Use:

1
class Employment(BaseModel):
2
    """User's current employment"""
3
    role: str | None = Field(..., description="Job title")
4
    employer: str | None = Field(..., description="Employer name")
5
    ...

Migration Guide

If you’re upgrading from a previous version of Graphiti:

• You can add entity types to new episodes, even if existing episodes in the graph did not have entity types. Existing nodes will continue to work without being classified.
• To add types to previously ingested data, you need to re-ingest it with entity types set into a new graph.