update

Author: JqRrt

cozeloop/decorator/decorator.py

@@ -4,7 +4,7 @@
 from typing import Optional, Callable, Any, overload, Dict, Generic, Iterator, TypeVar, List, cast, AsyncIterator
 from functools import wraps
 
-from langchain_core.runnables import RunnableLambda
+from langchain_core.runnables import RunnableLambda, RunnableConfig
 
 from cozeloop import Client, Span, start_span
 from cozeloop.decorator.utils import is_async_func, is_gen_func, is_async_gen_func, is_class_func
@@ -322,8 +322,7 @@ def to_runnable(
         Decorator to be RunnableLambda.
 
         :param func: The function to be decorated, Requirements are as follows：
-                     1. The input parameter should only have one field, not multiple fields. You can use dict or Class as the input parameter.
-                     2. When the func is called, parameter config(RunnableConfig) is required, you must use the config containing cozeloop callback handler of 'current request', otherwise, the trace may be lost!
+                     1. When the func is called, parameter config(RunnableConfig) is required, you must use the config containing cozeloop callback handler of 'current request', otherwise, the trace may be lost!
 
         Examples:
             @to_runnable
@@ -338,15 +337,14 @@ def decorator(func: Callable):
 
             @wraps(func)
             def sync_wrapper(*args: Any, **kwargs: Any):
-                config = kwargs.pop("config", None)
+                config = _get_config(**kwargs)
                 res = None
                 try:
-                    inp = None
-                    if len(args) == 1:
-                        inp = args[0]
-                    else:
-                        inp = kwargs
-                    res = RunnableLambda(func).invoke(input=inp, config=config, **kwargs)
+                    inp = {
+                        "args": args,
+                        "kwargs": kwargs
+                    }
+                    res = RunnableLambda(_param_wrapped_func).invoke(input=inp, config=config, **kwargs)
                     if hasattr(res, "__iter__"):
                         return res
                 except StopIteration:
@@ -359,15 +357,14 @@ def sync_wrapper(*args: Any, **kwargs: Any):
 
             @wraps(func)
             async def async_wrapper(*args: Any, **kwargs: Any):
-                config = kwargs.pop("config", None)
+                config = _get_config(**kwargs)
                 res = None
                 try:
-                    inp = None
-                    if len(args) == 1:
-                        inp = args[0]
-                    else:
-                        inp = kwargs
-                    res = await RunnableLambda(func).ainvoke(input=inp, config=config, **kwargs)
+                    inp = {
+                        "args": args,
+                        "kwargs": kwargs
+                    }
+                    res = await RunnableLambda(_param_wrapped_func_async).ainvoke(input=inp, config=config, **kwargs)
                     if hasattr(res, "__aiter__"):
                         return res
                 except StopIteration:
@@ -385,14 +382,13 @@ async def async_wrapper(*args: Any, **kwargs: Any):
 
             @wraps(func)
             def gen_wrapper(*args: Any, **kwargs: Any):
-                config = kwargs.pop("config", None)
+                config = _get_config(**kwargs)
                 try:
-                    inp = None
-                    if len(args) == 1:
-                        inp = args[0]
-                    else:
-                        inp = kwargs
-                    gen = RunnableLambda(func).invoke(input=inp, config=config, **kwargs)
+                    inp = {
+                        "args": args,
+                        "kwargs": kwargs
+                    }
+                    gen = RunnableLambda(_param_wrapped_func).invoke(input=inp, config=config, **kwargs)
                     try:
                         for item in gen:
                             yield item
@@ -403,14 +399,13 @@ def gen_wrapper(*args: Any, **kwargs: Any):
 
             @wraps(func)
             async def async_gen_wrapper(*args: Any, **kwargs: Any):
-                config = kwargs.pop("config", None)
+                config = _get_config(**kwargs)
                 try:
-                    inp = None
-                    if len(args) == 1:
-                        inp = args[0]
-                    else:
-                        inp = kwargs
-                    gen = RunnableLambda(func).invoke(input=inp, config=config, **kwargs)
+                    inp = {
+                        "args": args,
+                        "kwargs": kwargs
+                    }
+                    gen = RunnableLambda(_param_wrapped_func_async).invoke(input=inp, config=config, **kwargs)
                     items = []
                     try:
                         async for item in gen:
@@ -428,6 +423,25 @@ async def async_gen_wrapper(*args: Any, **kwargs: Any):
                     else:
                         raise e
 
+            # for convert parameter
+            def _param_wrapped_func(input_dict: dict) -> Any:
+                args = input_dict.get("args", ())
+                kwargs = input_dict.get("kwargs", {})
+                return func(*args, **kwargs)
+
+            async def _param_wrapped_func_async(input_dict: dict) -> Any:
+                args = input_dict.get("args", ())
+                kwargs = input_dict.get("kwargs", {})
+                return await func(*args, **kwargs)
+
+            def _get_config(**kwargs: Any) -> RunnableConfig | None:
+                config = kwargs.pop("config", None)
+                if config is None:
+                    config = RunnableConfig(run_name=func.__name__)
+                    config['run_name'] = func.__name__
+                elif isinstance(config, dict):
+                    config['run_name'] = func.__name__
+                return config
 
             if is_async_gen_func(func):
                 return async_gen_wrapper