simplify decorators

Author: Andrei Bratu

src/humanloop/client.py

@@ -1,7 +1,7 @@
+from contextlib import contextmanager
 import os
 import typing
-from typing import List, Optional, Sequence
-from typing_extensions import Unpack
+from typing import Any, List, Optional, Sequence
 
 import httpx
 from opentelemetry.sdk.resources import Resource
@@ -10,7 +10,6 @@
 
 from humanloop.core.client_wrapper import SyncClientWrapper
 from humanloop.eval_utils.run import prompt_call_evaluation_aware
-from humanloop.utilities.types import DecoratorPromptKernelRequestParams
 
 from humanloop.eval_utils import log_with_evaluation_context, run_eval
 from humanloop.eval_utils.types import Dataset, Evaluator, EvaluatorCheck, File
@@ -147,11 +146,11 @@ def __init__(
         else:
             self._opentelemetry_tracer = opentelemetry_tracer
 
+    @contextmanager
     def prompt(
         self,
         *,
-        path: Optional[str] = None,
-        **prompt_kernel: Unpack[DecoratorPromptKernelRequestParams],  # type: ignore
+        path: str,
     ):
         """Decorator for declaring a [Prompt](https://humanloop.com/docs/explanation/prompts) in code.
 
@@ -226,17 +225,14 @@ def call_llm(messages):
 
         :param prompt_kernel: Attributes that define the Prompt. See `class:DecoratorPromptKernelRequestParams`
         """
-        return prompt_decorator_factory(
-            opentelemetry_tracer=self._opentelemetry_tracer,
-            path=path,
-            **prompt_kernel,
-        )
+        return prompt_decorator_factory(path=path)
 
     def tool(
         self,
         *,
-        path: Optional[str] = None,
-        **tool_kernel: Unpack[ToolKernelRequestParams],  # type: ignore
+        path: str,
+        attributes: dict[str, Any] | None = None,
+        setup_values: dict[str, Any] | None = None,
     ):
         """Decorator for declaring a [Tool](https://humanloop.com/docs/explanation/tools) in code.
 
@@ -309,14 +305,15 @@ def calculator(a: int, b: Optional[int]) -> int:
         return tool_decorator_factory(
             opentelemetry_tracer=self._opentelemetry_tracer,
             path=path,
-            **tool_kernel,
+            attributes=attributes,
+            setup_values=setup_values,
         )
 
     def flow(
         self,
         *,
-        path: Optional[str] = None,
-        **flow_kernel: Unpack[FlowKernelRequestParams],  # type: ignore
+        path: str = None,
+        attributes: dict[str, Any] | None = None,
     ):
         """Decorator for declaring a [Flow](https://humanloop.com/docs/explanation/flows) in code.
 
@@ -371,7 +368,7 @@ def entrypoint():
         return flow_decorator_factory(
             opentelemetry_tracer=self._opentelemetry_tracer,
             path=path,
-            **flow_kernel,
+            attributes=attributes,
         )
 
 

src/humanloop/context.py

@@ -0,0 +1,61 @@
+from dataclasses import dataclass
+import threading
+from typing import Callable, Optional
+from opentelemetry import context as context_api
+
+from humanloop.otel.constants import (
+    HUMANLOOP_CONTEXT_EVALUATION,
+    HUMANLOOP_CONTEXT_PROMPT_PATH,
+    HUMANLOOP_CONTEXT_TRACE_ID,
+)
+
+
+ResetToken = object
+
+
+def get_trace_id() -> Optional[str]:
+    key = hash((HUMANLOOP_CONTEXT_TRACE_ID, threading.get_ident()))
+    context_api.get_value()
+    return context_api.get_value(key=key)
+
+
+def set_trace_id(flow_log_id: str) -> ResetToken:
+    key = hash((HUMANLOOP_CONTEXT_TRACE_ID, threading.get_ident()))
+    return context_api.attach(context_api.set_value(key=key, value=flow_log_id))
+
+
+def reset_trace_id_context(token: ResetToken):
+    context_api.detach(token=token)
+
+
+def set_prompt_path(path: str) -> ResetToken:
+    key = hash((HUMANLOOP_CONTEXT_PROMPT_PATH, threading.get_ident()))
+    return context_api.set_value(key=key, value=path)
+
+
+def reset_prompt_path(token: ResetToken):
+    context_api.detach(token=token)
+
+
+def get_prompt_path() -> Optional[str]:
+    key = hash((HUMANLOOP_CONTEXT_PROMPT_PATH, threading.get_ident()))
+    return context_api.get_value(key)
+
+
+@dataclass
+class EvaluationContext:
+    source_datapoint_id: str
+    run_id: str
+    callback: Callable[[str], None]
+    file_id: str
+    path: str
+
+
+def set_evaluation_context(evaluation_context: EvaluationContext):
+    key = hash((HUMANLOOP_CONTEXT_EVALUATION, threading.get_ident()))
+    context_api.set_value(key, evaluation_context)
+
+
+def get_evaluation_context() -> Optional[EvaluationContext]:
+    key = hash((HUMANLOOP_CONTEXT_EVALUATION, threading.get_ident()))
+    return context_api.get_value(key)

src/humanloop/context_variables.py

@@ -16,17 +16,16 @@
 import sys
 import threading
 import time
-import types
 import typing
 from concurrent.futures import ThreadPoolExecutor
 from datetime import datetime
 from functools import partial
 from logging import INFO
-from typing import Callable, Dict, List, Literal, Optional, Sequence, Tuple, TypeVar, Union
+from typing import Callable, Dict, List, Literal, Optional, Sequence, Tuple, Union
 
 from humanloop import EvaluatorResponse, FlowResponse, PromptResponse, ToolResponse
+from humanloop.context import EvaluationContext, set_evaluation_context
 from humanloop.core.api_error import ApiError
-from humanloop.context_variables import EvaluationContext, set_evaluation_context
 from humanloop.eval_utils.types import Dataset, Evaluator, EvaluatorCheck, File
 
 # We use TypedDicts for requests, which is consistent with the rest of the SDK

src/humanloop/eval_utils/run.py

@@ -2,9 +2,13 @@
 HUMANLOOP_FILE_KEY = "humanloop.file"
 # Attribute name prefix on Humanloop spans for log-related attributes
 HUMANLOOP_LOG_KEY = "humanloop.log"
+HUMANLOOP_LOG_ID_KEY = "humanloop.log_id"
 HUMANLOOP_FILE_TYPE_KEY = "humanloop.file.type"
 HUMANLOOP_PATH_KEY = "humanloop.file.path"
 # Required for the exporter to know when to mark the Flow Log as complete
 HUMANLOOP_FLOW_PREREQUISITES_KEY = "humanloop.flow.prerequisites"
-HUMANLOOP_INTERCEPTED_HL_CALL_SPAN_NAME = "humanloop_intercepted_hl_call"
-HUMANLOOP_INTERCEPTED_HL_CALL_RESPONSE = "intercepted_call_response"
+HUMANLOOP_INTERCEPTED_PROMPT_CALL_SPAN_NAME = "humanloop_intercepted_hl_call"
+HUMANLOOP_INTERCEPTED_PROMPT_CALL_RESPONSE = "intercepted_call_response"
+HUMANLOOP_CONTEXT_PROMPT_PATH = "humanloop.context.prompt.path"
+HUMANLOOP_CONTEXT_TRACE_ID = "humanloop.context.flow.trace_id"
+HUMANLOOP_CONTEXT_EVALUATION = "humanloop.context.evaluation"