graphiti-core 0.21.0rc13__py3-none-any.whl → 0.22.0rc1__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/prompts/extract_nodes.py

@@ -23,39 +23,44 @@ from .prompt_helpers import to_prompt_json
 
 
 class ExtractedEntity(BaseModel):
-    name: str = Field(..., description='Name of the extracted entity')
+    name: str = Field(..., description="Name of the extracted entity")
     entity_type_id: int = Field(
-        description='ID of the classified entity type. '
-        'Must be one of the provided entity_type_id integers.',
+        description="ID of the classified entity type. "
+        "Must be one of the provided entity_type_id integers.",
     )
 
 
 class ExtractedEntities(BaseModel):
-    extracted_entities: list[ExtractedEntity] = Field(..., description='List of extracted entities')
+    extracted_entities: list[ExtractedEntity] = Field(
+        ..., description="List of extracted entities"
+    )
 
 
 class MissedEntities(BaseModel):
-    missed_entities: list[str] = Field(..., description="Names of entities that weren't extracted")
+    missed_entities: list[str] = Field(
+        ..., description="Names of entities that weren't extracted"
+    )
 
 
 class EntityClassificationTriple(BaseModel):
-    uuid: str = Field(description='UUID of the entity')
-    name: str = Field(description='Name of the entity')
+    uuid: str = Field(description="UUID of the entity")
+    name: str = Field(description="Name of the entity")
     entity_type: str | None = Field(
-        default=None, description='Type of the entity. Must be one of the provided types or None'
+        default=None,
+        description="Type of the entity. Must be one of the provided types or None",
     )
 
 
 class EntityClassification(BaseModel):
     entity_classifications: list[EntityClassificationTriple] = Field(
-        ..., description='List of entities classification triples.'
+        ..., description="List of entities classification triples."
     )
 
 
 class EntitySummary(BaseModel):
     summary: str = Field(
         ...,
-        description='Summary containing the important information about the entity. Under 250 words',
+        description="Summary containing the important information about the entity. Under 8 sentences.",
     )
 
 
@@ -123,8 +128,8 @@ reference entities. Only extract distinct entities from the CURRENT MESSAGE. Don
 {context['custom_prompt']}
 """
     return [
-        Message(role='system', content=sys_prompt),
-        Message(role='user', content=user_prompt),
+        Message(role="system", content=sys_prompt),
+        Message(role="user", content=user_prompt),
     ]
 
 
@@ -156,8 +161,8 @@ Guidelines:
 3. Do NOT extract any properties that contain dates
 """
     return [
-        Message(role='system', content=sys_prompt),
-        Message(role='user', content=user_prompt),
+        Message(role="system", content=sys_prompt),
+        Message(role="user", content=user_prompt),
     ]
 
 
@@ -187,8 +192,8 @@ Guidelines:
 4. Be as explicit as possible in your node names, using full names and avoiding abbreviations.
 """
     return [
-        Message(role='system', content=sys_prompt),
-        Message(role='user', content=user_prompt),
+        Message(role="system", content=sys_prompt),
+        Message(role="user", content=user_prompt),
     ]
 
 
@@ -211,8 +216,8 @@ Given the above previous messages, current message, and list of extracted entiti
 extracted.
 """
     return [
-        Message(role='system', content=sys_prompt),
-        Message(role='user', content=user_prompt),
+        Message(role="system", content=sys_prompt),
+        Message(role="user", content=user_prompt),
     ]
 
 
@@ -243,19 +248,19 @@ def classify_nodes(context: dict[str, Any]) -> list[Message]:
     3. If none of the provided entity types accurately classify an extracted node, the type should be set to None
 """
     return [
-        Message(role='system', content=sys_prompt),
-        Message(role='user', content=user_prompt),
+        Message(role="system", content=sys_prompt),
+        Message(role="user", content=user_prompt),
     ]
 
 
 def extract_attributes(context: dict[str, Any]) -> list[Message]:
     return [
         Message(
-            role='system',
-            content='You are a helpful assistant that extracts entity properties from the provided text.',
+            role="system",
+            content="You are a helpful assistant that extracts entity properties from the provided text.",
         ),
         Message(
-            role='user',
+            role="user",
             content=f"""
 
         <MESSAGES>
@@ -281,11 +286,11 @@ def extract_attributes(context: dict[str, Any]) -> list[Message]:
 def extract_summary(context: dict[str, Any]) -> list[Message]:
     return [
         Message(
-            role='system',
-            content='You are a helpful assistant that extracts entity summaries from the provided text.',
+            role="system",
+            content="You are a helpful assistant that extracts entity summaries from the provided text.",
         ),
         Message(
-            role='user',
+            role="user",
             content=f"""
 
         <MESSAGES>
@@ -300,7 +305,7 @@ def extract_summary(context: dict[str, Any]) -> list[Message]:
         1. Do not hallucinate entity summary information if they cannot be found in the current context.
         2. Only use the provided MESSAGES and ENTITY to set attribute values.
         3. The summary attribute represents a summary of the ENTITY, and should be updated with new information about the Entity from the MESSAGES. 
-            Summaries must be no longer than 250 words.
+        4. Keep the summary concise and to the point. SUMMARIES MUST BE LESS THAN 8 SENTENCES.
 
         <ENTITY>
         {context['node']}
@@ -311,11 +316,11 @@ def extract_summary(context: dict[str, Any]) -> list[Message]:
 
 
 versions: Versions = {
-    'extract_message': extract_message,
-    'extract_json': extract_json,
-    'extract_text': extract_text,
-    'reflexion': reflexion,
-    'extract_summary': extract_summary,
-    'classify_nodes': classify_nodes,
-    'extract_attributes': extract_attributes,
+    "extract_message": extract_message,
+    "extract_json": extract_json,
+    "extract_text": extract_text,
+    "reflexion": reflexion,
+    "extract_summary": extract_summary,
+    "classify_nodes": classify_nodes,
+    "extract_attributes": extract_attributes,
 }

graphiti_core/prompts/summarize_nodes.py

@@ -25,12 +25,14 @@ from .prompt_helpers import to_prompt_json
 class Summary(BaseModel):
     summary: str = Field(
         ...,
-        description='Summary containing the important information about the entity. Under 250 words',
+        description="Summary containing the important information about the entity. Under 8 sentences",
     )
 
 
 class SummaryDescription(BaseModel):
-    description: str = Field(..., description='One sentence description of the provided summary')
+    description: str = Field(
+        ..., description="One sentence description of the provided summary"
+    )
 
 
 class Prompt(Protocol):
@@ -48,15 +50,15 @@ class Versions(TypedDict):
 def summarize_pair(context: dict[str, Any]) -> list[Message]:
     return [
         Message(
-            role='system',
-            content='You are a helpful assistant that combines summaries.',
+            role="system",
+            content="You are a helpful assistant that combines summaries.",
         ),
         Message(
-            role='user',
+            role="user",
             content=f"""
         Synthesize the information from the following two summaries into a single succinct summary.
         
-        Summaries must be under 250 words.
+        IMPORTANT: Keep the summary concise and to the point. SUMMARIES MUST BE LESS THAN 8 SENTENCES.
 
         Summaries:
         {to_prompt_json(context['node_summaries'], indent=2)}
@@ -68,11 +70,11 @@ def summarize_pair(context: dict[str, Any]) -> list[Message]:
 def summarize_context(context: dict[str, Any]) -> list[Message]:
     return [
         Message(
-            role='system',
-            content='You are a helpful assistant that extracts entity properties from the provided text.',
+            role="system",
+            content="You are a helpful assistant that generates a summary and attributes from provided text.",
         ),
         Message(
-            role='user',
+            role="user",
             content=f"""
             
         <MESSAGES>
@@ -82,7 +84,7 @@ def summarize_context(context: dict[str, Any]) -> list[Message]:
         
         Given the above MESSAGES and the following ENTITY name, create a summary for the ENTITY. Your summary must only use
         information from the provided MESSAGES. Your summary should also only contain information relevant to the
-        provided ENTITY. Summaries must be under 250 words.
+        provided ENTITY. 
         
         In addition, extract any values for the provided entity properties based on their descriptions.
         If the value of the entity property cannot be found in the current context, set the value of the property to the Python value None.
@@ -90,6 +92,7 @@ def summarize_context(context: dict[str, Any]) -> list[Message]:
         Guidelines:
         1. Do not hallucinate entity property values if they cannot be found in the current context.
         2. Only use the provided messages, entity, and entity context to set attribute values.
+        3. Keep the summary concise and to the point. SUMMARIES MUST BE LESS THAN 8 SENTENCES.
         
         <ENTITY>
         {context['node_name']}
@@ -110,14 +113,14 @@ def summarize_context(context: dict[str, Any]) -> list[Message]:
 def summary_description(context: dict[str, Any]) -> list[Message]:
     return [
         Message(
-            role='system',
-            content='You are a helpful assistant that describes provided contents in a single sentence.',
+            role="system",
+            content="You are a helpful assistant that describes provided contents in a single sentence.",
         ),
         Message(
-            role='user',
+            role="user",
             content=f"""
         Create a short one sentence description of the summary that explains what kind of information is summarized.
-        Summaries must be under 250 words.
+        Summaries must be under 8 sentences.
 
         Summary:
         {to_prompt_json(context['summary'], indent=2)}
@@ -127,7 +130,7 @@ def summary_description(context: dict[str, Any]) -> list[Message]:
 
 
 versions: Versions = {
-    'summarize_pair': summarize_pair,
-    'summarize_context': summarize_context,
-    'summary_description': summary_description,
+    "summarize_pair": summarize_pair,
+    "summarize_context": summarize_context,
+    "summary_description": summary_description,
 }

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', [])

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
@@ -478,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
+    )
 
-    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 []
-        ),
-    }
+    node.attributes.update(llm_response)
+
+    return node
 
-    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.0rc13.dist-info → graphiti_core-0.22.0rc1.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: graphiti-core
-Version: 0.21.0rc13
+Version: 0.22.0rc1
 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.0rc13.dist-info → graphiti_core-0.22.0rc1.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
@@ -48,12 +48,12 @@ graphiti_core/prompts/dedupe_nodes.py,sha256=YNNo19Cq8koLVoLCafpjYJOy5nmRZ-tEWhv
 graphiti_core/prompts/eval.py,sha256=GWFkfZoPfY8U7mV8Ngd_5a2S2fHS7KjajChntxv1UEY,5360
 graphiti_core/prompts/extract_edge_dates.py,sha256=3Drs3CmvP0gJN5BidWSxrNvLet3HPoTybU3BUIAoc0Y,4218
 graphiti_core/prompts/extract_edges.py,sha256=-yOIvCPwxIAXeqYpNCzouE6i3WfdsexzRXFmcXpQpAg,7113
-graphiti_core/prompts/extract_nodes.py,sha256=orbZiCqVL-4RNlckpUaQIq7Yb2JsIwT8e3ZAHRAEiLM,11281
+graphiti_core/prompts/extract_nodes.py,sha256=jMD-XRi4U3kjp9smHtA_kvnMBGWBfpBoKc45IoTIZs0,11360
 graphiti_core/prompts/invalidate_edges.py,sha256=yfpcs_pyctnoM77ULPZXEtKW0oHr1MeLsJzC5yrE-o4,3547
 graphiti_core/prompts/lib.py,sha256=DCyHePM4_q-CptTpEXGO_dBv9k7xDtclEaB1dGu7EcI,4092
 graphiti_core/prompts/models.py,sha256=NgxdbPHJpBEcpbXovKyScgpBc73Q-GIW-CBDlBtDjto,894
 graphiti_core/prompts/prompt_helpers.py,sha256=dpWbB8IYAqAZoU5qBx896jozKiQJTng4dGzWewZ_s4c,814
-graphiti_core/prompts/summarize_nodes.py,sha256=p_TNDG66uY71QNDo9hyk4crAfyzyEKlb4_lML3fxeWU,4197
+graphiti_core/prompts/summarize_nodes.py,sha256=7WnjRgYo1Z9bfnUWaUUXbiaLqygGLpemvB5inhhq44Y,4340
 graphiti_core/search/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 graphiti_core/search/search.py,sha256=2kj7fybSFv6Fnf_cfEUhJhrpfzNtmkPPZ0hV3BQCDqg,18387
 graphiti_core/search/search_config.py,sha256=v_rUHsu1yo5OuPfEm21lSuXexQs-o8qYwSSemW2QWhU,4165
@@ -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=_vgUUIE8bOX3fIT9MeBwPRi21KWgDSlFGG59m_zThKw,26498
+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=3Km0uimYxU2pcaT8a0jm4xTP01zR8bp2z8ZiFItRwr0,18623
+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.0rc13.dist-info/METADATA,sha256=DpoXBL7QKCncx28h4rgzZk76DOmqDXobyUlywI593Aw,27085
-graphiti_core-0.21.0rc13.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-graphiti_core-0.21.0rc13.dist-info/licenses/LICENSE,sha256=KCUwCyDXuVEgmDWkozHyniRyWjnWUWjkuDHfU6o3JlA,11325
-graphiti_core-0.21.0rc13.dist-info/RECORD,,
+graphiti_core-0.22.0rc1.dist-info/METADATA,sha256=NlIXn-TmrQ-_u-6CI6I7sEC7ioBKvQIKEl0oyqRq4YM,27084
+graphiti_core-0.22.0rc1.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+graphiti_core-0.22.0rc1.dist-info/licenses/LICENSE,sha256=KCUwCyDXuVEgmDWkozHyniRyWjnWUWjkuDHfU6o3JlA,11325
+graphiti_core-0.22.0rc1.dist-info/RECORD,,