diff --git a/runtime/prompty/prompty/tracer.py b/runtime/prompty/prompty/tracer.py index 5350df4..8c73df2 100644 --- a/runtime/prompty/prompty/tracer.py +++ b/runtime/prompty/prompty/tracer.py @@ -134,7 +134,17 @@ def _trace_sync( @wraps(func) def wrapper(*args, **kwargs): name, signature = _name(func, args) + altname: str = None + # special case + if "name" in okwargs: + altname = name + name = okwargs["name"] + del okwargs["name"] + with Tracer.start(name) as trace: + if altname != None: + trace("function", altname) + trace("signature", signature) # support arbitrary keyword @@ -178,10 +188,20 @@ def _trace_async( @wraps(func) async def wrapper(*args, **kwargs): name, signature = _name(func, args) + altname: str = None + # special case + if "name" in okwargs: + altname = name + name = okwargs["name"] + del okwargs["name"] + with Tracer.start(name) as trace: + if altname != None: + trace("function", altname) + trace("signature", signature) - # support arbitrary keyword + # support arbitrary keyword # arguments for trace decorator for k, v in okwargs.items(): trace(k, to_dict(v)) diff --git a/runtime/prompty/tests/test_tracing.py b/runtime/prompty/tests/test_tracing.py index 118e2aa..223380f 100644 --- a/runtime/prompty/tests/test_tracing.py +++ b/runtime/prompty/tests/test_tracing.py @@ -308,3 +308,18 @@ def test_tracing_attributes(): with Tracer.start("Test3", {"signature": "9", "ten": 10}) as trace: trace("inputs", 11) trace(Tracer.RESULT, 12) + + +@trace(name="OTHER") +def test_named_tracer(): + with Tracer.start("Test1", {Tracer.SIGNATURE: "test1", "two": 2}) as trace: + trace(Tracer.INPUTS, 3) + trace(Tracer.RESULT, 4) + + +@pytest.mark.asyncio +@trace(name="OTHER") +async def test_named_tracer_async(): + with Tracer.start("Test1", {Tracer.SIGNATURE: "test1", "two": 2}) as trace: + trace(Tracer.INPUTS, 3) + trace(Tracer.RESULT, 4)