content: {-# LANGUAGE OverloadedStrings #-}
  
  module Test.Langchain.PromptTemplate (tests) where
  
  import qualified Data.Map.Strict as HM
  import qualified Data.Text as T
  import Langchain.PromptTemplate
  import Langchain.Runnable.Core (invoke)
  import Test.Tasty
  import Test.Tasty.HUnit
  
  tests :: TestTree
  tests =
    testGroup
      "PromptTemplate Tests"
      [ testGroup
          "PromptTemplate"
          [ testCase "correctly interpolates all variables" $
              renderPrompt template vars @?= Right "Hello, Alice! Welcome to Wonderland."
          , testCase "handles templates with no variables" $
              let noVarTemplate = PromptTemplate "Hello, world!"
               in renderPrompt noVarTemplate HM.empty @?= Right "Hello, world!"
          , testCase "handles templates with repeated variables" $
              let repeatTemplate = PromptTemplate "{name} likes {food}. {name} eats {food} every day."
                  repeatVars = HM.fromList [("name", "Bob"), ("food", "pizza")]
               in renderPrompt repeatTemplate repeatVars @?= Right "Bob likes pizza. Bob eats pizza every day."
          , testCase "returns an error for missing variables" $
              let missingVars = HM.fromList [("name", "Charlie")]
               in case renderPrompt template missingVars of
                    Left err -> "place" `T.isInfixOf` T.pack (show err) @? "Expected error to contain 'place'"
                    Right _ -> assertFailure "Expected an error for missing variable"
                    {- TODO: Need to take care of incomplete brace cases
                    , testCase "handles unclosed braces" $
                        let invalidTemplate = PromptTemplate "Hello, {name! Welcome to {place}."
                         in case renderPrompt invalidTemplate vars of
                              Left err -> err @?= "Unclosed brace"
                              Right _ -> assertFailure "Expected an error for unclosed brace"
                    , testCase "handles complex nesting of placeholders" $
                        let complexTemplate = PromptTemplate "{{name}} is not a placeholder but {name} is."
                         in renderPrompt complexTemplate vars @?= Right "{Alice} is not a placeholder but Alice is."
                         -}
          ]
      , testCase "Runnable instance for PromptTemplate - invoke with variables" $ do
          let template1 = PromptTemplate "Hello, {name}!"
              vars1 = HM.fromList [("name", "Dave")]
          result <- invoke template1 vars1
          result @?= Right "Hello, Dave!"
      , testGroup
          "FewShotPromptTemplate"
          [ testCase "correctly formats a few-shot prompt" $
              let expected =
                    "Examples of {type}:\nInput: Hello\nOutput: Bonjour\n\nInput: Goodbye\nOutput: Au revoir\nNow translate: {query}"
               in renderFewShotPrompt fewShotTemplate @?= Right expected
          , testCase "handles empty examples list" $
              let emptyExamples = fewShotTemplate {fsExamples = []}
               in renderFewShotPrompt emptyExamples @?= Right "Examples of {type}:\n\nNow translate: {query}"
          , testCase "handles empty prefix and suffix" $
              let noPreSuf = fewShotTemplate {fsPrefix = "", fsSuffix = ""}
               in renderFewShotPrompt noPreSuf
                    @?= Right "Input: Hello\nOutput: Bonjour\n\nInput: Goodbye\nOutput: Au revoir"
          , testCase "returns an error when example variables are missing" $
              let badExamples =
                    fewShotTemplate
                      { fsExamples = [HM.fromList [("wrong", "value")]]
                      , fsExampleTemplate = "{input} translates to {output}"
                      }
               in case renderFewShotPrompt badExamples of
                    Left err ->
                      "input" `T.isInfixOf` T.pack (show err)
                        @? "Expected error to contain 'input'"
                    Right _ ->
                      assertFailure
                        "Expected an error for missing example variable"
          , testCase "correctly uses the example separator" $
              let customSep = fewShotTemplate {fsExampleSeparator = " ### "}
               in renderFewShotPrompt customSep
                    @?= Right
                      "Examples of {type}:\nInput: Hello\nOutput: Bonjour ### Input: Goodbye\nOutput: Au revoir\nNow translate: {query}"
          ]
      ]
    where
      template = PromptTemplate "Hello, {name}! Welcome to {place}."
      vars = HM.fromList [("name", "Alice"), ("place", "Wonderland")]
      fewShotTemplate =
        FewShotPromptTemplate
          { fsPrefix = "Examples of {type}:\n"
          , fsExamples =
              [ HM.fromList [("input", "Hello"), ("output", "Bonjour")]
              , HM.fromList [("input", "Goodbye"), ("output", "Au revoir")]
              ]
          , fsExampleTemplate = "Input: {input}\nOutput: {output}"
          , fsExampleSeparator = "\n\n"
          , fsSuffix = "\nNow translate: {query}"
          }
  