@@ -13,6 +13,7 @@ import { PromptKernelRequest } from "../api/types/PromptKernelRequest";
1313import {
1414 HUMANLOOP_FILE_KEY ,
1515 HUMANLOOP_FILE_TYPE_KEY ,
16+ HUMANLOOP_FLOW_SPAN_NAME ,
1617 HUMANLOOP_LOG_KEY ,
1718 HUMANLOOP_META_FUNCTION_NAME ,
1819} from "./constants" ;
@@ -36,15 +37,37 @@ interface CompletableSpan {
3637export class HumanloopSpanProcessor implements SpanProcessor {
3738 private spanExporter : SpanExporter ;
3839 private children : Map < string , CompletableSpan [ ] > ;
40+ // List of all span IDs that are contained in a Flow trace
41+ // They are passed to the Exporter as a span attribute
42+ // so the Exporter knows when to complete a trace
43+ private prerequisites : Map < string , string [ ] > ;
3944
4045 constructor ( exporter : SpanExporter ) {
4146 this . spanExporter = exporter ;
4247 this . children = new Map ( ) ;
48+ this . prerequisites = new Map ( ) ;
4349 }
4450
4551 async forceFlush ( ) : Promise < void > { }
4652
4753 onStart ( span : Span , _ : Context ) : void {
54+ const spanId = span . spanContext ( ) . spanId ;
55+ const parentSpanId = span . parentSpanId ;
56+ if ( span . name === HUMANLOOP_FLOW_SPAN_NAME ) {
57+ this . prerequisites . set ( spanId , [ ] ) ;
58+ }
59+ if ( parentSpanId !== undefined && isHumanloopSpan ( span ) ) {
60+ for ( const [ traceHead , allTraceNodes ] of this . prerequisites ) {
61+ if (
62+ parentSpanId === traceHead ||
63+ allTraceNodes . includes ( parentSpanId )
64+ ) {
65+ allTraceNodes . push ( spanId ) ;
66+ this . prerequisites . set ( traceHead , allTraceNodes ) ;
67+ break ;
68+ }
69+ }
70+ }
4871 // Handle stream case: when Prompt instrumented function calls a provider with streaming: true
4972 // The instrumentor span will end only when the ChunksResponse is consumed, which can happen
5073 // after the span created by the Prompt utility finishes. To handle this, we register all instrumentor
@@ -66,6 +89,7 @@ export class HumanloopSpanProcessor implements SpanProcessor {
6689 */
6790 onEnd ( span : ReadableSpan ) : void {
6891 if ( isHumanloopSpan ( span ) ) {
92+ // Wait for children to complete asynchronously
6993 new Promise < void > ( ( resolve ) => {
7094 const checkChildrenSpans = ( ) => {
7195 const childrenSpans = this . children . get ( span . spanContext ( ) . spanId ) ;
@@ -79,15 +103,28 @@ export class HumanloopSpanProcessor implements SpanProcessor {
79103 } ;
80104 checkChildrenSpans ( ) ;
81105 } ) . then ( ( _ ) => {
82- // All children/ instrumentor spans have arrived, we can process the
106+ // All instrumentor spans have arrived, we can process the
83107 // Humanloop parent span owning them
108+ if ( span . name === HUMANLOOP_FLOW_SPAN_NAME ) {
109+ // If the span if a Flow Log, add attribute with all span IDs it
110+ // needs to wait before completion
111+ writeToOpenTelemetrySpan (
112+ span ,
113+ this . prerequisites . get ( span . spanContext ( ) . spanId ) || [ ] ,
114+ HUMANLOOP_LOG_KEY ,
115+ ) ;
116+ this . prerequisites . delete ( span . spanContext ( ) . spanId ) ;
117+ }
118+
84119 this . processSpanDispatch (
85120 span ,
86121 this . children . get ( span . spanContext ( ) . spanId ) || [ ] ,
87122 ) ;
123+
88124 // Release references
89125 this . children . delete ( span . spanContext ( ) . spanId ) ;
90- // Export the Humanloop span
126+
127+ // Pass Humanloop span to Exporter
91128 this . spanExporter . export ( [ span ] , ( result : ExportResult ) => {
92129 if ( result . code !== ExportResultCode . SUCCESS ) {
93130 console . error ( "Failed to export span:" , result . error ) ;
@@ -182,7 +219,7 @@ export class HumanloopSpanProcessor implements SpanProcessor {
182219 // Placeholder for processing other file types
183220 break ;
184221 default :
185- console . error ( "Unknown Humanloop File Span " , span ) ;
222+ console . error ( "Unknown Humanloop File span " , span ) ;
186223 }
187224 }
188225
0 commit comments