content: {-# LANGUAGE OverloadedStrings #-}
  {-# LANGUAGE TypeFamilies #-}
  
  {- |
  Module      : Langchain.Retriever.MultiQueryRetriever
  Description : Multi-query retrieval implementation for LangChain Haskell
  Copyright   : (c) 2025 Tushar Adhatrao
  License     : MIT
  Maintainer  : Tushar Adhatrao <tusharadhatrao@gmail.com>
  Stability   : experimental
  
  Advanced retriever implementation that generates multiple queries from a single
  input to improve document retrieval. Integrates with LLMs for query expansion
  and vector stores for document retrieval
  
  Example usage:
  
  @
  -- Create components
  ollamaLLM = Ollama "llama3" []
  vs = VectorStoreRetriever (createVectorStore ...)
  
  -- Create retriever with default config
  mqRetriever = newMultiQueryRetriever vs ollamaLLM
  
  -- Retrieve documents
  docs <- _get_relevant_documents mqRetriever "Haskell features"
  -- Returns combined results from multiple generated queries
  @
  -}
  module Langchain.Retriever.MultiQueryRetriever
    ( MultiQueryRetriever (..)
    , QueryGenerationPrompt (..)
    , newMultiQueryRetriever
    , defaultQueryGenerationPrompt
    , newMultiQueryRetrieverWithConfig
    , defaultMultiQueryRetrieverConfig
    , generateQueries
    ) where
  
  import Langchain.DocumentLoader.Core (Document)
  import Langchain.LLM.Core (LLM (..))
  import Langchain.OutputParser.Core (NumberSeparatedList (..), OutputParser (..))
  import Langchain.PromptTemplate (PromptTemplate (..), renderPrompt)
  import Langchain.Retriever.Core (Retriever (..))
  import qualified Langchain.Runnable.Core as Run
  
  import Data.Either (rights)
  import Data.List (nub)
  import qualified Data.Map.Strict as HM
  import Data.Text (Text)
  import qualified Data.Text as T
  import Langchain.Error (LangchainError, llmError)
  
  {- | Query generation prompt template
  Controls how the LLM generates multiple query variants from the original query.
  
  Example prompt structure:
  
  @
  "You are an AI assistant... Original query: {query}... Generate {num_queries} versions..."
  @
  -}
  newtype QueryGenerationPrompt = QueryGenerationPrompt PromptTemplate
    deriving (Show, Eq)
  
  {- | Default query generation prompt
  Generates 3 query variants in numbered list format. Includes instructions for
  query diversity and formatting.
  -}
  defaultQueryGenerationPrompt :: QueryGenerationPrompt
  defaultQueryGenerationPrompt =
    QueryGenerationPrompt $
      PromptTemplate
        { templateString =
            T.unlines
              [ "You are an AI language model assistant that helps users by generating multiple search queries based on their initial query."
              , "These queries should help retrieve relevant documents or information from a vector database."
              , ""
              , "Original query: {query}"
              , ""
              , "Please generate {num_queries} different versions of this query that will help the user find the most relevant information."
              , "The queries should be different but related to the original query."
              , "Return these queries in the following format: 1. query 1 \n 2. query 2 \n 3. query 3"
              , "Only return queries and nothing else"
              ]
        }
  
  -- | Configuration for multi-query retrieval
  data MultiQueryRetrieverConfig = MultiQueryRetrieverConfig
    { numQueries :: Int
    -- ^ Number of queries to generate
    , queryGenerationPrompt :: QueryGenerationPrompt
    -- ^ Prompt template for query generation
    , includeMergeDocs :: Bool
    -- ^ Whether to include merged documents
    , includeOriginalQuery :: Bool
    -- ^ Whether to include results from original query
    }
  
  {- | Default configuration
  - 3 generated queries
  - Includes original query results
  - Uses default query generation prompt
  -}
  defaultMultiQueryRetrieverConfig :: MultiQueryRetrieverConfig
  defaultMultiQueryRetrieverConfig =
    MultiQueryRetrieverConfig
      { numQueries = 3
      , queryGenerationPrompt = defaultQueryGenerationPrompt
      , includeMergeDocs = True
      , includeOriginalQuery = True
      }
  
  {- | Multi-query retriever implementation
  Generates multiple queries using an LLM, retrieves documents for each query,
  and combines results. Improves recall by exploring different query formulations.
  
  Example instance:
  
  @
  mqRetriever = MultiQueryRetriever
    { retriever = vectorStoreRetriever
    , llm = ollamaLLM
    , config = defaultMultiQueryRetrieverConfig
    }
  @
  -}
  data (Retriever a, LLM m) => MultiQueryRetriever a m = MultiQueryRetriever
    { retriever :: a
    -- ^ The base retriever
    , llm :: m
    -- ^ The language model for generating queries
    , config :: MultiQueryRetrieverConfig
    -- ^ Configuration
    }
  
  {- | Create retriever with default settings
  Example:
  
  >>> newMultiQueryRetriever vsRetriever ollamaLLM
  MultiQueryRetriever {numQueries = 3, ...}
  -}
  newMultiQueryRetriever :: (Retriever a, LLM m) => a -> m -> MultiQueryRetriever a m
  newMultiQueryRetriever r l =
    MultiQueryRetriever
      { retriever = r
      , llm = l
      , config = defaultMultiQueryRetrieverConfig
      }
  
  {- | Create retriever with custom configuration
  Example:
  
  >>> let customCfg = defaultMultiQueryRetrieverConfig { numQueries = 5 }
  >>> newMultiQueryRetrieverWithConfig vsRetriever ollamaLLM customCfg
  MultiQueryRetriever {numQueries = 5, ...}
  -}
  newMultiQueryRetrieverWithConfig ::
    (Retriever a, LLM m) =>
    a ->
    m ->
    MultiQueryRetrieverConfig ->
    MultiQueryRetriever a m
  newMultiQueryRetrieverWithConfig r l c =
    MultiQueryRetriever
      { retriever = r
      , llm = l
      , config = c
      }
  
  {- | Generate multiple query variants using LLM
  Example:
  
  >>> generateQueries ollamaLLM prompt "Haskell" 3 True
  Right ["Haskell", "Haskell features", "Haskell applications"]
  -}
  generateQueries ::
    LLM m => m -> QueryGenerationPrompt -> Text -> Int -> Bool -> IO (Either LangchainError [Text])
  generateQueries model (QueryGenerationPrompt promptTemplate) query n includeOriginal = do
    let vars = HM.fromList [("query", query), ("num_queries", T.pack $ show n)]
    case renderPrompt promptTemplate vars of
      Left err -> return $ Left err
      Right prompt -> do
        result <- generate model prompt Nothing
        case result of
          Left err -> return $ Left err
          Right response -> do
            case parse response :: Either LangchainError NumberSeparatedList of
              Left err -> return $ Left err
              Right (NumberSeparatedList queries) -> do
                let uniqueQueries = nub $ filter (not . T.null) queries
                return $
                  Right $
                    if includeOriginal
                      then query : uniqueQueries
                      else uniqueQueries
  
  {- | Combine documents from multiple queries
  Removes duplicates while maintaining order (simplified approach).
  -}
  combineDocuments :: [[Document]] -> [Document]
  combineDocuments docLists =
    -- This is a simplified approach. In a production system, you'd want a more
    -- sophisticated way to identify and rank duplicate documents
    nub $ concat docLists
  
  {- | Retriever instance implementation
  1. Generates multiple queries using LLM
  2. Retrieves documents for each query
  3. Combines and deduplicates results
  
  Example retrieval:
  
  >>> _get_relevant_documents mqRetriever "Haskell"
  Right [Document "Haskell is...", Document "Functional programming...", ...]
  -}
  instance (Retriever a, LLM m) => Retriever (MultiQueryRetriever a m) where
    _get_relevant_documents r query = do
      let baseRetriever = retriever r
          model = llm r
          cfg = config r
  
      -- Generate multiple queries
      queriesResult <-
        generateQueries
          model
          (queryGenerationPrompt cfg)
          query
          (numQueries cfg)
          (includeOriginalQuery cfg)
  
      case queriesResult of
        Left err -> return $ Left err
        Right queries -> do
          -- Get documents for each query
          results <- mapM (_get_relevant_documents baseRetriever) queries
  
          -- Filter successful results
          let validResults = rights results
  
          if null validResults
            then return $ Left (llmError "No valid results from any query" Nothing Nothing)
            else return $ Right $ combineDocuments validResults
  
  {-
   ghci> :set -XOverloadedStrings
   ghci> let ollamaEmbed = OllamaEmbeddings "nomic-embed-text:latest" Nothing Nothing
   ghci> let vs = emptyInMemoryVectorStore ollamaEmbed
   ghci> import Data.Map (empty)
   ghci> import Data.Either
   ghci> newVs <- addDocuments vs [Document "Tushar is 25 years old." empty]
   ghci> let newVs_ = fromRight vs newVs
   ghci> let vRet = VectorStoreRetriever newVs_
   ghci> let ollamLLM = Ollama "llama3.2" []
   ghci> let mqRet = newMultiQueryRetriever vRet ollamLLM
   ghci> documents <- _get_relevant_documents mqRet "How old is Tushar?"
   ghci> documents
      Right [Document {pageContent = "Tushar is 25 years old.", metadata = fromList []}]
   -}
  
  {- | Runnable interface implementation
  Allows integration with LangChain workflows:
  
  >>> invoke mqRetriever "AI applications"
  Right [Document "Machine learning...", ...]
  -}
  instance (Retriever a, LLM m) => Run.Runnable (MultiQueryRetriever a m) where
    type RunnableInput (MultiQueryRetriever a m) = Text
    type RunnableOutput (MultiQueryRetriever a m) = [Document]
  
    invoke = _get_relevant_documents
  
  {- $examples
  Test case patterns:
  1. Query generation
     >>> generateQueries ollamaLLM prompt "Test" 2 False
     Right ["Test case", "Test example"]
  
  2. Full retrieval flow
     >>> _get_relevant_documents mqRetriever "Haskell"
     Right [Document "Functional...", Document "Type system..."]
  
  3. Configuration variants
     >>> let cfg = defaultMultiQueryRetrieverConfig { numQueries = 5 }
     >>> newMultiQueryRetrieverWithConfig vsRetriever ollamaLLM cfg
     MultiQueryRetriever {numQueries = 5, ...}
  -}
  