graphiti-core 0.15.1__py3-none-any.whl → 0.17.0__py3-none-any.whl

graphiti_core/cross_encoder/gemini_reranker_client.py

@@ -41,6 +41,10 @@ DEFAULT_MODEL = 'gemini-2.5-flash-lite-preview-06-17'
 
 
 class GeminiRerankerClient(CrossEncoderClient):
+    """
+    Google Gemini Reranker Client
+    """
+
     def __init__(
         self,
         config: LLMConfig | None = None,
@@ -57,7 +61,6 @@ class GeminiRerankerClient(CrossEncoderClient):
             config (LLMConfig | None): The configuration for the LLM client, including API key, model, base URL, temperature, and max tokens.
             client (genai.Client | None): An optional async client instance to use. If not provided, a new genai.Client is created.
         """
-
         if config is None:
             config = LLMConfig()
 

graphiti_core/driver/driver.py

@@ -19,8 +19,6 @@ from abc import ABC, abstractmethod
 from collections.abc import Coroutine
 from typing import Any
 
-from graphiti_core.helpers import DEFAULT_DATABASE
-
 logger = logging.getLogger(__name__)
 
 
@@ -54,7 +52,7 @@ class GraphDriver(ABC):
         raise NotImplementedError()
 
     @abstractmethod
-    def session(self, database: str) -> GraphDriverSession:
+    def session(self, database: str | None = None) -> GraphDriverSession:
         raise NotImplementedError()
 
     @abstractmethod
@@ -62,5 +60,5 @@ class GraphDriver(ABC):
         raise NotImplementedError()
 
     @abstractmethod
-    def delete_all_indexes(self, database_: str = DEFAULT_DATABASE) -> Coroutine:
+    def delete_all_indexes(self, database_: str | None = None) -> Coroutine:
         raise NotImplementedError()

graphiti_core/driver/falkordb_driver.py

@@ -33,7 +33,6 @@ else:
         ) from None
 
 from graphiti_core.driver.driver import GraphDriver, GraphDriverSession
-from graphiti_core.helpers import DEFAULT_DATABASE
 
 logger = logging.getLogger(__name__)
 
@@ -81,6 +80,7 @@ class FalkorDriver(GraphDriver):
         username: str | None = None,
         password: str | None = None,
         falkor_db: FalkorDB | None = None,
+        database: str = 'default_db',
     ):
         """
         Initialize the FalkorDB driver.
@@ -95,15 +95,16 @@ class FalkorDriver(GraphDriver):
             self.client = falkor_db
         else:
             self.client = FalkorDB(host=host, port=port, username=username, password=password)
+            self._database = database
 
     def _get_graph(self, graph_name: str | None) -> FalkorGraph:
-        # FalkorDB requires a non-None database name for multi-tenant graphs; the default is DEFAULT_DATABASE
+        # FalkorDB requires a non-None database name for multi-tenant graphs; the default is "default_db"
         if graph_name is None:
-            graph_name = DEFAULT_DATABASE
+            graph_name = self._database
         return self.client.select_graph(graph_name)
 
     async def execute_query(self, cypher_query_, **kwargs: Any):
-        graph_name = kwargs.pop('database_', DEFAULT_DATABASE)
+        graph_name = kwargs.pop('database_', self._database)
         graph = self._get_graph(graph_name)
 
         # Convert datetime objects to ISO strings (FalkorDB does not support datetime objects directly)
@@ -136,7 +137,7 @@ class FalkorDriver(GraphDriver):
 
         return records, header, None
 
-    def session(self, database: str | None) -> GraphDriverSession:
+    def session(self, database: str | None = None) -> GraphDriverSession:
         return FalkorDriverSession(self._get_graph(database))
 
     async def close(self) -> None:
@@ -148,10 +149,11 @@ class FalkorDriver(GraphDriver):
         elif hasattr(self.client.connection, 'close'):
             await self.client.connection.close()
 
-    async def delete_all_indexes(self, database_: str = DEFAULT_DATABASE) -> None:
+    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_,
+            database_=database,
         )
 
 

graphiti_core/driver/neo4j_driver.py

@@ -22,7 +22,6 @@ from neo4j import AsyncGraphDatabase, EagerResult
 from typing_extensions import LiteralString
 
 from graphiti_core.driver.driver import GraphDriver, GraphDriverSession
-from graphiti_core.helpers import DEFAULT_DATABASE
 
 logger = logging.getLogger(__name__)
 
@@ -30,34 +29,36 @@ logger = logging.getLogger(__name__)
 class Neo4jDriver(GraphDriver):
     provider: str = 'neo4j'
 
-    def __init__(
-        self,
-        uri: str,
-        user: str | None,
-        password: str | None,
-    ):
+    def __init__(self, uri: str, user: str | None, password: str | None, database: str = 'neo4j'):
         super().__init__()
         self.client = AsyncGraphDatabase.driver(
             uri=uri,
             auth=(user or '', password or ''),
         )
+        self._database = database
 
     async def execute_query(self, cypher_query_: LiteralString, **kwargs: Any) -> EagerResult:
+        # Check if database_ is provided in kwargs.
+        # If not populated, set the value to retain backwards compatibility
         params = kwargs.pop('params', None)
+        if params is None:
+            params = {}
+        params.setdefault('database_', self._database)
+
         result = await self.client.execute_query(cypher_query_, parameters_=params, **kwargs)
 
         return result
 
-    def session(self, database: str) -> GraphDriverSession:
-        return self.client.session(database=database)  # type: ignore
+    def session(self, database: str | None = None) -> GraphDriverSession:
+        _database = database or self._database
+        return self.client.session(database=_database)  # type: ignore
 
     async def close(self) -> None:
         return await self.client.close()
 
-    def delete_all_indexes(
-        self, database_: str = DEFAULT_DATABASE
-    ) -> Coroutine[Any, Any, EagerResult]:
+    def delete_all_indexes(self, database_: str | None = None) -> Coroutine[Any, Any, EagerResult]:
+        database = database_ or self._database
         return self.client.execute_query(
             'CALL db.indexes() YIELD name DROP INDEX name',
-            database_=database_,
+            database_=database,
         )

graphiti_core/edges.py

@@ -27,7 +27,7 @@ from typing_extensions import LiteralString
 from graphiti_core.driver.driver import GraphDriver
 from graphiti_core.embedder import EmbedderClient
 from graphiti_core.errors import EdgeNotFoundError, GroupsEdgesNotFoundError
-from graphiti_core.helpers import DEFAULT_DATABASE, parse_db_date
+from graphiti_core.helpers import parse_db_date
 from graphiti_core.models.edges.edge_db_queries import (
     COMMUNITY_EDGE_SAVE,
     ENTITY_EDGE_SAVE,
@@ -71,7 +71,6 @@ class Edge(BaseModel, ABC):
         DELETE e
         """,
             uuid=self.uuid,
-            database_=DEFAULT_DATABASE,
         )
 
         logger.debug(f'Deleted Edge: {self.uuid}')
@@ -99,7 +98,6 @@ class EpisodicEdge(Edge):
             uuid=self.uuid,
             group_id=self.group_id,
             created_at=self.created_at,
-            database_=DEFAULT_DATABASE,
         )
 
         logger.debug(f'Saved edge to Graph: {self.uuid}')
@@ -119,7 +117,6 @@ class EpisodicEdge(Edge):
             e.created_at AS created_at
         """,
             uuid=uuid,
-            database_=DEFAULT_DATABASE,
             routing_='r',
         )
 
@@ -143,7 +140,6 @@ class EpisodicEdge(Edge):
             e.created_at AS created_at
         """,
             uuids=uuids,
-            database_=DEFAULT_DATABASE,
             routing_='r',
         )
 
@@ -183,7 +179,6 @@ class EpisodicEdge(Edge):
             group_ids=group_ids,
             uuid=uuid_cursor,
             limit=limit,
-            database_=DEFAULT_DATABASE,
             routing_='r',
         )
 
@@ -231,9 +226,7 @@ class EntityEdge(Edge):
             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, database_=DEFAULT_DATABASE, routing_='r'
-        )
+        records, _, _ = await driver.execute_query(query, uuid=self.uuid, routing_='r')
 
         if len(records) == 0:
             raise EdgeNotFoundError(self.uuid)
@@ -261,7 +254,6 @@ class EntityEdge(Edge):
         result = await driver.execute_query(
             ENTITY_EDGE_SAVE,
             edge_data=edge_data,
-            database_=DEFAULT_DATABASE,
         )
 
         logger.debug(f'Saved edge to Graph: {self.uuid}')
@@ -276,7 +268,6 @@ class EntityEdge(Edge):
         """
             + ENTITY_EDGE_RETURN,
             uuid=uuid,
-            database_=DEFAULT_DATABASE,
             routing_='r',
         )
 
@@ -298,7 +289,6 @@ class EntityEdge(Edge):
         """
             + ENTITY_EDGE_RETURN,
             uuids=uuids,
-            database_=DEFAULT_DATABASE,
             routing_='r',
         )
 
@@ -331,7 +321,6 @@ class EntityEdge(Edge):
             group_ids=group_ids,
             uuid=uuid_cursor,
             limit=limit,
-            database_=DEFAULT_DATABASE,
             routing_='r',
         )
 
@@ -349,9 +338,7 @@ class EntityEdge(Edge):
                                                         """
             + ENTITY_EDGE_RETURN
         )
-        records, _, _ = await driver.execute_query(
-            query, node_uuid=node_uuid, database_=DEFAULT_DATABASE, routing_='r'
-        )
+        records, _, _ = await driver.execute_query(query, node_uuid=node_uuid, routing_='r')
 
         edges = [get_entity_edge_from_record(record) for record in records]
 
@@ -367,7 +354,6 @@ class CommunityEdge(Edge):
             uuid=self.uuid,
             group_id=self.group_id,
             created_at=self.created_at,
-            database_=DEFAULT_DATABASE,
         )
 
         logger.debug(f'Saved edge to Graph: {self.uuid}')
@@ -387,7 +373,6 @@ class CommunityEdge(Edge):
             e.created_at AS created_at
         """,
             uuid=uuid,
-            database_=DEFAULT_DATABASE,
             routing_='r',
         )
 
@@ -409,7 +394,6 @@ class CommunityEdge(Edge):
             e.created_at AS created_at
         """,
             uuids=uuids,
-            database_=DEFAULT_DATABASE,
             routing_='r',
         )
 
@@ -447,7 +431,6 @@ class CommunityEdge(Edge):
             group_ids=group_ids,
             uuid=uuid_cursor,
             limit=limit,
-            database_=DEFAULT_DATABASE,
             routing_='r',
         )
 

graphiti_core/embedder/gemini.py

@@ -47,15 +47,27 @@ class GeminiEmbedder(EmbedderClient):
     Google Gemini Embedder Client
     """
 
-    def __init__(self, config: GeminiEmbedderConfig | None = None):
+    def __init__(
+        self,
+        config: GeminiEmbedderConfig | None = None,
+        client: 'genai.Client | None' = None,
+    ):
+        """
+        Initialize the GeminiEmbedder with the provided configuration and client.
+
+        Args:
+            config (GeminiEmbedderConfig | None): The configuration for the GeminiEmbedder, including API key, model, base URL, temperature, and max tokens.
+            client (genai.Client | None): An optional async client instance to use. If not provided, a new genai.Client is created.
+        """
         if config is None:
             config = GeminiEmbedderConfig()
+
         self.config = config
 
-        # Configure the Gemini API
-        self.client = genai.Client(
-            api_key=config.api_key,
-        )
+        if client is None:
+            self.client = genai.Client(api_key=config.api_key)
+        else:
+            self.client = client
 
     async def create(
         self, input_data: str | list[str] | Iterable[int] | Iterable[Iterable[int]]

graphiti_core/graphiti.py

@@ -30,7 +30,6 @@ from graphiti_core.edges import EntityEdge, EpisodicEdge
 from graphiti_core.embedder import EmbedderClient, OpenAIEmbedder
 from graphiti_core.graphiti_types import GraphitiClients
 from graphiti_core.helpers import (
-    DEFAULT_DATABASE,
     semaphore_gather,
     validate_excluded_entity_types,
     validate_group_id,
@@ -57,7 +56,6 @@ from graphiti_core.utils.bulk_utils import (
     add_nodes_and_edges_bulk,
     dedupe_edges_bulk,
     dedupe_nodes_bulk,
-    extract_edge_dates_bulk,
     extract_nodes_and_edges_bulk,
     resolve_edge_pointers,
     retrieve_previous_episodes_bulk,
@@ -169,7 +167,6 @@ class Graphiti:
                 raise ValueError('uri must be provided when graph_driver is None')
             self.driver = Neo4jDriver(uri, user, password)
 
-        self.database = DEFAULT_DATABASE
         self.store_raw_episode_content = store_raw_episode_content
         self.max_coroutines = max_coroutines
         if llm_client:
@@ -508,7 +505,7 @@ class Graphiti:
 
             entity_edges = resolved_edges + invalidated_edges + duplicate_of_edges
 
-            episodic_edges = build_episodic_edges(nodes, episode, now)
+            episodic_edges = build_episodic_edges(nodes, episode.uuid, now)
 
             episode.entity_edges = [edge.uuid for edge in entity_edges]
 
@@ -536,8 +533,16 @@ class Graphiti:
         except Exception as e:
             raise e
 
-    #### WIP: USE AT YOUR OWN RISK ####
-    async def add_episode_bulk(self, bulk_episodes: list[RawEpisode], group_id: str = ''):
+    ##### EXPERIMENTAL #####
+    async def add_episode_bulk(
+        self,
+        bulk_episodes: list[RawEpisode],
+        group_id: str = '',
+        entity_types: dict[str, BaseModel] | None = None,
+        excluded_entity_types: list[str] | None = None,
+        edge_types: dict[str, BaseModel] | None = None,
+        edge_type_map: dict[tuple[str, str], list[str]] | None = None,
+    ):
         """
         Process multiple episodes in bulk and update the graph.
 
@@ -580,8 +585,17 @@ class Graphiti:
 
             validate_group_id(group_id)
 
+            # Create default edge type map
+            edge_type_map_default = (
+                {('Entity', 'Entity'): list(edge_types.keys())}
+                if edge_types is not None
+                else {('Entity', 'Entity'): []}
+            )
+
             episodes = [
-                EpisodicNode(
+                await EpisodicNode.get_by_uuid(self.driver, episode.uuid)
+                if episode.uuid is not None
+                else EpisodicNode(
                     name=episode.name,
                     labels=[],
                     source=episode.source,
@@ -594,68 +608,106 @@ class Graphiti:
                 for episode in bulk_episodes
             ]
 
-            # Save all the episodes
-            await semaphore_gather(
-                *[episode.save(self.driver) for episode in episodes],
-                max_coroutines=self.max_coroutines,
+            episodes_by_uuid: dict[str, EpisodicNode] = {
+                episode.uuid: episode for episode in episodes
+            }
+
+            # Save all episodes
+            await add_nodes_and_edges_bulk(
+                driver=self.driver,
+                episodic_nodes=episodes,
+                episodic_edges=[],
+                entity_nodes=[],
+                entity_edges=[],
+                embedder=self.embedder,
             )
 
             # Get previous episode context for each episode
-            episode_pairs = await retrieve_previous_episodes_bulk(self.driver, episodes)
+            episode_context = await retrieve_previous_episodes_bulk(self.driver, episodes)
 
-            # Extract all nodes and edges
-            (
-                extracted_nodes,
-                extracted_edges,
-                episodic_edges,
-            ) = await extract_nodes_and_edges_bulk(self.clients, episode_pairs, None, None)
-
-            # Generate embeddings
-            await semaphore_gather(
-                *[node.generate_name_embedding(self.embedder) for node in extracted_nodes],
-                *[edge.generate_embedding(self.embedder) for edge in extracted_edges],
-                max_coroutines=self.max_coroutines,
+            # Extract all nodes and edges for each episode
+            extracted_nodes_bulk, extracted_edges_bulk = await extract_nodes_and_edges_bulk(
+                self.clients,
+                episode_context,
+                edge_type_map=edge_type_map or edge_type_map_default,
+                edge_types=edge_types,
+                entity_types=entity_types,
+                excluded_entity_types=excluded_entity_types,
             )
 
-            # Dedupe extracted nodes, compress extracted edges
-            (nodes, uuid_map), extracted_edges_timestamped = await semaphore_gather(
-                dedupe_nodes_bulk(self.driver, self.llm_client, extracted_nodes),
-                extract_edge_dates_bulk(self.llm_client, extracted_edges, episode_pairs),
-                max_coroutines=self.max_coroutines,
+            # Dedupe extracted nodes in memory
+            nodes_by_episode, uuid_map = await dedupe_nodes_bulk(
+                self.clients, extracted_nodes_bulk, episode_context, entity_types
             )
 
-            # save nodes to KG
-            await semaphore_gather(
-                *[node.save(self.driver) for node in nodes],
-                max_coroutines=self.max_coroutines,
-            )
+            episodic_edges: list[EpisodicEdge] = []
+            for episode_uuid, nodes in nodes_by_episode.items():
+                episodic_edges.extend(build_episodic_edges(nodes, episode_uuid, now))
 
             # re-map edge pointers so that they don't point to discard dupe nodes
-            extracted_edges_with_resolved_pointers: list[EntityEdge] = resolve_edge_pointers(
-                extracted_edges_timestamped, uuid_map
-            )
-            episodic_edges_with_resolved_pointers: list[EpisodicEdge] = resolve_edge_pointers(
-                episodic_edges, uuid_map
-            )
+            extracted_edges_bulk_updated: list[list[EntityEdge]] = [
+                resolve_edge_pointers(edges, uuid_map) for edges in extracted_edges_bulk
+            ]
 
-            # save episodic edges to KG
-            await semaphore_gather(
-                *[edge.save(self.driver) for edge in episodic_edges_with_resolved_pointers],
-                max_coroutines=self.max_coroutines,
+            # Dedupe extracted edges in memory
+            edges_by_episode = await dedupe_edges_bulk(
+                self.clients,
+                extracted_edges_bulk_updated,
+                episode_context,
+                [],
+                edge_types or {},
+                edge_type_map or edge_type_map_default,
             )
 
-            # Dedupe extracted edges
-            edges = await dedupe_edges_bulk(
-                self.driver, self.llm_client, extracted_edges_with_resolved_pointers
+            # Extract node attributes
+            nodes_by_uuid: dict[str, EntityNode] = {
+                node.uuid: node for nodes in nodes_by_episode.values() for node in nodes
+            }
+
+            extract_attributes_params: list[tuple[EntityNode, list[EpisodicNode]]] = []
+            for node in nodes_by_uuid.values():
+                episode_uuids: list[str] = []
+                for episode_uuid, mentioned_nodes in nodes_by_episode.items():
+                    for mentioned_node in mentioned_nodes:
+                        if node.uuid == mentioned_node.uuid:
+                            episode_uuids.append(episode_uuid)
+                            break
+
+                episode_mentions: list[EpisodicNode] = [
+                    episodes_by_uuid[episode_uuid] for episode_uuid in episode_uuids
+                ]
+                episode_mentions.sort(key=lambda x: x.valid_at, reverse=True)
+
+                extract_attributes_params.append((node, episode_mentions))
+
+            new_hydrated_nodes: list[list[EntityNode]] = await semaphore_gather(
+                *[
+                    extract_attributes_from_nodes(
+                        self.clients,
+                        [params[0]],
+                        params[1][0],
+                        params[1][0:],
+                        entity_types,
+                    )
+                    for params in extract_attributes_params
+                ]
             )
-            logger.debug(f'extracted edge length: {len(edges)}')
 
-            # invalidate edges
+            hydrated_nodes = [node for nodes in new_hydrated_nodes for node in nodes]
 
-            # save edges to KG
-            await semaphore_gather(
-                *[edge.save(self.driver) for edge in edges],
-                max_coroutines=self.max_coroutines,
+            # TODO: Resolve nodes and edges against the existing graph
+            edges_by_uuid: dict[str, EntityEdge] = {
+                edge.uuid: edge for edges in edges_by_episode.values() for edge in edges
+            }
+
+            # save data to KG
+            await add_nodes_and_edges_bulk(
+                self.driver,
+                episodes,
+                episodic_edges,
+                hydrated_nodes,
+                list(edges_by_uuid.values()),
+                self.embedder,
             )
 
             end = time()
@@ -828,7 +880,7 @@ class Graphiti:
             await get_edge_invalidation_candidates(self.driver, [updated_edge], SearchFilters())
         )[0]
 
-        resolved_edge, invalidated_edges = await resolve_extracted_edge(
+        resolved_edge, invalidated_edges, _ = await resolve_extracted_edge(
             self.llm_client,
             updated_edge,
             related_edges,
@@ -867,9 +919,7 @@ class Graphiti:
         nodes_to_delete: list[EntityNode] = []
         for node in nodes:
             query: LiteralString = 'MATCH (e:Episodic)-[:MENTIONS]->(n:Entity {uuid: $uuid}) RETURN count(*) AS episode_count'
-            records, _, _ = await self.driver.execute_query(
-                query, uuid=node.uuid, database_=DEFAULT_DATABASE, routing_='r'
-            )
+            records, _, _ = await self.driver.execute_query(query, uuid=node.uuid, routing_='r')
 
             for record in records:
                 if record['episode_count'] == 1:

graphiti_core/helpers.py

@@ -32,7 +32,6 @@ from graphiti_core.errors import GroupIdValidationError
 
 load_dotenv()
 
-DEFAULT_DATABASE = os.getenv('DEFAULT_DATABASE', 'default_db')
 USE_PARALLEL_RUNTIME = bool(os.getenv('USE_PARALLEL_RUNTIME', False))
 SEMAPHORE_LIMIT = int(os.getenv('SEMAPHORE_LIMIT', 20))
 MAX_REFLEXION_ITERATIONS = int(os.getenv('MAX_REFLEXION_ITERATIONS', 0))

graphiti_core/llm_client/gemini_client.py

@@ -76,6 +76,7 @@ class GeminiClient(LLMClient):
         cache: bool = False,
         max_tokens: int = DEFAULT_MAX_TOKENS,
         thinking_config: types.ThinkingConfig | None = None,
+        client: 'genai.Client | None' = None,
     ):
         """
         Initialize the GeminiClient with the provided configuration, cache setting, and optional thinking config.
@@ -85,7 +86,7 @@ class GeminiClient(LLMClient):
             cache (bool): Whether to use caching for responses. Defaults to False.
             thinking_config (types.ThinkingConfig | None): Optional thinking configuration for models that support it.
                 Only use with models that support thinking (gemini-2.5+). Defaults to None.
-
+            client (genai.Client | None): An optional async client instance to use. If not provided, a new genai.Client is created.
         """
         if config is None:
             config = LLMConfig()
@@ -93,10 +94,12 @@ class GeminiClient(LLMClient):
         super().__init__(config, cache)
 
         self.model = config.model
-        # Configure the Gemini API
-        self.client = genai.Client(
-            api_key=config.api_key,
-        )
+
+        if client is None:
+            self.client = genai.Client(api_key=config.api_key)
+        else:
+            self.client = client
+
         self.max_tokens = max_tokens
         self.thinking_config = thinking_config