Embeddings CLI

src/MyLib.hs

@@ -14,7 +14,9 @@ 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)
@@ -22,7 +24,11 @@ import Options.Generic (Generic, ParseRecord (..), getRecord, unDefValue, unHelp
 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)
@@ -36,10 +42,23 @@ data Options
   = Quick {verbose :: Bool}
   | Full {verbose :: Bool}
   | Version
+  | Index
+      { verbose :: Bool,
+        credentials :: FilePath <?> "Path to the credentials.json",
+        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"
+        port :: Int <?> "Port to run the diagnostics" <!> "8094",
+        embeddingsDb :: Maybe FilePath <?> "Path to embeddings database"
       }
   deriving (Generic, Show)
 
@@ -166,6 +185,28 @@ hardLabels =
       "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
+
+printResult :: Int -> Text -> Double -> Text
+printResult n url score =
+  T.pack $ printf "%3d. %s (score: %.4f)" n (T.unpack url) score
+
 defaultMain :: IO ()
 defaultMain = do
   opts <- getRecord "Wallabag tagger"
@@ -173,7 +214,14 @@ defaultMain = do
     Quick {verbose} -> run (severity verbose) Unarchived
     Full {verbose} -> run (severity verbose) Everything
     Version -> print Version.compileTimeVersion
-    Daemon {verbose, credentials, port} ->
+    Index {verbose, credentials, 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 $ mapM_ (\(score, _eid, url, _title) -> T.putStrLn $ T.pack $ printf "%.4f  %s" (score :: Double) (T.unpack url)) entries
+    Daemon {verbose, credentials, port, embeddingsDb} ->
       runDaemon (unHelpful $ unDefValue port) (unHelpful credentials) (severity verbose)
   where
     severity True = K.DebugS