diff --git a/langfuse/_client/observe.py b/langfuse/_client/observe.py index 5aa88ce2e..73e11782e 100644 --- a/langfuse/_client/observe.py +++ b/langfuse/_client/observe.py @@ -2,7 +2,7 @@ import contextvars import inspect import os -from functools import wraps +from functools import partial, wraps from typing import ( Any, AsyncGenerator, @@ -60,6 +60,12 @@ def _is_coroutine_function(func: Any) -> bool: ) +def _get_function_name(func: Callable[..., Any]) -> str: + while isinstance(func, partial) and not hasattr(func, "__name__"): + func = func.func + return func.__name__ + + class LangfuseDecorator: """Implementation of the @observe decorator for seamless Langfuse tracing integration. @@ -275,7 +281,7 @@ async def async_wrapper(*args: Tuple[Any], **kwargs: Dict[str, Any]) -> Any: if trace_id else None ) - final_name = name or func.__name__ + final_name = name or _get_function_name(func) input = ( self._get_input_from_func_args( is_method=self._is_method(func), @@ -366,7 +372,7 @@ def sync_wrapper(*args: Any, **kwargs: Any) -> Any: if trace_id else None ) - final_name = name or func.__name__ + final_name = name or _get_function_name(func) input = ( self._get_input_from_func_args( is_method=self._is_method(func), diff --git a/tests/unit/test_observe.py b/tests/unit/test_observe.py index 21ff1418d..1d3ad1ce6 100644 --- a/tests/unit/test_observe.py +++ b/tests/unit/test_observe.py @@ -3,6 +3,7 @@ import gc import inspect import json +from functools import partial from typing import Any, AsyncGenerator, Generator, cast import pytest @@ -33,6 +34,76 @@ def _finished_spans_by_name(memory_exporter: Any, name: str) -> list[Any]: return [span for span in memory_exporter.get_finished_spans() if span.name == name] +@pytest.mark.parametrize("name", [None, "custom_name"]) +def test_observe_partial_function( + langfuse_memory_client: Any, memory_exporter: Any, name: str | None +) -> None: + def add(first: int, second: int, *, extra: int = 0) -> int: + return first + second + extra + + observed = observe(name=name)(partial(add, 2, extra=4)) + + assert observed(3) == 9 + + langfuse_memory_client.flush() + + spans = memory_exporter.get_finished_spans() + assert len(spans) == 1 + assert spans[0].name == (name or "add") + assert ( + json.loads(spans[0].attributes[LangfuseOtelSpanAttributes.OBSERVATION_OUTPUT]) + == 9 + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("name", [None, "custom_name"]) +async def test_observe_async_partial_function( + langfuse_memory_client: Any, memory_exporter: Any, name: str | None +) -> None: + async def add(first: int, second: int, *, extra: int = 0) -> int: + await asyncio.sleep(0) + return first + second + extra + + observed = observe(name=name)(partial(add, 2, extra=4)) + + assert inspect.iscoroutinefunction(observed) + assert await observed(3) == 9 + + langfuse_memory_client.flush() + + spans = memory_exporter.get_finished_spans() + assert len(spans) == 1 + assert spans[0].name == (name or "add") + assert ( + json.loads(spans[0].attributes[LangfuseOtelSpanAttributes.OBSERVATION_OUTPUT]) + == 9 + ) + + +@pytest.mark.parametrize("function_name", [None, "custom_add"]) +def test_observe_nested_partial_function( + langfuse_memory_client: Any, memory_exporter: Any, function_name: str | None +) -> None: + def add(first: int, second: int) -> int: + return first + second + + inner = partial(add, 2) + # Partials with attributes are not flattened when wrapped in another partial. + cast(Any, inner).description = "bound first argument" + if function_name: + cast(Any, inner).__name__ = function_name + observed = observe(partial(inner, 3)) + + assert observed() == 5 + + langfuse_memory_client.flush() + + spans = memory_exporter.get_finished_spans() + assert len(spans) == 1 + assert spans[0].name == (function_name or "add") + + @pytest.mark.asyncio async def test_capture_output_false_preserves_type_when_current_span_is_updated( langfuse_memory_client: Any, memory_exporter: Any