graphiti-core 0.21.0rc12__py3-none-any.whl → 0.22.0rc0__py3-none-any.whl

graphiti_core/driver/driver.py

@@ -14,29 +14,16 @@ See the License for the specific language governing permissions and
 limitations under the License.
 """
 
-import asyncio
 import copy
 import logging
 import os
 from abc import ABC, abstractmethod
 from collections.abc import Coroutine
-from datetime import datetime
 from enum import Enum
 from typing import Any
 
 from dotenv import load_dotenv
 
-from graphiti_core.embedder.client import EMBEDDING_DIM
-
-try:
-    from opensearchpy import AsyncOpenSearch, helpers
-
-    _HAS_OPENSEARCH = True
-except ImportError:
-    OpenSearch = None
-    helpers = None
-    _HAS_OPENSEARCH = False
-
 logger = logging.getLogger(__name__)
 
 DEFAULT_SIZE = 10
@@ -56,91 +43,6 @@ class GraphProvider(Enum):
     NEPTUNE = 'neptune'
 
 
-aoss_indices = [
-    {
-        'index_name': ENTITY_INDEX_NAME,
-        'body': {
-            'settings': {'index': {'knn': True}},
-            'mappings': {
-                'properties': {
-                    'uuid': {'type': 'keyword'},
-                    'name': {'type': 'text'},
-                    'summary': {'type': 'text'},
-                    'group_id': {'type': 'keyword'},
-                    'created_at': {'type': 'date', 'format': 'strict_date_optional_time_nanos'},
-                    'name_embedding': {
-                        'type': 'knn_vector',
-                        'dimension': EMBEDDING_DIM,
-                        'method': {
-                            'engine': 'faiss',
-                            'space_type': 'cosinesimil',
-                            'name': 'hnsw',
-                            'parameters': {'ef_construction': 128, 'm': 16},
-                        },
-                    },
-                }
-            },
-        },
-    },
-    {
-        'index_name': COMMUNITY_INDEX_NAME,
-        'body': {
-            'mappings': {
-                'properties': {
-                    'uuid': {'type': 'keyword'},
-                    'name': {'type': 'text'},
-                    'group_id': {'type': 'keyword'},
-                }
-            }
-        },
-    },
-    {
-        'index_name': EPISODE_INDEX_NAME,
-        'body': {
-            'mappings': {
-                'properties': {
-                    'uuid': {'type': 'keyword'},
-                    'content': {'type': 'text'},
-                    'source': {'type': 'text'},
-                    'source_description': {'type': 'text'},
-                    'group_id': {'type': 'keyword'},
-                    'created_at': {'type': 'date', 'format': 'strict_date_optional_time_nanos'},
-                    'valid_at': {'type': 'date', 'format': 'strict_date_optional_time_nanos'},
-                }
-            }
-        },
-    },
-    {
-        'index_name': ENTITY_EDGE_INDEX_NAME,
-        'body': {
-            'settings': {'index': {'knn': True}},
-            'mappings': {
-                'properties': {
-                    'uuid': {'type': 'keyword'},
-                    'name': {'type': 'text'},
-                    'fact': {'type': 'text'},
-                    'group_id': {'type': 'keyword'},
-                    'created_at': {'type': 'date', 'format': 'strict_date_optional_time_nanos'},
-                    'valid_at': {'type': 'date', 'format': 'strict_date_optional_time_nanos'},
-                    'expired_at': {'type': 'date', 'format': 'strict_date_optional_time_nanos'},
-                    'invalid_at': {'type': 'date', 'format': 'strict_date_optional_time_nanos'},
-                    'fact_embedding': {
-                        'type': 'knn_vector',
-                        'dimension': EMBEDDING_DIM,
-                        'method': {
-                            'engine': 'faiss',
-                            'space_type': 'cosinesimil',
-                            'name': 'hnsw',
-                            'parameters': {'ef_construction': 128, 'm': 16},
-                        },
-                    },
-                }
-            },
-        },
-    },
-]
-
-
 class GraphDriverSession(ABC):
     provider: GraphProvider
 
@@ -171,7 +73,7 @@ class GraphDriver(ABC):
         ''  # Neo4j (default) syntax does not require a prefix for fulltext queries
     )
     _database: str
-    aoss_client: AsyncOpenSearch | None  # type: ignore
+    aoss_client: Any  # type: ignore
 
     @abstractmethod
     def execute_query(self, cypher_query_: str, **kwargs: Any) -> Coroutine:
@@ -199,119 +101,6 @@ class GraphDriver(ABC):
 
         return cloned
 
-    async def delete_all_indexes_impl(self) -> Coroutine[Any, Any, Any]:
-        # No matter what happens above, always return True
-        return self.delete_aoss_indices()
-
-    async def create_aoss_indices(self):
-        client = self.aoss_client
-        if not client:
-            logger.warning('No OpenSearch client found')
-            return
-
-        for index in aoss_indices:
-            alias_name = index['index_name']
-
-            # If alias already exists, skip (idempotent behavior)
-            if await client.indices.exists_alias(name=alias_name):
-                continue
-
-            # Build a physical index name with timestamp
-            ts_suffix = datetime.utcnow().strftime('%Y%m%d%H%M%S')
-            physical_index_name = f'{alias_name}_{ts_suffix}'
-
-            # Create the index
-            await client.indices.create(index=physical_index_name, body=index['body'])
-
-            # Point alias to it
-            await client.indices.put_alias(index=physical_index_name, name=alias_name)
-
-        # Allow some time for index creation
-        await asyncio.sleep(1)
-
-    async def delete_aoss_indices(self):
-        client = self.aoss_client
-
-        if not client:
-            logger.warning('No OpenSearch client found')
-            return
-
-        for entry in aoss_indices:
-            alias_name = entry['index_name']
-
-            try:
-                # Resolve alias → indices
-                alias_info = await client.indices.get_alias(name=alias_name)
-                indices = list(alias_info.keys())
-
-                if not indices:
-                    logger.info(f"No indices found for alias '{alias_name}'")
-                    continue
-
-                for index in indices:
-                    if await client.indices.exists(index=index):
-                        await client.indices.delete(index=index)
-                        logger.info(f"Deleted index '{index}' (alias: {alias_name})")
-                    else:
-                        logger.warning(f"Index '{index}' not found for alias '{alias_name}'")
-
-            except Exception as e:
-                logger.error(f"Error deleting indices for alias '{alias_name}': {e}")
-
-    async def clear_aoss_indices(self):
-        client = self.aoss_client
-
-        if not client:
-            logger.warning('No OpenSearch client found')
-            return
-
-        for index in aoss_indices:
-            index_name = index['index_name']
-
-            if await client.indices.exists(index=index_name):
-                try:
-                    # Delete all documents but keep the index
-                    response = await client.delete_by_query(
-                        index=index_name,
-                        body={'query': {'match_all': {}}},
-                    )
-                    logger.info(f"Cleared index '{index_name}': {response}")
-                except Exception as e:
-                    logger.error(f"Error clearing index '{index_name}': {e}")
-            else:
-                logger.warning(f"Index '{index_name}' does not exist")
-
-    async def save_to_aoss(self, name: str, data: list[dict]) -> int:
-        client = self.aoss_client
-        if not client or not helpers:
-            logger.warning('No OpenSearch client found')
-            return 0
-
-        for index in aoss_indices:
-            if name.lower() == index['index_name']:
-                to_index = []
-                for d in data:
-                    doc = {}
-                    for p in index['body']['mappings']['properties']:
-                        if p in d:  # protect against missing fields
-                            doc[p] = d[p]
-
-                    item = {
-                        '_index': name,
-                        '_id': d['uuid'],
-                        '_routing': d.get('group_id'),
-                        '_source': doc,
-                    }
-                    to_index.append(item)
-
-                success, failed = await helpers.async_bulk(
-                    client, to_index, stats_only=True, request_timeout=60
-                )
-
-                return success if failed == 0 else success
-
-        return 0
-
     def build_fulltext_query(
         self, query: str, group_ids: list[str] | None = None, max_query_length: int = 128
     ) -> str:
@@ -320,3 +109,9 @@ class GraphDriver(ABC):
         Only implemented by providers that need custom fulltext query building.
         """
         raise NotImplementedError(f'build_fulltext_query not implemented for {self.provider}')
+
+    async def save_to_aoss(self, name: str, data: list[dict]) -> int:
+        return 0
+
+    async def clear_aoss_indices(self):
+        return 1

graphiti_core/driver/neo4j_driver.py

@@ -22,28 +22,9 @@ from neo4j import AsyncGraphDatabase, EagerResult
 from typing_extensions import LiteralString
 
 from graphiti_core.driver.driver import GraphDriver, GraphDriverSession, GraphProvider
-from graphiti_core.helpers import semaphore_gather
 
 logger = logging.getLogger(__name__)
 
-try:
-    import boto3
-    from opensearchpy import (
-        AIOHttpConnection,
-        AsyncOpenSearch,
-        AWSV4SignerAuth,
-        Urllib3AWSV4SignerAuth,
-        Urllib3HttpConnection,
-    )
-
-    _HAS_OPENSEARCH = True
-except ImportError:
-    boto3 = None
-    OpenSearch = None
-    Urllib3AWSV4SignerAuth = None
-    Urllib3HttpConnection = None
-    _HAS_OPENSEARCH = False
-
 
 class Neo4jDriver(GraphDriver):
     provider = GraphProvider.NEO4J
@@ -54,11 +35,6 @@ class Neo4jDriver(GraphDriver):
         user: str | None,
         password: str | None,
         database: str = 'neo4j',
-        aoss_host: str | None = None,
-        aoss_port: int | None = None,
-        aws_profile_name: str | None = None,
-        aws_region: str | None = None,
-        aws_service: str | None = None,
     ):
         super().__init__()
         self.client = AsyncGraphDatabase.driver(
@@ -68,24 +44,6 @@ class Neo4jDriver(GraphDriver):
         self._database = database
 
         self.aoss_client = None
-        if aoss_host and aoss_port and boto3 is not None:
-            try:
-                region = aws_region
-                service = aws_service
-                credentials = boto3.Session(profile_name=aws_profile_name).get_credentials()
-                auth = AWSV4SignerAuth(credentials, region or '', service or '')
-
-                self.aoss_client = AsyncOpenSearch(
-                    hosts=[{'host': aoss_host, 'port': aoss_port}],
-                    auth=auth,
-                    use_ssl=True,
-                    verify_certs=True,
-                    connection_class=AIOHttpConnection,
-                    pool_maxsize=20,
-                )  # type: ignore
-            except Exception as e:
-                logger.warning(f'Failed to initialize OpenSearch client: {e}')
-                self.aoss_client = None
 
     async def execute_query(self, cypher_query_: LiteralString, **kwargs: Any) -> EagerResult:
         # Check if database_ is provided in kwargs.
@@ -111,13 +69,6 @@ class Neo4jDriver(GraphDriver):
         return await self.client.close()
 
     def delete_all_indexes(self) -> Coroutine:
-        if self.aoss_client:
-            return semaphore_gather(
-                self.client.execute_query(
-                    'CALL db.indexes() YIELD name DROP INDEX name',
-                ),
-                self.delete_aoss_indices(),
-            )
         return self.client.execute_query(
             'CALL db.indexes() YIELD name DROP INDEX name',
         )

graphiti_core/driver/neptune_driver.py

@@ -22,21 +22,16 @@ from typing import Any
 
 import boto3
 from langchain_aws.graphs import NeptuneAnalyticsGraph, NeptuneGraph
-from opensearchpy import OpenSearch, Urllib3AWSV4SignerAuth, Urllib3HttpConnection
+from opensearchpy import OpenSearch, Urllib3AWSV4SignerAuth, Urllib3HttpConnection, helpers
 
-from graphiti_core.driver.driver import (
-    DEFAULT_SIZE,
-    GraphDriver,
-    GraphDriverSession,
-    GraphProvider,
-)
+from graphiti_core.driver.driver import GraphDriver, GraphDriverSession, GraphProvider
 
 logger = logging.getLogger(__name__)
+DEFAULT_SIZE = 10
 
-neptune_aoss_indices = [
+aoss_indices = [
     {
         'index_name': 'node_name_and_summary',
-        'alias_name': 'entities',
         'body': {
             'mappings': {
                 'properties': {
@@ -54,7 +49,6 @@ neptune_aoss_indices = [
     },
     {
         'index_name': 'community_name',
-        'alias_name': 'communities',
         'body': {
             'mappings': {
                 'properties': {
@@ -71,7 +65,6 @@ neptune_aoss_indices = [
     },
     {
         'index_name': 'episode_content',
-        'alias_name': 'episodes',
         'body': {
             'mappings': {
                 'properties': {
@@ -95,7 +88,6 @@ neptune_aoss_indices = [
     },
     {
         'index_name': 'edge_name_and_fact',
-        'alias_name': 'facts',
         'body': {
             'mappings': {
                 'properties': {
@@ -228,27 +220,52 @@ class NeptuneDriver(GraphDriver):
     async def _delete_all_data(self) -> Any:
         return await self.execute_query('MATCH (n) DETACH DELETE n')
 
+    def delete_all_indexes(self) -> Coroutine[Any, Any, Any]:
+        return self.delete_all_indexes_impl()
+
+    async def delete_all_indexes_impl(self) -> Coroutine[Any, Any, Any]:
+        # No matter what happens above, always return True
+        return self.delete_aoss_indices()
+
     async def create_aoss_indices(self):
-        for index in neptune_aoss_indices:
+        for index in aoss_indices:
             index_name = index['index_name']
             client = self.aoss_client
-            if not client:
-                raise ValueError(
-                    'You must provide an AOSS endpoint to create an OpenSearch driver.'
-                )
             if not client.indices.exists(index=index_name):
-                await client.indices.create(index=index_name, body=index['body'])
-
-            alias_name = index.get('alias_name', index_name)
-
-            if not client.indices.exists_alias(name=alias_name, index=index_name):
-                await client.indices.put_alias(index=index_name, name=alias_name)
-
+                client.indices.create(index=index_name, body=index['body'])
         # Sleep for 1 minute to let the index creation complete
         await asyncio.sleep(60)
 
-    def delete_all_indexes(self) -> Coroutine[Any, Any, Any]:
-        return self.delete_all_indexes_impl()
+    async def delete_aoss_indices(self):
+        for index in aoss_indices:
+            index_name = index['index_name']
+            client = self.aoss_client
+            if client.indices.exists(index=index_name):
+                client.indices.delete(index=index_name)
+
+    def run_aoss_query(self, name: str, query_text: str, limit: int = 10) -> dict[str, Any]:
+        for index in aoss_indices:
+            if name.lower() == index['index_name']:
+                index['query']['query']['multi_match']['query'] = query_text
+                query = {'size': limit, 'query': index['query']}
+                resp = self.aoss_client.search(body=query['query'], index=index['index_name'])
+                return resp
+        return {}
+
+    def save_to_aoss(self, name: str, data: list[dict]) -> int:
+        for index in aoss_indices:
+            if name.lower() == index['index_name']:
+                to_index = []
+                for d in data:
+                    item = {'_index': name, '_id': d['uuid']}
+                    for p in index['body']['mappings']['properties']:
+                        if p in d:
+                            item[p] = d[p]
+                    to_index.append(item)
+                success, failed = helpers.bulk(self.aoss_client, to_index, stats_only=True)
+                return success
+
+        return 0
 
 
 class NeptuneDriverSession(GraphDriverSession):

graphiti_core/llm_client/client.py

@@ -33,12 +33,16 @@ DEFAULT_TEMPERATURE = 0
 DEFAULT_CACHE_DIR = './llm_cache'
 
 
-def get_extraction_language_instruction() -> str:
+def get_extraction_language_instruction(group_id: str | None = None) -> str:
     """Returns instruction for language extraction behavior.
 
     Override this function to customize language extraction:
     - Return empty string to disable multilingual instructions
     - Return custom instructions for specific language requirements
+    - Use group_id to provide different instructions per group/partition
+
+    Args:
+        group_id: Optional partition identifier for the graph
 
     Returns:
         str: Language instruction to append to system messages
@@ -142,6 +146,7 @@ class LLMClient(ABC):
         response_model: type[BaseModel] | None = None,
         max_tokens: int | None = None,
         model_size: ModelSize = ModelSize.medium,
+        group_id: str | None = None,
     ) -> dict[str, typing.Any]:
         if max_tokens is None:
             max_tokens = self.max_tokens
@@ -155,7 +160,7 @@ class LLMClient(ABC):
             )
 
         # Add multilingual extraction instructions
-        messages[0].content += get_extraction_language_instruction()
+        messages[0].content += get_extraction_language_instruction(group_id)
 
         if self.cache_enabled and self.cache_dir is not None:
             cache_key = self._get_cache_key(messages)

graphiti_core/llm_client/gemini_client.py

@@ -357,6 +357,7 @@ class GeminiClient(LLMClient):
         response_model: type[BaseModel] | None = None,
         max_tokens: int | None = None,
         model_size: ModelSize = ModelSize.medium,
+        group_id: str | None = None,
     ) -> dict[str, typing.Any]:
         """
         Generate a response from the Gemini language model with retry logic and error handling.
@@ -367,6 +368,7 @@ class GeminiClient(LLMClient):
             response_model (type[BaseModel] | None): An optional Pydantic model to parse the response into.
             max_tokens (int | None): The maximum number of tokens to generate in the response.
             model_size (ModelSize): The size of the model to use (small or medium).
+            group_id (str | None): Optional partition identifier for the graph.
 
         Returns:
             dict[str, typing.Any]: The response from the language model.
@@ -376,7 +378,7 @@ class GeminiClient(LLMClient):
         last_output = None
 
         # Add multilingual extraction instructions
-        messages[0].content += get_extraction_language_instruction()
+        messages[0].content += get_extraction_language_instruction(group_id)
 
         while retry_count < self.MAX_RETRIES:
             try:

graphiti_core/llm_client/openai_base_client.py

@@ -175,6 +175,7 @@ class BaseOpenAIClient(LLMClient):
         response_model: type[BaseModel] | None = None,
         max_tokens: int | None = None,
         model_size: ModelSize = ModelSize.medium,
+        group_id: str | None = None,
     ) -> dict[str, typing.Any]:
         """Generate a response with retry logic and error handling."""
         if max_tokens is None:
@@ -184,7 +185,7 @@ class BaseOpenAIClient(LLMClient):
         last_error = None
 
         # Add multilingual extraction instructions
-        messages[0].content += get_extraction_language_instruction()
+        messages[0].content += get_extraction_language_instruction(group_id)
 
         while retry_count <= self.MAX_RETRIES:
             try:

graphiti_core/llm_client/openai_generic_client.py

@@ -120,6 +120,7 @@ class OpenAIGenericClient(LLMClient):
         response_model: type[BaseModel] | None = None,
         max_tokens: int | None = None,
         model_size: ModelSize = ModelSize.medium,
+        group_id: str | None = None,
     ) -> dict[str, typing.Any]:
         if max_tokens is None:
             max_tokens = self.max_tokens
@@ -136,7 +137,7 @@ class OpenAIGenericClient(LLMClient):
             )
 
         # Add multilingual extraction instructions
-        messages[0].content += get_extraction_language_instruction()
+        messages[0].content += get_extraction_language_instruction(group_id)
 
         while retry_count <= self.MAX_RETRIES:
             try:

graphiti_core/utils/maintenance/edge_operations.py

@@ -139,6 +139,7 @@ async def extract_edges(
             prompt_library.extract_edges.edge(context),
             response_model=ExtractedEdges,
             max_tokens=extract_edges_max_tokens,
+            group_id=group_id,
         )
         edges_data = ExtractedEdges(**llm_response).edges
 
@@ -150,6 +151,7 @@ async def extract_edges(
                 prompt_library.extract_edges.reflexion(context),
                 response_model=MissingFacts,
                 max_tokens=extract_edges_max_tokens,
+                group_id=group_id,
             )
 
             missing_facts = reflexion_response.get('missing_facts', [])
@@ -177,6 +179,10 @@ async def extract_edges(
         valid_at_datetime = None
         invalid_at_datetime = None
 
+        # Filter out empty edges
+        if not edge_data.fact.strip():
+            continue
+
         source_node_idx = edge_data.source_entity_id
         target_node_idx = edge_data.target_entity_id
 

graphiti_core/utils/maintenance/node_operations.py

@@ -64,6 +64,7 @@ async def extract_nodes_reflexion(
     episode: EpisodicNode,
     previous_episodes: list[EpisodicNode],
     node_names: list[str],
+    group_id: str | None = None,
 ) -> list[str]:
     # Prepare context for LLM
     context = {
@@ -73,7 +74,9 @@ async def extract_nodes_reflexion(
     }
 
     llm_response = await llm_client.generate_response(
-        prompt_library.extract_nodes.reflexion(context), MissedEntities
+        prompt_library.extract_nodes.reflexion(context),
+        MissedEntities,
+        group_id=group_id,
     )
     missed_entities = llm_response.get('missed_entities', [])
 
@@ -129,16 +132,19 @@ async def extract_nodes(
             llm_response = await llm_client.generate_response(
                 prompt_library.extract_nodes.extract_message(context),
                 response_model=ExtractedEntities,
+                group_id=episode.group_id,
             )
         elif episode.source == EpisodeType.text:
             llm_response = await llm_client.generate_response(
                 prompt_library.extract_nodes.extract_text(context),
                 response_model=ExtractedEntities,
+                group_id=episode.group_id,
             )
         elif episode.source == EpisodeType.json:
             llm_response = await llm_client.generate_response(
                 prompt_library.extract_nodes.extract_json(context),
                 response_model=ExtractedEntities,
+                group_id=episode.group_id,
             )
 
         response_object = ExtractedEntities(**llm_response)
@@ -152,6 +158,7 @@ async def extract_nodes(
                 episode,
                 previous_episodes,
                 [entity.name for entity in extracted_entities],
+                episode.group_id,
             )
 
             entities_missed = len(missing_entities) != 0
@@ -192,6 +199,7 @@ async def extract_nodes(
         logger.debug(f'Created new node: {new_node.name} (UUID: {new_node.uuid})')
 
     logger.debug(f'Extracted nodes: {[(n.name, n.uuid) for n in extracted_nodes]}')
+
     return extracted_nodes
 
 
@@ -477,63 +485,95 @@ async def extract_attributes_from_node(
     entity_type: type[BaseModel] | None = None,
     should_summarize_node: NodeSummaryFilter | None = None,
 ) -> EntityNode:
-    node_context: dict[str, Any] = {
-        'name': node.name,
-        'summary': node.summary,
-        'entity_types': node.labels,
-        'attributes': node.attributes,
-    }
+    # Extract attributes if entity type is defined and has attributes
+    llm_response = await _extract_entity_attributes(
+        llm_client, node, episode, previous_episodes, entity_type
+    )
 
-    attributes_context: dict[str, Any] = {
-        'node': node_context,
-        'episode_content': episode.content if episode is not None else '',
-        'previous_episodes': (
-            [ep.content for ep in previous_episodes] if previous_episodes is not None else []
-        ),
-    }
+    # Extract summary if needed
+    await _extract_entity_summary(
+        llm_client, node, episode, previous_episodes, should_summarize_node
+    )
+
+    node.attributes.update(llm_response)
+
+    return node
 
-    summary_context: dict[str, Any] = {
-        'node': node_context,
-        'episode_content': episode.content if episode is not None else '',
-        'previous_episodes': (
-            [ep.content for ep in previous_episodes] if previous_episodes is not None else []
-        ),
-    }
 
-    has_entity_attributes: bool = bool(
-        entity_type is not None and len(entity_type.model_fields) != 0
+async def _extract_entity_attributes(
+    llm_client: LLMClient,
+    node: EntityNode,
+    episode: EpisodicNode | None,
+    previous_episodes: list[EpisodicNode] | None,
+    entity_type: type[BaseModel] | None,
+) -> dict[str, Any]:
+    if entity_type is None or len(entity_type.model_fields) == 0:
+        return {}
+
+    attributes_context = _build_episode_context(
+        # should not include summary
+        node_data={
+            'name': node.name,
+            'entity_types': node.labels,
+            'attributes': node.attributes,
+        },
+        episode=episode,
+        previous_episodes=previous_episodes,
     )
 
-    llm_response = (
-        (
-            await llm_client.generate_response(
-                prompt_library.extract_nodes.extract_attributes(attributes_context),
-                response_model=entity_type,
-                model_size=ModelSize.small,
-            )
-        )
-        if has_entity_attributes
-        else {}
+    llm_response = await llm_client.generate_response(
+        prompt_library.extract_nodes.extract_attributes(attributes_context),
+        response_model=entity_type,
+        model_size=ModelSize.small,
+        group_id=node.group_id,
     )
 
-    # Determine if summary should be generated
-    generate_summary = True
-    if should_summarize_node is not None:
-        generate_summary = await should_summarize_node(node)
-
-    # Conditionally generate summary
-    if generate_summary:
-        summary_response = await llm_client.generate_response(
-            prompt_library.extract_nodes.extract_summary(summary_context),
-            response_model=EntitySummary,
-            model_size=ModelSize.small,
-        )
-        node.summary = summary_response.get('summary', '')
+    # validate response
+    entity_type(**llm_response)
 
-    if has_entity_attributes and entity_type is not None:
-        entity_type(**llm_response)
-    node_attributes = {key: value for key, value in llm_response.items()}
+    return llm_response
 
-    node.attributes.update(node_attributes)
 
-    return node
+async def _extract_entity_summary(
+    llm_client: LLMClient,
+    node: EntityNode,
+    episode: EpisodicNode | None,
+    previous_episodes: list[EpisodicNode] | None,
+    should_summarize_node: NodeSummaryFilter | None,
+) -> None:
+    if should_summarize_node is not None and not await should_summarize_node(node):
+        return
+
+    summary_context = _build_episode_context(
+        node_data={
+            'name': node.name,
+            'summary': node.summary,
+            'entity_types': node.labels,
+            'attributes': node.attributes,
+        },
+        episode=episode,
+        previous_episodes=previous_episodes,
+    )
+
+    summary_response = await llm_client.generate_response(
+        prompt_library.extract_nodes.extract_summary(summary_context),
+        response_model=EntitySummary,
+        model_size=ModelSize.small,
+        group_id=node.group_id,
+    )
+
+    node.summary = summary_response.get('summary', '')
+
+
+def _build_episode_context(
+    node_data: dict[str, Any],
+    episode: EpisodicNode | None,
+    previous_episodes: list[EpisodicNode] | None,
+) -> dict[str, Any]:
+    return {
+        'node': node_data,
+        'episode_content': episode.content if episode is not None else '',
+        'previous_episodes': (
+            [ep.content for ep in previous_episodes] if previous_episodes is not None else []
+        ),
+    }

{graphiti_core-0.21.0rc12.dist-info → graphiti_core-0.22.0rc0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: graphiti-core
-Version: 0.21.0rc12
+Version: 0.22.0rc0
 Summary: A temporal graph building library
 Project-URL: Homepage, https://help.getzep.com/graphiti/graphiti/overview
 Project-URL: Repository, https://github.com/getzep/graphiti

{graphiti_core-0.21.0rc12.dist-info → graphiti_core-0.22.0rc0.dist-info}/RECORD

@@ -13,11 +13,11 @@ graphiti_core/cross_encoder/client.py,sha256=KLsbfWKOEaAV3adFe3XZlAeb-gje9_sVKCV
 graphiti_core/cross_encoder/gemini_reranker_client.py,sha256=hmITG5YIib52nrKvINwRi4xTfAO1U4jCCaEVIwImHw0,6208
 graphiti_core/cross_encoder/openai_reranker_client.py,sha256=WHMl6Q6gEslR2EzjwpFSZt2Kh6bnu8alkLvzmi0MDtg,4674
 graphiti_core/driver/__init__.py,sha256=kCWimqQU19airu5gKwCmZtZuXkDfaQfKSUhMDoL-rTA,626
-graphiti_core/driver/driver.py,sha256=EO9Aj5O2vpH7iyvQQcE5uJGQ8eA-_i6f8NwfAlW8r74,10831
+graphiti_core/driver/driver.py,sha256=sF6CkGLNPIvUgrmWkVws7TvQCskRHiQKJze4Y4ibMmI,3357
 graphiti_core/driver/falkordb_driver.py,sha256=Q-dImfK4O2bkikqFzo0Wg2g7iFFRSuzy_c6u82tX6-M,9361
 graphiti_core/driver/kuzu_driver.py,sha256=RcWu8E0CCdofrFe34NmCeqfuhaZr_7ZN5jqDkI3VQMI,5453
-graphiti_core/driver/neo4j_driver.py,sha256=E93PdOZaH7wzEbIfoiDSYht49jr6zSzvMMyo1INGEOw,4096
-graphiti_core/driver/neptune_driver.py,sha256=akNLHhFHPEeQu-xO3PM51RomklntT6k5eA2CQ4AFbCc,10311
+graphiti_core/driver/neo4j_driver.py,sha256=xiMUvGpW-XFM_2ab5nJJTHoi_LM7CvVZVq6ZO0BbNwc,2380
+graphiti_core/driver/neptune_driver.py,sha256=dyQcaA5VnpNA_XkaWdvgGN3Q0QqbxWcVIud--yT8qhE,11266
 graphiti_core/embedder/__init__.py,sha256=EL564ZuE-DZjcuKNUK_exMn_XHXm2LdO9fzdXePVKL4,179
 graphiti_core/embedder/azure_openai.py,sha256=OyomPwC1fIsddI-3n6g00kQFdQznZorBhHwkQKCLUok,2384
 graphiti_core/embedder/client.py,sha256=BXFMXvuPWxaAzPaPILnxtqQQ4JWBFQv9GdBLOXUWgwE,1158
@@ -27,14 +27,14 @@ graphiti_core/embedder/voyage.py,sha256=oJHAZiNqjdEJOKgoKfGWcxK2-Ewqn5UB3vrBwIwP
 graphiti_core/llm_client/__init__.py,sha256=QgBWUiCeBp6YiA_xqyrDvJ9jIyy1hngH8g7FWahN3nw,776
 graphiti_core/llm_client/anthropic_client.py,sha256=xTFcrgMDK77BwnChBhYj51Jaa2mRNI850oJv2pKZI0A,12892
 graphiti_core/llm_client/azure_openai_client.py,sha256=ekERggAekbb7enes1RJqdRChf_mjaZTFXsnMbxO7azQ,2497
-graphiti_core/llm_client/client.py,sha256=KUWq7Gq9J4PdP06lLCBEb8OSZOE6luPqaQ3xgtpZwWg,6835
+graphiti_core/llm_client/client.py,sha256=xF3KtXbgP0jC6nKHtIiP5m9dNzxuZaqqQHCKiexijjU,7053
 graphiti_core/llm_client/config.py,sha256=pivp29CDIbDPqgw5NF9Ok2AwcqTV5z5_Q1bgNs1CDGs,2560
 graphiti_core/llm_client/errors.py,sha256=pn6brRiLW60DAUIXJYKBT6MInrS4ueuH1hNLbn_JbQo,1243
-graphiti_core/llm_client/gemini_client.py,sha256=AxD7sqsPQdgfcZCBIGN302s1hFYlBN9FOQcDEV0tw08,17725
+graphiti_core/llm_client/gemini_client.py,sha256=ohwuvJ-YTJ67xr6t5UYwSFo87WsyHeMiu8vNCifHod0,17850
 graphiti_core/llm_client/groq_client.py,sha256=bYLE_cg1QEhugsJOXh4b1vPbxagKeMWqk48240GCzMs,2922
-graphiti_core/llm_client/openai_base_client.py,sha256=LeEBZ33Y_bIz-YSr6aCbYKMI9r0SNPeZkALXQ0iFsSE,8488
+graphiti_core/llm_client/openai_base_client.py,sha256=HGt4CyyFCSZyCBwR__IbUUBF0V6Qwr9Ydu_XLtXPIA8,8533
 graphiti_core/llm_client/openai_client.py,sha256=AuaCFQFMJEGzBkFVouccq3XentmWRIKW0RLRBCUMm7Y,3763
-graphiti_core/llm_client/openai_generic_client.py,sha256=lyOQwzIMVb9pk3WWrU5zsG38J26QGKebxC40-lRYMJg,7007
+graphiti_core/llm_client/openai_generic_client.py,sha256=UseKg9rCqXizAdG1xGGU-jnfwuWJCvVkf-legT0MqjQ,7052
 graphiti_core/llm_client/utils.py,sha256=zKpxXEbKa369m4W7RDEf-m56kH46V1Mx3RowcWZEWWs,1000
 graphiti_core/migrations/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 graphiti_core/models/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
@@ -69,13 +69,13 @@ graphiti_core/utils/datetime_utils.py,sha256=J-zYSq7-H-2n9hYOXNIun12kM10vNX9mMAT
 graphiti_core/utils/maintenance/__init__.py,sha256=vW4H1KyapTl-OOz578uZABYcpND4wPx3Vt6aAPaXh78,301
 graphiti_core/utils/maintenance/community_operations.py,sha256=3IMxfOacZAYtZKebyYtWJYNZPLOPlS8Il-lzitEkoos,10681
 graphiti_core/utils/maintenance/dedup_helpers.py,sha256=B7k6KkB6Sii8PZCWNNTvsNiy4BNTNWpoLeGgrPLq6BE,9220
-graphiti_core/utils/maintenance/edge_operations.py,sha256=9jbFNM1Qm0wJJr9BR6gXyMiRuDgClim0MspDMBQmW40,26404
+graphiti_core/utils/maintenance/edge_operations.py,sha256=1hlcJRFnxthGkSr07QyDcOVug7N8dQj5aIENJ17JrpA,26564
 graphiti_core/utils/maintenance/graph_data_operations.py,sha256=42icj3S_ELAJ-NK3jVS_rg_243dmnaZOyUitJj_uJ-M,6085
-graphiti_core/utils/maintenance/node_operations.py,sha256=IKiqRqTeePTVFsl1X_N8DVRVAIhrSob7YPkuLvRM_Rk,18622
+graphiti_core/utils/maintenance/node_operations.py,sha256=ARng4x_pCpfA3g4bM0BncOkxBPaQ2IsdIaYfVq3V3X0,19603
 graphiti_core/utils/maintenance/temporal_operations.py,sha256=wq1I4kqeIoswit6sPohug91FEwrGaVnJ06g1vkJjSLY,3442
 graphiti_core/utils/maintenance/utils.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 graphiti_core/utils/ontology_utils/entity_types_utils.py,sha256=4eVgxLWY6Q8k9cRJ5pW59IYF--U4nXZsZIGOVb_yHfQ,1285
-graphiti_core-0.21.0rc12.dist-info/METADATA,sha256=5lFgZ88TQ2wk5EMyx_S3eljGZ9RL8jdLvqxzUTukVgA,27085
-graphiti_core-0.21.0rc12.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-graphiti_core-0.21.0rc12.dist-info/licenses/LICENSE,sha256=KCUwCyDXuVEgmDWkozHyniRyWjnWUWjkuDHfU6o3JlA,11325
-graphiti_core-0.21.0rc12.dist-info/RECORD,,
+graphiti_core-0.22.0rc0.dist-info/METADATA,sha256=uYwosUYjFpfCidLgBO8OuZ_P2fcrwNkRIkXn33lIwXk,27084
+graphiti_core-0.22.0rc0.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+graphiti_core-0.22.0rc0.dist-info/licenses/LICENSE,sha256=KCUwCyDXuVEgmDWkozHyniRyWjnWUWjkuDHfU6o3JlA,11325
+graphiti_core-0.22.0rc0.dist-info/RECORD,,