graphiti-core 0.10.4__py3-none-any.whl → 0.11.0__py3-none-any.whl

graphiti_core/edges.py

@@ -37,6 +37,21 @@ from graphiti_core.nodes import Node
 
 logger = logging.getLogger(__name__)
 
+ENTITY_EDGE_RETURN: LiteralString = """
+        RETURN
+            e.uuid AS uuid,
+            startNode(e).uuid AS source_node_uuid,
+            endNode(e).uuid AS target_node_uuid,
+            e.created_at AS created_at,
+            e.name AS name,
+            e.group_id AS group_id,
+            e.fact AS fact,
+            e.fact_embedding AS fact_embedding,
+            e.episodes AS episodes,
+            e.expired_at AS expired_at,
+            e.valid_at AS valid_at,
+            e.invalid_at AS invalid_at"""
+
 
 class Edge(BaseModel, ABC):
     uuid: str = Field(default_factory=lambda: str(uuid4()))
@@ -234,20 +249,8 @@ class EntityEdge(Edge):
         records, _, _ = await driver.execute_query(
             """
         MATCH (n:Entity)-[e:RELATES_TO {uuid: $uuid}]->(m:Entity)
-        RETURN
-            e.uuid AS uuid,
-            n.uuid AS source_node_uuid,
-            m.uuid AS target_node_uuid,
-            e.created_at AS created_at,
-            e.name AS name,
-            e.group_id AS group_id,
-            e.fact AS fact,
-            e.fact_embedding AS fact_embedding,
-            e.episodes AS episodes,
-            e.expired_at AS expired_at,
-            e.valid_at AS valid_at,
-            e.invalid_at AS invalid_at
-        """,
+        """
+            + ENTITY_EDGE_RETURN,
             uuid=uuid,
             database_=DEFAULT_DATABASE,
             routing_='r',
@@ -268,20 +271,8 @@ class EntityEdge(Edge):
             """
         MATCH (n:Entity)-[e:RELATES_TO]->(m:Entity)
         WHERE e.uuid IN $uuids
-        RETURN
-            e.uuid AS uuid,
-            n.uuid AS source_node_uuid,
-            m.uuid AS target_node_uuid,
-            e.created_at AS created_at,
-            e.name AS name,
-            e.group_id AS group_id,
-            e.fact AS fact,
-            e.fact_embedding AS fact_embedding,
-            e.episodes AS episodes,
-            e.expired_at AS expired_at,
-            e.valid_at AS valid_at,
-            e.invalid_at AS invalid_at
-        """,
+        """
+            + ENTITY_EDGE_RETURN,
             uuids=uuids,
             database_=DEFAULT_DATABASE,
             routing_='r',
@@ -308,20 +299,8 @@ class EntityEdge(Edge):
         WHERE e.group_id IN $group_ids
         """
             + cursor_query
+            + ENTITY_EDGE_RETURN
             + """
-        RETURN
-            e.uuid AS uuid,
-            n.uuid AS source_node_uuid,
-            m.uuid AS target_node_uuid,
-            e.created_at AS created_at,
-            e.name AS name,
-            e.group_id AS group_id,
-            e.fact AS fact,
-            e.fact_embedding AS fact_embedding,
-            e.episodes AS episodes,
-            e.expired_at AS expired_at,
-            e.valid_at AS valid_at,
-            e.invalid_at AS invalid_at
         ORDER BY e.uuid DESC 
         """
             + limit_query,
@@ -340,22 +319,12 @@ class EntityEdge(Edge):
 
     @classmethod
     async def get_by_node_uuid(cls, driver: AsyncDriver, node_uuid: str):
-        query: LiteralString = """
-        MATCH (n:Entity {uuid: $node_uuid})-[e:RELATES_TO]-(m:Entity)
-        RETURN DISTINCT
-            e.uuid AS uuid,
-            n.uuid AS source_node_uuid,
-            m.uuid AS target_node_uuid,
-            e.created_at AS created_at,
-            e.name AS name,
-            e.group_id AS group_id,
-            e.fact AS fact,
-            e.fact_embedding AS fact_embedding,
-            e.episodes AS episodes,
-            e.expired_at AS expired_at,
-            e.valid_at AS valid_at,
-            e.invalid_at AS invalid_at
-        """
+        query: LiteralString = (
+            """
+                        MATCH (n:Entity {uuid: $node_uuid})-[e:RELATES_TO]-(m:Entity)
+                        """
+            + ENTITY_EDGE_RETURN
+        )
         records, _, _ = await driver.execute_query(
             query, node_uuid=node_uuid, database_=DEFAULT_DATABASE, routing_='r'
         )
@@ -499,3 +468,9 @@ def get_community_edge_from_record(record: Any):
         target_node_uuid=record['target_node_uuid'],
         created_at=record['created_at'].to_native(),
     )
+
+
+async def create_entity_edge_embeddings(embedder: EmbedderClient, edges: list[EntityEdge]):
+    fact_embeddings = await embedder.create_batch([edge.fact for edge in edges])
+    for edge, fact_embedding in zip(edges, fact_embeddings, strict=True):
+        edge.fact_embedding = fact_embedding

graphiti_core/embedder/client.py

@@ -32,3 +32,6 @@ class EmbedderClient(ABC):
         self, input_data: str | list[str] | Iterable[int] | Iterable[Iterable[int]]
     ) -> list[float]:
         pass
+
+    async def create_batch(self, input_data_list: list[str]) -> list[list[float]]:
+        raise NotImplementedError()

graphiti_core/embedder/gemini.py

@@ -66,3 +66,13 @@ class GeminiEmbedder(EmbedderClient):
         )
 
         return result.embeddings[0].values
+
+    async def create_batch(self, input_data_list: list[str]) -> list[list[float]]:
+        # Generate embeddings
+        result = await self.client.aio.models.embed_content(
+            model=self.config.embedding_model or DEFAULT_EMBEDDING_MODEL,
+            contents=input_data_list,
+            config=types.EmbedContentConfig(output_dimensionality=self.config.embedding_dim),
+        )
+
+        return [embedding.values for embedding in result.embeddings]

graphiti_core/embedder/openai.py

@@ -58,3 +58,9 @@ class OpenAIEmbedder(EmbedderClient):
             input=input_data, model=self.config.embedding_model
         )
         return result.data[0].embedding[: self.config.embedding_dim]
+
+    async def create_batch(self, input_data_list: list[str]) -> list[list[float]]:
+        result = await self.client.embeddings.create(
+            input=input_data_list, model=self.config.embedding_model
+        )
+        return [embedding.embedding[: self.config.embedding_dim] for embedding in result.data]

graphiti_core/embedder/voyage.py

@@ -56,3 +56,10 @@ class VoyageAIEmbedder(EmbedderClient):
 
         result = await self.client.embed(input_list, model=self.config.embedding_model)
         return [float(x) for x in result.embeddings[0][: self.config.embedding_dim]]
+
+    async def create_batch(self, input_data_list: list[str]) -> list[list[float]]:
+        result = await self.client.embed(input_data_list, model=self.config.embedding_model)
+        return [
+            [float(x) for x in embedding[: self.config.embedding_dim]]
+            for embedding in result.embeddings
+        ]

graphiti_core/graphiti.py

@@ -27,6 +27,7 @@ from graphiti_core.cross_encoder.client import CrossEncoderClient
 from graphiti_core.cross_encoder.openai_reranker_client import OpenAIRerankerClient
 from graphiti_core.edges import EntityEdge, EpisodicEdge
 from graphiti_core.embedder import EmbedderClient, OpenAIEmbedder
+from graphiti_core.graphiti_types import GraphitiClients
 from graphiti_core.helpers import DEFAULT_DATABASE, semaphore_gather
 from graphiti_core.llm_client import LLMClient, OpenAIClient
 from graphiti_core.nodes import CommunityNode, EntityNode, EpisodeType, EpisodicNode
@@ -42,7 +43,6 @@ from graphiti_core.search.search_utils import (
     RELEVANT_SCHEMA_LIMIT,
     get_mentioned_nodes,
     get_relevant_edges,
-    get_relevant_nodes,
 )
 from graphiti_core.utils.bulk_utils import (
     RawEpisode,
@@ -72,7 +72,11 @@ from graphiti_core.utils.maintenance.graph_data_operations import (
     build_indices_and_constraints,
     retrieve_episodes,
 )
-from graphiti_core.utils.maintenance.node_operations import extract_nodes, resolve_extracted_nodes
+from graphiti_core.utils.maintenance.node_operations import (
+    extract_attributes_from_nodes,
+    extract_nodes,
+    resolve_extracted_nodes,
+)
 from graphiti_core.utils.maintenance.temporal_operations import get_edge_contradictions
 from graphiti_core.utils.ontology_utils.entity_types_utils import validate_entity_types
 
@@ -150,6 +154,13 @@ class Graphiti:
         else:
             self.cross_encoder = OpenAIRerankerClient()
 
+        self.clients = GraphitiClients(
+            driver=self.driver,
+            llm_client=self.llm_client,
+            embedder=self.embedder,
+            cross_encoder=self.cross_encoder,
+        )
+
     async def close(self):
         """
         Close the connection to the Neo4j database.
@@ -222,6 +233,7 @@ class Graphiti:
         reference_time: datetime,
         last_n: int = EPISODE_WINDOW_LEN,
         group_ids: list[str] | None = None,
+        source: EpisodeType | None = None,
     ) -> list[EpisodicNode]:
         """
         Retrieve the last n episodic nodes from the graph.
@@ -248,7 +260,7 @@ class Graphiti:
         The actual retrieval is performed by the `retrieve_episodes` function
         from the `graphiti_core.utils` module.
         """
-        return await retrieve_episodes(self.driver, reference_time, last_n, group_ids)
+        return await retrieve_episodes(self.driver, reference_time, last_n, group_ids, source)
 
     async def add_episode(
         self,
@@ -314,15 +326,16 @@ class Graphiti:
         """
         try:
             start = time()
-
-            entity_edges: list[EntityEdge] = []
             now = utc_now()
 
             validate_entity_types(entity_types)
 
             previous_episodes = (
                 await self.retrieve_episodes(
-                    reference_time, last_n=RELEVANT_SCHEMA_LIMIT, group_ids=[group_id]
+                    reference_time,
+                    last_n=RELEVANT_SCHEMA_LIMIT,
+                    group_ids=[group_id],
+                    source=source,
                 )
                 if previous_episode_uuids is None
                 else await EpisodicNode.get_by_uuids(self.driver, previous_episode_uuids)
@@ -346,132 +359,36 @@ class Graphiti:
             # Extract entities as nodes
 
             extracted_nodes = await extract_nodes(
-                self.llm_client, episode, previous_episodes, entity_types
+                self.clients, episode, previous_episodes, entity_types
             )
-            logger.debug(f'Extracted nodes: {[(n.name, n.uuid) for n in extracted_nodes]}')
-
-            # Calculate Embeddings
 
-            await semaphore_gather(
-                *[node.generate_name_embedding(self.embedder) for node in extracted_nodes]
-            )
-
-            # Find relevant nodes already in the graph
-            existing_nodes_lists: list[list[EntityNode]] = list(
-                await semaphore_gather(
-                    *[
-                        get_relevant_nodes(self.driver, SearchFilters(), [node])
-                        for node in extracted_nodes
-                    ]
-                )
-            )
-
-            # Resolve extracted nodes with nodes already in the graph and extract facts
-            logger.debug(f'Extracted nodes: {[(n.name, n.uuid) for n in extracted_nodes]}')
-
-            (mentioned_nodes, uuid_map), extracted_edges = await semaphore_gather(
+            # Extract edges and resolve nodes
+            (nodes, uuid_map), extracted_edges = await semaphore_gather(
                 resolve_extracted_nodes(
-                    self.llm_client,
+                    self.clients,
                     extracted_nodes,
-                    existing_nodes_lists,
                     episode,
                     previous_episodes,
                     entity_types,
                 ),
-                extract_edges(
-                    self.llm_client, episode, extracted_nodes, previous_episodes, group_id
-                ),
-            )
-            logger.debug(f'Adjusted mentioned nodes: {[(n.name, n.uuid) for n in mentioned_nodes]}')
-            nodes = mentioned_nodes
-
-            extracted_edges_with_resolved_pointers = resolve_edge_pointers(
-                extracted_edges, uuid_map
+                extract_edges(self.clients, episode, extracted_nodes, previous_episodes, group_id),
             )
 
-            # calculate embeddings
-            await semaphore_gather(
-                *[
-                    edge.generate_embedding(self.embedder)
-                    for edge in extracted_edges_with_resolved_pointers
-                ]
-            )
+            edges = resolve_edge_pointers(extracted_edges, uuid_map)
 
-            # Resolve extracted edges with related edges already in the graph
-            related_edges_list: list[list[EntityEdge]] = list(
-                await semaphore_gather(
-                    *[
-                        get_relevant_edges(
-                            self.driver,
-                            [edge],
-                            edge.source_node_uuid,
-                            edge.target_node_uuid,
-                            RELEVANT_SCHEMA_LIMIT,
-                        )
-                        for edge in extracted_edges_with_resolved_pointers
-                    ]
-                )
-            )
-            logger.debug(
-                f'Related edges lists: {[(e.name, e.uuid) for edges_lst in related_edges_list for e in edges_lst]}'
-            )
-            logger.debug(
-                f'Extracted edges: {[(e.name, e.uuid) for e in extracted_edges_with_resolved_pointers]}'
-            )
-
-            existing_source_edges_list: list[list[EntityEdge]] = list(
-                await semaphore_gather(
-                    *[
-                        get_relevant_edges(
-                            self.driver,
-                            [edge],
-                            edge.source_node_uuid,
-                            None,
-                            RELEVANT_SCHEMA_LIMIT,
-                        )
-                        for edge in extracted_edges_with_resolved_pointers
-                    ]
-                )
-            )
-
-            existing_target_edges_list: list[list[EntityEdge]] = list(
-                await semaphore_gather(
-                    *[
-                        get_relevant_edges(
-                            self.driver,
-                            [edge],
-                            None,
-                            edge.target_node_uuid,
-                            RELEVANT_SCHEMA_LIMIT,
-                        )
-                        for edge in extracted_edges_with_resolved_pointers
-                    ]
-                )
-            )
-
-            existing_edges_list: list[list[EntityEdge]] = [
-                source_lst + target_lst
-                for source_lst, target_lst in zip(
-                    existing_source_edges_list, existing_target_edges_list, strict=False
-                )
-            ]
-
-            resolved_edges, invalidated_edges = await resolve_extracted_edges(
-                self.llm_client,
-                extracted_edges_with_resolved_pointers,
-                related_edges_list,
-                existing_edges_list,
-                episode,
-                previous_episodes,
+            (resolved_edges, invalidated_edges), hydrated_nodes = await semaphore_gather(
+                resolve_extracted_edges(
+                    self.clients,
+                    edges,
+                ),
+                extract_attributes_from_nodes(
+                    self.clients, nodes, episode, previous_episodes, entity_types
+                ),
             )
 
-            entity_edges.extend(resolved_edges + invalidated_edges)
+            entity_edges = resolved_edges + invalidated_edges
 
-            logger.debug(f'Resolved edges: {[(e.name, e.uuid) for e in resolved_edges]}')
-
-            episodic_edges: list[EpisodicEdge] = build_episodic_edges(mentioned_nodes, episode, now)
-
-            logger.debug(f'Built episodic edges: {episodic_edges}')
+            episodic_edges = build_episodic_edges(nodes, episode, now)
 
             episode.entity_edges = [edge.uuid for edge in entity_edges]
 
@@ -565,7 +482,7 @@ class Graphiti:
                 extracted_nodes,
                 extracted_edges,
                 episodic_edges,
-            ) = await extract_nodes_and_edges_bulk(self.llm_client, episode_pairs)
+            ) = await extract_nodes_and_edges_bulk(self.clients, episode_pairs)
 
             # Generate embeddings
             await semaphore_gather(
@@ -684,9 +601,7 @@ class Graphiti:
 
         edges = (
             await search(
-                self.driver,
-                self.embedder,
-                self.cross_encoder,
+                self.clients,
                 query,
                 group_ids,
                 search_config,
@@ -728,9 +643,7 @@ class Graphiti:
         """
 
         return await search(
-            self.driver,
-            self.embedder,
-            self.cross_encoder,
+            self.clients,
             query,
             group_ids,
             config,
@@ -761,26 +674,17 @@ class Graphiti:
             await edge.generate_embedding(self.embedder)
 
         resolved_nodes, uuid_map = await resolve_extracted_nodes(
-            self.llm_client,
+            self.clients,
             [source_node, target_node],
-            [
-                await get_relevant_nodes(self.driver, SearchFilters(), [source_node]),
-                await get_relevant_nodes(self.driver, SearchFilters(), [target_node]),
-            ],
         )
 
         updated_edge = resolve_edge_pointers([edge], uuid_map)[0]
 
-        related_edges = await get_relevant_edges(
-            self.driver,
-            [updated_edge],
-            source_node_uuid=resolved_nodes[0].uuid,
-            target_node_uuid=resolved_nodes[1].uuid,
-        )
+        related_edges = await get_relevant_edges(self.driver, [updated_edge], SearchFilters(), 0.8)
 
-        resolved_edge = await dedupe_extracted_edge(self.llm_client, updated_edge, related_edges)
+        resolved_edge = await dedupe_extracted_edge(self.llm_client, updated_edge, related_edges[0])
 
-        contradicting_edges = await get_edge_contradictions(self.llm_client, edge, related_edges)
+        contradicting_edges = await get_edge_contradictions(self.llm_client, edge, related_edges[0])
         invalidated_edges = resolve_edge_contradictions(resolved_edge, contradicting_edges)
 
         await add_nodes_and_edges_bulk(

graphiti_core/graphiti_types.py

@@ -0,0 +1,31 @@
+"""
+Copyright 2024, Zep Software, Inc.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+"""
+
+from neo4j import AsyncDriver
+from pydantic import BaseModel, ConfigDict
+
+from graphiti_core.cross_encoder import CrossEncoderClient
+from graphiti_core.embedder import EmbedderClient
+from graphiti_core.llm_client import LLMClient
+
+
+class GraphitiClients(BaseModel):
+    driver: AsyncDriver
+    llm_client: LLMClient
+    embedder: EmbedderClient
+    cross_encoder: CrossEncoderClient
+
+    model_config = ConfigDict(arbitrary_types_allowed=True)

graphiti_core/helpers.py

@@ -22,15 +22,20 @@ from datetime import datetime
 import numpy as np
 from dotenv import load_dotenv
 from neo4j import time as neo4j_time
+from typing_extensions import LiteralString
 
 load_dotenv()
 
 DEFAULT_DATABASE = os.getenv('DEFAULT_DATABASE', None)
 USE_PARALLEL_RUNTIME = bool(os.getenv('USE_PARALLEL_RUNTIME', False))
 SEMAPHORE_LIMIT = int(os.getenv('SEMAPHORE_LIMIT', 20))
-MAX_REFLEXION_ITERATIONS = int(os.getenv('MAX_REFLEXION_ITERATIONS', 2))
+MAX_REFLEXION_ITERATIONS = int(os.getenv('MAX_REFLEXION_ITERATIONS', 0))
 DEFAULT_PAGE_LIMIT = 20
 
+RUNTIME_QUERY: LiteralString = (
+    'CYPHER runtime = parallel parallelRuntimeSupport=all\n' if USE_PARALLEL_RUNTIME else ''
+)
+
 
 def parse_db_date(neo_date: neo4j_time.DateTime | None) -> datetime | None:
     return neo_date.to_native() if neo_date else None

graphiti_core/llm_client/anthropic_client.py

@@ -262,7 +262,7 @@ class AnthropicClient(LLMClient):
         self,
         messages: list[Message],
         response_model: type[BaseModel] | None = None,
-        max_tokens: int = DEFAULT_MAX_TOKENS,
+        max_tokens: int | None = None,
     ) -> dict[str, typing.Any]:
         """
         Generate a response from the LLM.
@@ -280,6 +280,9 @@ class AnthropicClient(LLMClient):
             RefusalError: If the LLM refuses to respond.
             Exception: If an error occurs during the generation process.
         """
+        if max_tokens is None:
+            max_tokens = self.max_tokens
+
         retry_count = 0
         max_retries = 2
         last_error: Exception | None = None

graphiti_core/llm_client/client.py

@@ -127,8 +127,11 @@ class LLMClient(ABC):
         self,
         messages: list[Message],
         response_model: type[BaseModel] | None = None,
-        max_tokens: int = DEFAULT_MAX_TOKENS,
+        max_tokens: int | None = None,
     ) -> dict[str, typing.Any]:
+        if max_tokens is None:
+            max_tokens = self.max_tokens
+
         if response_model is not None:
             serialized_model = json.dumps(response_model.model_json_schema())
             messages[

graphiti_core/llm_client/gemini_client.py

@@ -166,7 +166,7 @@ class GeminiClient(LLMClient):
         self,
         messages: list[Message],
         response_model: type[BaseModel] | None = None,
-        max_tokens: int = DEFAULT_MAX_TOKENS,
+        max_tokens: int | None = None,
     ) -> dict[str, typing.Any]:
         """
         Generate a response from the Gemini language model.
@@ -180,6 +180,9 @@ class GeminiClient(LLMClient):
         Returns:
             dict[str, typing.Any]: The response from the language model.
         """
+        if max_tokens is None:
+            max_tokens = self.max_tokens
+
         # Call the internal _generate_response method
         return await self._generate_response(
             messages=messages, response_model=response_model, max_tokens=max_tokens

graphiti_core/llm_client/openai_client.py

@@ -131,8 +131,11 @@ class OpenAIClient(LLMClient):
         self,
         messages: list[Message],
         response_model: type[BaseModel] | None = None,
-        max_tokens: int = DEFAULT_MAX_TOKENS,
+        max_tokens: int | None = None,
     ) -> dict[str, typing.Any]:
+        if max_tokens is None:
+            max_tokens = self.max_tokens
+
         retry_count = 0
         last_error = None
 

graphiti_core/llm_client/openai_generic_client.py

@@ -117,8 +117,11 @@ class OpenAIGenericClient(LLMClient):
         self,
         messages: list[Message],
         response_model: type[BaseModel] | None = None,
-        max_tokens: int = DEFAULT_MAX_TOKENS,
+        max_tokens: int | None = None,
     ) -> dict[str, typing.Any]:
+        if max_tokens is None:
+            max_tokens = self.max_tokens
+
         retry_count = 0
         last_error = None
 

graphiti_core/models/edges/edge_db_queries.py

@@ -47,7 +47,7 @@ ENTITY_EDGE_SAVE_BULK = """
     SET r = {uuid: edge.uuid, name: edge.name, group_id: edge.group_id, fact: edge.fact, episodes: edge.episodes, 
     created_at: edge.created_at, expired_at: edge.expired_at, valid_at: edge.valid_at, invalid_at: edge.invalid_at}
     WITH r, edge CALL db.create.setRelationshipVectorProperty(r, "fact_embedding", edge.fact_embedding)
-    RETURN r.uuid AS uuid
+    RETURN edge.uuid AS uuid
 """
 
 COMMUNITY_EDGE_SAVE = """

graphiti_core/nodes.py

@@ -39,8 +39,6 @@ from graphiti_core.utils.datetime_utils import utc_now
 logger = logging.getLogger(__name__)
 
 ENTITY_NODE_RETURN: LiteralString = """
-        OPTIONAL MATCH (e:Episodic)-[r:MENTIONS]->(n)
-        WITH n, collect(e.uuid) AS episodes
         RETURN
             n.uuid As uuid, 
             n.name AS name,
@@ -49,8 +47,8 @@ ENTITY_NODE_RETURN: LiteralString = """
             n.created_at AS created_at, 
             n.summary AS summary,
             labels(n) AS labels,
-            properties(n) AS attributes,
-            episodes"""
+            properties(n) AS attributes
+            """
 
 
 class EpisodeType(Enum):
@@ -294,9 +292,6 @@ class EpisodicNode(Node):
 class EntityNode(Node):
     name_embedding: list[float] | None = Field(default=None, description='embedding of the name')
     summary: str = Field(description='regional summary of surrounding edges', default_factory=str)
-    episodes: list[str] | None = Field(
-        default=None, description='List of episode uuids that mention this node.'
-    )
     attributes: dict[str, Any] = Field(
         default={}, description='Additional attributes of the node. Dependent on node labels'
     )
@@ -337,8 +332,8 @@ class EntityNode(Node):
     async def get_by_uuid(cls, driver: AsyncDriver, uuid: str):
         query = (
             """
-                MATCH (n:Entity {uuid: $uuid})
-                """
+                                    MATCH (n:Entity {uuid: $uuid})
+                                    """
             + ENTITY_NODE_RETURN
         )
         records, _, _ = await driver.execute_query(
@@ -544,7 +539,6 @@ def get_entity_node_from_record(record: Any) -> EntityNode:
         created_at=record['created_at'].to_native(),
         summary=record['summary'],
         attributes=record['attributes'],
-        episodes=record['episodes'],
     )
 
     entity_node.attributes.pop('uuid', None)
@@ -566,3 +560,9 @@ def get_community_node_from_record(record: Any) -> CommunityNode:
         created_at=record['created_at'].to_native(),
         summary=record['summary'],
     )
+
+
+async def create_entity_node_embeddings(embedder: EmbedderClient, nodes: list[EntityNode]):
+    name_embeddings = await embedder.create_batch([node.name for node in nodes])
+    for node, name_embedding in zip(nodes, name_embeddings, strict=True):
+        node.name_embedding = name_embedding