wallabag-tagger/src/MyLib.hs

{-# LANGUAGE DataKinds #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE TypeOperators #-}

module MyLib (defaultMain, test) where

import AI
import Control.Exception.Annotated.UnliftIO (Exception, SomeException, catch, throw)
import Control.Lens (over, re, strict, view, (&), (.~), (^.))
import Control.Monad (forever)
import Data.Aeson (decode, encode)
import Data.Either (fromRight)
import Data.Maybe (fromMaybe)
import Data.Set qualified as S
import Data.Text (Text)
import Data.Text qualified as T
import Data.Text.Strict.Lens (utf8)
import Katip qualified as K
import M
import Network.Wai.Handler.Warp (defaultSettings, runSettings, setPort)
import Options.Generic (Generic, ParseRecord (..), getRecord, unDefValue, unHelpful, type (<!>), type (<?>))
import Pipes (Pipe, for, lift, runEffect, (>->))
import Pipes.Prelude qualified as P
import Server (Health (..), app)
import Text.Pandoc (PandocError, def, readHtml, runPure, writeMarkdown)
import UnliftIO (MonadUnliftIO, liftIO)
import UnliftIO.Async (mapConcurrently_)
import UnliftIO.Concurrent (threadDelay)
import UnliftIO.Exception (displayException, isAsyncException)
import UnliftIO.IORef (IORef, newIORef, writeIORef)
import Wallabag
import Pipes.Concurrent (performGC)
import qualified Version

data Options
  = Quick {verbose :: Bool}
  | Full {verbose :: Bool}
  | Version
  | Daemon
      { verbose :: Bool,
        credentials :: FilePath <?> "Path to the credentials.json",
        port :: Int <?> "Port to run the diagnostics" <!> "8094"
      }
  deriving (Generic, Show)

instance ParseRecord Options

tooShort :: Entry -> Bool
tooShort e = length (T.lines (view content e)) < 5

missingTags :: Entry -> Bool
missingTags Entry {_entryTags = []} = True
missingTags Entry {} = False

htmlToMarkdown :: Text -> Either PandocError Text
htmlToMarkdown input = runPure $ do
  doc <- readHtml def input
  writeMarkdown def doc

processableEntries :: (Monad m) => Pipe Entry Entry m ()
processableEntries = P.filter (not . tooShort) >-> P.filter missingTags

-- Convert html into markdown
markdownEntry :: (Monad m) => Pipe Entry Entry m ()
markdownEntry = P.map (over content (\x -> fromRight x $ htmlToMarkdown x))

run :: K.Severity -> Scan -> IO ()
run sev scan = initialS Nothing sev $ \s -> do
  runM s $ runScan scan

runScan :: Scan -> M ()
runScan scan = do
  Pipes.runEffect $
    Pipes.for
      (sourceEntries scan >-> markdownEntry >-> processableEntries)
      (lift . doWork)
  liftIO performGC

runDaemon :: Int -> FilePath -> K.Severity -> IO ()
runDaemon diagnosticsPort credsPath sev = initialS (Just credsPath) sev $ \s -> runM s $ do
  let serverSettings = setPort diagnosticsPort defaultSettings
  health <- newIORef OK

  mapConcurrently_
    id
    [ backgroundRunner hourly $ worker health Unarchived,
      backgroundRunner weekly $ worker health Everything,
      worker health Unarchived, -- Do a quick scan on startup
      liftIO $ runSettings serverSettings (app s health)
    ]
  K.logFM K.InfoS "finished"
  where
    worker health scan = handle (reportError health) $ do
      runScan scan
      writeIORef health OK
    weekly = 7 * 86400 * 1_000_000
    hourly = 60 * 60 * 1_000_000
    backgroundRunner delay action = forever $ do
      threadDelay delay
      action
    reportError :: (K.KatipContext m, K.Katip m, MonadUnliftIO m) => IORef Health -> SomeException -> m ()
    reportError health e
      | isAsyncException e = throw e
      | otherwise = do
          writeIORef health Degraded
          K.logFM K.ErrorS $ K.ls $ displayException e
    handle :: (Exception e, MonadUnliftIO m) => (e -> m a) -> m a -> m a
    handle = flip catch

test :: IO ()
test = run K.DebugS Unarchived

suggestTags :: Text -> M [Text]
suggestTags txt = do
  c <- view creds
  let prompt =
        T.unlines
          [ "Assign tags for the given article from this set of tags: " <> encode hardLabels ^. strict . utf8,
            "Respond in JSON. Respond only in JSON, nothing more",
            "Example response: [\"Politics\"]",
            "",
            "The article contents:",
            txt
          ]
  unparsed <- completions c prompt
  pure $ maybe [] parseTags unparsed
  where
    parseTags :: Text -> [Text]
    parseTags = fromMaybe [] . decode . view (re utf8 . re strict) . T.dropWhileEnd (/= ']') . T.dropWhile (/= '[')

suggestTagsForEntry :: Entry -> M Entry
suggestTagsForEntry e = do
  newTags <- suggestTags (e ^. content)
  pure $ e & tags .~ newTags

doWork :: Entry -> M ()
doWork e = do
  entryWithTags <- suggestTagsForEntry e
  applyTags entryWithTags

hardLabels :: S.Set Text
hardLabels =
  S.fromList
    [ "Politics",
      "Writing",
      "Nature",
      "Camping",
      "Hiking",
      "Economics",
      "Climate",
      "Software",
      "Programming",
      "Maths",
      "Animals",
      "Health",
      "Sexuality",
      "Relationships",
      "Media and television",
      "Social media",
      "Gaming",
      "Plants",
      "Gadgets",
      "AI",
      "Emergency",
      "War",
      "Open Source"
    ]

defaultMain :: IO ()
defaultMain = do
  opts <- getRecord "Wallabag tagger"
  case opts of
    Quick {verbose} -> run (severity verbose) Unarchived
    Full {verbose} -> run (severity verbose) Everything
    Version -> print Version.compileTimeVersion
    Daemon {verbose, credentials, port} ->
      runDaemon (unHelpful $ unDefValue port) (unHelpful credentials) (severity verbose)
  where
    severity True = K.DebugS
    severity _ = K.InfoS