From f54aaa6e955d48d0f4f003b58dc3f12f0e790784 Mon Sep 17 00:00:00 2001 From: Mats Rauhala Date: Thu, 11 Jun 2026 09:49:09 +0300 Subject: [PATCH] Embeddings CLI --- src/Embeddings.hs | 216 ++++++++++++++++++++++++++++++++++++++++++ src/M.hs | 12 ++- src/MyLib.hs | 52 +++++++++- src/Wallabag.hs | 4 +- wallabag-tagger.cabal | 3 + 5 files changed, 280 insertions(+), 7 deletions(-) create mode 100644 src/Embeddings.hs diff --git a/src/Embeddings.hs b/src/Embeddings.hs new file mode 100644 index 0000000..099b564 --- /dev/null +++ b/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 diff --git a/src/M.hs b/src/M.hs index 224a1af..d51678b 100644 --- a/src/M.hs +++ b/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 diff --git a/src/MyLib.hs b/src/MyLib.hs index f5ae6dc..808fbfd 100644 --- a/src/MyLib.hs +++ b/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 diff --git a/src/Wallabag.hs b/src/Wallabag.hs index e5484b5..7b788ea 100644 --- a/src/Wallabag.hs +++ b/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, diff --git a/wallabag-tagger.cabal b/wallabag-tagger.cabal index dfc76f3..a58e83a 100644 --- a/wallabag-tagger.cabal +++ b/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