Working on cool stuff

test_langgraph/tools.py

@@ -17,18 +17,21 @@ def search(query: str):
         return "It's 60 degrees and foggy."
     return "It's 90 degrees and sunny."
 
-def wrap(class_, method):
-    logger.debug("Wrapping %s.%s", class_.__name__, method.__name__)
+def wrap_method(class_, method):
+    logger.info("Wrapping %s.%s", class_.__name__, method.__name__)
     is_iterator = str(method.__annotations__.get('return', '')).startswith('collections.abc.Iterator')
-    def wrapper(input):
-        logger.info("AI called %s.%s(%s)", class_.__name__, method.__name__, input)
-        result = method(input)
+
+    def wrapper(input_value):
+        if isinstance(input_value, dict):
+            logger.warning("Silently converting from dict to plain value!")
+            input_value = next(input_value.values())
+        logger.info("AI called %s.%s(%s)", class_.__name__, method.__name__, repr(input_value))
+        result = method(input_value)
         if is_iterator:
             result = list(result)
         return result
 
-    #wrapper.__name__ = f'{class_.__name__}.{method.__name__}'
-    wrapper.__name__ = f'{method.__name__}'
+    wrapper.__name__ = f'{class_.__name__}.{method.__name__}'
     wrapper.__doc__ = method.__doc__
     wrapper.__annotations__ = method.__annotations__
     return tool(wrapper)
@@ -42,7 +45,7 @@ def wrap_all_methods_on_client(obj):
             continue
         if method.__doc__ is None:
             continue
-        yield wrap(obj.__class__, method)
+        yield wrap_method(obj.__class__, method)
 
 def get_tools():
     session = requests_cache.CachedSession('output/test.sqlite')