graphiti-core 0.17.4__py3-none-any.whl → 0.24.3__py3-none-any.whl

graphiti_core/prompts/prompt_helpers.py

@@ -1 +1,40 @@
+"""
+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
+from typing import Any
+
 DO_NOT_ESCAPE_UNICODE = '\nDo not escape unicode characters.\n'
+
+
+def to_prompt_json(data: Any, ensure_ascii: bool = False, indent: int | None = None) -> str:
+    """
+    Serialize data to JSON for use in prompts.
+
+    Args:
+        data: The data to serialize
+        ensure_ascii: If True, escape non-ASCII characters. If False (default), preserve them.
+        indent: Number of spaces for indentation. Defaults to None (minified).
+
+    Returns:
+        JSON string representation of the data
+
+    Notes:
+        By default (ensure_ascii=False), non-ASCII characters (e.g., Korean, Japanese, Chinese)
+        are preserved in their original form in the prompt, making them readable
+        in LLM logs and improving model understanding.
+    """
+    return json.dumps(data, ensure_ascii=ensure_ascii, indent=indent)

graphiti_core/prompts/snippets.py

@@ -0,0 +1,29 @@
+"""
+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.
+"""
+
+summary_instructions = """Guidelines:
+        1. Output only factual content. Never explain what you're doing, why, or mention limitations/constraints. 
+        2. Only use the provided messages, entity, and entity context to set attribute values.
+        3. Keep the summary concise and to the point. STATE FACTS DIRECTLY IN UNDER 250 CHARACTERS.
+
+        Example summaries:
+        BAD: "This is the only activity in the context. The user listened to this song. No other details were provided to include in this summary."
+        GOOD: "User played 'Blue Monday' by New Order (electronic genre) on 2024-12-03 at 14:22 UTC."
+        BAD: "Based on the messages provided, the user attended a meeting. This summary focuses on that event as it was the main topic discussed."
+        GOOD: "User attended Q3 planning meeting with sales team on March 15."
+        BAD: "The context shows John ordered pizza. Due to length constraints, other details are omitted from this summary."
+        GOOD: "John ordered pepperoni pizza from Mario's at 7:30 PM, delivered to office."
+        """

graphiti_core/prompts/summarize_nodes.py

@@ -14,18 +14,19 @@ See the License for the specific language governing permissions and
 limitations under the License.
 """
 
-import json
 from typing import Any, Protocol, TypedDict
 
 from pydantic import BaseModel, Field
 
 from .models import Message, PromptFunction, PromptVersion
+from .prompt_helpers import to_prompt_json
+from .snippets import summary_instructions
 
 
 class Summary(BaseModel):
     summary: str = Field(
         ...,
-        description='Summary containing the important information about the entity. Under 250 words',
+        description='Summary containing the important information about the entity. Under 250 characters',
     )
 
 
@@ -55,11 +56,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 250 words.
+
+        IMPORTANT: Keep the summary concise and to the point. SUMMARIES MUST BE LESS THAN 250 CHARACTERS.
 
         Summaries:
-        {json.dumps(context['node_summaries'], indent=2)}
+        {to_prompt_json(context['node_summaries'])}
         """,
         ),
     ]
@@ -69,38 +70,35 @@ def summarize_context(context: dict[str, Any]) -> list[Message]:
     return [
         Message(
             role='system',
-            content='You are a helpful assistant that extracts entity properties from the provided text.',
+            content='You are a helpful assistant that generates a summary and attributes from provided text.',
         ),
         Message(
             role='user',
             content=f"""
-            
-        <MESSAGES>
-        {json.dumps(context['previous_episodes'], indent=2)}
-        {json.dumps(context['episode_content'], indent=2)}
-        </MESSAGES>
-        
-        Given the above MESSAGES and the following ENTITY name, create a summary for the ENTITY. Your summary must only use
+        Given the MESSAGES and the ENTITY name, create a summary for the ENTITY. Your summary must only use
         information from the provided MESSAGES. Your summary should also only contain information relevant to the
-        provided ENTITY. Summaries must be under 250 words.
-        
+        provided ENTITY.
+
         In addition, extract any values for the provided entity properties based on their descriptions.
         If the value of the entity property cannot be found in the current context, set the value of the property to the Python value None.
-        
-        Guidelines:
-        1. Do not hallucinate entity property values if they cannot be found in the current context.
-        2. Only use the provided messages, entity, and entity context to set attribute values.
-        
+
+        {summary_instructions}
+
+        <MESSAGES>
+        {to_prompt_json(context['previous_episodes'])}
+        {to_prompt_json(context['episode_content'])}
+        </MESSAGES>
+
         <ENTITY>
         {context['node_name']}
         </ENTITY>
-        
+
         <ENTITY CONTEXT>
         {context['node_summary']}
         </ENTITY CONTEXT>
-        
+
         <ATTRIBUTES>
-        {json.dumps(context['attributes'], indent=2)}
+        {to_prompt_json(context['attributes'])}
         </ATTRIBUTES>
         """,
         ),
@@ -117,10 +115,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 250 words.
+        Summaries must be under 250 characters.
 
         Summary:
-        {json.dumps(context['summary'], indent=2)}
+        {to_prompt_json(context['summary'])}
         """,
         ),
     ]

graphiti_core/search/search.py

@@ -21,6 +21,7 @@ from time import time
 from graphiti_core.cross_encoder.client import CrossEncoderClient
 from graphiti_core.driver.driver import GraphDriver
 from graphiti_core.edges import EntityEdge
+from graphiti_core.embedder.client import EMBEDDING_DIM
 from graphiti_core.errors import SearchRerankerError
 from graphiti_core.graphiti_types import GraphitiClients
 from graphiti_core.helpers import semaphore_gather
@@ -29,6 +30,7 @@ from graphiti_core.search.search_config import (
     DEFAULT_SEARCH_LIMIT,
     CommunityReranker,
     CommunitySearchConfig,
+    CommunitySearchMethod,
     EdgeReranker,
     EdgeSearchConfig,
     EdgeSearchMethod,
@@ -72,34 +74,53 @@ async def search(
     center_node_uuid: str | None = None,
     bfs_origin_node_uuids: list[str] | None = None,
     query_vector: list[float] | None = None,
+    driver: GraphDriver | None = None,
 ) -> SearchResults:
     start = time()
 
-    driver = clients.driver
+    driver = driver or clients.driver
     embedder = clients.embedder
     cross_encoder = clients.cross_encoder
 
     if query.strip() == '':
-        return SearchResults(
-            edges=[],
-            nodes=[],
-            episodes=[],
-            communities=[],
+        return SearchResults()
+
+    if (
+        config.edge_config
+        and EdgeSearchMethod.cosine_similarity in config.edge_config.search_methods
+        or config.edge_config
+        and EdgeReranker.mmr == config.edge_config.reranker
+        or config.node_config
+        and NodeSearchMethod.cosine_similarity in config.node_config.search_methods
+        or config.node_config
+        and NodeReranker.mmr == config.node_config.reranker
+        or (
+            config.community_config
+            and CommunitySearchMethod.cosine_similarity in config.community_config.search_methods
         )
-    query_vector = (
-        query_vector
-        if query_vector is not None
-        else await embedder.create(input_data=[query.replace('\n', ' ')])
-    )
+        or (config.community_config and CommunityReranker.mmr == config.community_config.reranker)
+    ):
+        search_vector = (
+            query_vector
+            if query_vector is not None
+            else await embedder.create(input_data=[query.replace('\n', ' ')])
+        )
+    else:
+        search_vector = [0.0] * EMBEDDING_DIM
 
     # if group_ids is empty, set it to None
     group_ids = group_ids if group_ids and group_ids != [''] else None
-    edges, nodes, episodes, communities = await semaphore_gather(
+    (
+        (edges, edge_reranker_scores),
+        (nodes, node_reranker_scores),
+        (episodes, episode_reranker_scores),
+        (communities, community_reranker_scores),
+    ) = await semaphore_gather(
         edge_search(
             driver,
             cross_encoder,
             query,
-            query_vector,
+            search_vector,
             group_ids,
             config.edge_config,
             search_filter,
@@ -112,7 +133,7 @@ async def search(
             driver,
             cross_encoder,
             query,
-            query_vector,
+            search_vector,
             group_ids,
             config.node_config,
             search_filter,
@@ -125,7 +146,7 @@ async def search(
             driver,
             cross_encoder,
             query,
-            query_vector,
+            search_vector,
             group_ids,
             config.episode_config,
             search_filter,
@@ -136,7 +157,7 @@ async def search(
             driver,
             cross_encoder,
             query,
-            query_vector,
+            search_vector,
             group_ids,
             config.community_config,
             config.limit,
@@ -146,9 +167,13 @@ async def search(
 
     results = SearchResults(
         edges=edges,
+        edge_reranker_scores=edge_reranker_scores,
         nodes=nodes,
+        node_reranker_scores=node_reranker_scores,
         episodes=episodes,
+        episode_reranker_scores=episode_reranker_scores,
         communities=communities,
+        community_reranker_scores=community_reranker_scores,
     )
 
     latency = (time() - start) * 1000
@@ -170,50 +195,72 @@ async def edge_search(
     bfs_origin_node_uuids: list[str] | None = None,
     limit=DEFAULT_SEARCH_LIMIT,
     reranker_min_score: float = 0,
-) -> list[EntityEdge]:
+) -> tuple[list[EntityEdge], list[float]]:
     if config is None:
-        return []
-    search_results: list[list[EntityEdge]] = list(
-        await semaphore_gather(
-            *[
-                edge_fulltext_search(driver, query, search_filter, group_ids, 2 * limit),
-                edge_similarity_search(
-                    driver,
-                    query_vector,
-                    None,
-                    None,
-                    search_filter,
-                    group_ids,
-                    2 * limit,
-                    config.sim_min_score,
-                ),
-                edge_bfs_search(
-                    driver, bfs_origin_node_uuids, config.bfs_max_depth, search_filter, 2 * limit
-                ),
-            ]
+        return [], []
+
+    # Build search tasks based on configured search methods
+    search_tasks = []
+    if EdgeSearchMethod.bm25 in config.search_methods:
+        search_tasks.append(
+            edge_fulltext_search(driver, query, search_filter, group_ids, 2 * limit)
+        )
+    if EdgeSearchMethod.cosine_similarity in config.search_methods:
+        search_tasks.append(
+            edge_similarity_search(
+                driver,
+                query_vector,
+                None,
+                None,
+                search_filter,
+                group_ids,
+                2 * limit,
+                config.sim_min_score,
+            )
+        )
+    if EdgeSearchMethod.bfs in config.search_methods:
+        search_tasks.append(
+            edge_bfs_search(
+                driver,
+                bfs_origin_node_uuids,
+                config.bfs_max_depth,
+                search_filter,
+                group_ids,
+                2 * limit,
+            )
         )
-    )
+
+    # Execute only the configured search methods
+    search_results: list[list[EntityEdge]] = []
+    if search_tasks:
+        search_results = list(await semaphore_gather(*search_tasks))
 
     if EdgeSearchMethod.bfs in config.search_methods and bfs_origin_node_uuids is None:
         source_node_uuids = [edge.source_node_uuid for result in search_results for edge in result]
         search_results.append(
             await edge_bfs_search(
-                driver, source_node_uuids, config.bfs_max_depth, search_filter, 2 * limit
+                driver,
+                source_node_uuids,
+                config.bfs_max_depth,
+                search_filter,
+                group_ids,
+                2 * limit,
             )
         )
 
     edge_uuid_map = {edge.uuid: edge for result in search_results for edge in result}
 
     reranked_uuids: list[str] = []
+    edge_scores: list[float] = []
     if config.reranker == EdgeReranker.rrf or config.reranker == EdgeReranker.episode_mentions:
         search_result_uuids = [[edge.uuid for edge in result] for result in search_results]
 
-        reranked_uuids = rrf(search_result_uuids, min_score=reranker_min_score)
+        reranked_uuids, edge_scores = rrf(search_result_uuids, min_score=reranker_min_score)
     elif config.reranker == EdgeReranker.mmr:
         search_result_uuids_and_vectors = await get_embeddings_for_edges(
             driver, list(edge_uuid_map.values())
         )
-        reranked_uuids = maximal_marginal_relevance(
+        reranked_uuids, edge_scores = maximal_marginal_relevance(
             query_vector,
             search_result_uuids_and_vectors,
             config.mmr_lambda,
@@ -225,12 +272,13 @@ async def edge_search(
         reranked_uuids = [
             fact_to_uuid_map[fact] for fact, score in reranked_facts if score >= reranker_min_score
         ]
+        edge_scores = [score for _, score in reranked_facts if score >= reranker_min_score]
     elif config.reranker == EdgeReranker.node_distance:
         if center_node_uuid is None:
             raise SearchRerankerError('No center node provided for Node Distance reranker')
 
         # use rrf as a preliminary sort
-        sorted_result_uuids = rrf(
+        sorted_result_uuids, node_scores = rrf(
             [[edge.uuid for edge in result] for result in search_results],
             min_score=reranker_min_score,
         )
@@ -243,7 +291,7 @@ async def edge_search(
 
         source_uuids = [source_node_uuid for source_node_uuid in source_to_edge_uuid_map]
 
-        reranked_node_uuids = await node_distance_reranker(
+        reranked_node_uuids, edge_scores = await node_distance_reranker(
             driver, source_uuids, center_node_uuid, min_score=reranker_min_score
         )
 
@@ -255,7 +303,7 @@ async def edge_search(
     if config.reranker == EdgeReranker.episode_mentions:
         reranked_edges.sort(reverse=True, key=lambda edge: len(edge.episodes))
 
-    return reranked_edges[:limit]
+    return reranked_edges[:limit], edge_scores[:limit]
 
 
 async def node_search(
@@ -270,28 +318,54 @@ async def node_search(
     bfs_origin_node_uuids: list[str] | None = None,
     limit=DEFAULT_SEARCH_LIMIT,
     reranker_min_score: float = 0,
-) -> list[EntityNode]:
+) -> tuple[list[EntityNode], list[float]]:
     if config is None:
-        return []
-    search_results: list[list[EntityNode]] = list(
-        await semaphore_gather(
-            *[
-                node_fulltext_search(driver, query, search_filter, group_ids, 2 * limit),
-                node_similarity_search(
-                    driver, query_vector, search_filter, group_ids, 2 * limit, config.sim_min_score
-                ),
-                node_bfs_search(
-                    driver, bfs_origin_node_uuids, search_filter, config.bfs_max_depth, 2 * limit
-                ),
-            ]
+        return [], []
+
+    # Build search tasks based on configured search methods
+    search_tasks = []
+    if NodeSearchMethod.bm25 in config.search_methods:
+        search_tasks.append(
+            node_fulltext_search(driver, query, search_filter, group_ids, 2 * limit)
         )
-    )
+    if NodeSearchMethod.cosine_similarity in config.search_methods:
+        search_tasks.append(
+            node_similarity_search(
+                driver,
+                query_vector,
+                search_filter,
+                group_ids,
+                2 * limit,
+                config.sim_min_score,
+            )
+        )
+    if NodeSearchMethod.bfs in config.search_methods:
+        search_tasks.append(
+            node_bfs_search(
+                driver,
+                bfs_origin_node_uuids,
+                search_filter,
+                config.bfs_max_depth,
+                group_ids,
+                2 * limit,
+            )
+        )
+
+    # Execute only the configured search methods
+    search_results: list[list[EntityNode]] = []
+    if search_tasks:
+        search_results = list(await semaphore_gather(*search_tasks))
 
     if NodeSearchMethod.bfs in config.search_methods and bfs_origin_node_uuids is None:
         origin_node_uuids = [node.uuid for result in search_results for node in result]
         search_results.append(
             await node_bfs_search(
-                driver, origin_node_uuids, search_filter, config.bfs_max_depth, 2 * limit
+                driver,
+                origin_node_uuids,
+                search_filter,
+                config.bfs_max_depth,
+                group_ids,
+                2 * limit,
             )
         )
 
@@ -299,14 +373,15 @@ async def node_search(
     node_uuid_map = {node.uuid: node for result in search_results for node in result}
 
     reranked_uuids: list[str] = []
+    node_scores: list[float] = []
     if config.reranker == NodeReranker.rrf:
-        reranked_uuids = rrf(search_result_uuids, min_score=reranker_min_score)
+        reranked_uuids, node_scores = rrf(search_result_uuids, min_score=reranker_min_score)
     elif config.reranker == NodeReranker.mmr:
         search_result_uuids_and_vectors = await get_embeddings_for_nodes(
             driver, list(node_uuid_map.values())
         )
 
-        reranked_uuids = maximal_marginal_relevance(
+        reranked_uuids, node_scores = maximal_marginal_relevance(
             query_vector,
             search_result_uuids_and_vectors,
             config.mmr_lambda,
@@ -321,23 +396,24 @@ async def node_search(
             for name, score in reranked_node_names
             if score >= reranker_min_score
         ]
+        node_scores = [score for _, score in reranked_node_names if score >= reranker_min_score]
     elif config.reranker == NodeReranker.episode_mentions:
-        reranked_uuids = await episode_mentions_reranker(
+        reranked_uuids, node_scores = await episode_mentions_reranker(
             driver, search_result_uuids, min_score=reranker_min_score
         )
     elif config.reranker == NodeReranker.node_distance:
         if center_node_uuid is None:
             raise SearchRerankerError('No center node provided for Node Distance reranker')
-        reranked_uuids = await node_distance_reranker(
+        reranked_uuids, node_scores = await node_distance_reranker(
             driver,
-            rrf(search_result_uuids, min_score=reranker_min_score),
+            rrf(search_result_uuids, min_score=reranker_min_score)[0],
             center_node_uuid,
             min_score=reranker_min_score,
         )
 
     reranked_nodes = [node_uuid_map[uuid] for uuid in reranked_uuids]
 
-    return reranked_nodes[:limit]
+    return reranked_nodes[:limit], node_scores[:limit]
 
 
 async def episode_search(
@@ -350,9 +426,9 @@ async def episode_search(
     search_filter: SearchFilters,
     limit=DEFAULT_SEARCH_LIMIT,
     reranker_min_score: float = 0,
-) -> list[EpisodicNode]:
+) -> tuple[list[EpisodicNode], list[float]]:
     if config is None:
-        return []
+        return [], []
     search_results: list[list[EpisodicNode]] = list(
         await semaphore_gather(
             *[
@@ -365,12 +441,13 @@ async def episode_search(
     episode_uuid_map = {episode.uuid: episode for result in search_results for episode in result}
 
     reranked_uuids: list[str] = []
+    episode_scores: list[float] = []
     if config.reranker == EpisodeReranker.rrf:
-        reranked_uuids = rrf(search_result_uuids, min_score=reranker_min_score)
+        reranked_uuids, episode_scores = rrf(search_result_uuids, min_score=reranker_min_score)
 
     elif config.reranker == EpisodeReranker.cross_encoder:
         # use rrf as a preliminary reranker
-        rrf_result_uuids = rrf(search_result_uuids, min_score=reranker_min_score)
+        rrf_result_uuids, episode_scores = rrf(search_result_uuids, min_score=reranker_min_score)
         rrf_results = [episode_uuid_map[uuid] for uuid in rrf_result_uuids][:limit]
 
         content_to_uuid_map = {episode.content: episode.uuid for episode in rrf_results}
@@ -381,10 +458,11 @@ async def episode_search(
             for content, score in reranked_contents
             if score >= reranker_min_score
         ]
+        episode_scores = [score for _, score in reranked_contents if score >= reranker_min_score]
 
     reranked_episodes = [episode_uuid_map[uuid] for uuid in reranked_uuids]
 
-    return reranked_episodes[:limit]
+    return reranked_episodes[:limit], episode_scores[:limit]
 
 
 async def community_search(
@@ -396,9 +474,9 @@ async def community_search(
     config: CommunitySearchConfig | None,
     limit=DEFAULT_SEARCH_LIMIT,
     reranker_min_score: float = 0,
-) -> list[CommunityNode]:
+) -> tuple[list[CommunityNode], list[float]]:
     if config is None:
-        return []
+        return [], []
 
     search_results: list[list[CommunityNode]] = list(
         await semaphore_gather(
@@ -417,14 +495,15 @@ async def community_search(
     }
 
     reranked_uuids: list[str] = []
+    community_scores: list[float] = []
     if config.reranker == CommunityReranker.rrf:
-        reranked_uuids = rrf(search_result_uuids, min_score=reranker_min_score)
+        reranked_uuids, community_scores = rrf(search_result_uuids, min_score=reranker_min_score)
     elif config.reranker == CommunityReranker.mmr:
         search_result_uuids_and_vectors = await get_embeddings_for_communities(
             driver, list(community_uuid_map.values())
         )
 
-        reranked_uuids = maximal_marginal_relevance(
+        reranked_uuids, community_scores = maximal_marginal_relevance(
             query_vector, search_result_uuids_and_vectors, config.mmr_lambda, reranker_min_score
         )
     elif config.reranker == CommunityReranker.cross_encoder:
@@ -433,7 +512,8 @@ async def community_search(
         reranked_uuids = [
             name_to_uuid_map[name] for name, score in reranked_nodes if score >= reranker_min_score
         ]
+        community_scores = [score for _, score in reranked_nodes if score >= reranker_min_score]
 
     reranked_communities = [community_uuid_map[uuid] for uuid in reranked_uuids]
 
-    return reranked_communities[:limit]
+    return reranked_communities[:limit], community_scores[:limit]

graphiti_core/search/search_config.py

@@ -119,7 +119,42 @@ class SearchConfig(BaseModel):
 
 
 class SearchResults(BaseModel):
-    edges: list[EntityEdge]
-    nodes: list[EntityNode]
-    episodes: list[EpisodicNode]
-    communities: list[CommunityNode]
+    edges: list[EntityEdge] = Field(default_factory=list)
+    edge_reranker_scores: list[float] = Field(default_factory=list)
+    nodes: list[EntityNode] = Field(default_factory=list)
+    node_reranker_scores: list[float] = Field(default_factory=list)
+    episodes: list[EpisodicNode] = Field(default_factory=list)
+    episode_reranker_scores: list[float] = Field(default_factory=list)
+    communities: list[CommunityNode] = Field(default_factory=list)
+    community_reranker_scores: list[float] = Field(default_factory=list)
+
+    @classmethod
+    def merge(cls, results_list: list['SearchResults']) -> 'SearchResults':
+        """
+        Merge multiple SearchResults objects into a single SearchResults object.
+
+        Parameters
+        ----------
+        results_list : list[SearchResults]
+            List of SearchResults objects to merge
+
+        Returns
+        -------
+        SearchResults
+            A single SearchResults object containing all results
+        """
+        if not results_list:
+            return cls()
+
+        merged = cls()
+        for result in results_list:
+            merged.edges.extend(result.edges)
+            merged.edge_reranker_scores.extend(result.edge_reranker_scores)
+            merged.nodes.extend(result.nodes)
+            merged.node_reranker_scores.extend(result.node_reranker_scores)
+            merged.episodes.extend(result.episodes)
+            merged.episode_reranker_scores.extend(result.episode_reranker_scores)
+            merged.communities.extend(result.communities)
+            merged.community_reranker_scores.extend(result.community_reranker_scores)
+
+        return merged