feature(stream): Adjusted base handler class for expected types, adjusted mock version of logger class.

Author: Murike

src/types/logging/span.types.ts

@@ -482,7 +482,12 @@ export class ToolSpan extends BaseStep {
 }
 
 // Type for all span types
-export type Span = WorkflowSpan | LlmSpan | RetrieverSpan | ToolSpan;
+export type Span =
+  | WorkflowSpan
+  | AgentSpan
+  | LlmSpan
+  | RetrieverSpan
+  | ToolSpan;
 
 /**
  * Type guard to validate if a value is a valid AgentType

src/types/streaming-adapter.types.ts

@@ -15,4 +15,3 @@ export interface StreamingMetricsResult {
   durationNs: number;
   tokenUsage: TokenUsage | null;
 }
-

src/utils/streaming/base-streaming-adapter.ts

@@ -2,7 +2,10 @@ import { GalileoLogger } from '../galileo-logger';
 import { ErrorManager } from '../errors';
 import { StreamingMetrics } from './metrics';
 import { StreamingFinalizer, StreamingFinalizerConfig } from './finalizer';
-import { TokenUsage, StreamingMetricsResult } from '../../types/streaming-adapter.types';
+import {
+  TokenUsage,
+  StreamingMetricsResult
+} from '../../types/streaming-adapter.types';
 import {
   LlmSpanAllowedInputType,
   LlmSpanAllowedOutputType

src/utils/streaming/finalizer.ts

@@ -1,26 +1,13 @@
-import { GalileoLogger } from '../galileo-logger';
-import { StreamingMetrics } from './metrics';
-import {
+import type { GalileoLogger } from '../galileo-logger';
+import type { StreamingMetrics } from './metrics';
+import type {
   LlmSpanAllowedInputType,
   LlmSpanAllowedOutputType
 } from '../../types/logging/step.types';
-import { Span } from '../../types/logging/span.types';
+import type { Span } from '../../types/logging/span.types';
 import { Trace } from '../../types/logging/trace.types';
 import type { ToolDefinition } from './base-streaming-adapter';
 
-/**
- * Extended logger interface with streaming update methods
- * These methods are mocked and will be replaced when streaming mode PR is merged
- */
-export interface GalileoLoggerWithStreaming extends GalileoLogger {
-  _updateSpanStreaming?(span: Span, output: LlmSpanAllowedOutputType): void;
-  _updateTraceStreaming?(
-    trace: Trace,
-    output: LlmSpanAllowedOutputType,
-    isComplete?: boolean
-  ): void;
-}
-
 /**
  * Configuration for streaming finalization
  */
@@ -114,21 +101,13 @@ export class StreamingFinalizer {
       return;
     }
 
-    // Type guard to check if logger has streaming methods
-    const loggerWithStreaming = this.logger as GalileoLoggerWithStreaming;
-
-    // TO-DO: After logger PR merged, substitute for actual implementation.
+    // Call streaming update methods (mock implementations in GalileoLogger)
     if (currentParent instanceof Trace) {
-      // Mock implementation - assumes _updateTraceStreaming exists in another branch/PR
-      loggerWithStreaming._updateTraceStreaming?.(
-        currentParent,
-        output,
-        false
-      );
+      // Update trace with partial output during streaming
+      this.logger._updateTraceStreaming(currentParent, output, false);
     } else {
-      // Mock implementation - assumes _updateSpanStreaming exists in another branch/PR
-      // currentParent is a Span (WorkflowSpan, AgentSpan, etc.)
-      loggerWithStreaming._updateSpanStreaming?.(currentParent as Span, output);
+      // currentParent is a Span (WorkflowSpan, AgentSpan, LlmSpan, etc.)
+      this.logger._updateSpanStreaming(currentParent as Span, output);
     }
   }
 }

src/utils/streaming/index.ts

@@ -12,7 +12,10 @@ export { StreamingMetrics } from './metrics';
 export { StreamingFinalizer } from './finalizer';
 export { BaseStreamingAdapter } from './base-streaming-adapter';
 export { ErrorManager, ErrorInfo } from '../errors';
-export type { TokenUsage, StreamingMetricsResult } from '../../types/streaming-adapter.types';
+export type {
+  TokenUsage,
+  StreamingMetricsResult
+} from '../../types/streaming-adapter.types';
 export type { StreamingFinalizerConfig } from './finalizer';
 export type {
   BaseStreamingAdapterConfig,

src/utils/streaming/metrics.ts

@@ -1,4 +1,7 @@
-import type { TokenUsage, StreamingMetricsResult } from '../../types/streaming-adapter.types';
+import type {
+  TokenUsage,
+  StreamingMetricsResult
+} from '../../types/streaming-adapter.types';
 
 /**
  * Tracks streaming metrics including TTFT, duration, and token usage
@@ -77,4 +80,3 @@ export class StreamingMetrics {
     this.tokenUsage = null;
   }
 }
-