graphiti-core 0.18.9__py3-none-any.whl → 0.19.0rc2__py3-none-any.whl

graphiti_core/driver/driver.py

@@ -27,6 +27,7 @@ logger = logging.getLogger(__name__)
 class GraphProvider(Enum):
     NEO4J = 'neo4j'
     FALKORDB = 'falkordb'
+    NEPTUNE = 'neptune'
 
 
 class GraphDriverSession(ABC):

graphiti_core/driver/neptune_driver.py

@@ -0,0 +1,299 @@
+"""
+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.
+"""
+
+import asyncio
+import datetime
+import logging
+from collections.abc import Coroutine
+from typing import Any
+
+import boto3
+from langchain_aws.graphs import NeptuneAnalyticsGraph, NeptuneGraph
+from opensearchpy import OpenSearch, Urllib3AWSV4SignerAuth, Urllib3HttpConnection, helpers
+
+from graphiti_core.driver.driver import GraphDriver, GraphDriverSession, GraphProvider
+
+logger = logging.getLogger(__name__)
+DEFAULT_SIZE = 10
+
+aoss_indices = [
+    {
+        'index_name': 'node_name_and_summary',
+        'body': {
+            'mappings': {
+                'properties': {
+                    'uuid': {'type': 'keyword'},
+                    'name': {'type': 'text'},
+                    'summary': {'type': 'text'},
+                    'group_id': {'type': 'text'},
+                }
+            }
+        },
+        'query': {
+            'query': {'multi_match': {'query': '', 'fields': ['name', 'summary', 'group_id']}},
+            'size': DEFAULT_SIZE,
+        },
+    },
+    {
+        'index_name': 'community_name',
+        'body': {
+            'mappings': {
+                'properties': {
+                    'uuid': {'type': 'keyword'},
+                    'name': {'type': 'text'},
+                    'group_id': {'type': 'text'},
+                }
+            }
+        },
+        'query': {
+            'query': {'multi_match': {'query': '', 'fields': ['name', 'group_id']}},
+            'size': DEFAULT_SIZE,
+        },
+    },
+    {
+        'index_name': 'episode_content',
+        'body': {
+            'mappings': {
+                'properties': {
+                    'uuid': {'type': 'keyword'},
+                    'content': {'type': 'text'},
+                    'source': {'type': 'text'},
+                    'source_description': {'type': 'text'},
+                    'group_id': {'type': 'text'},
+                }
+            }
+        },
+        'query': {
+            'query': {
+                'multi_match': {
+                    'query': '',
+                    'fields': ['content', 'source', 'source_description', 'group_id'],
+                }
+            },
+            'size': DEFAULT_SIZE,
+        },
+    },
+    {
+        'index_name': 'edge_name_and_fact',
+        'body': {
+            'mappings': {
+                'properties': {
+                    'uuid': {'type': 'keyword'},
+                    'name': {'type': 'text'},
+                    'fact': {'type': 'text'},
+                    'group_id': {'type': 'text'},
+                }
+            }
+        },
+        'query': {
+            'query': {'multi_match': {'query': '', 'fields': ['name', 'fact', 'group_id']}},
+            'size': DEFAULT_SIZE,
+        },
+    },
+]
+
+
+class NeptuneDriver(GraphDriver):
+    provider: GraphProvider = GraphProvider.NEPTUNE
+
+    def __init__(self, host: str, aoss_host: str, port: int = 8182, aoss_port: int = 443):
+        """This initializes a NeptuneDriver for use with Neptune as a backend
+
+        Args:
+            host (str): The Neptune Database or Neptune Analytics host
+            aoss_host (str): The OpenSearch host value
+            port (int, optional): The Neptune Database port, ignored for Neptune Analytics. Defaults to 8182.
+            aoss_port (int, optional): The OpenSearch port. Defaults to 443.
+        """
+        if not host:
+            raise ValueError('You must provide an endpoint to create a NeptuneDriver')
+
+        if host.startswith('neptune-db://'):
+            # This is a Neptune Database Cluster
+            endpoint = host.replace('neptune-db://', '')
+            self.client = NeptuneGraph(endpoint, port)
+            logger.debug('Creating Neptune Database session for %s', host)
+        elif host.startswith('neptune-graph://'):
+            # This is a Neptune Analytics Graph
+            graphId = host.replace('neptune-graph://', '')
+            self.client = NeptuneAnalyticsGraph(graphId)
+            logger.debug('Creating Neptune Graph session for %s', host)
+        else:
+            raise ValueError(
+                'You must provide an endpoint to create a NeptuneDriver as either neptune-db://<endpoint> or neptune-graph://<graphid>'
+            )
+
+        if not aoss_host:
+            raise ValueError('You must provide an AOSS endpoint to create an OpenSearch driver.')
+
+        session = boto3.Session()
+        self.aoss_client = OpenSearch(
+            hosts=[{'host': aoss_host, 'port': aoss_port}],
+            http_auth=Urllib3AWSV4SignerAuth(
+                session.get_credentials(), session.region_name, 'aoss'
+            ),
+            use_ssl=True,
+            verify_certs=True,
+            connection_class=Urllib3HttpConnection,
+            pool_maxsize=20,
+        )
+
+    def _sanitize_parameters(self, query, params: dict):
+        if isinstance(query, list):
+            queries = []
+            for q in query:
+                queries.append(self._sanitize_parameters(q, params))
+            return queries
+        else:
+            for k, v in params.items():
+                if isinstance(v, datetime.datetime):
+                    params[k] = v.isoformat()
+                elif isinstance(v, list):
+                    # Handle lists that might contain datetime objects
+                    for i, item in enumerate(v):
+                        if isinstance(item, datetime.datetime):
+                            v[i] = item.isoformat()
+                            query = str(query).replace(f'${k}', f'datetime(${k})')
+                        if isinstance(item, dict):
+                            query = self._sanitize_parameters(query, v[i])
+
+                    # If the list contains datetime objects, we need to wrap each element with datetime()
+                    if any(isinstance(item, str) and 'T' in item for item in v):
+                        # Create a new list expression with datetime() wrapped around each element
+                        datetime_list = (
+                            '['
+                            + ', '.join(
+                                f'datetime("{item}")'
+                                if isinstance(item, str) and 'T' in item
+                                else repr(item)
+                                for item in v
+                            )
+                            + ']'
+                        )
+                        query = str(query).replace(f'${k}', datetime_list)
+                elif isinstance(v, dict):
+                    query = self._sanitize_parameters(query, v)
+            return query
+
+    async def execute_query(
+        self, cypher_query_, **kwargs: Any
+    ) -> tuple[dict[str, Any], None, None]:
+        params = dict(kwargs)
+        if isinstance(cypher_query_, list):
+            for q in cypher_query_:
+                result, _, _ = self._run_query(q[0], q[1])
+            return result, None, None
+        else:
+            return self._run_query(cypher_query_, params)
+
+    def _run_query(self, cypher_query_, params):
+        cypher_query_ = str(self._sanitize_parameters(cypher_query_, params))
+        try:
+            result = self.client.query(cypher_query_, params=params)
+        except Exception as e:
+            logger.error('Query: %s', cypher_query_)
+            logger.error('Parameters: %s', params)
+            logger.error('Error executing query: %s', e)
+            raise e
+
+        return result, None, None
+
+    def session(self, database: str | None = None) -> GraphDriverSession:
+        return NeptuneDriverSession(driver=self)
+
+    async def close(self) -> None:
+        return self.client.client.close()
+
+    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 aoss_indices:
+            index_name = index['index_name']
+            client = self.aoss_client
+            if not client.indices.exists(index=index_name):
+                client.indices.create(index=index_name, body=index['body'])
+        # Sleep for 1 minute to let the index creation complete
+        await asyncio.sleep(60)
+
+    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}
+                    for p in index['body']['mappings']['properties']:
+                        item[p] = d[p]
+                    to_index.append(item)
+                success, failed = helpers.bulk(self.aoss_client, to_index, stats_only=True)
+                if failed > 0:
+                    return success
+                else:
+                    return 0
+
+        return 0
+
+
+class NeptuneDriverSession(GraphDriverSession):
+    def __init__(self, driver: NeptuneDriver):  # type: ignore[reportUnknownArgumentType]
+        self.driver = driver
+
+    async def __aenter__(self):
+        return self
+
+    async def __aexit__(self, exc_type, exc, tb):
+        # No cleanup needed for Neptune, but method must exist
+        pass
+
+    async def close(self):
+        # No explicit close needed for Neptune, but method must exist
+        pass
+
+    async def execute_write(self, func, *args, **kwargs):
+        # Directly await the provided async function with `self` as the transaction/session
+        return await func(self, *args, **kwargs)
+
+    async def run(self, query: str | list, **kwargs: Any) -> Any:
+        if isinstance(query, list):
+            res = None
+            for q in query:
+                res = await self.driver.execute_query(q, **kwargs)
+            return res
+        else:
+            return await self.driver.execute_query(str(query), **kwargs)

graphiti_core/edges.py

@@ -24,13 +24,14 @@ from uuid import uuid4
 from pydantic import BaseModel, Field
 from typing_extensions import LiteralString
 
-from graphiti_core.driver.driver import GraphDriver
+from graphiti_core.driver.driver import GraphDriver, GraphProvider
 from graphiti_core.embedder import EmbedderClient
 from graphiti_core.errors import EdgeNotFoundError, GroupsEdgesNotFoundError
 from graphiti_core.helpers import parse_db_date
 from graphiti_core.models.edges.edge_db_queries import (
     COMMUNITY_EDGE_RETURN,
     ENTITY_EDGE_RETURN,
+    ENTITY_EDGE_RETURN_NEPTUNE,
     EPISODIC_EDGE_RETURN,
     EPISODIC_EDGE_SAVE,
     get_community_edge_save_query,
@@ -214,11 +215,19 @@ class EntityEdge(Edge):
         return self.fact_embedding
 
     async def load_fact_embedding(self, driver: GraphDriver):
-        records, _, _ = await driver.execute_query(
+        if driver.provider == GraphProvider.NEPTUNE:
+            query: LiteralString = """
+            MATCH (n:Entity)-[e:RELATES_TO {uuid: $uuid}]->(m:Entity)
+                RETURN [x IN split(e.fact_embedding, ",") | toFloat(x)] as fact_embedding
             """
+        else:
+            query: LiteralString = """
             MATCH (n:Entity)-[e:RELATES_TO {uuid: $uuid}]->(m:Entity)
             RETURN e.fact_embedding AS fact_embedding
-            """,
+            """
+
+        records, _, _ = await driver.execute_query(
+            query,
             uuid=self.uuid,
             routing_='r',
         )
@@ -246,6 +255,9 @@ class EntityEdge(Edge):
 
         edge_data.update(self.attributes or {})
 
+        if driver.provider == GraphProvider.NEPTUNE:
+            driver.save_to_aoss('edge_name_and_fact', [edge_data])  # pyright: ignore reportAttributeAccessIssue
+
         result = await driver.execute_query(
             get_entity_edge_save_query(driver.provider),
             edge_data=edge_data,
@@ -262,7 +274,11 @@ class EntityEdge(Edge):
             MATCH (n:Entity)-[e:RELATES_TO {uuid: $uuid}]->(m:Entity)
             RETURN
             """
-            + ENTITY_EDGE_RETURN,
+            + (
+                ENTITY_EDGE_RETURN_NEPTUNE
+                if driver.provider == GraphProvider.NEPTUNE
+                else ENTITY_EDGE_RETURN
+            ),
             uuid=uuid,
             routing_='r',
         )
@@ -284,7 +300,11 @@ class EntityEdge(Edge):
             WHERE e.uuid IN $uuids
             RETURN
             """
-            + ENTITY_EDGE_RETURN,
+            + (
+                ENTITY_EDGE_RETURN_NEPTUNE
+                if driver.provider == GraphProvider.NEPTUNE
+                else ENTITY_EDGE_RETURN
+            ),
             uuids=uuids,
             routing_='r',
         )
@@ -321,7 +341,11 @@ class EntityEdge(Edge):
             + """
             RETURN
             """
-            + ENTITY_EDGE_RETURN
+            + (
+                ENTITY_EDGE_RETURN_NEPTUNE
+                if driver.provider == GraphProvider.NEPTUNE
+                else ENTITY_EDGE_RETURN
+            )
             + with_embeddings_query
             + """
             ORDER BY e.uuid DESC 
@@ -346,7 +370,11 @@ class EntityEdge(Edge):
             MATCH (n:Entity {uuid: $node_uuid})-[e:RELATES_TO]-(m:Entity)
             RETURN
             """
-            + ENTITY_EDGE_RETURN,
+            + (
+                ENTITY_EDGE_RETURN_NEPTUNE
+                if driver.provider == GraphProvider.NEPTUNE
+                else ENTITY_EDGE_RETURN
+            ),
             node_uuid=node_uuid,
             routing_='r',
         )

graphiti_core/graphiti.py

@@ -89,6 +89,7 @@ from graphiti_core.utils.maintenance.edge_operations import (
 )
 from graphiti_core.utils.maintenance.graph_data_operations import (
     EPISODE_WINDOW_LEN,
+    build_dynamic_indexes,
     build_indices_and_constraints,
     retrieve_episodes,
 )
@@ -450,6 +451,7 @@ class Graphiti:
 
             validate_excluded_entity_types(excluded_entity_types, entity_types)
             validate_group_id(group_id)
+            await build_dynamic_indexes(self.driver, group_id)
 
             previous_episodes = (
                 await self.retrieve_episodes(

graphiti_core/llm_client/config.py

@@ -17,7 +17,7 @@ limitations under the License.
 from enum import Enum
 
 DEFAULT_MAX_TOKENS = 8192
-DEFAULT_TEMPERATURE = 0
+DEFAULT_TEMPERATURE = 1
 
 
 class ModelSize(Enum):

graphiti_core/llm_client/openai_base_client.py

@@ -31,8 +31,10 @@ from .errors import RateLimitError, RefusalError
 
 logger = logging.getLogger(__name__)
 
-DEFAULT_MODEL = 'gpt-4.1-mini'
-DEFAULT_SMALL_MODEL = 'gpt-4.1-nano'
+DEFAULT_MODEL = 'gpt-5-mini'
+DEFAULT_SMALL_MODEL = 'gpt-5-nano'
+DEFAULT_REASONING = 'minimal'
+DEFAULT_VERBOSITY = 'low'
 
 
 class BaseOpenAIClient(LLMClient):
@@ -51,6 +53,8 @@ class BaseOpenAIClient(LLMClient):
         config: LLMConfig | None = None,
         cache: bool = False,
         max_tokens: int = DEFAULT_MAX_TOKENS,
+        reasoning: str | None = DEFAULT_REASONING,
+        verbosity: str | None = DEFAULT_VERBOSITY,
     ):
         if cache:
             raise NotImplementedError('Caching is not implemented for OpenAI-based clients')
@@ -60,6 +64,8 @@ class BaseOpenAIClient(LLMClient):
 
         super().__init__(config, cache)
         self.max_tokens = max_tokens
+        self.reasoning = reasoning
+        self.verbosity = verbosity
 
     @abstractmethod
     async def _create_completion(
@@ -81,6 +87,8 @@ class BaseOpenAIClient(LLMClient):
         temperature: float | None,
         max_tokens: int,
         response_model: type[BaseModel],
+        reasoning: str | None,
+        verbosity: str | None,
     ) -> Any:
         """Create a structured completion using the specific client implementation."""
         pass
@@ -140,6 +148,8 @@ class BaseOpenAIClient(LLMClient):
                     temperature=self.temperature,
                     max_tokens=max_tokens or self.max_tokens,
                     response_model=response_model,
+                    reasoning=self.reasoning,
+                    verbosity=self.verbosity,
                 )
                 return self._handle_structured_response(response)
             else:

graphiti_core/llm_client/openai_client.py

@@ -21,7 +21,7 @@ from openai.types.chat import ChatCompletionMessageParam
 from pydantic import BaseModel
 
 from .config import DEFAULT_MAX_TOKENS, LLMConfig
-from .openai_base_client import BaseOpenAIClient
+from .openai_base_client import DEFAULT_REASONING, DEFAULT_VERBOSITY, BaseOpenAIClient
 
 
 class OpenAIClient(BaseOpenAIClient):
@@ -41,6 +41,8 @@ class OpenAIClient(BaseOpenAIClient):
         cache: bool = False,
         client: typing.Any = None,
         max_tokens: int = DEFAULT_MAX_TOKENS,
+        reasoning: str = DEFAULT_REASONING,
+        verbosity: str = DEFAULT_VERBOSITY,
     ):
         """
         Initialize the OpenAIClient with the provided configuration, cache setting, and client.
@@ -50,7 +52,7 @@ class OpenAIClient(BaseOpenAIClient):
             cache (bool): Whether to use caching for responses. Defaults to False.
             client (Any | None): An optional async client instance to use. If not provided, a new AsyncOpenAI client is created.
         """
-        super().__init__(config, cache, max_tokens)
+        super().__init__(config, cache, max_tokens, reasoning, verbosity)
 
         if config is None:
             config = LLMConfig()
@@ -67,6 +69,8 @@ class OpenAIClient(BaseOpenAIClient):
         temperature: float | None,
         max_tokens: int,
         response_model: type[BaseModel],
+        reasoning: str | None = None,
+        verbosity: str | None = None,
     ):
         """Create a structured completion using OpenAI's beta parse API."""
         response = await self.client.responses.parse(
@@ -75,6 +79,8 @@ class OpenAIClient(BaseOpenAIClient):
             temperature=temperature,
             max_output_tokens=max_tokens,
             text_format=response_model,  # type: ignore
+            reasoning={'effort': reasoning} if reasoning is not None else None,  # type: ignore
+            text={'verbosity': verbosity} if verbosity is not None else None,  # type: ignore
         )
 
         return response
@@ -86,6 +92,8 @@ class OpenAIClient(BaseOpenAIClient):
         temperature: float | None,
         max_tokens: int,
         response_model: type[BaseModel] | None = None,
+        reasoning: str | None = None,
+        verbosity: str | None = None,
     ):
         """Create a regular completion with JSON format."""
         return await self.client.chat.completions.create(

graphiti_core/migrations/neo4j_node_group_labels.py

@@ -0,0 +1,53 @@
+from graphiti_core.driver.driver import GraphDriver
+from graphiti_core.helpers import validate_group_id
+from graphiti_core.utils.maintenance.graph_data_operations import build_dynamic_indexes
+
+
+async def neo4j_node_group_labels(driver: GraphDriver, group_id: str, batch_size: int = 100):
+    validate_group_id(group_id)
+    await build_dynamic_indexes(driver, group_id)
+
+    episode_query = """
+                        MATCH (n:Episodic {group_id: $group_id})
+                        CALL {
+                            WITH n
+                            SET n:$group_label
+                        } IN TRANSACTIONS OF $batch_size ROWS"""
+
+    entity_query = """
+                        MATCH (n:Entity {group_id: $group_id})
+                        CALL {
+                            WITH n
+                            SET n:$group_label
+                        } IN TRANSACTIONS OF $batch_size ROWS"""
+
+    community_query = """
+                        MATCH (n:Community {group_id: $group_id})
+                        CALL {
+                            WITH n
+                            SET n:$group_label
+                        } IN TRANSACTIONS OF $batch_size ROWS"""
+
+    async with driver.session() as session:
+        await session.run(
+            episode_query,
+            group_id=group_id,
+            group_label='Episodic_' + group_id.replace('-', ''),
+            batch_size=batch_size,
+        )
+
+    async with driver.session() as session:
+        await session.run(
+            entity_query,
+            group_id=group_id,
+            group_label='Entity_' + group_id.replace('-', ''),
+            batch_size=batch_size,
+        )
+
+    async with driver.session() as session:
+        await session.run(
+            community_query,
+            group_id=group_id,
+            group_label='Community_' + group_id.replace('-', ''),
+            batch_size=batch_size,
+        )