baikai-trace-otel 0.2.0.0 → 0.3.0.0
raw patch · 3 files changed
+66/−16 lines, 3 filesdep ~baikaiPVP ok
version bump matches the API change (PVP)
Dependency ranges changed: baikai
API changes (from Hackage documentation)
Files
- baikai-trace-otel.cabal +4/−3
- src/Baikai/Trace/Sink/OpenTelemetry.hs +8/−2
- test/Main.hs +54/−11
baikai-trace-otel.cabal view
@@ -1,6 +1,6 @@ cabal-version: 3.4 name: baikai-trace-otel-version: 0.2.0.0+version: 0.3.0.0 synopsis: OpenTelemetry TraceSink for baikai. description: Provides an opt-in OpenTelemetry adapter for the baikai 'TraceSink'@@ -35,7 +35,7 @@ hs-source-dirs: src exposed-modules: Baikai.Trace.Sink.OpenTelemetry build-depends:- , baikai ^>=0.2.0+ , baikai ^>=0.3.0 , base >=4.20 && <5 , containers , hs-opentelemetry-api >=1.0 && <1.1@@ -53,7 +53,7 @@ main-is: Main.hs ghc-options: -threaded -with-rtsopts=-N build-depends:- , baikai ^>=0.2.0+ , baikai ^>=0.3.0 , baikai-trace-otel , base , generic-lens@@ -61,6 +61,7 @@ , hs-opentelemetry-exporter-in-memory >=1.0 && <1.1 , hs-opentelemetry-sdk >=1.0 && <1.1 , lens ^>=5.3+ , streamly-core ^>=0.3 , tasty , tasty-hunit , text
src/Baikai/Trace/Sink/OpenTelemetry.hs view
@@ -107,6 +107,9 @@ pure (Map.insert eventId sp m) CallFinished {eventId, timestamp, model, latencyMs, inputTokens, outputTokens, usd} -> case Map.lookup eventId m of+ -- A terminal without a live span is unreachable for normal withTraceStream+ -- usage because each traced call drives a fresh fold. Keep the silent drop so+ -- hand-fed or replayed event streams cannot crash the sink. Nothing -> pure m Just sp -> do let attrs :: HashMap.HashMap Text Attr.Attribute@@ -116,7 +119,7 @@ maybe id (\s -> HashMap.insert "baikai.cost.usd" (Attr.toAttribute (Scientific.toRealFloat s :: Double))) usd $ AttrMap.insertByKey SC.genAi_response_model model $ HashMap.fromList- [ ("baikai.latency_ms", Attr.toAttribute (fromIntegral latencyMs :: Int))+ [ ("baikai.latency_ms", Attr.toAttribute latencyMs) ] Otel.addAttributes sp attrs Otel.setStatus sp Otel.Ok@@ -124,11 +127,14 @@ pure (Map.delete eventId m) CallFailed {eventId, timestamp, latencyMs, errorMessage} -> case Map.lookup eventId m of+ -- A terminal without a live span is unreachable for normal withTraceStream+ -- usage because each traced call drives a fresh fold. Keep the silent drop so+ -- hand-fed or replayed event streams cannot crash the sink. Nothing -> pure m Just sp -> do Otel.addAttributes sp $ HashMap.fromList- [ ("baikai.latency_ms", Attr.toAttribute (fromIntegral latencyMs :: Int)),+ [ ("baikai.latency_ms", Attr.toAttribute latencyMs), ("baikai.error", Attr.toAttribute errorMessage) ] Otel.setStatus sp (Otel.Error errorMessage)
test/Main.hs view
@@ -4,28 +4,32 @@ import Baikai.Api (Api (..)) import Baikai.Content (AssistantContent (..), TextContent (..))-import Baikai.Context (Context (..), _Context)+import Baikai.Context (Context (..), emptyContext) import Baikai.Error (BaikaiError, providerError) import Baikai.Message (AssistantPayload (..), user)-import Baikai.Model (Model (..), _Model)-import Baikai.Options (Options, _Options)+import Baikai.Model (Model (..), emptyModel)+import Baikai.Options (Options, emptyOptions) import Baikai.Provider (ApiProvider (..), registerApiProvider) import Baikai.Response (Response (..)) import Baikai.StopReason (StopReason (..)) import Baikai.Stream (liftCompleteToStream)-import Baikai.Trace (withTrace)+import Baikai.Trace (withTrace, withTraceStream) import Baikai.Trace.Sink.OpenTelemetry (otelSink)-import Baikai.Usage (Usage, _Usage)+import Baikai.Usage (Usage, zeroUsage)+import Control.Concurrent (threadDelay) import Control.Exception (throwIO) import Control.Lens ((&), (.~), (^.)) import Data.Generics.Labels () import Data.HashMap.Strict qualified as HashMap import Data.IORef (IORef, readIORef) import Data.Text (Text)+import Data.Text qualified as Text import Data.Vector qualified as V import OpenTelemetry.Attributes qualified as Attr import OpenTelemetry.Exporter.InMemory.Span (inMemoryListExporter) import OpenTelemetry.Trace.Core qualified as Otel+import Streamly.Data.Stream qualified as Stream+import System.Mem (performMajorGC) import Test.Tasty (TestTree, defaultMain, testGroup) import Test.Tasty.HUnit (assertBool, assertEqual, assertFailure, testCase, (@?=)) @@ -35,7 +39,8 @@ testGroup "baikai-trace-otel" [ successSpanTest,- failureSpanTest+ failureSpanTest,+ abortSpanTest ] -- | Build a stub 'Model' under a private 'Api' tag. Each test uses@@ -43,21 +48,21 @@ -- registry between tests. stubModel :: Api -> Model stubModel a =- _Model+ emptyModel & #modelId .~ "stub-1" & #api .~ a & #provider .~ "stub.otel" & #maxOutputTokens .~ 16 stubContext :: Context-stubContext = _Context & #messages .~ V.fromList [user "hello"]+stubContext = emptyContext & #messages .~ V.fromList [user "hello"] stubOptions :: Options-stubOptions = _Options & #maxTokens .~ Just 16+stubOptions = emptyOptions & #maxTokens .~ Just 16 sampleUsage :: Usage sampleUsage =- _Usage+ zeroUsage & #inputTokens .~ 12 & #outputTokens .~ 3 & #totalTokens .~ 15@@ -71,7 +76,7 @@ usage = sampleUsage, stopReason = Stop, errorMessage = Nothing,- timestamp = read "2026-05-14 00:00:00 UTC"+ timestamp = Just (read "2026-05-14 00:00:00 UTC") }, model = stubModel a, api = a,@@ -177,3 +182,41 @@ assertBool "has baikai.error" (HashMap.member "baikai.error" attrs) assertBool "has baikai.latency_ms" (HashMap.member "baikai.latency_ms" attrs) _ -> assertFailure "expected exactly one span"++abortSpanTest :: TestTree+abortSpanTest =+ testCase "early abort closes the span with Error status" $ do+ let a = Custom "baikai-otel-abort"+ registerOk a+ (tracer, getSpans) <- newTracerWithInMemory+ let sink = otelSink tracer+ emitted <-+ Stream.toList+ (Stream.take 1 (withTraceStream sink (stubModel a) stubContext stubOptions))+ length emitted @?= 1+ spans <- awaitSpans getSpans 1+ assertEqual "exactly one span recorded" 1 (length spans)+ case spans of+ [sp] -> do+ hot <- spanHotSnapshot sp+ case Otel.hotStatus hot of+ Otel.Error msg ->+ assertBool+ ("expected abort message, got: " <> show msg)+ ("aborted" `Text.isInfixOf` msg)+ other -> assertFailure ("expected Error status, got: " <> show other)+ _ -> assertFailure "expected exactly one span"++-- The trace finalizer on an abandoned stream runs from streamly's GC hook.+awaitSpans :: IO [Otel.ImmutableSpan] -> Int -> IO [Otel.ImmutableSpan]+awaitSpans getSpans n = go (100 :: Int)+ where+ go 0 = do+ spans <- getSpans+ assertFailure ("timed out waiting for spans; got: " <> show (length spans))+ go k = do+ performMajorGC+ spans <- getSpans+ if length spans >= n+ then pure spans+ else threadDelay 50000 >> go (k - 1)