From 6aa25a19016660121be48dc6bc8e094611a08ef8 Mon Sep 17 00:00:00 2001
From: Preston Rasmussen <109292228+prasmussen15@users.noreply.github.com>
Date: Thu, 10 Apr 2025 06:57:58 -0400
Subject: [PATCH] update context string (#346)

* update context string

* Update graphiti_core/search/search_helpers.py

Co-authored-by: ellipsis-dev[bot] <65095814+ellipsis-dev[bot]@users.noreply.github.com>

* remove unused imports

* bump version

---------

Co-authored-by: ellipsis-dev[bot] <65095814+ellipsis-dev[bot]@users.noreply.github.com>
---
 graphiti_core/search/search_config.py  |  2 +-
 graphiti_core/search/search_helpers.py | 50 +++++++++++++++++---------
 pyproject.toml                         |  2 +-
 tests/test_graphiti_int.py             | 16 +++------
 4 files changed, 39 insertions(+), 31 deletions(-)

diff --git a/graphiti_core/search/search_config.py b/graphiti_core/search/search_config.py
index f5148842..9aa23daa 100644
--- a/graphiti_core/search/search_config.py
+++ b/graphiti_core/search/search_config.py
@@ -26,7 +26,7 @@ from graphiti_core.search.search_utils import (
     MAX_SEARCH_DEPTH,
 )
 
-DEFAULT_SEARCH_LIMIT = 20
+DEFAULT_SEARCH_LIMIT = 10
 
 
 class EdgeSearchMethod(Enum):
diff --git a/graphiti_core/search/search_helpers.py b/graphiti_core/search/search_helpers.py
index aae66bc3..a034847c 100644
--- a/graphiti_core/search/search_helpers.py
+++ b/graphiti_core/search/search_helpers.py
@@ -14,6 +14,8 @@ See the License for the specific language governing permissions and
 limitations under the License.
 """
 
+import json
+
 from graphiti_core.edges import EntityEdge
 from graphiti_core.search.search_config import SearchResults
 
@@ -25,23 +27,37 @@ def format_edge_date_range(edge: EntityEdge) -> str:
 
 def search_results_to_context_string(search_results: SearchResults) -> str:
     """Reformats a set of SearchResults into a single string to pass directly to an LLM as context"""
-    context_string = """FACTS and ENTITIES represent relevant context to the current conversation.
-                        COMMUNITIES represent a cluster of closely related entities.
+    fact_json = [
+        {
+            'fact': edge.fact,
+            'valid_at': str(edge.valid_at),
+            'invalid_at': str(edge.invalid_at or 'Present'),
+        }
+        for edge in search_results.edges
+    ]
+    entity_json = [
+        {'entity_name': node.name, 'summary': node.summary} for node in search_results.nodes
+    ]
+    community_json = [
+        {'community_name': community.name, 'summary': community.summary}
+        for community in search_results.communities
+    ]
 
-                        # These are the most relevant facts and their valid date ranges
-                        # format: FACT (Date range: from - to)
-                    """
-    context_string += '<FACTS>\n'
-    for edge in search_results.edges:
-        context_string += f'- {edge.fact} ({format_edge_date_range(edge)})\n'
-    context_string += '</FACTS>\n'
-    context_string += '<ENTITIES>\n'
-    for node in search_results.nodes:
-        context_string += f'- {node.name}: {node.summary}\n'
-    context_string += '</ENTITIES>\n'
-    context_string += '<COMMUNITIES>\n'
-    for community in search_results.communities:
-        context_string += f'- {community.name}: {community.summary}\n'
-    context_string += '</COMMUNITIES>\n'
+    context_string = f"""
+    FACTS and ENTITIES represent relevant context to the current conversation.
+    COMMUNITIES represent a cluster of closely related entities.
+
+    These are the most relevant facts and their valid and invalid dates. Facts are considered valid
+    between their valid_at and invalid_at dates. Facts with an invalid_at date of "Present" are considered valid.
+    <FACTS>
+    {json.dumps(fact_json, indent=12)}
+    </FACTS>
+    <ENTITIES>
+    {json.dumps(entity_json, indent=12)}
+    </ENTITIES>
+    <COMMUNITIES>
+    {json.dumps(community_json, indent=12)}
+    </COMMUNITIES>
+"""
 
     return context_string
diff --git a/pyproject.toml b/pyproject.toml
index 49ba6c53..1bd9e830 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -1,7 +1,7 @@
 [project]
 name = "graphiti-core"
 description = "A temporal graph building library"
-version = "0.9.5"
+version = "0.9.6"
 authors = [
     { "name" = "Paul Paliychuk", "email" = "paul@getzep.com" },
     { "name" = "Preston Rasmussen", "email" = "preston@getzep.com" },
diff --git a/tests/test_graphiti_int.py b/tests/test_graphiti_int.py
index cde8c6ed..6fd93900 100644
--- a/tests/test_graphiti_int.py
+++ b/tests/test_graphiti_int.py
@@ -26,8 +26,7 @@ from graphiti_core.edges import EntityEdge, EpisodicEdge
 from graphiti_core.graphiti import Graphiti
 from graphiti_core.helpers import semaphore_gather
 from graphiti_core.nodes import EntityNode, EpisodicNode
-from graphiti_core.search.search_config_recipes import COMBINED_HYBRID_SEARCH_CROSS_ENCODER
-from graphiti_core.search.search_filters import SearchFilters
+from graphiti_core.search.search_helpers import search_results_to_context_string
 
 pytestmark = pytest.mark.integration
 
@@ -66,18 +65,11 @@ async def test_graphiti_init():
     logger = setup_logging()
     graphiti = Graphiti(NEO4J_URI, NEO4j_USER, NEO4j_PASSWORD)
 
-    results = await graphiti._search(
-        'My name is Alice',
-        COMBINED_HYBRID_SEARCH_CROSS_ENCODER,
-        group_ids=['test'],
-        search_filter=SearchFilters(node_labels=['Entity']),
+    results = await graphiti.search_(
+        query="Who is Alice's friend?",
     )
 
-    pretty_results = {
-        'edges': [edge.fact for edge in results.edges],
-        'nodes': [node.name for node in results.nodes],
-        'communities': [community.name for community in results.communities],
-    }
+    pretty_results = search_results_to_context_string(results)
 
     logger.info(pretty_results)