Custom Entity and Edge Types

1
from pydantic import BaseModel, Field
2
from datetime import datetime
3
from typing import Optional
4
5
# Custom Entity Types
6
class Person(BaseModel):
7
    """A person entity with biographical information."""
8
    age: Optional[int] = Field(None, description="Age of the person")
9
    occupation: Optional[str] = Field(None, description="Current occupation")
10
    location: Optional[str] = Field(None, description="Current location")
11
    birth_date: Optional[datetime] = Field(None, description="Date of birth")
12
13
class Company(BaseModel):
14
    """A business organization."""
15
    industry: Optional[str] = Field(None, description="Primary industry")
16
    founded_year: Optional[int] = Field(None, description="Year company was founded")
17
    headquarters: Optional[str] = Field(None, description="Location of headquarters")
18
    employee_count: Optional[int] = Field(None, description="Number of employees")
19
20
class Product(BaseModel):
21
    """A product or service."""
22
    category: Optional[str] = Field(None, description="Product category")
23
    price: Optional[float] = Field(None, description="Price in USD")
24
    release_date: Optional[datetime] = Field(None, description="Product release date")
25
26
# Custom Edge Types
27
class Employment(BaseModel):
28
    """Employment relationship between a person and company."""
29
    position: Optional[str] = Field(None, description="Job title or position")
30
    start_date: Optional[datetime] = Field(None, description="Employment start date")
31
    end_date: Optional[datetime] = Field(None, description="Employment end date")
32
    salary: Optional[float] = Field(None, description="Annual salary in USD")
33
    is_current: Optional[bool] = Field(None, description="Whether employment is current")
34
35
class Investment(BaseModel):
36
    """Investment relationship between entities."""
37
    amount: Optional[float] = Field(None, description="Investment amount in USD")
38
    investment_type: Optional[str] = Field(None, description="Type of investment (equity, debt, etc.)")
39
    stake_percentage: Optional[float] = Field(None, description="Percentage ownership")
40
    investment_date: Optional[datetime] = Field(None, description="Date of investment")
41
42
class Partnership(BaseModel):
43
    """Partnership relationship between companies."""
44
    partnership_type: Optional[str] = Field(None, description="Type of partnership")
45
    duration: Optional[str] = Field(None, description="Expected duration")
46
    deal_value: Optional[float] = Field(None, description="Financial value of partnership")

1
entity_types = {
2
    "Person": Person,
3
    "Company": Company,
4
    "Product": Product
5
}
6
7
edge_types = {
8
    "Employment": Employment,
9
    "Investment": Investment,
10
    "Partnership": Partnership
11
}
12
13
edge_type_map = {
14
    ("Person", "Company"): ["Employment"],
15
    ("Company", "Company"): ["Partnership", "Investment"],
16
    ("Person", "Person"): ["Partnership"],
17
    ("Entity", "Entity"): ["Investment"],  # Apply to any entity type
18
}
19
20
await graphiti.add_episode(
21
    name="Business Update",
22
    episode_body="Sarah joined TechCorp as CTO in January 2023 with a $200K salary. TechCorp partnered with DataCorp in a $5M deal.",
23
    source_description="Business news",
24
    reference_time=datetime.now(),
25
    entity_types=entity_types,
26
    edge_types=edge_types,
27
    edge_type_map=edge_type_map
28
)

1
from graphiti_core.search.search_filters import SearchFilters
2
3
# Search for only specific entity types
4
search_filter = SearchFilters(
5
    node_labels=["Person", "Company"]  # Only return Person and Company entities
6
)
7
8
results = await graphiti.search_(
9
    query="Who works at tech companies?",
10
    search_filter=search_filter
11
)
12
13
# Search for only specific edge types
14
search_filter = SearchFilters(
15
    edge_types=["Employment", "Partnership"]  # Only return Employment and Partnership edges
16
)
17
18
results = await graphiti.search_(
19
    query="Tell me about business relationships",
20
    search_filter=search_filter
21
)

1
edge_type_map = {
2
    ("Person", "Company"): ["Employment"],
3
    ("Company", "Company"): ["Partnership", "Investment"],
4
    ("Person", "Person"): ["Partnership"],
5
    ("Entity", "Entity"): ["Investment"],  # Apply to any entity type
6
}

1
from pydantic import validator
2
3
class Person(BaseModel):
4
    """A person entity."""
5
    age: Optional[int] = Field(None, description="Age in years")
6
    
7
    @validator('age')
8
    def validate_age(cls, v):
9
        if v is not None and (v < 0 or v > 150):
10
            raise ValueError('Age must be between 0 and 150')
11
        return v

1
class Customer(BaseModel):
2
    contact_info: Optional[str] = Field(None, description="Name and email")  # Don't do this

1
class Customer(BaseModel):
2
    name: Optional[str] = Field(None, description="Customer name")
3
    email: Optional[str] = Field(None, description="Customer email address")

1
# Good: Specific and meaningful
2
edge_type_map = {
3
    ("Person", "Company"): ["Employment", "Investment"],
4
    ("Company", "Company"): ["Partnership", "Acquisition"],
5
    ("Person", "Product"): ["Usage", "Review"],
6
    ("Entity", "Entity"): ["RELATES_TO"]  # Fallback for unexpected relationships
7
}
8
9
# Avoid: Too granular
10
edge_type_map = {
11
    ("CEO", "TechCompany"): ["CEOEmployment"],
12
    ("Engineer", "TechCompany"): ["EngineerEmployment"],
13
    # This creates too many specific types
14
}

1
await graphiti.add_episode(
2
    name="Business Update",
3
    episode_body="The meeting discussed various topics including weather and sports.",
4
    source_description="Meeting notes",
5
    reference_time=datetime.now(),
6
    entity_types=entity_types,
7
    excluded_entity_types=["Person"]  # Won't extract Person entities
8
)