361 lines
11 KiB
Haskell
361 lines
11 KiB
Haskell
{-# 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, MonadMask)
|
|
import Control.Monad.Reader (MonadReader, liftIO)
|
|
import Creds (Creds, openrouterKey)
|
|
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.Ord (Down (..))
|
|
import Data.Text (Text)
|
|
import Data.Text qualified as T
|
|
import Data.Text.Strict.Lens (utf8)
|
|
import Data.Vector.Unboxed qualified as UV
|
|
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)
|
|
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)
|
|
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"
|
|
|
|
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 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 $
|
|
Pipes.for
|
|
(Wallabag.sourceEntries Everything)
|
|
(lift . processEntry dbPath)
|
|
|
|
chunkMaxTokens :: Int
|
|
chunkMaxTokens = 6000
|
|
|
|
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 :: Int -> [Text] -> [Text]
|
|
coalesceBlocks titleSize blocks = reverse (go [] 0 blocks)
|
|
where
|
|
chunkMaxChars = (chunkMaxTokens * 3) - titleSize
|
|
|
|
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 titleSize (extractBlocks html))
|
|
where
|
|
prefixTitle t chunk = "Title: " <> t <> "\n\n" <> chunk
|
|
titleSize = T.length ("Title: " <> title_)
|
|
|
|
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
|
|
let chunks = chunkArticle (e ^. title) (e ^. content)
|
|
if null chunks
|
|
then do
|
|
K.logFM K.WarningS $ K.ls $ "Skipping (no extractable content): " <> e ^. title
|
|
liftIO $ close conn
|
|
else do
|
|
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,
|
|
HasType Creds r,
|
|
MonadReader r m,
|
|
HasEmbeddingMetrics r HTTPMetric,
|
|
HasKatipTracer r KatipTracer,
|
|
MonadCatch m,
|
|
K.KatipContext m,
|
|
MonadMask m
|
|
) =>
|
|
FilePath ->
|
|
Text ->
|
|
Int ->
|
|
m [(Double, Integer, Text, Text)]
|
|
search dbPath q limit_ = do
|
|
queryVec <- generateEmbedding q
|
|
case queryVec of
|
|
Nothing -> pure []
|
|
Just qv -> do
|
|
conn <- liftIO $ open dbPath
|
|
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 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 ->
|
|
EmbeddingRow ->
|
|
Maybe (Float, Integer, Text, Text)
|
|
scoreRow qv EmbeddingRow{embedding = Embedding uv, ..} =
|
|
let sim = cosineSimilarity qv uv
|
|
in Just (sim, entryId, entryUrl, entryTitle)
|