diff --git a/LICENSE b/LICENSE new file mode 100644 --- /dev/null +++ b/LICENSE @@ -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. diff --git a/baikai-trace-otel.cabal b/baikai-trace-otel.cabal new file mode 100644 --- /dev/null +++ b/baikai-trace-otel.cabal @@ -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 diff --git a/src/Baikai/Trace/Sink/OpenTelemetry.hs b/src/Baikai/Trace/Sink/OpenTelemetry.hs new file mode 100644 --- /dev/null +++ b/src/Baikai/Trace/Sink/OpenTelemetry.hs @@ -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 diff --git a/test/Main.hs b/test/Main.hs new file mode 100644 --- /dev/null +++ b/test/Main.hs @@ -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"