feat: async close and multi-group search support

graphiti_core/graphiti.py

@@ -129,7 +129,7 @@ def __init__(
         else:
             self.llm_client = OpenAIClient()
 
-    def close(self):
+    async def close(self):
         """
         Close the connection to the Neo4j database.
 
@@ -159,7 +159,7 @@ def close(self):
             finally:
                 graphiti.close()
         """
-        self.driver.close()
+        await self.driver.close()
 
     async def build_indices_and_constraints(self):
         """

graphiti_core/utils/maintenance/community_operations.py

@@ -154,7 +154,7 @@ async def generate_summary_description(llm_client: LLMClient, summary: str) -> s
 
 
 async def build_community(
-        llm_client: LLMClient, community_cluster: list[EntityNode]
+    llm_client: LLMClient, community_cluster: list[EntityNode]
 ) -> tuple[CommunityNode, list[CommunityEdge]]:
     summaries = [entity.summary for entity in community_cluster]
     length = len(summaries)
@@ -168,7 +168,7 @@ async def build_community(
                 *[
                     summarize_pair(llm_client, (str(left_summary), str(right_summary)))
                     for left_summary, right_summary in zip(
-                        summaries[: int(length / 2)], summaries[int(length / 2):]
+                        summaries[: int(length / 2)], summaries[int(length / 2) :]
                     )
                 ]
             )
@@ -196,7 +196,7 @@ async def build_community(
 
 
 async def build_communities(
-        driver: AsyncDriver, llm_client: LLMClient
+    driver: AsyncDriver, llm_client: LLMClient
 ) -> tuple[list[CommunityNode], list[CommunityEdge]]:
     community_clusters = await get_community_clusters(driver)
 
@@ -227,7 +227,7 @@ async def remove_communities(driver: AsyncDriver):
 
 
 async def determine_entity_community(
-        driver: AsyncDriver, entity: EntityNode
+    driver: AsyncDriver, entity: EntityNode
 ) -> tuple[CommunityNode | None, bool]:
     # Check if the node is already part of a community
     records, _, _ = await driver.execute_query(
@@ -288,7 +288,7 @@ async def determine_entity_community(
 
 
 async def update_community(
-        driver: AsyncDriver, llm_client: LLMClient, embedder, entity: EntityNode
+    driver: AsyncDriver, llm_client: LLMClient, embedder, entity: EntityNode
 ):
     community, is_new = await determine_entity_community(driver, entity)
 
@@ -307,4 +307,4 @@ async def update_community(
 
     await community.generate_name_embedding(embedder)
 
-    await community.save(driver)
+    await community.save(driver)

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "graphiti-core"
-version = "0.3.4"
+version = "0.3.5"
 description = "A temporal graph building library"
 authors = [
     "Paul Paliychuk <paul@getzep.com>",

server/graph_service/dto/retrieve.py

@@ -6,7 +6,7 @@
 
 
 class SearchQuery(BaseModel):
-    group_id: str = Field(..., description='The group id of the memory to get')
+    group_ids: list[str] = Field(description='The group ids for the memories to search')
     query: str
     max_facts: int = Field(default=10, description='The maximum number of facts to retrieve')
 

server/graph_service/routers/retrieve.py

@@ -1,4 +1,5 @@
 from datetime import datetime
+from typing import List, Optional, cast
 
 from fastapi import APIRouter, status
 
@@ -17,7 +18,9 @@
 @router.post('/search', status_code=status.HTTP_200_OK)
 async def search(query: SearchQuery, graphiti: ZepGraphitiDep):
     relevant_edges = await graphiti.search(
-        group_ids=[query.group_id],
+        group_ids=cast(
+            Optional[List[Optional[str]]], query.group_ids
+        ),  # Cast query.group_ids to match the expected type in graphiti.search
         query=query.query,
         num_results=query.max_facts,
     )

server/graph_service/zep_graphiti.py

@@ -79,7 +79,7 @@ async def get_graphiti(settings: ZepEnvDep):
     try:
         yield client
     finally:
-        client.close()
+        await client.close()
 
 
 async def initialize_graphiti(settings: ZepEnvDep):