Error handling in decorators

Author: Andrei Bratu

src/humanloop/context.py

@@ -1,12 +1,12 @@
 from contextlib import contextmanager
 from dataclasses import dataclass
 import threading
-from typing import Callable, Generator, Optional
+from typing import Any, Callable, Generator, Literal, Optional
 from opentelemetry import context as context_api
 
 from humanloop.otel.constants import (
     HUMANLOOP_CONTEXT_EVALUATION,
-    HUMANLOOP_CONTEXT_PROMPT,
+    HUMANLOOP_CONTEXT_DECORATOR,
     HUMANLOOP_CONTEXT_TRACE_ID,
 )
 
@@ -25,14 +25,15 @@ def set_trace_id(flow_log_id: str) -> Generator[None, None, None]:
 
 
 @dataclass
-class PromptContext:
+class DecoratorContext:
     path: str
-    template: Optional[str]
+    type: Literal["prompt", "tool", "flow"]
+    version: dict[str, Optional[Any]]
 
 
 @contextmanager
-def set_prompt_context(prompt_context: PromptContext) -> Generator[None, None, None]:
-    key = hash((HUMANLOOP_CONTEXT_PROMPT, threading.get_ident()))
+def set_decorator_context(prompt_context: DecoratorContext) -> Generator[None, None, None]:
+    key = hash((HUMANLOOP_CONTEXT_DECORATOR, threading.get_ident()))
     reset_token = context_api.attach(
         context_api.set_value(
             key=key,
@@ -43,15 +44,15 @@ def set_prompt_context(prompt_context: PromptContext) -> Generator[None, None, N
     context_api.detach(token=reset_token)
 
 
-def get_prompt_context() -> Optional[PromptContext]:
-    key = hash((HUMANLOOP_CONTEXT_PROMPT, threading.get_ident()))
+def get_decorator_context() -> Optional[DecoratorContext]:
+    key = hash((HUMANLOOP_CONTEXT_DECORATOR, threading.get_ident()))
     return context_api.get_value(key)
 
 
 class EvaluationContext:
     source_datapoint_id: str
     run_id: str
-    callback: Callable[[str], None]
+    logging_callback: Callable[[str], None]
     file_id: str
     path: str
     logging_counter: int
@@ -60,13 +61,13 @@ def __init__(
         self,
         source_datapoint_id: str,
         run_id: str,
-        callback: Callable[[str], None],
+        logging_callback: Callable[[str], None],
         file_id: str,
         path: str,
     ):
         self.source_datapoint_id = source_datapoint_id
         self.run_id = run_id
-        self.callback = callback
+        self.logging_callback = logging_callback
         self.file_id = file_id
         self.path = path
         self.logging_counter = 0

src/humanloop/eval_utils/run.py

@@ -67,9 +67,15 @@
 from humanloop.types.datapoint_response import DatapointResponse
 from humanloop.types.dataset_response import DatasetResponse
 from humanloop.types.evaluation_run_response import EvaluationRunResponse
+from humanloop.types.evaluator_log_response import EvaluatorLogResponse
+from humanloop.types.flow_log_response import FlowLogResponse
+from humanloop.types.log_response import LogResponse
+from humanloop.types.prompt_log_response import PromptLogResponse
 from humanloop.types.run_stats_response import RunStatsResponse
 from pydantic import ValidationError
 
+from humanloop.types.tool_log_response import ToolLogResponse
+
 if typing.TYPE_CHECKING:
     from humanloop.client import BaseHumanloop
 
@@ -99,11 +105,13 @@
 CLIENT_TYPE = TypeVar("CLIENT_TYPE", PromptsClient, ToolsClient, FlowsClient, EvaluatorsClient)
 
 
-class HumanloopUtilityError(Exception):
-    def __init__(self, message):
+class HumanloopDecoratorError(Exception):
+    def __init__(self, message: Optional[str] = None):
         self.message = message
 
     def __str__(self):
+        if self.message is None:
+            return super().__str__()
         return self.message
 
 
@@ -202,7 +210,7 @@ def upload_callback(log_id: str):
             with set_evaluation_context(
                 EvaluationContext(
                     source_datapoint_id=dp.id,
-                    callback=upload_callback,
+                    logging_callback=upload_callback,
                     file_id=hl_file.id,
                     run_id=run.id,
                     path=hl_file.path,
@@ -219,17 +227,25 @@ def upload_callback(log_id: str):
                 try:
                     output = _call_function(function_, hl_file.type, dp)
                     evaluation_context = get_evaluation_context()
-                    if not evaluation_context.logging_counter == 0:
+                    if evaluation_context is None:
+                        raise HumanloopDecoratorError(
+                            "Internal error: evaluation context is not set while processing a datapoint."
+                        )
+                    if evaluation_context.logging_counter == 0:
                         # function_ did not Log against the source_datapoint_id/ run_id pair
                         # so we need to create a Log
-                        log_func(
+                        log = log_func(
                             inputs=dp.inputs,
                             output=output,
                             start_time=start_time,
                             end_time=datetime.now(),
                             source_datapoint_id=dp.id,
                             run_id=run.id,
                         )
+                        evaluation_context.logging_counter += 1
+                        evaluation_context.logging_callback(log.id)
+                except HumanloopDecoratorError as e:
+                    raise e
                 except Exception as e:
                     log_func(
                         inputs=dp.inputs,
@@ -648,13 +664,49 @@ def _call_function(
     return output
 
 
+def _get_log_func(
+    client: "BaseHumanloop",
+    file_type: Literal["flow"],
+    file_id: str,
+    version_id: str,
+    run_id: str,
+) -> Callable[..., FlowLogResponse]: ...
+
+
+def _get_log_func(
+    client: "BaseHumanloop",
+    file_type: Literal["prompt"],
+    file_id: str,
+    version_id: str,
+    run_id: str,
+) -> Callable[..., PromptLogResponse]: ...
+
+
+def _get_log_func(
+    client: "BaseHumanloop",
+    file_type: Literal["tool"],
+    file_id: str,
+    version_id: str,
+    run_id: str,
+) -> Callable[..., ToolLogResponse]: ...
+
+
+def _get_log_func(
+    client: "BaseHumanloop",
+    file_type: Literal["evaluator"],
+    file_id: str,
+    version_id: str,
+    run_id: str,
+) -> Callable[..., EvaluatorLogResponse]: ...
+
+
 def _get_log_func(
     client: "BaseHumanloop",
     file_type: FileType,
     file_id: str,
     version_id: str,
     run_id: str,
-) -> Callable:
+) -> Callable[..., LogResponse]:
     """Returns the appropriate log function pre-filled with common parameters."""
     log_request = {
         # TODO: why does the Log `id` field refer to the file ID in the API?

src/humanloop/otel/constants.py

@@ -6,6 +6,6 @@
 HUMANLOOP_FILE_TYPE_KEY = "humanloop.file.type"
 HUMANLOOP_PATH_KEY = "humanloop.file.path"
 # Opentelemetry context
-HUMANLOOP_CONTEXT_PROMPT = "humanloop.context.prompt"
+HUMANLOOP_CONTEXT_DECORATOR = "humanloop.context.decorator"
 HUMANLOOP_CONTEXT_TRACE_ID = "humanloop.context.flow.trace_id"
 HUMANLOOP_CONTEXT_EVALUATION = "humanloop.context.evaluation"

src/humanloop/otel/exporter.py

@@ -201,6 +201,6 @@ def _do_work(self):
             else:
                 if evaluation_context and file_path == evaluation_context.path:
                     log_id = response.json()["records"][0]["log_id"]
-                    evaluation_context.callback(log_id)
+                    evaluation_context.logging_callback(log_id)
 
             self._upload_queue.task_done()

src/humanloop/overload.py

@@ -1,12 +1,16 @@
+import inspect
 import logging
 import types
 from typing import TypeVar, Union
 import typing
 
-from humanloop.context import get_trace_id
-from humanloop.eval_utils.run import HumanloopUtilityError
+from humanloop.context import get_decorator_context, get_trace_id
+from humanloop.eval_utils.run import HumanloopDecoratorError
 
+from humanloop.evaluators.client import EvaluatorsClient
+from humanloop.flows.client import FlowsClient
 from humanloop.prompts.client import PromptsClient
+from humanloop.tools.client import ToolsClient
 from humanloop.types.create_evaluator_log_response import CreateEvaluatorLogResponse
 from humanloop.types.create_flow_log_response import CreateFlowLogResponse
 from humanloop.types.create_prompt_log_response import CreatePromptLogResponse
@@ -16,7 +20,7 @@
 logger = logging.getLogger("humanloop.sdk")
 
 
-CLIENT_TYPE = TypeVar("CLIENT_TYPE")
+CLIENT_TYPE = TypeVar("CLIENT_TYPE", PromptsClient, FlowsClient, EvaluatorsClient, ToolsClient)
 
 
 def overload_log(client: CLIENT_TYPE) -> CLIENT_TYPE:
@@ -41,20 +45,29 @@ def _overload_log(
         CreateEvaluatorLogResponse,
     ]:
         trace_id = get_trace_id()
+        if trace_id is not None and type(client) is FlowsClient:
+            context = get_decorator_context()
+            if context is None:
+                raise HumanloopDecoratorError("Internal error: trace_id context is set outside a decorator context.")
+            raise HumanloopDecoratorError(
+                f"Using flows.log() in this context is not allowed at line {inspect.currentframe().f_lineno}: "
+                f"Flow decorator for File {context.path} manages the tracing and trace completion."
+            )
         if trace_id is not None:
             if "trace_parent_id" in kwargs:
-                # TODO: revisit
-                logger.warning("Overriding trace_parent_id argument")
+                logger.warning(
+                    "Ignoring trace_parent_id argument at line %d: the Flow decorator manages tracing.",
+                    inspect.currentframe().f_lineno,
+                )
             kwargs = {
                 **kwargs,
                 "trace_parent_id": trace_id,
             }
         try:
             response = self._log(**kwargs)
         except Exception as e:
-            # TODO handle
-            # TODO: Bug found in backend: not specifying a model 400s but creates a File
-            raise HumanloopUtilityError(message=str(e)) from e
+            # Re-raising as HumanloopDecoratorError so the decorators don't catch it
+            raise HumanloopDecoratorError from e
 
         return response
 
@@ -73,8 +86,10 @@ def _overload_call(self, **kwargs) -> PromptCallResponse:
         trace_id = get_trace_id()
         if trace_id is not None:
             if "trace_parent_id" in kwargs:
-                # TODO: revisit
-                logger.warning("Overriding trace_parent_id argument")
+                logger.warning(
+                    "Ignoring trace_parent_id argument at line %d: the Flow decorator manages tracing.",
+                    inspect.currentframe().f_lineno,
+                )
             kwargs = {
                 **kwargs,
                 "trace_parent_id": trace_id,
@@ -84,9 +99,8 @@ def _overload_call(self, **kwargs) -> PromptCallResponse:
             response = self._call(**kwargs)
             response = typing.cast(PromptCallResponse, response)
         except Exception as e:
-            # TODO handle
-            # TODO: Bug found in backend: not specifying a model 400s but creates a File
-            raise HumanloopUtilityError(message=str(e)) from e
+            # Re-raising as HumanloopDecoratorError so the decorators don't catch it
+            raise HumanloopDecoratorError from e
 
         return response
 

src/humanloop/utilities/flow.py

@@ -4,11 +4,10 @@
 from typing_extensions import ParamSpec
 
 from opentelemetry.trace import Span, Tracer
-from opentelemetry import context as context_api
 import requests
 
 from humanloop.base_client import BaseHumanloop
-from humanloop.context import get_evaluation_context, get_trace_id, set_trace_id
+from humanloop.context import get_trace_id, set_trace_id
 from humanloop.types.chat_message import ChatMessage
 from humanloop.utilities.helpers import bind_args
 from humanloop.eval_utils.types import File

src/humanloop/utilities/prompt.py

@@ -4,7 +4,7 @@
 
 from typing import Callable, Optional
 
-from humanloop.context import PromptContext, set_prompt_context
+from humanloop.context import PromptContext, set_decorator_context
 
 logger = logging.getLogger("humanloop.sdk")
 
@@ -13,7 +13,7 @@ def prompt_decorator_factory(path: str, template: Optional[str]):
     def decorator(func: Callable):
         @wraps(func)
         def wrapper(*args, **kwargs):
-            with set_prompt_context(
+            with set_decorator_context(
                 PromptContext(
                     path=path,
                     template=template,

tests/utilities/test_prompt.py

@@ -13,7 +13,7 @@
 from groq import Groq
 from groq import NotFoundError as GroqNotFoundError
 from humanloop.client import Humanloop
-from humanloop.eval_utils.run import HumanloopUtilityError
+from humanloop.eval_utils.run import HumanloopDecoratorError
 from humanloop.utilities.prompt import prompt_decorator_factory
 from humanloop.otel.constants import HUMANLOOP_FILE_KEY
 from humanloop.otel.helpers import is_humanloop_span, read_from_opentelemetry_span
@@ -449,7 +449,7 @@ def call_llm_with_hl_call():
         )
         return response.logs[0].output_message.content  # type: ignore [union-attr]
 
-    with pytest.raises(HumanloopUtilityError):
+    with pytest.raises(HumanloopDecoratorError):
         call_llm_with_hl_call()
 
     response = humanloop_client.directories.get(id=test_directory.id)
@@ -481,7 +481,7 @@ def call_llm_with_hl_call():
 
         return response.logs[0].output_message.content
 
-    with pytest.raises(HumanloopUtilityError):
+    with pytest.raises(HumanloopDecoratorError):
         call_llm_with_hl_call()
 
     response = humanloop_client.directories.get(id=test_directory.id)
@@ -532,7 +532,7 @@ def call_llm_with_hl_call():
 
         return response.logs[0].output_message.content
 
-    with pytest.raises(HumanloopUtilityError):
+    with pytest.raises(HumanloopDecoratorError):
         call_llm_with_hl_call()
 
     response = humanloop_client.directories.get(id=test_directory.id)