Chunking

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 "<p>This is a paragraph with some text content that should be long enough to create multiple chunks when repeated many times.</p>"
+      chunks = chunkArticle "Test Article" longHtml
+      n = length chunks
+  n `shouldBe` 7
+
+testShortArticleOneChunk :: IO ()
+testShortArticleOneChunk = do
+  let shortHtml = "<p>A short paragraph.</p>"
+      chunks = chunkArticle "Short" shortHtml
+  length chunks `shouldBe` 1
+
+testTitlePrefix :: IO ()
+testTitlePrefix = do
+  let chunks = chunkArticle "My Title" "<p>Some content.</p>"
+  forM_ chunks $ \chunk ->
+    chunk `shouldSatisfy` ("Title: My Title\n\n" `T.isPrefixOf`)
+
+testTokenLimit :: IO ()
+testTokenLimit = do
+  let longHtml = T.replicate 1000 "<p>This is a paragraph with some text content that should be long enough to create multiple chunks when repeated many times.</p>"
+      chunks = chunkArticle "Test Article" longHtml
+      counts = map tokensApprox chunks
+  counts `shouldSatisfy` all (<= chunkMaxTokens)