Chunking

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 =
+  ["</p>", "</div>", "</blockquote>", "</section>", "</li>"]
+    <> map (\n -> "</h" <> 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)