graphiti-core 0.17.4__py3-none-any.whl → 0.25.3__py3-none-any.whl

graphiti_core/cross_encoder/gemini_reranker_client.py

@@ -37,7 +37,7 @@ else:
 
 logger = logging.getLogger(__name__)
 
-DEFAULT_MODEL = 'gemini-2.5-flash-lite-preview-06-17'
+DEFAULT_MODEL = 'gemini-2.5-flash-lite'
 
 
 class GeminiRerankerClient(CrossEncoderClient):

graphiti_core/cross_encoder/openai_reranker_client.py

@@ -84,7 +84,7 @@ class OpenAIRerankerClient(CrossEncoderClient):
             responses = await semaphore_gather(
                 *[
                     self.client.chat.completions.create(
-                        model=DEFAULT_MODEL,
+                        model=self.config.model or DEFAULT_MODEL,
                         messages=openai_messages,
                         temperature=0,
                         max_tokens=1,

graphiti_core/decorators.py

@@ -0,0 +1,110 @@
+"""
+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 functools
+import inspect
+from collections.abc import Awaitable, Callable
+from typing import Any, TypeVar
+
+from graphiti_core.driver.driver import GraphProvider
+from graphiti_core.helpers import semaphore_gather
+from graphiti_core.search.search_config import SearchResults
+
+F = TypeVar('F', bound=Callable[..., Awaitable[Any]])
+
+
+def handle_multiple_group_ids(func: F) -> F:
+    """
+    Decorator for FalkorDB methods that need to handle multiple group_ids.
+    Runs the function for each group_id separately and merges results.
+    """
+
+    @functools.wraps(func)
+    async def wrapper(self, *args, **kwargs):
+        group_ids_func_pos = get_parameter_position(func, 'group_ids')
+        group_ids_pos = (
+            group_ids_func_pos - 1 if group_ids_func_pos is not None else None
+        )  # Adjust for zero-based index
+        group_ids = kwargs.get('group_ids')
+
+        # If not in kwargs and position exists, get from args
+        if group_ids is None and group_ids_pos is not None and len(args) > group_ids_pos:
+            group_ids = args[group_ids_pos]
+
+        # Only handle FalkorDB with multiple group_ids
+        if (
+            hasattr(self, 'clients')
+            and hasattr(self.clients, 'driver')
+            and self.clients.driver.provider == GraphProvider.FALKORDB
+            and group_ids
+            and len(group_ids) > 1
+        ):
+            # Execute for each group_id concurrently
+            driver = self.clients.driver
+
+            async def execute_for_group(gid: str):
+                # Remove group_ids from args if it was passed positionally
+                filtered_args = list(args)
+                if group_ids_pos is not None and len(args) > group_ids_pos:
+                    filtered_args.pop(group_ids_pos)
+
+                return await func(
+                    self,
+                    *filtered_args,
+                    **{**kwargs, 'group_ids': [gid], 'driver': driver.clone(database=gid)},
+                )
+
+            results = await semaphore_gather(
+                *[execute_for_group(gid) for gid in group_ids],
+                max_coroutines=getattr(self, 'max_coroutines', None),
+            )
+
+            # Merge results based on type
+            if isinstance(results[0], SearchResults):
+                return SearchResults.merge(results)
+            elif isinstance(results[0], list):
+                return [item for result in results for item in result]
+            elif isinstance(results[0], tuple):
+                # Handle tuple outputs (like build_communities returning (nodes, edges))
+                merged_tuple = []
+                for i in range(len(results[0])):
+                    component_results = [result[i] for result in results]
+                    if isinstance(component_results[0], list):
+                        merged_tuple.append(
+                            [item for component in component_results for item in component]
+                        )
+                    else:
+                        merged_tuple.append(component_results)
+                return tuple(merged_tuple)
+            else:
+                return results
+
+        # Normal execution
+        return await func(self, *args, **kwargs)
+
+    return wrapper  # type: ignore
+
+
+def get_parameter_position(func: Callable, param_name: str) -> int | None:
+    """
+    Returns the positional index of a parameter in the function signature.
+    If the parameter is not found, returns None.
+    """
+    sig = inspect.signature(func)
+    for idx, (name, _param) in enumerate(sig.parameters.items()):
+        if name == param_name:
+            return idx
+    return None

graphiti_core/driver/driver.py

@@ -14,15 +14,41 @@ See the License for the specific language governing permissions and
 limitations under the License.
 """
 
+import copy
 import logging
+import os
 from abc import ABC, abstractmethod
 from collections.abc import Coroutine
+from enum import Enum
 from typing import Any
 
+from dotenv import load_dotenv
+
+from graphiti_core.driver.graph_operations.graph_operations import GraphOperationsInterface
+from graphiti_core.driver.search_interface.search_interface import SearchInterface
+
 logger = logging.getLogger(__name__)
 
+DEFAULT_SIZE = 10
+
+load_dotenv()
+
+ENTITY_INDEX_NAME = os.environ.get('ENTITY_INDEX_NAME', 'entities')
+EPISODE_INDEX_NAME = os.environ.get('EPISODE_INDEX_NAME', 'episodes')
+COMMUNITY_INDEX_NAME = os.environ.get('COMMUNITY_INDEX_NAME', 'communities')
+ENTITY_EDGE_INDEX_NAME = os.environ.get('ENTITY_EDGE_INDEX_NAME', 'entity_edges')
+
+
+class GraphProvider(Enum):
+    NEO4J = 'neo4j'
+    FALKORDB = 'falkordb'
+    KUZU = 'kuzu'
+    NEPTUNE = 'neptune'
+
 
 class GraphDriverSession(ABC):
+    provider: GraphProvider
+
     async def __aenter__(self):
         return self
 
@@ -45,7 +71,14 @@ class GraphDriverSession(ABC):
 
 
 class GraphDriver(ABC):
-    provider: str
+    provider: GraphProvider
+    fulltext_syntax: str = (
+        ''  # Neo4j (default) syntax does not require a prefix for fulltext queries
+    )
+    _database: str
+    default_group_id: str = ''
+    search_interface: SearchInterface | None = None
+    graph_operations_interface: GraphOperationsInterface | None = None
 
     @abstractmethod
     def execute_query(self, cypher_query_: str, **kwargs: Any) -> Coroutine:
@@ -60,5 +93,32 @@ class GraphDriver(ABC):
         raise NotImplementedError()
 
     @abstractmethod
-    def delete_all_indexes(self, database_: str | None = None) -> Coroutine:
+    def delete_all_indexes(self) -> Coroutine:
         raise NotImplementedError()
+
+    def with_database(self, database: str) -> 'GraphDriver':
+        """
+        Returns a shallow copy of this driver with a different default database.
+        Reuses the same connection (e.g. FalkorDB, Neo4j).
+        """
+        cloned = copy.copy(self)
+        cloned._database = database
+
+        return cloned
+
+    @abstractmethod
+    async def build_indices_and_constraints(self, delete_existing: bool = False):
+        raise NotImplementedError()
+
+    def clone(self, database: str) -> 'GraphDriver':
+        """Clone the driver with a different database or graph name."""
+        return self
+
+    def build_fulltext_query(
+        self, query: str, group_ids: list[str] | None = None, max_query_length: int = 128
+    ) -> str:
+        """
+        Specific fulltext query builder for database providers.
+        Only implemented by providers that need custom fulltext query building.
+        """
+        raise NotImplementedError(f'build_fulltext_query not implemented for {self.provider}')

graphiti_core/driver/falkordb_driver.py

@@ -14,8 +14,9 @@ See the License for the specific language governing permissions and
 limitations under the License.
 """
 
+import asyncio
+import datetime
 import logging
-from datetime import datetime
 from typing import TYPE_CHECKING, Any
 
 if TYPE_CHECKING:
@@ -32,12 +33,52 @@ else:
             'Install it with: pip install graphiti-core[falkordb]'
         ) from None
 
-from graphiti_core.driver.driver import GraphDriver, GraphDriverSession
+from graphiti_core.driver.driver import GraphDriver, GraphDriverSession, GraphProvider
+from graphiti_core.graph_queries import get_fulltext_indices, get_range_indices
+from graphiti_core.utils.datetime_utils import convert_datetimes_to_strings
 
 logger = logging.getLogger(__name__)
 
+STOPWORDS = [
+    'a',
+    'is',
+    'the',
+    'an',
+    'and',
+    'are',
+    'as',
+    'at',
+    'be',
+    'but',
+    'by',
+    'for',
+    'if',
+    'in',
+    'into',
+    'it',
+    'no',
+    'not',
+    'of',
+    'on',
+    'or',
+    'such',
+    'that',
+    'their',
+    'then',
+    'there',
+    'these',
+    'they',
+    'this',
+    'to',
+    'was',
+    'will',
+    'with',
+]
+
 
 class FalkorDriverSession(GraphDriverSession):
+    provider = GraphProvider.FALKORDB
+
     def __init__(self, graph: FalkorGraph):
         self.graph = graph
 
@@ -71,7 +112,10 @@ class FalkorDriverSession(GraphDriverSession):
 
 
 class FalkorDriver(GraphDriver):
-    provider: str = 'falkordb'
+    provider = GraphProvider.FALKORDB
+    default_group_id: str = '\\_'
+    fulltext_syntax: str = '@'  # FalkorDB uses a redisearch-like syntax for fulltext queries
+    aoss_client: None = None
 
     def __init__(
         self,
@@ -88,14 +132,32 @@ class FalkorDriver(GraphDriver):
         FalkorDB is a multi-tenant graph database.
         To connect, provide the host and port.
         The default parameters assume a local (on-premises) FalkorDB instance.
+
+        Args:
+        host (str): The host where FalkorDB is running.
+        port (int): The port on which FalkorDB is listening.
+        username (str | None): The username for authentication (if required).
+        password (str | None): The password for authentication (if required).
+        falkor_db (FalkorDB | None): An existing FalkorDB instance to use instead of creating a new one.
+        database (str): The name of the database to connect to. Defaults to 'default_db'.
         """
         super().__init__()
+        self._database = database
         if falkor_db is not None:
             # If a FalkorDB instance is provided, use it directly
             self.client = falkor_db
         else:
             self.client = FalkorDB(host=host, port=port, username=username, password=password)
-            self._database = database
+
+        # Schedule the indices and constraints to be built
+        try:
+            # Try to get the current event loop
+            loop = asyncio.get_running_loop()
+            # Schedule the build_indices_and_constraints to run
+            loop.create_task(self.build_indices_and_constraints())
+        except RuntimeError:
+            # No event loop running, this will be handled later
+            pass
 
     def _get_graph(self, graph_name: str | None) -> FalkorGraph:
         # FalkorDB requires a non-None database name for multi-tenant graphs; the default is "default_db"
@@ -104,8 +166,7 @@ class FalkorDriver(GraphDriver):
         return self.client.select_graph(graph_name)
 
     async def execute_query(self, cypher_query_, **kwargs: Any):
-        graph_name = kwargs.pop('database_', self._database)
-        graph = self._get_graph(graph_name)
+        graph = self._get_graph(self._database)
 
         # Convert datetime objects to ISO strings (FalkorDB does not support datetime objects directly)
         params = convert_datetimes_to_strings(dict(kwargs))
@@ -117,7 +178,7 @@ class FalkorDriver(GraphDriver):
                 # check if index already exists
                 logger.info(f'Index already exists: {e}')
                 return None
-            logger.error(f'Error executing FalkorDB query: {e}')
+            logger.error(f'Error executing FalkorDB query: {e}\n{cypher_query_}\n{params}')
             raise
 
         # Convert the result header to a list of strings
@@ -149,22 +210,153 @@ class FalkorDriver(GraphDriver):
         elif hasattr(self.client.connection, 'close'):
             await self.client.connection.close()
 
-    async def delete_all_indexes(self, database_: str | None = None) -> None:
-        database = database_ or self._database
-        await self.execute_query(
-            'CALL db.indexes() YIELD name DROP INDEX name',
-            database_=database,
+    async def delete_all_indexes(self) -> None:
+        result = await self.execute_query('CALL db.indexes()')
+        if not result:
+            return
+
+        records, _, _ = result
+        drop_tasks = []
+
+        for record in records:
+            label = record['label']
+            entity_type = record['entitytype']
+
+            for field_name, index_type in record['types'].items():
+                if 'RANGE' in index_type:
+                    drop_tasks.append(self.execute_query(f'DROP INDEX ON :{label}({field_name})'))
+                elif 'FULLTEXT' in index_type:
+                    if entity_type == 'NODE':
+                        drop_tasks.append(
+                            self.execute_query(
+                                f'DROP FULLTEXT INDEX FOR (n:{label}) ON (n.{field_name})'
+                            )
+                        )
+                    elif entity_type == 'RELATIONSHIP':
+                        drop_tasks.append(
+                            self.execute_query(
+                                f'DROP FULLTEXT INDEX FOR ()-[e:{label}]-() ON (e.{field_name})'
+                            )
+                        )
+
+        if drop_tasks:
+            await asyncio.gather(*drop_tasks)
+
+    async def build_indices_and_constraints(self, delete_existing=False):
+        if delete_existing:
+            await self.delete_all_indexes()
+        index_queries = get_range_indices(self.provider) + get_fulltext_indices(self.provider)
+        for query in index_queries:
+            await self.execute_query(query)
+
+    def clone(self, database: str) -> 'GraphDriver':
+        """
+        Returns a shallow copy of this driver with a different default database.
+        Reuses the same connection (e.g. FalkorDB, Neo4j).
+        """
+        if database == self._database:
+            cloned = self
+        elif database == self.default_group_id:
+            cloned = FalkorDriver(falkor_db=self.client)
+        else:
+            # Create a new instance of FalkorDriver with the same connection but a different database
+            cloned = FalkorDriver(falkor_db=self.client, database=database)
+
+        return cloned
+
+    async def health_check(self) -> None:
+        """Check FalkorDB connectivity by running a simple query."""
+        try:
+            await self.execute_query('MATCH (n) RETURN 1 LIMIT 1')
+            return None
+        except Exception as e:
+            print(f'FalkorDB health check failed: {e}')
+            raise
+
+    @staticmethod
+    def convert_datetimes_to_strings(obj):
+        if isinstance(obj, dict):
+            return {k: FalkorDriver.convert_datetimes_to_strings(v) for k, v in obj.items()}
+        elif isinstance(obj, list):
+            return [FalkorDriver.convert_datetimes_to_strings(item) for item in obj]
+        elif isinstance(obj, tuple):
+            return tuple(FalkorDriver.convert_datetimes_to_strings(item) for item in obj)
+        elif isinstance(obj, datetime):
+            return obj.isoformat()
+        else:
+            return obj
+
+    def sanitize(self, query: str) -> str:
+        """
+        Replace FalkorDB special characters with whitespace.
+        Based on FalkorDB tokenization rules: ,.<>{}[]"':;!@#$%^&*()-+=~
+        """
+        # FalkorDB separator characters that break text into tokens
+        separator_map = str.maketrans(
+            {
+                ',': ' ',
+                '.': ' ',
+                '<': ' ',
+                '>': ' ',
+                '{': ' ',
+                '}': ' ',
+                '[': ' ',
+                ']': ' ',
+                '"': ' ',
+                "'": ' ',
+                ':': ' ',
+                ';': ' ',
+                '!': ' ',
+                '@': ' ',
+                '#': ' ',
+                '$': ' ',
+                '%': ' ',
+                '^': ' ',
+                '&': ' ',
+                '*': ' ',
+                '(': ' ',
+                ')': ' ',
+                '-': ' ',
+                '+': ' ',
+                '=': ' ',
+                '~': ' ',
+                '?': ' ',
+            }
         )
+        sanitized = query.translate(separator_map)
+        # Clean up multiple spaces
+        sanitized = ' '.join(sanitized.split())
+        return sanitized
+
+    def build_fulltext_query(
+        self, query: str, group_ids: list[str] | None = None, max_query_length: int = 128
+    ) -> str:
+        """
+        Build a fulltext query string for FalkorDB using RedisSearch syntax.
+        FalkorDB uses RedisSearch-like syntax where:
+        - Field queries use @ prefix: @field:value
+        - Multiple values for same field: (@field:value1|value2)
+        - Text search doesn't need @ prefix for content fields
+        - AND is implicit with space: (@group_id:value) (text)
+        - OR uses pipe within parentheses: (@group_id:value1|value2)
+        """
+        if group_ids is None or len(group_ids) == 0:
+            group_filter = ''
+        else:
+            group_values = '|'.join(group_ids)
+            group_filter = f'(@group_id:{group_values})'
+
+        sanitized_query = self.sanitize(query)
+
+        # Remove stopwords from the sanitized query
+        query_words = sanitized_query.split()
+        filtered_words = [word for word in query_words if word.lower() not in STOPWORDS]
+        sanitized_query = ' | '.join(filtered_words)
+
+        # If the query is too long return no query
+        if len(sanitized_query.split(' ')) + len(group_ids or '') >= max_query_length:
+            return ''
 
+        full_query = group_filter + ' (' + sanitized_query + ')'
 
-def convert_datetimes_to_strings(obj):
-    if isinstance(obj, dict):
-        return {k: convert_datetimes_to_strings(v) for k, v in obj.items()}
-    elif isinstance(obj, list):
-        return [convert_datetimes_to_strings(item) for item in obj]
-    elif isinstance(obj, tuple):
-        return tuple(convert_datetimes_to_strings(item) for item in obj)
-    elif isinstance(obj, datetime):
-        return obj.isoformat()
-    else:
-        return obj
+        return full_query