chore: Support a list of group_ids on search + await driver.close()

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):
         """

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

@@ -17,7 +17,7 @@
 @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=query.group_ids,
         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):