graphiti-core 0.17.11__py3-none-any.whl → 0.18.1__py3-none-any.whl

graphiti_core/driver/driver.py

@@ -14,14 +14,21 @@ See the License for the specific language governing permissions and
 limitations under the License.
 """
 
+import copy
 import logging
 from abc import ABC, abstractmethod
 from collections.abc import Coroutine
+from enum import Enum
 from typing import Any
 
 logger = logging.getLogger(__name__)
 
 
+class GraphProvider(Enum):
+    NEO4J = 'neo4j'
+    FALKORDB = 'falkordb'
+
+
 class GraphDriverSession(ABC):
     async def __aenter__(self):
         return self
@@ -45,10 +52,11 @@ 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
 
     @abstractmethod
     def execute_query(self, cypher_query_: str, **kwargs: Any) -> Coroutine:
@@ -63,5 +71,15 @@ 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

graphiti_core/driver/falkordb_driver.py

@@ -32,7 +32,7 @@ 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
 
 logger = logging.getLogger(__name__)
 
@@ -71,7 +71,7 @@ class FalkorDriverSession(GraphDriverSession):
 
 
 class FalkorDriver(GraphDriver):
-    provider: str = 'falkordb'
+    provider = GraphProvider.FALKORDB
 
     def __init__(
         self,
@@ -90,12 +90,13 @@ class FalkorDriver(GraphDriver):
         The default parameters assume a local (on-premises) FalkorDB instance.
         """
         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
 
         self.fulltext_syntax = '@'  # FalkorDB uses a redisearch-like syntax for fulltext queries see https://redis.io/docs/latest/develop/ai/search-and-query/query/full-text/
 
@@ -106,8 +107,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))
@@ -119,7 +119,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
@@ -151,13 +151,20 @@ 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
+    async def delete_all_indexes(self) -> None:
         await self.execute_query(
             'CALL db.indexes() YIELD name DROP INDEX name',
-            database_=database,
         )
 
+    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).
+        """
+        cloned = FalkorDriver(falkor_db=self.client, database=database)
+
+        return cloned
+
 
 def convert_datetimes_to_strings(obj):
     if isinstance(obj, dict):

graphiti_core/driver/neo4j_driver.py

@@ -21,13 +21,13 @@ from typing import Any
 from neo4j import AsyncGraphDatabase, EagerResult
 from typing_extensions import LiteralString
 
-from graphiti_core.driver.driver import GraphDriver, GraphDriverSession
+from graphiti_core.driver.driver import GraphDriver, GraphDriverSession, GraphProvider
 
 logger = logging.getLogger(__name__)
 
 
 class Neo4jDriver(GraphDriver):
-    provider: str = 'neo4j'
+    provider = GraphProvider.NEO4J
 
     def __init__(self, uri: str, user: str | None, password: str | None, database: str = 'neo4j'):
         super().__init__()
@@ -45,7 +45,11 @@ class Neo4jDriver(GraphDriver):
             params = {}
         params.setdefault('database_', self._database)
 
-        result = await self.client.execute_query(cypher_query_, parameters_=params, **kwargs)
+        try:
+            result = await self.client.execute_query(cypher_query_, parameters_=params, **kwargs)
+        except Exception as e:
+            logger.error(f'Error executing Neo4j query: {e}\n{cypher_query_}\n{params}')
+            raise
 
         return result
 
@@ -56,9 +60,7 @@ class Neo4jDriver(GraphDriver):
     async def close(self) -> None:
         return await self.client.close()
 
-    def delete_all_indexes(self, database_: str | None = None) -> Coroutine[Any, Any, EagerResult]:
-        database = database_ or self._database
+    def delete_all_indexes(self) -> Coroutine[Any, Any, EagerResult]:
         return self.client.execute_query(
             'CALL db.indexes() YIELD name DROP INDEX name',
-            database_=database,
         )

graphiti_core/edges.py

@@ -29,29 +29,17 @@ 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_SAVE,
-    ENTITY_EDGE_SAVE,
+    COMMUNITY_EDGE_RETURN,
+    ENTITY_EDGE_RETURN,
+    EPISODIC_EDGE_RETURN,
     EPISODIC_EDGE_SAVE,
+    get_community_edge_save_query,
+    get_entity_edge_save_query,
 )
 from graphiti_core.nodes import Node
 
 logger = logging.getLogger(__name__)
 
-ENTITY_EDGE_RETURN: LiteralString = """
-        RETURN
-            e.uuid AS uuid,
-            startNode(e).uuid AS source_node_uuid,
-            endNode(e).uuid AS target_node_uuid,
-            e.created_at AS created_at,
-            e.name AS name,
-            e.group_id AS group_id,
-            e.fact AS fact,
-            e.episodes AS episodes,
-            e.expired_at AS expired_at,
-            e.valid_at AS valid_at,
-            e.invalid_at AS invalid_at,
-            properties(e) AS attributes"""
-
 
 class Edge(BaseModel, ABC):
     uuid: str = Field(default_factory=lambda: str(uuid4()))
@@ -66,9 +54,9 @@ class Edge(BaseModel, ABC):
     async def delete(self, driver: GraphDriver):
         result = await driver.execute_query(
             """
-        MATCH (n)-[e:MENTIONS|RELATES_TO|HAS_MEMBER {uuid: $uuid}]->(m)
-        DELETE e
-        """,
+            MATCH (n)-[e:MENTIONS|RELATES_TO|HAS_MEMBER {uuid: $uuid}]->(m)
+            DELETE e
+            """,
             uuid=self.uuid,
         )
 
@@ -107,14 +95,10 @@ class EpisodicEdge(Edge):
     async def get_by_uuid(cls, driver: GraphDriver, uuid: str):
         records, _, _ = await driver.execute_query(
             """
-        MATCH (n:Episodic)-[e:MENTIONS {uuid: $uuid}]->(m:Entity)
-        RETURN
-            e.uuid As uuid,
-            e.group_id AS group_id,
-            n.uuid AS source_node_uuid, 
-            m.uuid AS target_node_uuid, 
-            e.created_at AS created_at
-        """,
+            MATCH (n:Episodic)-[e:MENTIONS {uuid: $uuid}]->(m:Entity)
+            RETURN
+            """
+            + EPISODIC_EDGE_RETURN,
             uuid=uuid,
             routing_='r',
         )
@@ -129,15 +113,11 @@ class EpisodicEdge(Edge):
     async def get_by_uuids(cls, driver: GraphDriver, uuids: list[str]):
         records, _, _ = await driver.execute_query(
             """
-        MATCH (n:Episodic)-[e:MENTIONS]->(m:Entity)
-        WHERE e.uuid IN $uuids
-        RETURN
-            e.uuid As uuid,
-            e.group_id AS group_id,
-            n.uuid AS source_node_uuid, 
-            m.uuid AS target_node_uuid, 
-            e.created_at AS created_at
-        """,
+            MATCH (n:Episodic)-[e:MENTIONS]->(m:Entity)
+            WHERE e.uuid IN $uuids
+            RETURN
+            """
+            + EPISODIC_EDGE_RETURN,
             uuids=uuids,
             routing_='r',
         )
@@ -161,19 +141,17 @@ class EpisodicEdge(Edge):
 
         records, _, _ = await driver.execute_query(
             """
-        MATCH (n:Episodic)-[e:MENTIONS]->(m:Entity)
-        WHERE e.group_id IN $group_ids
-        """
+            MATCH (n:Episodic)-[e:MENTIONS]->(m:Entity)
+            WHERE e.group_id IN $group_ids
+            """
             + cursor_query
             + """
-        RETURN
-            e.uuid As uuid,
-            e.group_id AS group_id,
-            n.uuid AS source_node_uuid, 
-            m.uuid AS target_node_uuid, 
-            e.created_at AS created_at
-        ORDER BY e.uuid DESC 
-        """
+            RETURN
+            """
+            + EPISODIC_EDGE_RETURN
+            + """
+            ORDER BY e.uuid DESC 
+            """
             + limit_query,
             group_ids=group_ids,
             uuid=uuid_cursor,
@@ -221,11 +199,14 @@ class EntityEdge(Edge):
         return self.fact_embedding
 
     async def load_fact_embedding(self, driver: GraphDriver):
-        query: LiteralString = """
+        records, _, _ = await driver.execute_query(
+            """
             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')
+            """,
+            uuid=self.uuid,
+            routing_='r',
+        )
 
         if len(records) == 0:
             raise EdgeNotFoundError(self.uuid)
@@ -251,7 +232,7 @@ class EntityEdge(Edge):
         edge_data.update(self.attributes or {})
 
         result = await driver.execute_query(
-            ENTITY_EDGE_SAVE,
+            get_entity_edge_save_query(driver.provider),
             edge_data=edge_data,
         )
 
@@ -263,8 +244,9 @@ class EntityEdge(Edge):
     async def get_by_uuid(cls, driver: GraphDriver, uuid: str):
         records, _, _ = await driver.execute_query(
             """
-        MATCH (n:Entity)-[e:RELATES_TO {uuid: $uuid}]->(m:Entity)
-        """
+            MATCH (n:Entity)-[e:RELATES_TO {uuid: $uuid}]->(m:Entity)
+            RETURN
+            """
             + ENTITY_EDGE_RETURN,
             uuid=uuid,
             routing_='r',
@@ -283,9 +265,10 @@ class EntityEdge(Edge):
 
         records, _, _ = await driver.execute_query(
             """
-        MATCH (n:Entity)-[e:RELATES_TO]->(m:Entity)
-        WHERE e.uuid IN $uuids
-        """
+            MATCH (n:Entity)-[e:RELATES_TO]->(m:Entity)
+            WHERE e.uuid IN $uuids
+            RETURN
+            """
             + ENTITY_EDGE_RETURN,
             uuids=uuids,
             routing_='r',
@@ -314,22 +297,21 @@ class EntityEdge(Edge):
             else ''
         )
 
-        query: LiteralString = (
+        records, _, _ = await driver.execute_query(
             """
             MATCH (n:Entity)-[e:RELATES_TO]->(m:Entity)
             WHERE e.group_id IN $group_ids
             """
             + cursor_query
+            + """
+            RETURN
+            """
             + ENTITY_EDGE_RETURN
             + with_embeddings_query
             + """
-        ORDER BY e.uuid DESC 
-        """
-            + limit_query
-        )
-
-        records, _, _ = await driver.execute_query(
-            query,
+            ORDER BY e.uuid DESC 
+            """
+            + limit_query,
             group_ids=group_ids,
             uuid=uuid_cursor,
             limit=limit,
@@ -344,13 +326,15 @@ class EntityEdge(Edge):
 
     @classmethod
     async def get_by_node_uuid(cls, driver: GraphDriver, node_uuid: str):
-        query: LiteralString = (
+        records, _, _ = await driver.execute_query(
             """
-                                                                    MATCH (n:Entity {uuid: $node_uuid})-[e:RELATES_TO]-(m:Entity)
-                                                                    """
-            + ENTITY_EDGE_RETURN
+            MATCH (n:Entity {uuid: $node_uuid})-[e:RELATES_TO]-(m:Entity)
+            RETURN
+            """
+            + ENTITY_EDGE_RETURN,
+            node_uuid=node_uuid,
+            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]
 
@@ -360,7 +344,7 @@ class EntityEdge(Edge):
 class CommunityEdge(Edge):
     async def save(self, driver: GraphDriver):
         result = await driver.execute_query(
-            COMMUNITY_EDGE_SAVE,
+            get_community_edge_save_query(driver.provider),
             community_uuid=self.source_node_uuid,
             entity_uuid=self.target_node_uuid,
             uuid=self.uuid,
@@ -376,14 +360,10 @@ class CommunityEdge(Edge):
     async def get_by_uuid(cls, driver: GraphDriver, uuid: str):
         records, _, _ = await driver.execute_query(
             """
-        MATCH (n:Community)-[e:HAS_MEMBER {uuid: $uuid}]->(m:Entity | Community)
-        RETURN
-            e.uuid As uuid,
-            e.group_id AS group_id,
-            n.uuid AS source_node_uuid, 
-            m.uuid AS target_node_uuid, 
-            e.created_at AS created_at
-        """,
+            MATCH (n:Community)-[e:HAS_MEMBER {uuid: $uuid}]->(m)
+            RETURN
+            """
+            + COMMUNITY_EDGE_RETURN,
             uuid=uuid,
             routing_='r',
         )
@@ -396,15 +376,11 @@ class CommunityEdge(Edge):
     async def get_by_uuids(cls, driver: GraphDriver, uuids: list[str]):
         records, _, _ = await driver.execute_query(
             """
-        MATCH (n:Community)-[e:HAS_MEMBER]->(m:Entity | Community)
-        WHERE e.uuid IN $uuids
-        RETURN
-            e.uuid As uuid,
-            e.group_id AS group_id,
-            n.uuid AS source_node_uuid, 
-            m.uuid AS target_node_uuid, 
-            e.created_at AS created_at
-        """,
+            MATCH (n:Community)-[e:HAS_MEMBER]->(m)
+            WHERE e.uuid IN $uuids
+            RETURN
+            """
+            + COMMUNITY_EDGE_RETURN,
             uuids=uuids,
             routing_='r',
         )
@@ -426,19 +402,17 @@ class CommunityEdge(Edge):
 
         records, _, _ = await driver.execute_query(
             """
-        MATCH (n:Community)-[e:HAS_MEMBER]->(m:Entity | Community)
-        WHERE e.group_id IN $group_ids
-        """
+            MATCH (n:Community)-[e:HAS_MEMBER]->(m)
+            WHERE e.group_id IN $group_ids
+            """
             + cursor_query
             + """
-        RETURN
-            e.uuid As uuid,
-            e.group_id AS group_id,
-            n.uuid AS source_node_uuid, 
-            m.uuid AS target_node_uuid, 
-            e.created_at AS created_at
-        ORDER BY e.uuid DESC
-        """
+            RETURN
+            """
+            + COMMUNITY_EDGE_RETURN
+            + """
+            ORDER BY e.uuid DESC
+            """
             + limit_query,
             group_ids=group_ids,
             uuid=uuid_cursor,

graphiti_core/graph_queries.py

@@ -5,16 +5,9 @@ This module provides database-agnostic query generation for Neo4j and FalkorDB,
 supporting index creation, fulltext search, and bulk operations.
 """
 
-from typing import Any
-
 from typing_extensions import LiteralString
 
-from graphiti_core.models.edges.edge_db_queries import (
-    ENTITY_EDGE_SAVE_BULK,
-)
-from graphiti_core.models.nodes.node_db_queries import (
-    ENTITY_NODE_SAVE_BULK,
-)
+from graphiti_core.driver.driver import GraphProvider
 
 # Mapping from Neo4j fulltext index names to FalkorDB node labels
 NEO4J_TO_FALKORDB_MAPPING = {
@@ -25,8 +18,8 @@ NEO4J_TO_FALKORDB_MAPPING = {
 }
 
 
-def get_range_indices(db_type: str = 'neo4j') -> list[LiteralString]:
-    if db_type == 'falkordb':
+def get_range_indices(provider: GraphProvider) -> list[LiteralString]:
+    if provider == GraphProvider.FALKORDB:
         return [
             # Entity node
             'CREATE INDEX FOR (n:Entity) ON (n.uuid, n.group_id, n.name, n.created_at)',
@@ -41,109 +34,70 @@ def get_range_indices(db_type: str = 'neo4j') -> list[LiteralString]:
             # HAS_MEMBER edge
             'CREATE INDEX FOR ()-[e:HAS_MEMBER]-() ON (e.uuid)',
         ]
-    else:
-        return [
-            'CREATE INDEX entity_uuid IF NOT EXISTS FOR (n:Entity) ON (n.uuid)',
-            'CREATE INDEX episode_uuid IF NOT EXISTS FOR (n:Episodic) ON (n.uuid)',
-            'CREATE INDEX community_uuid IF NOT EXISTS FOR (n:Community) ON (n.uuid)',
-            'CREATE INDEX relation_uuid IF NOT EXISTS FOR ()-[e:RELATES_TO]-() ON (e.uuid)',
-            'CREATE INDEX mention_uuid IF NOT EXISTS FOR ()-[e:MENTIONS]-() ON (e.uuid)',
-            'CREATE INDEX has_member_uuid IF NOT EXISTS FOR ()-[e:HAS_MEMBER]-() ON (e.uuid)',
-            'CREATE INDEX entity_group_id IF NOT EXISTS FOR (n:Entity) ON (n.group_id)',
-            'CREATE INDEX episode_group_id IF NOT EXISTS FOR (n:Episodic) ON (n.group_id)',
-            'CREATE INDEX relation_group_id IF NOT EXISTS FOR ()-[e:RELATES_TO]-() ON (e.group_id)',
-            'CREATE INDEX mention_group_id IF NOT EXISTS FOR ()-[e:MENTIONS]-() ON (e.group_id)',
-            'CREATE INDEX name_entity_index IF NOT EXISTS FOR (n:Entity) ON (n.name)',
-            'CREATE INDEX created_at_entity_index IF NOT EXISTS FOR (n:Entity) ON (n.created_at)',
-            'CREATE INDEX created_at_episodic_index IF NOT EXISTS FOR (n:Episodic) ON (n.created_at)',
-            'CREATE INDEX valid_at_episodic_index IF NOT EXISTS FOR (n:Episodic) ON (n.valid_at)',
-            'CREATE INDEX name_edge_index IF NOT EXISTS FOR ()-[e:RELATES_TO]-() ON (e.name)',
-            'CREATE INDEX created_at_edge_index IF NOT EXISTS FOR ()-[e:RELATES_TO]-() ON (e.created_at)',
-            'CREATE INDEX expired_at_edge_index IF NOT EXISTS FOR ()-[e:RELATES_TO]-() ON (e.expired_at)',
-            'CREATE INDEX valid_at_edge_index IF NOT EXISTS FOR ()-[e:RELATES_TO]-() ON (e.valid_at)',
-            'CREATE INDEX invalid_at_edge_index IF NOT EXISTS FOR ()-[e:RELATES_TO]-() ON (e.invalid_at)',
-        ]
 
-
-def get_fulltext_indices(db_type: str = 'neo4j') -> list[LiteralString]:
-    if db_type == 'falkordb':
+    return [
+        'CREATE INDEX entity_uuid IF NOT EXISTS FOR (n:Entity) ON (n.uuid)',
+        'CREATE INDEX episode_uuid IF NOT EXISTS FOR (n:Episodic) ON (n.uuid)',
+        'CREATE INDEX community_uuid IF NOT EXISTS FOR (n:Community) ON (n.uuid)',
+        'CREATE INDEX relation_uuid IF NOT EXISTS FOR ()-[e:RELATES_TO]-() ON (e.uuid)',
+        'CREATE INDEX mention_uuid IF NOT EXISTS FOR ()-[e:MENTIONS]-() ON (e.uuid)',
+        'CREATE INDEX has_member_uuid IF NOT EXISTS FOR ()-[e:HAS_MEMBER]-() ON (e.uuid)',
+        'CREATE INDEX entity_group_id IF NOT EXISTS FOR (n:Entity) ON (n.group_id)',
+        'CREATE INDEX episode_group_id IF NOT EXISTS FOR (n:Episodic) ON (n.group_id)',
+        'CREATE INDEX relation_group_id IF NOT EXISTS FOR ()-[e:RELATES_TO]-() ON (e.group_id)',
+        'CREATE INDEX mention_group_id IF NOT EXISTS FOR ()-[e:MENTIONS]-() ON (e.group_id)',
+        'CREATE INDEX name_entity_index IF NOT EXISTS FOR (n:Entity) ON (n.name)',
+        'CREATE INDEX created_at_entity_index IF NOT EXISTS FOR (n:Entity) ON (n.created_at)',
+        'CREATE INDEX created_at_episodic_index IF NOT EXISTS FOR (n:Episodic) ON (n.created_at)',
+        'CREATE INDEX valid_at_episodic_index IF NOT EXISTS FOR (n:Episodic) ON (n.valid_at)',
+        'CREATE INDEX name_edge_index IF NOT EXISTS FOR ()-[e:RELATES_TO]-() ON (e.name)',
+        'CREATE INDEX created_at_edge_index IF NOT EXISTS FOR ()-[e:RELATES_TO]-() ON (e.created_at)',
+        'CREATE INDEX expired_at_edge_index IF NOT EXISTS FOR ()-[e:RELATES_TO]-() ON (e.expired_at)',
+        'CREATE INDEX valid_at_edge_index IF NOT EXISTS FOR ()-[e:RELATES_TO]-() ON (e.valid_at)',
+        'CREATE INDEX invalid_at_edge_index IF NOT EXISTS FOR ()-[e:RELATES_TO]-() ON (e.invalid_at)',
+    ]
+
+
+def get_fulltext_indices(provider: GraphProvider) -> list[LiteralString]:
+    if provider == GraphProvider.FALKORDB:
         return [
             """CREATE FULLTEXT INDEX FOR (e:Episodic) ON (e.content, e.source, e.source_description, e.group_id)""",
             """CREATE FULLTEXT INDEX FOR (n:Entity) ON (n.name, n.summary, n.group_id)""",
             """CREATE FULLTEXT INDEX FOR (n:Community) ON (n.name, n.group_id)""",
             """CREATE FULLTEXT INDEX FOR ()-[e:RELATES_TO]-() ON (e.name, e.fact, e.group_id)""",
         ]
-    else:
-        return [
-            """CREATE FULLTEXT INDEX episode_content IF NOT EXISTS 
-            FOR (e:Episodic) ON EACH [e.content, e.source, e.source_description, e.group_id]""",
-            """CREATE FULLTEXT INDEX node_name_and_summary IF NOT EXISTS 
-            FOR (n:Entity) ON EACH [n.name, n.summary, n.group_id]""",
-            """CREATE FULLTEXT INDEX community_name IF NOT EXISTS 
-            FOR (n:Community) ON EACH [n.name, n.group_id]""",
-            """CREATE FULLTEXT INDEX edge_name_and_fact IF NOT EXISTS 
-            FOR ()-[e:RELATES_TO]-() ON EACH [e.name, e.fact, e.group_id]""",
-        ]
+
+    return [
+        """CREATE FULLTEXT INDEX episode_content IF NOT EXISTS
+        FOR (e:Episodic) ON EACH [e.content, e.source, e.source_description, e.group_id]""",
+        """CREATE FULLTEXT INDEX node_name_and_summary IF NOT EXISTS
+        FOR (n:Entity) ON EACH [n.name, n.summary, n.group_id]""",
+        """CREATE FULLTEXT INDEX community_name IF NOT EXISTS
+        FOR (n:Community) ON EACH [n.name, n.group_id]""",
+        """CREATE FULLTEXT INDEX edge_name_and_fact IF NOT EXISTS
+        FOR ()-[e:RELATES_TO]-() ON EACH [e.name, e.fact, e.group_id]""",
+    ]
 
 
-def get_nodes_query(db_type: str = 'neo4j', name: str = '', query: str | None = None) -> str:
-    if db_type == 'falkordb':
+def get_nodes_query(provider: GraphProvider, name: str = '', query: str | None = None) -> str:
+    if provider == GraphProvider.FALKORDB:
         label = NEO4J_TO_FALKORDB_MAPPING[name]
         return f"CALL db.idx.fulltext.queryNodes('{label}', {query})"
-    else:
-        return f'CALL db.index.fulltext.queryNodes("{name}", {query}, {{limit: $limit}})'
+
+    return f'CALL db.index.fulltext.queryNodes("{name}", {query}, {{limit: $limit}})'
 
 
-def get_vector_cosine_func_query(vec1, vec2, db_type: str = 'neo4j') -> str:
-    if db_type == 'falkordb':
+def get_vector_cosine_func_query(vec1, vec2, provider: GraphProvider) -> str:
+    if provider == GraphProvider.FALKORDB:
         # FalkorDB uses a different syntax for regular cosine similarity and Neo4j uses normalized cosine similarity
         return f'(2 - vec.cosineDistance({vec1}, vecf32({vec2})))/2'
-    else:
-        return f'vector.similarity.cosine({vec1}, {vec2})'
 
+    return f'vector.similarity.cosine({vec1}, {vec2})'
 
-def get_relationships_query(name: str, db_type: str = 'neo4j') -> str:
-    if db_type == 'falkordb':
+
+def get_relationships_query(name: str, provider: GraphProvider) -> str:
+    if provider == GraphProvider.FALKORDB:
         label = NEO4J_TO_FALKORDB_MAPPING[name]
         return f"CALL db.idx.fulltext.queryRelationships('{label}', $query)"
-    else:
-        return f'CALL db.index.fulltext.queryRelationships("{name}", $query, {{limit: $limit}})'
-
-
-def get_entity_node_save_bulk_query(nodes, db_type: str = 'neo4j') -> str | Any:
-    if db_type == 'falkordb':
-        queries = []
-        for node in nodes:
-            for label in node['labels']:
-                queries.append(
-                    (
-                        f"""
-                    UNWIND $nodes AS node
-                    MERGE (n:Entity {{uuid: node.uuid}})
-                    SET n:{label}
-                    SET n = node
-                    WITH n, node
-                    SET n.name_embedding = vecf32(node.name_embedding)
-                    RETURN n.uuid AS uuid
-                """,
-                        {'nodes': [node]},
-                    )
-                )
-        return queries
-    else:
-        return ENTITY_NODE_SAVE_BULK
-
-
-def get_entity_edge_save_bulk_query(db_type: str = 'neo4j') -> str:
-    if db_type == 'falkordb':
-        return """
-        UNWIND $entity_edges AS edge
-        MATCH (source:Entity {uuid: edge.source_node_uuid}) 
-        MATCH (target:Entity {uuid: edge.target_node_uuid}) 
-        MERGE (source)-[r:RELATES_TO {uuid: edge.uuid}]->(target)
-        SET r = {uuid: edge.uuid, name: edge.name, group_id: edge.group_id, fact: edge.fact, episodes: edge.episodes, 
-        created_at: edge.created_at, expired_at: edge.expired_at, valid_at: edge.valid_at, invalid_at: edge.invalid_at, fact_embedding: vecf32(edge.fact_embedding)}
-        WITH r, edge
-        RETURN edge.uuid AS uuid"""
-    else:
-        return ENTITY_EDGE_SAVE_BULK
+
+    return f'CALL db.index.fulltext.queryRelationships("{name}", $query, {{limit: $limit}})'