type checks

Author: Andrei Bratu

.gitignore

@@ -5,3 +5,5 @@ poetry.toml
 .ruff_cache/
 .vscode
 .env
+tests/assets/*.jsonl
+tests/assets/*.parquet

src/humanloop/client.py

@@ -38,10 +38,8 @@ def __init__(
         self,
         *,
         client_wrapper: SyncClientWrapper,
-        evaluation_context_variable: ContextVar[Optional[EvaluationContext]],
     ):
         super().__init__(client_wrapper=client_wrapper)
-        self._evaluation_context_variable = evaluation_context_variable
 
     def run(
         self,
@@ -70,7 +68,6 @@ def run(
             dataset=dataset,
             evaluators=evaluators,
             workers=workers,
-            evaluation_context_variable=self._evaluation_context_variable,
         )
 
 
@@ -118,31 +115,14 @@ def __init__(
             httpx_client=httpx_client,
         )
 
-        self.evaluation_context_variable: ContextVar[Optional[EvaluationContext]] = ContextVar(
-            EVALUATION_CONTEXT_VARIABLE_NAME
-        )
-
-        eval_client = ExtendedEvalsClient(
-            client_wrapper=self._client_wrapper,
-            evaluation_context_variable=self.evaluation_context_variable,
-        )
+        eval_client = ExtendedEvalsClient(client_wrapper=self._client_wrapper)
         eval_client.client = self
         self.evaluations = eval_client
         self.prompts = ExtendedPromptsClient(client_wrapper=self._client_wrapper)
 
         # Overload the .log method of the clients to be aware of Evaluation Context
-        # TODO: Overload the log for Evaluators and Tools once run_id is added
-        # to them.
-        self.prompts = log_with_evaluation_context(
-            client=self.prompts,
-            evaluation_context_variable=self.evaluation_context_variable,
-        )
-        # self.evaluators = log_with_evaluation_context(client=self.evaluators)
-        # self.tools = log_with_evaluation_context(client=self.tools)
-        self.flows = log_with_evaluation_context(
-            client=self.flows,
-            evaluation_context_variable=self.evaluation_context_variable,
-        )
+        self.prompts = log_with_evaluation_context(client=self.prompts)
+        self.flows = log_with_evaluation_context(client=self.flows)
 
         if opentelemetry_tracer_provider is not None:
             self._tracer_provider = opentelemetry_tracer_provider
@@ -157,9 +137,7 @@ def __init__(
         instrument_provider(provider=self._tracer_provider)
         self._tracer_provider.add_span_processor(
             HumanloopSpanProcessor(
-                exporter=HumanloopSpanExporter(
-                    client=self,
-                )
+                exporter=HumanloopSpanExporter(client=self),
             ),
         )
 

src/humanloop/eval_utils/context.py

@@ -1,7 +1,10 @@
-from typing import Callable, TypedDict
+from contextvars import ContextVar
+from dataclasses import dataclass
+from typing import Any, Callable
 
 
-class EvaluationContext(TypedDict):
+@dataclass
+class EvaluationContext:
     """Context Log to Humanloop.
 
     Per datapoint state that is set when an Evaluation is ran.
@@ -24,3 +27,43 @@ class EvaluationContext(TypedDict):
 
 
 EVALUATION_CONTEXT_VARIABLE_NAME = "__EVALUATION_CONTEXT"
+
+_EVALUATION_CONTEXT_VAR: ContextVar[EvaluationContext] = ContextVar(EVALUATION_CONTEXT_VARIABLE_NAME)
+
+_UnsafeEvaluationContextRead = RuntimeError("EvaluationContext not set in the current thread.")
+
+
+def set_evaluation_context(context: EvaluationContext):
+    _EVALUATION_CONTEXT_VAR.set(context)
+
+
+def get_evaluation_context() -> EvaluationContext:
+    try:
+        return _EVALUATION_CONTEXT_VAR.get()
+    except LookupError:
+        raise _UnsafeEvaluationContextRead
+
+
+def evaluation_context_set() -> bool:
+    try:
+        _EVALUATION_CONTEXT_VAR.get()
+        return True
+    except LookupError:
+        return False
+
+
+def log_belongs_to_evaluated_file(log_args: dict[str, Any]) -> bool:
+    try:
+        evaluation_context: EvaluationContext = _EVALUATION_CONTEXT_VAR.get()
+        return evaluation_context.file_id == log_args.get("id") or evaluation_context.path == log_args.get("path")
+    except LookupError:
+        # Not in an evaluation context
+        return False
+
+
+def is_evaluated_file(file_path) -> bool:
+    try:
+        evaluation_context = _EVALUATION_CONTEXT_VAR.get()
+        return evaluation_context.path == file_path
+    except LookupError:
+        raise _UnsafeEvaluationContextRead