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