43 lines
1.5 KiB
Haskell
43 lines
1.5 KiB
Haskell
{-# 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 = 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)
|