Decorators support async functions

Andrei Bratu · Andrei Bratu · commit 7becdeba9aff · 2025-03-19T10:25:27.000Z
diff --git a/src/humanloop/client.py b/src/humanloop/client.py
@@ -1,3 +1,4 @@
+from importlib import metadata
 import os
 import typing
 from typing import Any, List, Optional, Sequence
@@ -9,6 +10,7 @@
 
 from humanloop.core.client_wrapper import SyncClientWrapper
 
+from humanloop.error import HumanloopRuntimeError
 from humanloop.evals import run_eval
 from humanloop.evals.types import Dataset, Evaluator, EvaluatorCheck, File
 
@@ -68,6 +70,41 @@ def run(
         )
 
 
+class HumanloopTracerSingleton:
+    _instance = None
+
+    def __init__(self, hl_client_headers: dict[str, str], hl_client_base_url: str):
+        if HumanloopTracerSingleton._instance is not None:
+            raise HumanloopRuntimeError("Internal error: HumanloopTracerSingleton already initialized")
+
+        self.tracer_provider = TracerProvider(
+            resource=Resource(
+                attributes={
+                    "service.name": "humanloop-python-sdk",
+                    "service.version": metadata.version("humanloop"),
+                }
+            )
+        )
+        self.tracer_provider.add_span_processor(
+            HumanloopSpanProcessor(
+                exporter=HumanloopSpanExporter(
+                    hl_client_headers=hl_client_headers,
+                    hl_client_base_url=hl_client_base_url,
+                )
+            )
+        )
+
+        instrument_provider(provider=self.tracer_provider)
+
+        self.tracer = self.tracer_provider.get_tracer("humanloop.sdk")
+
+    @classmethod
+    def get_instance(cls, hl_client_headers: dict[str, str], hl_client_base_url: str):
+        if cls._instance is None:
+            cls._instance = cls(hl_client_headers, hl_client_base_url)
+        return cls._instance
+
+
 class Humanloop(BaseHumanloop):
     """
     See docstring of :class:`BaseHumanloop`.
@@ -117,27 +154,12 @@ def __init__(
         self.flows = overload_log(client=self.flows)
         self.tools = overload_log(client=self.tools)
 
-        if opentelemetry_tracer_provider is not None:
-            self._tracer_provider = opentelemetry_tracer_provider
-        else:
-            self._tracer_provider = TracerProvider(
-                resource=Resource(
-                    attributes={
-                        "instrumentor": "humanloop.sdk",
-                    }
-                ),
-            )
-        instrument_provider(provider=self._tracer_provider)
-        self._tracer_provider.add_span_processor(
-            HumanloopSpanProcessor(exporter=HumanloopSpanExporter(client=self)),
+        # Initialize the tracer singleton
+        self._tracer_singleton = HumanloopTracerSingleton.get_instance(
+            hl_client_headers=self._client_wrapper.get_headers(),
+            hl_client_base_url=self._client_wrapper._base_url,
         )
-
-        if opentelemetry_tracer is None:
-            self._opentelemetry_tracer = self._tracer_provider.get_tracer(
-                "humanloop.sdk"
-            )
-        else:
-            self._opentelemetry_tracer = opentelemetry_tracer
+        self._opentelemetry_tracer = self._tracer_singleton.tracer
 
     def prompt(
         self,
diff --git a/src/humanloop/decorators/flow.py b/src/humanloop/decorators/flow.py
@@ -1,3 +1,4 @@
+import asyncio
 import logging
 from functools import wraps
 from typing import Any, Callable, Optional, TypeVar
@@ -121,13 +122,99 @@ def wrapper(*args: P.args, **kwargs: P.kwargs) -> Optional[R]:
                         # Return the output of the decorated function
                         return func_output  # type: ignore [return-value]
 
+        @wraps(func)
+        async def async_wrapper(*args: P.args, **kwargs: P.kwargs) -> Optional[R]:
+            span: Span
+            with set_decorator_context(
+                DecoratorContext(
+                    path=decorator_path,
+                    type="flow",
+                    version=flow_kernel,
+                )
+            ) as decorator_context:
+                with opentelemetry_tracer.start_as_current_span(HUMANLOOP_FLOW_SPAN_NAME) as span:  # type: ignore
+                    span.set_attribute(HUMANLOOP_FILE_PATH_KEY, decorator_path)
+                    span.set_attribute(HUMANLOOP_FILE_TYPE_KEY, file_type)
+                    trace_id = get_trace_id()
+                    func_args = bind_args(func, args, kwargs)
+
+                    # Create the trace ahead so we have a parent ID to reference
+                    init_log_inputs = {
+                        "inputs": {k: v for k, v in func_args.items() if k != "messages"},
+                        "messages": func_args.get("messages"),
+                        "trace_parent_id": trace_id,
+                    }
+                    this_flow_log: FlowLogResponse = client.flows._log(  # type: ignore [attr-defined]
+                        path=decorator_context.path,
+                        flow=decorator_context.version,
+                        log_status="incomplete",
+                        **init_log_inputs,
+                    )
+
+                    with set_trace_id(this_flow_log.id):
+                        func_output: Optional[R]
+                        log_output: Optional[str]
+                        log_error: Optional[str]
+                        log_output_message: Optional[ChatMessage]
+                        try:
+                            func_output = await func(*args, **kwargs)
+                            if (
+                                isinstance(func_output, dict)
+                                and len(func_output.keys()) == 2
+                                and "role" in func_output
+                                and "content" in func_output
+                            ):
+                                log_output_message = func_output  # type: ignore [assignment]
+                                log_output = None
+                            else:
+                                log_output = process_output(func=func, output=func_output)
+                                log_output_message = None
+                            log_error = None
+                        except HumanloopRuntimeError as e:
+                            # Critical error, re-raise
+                            client.logs.delete(id=this_flow_log.id)
+                            span.record_exception(e)
+                            raise e
+                        except Exception as e:
+                            logger.error(f"Error calling {func.__name__}: {e}")
+                            log_output = None
+                            log_output_message = None
+                            log_error = str(e)
+                            func_output = None
+
+                        updated_flow_log = {
+                            "log_status": "complete",
+                            "output": log_output,
+                            "error": log_error,
+                            "output_message": log_output_message,
+                            "id": this_flow_log.id,
+                        }
+                        # Write the Flow Log to the Span on HL_LOG_OT_KEY
+                        write_to_opentelemetry_span(
+                            span=span,  # type: ignore [arg-type]
+                            key=HUMANLOOP_LOG_KEY,
+                            value=updated_flow_log,  # type: ignore
+                        )
+                        # Return the output of the decorated function
+                        return func_output  # type: ignore [return-value]
+
+        # If the decorated function is an async function, return the async wrapper
+        if asyncio.iscoroutinefunction(func):
+            async_wrapper.file = File(  # type: ignore
+                path=decorator_path,
+                type=file_type,  # type: ignore [arg-type, typeddict-item]
+                version=FlowDict(**flow_kernel),  # type: ignore
+                callable=async_wrapper,
+            )
+            return async_wrapper
+
+        # If the decorated function is a sync function, return the sync wrapper
         wrapper.file = File(  # type: ignore
             path=decorator_path,
             type=file_type,  # type: ignore [arg-type, typeddict-item]
             version=FlowDict(**flow_kernel),  # type: ignore
             callable=wrapper,
         )
-
         return wrapper
 
     return decorator
diff --git a/src/humanloop/decorators/tool.py b/src/humanloop/decorators/tool.py
@@ -1,3 +1,4 @@
+import asyncio
 import builtins
 import inspect
 import logging
@@ -112,13 +113,80 @@ def wrapper(*args: P.args, **kwargs: P.kwargs) -> Optional[R]:
                 # Return the output of the decorated function
                 return func_output
 
+        @wraps(func)
+        async def async_wrapper(*args: P.args, **kwargs: P.kwargs) -> Optional[R]:
+            evaluation_context = get_evaluation_context()
+            if evaluation_context is not None:
+                if evaluation_context.path == path:
+                    raise HumanloopRuntimeError("Tools cannot be evaluated with the `evaluations.run()` utility.")
+            with opentelemetry_tracer.start_as_current_span("humanloop.tool") as span:
+                # Write the Tool Kernel to the Span on HL_FILE_OT_KEY
+                write_to_opentelemetry_span(
+                    span=span,  # type: ignore [arg-type]
+                    key=HUMANLOOP_FILE_KEY,
+                    value=tool_kernel,  # type: ignore [arg-type]
+                )
+                span.set_attribute(HUMANLOOP_FILE_PATH_KEY, path)
+                span.set_attribute(HUMANLOOP_FILE_TYPE_KEY, file_type)
+
+                log_inputs: dict[str, Any] = bind_args(func, args, kwargs)
+                log_error: Optional[str]
+                log_output: str
+
+                func_output: Optional[R]
+                try:
+                    func_output = await func(*args, **kwargs)
+                    log_output = process_output(
+                        func=func,
+                        output=func_output,
+                    )
+                    log_error = None
+                except HumanloopRuntimeError as e:
+                    # Critical error, re-raise
+                    raise e
+                except Exception as e:
+                    logger.error(f"Error calling {func.__name__}: {e}")
+                    output = None
+                    log_output = process_output(
+                        func=func,
+                        output=output,
+                    )
+                    log_error = str(e)
+
+                # Populate Tool Log attributes
+                tool_log = {
+                    "inputs": log_inputs,
+                    "output": log_output,
+                    "error": log_error,
+                    "trace_parent_id": get_trace_id(),
+                }
+                # Write the Tool Log to the Span on HL_LOG_OT_KEY
+                write_to_opentelemetry_span(
+                    span=span,  # type: ignore [arg-type]
+                    key=HUMANLOOP_LOG_KEY,
+                    value=tool_log,  # type: ignore [arg-type]
+                )
+
+                # Return the output of the decorated function
+                return func_output
+
+        # If the decorated function is an async function, return the async wrapper
+        if asyncio.iscoroutinefunction(func):
+            async_wrapper.file = File(  # type: ignore
+                path=path,
+                type=file_type,  # type: ignore [arg-type, typeddict-item]
+                version=tool_kernel,
+                callable=async_wrapper,
+            )
+            return async_wrapper
+
+        # If the decorated function is a sync function, return the sync wrapper
         wrapper.file = File(  # type: ignore
             path=path,
             type=file_type,  # type: ignore [arg-type, typeddict-item]
             version=tool_kernel,
             callable=wrapper,
         )
-
         return wrapper
 
     return decorator
diff --git a/src/humanloop/otel/exporter/__init__.py b/src/humanloop/otel/exporter/__init__.py
@@ -1,7 +1,6 @@
 import logging
 
 import time
-import typing
 from queue import Empty as EmptyQueue
 from queue import Queue
 from threading import Thread
@@ -26,27 +25,25 @@
 )
 
 
-if typing.TYPE_CHECKING:
-    from humanloop.client import Humanloop
-
-
 logger = logging.getLogger("humanloop.sdk")
 
 
 class HumanloopSpanExporter(SpanExporter):
-    DEFAULT_NUMBER_THREADS = 1
+    DEFAULT_NUMBER_THREADS = 4
 
     def __init__(
         self,
-        client: "Humanloop",
+        hl_client_headers: dict[str, str],
+        hl_client_base_url: str,
         worker_threads: Optional[int] = None,
     ) -> None:
         """Upload Spans created by SDK decorators to Humanloop.
 
         Spans not created by Humanloop SDK decorators will be ignored.
         """
         super().__init__()
-        self._client = client
+        self._hl_client_headers = hl_client_headers
+        self._base_url = hl_client_base_url
         # Work queue for the threads uploading the spans
         self._upload_queue: Queue = Queue()
         # Worker threads to export the spans
@@ -135,12 +132,9 @@ def _do_work(self):
                 continue
 
             span_to_export, eval_context_callback = thread_args
-
             response = requests.post(
-                f"{self._client._client_wrapper.get_base_url()}/import/otel/v1/traces",
-                headers={
-                    **self._client._client_wrapper.get_headers(),
-                },
+                f"{self._base_url}/import/otel/v1/traces",
+                headers=self._hl_client_headers,
                 data=serialize_span(span_to_export),
             )
             if response.status_code != 200:
