address feedback

llama_stack/distribution/routers/routers.py

@@ -222,8 +222,8 @@ class DatasetIORouter(DatasetIO):
             filter_condition=filter_condition,
         )
 
-    async def upload_rows(self, dataset_id: str, rows: List[Dict[str, Any]]) -> None:
-        return await self.routing_table.get_provider_impl(dataset_id).upload_rows(
+    async def append_rows(self, dataset_id: str, rows: List[Dict[str, Any]]) -> None:
+        return await self.routing_table.get_provider_impl(dataset_id).append_rows(
             dataset_id=dataset_id,
             rows=rows,
         )

llama_stack/distribution/tracing.py

@@ -24,7 +24,7 @@ def serialize_value(value: Any) -> str:
             return value.model_dump_json()
         elif isinstance(value, list) and value and isinstance(value[0], BaseModel):
             return json.dumps([item.model_dump_json() for item in value])
-        elif hasattr(value, "to_dict"):  # For objects with to_dict method
+        elif hasattr(value, "to_dict"):
             return json.dumps(value.to_dict())
         elif isinstance(value, (dict, list, int, float, str, bool)):
             return json.dumps(value)
@@ -34,21 +34,6 @@ def serialize_value(value: Any) -> str:
         return str(value)
 
 
-def traced(input: str = None):
-    """
-    A method decorator that enables tracing with input and output capture.
-
-    Args:
-        input: Name of the input parameter to capture in traces
-    """
-
-    def decorator(method: Callable) -> Callable:
-        method._trace_input = input
-        return method
-
-    return decorator
-
-
 def trace_protocol(cls: Type[T]) -> Type[T]:
     """
     A class decorator that automatically traces all methods in a protocol/base class
@@ -59,22 +44,6 @@ def trace_protocol(cls: Type[T]) -> Type[T]:
         is_async = asyncio.iscoroutinefunction(method)
         is_async_gen = inspect.isasyncgenfunction(method)
 
-        def get_traced_input(args: tuple, kwargs: dict) -> dict:
-            trace_input = getattr(method, "_trace_input", None)
-            if not trace_input:
-                return {}
-
-            # Get the mapping of parameter names to values
-            sig = inspect.signature(method)
-            bound_args = sig.bind(None, *args, **kwargs)  # None for self
-            bound_args.apply_defaults()
-            params = dict(list(bound_args.arguments.items())[1:])  # Skip 'self'
-
-            # Return the input value if the key exists
-            if trace_input in params:
-                return {"input": serialize_value(params[trace_input])}
-            return {}
-
         def create_span_context(self: Any, *args: Any, **kwargs: Any) -> tuple:
             class_name = self.__class__.__name__
             method_name = method.__name__
@@ -87,7 +56,6 @@ def trace_protocol(cls: Type[T]) -> Type[T]:
                 "method": method_name,
                 "type": span_type,
                 "args": serialize_value(args),
-                **get_traced_input(args, kwargs),
             }
 
             return class_name, method_name, span_attributes
@@ -145,33 +113,16 @@ def trace_protocol(cls: Type[T]) -> Type[T]:
         else:
             return sync_wrapper
 
-    # Store the original __init_subclass__ if it exists
     original_init_subclass = getattr(cls, "__init_subclass__", None)
 
-    # Define a new __init_subclass__ to handle child classes
     def __init_subclass__(cls_child, **kwargs):  # noqa: N807
-        # Call original __init_subclass__ if it exists
         if original_init_subclass:
             original_init_subclass(**kwargs)
 
-        traced_methods = {}
-        for parent in cls_child.__mro__[1:]:  # Skip the class itself
-            for name, method in vars(parent).items():
-                if inspect.isfunction(method) and getattr(
-                    method, "_trace_input", None
-                ):  # noqa: B009
-                    traced_methods[name] = getattr(method, "_trace_input")  # noqa: B009
-
-        # Trace child class methods if their name matches a traced parent method
         for name, method in vars(cls_child).items():
             if inspect.isfunction(method) and not name.startswith("_"):
-                if name in traced_methods:
-                    # Copy the trace configuration from the parent
-                    setattr(method, "_trace_input", traced_methods[name])  # noqa: B010
-
                 setattr(cls_child, name, trace_method(method))  # noqa: B010
 
-    # Set the new __init_subclass__
     cls.__init_subclass__ = classmethod(__init_subclass__)
 
     return cls