Model.hs

{-# LANGUAGE OverloadedStrings #-}

module Test.Langchain.Core.Model (tests) where

import Test.Tasty
import Test.Tasty.HUnit

import Control.Monad.Except (runExceptT)
import Data.List.NonEmpty (NonEmpty (..))
import qualified Data.Map.Strict as Map

import Langchain.Core.Model
import Test.Langchain.Core.TestModel (TestChatModel (..))

tests :: TestTree
tests =
  testGroup
    "Langchain.Core.Model"
    [ testGroup
        "Multi-Modal ContentBlock & Message"
        [ testCase "textMessage creates User message with TextBlock" $ do
            let msg = userMessage "Hello AI"
            messageRole msg @?= User
            extractMessageText msg @?= "Hello AI"
            messageMetadata msg @?= Map.empty
        , testCase "systemMessage creates System message" $ do
            let msg = systemMessage "You are a assistant"
            messageRole msg @?= System
            extractMessageText msg @?= "You are a assistant"
        , testCase "imageMessage creates ImageBlock message" $ do
            let msg = imageMessage User "image/png" "base64data=="
            messageRole msg @?= User
            case messageContents msg of
              (ImageBlock (ImageContent (ImageBase64 (Just mime) b64) Nothing Nothing) :| []) -> do
                mime @?= "image/png"
                b64 @?= "base64data=="
              _ -> assertFailure "Expected ImageBlock"
        ]
    , testGroup
        "Effect-Polymorphic ChatModel"
        [ testCase "invoke returns Assistant response" $ do
            let model = TestChatModel "Hello human" "mock-gpt"
                input = [userMessage "Hi"]
            res <- runExceptT $ invoke model input Nothing
            case res of
              Left err -> assertFailure $ "Unexpected error: " ++ show err
              Right msg -> do
                messageRole msg @?= Assistant
                extractMessageText msg @?= "Hello human"
        , testCase "batch processes multiple inputs sequentially" $ do
            let model = TestChatModel "Pong" "mock-gpt"
                inputs = [[userMessage "Ping 1"], [userMessage "Ping 2"]]
            res <- runExceptT $ batch model inputs Nothing
            case res of
              Left err -> assertFailure $ "Unexpected error: " ++ show err
              Right msgs -> do
                length msgs @?= 2
                map extractMessageText msgs @?= ["Pong", "Pong"]
        ]
    ]