226 lines
7.3 KiB
Haskell
226 lines
7.3 KiB
Haskell
{-# 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 (forM_, 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.IO qualified as T
|
|
import Data.Text.Strict.Lens (utf8)
|
|
import Embeddings (indexEntries, search)
|
|
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 System.Directory (createDirectoryIfMissing, getHomeDirectory)
|
|
import System.Environment (lookupEnv)
|
|
import System.FilePath (takeDirectory)
|
|
import Text.Pandoc (PandocError, def, readHtml, runPure, writeMarkdown)
|
|
import Text.Printf (printf)
|
|
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
|
|
| Index
|
|
{ verbose :: Bool,
|
|
rebuild :: Bool <?> "Rebuild index from scratch" <!> "False",
|
|
embeddingsDb :: Maybe FilePath <?> "Path to embeddings database"
|
|
}
|
|
| Search
|
|
{ verbose :: Bool,
|
|
query :: Text <?> "Search query",
|
|
limit :: Int <?> "Maximum results" <!> "10",
|
|
embeddingsDb :: Maybe FilePath <?> "Path to embeddings database"
|
|
}
|
|
| Daemon
|
|
{ verbose :: Bool,
|
|
credentials :: FilePath <?> "Path to the credentials.json",
|
|
port :: Int <?> "Port to run the diagnostics" <!> "8094",
|
|
embeddingsDb :: Maybe FilePath <?> "Path to embeddings database"
|
|
}
|
|
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"
|
|
]
|
|
|
|
resolveEmbeddingsDb :: Maybe FilePath -> IO FilePath
|
|
resolveEmbeddingsDb (Just fp) = pure fp
|
|
resolveEmbeddingsDb Nothing = do
|
|
env <- lookupEnv "WALLABAG_EMBEDDINGS_DB"
|
|
case env of
|
|
Just fp -> pure fp
|
|
Nothing -> do
|
|
home <- getHomeDirectory
|
|
let dir = home <> "/.wallabag-tagger"
|
|
createDirectoryIfMissing True dir
|
|
pure $ dir <> "/embeddings.db"
|
|
|
|
runWithDb :: Maybe FilePath -> K.Severity -> (FilePath -> M ()) -> IO ()
|
|
runWithDb mbDb sev f = initialS Nothing sev $ \s -> runM s $ do
|
|
dbPath <- liftIO $ resolveEmbeddingsDb mbDb
|
|
liftIO $ createDirectoryIfMissing True (takeDirectory dbPath)
|
|
f dbPath
|
|
|
|
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
|
|
Index {verbose, rebuild, embeddingsDb} ->
|
|
runWithDb (unHelpful embeddingsDb) (severity verbose) $ \db ->
|
|
indexEntries db (unHelpful $ unDefValue rebuild)
|
|
Search {verbose, query, limit, embeddingsDb} ->
|
|
runWithDb (unHelpful embeddingsDb) (severity verbose) $ \db -> do
|
|
entries <- search db (unHelpful query) (unHelpful $ unDefValue limit)
|
|
liftIO $ forM_ entries $ \(score, _eid, u, t) -> do
|
|
T.putStrLn $ T.pack $ printf "%.4f %s" (score :: Double) (T.unpack t)
|
|
T.putStrLn $ T.pack $ printf " %s" (T.unpack u)
|
|
Daemon {verbose, credentials, port, embeddingsDb} ->
|
|
runDaemon (unHelpful $ unDefValue port) (unHelpful credentials) (severity verbose)
|
|
where
|
|
severity True = K.DebugS
|
|
severity _ = K.InfoS
|