Embeddings CLI

src/Embeddings.hs

@@ -0,0 +1,216 @@
+{-# LANGUAGE LambdaCase #-}
+{-# LANGUAGE OverloadedStrings #-}
+{-# OPTIONS_GHC -Wall #-}
+
+module Embeddings
+  ( generateEmbedding,
+    indexEntries,
+    search,
+  ) where
+
+import Control.Exception.Annotated (checkpointCallStack)
+import Control.Lens (re, traversed, view, (&), (.~), (^.), (^..), (^?))
+import Control.Monad (when)
+import Control.Monad.Catch (MonadCatch, MonadThrow)
+import Control.Monad.Reader (MonadReader, liftIO)
+import Creds (Creds, openrouterKey)
+import Data.Aeson (decodeStrict, object, (.=))
+import Data.Aeson.Lens (key, _Array, _Double)
+import Data.Functor.Contravariant ((>$<))
+import Data.Generics.Product (HasType, typed)
+import Data.List (sortOn)
+import Data.Maybe (mapMaybe)
+import Data.Ord (Down (..))
+import Data.Text (Text)
+import Data.Text qualified as T
+import Data.Text.Encoding (encodeUtf8)
+import Data.Text.Strict.Lens (utf8)
+import Data.Vector.Unboxed qualified as UV
+import Database.SQLite.Simple (Only (..), open, close)
+import Database.SQLite.Simple qualified as SQL
+import GHC.Stack (HasCallStack)
+import HTTP (HTTPMetric, Retry (SafeTransient), Trace, logFormatter, postWith, tracerMetrics)
+import Katip qualified as K
+import M (HasEmbeddingMetrics, HasKatipTracer, HasWallabagMetrics, embeddingMetrics, katipTracer, KatipTracer)
+import Network.Wreq (defaults, header, responseBody)
+import Network.Wreq.Session (Session)
+import Pipes (lift)
+import Pipes qualified
+import Tracer (Tracer (..))
+import UnliftIO (MonadIO, MonadUnliftIO)
+import Wallabag (Entry, Scan (Everything), content, identifier, title, url)
+import qualified Wallabag
+import Wallabag.Auth (Auth)
+
+embeddingUrl :: String
+embeddingUrl = "https://openrouter.ai/api/v1/embeddings"
+
+mkTracerMetrics ::
+  ( MonadReader r m,
+    HasEmbeddingMetrics r HTTPMetric,
+    HasKatipTracer r KatipTracer
+  ) =>
+  m (Tracer Trace)
+mkTracerMetrics = do
+  l <- view katipTracer
+  m <- view embeddingMetrics
+  pure $ (logFormatter >$< l) <> tracerMetrics m
+
+generateEmbedding ::
+  ( MonadIO m,
+    HasType Session r,
+    HasType Creds r,
+    MonadReader r m,
+    HasEmbeddingMetrics r HTTPMetric,
+    HasKatipTracer r KatipTracer,
+    MonadCatch m,
+    K.KatipContext m
+  ) =>
+  Text ->
+  m (Maybe (UV.Vector Float))
+generateEmbedding input = checkpointCallStack $ do
+  c <- view $ typed @Creds
+  let opts = defaults & header "Authorization" .~ ["Bearer " <> c ^. openrouterKey . re utf8]
+  let j =
+        object
+          [ "model" .= ("text-embedding-3-small" :: Text),
+            "input" .= input
+          ]
+  tm <- mkTracerMetrics
+  session <- view $ typed @Session
+  r <- liftIO $ postWith tm session SafeTransient opts embeddingUrl j
+  case r ^? responseBody . key "data" . _Array . traversed . key "embedding" of
+    Just embeddingVal -> do
+      let floats = embeddingVal ^.. _Array . traversed . _Double
+      if null floats
+        then do
+          K.logFM K.ErrorS $ K.ls $ "Parsing gave empty embedding array from: " <> T.pack (show (r ^. responseBody))
+          pure Nothing
+        else pure $ Just $ UV.fromList (map realToFrac floats)
+    Nothing -> do
+      K.logFM K.ErrorS $ K.ls $ "Couldn't find embedding in response: " <> T.pack (show (r ^. responseBody))
+      pure Nothing
+
+cosineSimilarity :: UV.Vector Float -> UV.Vector Float -> Float
+cosineSimilarity a b = dot / (normA * normB)
+  where
+    dot = UV.sum $ UV.zipWith (*) a b
+    normA = sqrt $ UV.sum $ UV.map (^ (2 :: Int)) a
+    normB = sqrt $ UV.sum $ UV.map (^ (2 :: Int)) b
+
+indexEntries ::
+  ( MonadIO m,
+    HasType Session r,
+    HasType Auth r,
+    HasType Creds r,
+    MonadReader r m,
+    K.KatipContext m,
+    HasWallabagMetrics r HTTPMetric,
+    HasEmbeddingMetrics r HTTPMetric,
+    HasKatipTracer r KatipTracer,
+    MonadThrow m,
+    MonadCatch m,
+    HasCallStack,
+    MonadUnliftIO m
+  ) =>
+  FilePath ->
+  Bool ->
+  m ()
+indexEntries dbPath rebuild = do
+  liftIO $ do
+    conn <- open dbPath
+    SQL.execute_ conn
+      "CREATE TABLE IF NOT EXISTS embeddings (entry_id INTEGER PRIMARY KEY, url TEXT NOT NULL, title TEXT NOT NULL, embedding TEXT NOT NULL, created_at TEXT NOT NULL DEFAULT (datetime('now')))"
+    when rebuild $ SQL.execute_ conn "DELETE FROM embeddings"
+    close conn
+  Pipes.runEffect $
+    Pipes.for
+      (Wallabag.sourceEntries Everything)
+      (lift . processEntry dbPath)
+
+maxEmbeddingTokens :: Int
+maxEmbeddingTokens = 8192
+
+-- Rough token estimate: ~3 chars per token (conservative for non-English text)
+-- Used to pre-filter content before the API call rejects it.
+tokensApprox :: Text -> Int
+tokensApprox t = T.length t `div` 3
+
+processEntry ::
+  ( MonadIO m,
+    HasType Session r,
+    HasType Creds r,
+    MonadReader r m,
+    HasEmbeddingMetrics r HTTPMetric,
+    HasKatipTracer r KatipTracer,
+    K.KatipContext m,
+    MonadCatch m
+  ) =>
+  FilePath ->
+  Entry ->
+  m ()
+processEntry dbPath e = do
+  conn <- liftIO $ open dbPath
+  existing <- liftIO $ (SQL.query conn "SELECT 1 FROM embeddings WHERE entry_id = ?" (Only (e ^. identifier)) :: IO [Only Int])
+  if null existing
+    then do
+      K.logFM K.InfoS $ K.ls $ "Indexing: " <> e ^. title
+      let entryContent = e ^. content
+      if tokensApprox entryContent > maxEmbeddingTokens
+        then do
+          K.logFM K.WarningS $ K.ls $ "Skipping (content too long, ~" <> T.pack (show (tokensApprox entryContent)) <> " estimated tokens): " <> e ^. title
+          liftIO $ close conn
+        else do
+          mVec <- generateEmbedding entryContent
+          case mVec of
+            Just vec -> do
+              let embeddingJson = T.pack $ show $ UV.toList vec
+              liftIO $
+                SQL.execute conn
+                  "INSERT OR REPLACE INTO embeddings (entry_id, url, title, embedding) VALUES (?, ?, ?, ?)"
+                  (e ^. identifier, e ^. url, e ^. title, embeddingJson)
+            Nothing ->
+              K.logFM K.WarningS $ K.ls $ "Failed to generate embedding for: " <> e ^. title
+          liftIO $ close conn
+    else do
+      K.logFM K.DebugS $ K.ls $ "Skipping (already indexed): " <> e ^. title
+      liftIO $ close conn
+
+search ::
+  ( MonadIO m,
+    HasType Session r,
+    HasType Creds r,
+    MonadReader r m,
+    HasEmbeddingMetrics r HTTPMetric,
+    HasKatipTracer r KatipTracer,
+    MonadCatch m,
+    K.KatipContext m
+  ) =>
+  FilePath ->
+  Text ->
+  Int ->
+  m [(Double, Integer, Text, Text)]
+search dbPath query limit_ = do
+  queryVec <- generateEmbedding query
+  case queryVec of
+    Nothing -> pure []
+    Just qv -> do
+      conn <- liftIO $ open dbPath
+      rows <- liftIO $ SQL.query_ conn "SELECT entry_id, url, title, embedding FROM embeddings"
+      let scored = mapMaybe (scoreRow qv) rows
+          sorted = sortOn (Down . (\(s, _, _, _) -> s)) scored
+          topN = take limit_ sorted
+      liftIO $ close conn
+      pure topN
+
+scoreRow ::
+  UV.Vector Float ->
+  (Integer, Text, Text, Text) ->
+  Maybe (Double, Integer, Text, Text)
+scoreRow qv (eid, url', title', embeddingJson) =
+  case decodeStrict (encodeUtf8 embeddingJson) of
+    Just (vec :: [Double]) ->
+      let uv = UV.fromList (map realToFrac vec)
+          sim = cosineSimilarity qv uv
+       in Just (realToFrac sim, eid, url', title')
+    Nothing -> Nothing

src/M.hs

@@ -16,7 +16,7 @@ import Control.Monad.Catch (MonadCatch, MonadThrow)
 import Control.Monad.Reader (MonadIO, MonadReader, ReaderT, local, runReaderT)
 import Control.Monad.Trans.Control (MonadBaseControl)
 import Creds
-import Data.HashMap.Strict qualified as HM
+
 import GHC.Generics (Generic)
 import HTTP (HTTPMetric (HTTPMetric))
 import Katip qualified as K
@@ -37,6 +37,7 @@ data S = S
     _sCreds :: Creds,
     _sWallabagMetrics :: HTTPMetric,
     _sOpenrouterMetrics :: HTTPMetric,
+    _sEmbeddingMetrics :: HTTPMetric,
     _sKatipTracer :: KatipTracer,
     _sMetricsStore :: EKG.Store,
     _sSession :: Wreq.Session
@@ -92,6 +93,11 @@ initialS credsPath severity f = do
       <$> EKG.createDistribution "openrouter.requests" store
       <*> EKG.createCounter "openrouter.retries" store
       <*> EKG.createCounter "openrouter.failures" store
+  emMetrics <-
+    HTTPMetric
+      <$> EKG.createDistribution "embeddings.requests" store
+      <*> EKG.createCounter "embeddings.retries" store
+      <*> EKG.createCounter "embeddings.failures" store
   c <- maybe getDefaultCreds readCreds credsPath
   handleScribe <- K.mkHandleScribe K.ColorIfTerminal stdout (K.permitItem severity) K.V2
   let makeLogEnv = K.registerScribe "stdout" handleScribe K.defaultScribeSettings =<< K.initLogEnv "wallabag" "production"
@@ -99,7 +105,5 @@ initialS credsPath severity f = do
     let kTrace = mkKatipTrace le ns
         ns = "main"
     cache <- newTVarIO Unauthenticated
-    f (S cache ns le mempty c wbMetrics orMetrics kTrace store wsession)
-  -- TODO: This is just for debug purposes
-  EKG.sampleAll store >>= mapM_ print . HM.toList
+    f (S cache ns le mempty c wbMetrics orMetrics emMetrics kTrace store wsession)
   pure res

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

src/Wallabag.hs

@@ -46,7 +46,8 @@ data Entry = Entry
   { _entryIdentifier :: Integer,
     _entryContent :: Text,
     _entryTags :: [Text],
-    _entryTitle :: Text
+    _entryTitle :: Text,
+    _entryUrl :: Text
   }
   deriving (Show)
 
@@ -168,6 +169,7 @@ sourceEntries s = P.unfoldr (toEither <=< go) (Just startUrl) >-> P.mapFoldable
         <*> v ^? key "content" . _String
         <*> Just (v ^.. key "tags" . _Array . traversed . key "label" . _String)
         <*> v ^? key "title" . _String
+        <*> v ^? key "url" . _String
 
 mkTracerMetrics ::
   ( MonadReader r m,

wallabag-tagger.cabal

@@ -19,6 +19,7 @@ common warnings
 library
     import:           warnings
     exposed-modules:  AI
+                      Embeddings
                       Creds
                       M
                       Wallabag
@@ -31,6 +32,8 @@ library
     -- other-modules:
     -- other-extensions:
     build-depends:    base ^>=4.20
+                    , sqlite-simple
+                    , vector
                     , wreq
                     , lens
                     , katip