From f34a737d0c6d4bcfe941b65b5d83b4a5bc03fb8e Mon Sep 17 00:00:00 2001
From: Jon Michael Aanes <jonjmaa@gmail.com>
Date: Wed, 18 Jun 2025 22:24:01 +0200
Subject: [PATCH] Configure wrapping

---
 test_langgraph/tools.py | 16 ++++++++++------
 1 file changed, 10 insertions(+), 6 deletions(-)

diff --git a/test_langgraph/tools.py b/test_langgraph/tools.py
index f8bc63f..5a57379 100644
--- a/test_langgraph/tools.py
+++ b/test_langgraph/tools.py
@@ -60,10 +60,11 @@ RETURN_FORMATS = {
 
 RETURN_FORMAT = 'json'
 
-MAX_TOOL_RESULT_LEN = 1000
+
 APPEND_RESULT_TYPE_DOCS = True
 
-def wrap_method(class_, method):
+def wrap_method(class_, method, max_tool_result_len = 1_000_000_000,
+                provide_instructions: bool = False):
     logger.info('Wrapping %s.%s', class_.__name__, method.__name__)
     return_type = method.__annotations__.get('return', '')
     is_iterator = str(return_type).startswith(
@@ -91,10 +92,13 @@ def wrap_method(class_, method):
                 repr(input_value),
             )
             raise
-        if len(result_str) > MAX_TOOL_RESULT_LEN:
-            result_str = result_str[:MAX_TOOL_RESULT_LEN] + ' (remaining tool result elicited...)'
-        if APPEND_RESULT_TYPE_DOCS and (return_docs := getattr(return_type, '__doc__', None)):
-            result_str = result_str+'\n'+return_docs
+        if len(result_str) > max_tool_result_len:
+            result_str = result_str[:max_tool_result_len]
+            if provide_instructions:
+                result_str += '\nremaining tool result elicited...'
+        if provide_instructions:
+            if APPEND_RESULT_TYPE_DOCS and (return_docs := getattr(return_type, '__doc__', None)):
+                result_str += '\n'+return_docs
         return result_str
 
     wrapper.__name__ = f'{class_.__name__}.{method.__name__}'