From f5197caa7738543224eb0d600da676fa8a08a5e1 Mon Sep 17 00:00:00 2001 From: Mats Rauhala Date: Thu, 11 Jun 2026 15:13:29 +0300 Subject: [PATCH] Chunking --- default.nix | 28 +++--- src/Embeddings.hs | 226 ++++++++++++++++++++++++++++++++++-------- src/M.hs | 5 +- src/MyLib.hs | 9 +- src/Wallabag.hs | 10 +- test/Main.hs | 40 +++++++- wallabag-tagger.cabal | 10 +- 7 files changed, 259 insertions(+), 69 deletions(-) diff --git a/default.nix b/default.nix index 0bedb66..888746a 100644 --- a/default.nix +++ b/default.nix @@ -1,10 +1,11 @@ { mkDerivation, aeson, annotated-exception, base, bytestring -, containers, directory, ekg-core, exceptions, filepath -, generic-lens, http-client, http-types, katip, lens, lens-aeson -, lib, monad-control, mtl, optparse-generic, pandoc, pipes -, pipes-concurrency, retry, servant, servant-server, stm, streamly +, cereal, containers, directory, ekg-core, exceptions, filepath +, generic-lens, hedgehog, hspec, hspec-hedgehog, http-client +, http-types, katip, lens, lens-aeson, lib, monad-control, mtl +, optparse-generic, pandoc, pipes, pipes-concurrency, pipes-safe +, retry, servant, servant-server, sqlite-simple, stm, streamly , streamly-core, template-haskell, text, time, transformers-base -, unliftio, unordered-containers, wai, warp, wreq +, unliftio, unordered-containers, vector, wai, warp, wreq }: mkDerivation { pname = "wallabag-tagger"; @@ -13,15 +14,18 @@ mkDerivation { isLibrary = true; isExecutable = true; libraryHaskellDepends = [ - aeson annotated-exception base bytestring containers directory - ekg-core exceptions filepath generic-lens http-client http-types - katip lens lens-aeson monad-control mtl optparse-generic pandoc - pipes pipes-concurrency retry servant servant-server stm streamly - streamly-core template-haskell text time transformers-base unliftio - unordered-containers wai warp wreq + aeson annotated-exception base bytestring cereal containers + directory ekg-core exceptions filepath generic-lens http-client + http-types katip lens lens-aeson monad-control mtl optparse-generic + pandoc pipes pipes-concurrency pipes-safe retry servant + servant-server sqlite-simple stm streamly streamly-core + template-haskell text time transformers-base unliftio + unordered-containers vector wai warp wreq ]; executableHaskellDepends = [ base ]; - testHaskellDepends = [ base ]; + testHaskellDepends = [ + base hedgehog hspec hspec-hedgehog sqlite-simple text vector + ]; license = lib.meta.getLicenseFromSpdxId "BSD-3-Clause"; mainProgram = "wallabag-tagger"; } diff --git a/src/Embeddings.hs b/src/Embeddings.hs index 099b564..b0f746f 100644 --- a/src/Embeddings.hs +++ b/src/Embeddings.hs @@ -1,32 +1,38 @@ {-# LANGUAGE LambdaCase #-} {-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE StrictData #-} {-# OPTIONS_GHC -Wall #-} +{-# LANGUAGE RecordWildCards #-} module Embeddings ( generateEmbedding, indexEntries, search, + chunkArticle, + tokensApprox, + chunkMaxTokens, + Embedding(..), -- for testing ) 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.Catch (MonadCatch, MonadThrow, MonadMask) import Control.Monad.Reader (MonadReader, liftIO) import Creds (Creds, openrouterKey) -import Data.Aeson (decodeStrict, object, (.=)) +import Data.Aeson (object, (.=)) import Data.Aeson.Lens (key, _Array, _Double) import Data.Functor.Contravariant ((>$<)) import Data.Generics.Product (HasType, typed) import Data.List (sortOn) +import Data.Map.Strict qualified as Map 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 (Only (..), open, close, SQLData (SQLBlob), ResultError (ConversionFailed)) import Database.SQLite.Simple qualified as SQL import GHC.Stack (HasCallStack) import HTTP (HTTPMetric, Retry (SafeTransient), Trace, logFormatter, postWith, tracerMetrics) @@ -34,13 +40,20 @@ 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 (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) +import GHC.Generics (Generic) +import Database.SQLite.Simple.ToField (ToField(..)) +import Database.SQLite.Simple.FromField (FromField(..), fieldData, returnError) +import qualified Data.Serialize as Cereal +import qualified Pipes.Safe +import Pipes.Safe (MonadSafe) +import qualified Pipes.Prelude embeddingUrl :: String embeddingUrl = "https://openrouter.ai/api/v1/embeddings" @@ -120,7 +133,7 @@ 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')))" + "CREATE TABLE IF NOT EXISTS embeddings (entry_id INTEGER NOT NULL, chunk_idx INTEGER NOT NULL, url TEXT NOT NULL, title TEXT NOT NULL, embedding TEXT NOT NULL, created_at TEXT NOT NULL DEFAULT (datetime('now')), PRIMARY KEY (entry_id, chunk_idx))" when rebuild $ SQL.execute_ conn "DELETE FROM embeddings" close conn Pipes.runEffect $ @@ -128,14 +141,85 @@ indexEntries dbPath rebuild = do (Wallabag.sourceEntries Everything) (lift . processEntry dbPath) -maxEmbeddingTokens :: Int -maxEmbeddingTokens = 8192 +chunkMaxTokens :: Int +chunkMaxTokens = 6000 --- 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 +stripTags :: Text -> Text +stripTags t = + case T.breakOn "<" t of + (before, rest) + | T.null rest -> before + | otherwise -> + case T.breakOn ">" (T.drop 1 rest) of + (_, after) -> stripTags before <> stripTags (T.drop 1 after) + +normalize :: Text -> Text +normalize = T.unwords . T.words + +blockEndTags :: [Text] +blockEndTags = + ["

", "", "", "", ""] + <> map (\n -> " T.pack (show n) <> ">") [1 .. 6 :: Int] + +sentinel :: Text +sentinel = "\x00BLOCK\x00" + +replaceBlockEnds :: Text -> Text +replaceBlockEnds html = foldr (\tag -> T.replace tag sentinel) html blockEndTags + +extractBlocks :: Text -> [Text] +extractBlocks html = + map normalize + . map stripTags + . filter (not . T.null) + $ T.splitOn sentinel (replaceBlockEnds html) + +splitSentences :: Text -> [Text] +splitSentences t = + filter (not . T.null) + . map T.strip + $ T.splitOn sent + $ T.replace ". " (". " <> sent) + $ T.replace "? " ("? " <> sent) + $ T.replace "! " ("! " <> sent) + $ t + where + sent = "\x00SENT\x00" + +hardWrap :: Text -> [Text] +hardWrap t + | T.null t = [] + | otherwise = let (a, b) = T.splitAt (chunkMaxTokens * 3) t in a : hardWrap b + +coalesceBlocks :: [Text] -> [Text] +coalesceBlocks blocks = reverse (go [] 0 blocks) + where + chunkMaxChars = chunkMaxTokens * 3 + + go :: [Text] -> Int -> [Text] -> [Text] + go acc _ [] = [T.concat (reverse acc) | not (null acc)] + go acc len (x : xs) + | T.length x > chunkMaxChars && null acc = + let sentences = splitSentences x + splitChunks = + if sentences == [x] + then hardWrap x + else sentences + in go acc len (splitChunks ++ xs) + | not (null acc) && len + T.length x > chunkMaxChars = + T.concat (reverse acc) : go [] 0 (x : xs) + | otherwise = + go (x : acc) (len + T.length x) xs + +chunkArticle :: Text -> Text -> [Text] +chunkArticle title_ html = + map (prefixTitle title_) (coalesceBlocks (extractBlocks html)) + where + prefixTitle t chunk = "Title: " <> t <> "\n\n" <> chunk + processEntry :: ( MonadIO m, HasType Session r, @@ -154,28 +238,74 @@ processEntry dbPath e = do 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 + let chunks = chunkArticle (e ^. title) (e ^. content) + if null chunks then do - K.logFM K.WarningS $ K.ls $ "Skipping (content too long, ~" <> T.pack (show (tokensApprox entryContent)) <> " estimated tokens): " <> e ^. title + K.logFM K.WarningS $ K.ls $ "Skipping (no extractable content): " <> 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 + K.logFM K.InfoS $ K.ls $ "Indexing: " <> e ^. title <> " (" <> T.pack (show (length chunks)) <> " chunks)" + mapM_ (processChunk conn (e ^. identifier) (e ^. url) (e ^. title)) (zip [0 :: Int ..] chunks) liftIO $ close conn else do K.logFM K.DebugS $ K.ls $ "Skipping (already indexed): " <> e ^. title liftIO $ close conn + +newtype Embedding = Embedding (UV.Vector Float) + deriving (Show, Eq) + + +instance ToField Embedding where + toField (Embedding vec) = toField $ Cereal.encode $ UV.toList vec + +instance FromField Embedding where + fromField f = do + case fieldData f of + SQLBlob bs -> either (returnError ConversionFailed f) (pure . Embedding . UV.fromList) (Cereal.decode bs) + _ -> returnError ConversionFailed f "Failed to convert to embedding vector" + +data EmbeddingRow = EmbeddingRow + { entryId :: Integer + , chunkIdx :: Int + , entryUrl :: Text + , entryTitle :: Text + , embedding :: Embedding + } + deriving (Generic) + +instance SQL.ToRow EmbeddingRow + +instance SQL.FromRow EmbeddingRow + +processChunk :: + ( MonadIO m, + HasType Session r, + HasType Creds r, + MonadReader r m, + HasEmbeddingMetrics r HTTPMetric, + HasKatipTracer r KatipTracer, + K.KatipContext m, + MonadCatch m + ) => + SQL.Connection -> + Integer -> + Text -> + Text -> + (Int, Text) -> + m () +processChunk conn eid url_ title_ (idx, chunk) = do + mVec <- generateEmbedding chunk + case mVec of + Just vec -> do + let embeddingRow = EmbeddingRow eid idx url_ title_ (Embedding vec) + liftIO $ SQL.execute conn + "INSERT OR REPLACE INTO embeddings (entry_id, chunk_idx, url, title, embedding) VALUES (?, ?, ?, ?, ?)" + embeddingRow + Nothing -> + K.logFM K.WarningS $ K.ls $ "Failed to generate embedding for chunk " <> T.pack (show idx) <> " of: " <> title_ + + search :: ( MonadIO m, HasType Session r, @@ -184,33 +314,47 @@ search :: HasEmbeddingMetrics r HTTPMetric, HasKatipTracer r KatipTracer, MonadCatch m, - K.KatipContext m + K.KatipContext m, + MonadMask m ) => FilePath -> Text -> Int -> m [(Double, Integer, Text, Text)] -search dbPath query limit_ = do - queryVec <- generateEmbedding query +search dbPath q limit_ = do + queryVec <- generateEmbedding q 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 + results <- Pipes.Safe.runSafeT $ Pipes.runEffect (Pipes.Prelude.fold bestChunk Map.empty topResults $ query_ conn "SELECT entry_id, chunk_idx, url, title, embedding FROM embeddings" >-> Pipes.Prelude.mapMaybe (scoreRow qv)) liftIO $ close conn - pure topN + pure results + where + bestChunk m (s, eid, url_, title_) = + Map.insertWith + (\(s1, u, t) (s2, _, _) -> (max s1 s2, u, t)) + eid + (s, url_, title_) + m + sortByScore = sortOn (Down . (\(_, (s, _, _)) -> s)) . Map.toList + topResults = map (\(eid, (s, url_, title_)) -> (realToFrac s, eid, url_, title_)) . take limit_ . sortByScore + + +query_ :: (MonadSafe m, SQL.FromRow a) => SQL.Connection -> SQL.Query -> Pipes.Producer' a m () +query_ conn q = Pipes.Safe.bracket (liftIO openStmt) (liftIO . SQL.closeStatement) eachRow + where + openStmt = do + SQL.openStatement conn q + eachRow stmt = do + row <- liftIO $ SQL.nextRow stmt + maybe (pure ()) (\r -> Pipes.yield r >> eachRow stmt) row + 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 + EmbeddingRow -> + Maybe (Float, Integer, Text, Text) +scoreRow qv EmbeddingRow{embedding = Embedding uv, ..} = + let sim = cosineSimilarity qv uv + in Just (sim, entryId, entryUrl, entryTitle) diff --git a/src/M.hs b/src/M.hs index d51678b..e49b1f8 100644 --- a/src/M.hs +++ b/src/M.hs @@ -12,7 +12,7 @@ import Control.Exception (bracket) import Control.Lens (over, view) import Control.Lens.TH (makeFields) import Control.Monad.Base (MonadBase) -import Control.Monad.Catch (MonadCatch, MonadThrow) +import Control.Monad.Catch (MonadCatch, MonadThrow, MonadMask) import Control.Monad.Reader (MonadIO, MonadReader, ReaderT, local, runReaderT) import Control.Monad.Trans.Control (MonadBaseControl) import Creds @@ -57,7 +57,8 @@ newtype M a = M (ReaderT S IO a) MonadThrow, MonadCatch, MonadBaseControl IO, - MonadBase IO + MonadBase IO, + MonadMask ) runM :: S -> M a -> IO a diff --git a/src/MyLib.hs b/src/MyLib.hs index 808fbfd..c49ada0 100644 --- a/src/MyLib.hs +++ b/src/MyLib.hs @@ -44,7 +44,6 @@ data Options | 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" } @@ -203,10 +202,6 @@ runWithDb mbDb sev f = initialS Nothing sev $ \s -> runM s $ do 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" @@ -214,13 +209,13 @@ defaultMain = do Quick {verbose} -> run (severity verbose) Unarchived Full {verbose} -> run (severity verbose) Everything Version -> print Version.compileTimeVersion - Index {verbose, credentials, rebuild, embeddingsDb} -> + 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 $ mapM_ (\(score, _eid, url, _title) -> T.putStrLn $ T.pack $ printf "%.4f %s" (score :: Double) (T.unpack url)) entries + liftIO $ mapM_ (\(score, _eid, u, _title) -> T.putStrLn $ T.pack $ printf "%.4f %s" (score :: Double) (T.unpack u)) entries Daemon {verbose, credentials, port, embeddingsDb} -> runDaemon (unHelpful $ unDefValue port) (unHelpful credentials) (severity verbose) where diff --git a/src/Wallabag.hs b/src/Wallabag.hs index 7b788ea..5b73673 100644 --- a/src/Wallabag.hs +++ b/src/Wallabag.hs @@ -152,13 +152,13 @@ sourceEntries s = P.unfoldr (toEither <=< go) (Just startUrl) >-> P.mapFoldable Everything -> "https://app.wallabag.it/api/entries.json" go :: Maybe String -> m (Maybe ([Value], Maybe String)) go Nothing = pure Nothing - go (Just url) = checkpoint (Annotation url) $ do + go (Just u) = checkpoint (Annotation u) $ do metrics <- mkTracerMetrics t <- maybe (throw MissingTokenException) pure =<< getToken let opts = defaults & header "Authorization" .~ ["Bearer " <> t ^. re utf8] - K.logFM K.DebugS $ K.ls ("Fetching entries '" <> url <> "'") + K.logFM K.DebugS $ K.ls ("Fetching entries '" <> u <> "'") session <- view $ typed @Session - r <- liftIO $ getWith metrics session SafeTransient opts url + r <- liftIO $ getWith metrics session SafeTransient opts u let items = r ^.. responseBody . key "_embedded" . key "items" . _Array . traversed let next = r ^? responseBody . key "_links" . key "next" . key "href" . _String . unpacked pure $ Just (items, next) @@ -202,8 +202,8 @@ applyTags e = checkpoint (Annotation $ e ^. identifier) $ do t <- maybe (throw MissingTokenException) pure =<< getToken let opts = defaults & header "Authorization" .~ ["Bearer " <> t ^. re utf8] K.logFM K.InfoS (K.ls ("Applying tags for '" <> e ^. title <> "': " <> T.pack (show (e ^. tags)))) - let url = "https://app.wallabag.it/api/entries/" <> show (e ^. identifier) <> "/tags" + let u = "https://app.wallabag.it/api/entries/" <> show (e ^. identifier) <> "/tags" metrics <- mkTracerMetrics session <- view $ typed @Session - _ <- liftIO $ postWith metrics session SafeTransient opts url ["tags" := T.intercalate "," (e ^. tags)] + _ <- liftIO $ postWith metrics session SafeTransient opts u ["tags" := T.intercalate "," (e ^. tags)] pure () diff --git a/test/Main.hs b/test/Main.hs index 3e2059e..ab6a8f5 100644 --- a/test/Main.hs +++ b/test/Main.hs @@ -1,4 +1,42 @@ +{-# LANGUAGE OverloadedStrings #-} module Main (main) where +import Data.Foldable (forM_) +import Embeddings (chunkArticle, tokensApprox, chunkMaxTokens, Embedding(..)) +import Test.Hspec +import qualified Data.Text as T + main :: IO () -main = putStrLn "Test suite not yet implemented." +main = hspec $ do + describe "Chunking" $ do + it "Splits a long article into multiple chunks" testChunking + it "Retains a short article as one chunk" testShortArticleOneChunk + it "Prefixes each chunk with the title" testTitlePrefix + it "Chunks to small enough items" testTokenLimit + + +testChunking :: IO () +testChunking = do + let longHtml = T.replicate 1000 "

This is a paragraph with some text content that should be long enough to create multiple chunks when repeated many times.

" + chunks = chunkArticle "Test Article" longHtml + n = length chunks + n `shouldBe` 7 + +testShortArticleOneChunk :: IO () +testShortArticleOneChunk = do + let shortHtml = "

A short paragraph.

" + chunks = chunkArticle "Short" shortHtml + length chunks `shouldBe` 1 + +testTitlePrefix :: IO () +testTitlePrefix = do + let chunks = chunkArticle "My Title" "

Some content.

" + forM_ chunks $ \chunk -> + chunk `shouldSatisfy` ("Title: My Title\n\n" `T.isPrefixOf`) + +testTokenLimit :: IO () +testTokenLimit = do + let longHtml = T.replicate 1000 "

This is a paragraph with some text content that should be long enough to create multiple chunks when repeated many times.

" + chunks = chunkArticle "Test Article" longHtml + counts = map tokensApprox chunks + counts `shouldSatisfy` all (<= chunkMaxTokens) diff --git a/wallabag-tagger.cabal b/wallabag-tagger.cabal index a58e83a..6e4b2e0 100644 --- a/wallabag-tagger.cabal +++ b/wallabag-tagger.cabal @@ -42,6 +42,7 @@ library , ekg-core , retry , pipes + , pipes-safe , pipes-concurrency , streamly , streamly-core @@ -69,6 +70,7 @@ library , optparse-generic , containers , template-haskell + , cereal hs-source-dirs: src default-language: GHC2021 @@ -95,4 +97,10 @@ test-suite wallabag-tagger-test main-is: Main.hs build-depends: base ^>=4.20, - wallabag-tagger + wallabag-tagger, + text, + hspec, + hspec-hedgehog, + hedgehog, + vector, + sqlite-simple