graphiti-core 0.4.3__py3-none-any.whl → 0.5.0__py3-none-any.whl

graphiti_core/prompts/summarize_nodes.py

@@ -17,9 +17,21 @@ limitations under the License.
 import json
 from typing import Any, Protocol, TypedDict
 
+from pydantic import BaseModel, Field
+
 from .models import Message, PromptFunction, PromptVersion
 
 
+class Summary(BaseModel):
+    summary: str = Field(
+        ..., description='Summary containing the important information from both summaries'
+    )
+
+
+class SummaryDescription(BaseModel):
+    description: str = Field(..., description='One sentence description of the provided summary')
+
+
 class Prompt(Protocol):
     summarize_pair: PromptVersion
     summarize_context: PromptVersion
@@ -42,14 +54,11 @@ def summarize_pair(context: dict[str, Any]) -> list[Message]:
             role='user',
             content=f"""
         Synthesize the information from the following two summaries into a single succinct summary.
+        
+        Summaries must be under 500 words.
 
         Summaries:
         {json.dumps(context['node_summaries'], indent=2)}
-
-        Respond with a JSON object in the following format:
-            {{
-                "summary": "Summary containing the important information from both summaries"
-            }}
         """,
         ),
     ]
@@ -74,15 +83,11 @@ def summarize_context(context: dict[str, Any]) -> list[Message]:
         information from the provided MESSAGES. Your summary should also only contain information relevant to the
         provided ENTITY.
         
+        Summaries must be under 500 words.
+        
         <ENTITY>
         {context['node_name']}
         </ENTITY>
-        
-
-        Respond with a JSON object in the following format:
-            {{
-                "summary": "Entity summary"
-            }}
         """,
         ),
     ]
@@ -98,14 +103,10 @@ def summary_description(context: dict[str, Any]) -> list[Message]:
             role='user',
             content=f"""
         Create a short one sentence description of the summary that explains what kind of information is summarized.
+        Summaries must be under 500 words.
 
         Summary:
         {json.dumps(context['summary'], indent=2)}
-
-        Respond with a JSON object in the following format:
-            {{
-                "description": "One sentence description of the provided summary"
-            }}
         """,
         ),
     ]

graphiti_core/search/search.py

@@ -14,7 +14,6 @@ See the License for the specific language governing permissions and
 limitations under the License.
 """
 
-import asyncio
 import logging
 from collections import defaultdict
 from time import time
@@ -25,6 +24,7 @@ from graphiti_core.cross_encoder.client import CrossEncoderClient
 from graphiti_core.edges import EntityEdge
 from graphiti_core.embedder import EmbedderClient
 from graphiti_core.errors import SearchRerankerError
+from graphiti_core.helpers import semaphore_gather
 from graphiti_core.nodes import CommunityNode, EntityNode
 from graphiti_core.search.search_config import (
     DEFAULT_SEARCH_LIMIT,
@@ -78,7 +78,7 @@ async def search(
 
     # if group_ids is empty, set it to None
     group_ids = group_ids if group_ids else None
-    edges, nodes, communities = await asyncio.gather(
+    edges, nodes, communities = await semaphore_gather(
         edge_search(
             driver,
             cross_encoder,
@@ -141,7 +141,7 @@ async def edge_search(
         return []
 
     search_results: list[list[EntityEdge]] = list(
-        await asyncio.gather(
+        await semaphore_gather(
             *[
                 edge_fulltext_search(driver, query, group_ids, 2 * limit),
                 edge_similarity_search(
@@ -226,7 +226,7 @@ async def node_search(
         return []
 
     search_results: list[list[EntityNode]] = list(
-        await asyncio.gather(
+        await semaphore_gather(
             *[
                 node_fulltext_search(driver, query, group_ids, 2 * limit),
                 node_similarity_search(
@@ -295,7 +295,7 @@ async def community_search(
         return []
 
     search_results: list[list[CommunityNode]] = list(
-        await asyncio.gather(
+        await semaphore_gather(
             *[
                 community_fulltext_search(driver, query, group_ids, 2 * limit),
                 community_similarity_search(

graphiti_core/search/search_utils.py

@@ -14,10 +14,10 @@ See the License for the specific language governing permissions and
 limitations under the License.
 """
 
-import asyncio
 import logging
 from collections import defaultdict
 from time import time
+from typing import Any
 
 import numpy as np
 from neo4j import AsyncDriver, Query
@@ -29,6 +29,7 @@ from graphiti_core.helpers import (
     USE_PARALLEL_RUNTIME,
     lucene_sanitize,
     normalize_l2,
+    semaphore_gather,
 )
 from graphiti_core.nodes import (
     CommunityNode,
@@ -191,12 +192,27 @@ async def edge_similarity_search(
         'CYPHER runtime = parallel parallelRuntimeSupport=all\n' if USE_PARALLEL_RUNTIME else ''
     )
 
-    query: LiteralString = """
-                MATCH (n:Entity)-[r:RELATES_TO]->(m:Entity)
-                WHERE ($group_ids IS NULL OR r.group_id IN $group_ids)
-                AND ($source_uuid IS NULL OR n.uuid IN [$source_uuid, $target_uuid])
-                AND ($target_uuid IS NULL OR m.uuid IN [$source_uuid, $target_uuid])
-                WITH DISTINCT r, vector.similarity.cosine(r.fact_embedding, $search_vector) AS score
+    query_params: dict[str, Any] = {}
+
+    group_filter_query: LiteralString = ''
+    if group_ids is not None:
+        group_filter_query += 'WHERE r.group_id IN $group_ids'
+        query_params['group_ids'] = group_ids
+        query_params['source_node_uuid'] = source_node_uuid
+        query_params['target_node_uuid'] = target_node_uuid
+
+        if source_node_uuid is not None:
+            group_filter_query += '\nAND (n.uuid IN [$source_uuid, $target_uuid])'
+
+        if target_node_uuid is not None:
+            group_filter_query += '\nAND (m.uuid IN [$source_uuid, $target_uuid])'
+
+    query: LiteralString = (
+        """
+                        MATCH (n:Entity)-[r:RELATES_TO]->(m:Entity)
+                        """
+        + group_filter_query
+        + """\nWITH DISTINCT r, vector.similarity.cosine(r.fact_embedding, $search_vector) AS score
                 WHERE score > $min_score
                 RETURN
                     r.uuid AS uuid,
@@ -214,9 +230,11 @@ async def edge_similarity_search(
                 ORDER BY score DESC
                 LIMIT $limit
         """
+    )
 
     records, _, _ = await driver.execute_query(
         runtime_query + query,
+        query_params,
         search_vector=search_vector,
         source_uuid=source_node_uuid,
         target_uuid=target_node_uuid,
@@ -325,11 +343,20 @@ async def node_similarity_search(
         'CYPHER runtime = parallel parallelRuntimeSupport=all\n' if USE_PARALLEL_RUNTIME else ''
     )
 
+    query_params: dict[str, Any] = {}
+
+    group_filter_query: LiteralString = ''
+    if group_ids is not None:
+        group_filter_query += 'WHERE n.group_id IN $group_ids'
+        query_params['group_ids'] = group_ids
+
     records, _, _ = await driver.execute_query(
         runtime_query
         + """
             MATCH (n:Entity)
-            WHERE $group_ids IS NULL OR n.group_id IN $group_ids
+            """
+        + group_filter_query
+        + """
             WITH n, vector.similarity.cosine(n.name_embedding, $search_vector) AS score
             WHERE score > $min_score
             RETURN
@@ -342,6 +369,7 @@ async def node_similarity_search(
             ORDER BY score DESC
             LIMIT $limit
             """,
+        query_params,
         search_vector=search_vector,
         group_ids=group_ids,
         limit=limit,
@@ -436,11 +464,20 @@ async def community_similarity_search(
         'CYPHER runtime = parallel parallelRuntimeSupport=all\n' if USE_PARALLEL_RUNTIME else ''
     )
 
+    query_params: dict[str, Any] = {}
+
+    group_filter_query: LiteralString = ''
+    if group_ids is not None:
+        group_filter_query += 'WHERE comm.group_id IN $group_ids'
+        query_params['group_ids'] = group_ids
+
     records, _, _ = await driver.execute_query(
         runtime_query
         + """
            MATCH (comm:Community)
-           WHERE ($group_ids IS NULL OR comm.group_id IN $group_ids)
+           """
+        + group_filter_query
+        + """
            WITH comm, vector.similarity.cosine(comm.name_embedding, $search_vector) AS score
            WHERE score > $min_score
            RETURN
@@ -512,7 +549,7 @@ async def hybrid_node_search(
 
     start = time()
     results: list[list[EntityNode]] = list(
-        await asyncio.gather(
+        await semaphore_gather(
             *[node_fulltext_search(driver, q, group_ids, 2 * limit) for q in queries],
             *[node_similarity_search(driver, e, group_ids, 2 * limit) for e in embeddings],
         )
@@ -582,7 +619,7 @@ async def get_relevant_edges(
     relevant_edges: list[EntityEdge] = []
     relevant_edge_uuids = set()
 
-    results = await asyncio.gather(
+    results = await semaphore_gather(
         *[
             edge_similarity_search(
                 driver,
@@ -631,7 +668,7 @@ async def node_distance_reranker(
 ) -> list[str]:
     # filter out node_uuid center node node uuid
     filtered_uuids = list(filter(lambda node_uuid: node_uuid != center_node_uuid, node_uuids))
-    scores: dict[str, float] = {}
+    scores: dict[str, float] = {center_node_uuid: 0.0}
 
     # Find the shortest path to center node
     query = Query("""
@@ -649,9 +686,13 @@ async def node_distance_reranker(
 
     for result in path_results:
         uuid = result['uuid']
-        score = result['score'] if 'score' in result else float('inf')
+        score = result['score']
         scores[uuid] = score
 
+    for uuid in filtered_uuids:
+        if uuid not in scores:
+            scores[uuid] = float('inf')
+
     # rerank on shortest distance
     filtered_uuids.sort(key=lambda cur_uuid: scores[cur_uuid])
 

graphiti_core/utils/__init__.py

@@ -1,15 +0,0 @@
-from .maintenance import (
-    build_episodic_edges,
-    clear_data,
-    extract_edges,
-    extract_nodes,
-    retrieve_episodes,
-)
-
-__all__ = [
-    'extract_edges',
-    'build_episodic_edges',
-    'extract_nodes',
-    'clear_data',
-    'retrieve_episodes',
-]

graphiti_core/utils/bulk_utils.py

@@ -14,11 +14,10 @@ See the License for the specific language governing permissions and
 limitations under the License.
 """
 
-import asyncio
 import logging
 import typing
 from collections import defaultdict
-from datetime import datetime, timezone
+from datetime import datetime
 from math import ceil
 
 from neo4j import AsyncDriver, AsyncManagedTransaction
@@ -26,6 +25,7 @@ from numpy import dot, sqrt
 from pydantic import BaseModel
 
 from graphiti_core.edges import Edge, EntityEdge, EpisodicEdge
+from graphiti_core.helpers import semaphore_gather
 from graphiti_core.llm_client import LLMClient
 from graphiti_core.models.edges.edge_db_queries import (
     ENTITY_EDGE_SAVE_BULK,
@@ -37,14 +37,17 @@ from graphiti_core.models.nodes.node_db_queries import (
 )
 from graphiti_core.nodes import EntityNode, EpisodeType, EpisodicNode
 from graphiti_core.search.search_utils import get_relevant_edges, get_relevant_nodes
-from graphiti_core.utils import retrieve_episodes
+from graphiti_core.utils.datetime_utils import utc_now
 from graphiti_core.utils.maintenance.edge_operations import (
     build_episodic_edges,
     dedupe_edge_list,
     dedupe_extracted_edges,
     extract_edges,
 )
-from graphiti_core.utils.maintenance.graph_data_operations import EPISODE_WINDOW_LEN
+from graphiti_core.utils.maintenance.graph_data_operations import (
+    EPISODE_WINDOW_LEN,
+    retrieve_episodes,
+)
 from graphiti_core.utils.maintenance.node_operations import (
     dedupe_extracted_nodes,
     dedupe_node_list,
@@ -68,7 +71,7 @@ class RawEpisode(BaseModel):
 async def retrieve_previous_episodes_bulk(
     driver: AsyncDriver, episodes: list[EpisodicNode]
 ) -> list[tuple[EpisodicNode, list[EpisodicNode]]]:
-    previous_episodes_list = await asyncio.gather(
+    previous_episodes_list = await semaphore_gather(
         *[
             retrieve_episodes(
                 driver, episode.valid_at, last_n=EPISODE_WINDOW_LEN, group_ids=[episode.group_id]
@@ -115,7 +118,7 @@ async def add_nodes_and_edges_bulk_tx(
 async def extract_nodes_and_edges_bulk(
     llm_client: LLMClient, episode_tuples: list[tuple[EpisodicNode, list[EpisodicNode]]]
 ) -> tuple[list[EntityNode], list[EntityEdge], list[EpisodicEdge]]:
-    extracted_nodes_bulk = await asyncio.gather(
+    extracted_nodes_bulk = await semaphore_gather(
         *[
             extract_nodes(llm_client, episode, previous_episodes)
             for episode, previous_episodes in episode_tuples
@@ -127,7 +130,7 @@ async def extract_nodes_and_edges_bulk(
         [episode[1] for episode in episode_tuples],
     )
 
-    extracted_edges_bulk = await asyncio.gather(
+    extracted_edges_bulk = await semaphore_gather(
         *[
             extract_edges(
                 llm_client,
@@ -168,13 +171,13 @@ async def dedupe_nodes_bulk(
     node_chunks = [nodes[i : i + CHUNK_SIZE] for i in range(0, len(nodes), CHUNK_SIZE)]
 
     existing_nodes_chunks: list[list[EntityNode]] = list(
-        await asyncio.gather(
+        await semaphore_gather(
             *[get_relevant_nodes(driver, node_chunk) for node_chunk in node_chunks]
         )
     )
 
     results: list[tuple[list[EntityNode], dict[str, str]]] = list(
-        await asyncio.gather(
+        await semaphore_gather(
             *[
                 dedupe_extracted_nodes(llm_client, node_chunk, existing_nodes_chunks[i])
                 for i, node_chunk in enumerate(node_chunks)
@@ -202,13 +205,13 @@ async def dedupe_edges_bulk(
     ]
 
     relevant_edges_chunks: list[list[EntityEdge]] = list(
-        await asyncio.gather(
+        await semaphore_gather(
             *[get_relevant_edges(driver, edge_chunk, None, None) for edge_chunk in edge_chunks]
         )
     )
 
     resolved_edge_chunks: list[list[EntityEdge]] = list(
-        await asyncio.gather(
+        await semaphore_gather(
             *[
                 dedupe_extracted_edges(llm_client, edge_chunk, relevant_edges_chunks[i])
                 for i, edge_chunk in enumerate(edge_chunks)
@@ -289,7 +292,9 @@ async def compress_nodes(
             # add both nodes to the shortest chunk
             node_chunks[-1].extend([n, m])
 
-    results = await asyncio.gather(*[dedupe_node_list(llm_client, chunk) for chunk in node_chunks])
+    results = await semaphore_gather(
+        *[dedupe_node_list(llm_client, chunk) for chunk in node_chunks]
+    )
 
     extended_map = dict(uuid_map)
     compressed_nodes: list[EntityNode] = []
@@ -312,7 +317,9 @@ async def compress_edges(llm_client: LLMClient, edges: list[EntityEdge]) -> list
     # We build a map of the edges based on their source and target nodes.
     edge_chunks = chunk_edges_by_nodes(edges)
 
-    results = await asyncio.gather(*[dedupe_edge_list(llm_client, chunk) for chunk in edge_chunks])
+    results = await semaphore_gather(
+        *[dedupe_edge_list(llm_client, chunk) for chunk in edge_chunks]
+    )
 
     compressed_edges: list[EntityEdge] = []
     for edge_chunk in results:
@@ -365,7 +372,7 @@ async def extract_edge_dates_bulk(
         episode.uuid: (episode, previous_episodes) for episode, previous_episodes in episode_pairs
     }
 
-    results = await asyncio.gather(
+    results = await semaphore_gather(
         *[
             extract_edge_dates(
                 llm_client,
@@ -385,7 +392,7 @@ async def extract_edge_dates_bulk(
         edge.valid_at = valid_at
         edge.invalid_at = invalid_at
         if edge.invalid_at:
-            edge.expired_at = datetime.now(timezone.utc)
+            edge.expired_at = utc_now()
 
     return edges
 

graphiti_core/utils/datetime_utils.py

@@ -0,0 +1,42 @@
+"""
+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 datetime import datetime, timezone
+
+
+def utc_now() -> datetime:
+    """Returns the current UTC datetime with timezone information."""
+    return datetime.now(timezone.utc)
+
+
+def ensure_utc(dt: datetime | None) -> datetime | None:
+    """
+    Ensures a datetime is timezone-aware and in UTC.
+    If the datetime is naive (no timezone), assumes it's in UTC.
+    If the datetime has a different timezone, converts it to UTC.
+    Returns None if input is None.
+    """
+    if dt is None:
+        return None
+
+    if dt.tzinfo is None:
+        # If datetime is naive, assume it's UTC
+        return dt.replace(tzinfo=timezone.utc)
+    elif dt.tzinfo != timezone.utc:
+        # If datetime has a different timezone, convert to UTC
+        return dt.astimezone(timezone.utc)
+
+    return dt

graphiti_core/utils/maintenance/community_operations.py

@@ -1,14 +1,13 @@
 import asyncio
 import logging
 from collections import defaultdict
-from datetime import datetime, timezone
 
 from neo4j import AsyncDriver
 from pydantic import BaseModel
 
 from graphiti_core.edges import CommunityEdge
 from graphiti_core.embedder import EmbedderClient
-from graphiti_core.helpers import DEFAULT_DATABASE
+from graphiti_core.helpers import DEFAULT_DATABASE, semaphore_gather
 from graphiti_core.llm_client import LLMClient
 from graphiti_core.nodes import (
     CommunityNode,
@@ -16,6 +15,8 @@ from graphiti_core.nodes import (
     get_community_node_from_record,
 )
 from graphiti_core.prompts import prompt_library
+from graphiti_core.prompts.summarize_nodes import Summary, SummaryDescription
+from graphiti_core.utils.datetime_utils import utc_now
 from graphiti_core.utils.maintenance.edge_operations import build_community_edges
 
 MAX_COMMUNITY_BUILD_CONCURRENCY = 10
@@ -70,7 +71,7 @@ async def get_community_clusters(
 
         community_clusters.extend(
             list(
-                await asyncio.gather(
+                await semaphore_gather(
                     *[EntityNode.get_by_uuids(driver, cluster) for cluster in cluster_uuids]
                 )
             )
@@ -131,7 +132,7 @@ async def summarize_pair(llm_client: LLMClient, summary_pair: tuple[str, str]) -
     context = {'node_summaries': [{'summary': summary} for summary in summary_pair]}
 
     llm_response = await llm_client.generate_response(
-        prompt_library.summarize_nodes.summarize_pair(context)
+        prompt_library.summarize_nodes.summarize_pair(context), response_model=Summary
     )
 
     pair_summary = llm_response.get('summary', '')
@@ -143,7 +144,8 @@ async def generate_summary_description(llm_client: LLMClient, summary: str) -> s
     context = {'summary': summary}
 
     llm_response = await llm_client.generate_response(
-        prompt_library.summarize_nodes.summary_description(context)
+        prompt_library.summarize_nodes.summary_description(context),
+        response_model=SummaryDescription,
     )
 
     description = llm_response.get('description', '')
@@ -162,7 +164,7 @@ async def build_community(
             odd_one_out = summaries.pop()
             length -= 1
         new_summaries: list[str] = list(
-            await asyncio.gather(
+            await semaphore_gather(
                 *[
                     summarize_pair(llm_client, (str(left_summary), str(right_summary)))
                     for left_summary, right_summary in zip(
@@ -178,7 +180,7 @@ async def build_community(
 
     summary = summaries[0]
     name = await generate_summary_description(llm_client, summary)
-    now = datetime.now(timezone.utc)
+    now = utc_now()
     community_node = CommunityNode(
         name=name,
         group_id=community_cluster[0].group_id,
@@ -205,7 +207,9 @@ async def build_communities(
             return await build_community(llm_client, cluster)
 
     communities: list[tuple[CommunityNode, list[CommunityEdge]]] = list(
-        await asyncio.gather(*[limited_build_community(cluster) for cluster in community_clusters])
+        await semaphore_gather(
+            *[limited_build_community(cluster) for cluster in community_clusters]
+        )
     )
 
     community_nodes: list[CommunityNode] = []
@@ -305,7 +309,7 @@ async def update_community(
     community.name = new_name
 
     if is_new:
-        community_edge = (build_community_edges([entity], community, datetime.now(timezone.utc)))[0]
+        community_edge = (build_community_edges([entity], community, utc_now()))[0]
         await community_edge.save(driver)
 
     await community.generate_name_embedding(embedder)