{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE RecordWildCards #-}
{-# LANGUAGE TypeFamilies #-}

{- |
Module:      Langchain.PromptTemplate
Copyright:   (c) 2025 Tushar Adhatrao
License:     MIT
Maintainer:  Tushar Adhatrao <tusharadhatrao@gmail.com>
Stability:   experimental

This module provides types and functions for working with prompt templates in Langchain.
Prompt templates are used to structure inputs for language models, allowing for dynamic
insertion of variables into predefined text formats. They are essential for creating
flexible and reusable prompts that can be customized based on input data.

The main types are:

* 'PromptTemplate': A simple template with placeholders for variables.
* 'FewShotPromptTemplate': A template that includes few-shot examples for better context,
  useful in scenarios like few-shot learning.

These types are designed to be compatible with the Langchain Python library's prompt template
functionality: [Langchain PromptTemplate](https://python.langchain.com/docs/concepts/prompt_templates/).

== Examples

See the documentation for 'renderPrompt' and 'renderFewShotPrompt' for usage examples.
-}
module Langchain.PromptTemplate
  ( -- * Core Types
    PromptTemplate (..)
  , FewShotPromptTemplate (..)

    -- * Rendering Functions
  , renderPrompt
  , renderFewShotPrompt
  ) where

import qualified Data.Map.Strict as HM
import Data.Text (Text)
import qualified Data.Text as T
import Langchain.Runnable.Core (Runnable (..))

-- TODO: Add Mechanism for custom example selector

{- | Represents a prompt template with a template string.
The template string can contain placeholders of the form {key},
where key is a sequence of alphanumeric characters and underscores.
-}
newtype PromptTemplate = PromptTemplate
  { PromptTemplate -> Text
templateString :: Text
  }
  deriving (Int -> PromptTemplate -> ShowS
[PromptTemplate] -> ShowS
PromptTemplate -> String
(Int -> PromptTemplate -> ShowS)
-> (PromptTemplate -> String)
-> ([PromptTemplate] -> ShowS)
-> Show PromptTemplate
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> PromptTemplate -> ShowS
showsPrec :: Int -> PromptTemplate -> ShowS
$cshow :: PromptTemplate -> String
show :: PromptTemplate -> String
$cshowList :: [PromptTemplate] -> ShowS
showList :: [PromptTemplate] -> ShowS
Show, PromptTemplate -> PromptTemplate -> Bool
(PromptTemplate -> PromptTemplate -> Bool)
-> (PromptTemplate -> PromptTemplate -> Bool) -> Eq PromptTemplate
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: PromptTemplate -> PromptTemplate -> Bool
== :: PromptTemplate -> PromptTemplate -> Bool
$c/= :: PromptTemplate -> PromptTemplate -> Bool
/= :: PromptTemplate -> PromptTemplate -> Bool
Eq)

{- | Render a prompt template with the given variables.
Returns either an error message if a variable is missing or the rendered template.

=== Using 'renderPrompt'

To render a prompt template with variables:

@
let template = PromptTemplate "Hello, {name}! Welcome to {place}."
vars = HM.fromList [("name", "Alice"), ("place", "Wonderland")]
result <- renderPrompt template vars
-- Result: Right "Hello, Alice! Welcome to Wonderland."
@

If a variable is missing:

@
let vars = HM.fromList [("name", "Alice")]
result <- renderPrompt template vars
-- Result: Left "Missing variable: place"
@
-}
renderPrompt :: PromptTemplate -> HM.Map Text Text -> Either String Text
renderPrompt :: PromptTemplate -> Map Text Text -> Either String Text
renderPrompt (PromptTemplate Text
template) Map Text Text
vars = Map Text Text -> Text -> Either String Text
interpolate Map Text Text
vars Text
template

{- | Represents a few-shot prompt template with examples.
This type allows for creating prompts that include example inputs and outputs,
which can be useful for few-shot learning scenarios.
-}
data FewShotPromptTemplate = FewShotPromptTemplate
  { FewShotPromptTemplate -> Text
fsPrefix :: Text
  -- ^ Text before the examples
  , FewShotPromptTemplate -> [Map Text Text]
fsExamples :: [HM.Map Text Text]
  -- ^ List of example variable maps
  , FewShotPromptTemplate -> Text
fsExampleTemplate :: Text
  -- ^ Template for formatting each example
  , FewShotPromptTemplate -> Text
fsExampleSeparator :: Text
  -- ^ Separator between formatted examples
  , FewShotPromptTemplate -> Text
fsSuffix :: Text
  -- ^ Text after the examples, with placeholders
  }
  deriving (Int -> FewShotPromptTemplate -> ShowS
[FewShotPromptTemplate] -> ShowS
FewShotPromptTemplate -> String
(Int -> FewShotPromptTemplate -> ShowS)
-> (FewShotPromptTemplate -> String)
-> ([FewShotPromptTemplate] -> ShowS)
-> Show FewShotPromptTemplate
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> FewShotPromptTemplate -> ShowS
showsPrec :: Int -> FewShotPromptTemplate -> ShowS
$cshow :: FewShotPromptTemplate -> String
show :: FewShotPromptTemplate -> String
$cshowList :: [FewShotPromptTemplate] -> ShowS
showList :: [FewShotPromptTemplate] -> ShowS
Show, FewShotPromptTemplate -> FewShotPromptTemplate -> Bool
(FewShotPromptTemplate -> FewShotPromptTemplate -> Bool)
-> (FewShotPromptTemplate -> FewShotPromptTemplate -> Bool)
-> Eq FewShotPromptTemplate
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: FewShotPromptTemplate -> FewShotPromptTemplate -> Bool
== :: FewShotPromptTemplate -> FewShotPromptTemplate -> Bool
$c/= :: FewShotPromptTemplate -> FewShotPromptTemplate -> Bool
/= :: FewShotPromptTemplate -> FewShotPromptTemplate -> Bool
Eq)

{- | Render a few-shot prompt template with the given input variables.
Returns either an error message if interpolation fails or the fully rendered prompt.

=== Using 'renderFewShotPrompt'

To render a few-shot prompt template:

@
let 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}\n"
      , fsExampleSeparator = "\n"
      , fsSuffix = "Now translate: {query}"
      }
result <- renderFewShotPrompt fewShotTemplate
-- Result: Right "Examples of {type}:\nInput: Hello\nOutput: Bonjour\n\nInput: Goodbye\nOutput: Au revoir\nNow translate: {query}"
@
-}
renderFewShotPrompt :: FewShotPromptTemplate -> Either String Text
renderFewShotPrompt :: FewShotPromptTemplate -> Either String Text
renderFewShotPrompt FewShotPromptTemplate {[Map Text Text]
Text
fsPrefix :: FewShotPromptTemplate -> Text
fsExamples :: FewShotPromptTemplate -> [Map Text Text]
fsExampleTemplate :: FewShotPromptTemplate -> Text
fsExampleSeparator :: FewShotPromptTemplate -> Text
fsSuffix :: FewShotPromptTemplate -> Text
fsPrefix :: Text
fsExamples :: [Map Text Text]
fsExampleTemplate :: Text
fsExampleSeparator :: Text
fsSuffix :: Text
..} = do
  -- Format each example using the example template
  [Text]
formattedExamples <-
    (Map Text Text -> Either String Text)
-> [Map Text Text] -> Either String [Text]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM
      (\Map Text Text
ex -> Map Text Text -> Text -> Either String Text
interpolate Map Text Text
ex Text
fsExampleTemplate)
      [Map Text Text]
fsExamples
  -- Join the formatted examples with the separator
  let examplesText :: Text
examplesText = Text -> [Text] -> Text
T.intercalate Text
fsExampleSeparator [Text]
formattedExamples
  -- Combine prefix, examples, and suffix
  Text -> Either String Text
forall a. a -> Either String a
forall (m :: * -> *) a. Monad m => a -> m a
return (Text -> Either String Text) -> Text -> Either String Text
forall a b. (a -> b) -> a -> b
$ Text
fsPrefix Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
examplesText Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
fsSuffix

{- | Interpolate variables into a template string.
Placeholders are of the form {key}, where key is a sequence of alphanumeric characters and underscores.
-}
interpolate :: HM.Map Text Text -> Text -> Either String Text
interpolate :: Map Text Text -> Text -> Either String Text
interpolate Map Text Text
vars Text
template = Text -> Either String Text
go Text
template
  where
    go :: Text -> Either String Text
    go :: Text -> Either String Text
go Text
t =
      case HasCallStack => Text -> Text -> (Text, Text)
Text -> Text -> (Text, Text)
T.breakOn Text
"{" Text
t of
        (Text
before, Text
after) | Text -> Bool
T.null Text
after -> Text -> Either String Text
forall a b. b -> Either a b
Right Text
before
        (Text
before, Text
after') ->
          case HasCallStack => Text -> Text -> (Text, Text)
Text -> Text -> (Text, Text)
T.breakOn Text
"}" (Int -> Text -> Text
T.drop Int
1 Text
after') of
            (Text
_, Text
after'') | Text -> Bool
T.null Text
after'' -> String -> Either String Text
forall a b. a -> Either a b
Left String
"Unclosed brace"
            (Text
key, Text
after''') ->
              let key' :: Text
key' = Text -> Text
T.strip Text
key
               in case Text -> Map Text Text -> Maybe Text
forall k a. Ord k => k -> Map k a -> Maybe a
HM.lookup Text
key' Map Text Text
vars of
                    Just Text
val -> do
                      Text
rest <- Text -> Either String Text
go (Int -> Text -> Text
T.drop Int
1 Text
after''')
                      Text -> Either String Text
forall a. a -> Either String a
forall (m :: * -> *) a. Monad m => a -> m a
return (Text -> Either String Text) -> Text -> Either String Text
forall a b. (a -> b) -> a -> b
$ Text
before Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
val Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
rest
                    Maybe Text
Nothing -> String -> Either String Text
forall a b. a -> Either a b
Left (String -> Either String Text) -> String -> Either String Text
forall a b. (a -> b) -> a -> b
$ String
"Missing variable: " String -> ShowS
forall a. Semigroup a => a -> a -> a
<> Text -> String
T.unpack Text
key'

instance Runnable PromptTemplate where
  type RunnableInput PromptTemplate = HM.Map Text Text
  type RunnableOutput PromptTemplate = Text

  invoke :: PromptTemplate
-> RunnableInput PromptTemplate
-> IO (Either String (RunnableOutput PromptTemplate))
invoke PromptTemplate
template RunnableInput PromptTemplate
variables = Either String Text -> IO (Either String Text)
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Either String Text -> IO (Either String Text))
-> Either String Text -> IO (Either String Text)
forall a b. (a -> b) -> a -> b
$ PromptTemplate -> Map Text Text -> Either String Text
renderPrompt PromptTemplate
template Map Text Text
RunnableInput PromptTemplate
variables

{-
instance Runnable FewShotPromptTemplate where
  type RunnableInput FewShotPromptTemplate = Maybe [Text]
  type RunnableOutput FewShotPromptTemplate = Text

  invoke t m = pure $ renderFewShotPrompt t m
-}