Custom ontology

examples/podcast/podcast_runner.py

@@ -20,6 +20,7 @@
 import sys
 
 from dotenv import load_dotenv
+from pydantic import BaseModel, Field
 from transcript_parser import parse_podcast_messages
 
 from graphiti_core import Graphiti
@@ -53,6 +54,12 @@ def setup_logging():
     return logger
 
 
+class Person(BaseModel):
+    first_name: str | None = Field(..., description='First name')
+    last_name: str | None = Field(..., description='Last name')
+    occupation: str | None = Field(..., description="The person's work occupation")
+
+
 async def main():
     setup_logging()
     client = Graphiti(neo4j_uri, neo4j_user, neo4j_password)
@@ -67,6 +74,7 @@ async def main():
             reference_time=message.actual_timestamp,
             source_description='Podcast Transcript',
             group_id='podcast',
+            entity_types={'Person': Person},
         )
 
 

graphiti_core/graphiti.py

@@ -262,6 +262,7 @@ async def add_episode(
         group_id: str = '',
         uuid: str | None = None,
         update_communities: bool = False,
+        entity_types: dict[str, BaseModel] | None = None,
     ) -> AddEpisodeResults:
         """
         Process an episode and update the graph.
@@ -336,7 +337,9 @@ async def add_episode_endpoint(episode_data: EpisodeData):
 
             # Extract entities as nodes
 
-            extracted_nodes = await extract_nodes(self.llm_client, episode, previous_episodes)
+            extracted_nodes = await extract_nodes(
+                self.llm_client, episode, previous_episodes, entity_types
+            )
             logger.debug(f'Extracted nodes: {[(n.name, n.uuid) for n in extracted_nodes]}')
 
             # Calculate Embeddings
@@ -362,6 +365,7 @@ async def add_episode_endpoint(episode_data: EpisodeData):
                     existing_nodes_lists,
                     episode,
                     previous_episodes,
+                    entity_types,
                 ),
                 extract_edges(
                     self.llm_client, episode, extracted_nodes, previous_episodes, group_id

graphiti_core/models/nodes/node_db_queries.py

@@ -31,14 +31,16 @@
 
 ENTITY_NODE_SAVE = """
         MERGE (n:Entity {uuid: $uuid})
-        SET n = {uuid: $uuid, name: $name, group_id: $group_id, summary: $summary, created_at: $created_at}
+        SET n:$($labels)
+        SET n = $entity_data
         WITH n CALL db.create.setNodeVectorProperty(n, "name_embedding", $name_embedding)
         RETURN n.uuid AS uuid"""
 
 ENTITY_NODE_SAVE_BULK = """
     UNWIND $nodes AS node
     MERGE (n:Entity {uuid: node.uuid})
-    SET n = {uuid: node.uuid, name: node.name, group_id: node.group_id, summary: node.summary, created_at: node.created_at}
+    SET n:$(node.labels)
+    SET n = node
     WITH n, node CALL db.create.setNodeVectorProperty(n, "name_embedding", node.name_embedding)
     RETURN n.uuid AS uuid
 """