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__}'