1- import * as contextApi from "@opentelemetry/api" ;
1+ import * as otel_api from "@opentelemetry/api" ;
2+ import { AsyncHooksContextManager } from "@opentelemetry/context-async-hooks" ;
23
34import { HumanloopRuntimeError } from "./error" ;
45import {
@@ -7,15 +8,18 @@ import {
78 HUMANLOOP_CONTEXT_TRACE_ID ,
89} from "./otel/constants" ;
910
11+ export const HL_CONTEXT = new AsyncHooksContextManager ( ) ;
12+ HL_CONTEXT . enable ( ) ;
13+
1014export function getTraceId ( ) : string | undefined {
11- const key = contextApi . createContextKey ( HUMANLOOP_CONTEXT_TRACE_ID ) ;
12- const value = contextApi . context . active ( ) . getValue ( key ) ;
15+ const key = otel_api . createContextKey ( HUMANLOOP_CONTEXT_TRACE_ID ) ;
16+ const value = HL_CONTEXT . active ( ) . getValue ( key ) ;
1317 return ( value || undefined ) as string | undefined ;
1418}
1519
16- export function setTraceId ( flowLogId : string ) : contextApi . Context {
17- const key = contextApi . createContextKey ( HUMANLOOP_CONTEXT_TRACE_ID ) ;
18- return contextApi . context . active ( ) . setValue ( key , flowLogId ) ;
20+ export function setTraceId ( flowLogId : string ) : otel_api . Context {
21+ const key = otel_api . createContextKey ( HUMANLOOP_CONTEXT_TRACE_ID ) ;
22+ return HL_CONTEXT . active ( ) . setValue ( key , flowLogId ) ;
1923}
2024
2125export type DecoratorContext = {
@@ -26,14 +30,14 @@ export type DecoratorContext = {
2630
2731export function setDecoratorContext (
2832 decoratorContext : DecoratorContext ,
29- ) : contextApi . Context {
30- const key = contextApi . createContextKey ( HUMANLOOP_CONTEXT_DECORATOR ) ;
31- return contextApi . context . active ( ) . setValue ( key , decoratorContext ) ;
33+ ) : otel_api . Context {
34+ const key = otel_api . createContextKey ( HUMANLOOP_CONTEXT_DECORATOR ) ;
35+ return HL_CONTEXT . active ( ) . setValue ( key , decoratorContext ) ;
3236}
3337
3438export function getDecoratorContext ( ) : DecoratorContext | undefined {
35- const key = contextApi . createContextKey ( HUMANLOOP_CONTEXT_DECORATOR ) ;
36- return ( contextApi . context . active ( ) . getValue ( key ) || undefined ) as
39+ const key = otel_api . createContextKey ( HUMANLOOP_CONTEXT_DECORATOR ) ;
40+ return ( HL_CONTEXT . active ( ) . getValue ( key ) || undefined ) as
3741 | DecoratorContext
3842 | undefined ;
3943}
@@ -134,14 +138,14 @@ export class EvaluationContext {
134138
135139export function setEvaluationContext (
136140 evaluationContext : EvaluationContext ,
137- ) : contextApi . Context {
138- const key = contextApi . createContextKey ( HUMANLOOP_CONTEXT_EVALUATION ) ;
139- return contextApi . context . active ( ) . setValue ( key , evaluationContext ) ;
141+ ) : otel_api . Context {
142+ const key = otel_api . createContextKey ( HUMANLOOP_CONTEXT_EVALUATION ) ;
143+ return HL_CONTEXT . active ( ) . setValue ( key , evaluationContext ) ;
140144}
141145
142146export function getEvaluationContext ( ) : EvaluationContext | undefined {
143- const key = contextApi . createContextKey ( HUMANLOOP_CONTEXT_EVALUATION ) ;
144- return ( contextApi . context . active ( ) . getValue ( key ) || undefined ) as
147+ const key = otel_api . createContextKey ( HUMANLOOP_CONTEXT_EVALUATION ) ;
148+ return ( HL_CONTEXT . active ( ) . getValue ( key ) || undefined ) as
145149 | EvaluationContext
146150 | undefined ;
147151}
0 commit comments