PromptTemplate.hs

{-# 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.Error (LangchainResult, validationError)
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
  { templateString :: Text
  }
  deriving (Show, 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 -> LangchainResult Text
renderPrompt (PromptTemplate template) vars = interpolate vars 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
  { fsPrefix :: Text
  -- ^ Text before the examples
  , fsExamples :: [HM.Map Text Text]
  -- ^ List of example variable maps
  , fsExampleTemplate :: Text
  -- ^ Template for formatting each example
  , fsExampleSeparator :: Text
  -- ^ Separator between formatted examples
  , fsSuffix :: Text
  -- ^ Text after the examples, with placeholders
  }
  deriving (Show, 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 -> LangchainResult Text
renderFewShotPrompt FewShotPromptTemplate {..} = do
  -- Format each example using the example template
  formattedExamples <-
    mapM
      (`interpolate` fsExampleTemplate)
      fsExamples
  -- Join the formatted examples with the separator
  let examplesText = T.intercalate fsExampleSeparator formattedExamples
  -- Combine prefix, examples, and suffix
  return $ fsPrefix <> examplesText <> 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 -> LangchainResult Text
interpolate vars = go
  where
    go :: Text -> LangchainResult Text
    go t =
      case T.breakOn "{" t of
        (before, after) | T.null after -> Right before
        (before, after') ->
          case T.breakOn "}" (T.drop 1 after') of
            (_, after'') | T.null after'' -> Left $ validationError "Unclosed brace" Nothing Nothing
            (key, after''') ->
              let key' = T.strip key
               in case HM.lookup key' vars of
                    Just val -> do
                      rest <- go (T.drop 1 after''')
                      return $ before <> val <> rest
                    Nothing -> Left $ validationError ("Missing variable: " <> key') (Just key') Nothing

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

  invoke template variables = pure $ renderPrompt template variables

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

  invoke t m = pure $ renderFewShotPrompt t m
-}