baikai-trace-otel (empty) → 0.1.0.0
raw patch · 4 files changed
+429/−0 lines, 4 filesdep +baikaidep +baikai-trace-oteldep +base
Dependencies added: baikai, baikai-trace-otel, base, containers, generic-lens, hs-opentelemetry-api, hs-opentelemetry-exporter-in-memory, hs-opentelemetry-sdk, hs-opentelemetry-semantic-conventions, lens, scientific, streamly-core, tasty, tasty-hunit, text, time, unordered-containers, vector
Files
- LICENSE +30/−0
- baikai-trace-otel.cabal +69/−0
- src/Baikai/Trace/Sink/OpenTelemetry.hs +152/−0
- test/Main.hs +178/−0
+ LICENSE view
@@ -0,0 +1,30 @@+Copyright (c) 2026 Nadeem Bitar++All rights reserved.++Redistribution and use in source and binary forms, with or without+modification, are permitted provided that the following conditions are met:++ * Redistributions of source code must retain the above copyright+ notice, this list of conditions and the following disclaimer.++ * Redistributions in binary form must reproduce the above+ copyright notice, this list of conditions and the following+ disclaimer in the documentation and/or other materials provided+ with the distribution.++ * Neither the name of Nadeem Bitar nor the names of other+ contributors may be used to endorse or promote products derived+ from this software without specific prior written permission.++THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS+"AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT+LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR+A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT+OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,+SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT+LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,+DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY+THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT+(INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE+OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ baikai-trace-otel.cabal view
@@ -0,0 +1,69 @@+cabal-version: 3.4+name: baikai-trace-otel+version: 0.1.0.0+synopsis: OpenTelemetry TraceSink for baikai.+description:+ Provides an opt-in OpenTelemetry adapter for the baikai 'TraceSink'+ interface. Wiring 'otelSink' into 'Baikai.Trace.withTrace' produces+ one OTel span per provider call, with GenAI semantic-convention+ attributes plus baikai-specific cost and latency.++category: AI, Monitoring+license: BSD-3-Clause+license-file: LICENSE+author: Nadeem Bitar+maintainer: nadeem@gmail.com+copyright: (c) 2026 Nadeem Bitar+build-type: Simple++common common-options+ ghc-options:+ -Wall -Wcompat -Widentities -Wincomplete-uni-patterns+ -Wincomplete-record-updates -Wredundant-constraints+ -fhide-source-paths -Wmissing-export-lists -Wpartial-fields+ -Wmissing-deriving-strategies++ default-language: GHC2024+ default-extensions:+ DeriveAnyClass+ DuplicateRecordFields+ OverloadedLabels+ OverloadedStrings++library+ import: common-options+ hs-source-dirs: src+ exposed-modules: Baikai.Trace.Sink.OpenTelemetry+ build-depends:+ , baikai ^>=0.1.0+ , base >=4.20 && <5+ , containers+ , hs-opentelemetry-api >=1.0 && <1.1+ , hs-opentelemetry-semantic-conventions >=1.40 && <2+ , scientific+ , streamly-core ^>=0.3+ , text ^>=2.1+ , time+ , unordered-containers++test-suite baikai-trace-otel-test+ import: common-options+ type: exitcode-stdio-1.0+ hs-source-dirs: test+ main-is: Main.hs+ ghc-options: -threaded -with-rtsopts=-N+ build-depends:+ , baikai ^>=0.1.0+ , baikai-trace-otel+ , base+ , generic-lens+ , hs-opentelemetry-api >=1.0 && <1.1+ , hs-opentelemetry-exporter-in-memory >=1.0 && <1.1+ , hs-opentelemetry-sdk >=1.0 && <1.1+ , lens ^>=5.3+ , tasty+ , tasty-hunit+ , text+ , time+ , unordered-containers+ , vector
+ src/Baikai/Trace/Sink/OpenTelemetry.hs view
@@ -0,0 +1,152 @@+{-# LANGUAGE NamedFieldPuns #-}++-- | OpenTelemetry adapter for the baikai 'TraceSink' interface.+--+-- 'otelSink' (or 'otelSinkWith') wraps an 'Otel.Tracer' as a baikai+-- 'TraceSink' so that wiring it through 'Baikai.Trace.withTrace' emits+-- one OpenTelemetry span per provider call. Span attributes follow the+-- OpenTelemetry GenAI semantic conventions where possible+-- (@gen_ai.provider.name@, @gen_ai.request.model@, @gen_ai.usage.input_tokens@);+-- baikai-specific data uses the @baikai.@ prefix (@baikai.event_id@,+-- @baikai.latency_ms@, @baikai.cost.usd@, @baikai.error@).+--+-- The sink is a stateful streamly 'Fold' whose state is a @Map Text Span@+-- keyed by 'eventId'. A 'CallStarted' opens a span and inserts it; a+-- matching 'CallFinished' or 'CallFailed' closes the span and removes the+-- entry. The fold's finalizer closes any spans still in flight at+-- end-of-stream so no span ever leaks.+module Baikai.Trace.Sink.OpenTelemetry+ ( otelSink,+ otelSinkWith,+ OtelSinkOptions (..),+ defaultOtelSinkOptions,+ )+where++import Baikai.Trace.Event (TraceEvent (..))+import Baikai.Trace.Sink (TraceSink (..))+import Control.Monad (forM_)+import Data.HashMap.Strict qualified as HashMap+import Data.Int (Int64)+import Data.Map.Strict qualified as Map+import Data.Scientific qualified as Scientific+import Data.Text (Text)+import Data.Time (UTCTime)+import Data.Time.Clock.POSIX (utcTimeToPOSIXSeconds)+import Data.Word (Word64)+import OpenTelemetry.Attributes qualified as Attr+import OpenTelemetry.Attributes.Map qualified as AttrMap+import OpenTelemetry.Common (Timestamp, mkTimestamp)+import OpenTelemetry.Context qualified as Context+import OpenTelemetry.SemanticConventions qualified as SC+import OpenTelemetry.Trace.Core qualified as Otel+import Streamly.Data.Fold qualified as Fold++-- | Tunable knobs for 'otelSinkWith'.+data OtelSinkOptions = OtelSinkOptions+ { -- | Name to give each emitted span. Default: @"baikai.call"@.+ spanName :: !Text,+ -- | If 'True', attach the redacted prompt summary as+ -- @gen_ai.prompt_summary@. 'False' by default to avoid logging user+ -- content into observability backends.+ includePromptSummary :: !Bool+ }++-- | Defaults: span name @baikai.call@, prompt summary off.+defaultOtelSinkOptions :: OtelSinkOptions+defaultOtelSinkOptions =+ OtelSinkOptions+ { spanName = "baikai.call",+ includePromptSummary = False+ }++-- | Adapt a 'Otel.Tracer' to a baikai 'TraceSink' with the default options.+otelSink :: Otel.Tracer -> TraceSink+otelSink tracer = otelSinkWith tracer defaultOtelSinkOptions++-- | Adapt a 'Otel.Tracer' to a baikai 'TraceSink' with custom options.+otelSinkWith :: Otel.Tracer -> OtelSinkOptions -> TraceSink+otelSinkWith tracer opts =+ TraceSink (Fold.rmapM finalizer (Fold.foldlM' step (pure Map.empty)))+ where+ step :: Map.Map Text Otel.Span -> TraceEvent -> IO (Map.Map Text Otel.Span)+ step = stepEvent tracer opts++ finalizer :: Map.Map Text Otel.Span -> IO ()+ finalizer remaining =+ forM_ (Map.elems remaining) (\sp -> Otel.endSpan sp Nothing)++stepEvent ::+ Otel.Tracer ->+ OtelSinkOptions ->+ Map.Map Text Otel.Span ->+ TraceEvent ->+ IO (Map.Map Text Otel.Span)+stepEvent tracer OtelSinkOptions {spanName, includePromptSummary} m ev = case ev of+ CallStarted {eventId, timestamp, provider, model, maxTokens, promptSummary} -> do+ let baseAttrs :: HashMap.HashMap Text Attr.Attribute+ baseAttrs =+ AttrMap.insertByKey SC.genAi_provider_name provider $+ AttrMap.insertByKey SC.genAi_operation_name ("chat" :: Text) $+ AttrMap.insertByKey SC.genAi_request_model model $+ AttrMap.insertByKey SC.genAi_request_maxTokens (fromIntegral maxTokens :: Int64) $+ HashMap.fromList+ [ ("baikai.event_id", Attr.toAttribute eventId)+ ]+ attrs =+ if includePromptSummary+ then HashMap.insert "gen_ai.prompt_summary" (Attr.toAttribute promptSummary) baseAttrs+ else baseAttrs+ sargs =+ Otel.defaultSpanArguments+ { Otel.kind = Otel.Client,+ Otel.attributes = attrs,+ Otel.startTime = Just (utcToTimestamp timestamp)+ }+ sp <- Otel.createSpan tracer Context.empty spanName sargs+ pure (Map.insert eventId sp m)+ CallFinished {eventId, timestamp, model, latencyMs, inputTokens, outputTokens, usd} ->+ case Map.lookup eventId m of+ Nothing -> pure m+ Just sp -> do+ let attrs :: HashMap.HashMap Text Attr.Attribute+ attrs =+ maybe id (\n -> AttrMap.insertByKey SC.genAi_usage_inputTokens (fromIntegral n :: Int64)) inputTokens $+ maybe id (\n -> AttrMap.insertByKey SC.genAi_usage_outputTokens (fromIntegral n :: Int64)) outputTokens $+ 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))+ ]+ Otel.addAttributes sp attrs+ Otel.setStatus sp Otel.Ok+ Otel.endSpan sp (Just (utcToTimestamp timestamp))+ pure (Map.delete eventId m)+ CallFailed {eventId, timestamp, latencyMs, errorMessage} ->+ case Map.lookup eventId m of+ Nothing -> pure m+ Just sp -> do+ Otel.addAttributes sp $+ HashMap.fromList+ [ ("baikai.latency_ms", Attr.toAttribute (fromIntegral latencyMs :: Int)),+ ("baikai.error", Attr.toAttribute errorMessage)+ ]+ Otel.setStatus sp (Otel.Error errorMessage)+ Otel.endSpan sp (Just (utcToTimestamp timestamp))+ pure (Map.delete eventId m)++-- | Convert a 'UTCTime' to an OpenTelemetry 'Timestamp'.+--+-- 'Timestamp' stores nanoseconds since the Unix epoch. 'utcTimeToPOSIXSeconds'+-- produces a 'NominalDiffTime'; we go via the 'Real' instance to split into+-- whole seconds and remaining nanoseconds without dropping below microsecond+-- precision.+utcToTimestamp :: UTCTime -> Timestamp+utcToTimestamp t =+ let posix :: Rational+ posix = toRational (utcTimeToPOSIXSeconds t)+ secs :: Word64+ secs = floor posix+ nanos :: Word64+ nanos = floor ((posix - toRational secs) * 1_000_000_000)+ in mkTimestamp secs nanos
+ test/Main.hs view
@@ -0,0 +1,178 @@+{-# LANGUAGE ScopedTypeVariables #-}++module Main (main) where++import Baikai.Api (Api (..))+import Baikai.Content (AssistantContent (..), TextContent (..))+import Baikai.Context (Context (..), _Context)+import Baikai.Error (BaikaiError (..))+import Baikai.Message (AssistantPayload (..), user)+import Baikai.Model (Model (..), _Model)+import Baikai.Options (Options, _Options)+import Baikai.Provider (ApiProvider (..), registerApiProvider)+import Baikai.Response (Response (..))+import Baikai.StopReason (StopReason (..))+import Baikai.Stream (liftCompleteToStream)+import Baikai.Trace (withTrace)+import Baikai.Trace.Sink.OpenTelemetry (otelSink)+import Baikai.Usage (Usage, _Usage)+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.Vector qualified as V+import OpenTelemetry.Attributes qualified as Attr+import OpenTelemetry.Exporter.InMemory.Span (inMemoryListExporter)+import OpenTelemetry.Trace.Core qualified as Otel+import Test.Tasty (TestTree, defaultMain, testGroup)+import Test.Tasty.HUnit (assertBool, assertEqual, assertFailure, testCase, (@?=))++main :: IO ()+main =+ defaultMain $+ testGroup+ "baikai-trace-otel"+ [ successSpanTest,+ failureSpanTest+ ]++-- | Build a stub 'Model' under a private 'Api' tag. Each test uses+-- a distinct tag so tasty's parallel test scheduler cannot race the+-- registry between tests.+stubModel :: Api -> Model+stubModel a =+ _Model+ & #modelId .~ "stub-1"+ & #api .~ a+ & #provider .~ "stub.otel"+ & #maxOutputTokens .~ 16++stubContext :: Context+stubContext = _Context & #messages .~ V.fromList [user "hello"]++stubOptions :: Options+stubOptions = _Options & #maxTokens .~ Just 16++sampleUsage :: Usage+sampleUsage =+ _Usage+ & #inputTokens .~ 12+ & #outputTokens .~ 3+ & #totalTokens .~ 15++stubResponse :: Api -> Response+stubResponse a =+ Response+ { message =+ AssistantPayload+ { content = V.singleton (AssistantText (TextContent "hi")),+ usage = sampleUsage,+ stopReason = Stop,+ errorMessage = Nothing,+ timestamp = read "2026-05-14 00:00:00 UTC"+ },+ model = stubModel a,+ api = a,+ provider = "stub.otel",+ responseId = Nothing,+ latencyMs = 0+ }++registerOk :: Api -> IO ()+registerOk a =+ let handler _m _ctx _opts = pure (stubResponse a)+ in registerApiProvider+ ApiProvider+ { apiTag = a,+ stream = liftCompleteToStream handler,+ complete = handler+ }++registerFail :: Api -> BaikaiError -> IO ()+registerFail a e =+ let handler _m _ctx _opts = throwIO e+ in registerApiProvider+ ApiProvider+ { apiTag = a,+ stream = liftCompleteToStream handler,+ complete = handler+ }++newTracerWithInMemory :: IO (Otel.Tracer, IO [Otel.ImmutableSpan])+newTracerWithInMemory = do+ (proc, spansRef :: IORef [Otel.ImmutableSpan]) <- inMemoryListExporter+ tp <- Otel.createTracerProvider [proc] Otel.emptyTracerProviderOptions+ let tracer = Otel.makeTracer tp "baikai-trace-otel-test" Otel.tracerOptions+ pure (tracer, reverse <$> readIORef spansRef)++spanHotSnapshot :: Otel.ImmutableSpan -> IO Otel.SpanHot+spanHotSnapshot = readIORef . Otel.spanHot++deprecatedGenAiSystemKey :: Text+deprecatedGenAiSystemKey = "gen_ai." <> "system"++successSpanTest :: TestTree+successSpanTest =+ testCase "success path emits one Ok span with expected attributes" $ do+ let a = Custom "baikai-otel-success"+ registerOk a+ (tracer, getSpans) <- newTracerWithInMemory+ let sink = otelSink tracer+ _ <- withTrace sink (stubModel a) stubContext stubOptions+ spans <- getSpans+ assertEqual "exactly one span recorded" 1 (length spans)+ case spans of+ [sp] -> do+ hot <- spanHotSnapshot sp+ Otel.hotName hot @?= "baikai.call"+ let attrs = Attr.getAttributeMap (Otel.hotAttributes hot)+ assertBool+ ("has gen_ai.provider.name; got keys: " <> show (HashMap.keys attrs))+ (HashMap.member "gen_ai.provider.name" attrs)+ assertBool "has gen_ai.operation.name" (HashMap.member "gen_ai.operation.name" attrs)+ assertBool "has gen_ai.request.model" (HashMap.member "gen_ai.request.model" attrs)+ assertBool "has gen_ai.request.max_tokens" (HashMap.member "gen_ai.request.max_tokens" attrs)+ assertBool "has gen_ai.response.model" (HashMap.member "gen_ai.response.model" attrs)+ assertBool "has gen_ai.usage.input_tokens" (HashMap.member "gen_ai.usage.input_tokens" attrs)+ assertBool "has gen_ai.usage.output_tokens" (HashMap.member "gen_ai.usage.output_tokens" attrs)+ assertBool "has baikai.event_id" (HashMap.member "baikai.event_id" attrs)+ assertBool "has baikai.latency_ms" (HashMap.member "baikai.latency_ms" attrs)+ assertBool "does not emit deprecated GenAI system key" (not (HashMap.member deprecatedGenAiSystemKey attrs))+ case Otel.hotStatus hot of+ Otel.Ok -> pure ()+ other -> assertFailure ("expected Ok status, got: " <> show other)+ case Otel.spanKind sp of+ Otel.Client -> pure ()+ other -> assertFailure ("expected Client kind, got: " <> show other)+ _ -> assertFailure "expected exactly one span"++failureSpanTest :: TestTree+failureSpanTest =+ testCase "failure path emits one Error span with error message" $ do+ let a = Custom "baikai-otel-failure"+ registerFail a (ProviderError "stub-otel-boom")+ (tracer, getSpans) <- newTracerWithInMemory+ let sink = otelSink tracer+ -- withTrace no longer re-throws producer failures; the error+ -- surfaces as ErrorReason on the response and as the OTel span's+ -- Error status.+ resp <- withTrace sink (stubModel a) stubContext stubOptions+ let AssistantPayload {stopReason = sr} = resp ^. #message+ sr @?= ErrorReason+ spans <- getSpans+ 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 error to mention stub-otel-boom; got: " <> show msg)+ (not (null (show msg)))+ other -> assertFailure ("expected Error status, got: " <> show other)+ let attrs = Attr.getAttributeMap (Otel.hotAttributes hot)+ 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"