Metrics

src/AI.hs

@@ -3,19 +3,35 @@
 
 module AI where
 
-import Control.Lens (re, traversed, (&), (.~), (^.), (^?))
-import Control.Monad.Reader (liftIO)
+import Control.Exception.Annotated (checkpointCallStack)
+import Control.Lens (re, traversed, view, (&), (.~), (^.), (^?))
+import Control.Monad.Catch (MonadCatch)
+import Control.Monad.Reader (MonadReader, liftIO)
 import Creds
 import Data.Aeson (object, (.=))
 import Data.Aeson.Lens (key, _Array, _String)
 import Data.Text (Text)
 import Data.Text.Strict.Lens (utf8)
-import Network.Wreq (defaults, header, postWith, responseBody)
+import HTTP (HTTPMetric, Retry (SafeTransient), Trace, postWith, tracerMetrics)
+import M (HasOpenrouterMetrics, openrouterMetrics)
+import Network.Wreq (defaults, header, responseBody)
+import Tracer (Tracer (..))
 import UnliftIO (MonadIO)
-import Control.Exception.Annotated (checkpointCallStack)
-import Control.Monad.Catch (MonadCatch)
 
-completions :: (MonadIO m, MonadCatch m) => Creds -> Text -> m (Maybe Text)
+mkTracerMetrics :: (MonadReader r m, HasOpenrouterMetrics r HTTPMetric) => m (Tracer HTTP.Trace)
+mkTracerMetrics = do
+  m <- view openrouterMetrics
+  pure $ HTTP.tracerMetrics m
+
+completions ::
+  ( MonadIO m,
+    MonadCatch m,
+    MonadReader r0 m,
+    HasOpenrouterMetrics r0 HTTPMetric
+  ) =>
+  Creds ->
+  Text ->
+  m (Maybe Text)
 completions c prompt = checkpointCallStack $ do
   let opts = defaults & header "Authorization" .~ ["Bearer " <> c ^. openrouterKey . re utf8]
   let j =
@@ -24,5 +40,6 @@ completions c prompt = checkpointCallStack $ do
             "messages"
               .= [object ["role" .= ("user" :: Text), "content" .= prompt]]
           ]
-  r <- liftIO $ postWith opts "https://openrouter.ai/api/v1/chat/completions" j
+  metrics <- mkTracerMetrics
+  r <- liftIO $ postWith metrics SafeTransient opts "https://openrouter.ai/api/v1/chat/completions" j
   pure $ r ^? responseBody . key "choices" . _Array . traversed . key "message" . key "content" . _String