Harry PR feedback

Author: Andrei Bratu

src/humanloop/decorators/flow.py

@@ -1,7 +1,7 @@
 import logging
-from functools import wraps
 import os
 import sys
+from functools import wraps
 from typing import Any, Callable, Mapping, Optional, Sequence
 
 from opentelemetry.sdk.trace import Span

src/humanloop/decorators/prompt.py

@@ -6,7 +6,6 @@
 from opentelemetry.sdk.trace import Span
 from opentelemetry.trace import Tracer
 
-
 if typing.TYPE_CHECKING:
     from humanloop import ToolFunctionParams
 from humanloop.decorators.helpers import args_to_inputs

src/humanloop/eval_utils/context.py

@@ -1,4 +1,3 @@
-from contextvars import ContextVar, Token
 from typing import Callable, TypedDict
 
 

src/humanloop/eval_utils/run.py

@@ -8,7 +8,6 @@
 not be called directly.
 """
 
-from contextvars import ContextVar
 import copy
 import inspect
 import json
@@ -19,6 +18,7 @@
 import types
 import typing
 from concurrent.futures import ThreadPoolExecutor
+from contextvars import ContextVar
 from datetime import datetime
 from functools import partial
 from logging import INFO
@@ -144,6 +144,12 @@ def _overloaded_log(
                     kwargs[attribute] = evaluation_context[attribute]  # type: ignore
 
         # Call the original .log method
+        logger.debug(
+            "Logging %s inside _overloaded_log on Thread %s",
+            kwargs,
+            evaluation_context,
+            threading.get_ident(),
+        )
         response = self._log(**kwargs)
 
         # Call the callback so the Evaluation can be updated
@@ -174,6 +180,7 @@ def _overloaded_log(
     # Replace the original log method with the overloaded one
     client.log = types.MethodType(_overloaded_log, client)  # type: ignore
     # Return the client with the overloaded log method
+    logger.debug("Overloaded the .log method of %s", client)
     return client
 
 
@@ -246,7 +253,7 @@ def run_eval(
     global _PROGRESS_BAR
 
     if hasattr(file["callable"], "file"):
-        # When the decorator inside `file`` is a decorated function,
+        # When the decorator inside `file` is a decorated function,
         # we need to validate that the other parameters of `file`
         # match the attributes of the decorator
         inner_file: File = file["callable"].file
@@ -429,6 +436,12 @@ def run_eval(
     def process_datapoint(dp: Datapoint, file_id: str, file_path: str, run_id: str):
         def upload_callback(log: dict):
             """Logic ran after the Log has been created."""
+            logger.debug(
+                "upload_callback on Thread %s: log %s datapoint_target %s",
+                threading.get_ident(),
+                log,
+                dp.target,
+            )
             _run_local_evaluators(
                 client=client,
                 log=log,
@@ -448,6 +461,12 @@ def upload_callback(log: dict):
                 path=file_path,
             )
         )
+        logger.debug(
+            "process_datapoint on Thread %s: evaluating Datapoint %s with EvaluationContext %s",
+            threading.get_ident(),
+            datapoint_dict,
+            evaluation_context_variable.get(),
+        )
         log_func = _get_log_func(
             client=client,
             file_type=type_,
@@ -482,6 +501,11 @@ def upload_callback(log: dict):
             if context_variable is not None:
                 # Evaluation Context has not been consumed
                 # function_ is a plain callable so we need to create a Log
+                logger.debug(
+                    "process_datapoint on Thread %s: function_ %s is a simple callable, context was not consumed",
+                    threading.get_ident(),
+                    function_.__name__,
+                )
                 log_func(
                     inputs=datapoint.inputs,
                     output=output,

src/humanloop/otel/exporter.py

@@ -1,6 +1,7 @@
-import copy
+import contextvars
 import json
 import logging
+import threading
 import typing
 from queue import Empty as EmptyQueue
 from queue import Queue
@@ -62,6 +63,7 @@ def __init__(
         self._shutdown: bool = False
         for thread in self._threads:
             thread.start()
+            logger.debug("Exporter Thread %s started", thread.ident)
 
     def export(self, spans: trace.Sequence[ReadableSpan]) -> SpanExportResult:
         def is_evaluated_file(
@@ -89,15 +91,25 @@ def is_evaluated_file(
                     self._upload_queue.put(
                         (
                             span,
-                            copy.deepcopy(evaluation_context),
-                        )
+                            contextvars.copy_context()[self._client.evaluation_context_variable],
+                        ),
+                    )
+                    logger.debug(
+                        "Span %s with EvaluationContext %s added to upload queue",
+                        span.attributes,
+                        contextvars.copy_context()[self._client.evaluation_context_variable],
                     )
             # Reset the EvaluationContext so run eval does not
             # create a duplicate Log
             if evaluation_context is not None and is_evaluated_file(
                 spans[0],
                 evaluation_context,
             ):
+                logger.debug(
+                    "EvaluationContext %s marked as exhausted for Log in Span %s",
+                    evaluation_context,
+                    spans[0].attributes,
+                )
                 # Mark the EvaluationContext as used
                 self._client.evaluation_context_variable.set(None)
             return SpanExportResult.SUCCESS
@@ -109,6 +121,7 @@ def shutdown(self) -> None:
         self._shutdown = True
         for thread in self._threads:
             thread.join()
+            logger.debug("Exporter Thread %s joined", thread.ident)
 
     def force_flush(self, timeout_millis: int = 3000) -> bool:
         self._shutdown = True
@@ -146,21 +159,37 @@ def _do_work(self):
                 # Set the EvaluationContext for the thread so the .log action works as expected
                 # NOTE: Expecting the evaluation thread to send a single span so we are
                 #   not resetting the EvaluationContext in the scope of the export thread
-                self._client.evaluation_context_variable.set(
-                    copy.deepcopy(evaluation_context),
-                )
+                self._client.evaluation_context_variable.set(evaluation_context)
             except EmptyQueue:
                 continue
             trace_metadata = TRACE_FLOW_CONTEXT.get(span_to_export.get_span_context().span_id)
             if trace_metadata is None:
                 # Span is not part of a Flow Log
                 self._export_span_dispatch(span_to_export)
+                logger.debug(
+                    "_do_work on Thread %s: Dispatched span %s with FlowContext %s which is not part of a Flow",
+                    threading.get_ident(),
+                    span_to_export.attributes,
+                    trace_metadata,
+                )
             elif trace_metadata["trace_parent_id"] is None:
                 # Span is the head of a Flow Trace
                 self._export_span_dispatch(span_to_export)
+                logger.debug(
+                    "Dispatched span %s which is a Flow Log with FlowContext %s",
+                    span_to_export.attributes,
+                    trace_metadata,
+                )
             elif trace_metadata["trace_parent_id"] in self._span_id_to_uploaded_log_id:
                 # Span is part of a Flow and its parent has been uploaded
                 self._export_span_dispatch(span_to_export)
+                logger.debug(
+                    "_do_work on Thread %s: Dispatched span %s after its parent %s with FlowContext %s",
+                    threading.get_ident(),
+                    span_to_export.attributes,
+                    trace_metadata["trace_parent_id"],
+                    trace_metadata,
+                )
             else:
                 # Requeue the Span to be uploaded later
                 self._upload_queue.put((span_to_export, evaluation_context))

src/humanloop/otel/helpers.py

@@ -196,7 +196,7 @@ def read_from_opentelemetry_span(span: ReadableSpan, key: str = "") -> NestedDic
                 sub_result[part] = span_value
             else:
                 if part not in sub_result:
-                    # New dict since
+                    # Create new dict for a previously unseen part of the key
                     sub_result[part] = {}
                 sub_result = sub_result[part]  # type: ignore