graphiti-core 0.12.0rc5__py3-none-any.whl → 0.12.1__py3-none-any.whl

graphiti_core/graph_queries.py

@@ -0,0 +1,149 @@
+"""
+Database query utilities for different graph database backends.
+
+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,
+)
+
+# Mapping from Neo4j fulltext index names to FalkorDB node labels
+NEO4J_TO_FALKORDB_MAPPING = {
+    'node_name_and_summary': 'Entity',
+    'community_name': 'Community',
+    'episode_content': 'Episodic',
+    'edge_name_and_fact': 'RELATES_TO',
+}
+
+
+def get_range_indices(db_type: str = 'neo4j') -> list[LiteralString]:
+    if db_type == 'falkordb':
+        return [
+            # Entity node
+            'CREATE INDEX FOR (n:Entity) ON (n.uuid, n.group_id, n.name, n.created_at)',
+            # Episodic node
+            'CREATE INDEX FOR (n:Episodic) ON (n.uuid, n.group_id, n.created_at, n.valid_at)',
+            # Community node
+            'CREATE INDEX FOR (n:Community) ON (n.uuid)',
+            # RELATES_TO edge
+            'CREATE INDEX FOR ()-[e:RELATES_TO]-() ON (e.uuid, e.group_id, e.name, e.created_at, e.expired_at, e.valid_at, e.invalid_at)',
+            # MENTIONS edge
+            'CREATE INDEX FOR ()-[e:MENTIONS]-() ON (e.uuid, e.group_id)',
+            # 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 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]""",
+        ]
+
+
+def get_nodes_query(db_type: str = 'neo4j', name: str = '', query: str | None = None) -> str:
+    if db_type == '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}})'
+
+
+def get_vector_cosine_func_query(vec1, vec2, db_type: str = 'neo4j') -> str:
+    if db_type == '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})'
+
+
+def get_relationships_query(name: str, db_type: str = 'neo4j') -> str:
+    if db_type == '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

graphiti_core/graphiti.py

@@ -19,12 +19,13 @@ from datetime import datetime
 from time import time
 
 from dotenv import load_dotenv
-from neo4j import AsyncGraphDatabase
 from pydantic import BaseModel
 from typing_extensions import LiteralString
 
 from graphiti_core.cross_encoder.client import CrossEncoderClient
 from graphiti_core.cross_encoder.openai_reranker_client import OpenAIRerankerClient
+from graphiti_core.driver.driver import GraphDriver
+from graphiti_core.driver.neo4j_driver import Neo4jDriver
 from graphiti_core.edges import EntityEdge, EpisodicEdge
 from graphiti_core.embedder import EmbedderClient, OpenAIEmbedder
 from graphiti_core.graphiti_types import GraphitiClients
@@ -94,12 +95,13 @@ class Graphiti:
     def __init__(
         self,
         uri: str,
-        user: str,
-        password: str,
+        user: str | None = None,
+        password: str | None = None,
         llm_client: LLMClient | None = None,
         embedder: EmbedderClient | None = None,
         cross_encoder: CrossEncoderClient | None = None,
         store_raw_episode_content: bool = True,
+        graph_driver: GraphDriver | None = None,
     ):
         """
         Initialize a Graphiti instance.
@@ -137,7 +139,9 @@ class Graphiti:
         Make sure to set the OPENAI_API_KEY environment variable before initializing
         Graphiti if you're using the default OpenAIClient.
         """
-        self.driver = AsyncGraphDatabase.driver(uri, auth=(user, password))
+
+        self.driver = graph_driver if graph_driver else Neo4jDriver(uri, user, password)
+
         self.database = DEFAULT_DATABASE
         self.store_raw_episode_content = store_raw_episode_content
         if llm_client:

graphiti_core/graphiti_types.py

@@ -14,16 +14,16 @@ See the License for the specific language governing permissions and
 limitations under the License.
 """
 
-from neo4j import AsyncDriver
 from pydantic import BaseModel, ConfigDict
 
 from graphiti_core.cross_encoder import CrossEncoderClient
+from graphiti_core.driver.driver import GraphDriver
 from graphiti_core.embedder import EmbedderClient
 from graphiti_core.llm_client import LLMClient
 
 
 class GraphitiClients(BaseModel):
-    driver: AsyncDriver
+    driver: GraphDriver
     llm_client: LLMClient
     embedder: EmbedderClient
     cross_encoder: CrossEncoderClient

graphiti_core/helpers.py

@@ -27,7 +27,7 @@ from typing_extensions import LiteralString
 
 load_dotenv()
 
-DEFAULT_DATABASE = os.getenv('DEFAULT_DATABASE', None)
+DEFAULT_DATABASE = os.getenv('DEFAULT_DATABASE', 'neo4j')
 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))
@@ -38,8 +38,14 @@ RUNTIME_QUERY: LiteralString = (
 )
 
 
-def parse_db_date(neo_date: neo4j_time.DateTime | None) -> datetime | None:
-    return neo_date.to_native() if neo_date else None
+def parse_db_date(neo_date: neo4j_time.DateTime | str | None) -> datetime | None:
+    return (
+        neo_date.to_native()
+        if isinstance(neo_date, neo4j_time.DateTime)
+        else datetime.fromisoformat(neo_date)
+        if neo_date
+        else None
+    )
 
 
 def lucene_sanitize(query: str) -> str:

graphiti_core/llm_client/__init__.py

@@ -1,3 +1,19 @@
+"""
+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.
+"""
+
 from .client import LLMClient
 from .config import LLMConfig
 from .errors import RateLimitError

graphiti_core/llm_client/azure_openai_client.py

@@ -0,0 +1,73 @@
+"""
+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 json
+import logging
+from typing import Any
+
+from openai import AsyncAzureOpenAI
+from openai.types.chat import ChatCompletionMessageParam
+from pydantic import BaseModel
+
+from ..prompts.models import Message
+from .client import LLMClient
+from .config import LLMConfig, ModelSize
+
+logger = logging.getLogger(__name__)
+
+
+class AzureOpenAILLMClient(LLMClient):
+    """Wrapper class for AsyncAzureOpenAI that implements the LLMClient interface."""
+
+    def __init__(self, azure_client: AsyncAzureOpenAI, config: LLMConfig | None = None):
+        super().__init__(config, cache=False)
+        self.azure_client = azure_client
+
+    async def _generate_response(
+        self,
+        messages: list[Message],
+        response_model: type[BaseModel] | None = None,
+        max_tokens: int = 1024,
+        model_size: ModelSize = ModelSize.medium,
+    ) -> dict[str, Any]:
+        """Generate response using Azure OpenAI client."""
+        # Convert messages to OpenAI format
+        openai_messages: list[ChatCompletionMessageParam] = []
+        for message in messages:
+            message.content = self._clean_input(message.content)
+            if message.role == 'user':
+                openai_messages.append({'role': 'user', 'content': message.content})
+            elif message.role == 'system':
+                openai_messages.append({'role': 'system', 'content': message.content})
+
+        # Ensure model is a string
+        model_name = self.model if self.model else 'gpt-4o-mini'
+
+        try:
+            response = await self.azure_client.chat.completions.create(
+                model=model_name,
+                messages=openai_messages,
+                temperature=float(self.temperature) if self.temperature is not None else 0.7,
+                max_tokens=max_tokens,
+                response_format={'type': 'json_object'},
+            )
+            result = response.choices[0].message.content or '{}'
+
+            # Parse JSON response
+            return json.loads(result)
+        except Exception as e:
+            logger.error(f'Error in Azure OpenAI LLM response: {e}')
+            raise

graphiti_core/nodes.py

@@ -22,13 +22,13 @@ from time import time
 from typing import Any
 from uuid import uuid4
 
-from neo4j import AsyncDriver
 from pydantic import BaseModel, Field
 from typing_extensions import LiteralString
 
+from graphiti_core.driver.driver import GraphDriver
 from graphiti_core.embedder import EmbedderClient
 from graphiti_core.errors import NodeNotFoundError
-from graphiti_core.helpers import DEFAULT_DATABASE
+from graphiti_core.helpers import DEFAULT_DATABASE, parse_db_date
 from graphiti_core.models.nodes.node_db_queries import (
     COMMUNITY_NODE_SAVE,
     ENTITY_NODE_SAVE,
@@ -94,9 +94,9 @@ class Node(BaseModel, ABC):
     created_at: datetime = Field(default_factory=lambda: utc_now())
 
     @abstractmethod
-    async def save(self, driver: AsyncDriver): ...
+    async def save(self, driver: GraphDriver): ...
 
-    async def delete(self, driver: AsyncDriver):
+    async def delete(self, driver: GraphDriver):
         result = await driver.execute_query(
             """
         MATCH (n:Entity|Episodic|Community {uuid: $uuid})
@@ -119,7 +119,7 @@ class Node(BaseModel, ABC):
         return False
 
     @classmethod
-    async def delete_by_group_id(cls, driver: AsyncDriver, group_id: str):
+    async def delete_by_group_id(cls, driver: GraphDriver, group_id: str):
         await driver.execute_query(
             """
         MATCH (n:Entity|Episodic|Community {group_id: $group_id})
@@ -132,10 +132,10 @@ class Node(BaseModel, ABC):
         return 'SUCCESS'
 
     @classmethod
-    async def get_by_uuid(cls, driver: AsyncDriver, uuid: str): ...
+    async def get_by_uuid(cls, driver: GraphDriver, uuid: str): ...
 
     @classmethod
-    async def get_by_uuids(cls, driver: AsyncDriver, uuids: list[str]): ...
+    async def get_by_uuids(cls, driver: GraphDriver, uuids: list[str]): ...
 
 
 class EpisodicNode(Node):
@@ -150,7 +150,7 @@ class EpisodicNode(Node):
         default_factory=list,
     )
 
-    async def save(self, driver: AsyncDriver):
+    async def save(self, driver: GraphDriver):
         result = await driver.execute_query(
             EPISODIC_NODE_SAVE,
             uuid=self.uuid,
@@ -165,12 +165,12 @@ class EpisodicNode(Node):
             database_=DEFAULT_DATABASE,
         )
 
-        logger.debug(f'Saved Node to neo4j: {self.uuid}')
+        logger.debug(f'Saved Node to Graph: {self.uuid}')
 
         return result
 
     @classmethod
-    async def get_by_uuid(cls, driver: AsyncDriver, uuid: str):
+    async def get_by_uuid(cls, driver: GraphDriver, uuid: str):
         records, _, _ = await driver.execute_query(
             """
         MATCH (e:Episodic {uuid: $uuid})
@@ -197,7 +197,7 @@ class EpisodicNode(Node):
         return episodes[0]
 
     @classmethod
-    async def get_by_uuids(cls, driver: AsyncDriver, uuids: list[str]):
+    async def get_by_uuids(cls, driver: GraphDriver, uuids: list[str]):
         records, _, _ = await driver.execute_query(
             """
         MATCH (e:Episodic) WHERE e.uuid IN $uuids
@@ -224,7 +224,7 @@ class EpisodicNode(Node):
     @classmethod
     async def get_by_group_ids(
         cls,
-        driver: AsyncDriver,
+        driver: GraphDriver,
         group_ids: list[str],
         limit: int | None = None,
         uuid_cursor: str | None = None,
@@ -263,7 +263,7 @@ class EpisodicNode(Node):
         return episodes
 
     @classmethod
-    async def get_by_entity_node_uuid(cls, driver: AsyncDriver, entity_node_uuid: str):
+    async def get_by_entity_node_uuid(cls, driver: GraphDriver, entity_node_uuid: str):
         records, _, _ = await driver.execute_query(
             """
         MATCH (e:Episodic)-[r:MENTIONS]->(n:Entity {uuid: $entity_node_uuid})
@@ -304,7 +304,7 @@ class EntityNode(Node):
 
         return self.name_embedding
 
-    async def load_name_embedding(self, driver: AsyncDriver):
+    async def load_name_embedding(self, driver: GraphDriver):
         query: LiteralString = """
             MATCH (n:Entity {uuid: $uuid})
             RETURN n.name_embedding AS name_embedding
@@ -318,7 +318,7 @@ class EntityNode(Node):
 
         self.name_embedding = records[0]['name_embedding']
 
-    async def save(self, driver: AsyncDriver):
+    async def save(self, driver: GraphDriver):
         entity_data: dict[str, Any] = {
             'uuid': self.uuid,
             'name': self.name,
@@ -337,16 +337,16 @@ class EntityNode(Node):
             database_=DEFAULT_DATABASE,
         )
 
-        logger.debug(f'Saved Node to neo4j: {self.uuid}')
+        logger.debug(f'Saved Node to Graph: {self.uuid}')
 
         return result
 
     @classmethod
-    async def get_by_uuid(cls, driver: AsyncDriver, uuid: str):
+    async def get_by_uuid(cls, driver: GraphDriver, uuid: str):
         query = (
             """
-                                                        MATCH (n:Entity {uuid: $uuid})
-                                                        """
+                                                                    MATCH (n:Entity {uuid: $uuid})
+                                                                    """
             + ENTITY_NODE_RETURN
         )
         records, _, _ = await driver.execute_query(
@@ -364,7 +364,7 @@ class EntityNode(Node):
         return nodes[0]
 
     @classmethod
-    async def get_by_uuids(cls, driver: AsyncDriver, uuids: list[str]):
+    async def get_by_uuids(cls, driver: GraphDriver, uuids: list[str]):
         records, _, _ = await driver.execute_query(
             """
         MATCH (n:Entity) WHERE n.uuid IN $uuids
@@ -382,7 +382,7 @@ class EntityNode(Node):
     @classmethod
     async def get_by_group_ids(
         cls,
-        driver: AsyncDriver,
+        driver: GraphDriver,
         group_ids: list[str],
         limit: int | None = None,
         uuid_cursor: str | None = None,
@@ -416,7 +416,7 @@ class CommunityNode(Node):
     name_embedding: list[float] | None = Field(default=None, description='embedding of the name')
     summary: str = Field(description='region summary of member nodes', default_factory=str)
 
-    async def save(self, driver: AsyncDriver):
+    async def save(self, driver: GraphDriver):
         result = await driver.execute_query(
             COMMUNITY_NODE_SAVE,
             uuid=self.uuid,
@@ -428,7 +428,7 @@ class CommunityNode(Node):
             database_=DEFAULT_DATABASE,
         )
 
-        logger.debug(f'Saved Node to neo4j: {self.uuid}')
+        logger.debug(f'Saved Node to Graph: {self.uuid}')
 
         return result
 
@@ -441,7 +441,7 @@ class CommunityNode(Node):
 
         return self.name_embedding
 
-    async def load_name_embedding(self, driver: AsyncDriver):
+    async def load_name_embedding(self, driver: GraphDriver):
         query: LiteralString = """
             MATCH (c:Community {uuid: $uuid})
             RETURN c.name_embedding AS name_embedding
@@ -456,7 +456,7 @@ class CommunityNode(Node):
         self.name_embedding = records[0]['name_embedding']
 
     @classmethod
-    async def get_by_uuid(cls, driver: AsyncDriver, uuid: str):
+    async def get_by_uuid(cls, driver: GraphDriver, uuid: str):
         records, _, _ = await driver.execute_query(
             """
         MATCH (n:Community {uuid: $uuid})
@@ -480,7 +480,7 @@ class CommunityNode(Node):
         return nodes[0]
 
     @classmethod
-    async def get_by_uuids(cls, driver: AsyncDriver, uuids: list[str]):
+    async def get_by_uuids(cls, driver: GraphDriver, uuids: list[str]):
         records, _, _ = await driver.execute_query(
             """
         MATCH (n:Community) WHERE n.uuid IN $uuids
@@ -503,7 +503,7 @@ class CommunityNode(Node):
     @classmethod
     async def get_by_group_ids(
         cls,
-        driver: AsyncDriver,
+        driver: GraphDriver,
         group_ids: list[str],
         limit: int | None = None,
         uuid_cursor: str | None = None,
@@ -542,8 +542,8 @@ class CommunityNode(Node):
 def get_episodic_node_from_record(record: Any) -> EpisodicNode:
     return EpisodicNode(
         content=record['content'],
-        created_at=record['created_at'].to_native().timestamp(),
-        valid_at=(record['valid_at'].to_native()),
+        created_at=parse_db_date(record['created_at']),  # type: ignore
+        valid_at=parse_db_date(record['valid_at']),  # type: ignore
         uuid=record['uuid'],
         group_id=record['group_id'],
         source=EpisodeType.from_str(record['source']),
@@ -559,7 +559,7 @@ def get_entity_node_from_record(record: Any) -> EntityNode:
         name=record['name'],
         group_id=record['group_id'],
         labels=record['labels'],
-        created_at=record['created_at'].to_native(),
+        created_at=parse_db_date(record['created_at']),  # type: ignore
         summary=record['summary'],
         attributes=record['attributes'],
     )
@@ -580,7 +580,7 @@ def get_community_node_from_record(record: Any) -> CommunityNode:
         name=record['name'],
         group_id=record['group_id'],
         name_embedding=record['name_embedding'],
-        created_at=record['created_at'].to_native(),
+        created_at=parse_db_date(record['created_at']),  # type: ignore
         summary=record['summary'],
     )
 

graphiti_core/search/search.py

@@ -18,9 +18,8 @@ import logging
 from collections import defaultdict
 from time import time
 
-from neo4j import AsyncDriver
-
 from graphiti_core.cross_encoder.client import CrossEncoderClient
+from graphiti_core.driver.driver import GraphDriver
 from graphiti_core.edges import EntityEdge
 from graphiti_core.errors import SearchRerankerError
 from graphiti_core.graphiti_types import GraphitiClients
@@ -94,7 +93,7 @@ async def search(
     )
 
     # if group_ids is empty, set it to None
-    group_ids = group_ids if group_ids else None
+    group_ids = group_ids if group_ids and group_ids != [''] else None
     edges, nodes, episodes, communities = await semaphore_gather(
         edge_search(
             driver,
@@ -160,7 +159,7 @@ async def search(
 
 
 async def edge_search(
-    driver: AsyncDriver,
+    driver: GraphDriver,
     cross_encoder: CrossEncoderClient,
     query: str,
     query_vector: list[float],
@@ -174,7 +173,6 @@ async def edge_search(
 ) -> list[EntityEdge]:
     if config is None:
         return []
-
     search_results: list[list[EntityEdge]] = list(
         await semaphore_gather(
             *[
@@ -261,7 +259,7 @@ async def edge_search(
 
 
 async def node_search(
-    driver: AsyncDriver,
+    driver: GraphDriver,
     cross_encoder: CrossEncoderClient,
     query: str,
     query_vector: list[float],
@@ -275,7 +273,6 @@ async def node_search(
 ) -> list[EntityNode]:
     if config is None:
         return []
-
     search_results: list[list[EntityNode]] = list(
         await semaphore_gather(
             *[
@@ -344,7 +341,7 @@ async def node_search(
 
 
 async def episode_search(
-    driver: AsyncDriver,
+    driver: GraphDriver,
     cross_encoder: CrossEncoderClient,
     query: str,
     _query_vector: list[float],
@@ -356,7 +353,6 @@ async def episode_search(
 ) -> list[EpisodicNode]:
     if config is None:
         return []
-
     search_results: list[list[EpisodicNode]] = list(
         await semaphore_gather(
             *[
@@ -392,7 +388,7 @@ async def episode_search(
 
 
 async def community_search(
-    driver: AsyncDriver,
+    driver: GraphDriver,
     cross_encoder: CrossEncoderClient,
     query: str,
     query_vector: list[float],