{-# LANGUAGE OverloadedStrings #-}

{- |
Module      : Data.Ollama.Embeddings
Copyright   : (c) 2025 Tushar Adhatrao
License     : MIT
Maintainer  : Tushar Adhatrao <tusharadhatrao@gmail.com>
Stability   : experimental
Description : Functionality for generating text embeddings using the Ollama API.

This module provides functions to generate text embeddings from an Ollama model. It includes both
high-level ('embedding', 'embeddingM') and low-level ('embeddingOps', 'embeddingOpsM') APIs for
generating embeddings, with support for customizing model options, truncation, and keep-alive settings.
The embeddings are returned as a list of float vectors, suitable for tasks like semantic search or
text similarity analysis.

The 'EmbeddingOps' type configures the embedding request, and 'EmbeddingResp' represents the response
containing the model name and the generated embeddings. The 'defaultEmbeddingOps' provides a default
configuration for convenience.

Example:

>>> embedding "llama3.2" ["Hello, world!"]
Right (EmbeddingResp "llama3.2" [[0.1, 0.2, ...]])
-}
module Data.Ollama.Embeddings
  ( -- * Embedding API
    embedding
  , embeddingOps
  , embeddingM
  , embeddingOpsM

    -- * Configuration and Response Types
  , defaultEmbeddingOps
  , EmbeddingOps (..)
  , EmbeddingResp (..)

    -- * Model Options
  , ModelOptions (..)
  , defaultModelOptions
  ) where

import Control.Monad.IO.Class (MonadIO (liftIO))
import Data.Aeson
import Data.Ollama.Common.Config (OllamaConfig)
import Data.Ollama.Common.Error (OllamaError)
import Data.Ollama.Common.Types (ModelOptions (..))
import Data.Ollama.Common.Utils as CU
import Data.Text (Text)

{- | Default configuration for embedding requests.

Provides a default 'EmbeddingOps' with the "llama3.2" model, an empty input list, and no additional options.
Can be customized by modifying fields as needed.
-}
defaultEmbeddingOps :: EmbeddingOps
defaultEmbeddingOps :: EmbeddingOps
defaultEmbeddingOps =
  EmbeddingOps
    { model :: Text
model = Text
"llama3.2"
    , input :: [Text]
input = []
    , truncateInput :: Maybe Bool
truncateInput = Maybe Bool
forall a. Maybe a
Nothing
    , keepAliveEmbed :: Maybe Int
keepAliveEmbed = Maybe Int
forall a. Maybe a
Nothing
    , modelOptions :: Maybe ModelOptions
modelOptions = Maybe ModelOptions
forall a. Maybe a
Nothing
    }

-- | Configuration for an embedding request.
data EmbeddingOps = EmbeddingOps
  { EmbeddingOps -> Text
model :: !Text
  -- ^ The name of the model to use for generating embeddings (e.g., "llama3.2").
  , EmbeddingOps -> [Text]
input :: ![Text]
  -- ^ List of input texts to generate embeddings for.
  , EmbeddingOps -> Maybe Bool
truncateInput :: !(Maybe Bool)
  -- ^ Optional flag to truncate input if it exceeds model limits.
  , EmbeddingOps -> Maybe Int
keepAliveEmbed :: !(Maybe Int)
  -- ^ Optional override for the keep-alive timeout in minutes.
  , EmbeddingOps -> Maybe ModelOptions
modelOptions :: !(Maybe ModelOptions)
  -- ^ Optional model parameters (e.g., temperature) as specified in the Modelfile.
  --
  -- @since 0.2.0.0
  }
  deriving (Int -> EmbeddingOps -> ShowS
[EmbeddingOps] -> ShowS
EmbeddingOps -> String
(Int -> EmbeddingOps -> ShowS)
-> (EmbeddingOps -> String)
-> ([EmbeddingOps] -> ShowS)
-> Show EmbeddingOps
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> EmbeddingOps -> ShowS
showsPrec :: Int -> EmbeddingOps -> ShowS
$cshow :: EmbeddingOps -> String
show :: EmbeddingOps -> String
$cshowList :: [EmbeddingOps] -> ShowS
showList :: [EmbeddingOps] -> ShowS
Show, EmbeddingOps -> EmbeddingOps -> Bool
(EmbeddingOps -> EmbeddingOps -> Bool)
-> (EmbeddingOps -> EmbeddingOps -> Bool) -> Eq EmbeddingOps
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: EmbeddingOps -> EmbeddingOps -> Bool
== :: EmbeddingOps -> EmbeddingOps -> Bool
$c/= :: EmbeddingOps -> EmbeddingOps -> Bool
/= :: EmbeddingOps -> EmbeddingOps -> Bool
Eq)

-- | Response type for an embedding request.
data EmbeddingResp = EmbeddingResp
  { EmbeddingResp -> Text
respondedModel :: !Text
  -- ^ The name of the model that generated the embeddings.
  , EmbeddingResp -> [[Float]]
respondedEmbeddings :: ![[Float]]
  -- ^ List of embedding vectors, one for each input text.
  }
  deriving (Int -> EmbeddingResp -> ShowS
[EmbeddingResp] -> ShowS
EmbeddingResp -> String
(Int -> EmbeddingResp -> ShowS)
-> (EmbeddingResp -> String)
-> ([EmbeddingResp] -> ShowS)
-> Show EmbeddingResp
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> EmbeddingResp -> ShowS
showsPrec :: Int -> EmbeddingResp -> ShowS
$cshow :: EmbeddingResp -> String
show :: EmbeddingResp -> String
$cshowList :: [EmbeddingResp] -> ShowS
showList :: [EmbeddingResp] -> ShowS
Show, EmbeddingResp -> EmbeddingResp -> Bool
(EmbeddingResp -> EmbeddingResp -> Bool)
-> (EmbeddingResp -> EmbeddingResp -> Bool) -> Eq EmbeddingResp
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: EmbeddingResp -> EmbeddingResp -> Bool
== :: EmbeddingResp -> EmbeddingResp -> Bool
$c/= :: EmbeddingResp -> EmbeddingResp -> Bool
/= :: EmbeddingResp -> EmbeddingResp -> Bool
Eq)

instance FromJSON EmbeddingResp where
  parseJSON :: Value -> Parser EmbeddingResp
parseJSON = String
-> (Object -> Parser EmbeddingResp)
-> Value
-> Parser EmbeddingResp
forall a. String -> (Object -> Parser a) -> Value -> Parser a
withObject String
"EmbeddingResp" ((Object -> Parser EmbeddingResp) -> Value -> Parser EmbeddingResp)
-> (Object -> Parser EmbeddingResp)
-> Value
-> Parser EmbeddingResp
forall a b. (a -> b) -> a -> b
$ \Object
v ->
    Text -> [[Float]] -> EmbeddingResp
EmbeddingResp
      (Text -> [[Float]] -> EmbeddingResp)
-> Parser Text -> Parser ([[Float]] -> EmbeddingResp)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Object
v Object -> Key -> Parser Text
forall a. FromJSON a => Object -> Key -> Parser a
.: Key
"model"
      Parser ([[Float]] -> EmbeddingResp)
-> Parser [[Float]] -> Parser EmbeddingResp
forall a b. Parser (a -> b) -> Parser a -> Parser b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Object
v Object -> Key -> Parser [[Float]]
forall a. FromJSON a => Object -> Key -> Parser a
.: Key
"embeddings"

instance ToJSON EmbeddingOps where
  toJSON :: EmbeddingOps -> Value
toJSON (EmbeddingOps Text
model_ [Text]
input_ Maybe Bool
truncate' Maybe Int
keepAlive_ Maybe ModelOptions
ops) =
    [Pair] -> Value
object
      [ Key
"model" Key -> Text -> Pair
forall v. ToJSON v => Key -> v -> Pair
forall e kv v. (KeyValue e kv, ToJSON v) => Key -> v -> kv
.= Text
model_
      , Key
"input" Key -> [Text] -> Pair
forall v. ToJSON v => Key -> v -> Pair
forall e kv v. (KeyValue e kv, ToJSON v) => Key -> v -> kv
.= [Text]
input_
      , Key
"truncate" Key -> Maybe Bool -> Pair
forall v. ToJSON v => Key -> v -> Pair
forall e kv v. (KeyValue e kv, ToJSON v) => Key -> v -> kv
.= Maybe Bool
truncate'
      , Key
"keep_alive" Key -> Maybe Int -> Pair
forall v. ToJSON v => Key -> v -> Pair
forall e kv v. (KeyValue e kv, ToJSON v) => Key -> v -> kv
.= Maybe Int
keepAlive_
      , Key
"options" Key -> Maybe ModelOptions -> Pair
forall v. ToJSON v => Key -> v -> Pair
forall e kv v. (KeyValue e kv, ToJSON v) => Key -> v -> kv
.= Maybe ModelOptions
ops
      ]

{- | Generates embeddings for a list of input texts with full configuration.

Sends a POST request to the "/api//embed" endpoint to generate embeddings for the provided inputs.
Allows customization of truncation, keep-alive settings, model options, and Ollama configuration.
Returns 'Right' with an 'EmbeddingResp' on success or 'Left' with an 'OllamaError' on failure.
-}
embeddingOps ::
  -- | Model name
  Text ->
  -- | List of input texts
  [Text] ->
  -- | Optional truncation flag
  Maybe Bool ->
  -- | Optional keep-alive timeout in minutes
  Maybe Int ->
  -- | Optional model options
  Maybe ModelOptions ->
  -- | Optional 'OllamaConfig' (defaults to 'defaultOllamaConfig' if 'Nothing')
  Maybe OllamaConfig ->
  IO (Either OllamaError EmbeddingResp)
embeddingOps :: Text
-> [Text]
-> Maybe Bool
-> Maybe Int
-> Maybe ModelOptions
-> Maybe OllamaConfig
-> IO (Either OllamaError EmbeddingResp)
embeddingOps Text
modelName [Text]
input_ Maybe Bool
mTruncate Maybe Int
mKeepAlive Maybe ModelOptions
mbOptions Maybe OllamaConfig
mbConfig = do
  Text
-> ByteString
-> Maybe EmbeddingOps
-> Maybe OllamaConfig
-> (Response BodyReader -> IO (Either OllamaError EmbeddingResp))
-> IO (Either OllamaError EmbeddingResp)
forall payload response.
ToJSON payload =>
Text
-> ByteString
-> Maybe payload
-> Maybe OllamaConfig
-> (Response BodyReader -> IO (Either OllamaError response))
-> IO (Either OllamaError response)
withOllamaRequest
    Text
"/api//embed"
    ByteString
"POST"
    ( EmbeddingOps -> Maybe EmbeddingOps
forall a. a -> Maybe a
Just (EmbeddingOps -> Maybe EmbeddingOps)
-> EmbeddingOps -> Maybe EmbeddingOps
forall a b. (a -> b) -> a -> b
$
        EmbeddingOps
          { model :: Text
model = Text
modelName
          , input :: [Text]
input = [Text]
input_
          , truncateInput :: Maybe Bool
truncateInput = Maybe Bool
mTruncate
          , keepAliveEmbed :: Maybe Int
keepAliveEmbed = Maybe Int
mKeepAlive
          , modelOptions :: Maybe ModelOptions
modelOptions = Maybe ModelOptions
mbOptions
          }
    )
    Maybe OllamaConfig
mbConfig
    Response BodyReader -> IO (Either OllamaError EmbeddingResp)
forall a.
FromJSON a =>
Response BodyReader -> IO (Either OllamaError a)
commonNonStreamingHandler

{- | Simplified API for generating embeddings.

A higher-level function that generates embeddings using default settings for truncation, keep-alive,
model options, and Ollama configuration. Suitable for basic use cases.
-}
embedding ::
  -- | Model name
  Text ->
  -- | List of input texts
  [Text] ->
  IO (Either OllamaError EmbeddingResp)
embedding :: Text -> [Text] -> IO (Either OllamaError EmbeddingResp)
embedding Text
modelName [Text]
input_ =
  Text
-> [Text]
-> Maybe Bool
-> Maybe Int
-> Maybe ModelOptions
-> Maybe OllamaConfig
-> IO (Either OllamaError EmbeddingResp)
embeddingOps Text
modelName [Text]
input_ Maybe Bool
forall a. Maybe a
Nothing Maybe Int
forall a. Maybe a
Nothing Maybe ModelOptions
forall a. Maybe a
Nothing Maybe OllamaConfig
forall a. Maybe a
Nothing

{- | MonadIO version of 'embedding' for use in monadic contexts.

Lifts the 'embedding' function into a 'MonadIO' context, allowing it to be used in monadic computations.
-}
embeddingM :: MonadIO m => Text -> [Text] -> m (Either OllamaError EmbeddingResp)
embeddingM :: forall (m :: * -> *).
MonadIO m =>
Text -> [Text] -> m (Either OllamaError EmbeddingResp)
embeddingM Text
m [Text]
ip = IO (Either OllamaError EmbeddingResp)
-> m (Either OllamaError EmbeddingResp)
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO (Either OllamaError EmbeddingResp)
 -> m (Either OllamaError EmbeddingResp))
-> IO (Either OllamaError EmbeddingResp)
-> m (Either OllamaError EmbeddingResp)
forall a b. (a -> b) -> a -> b
$ Text -> [Text] -> IO (Either OllamaError EmbeddingResp)
embedding Text
m [Text]
ip

{- | MonadIO version of 'embeddingOps' for use in monadic contexts.

Lifts the 'embeddingOps' function into a 'MonadIO' context, allowing it to be used in monadic computations
with full configuration options.
-}
embeddingOpsM ::
  MonadIO m =>
  Text ->
  [Text] ->
  Maybe Bool ->
  Maybe Int ->
  Maybe ModelOptions ->
  Maybe OllamaConfig ->
  m (Either OllamaError EmbeddingResp)
embeddingOpsM :: forall (m :: * -> *).
MonadIO m =>
Text
-> [Text]
-> Maybe Bool
-> Maybe Int
-> Maybe ModelOptions
-> Maybe OllamaConfig
-> m (Either OllamaError EmbeddingResp)
embeddingOpsM Text
m [Text]
ip Maybe Bool
mbTruncate Maybe Int
mbKeepAlive Maybe ModelOptions
mbOptions Maybe OllamaConfig
mbCfg =
  IO (Either OllamaError EmbeddingResp)
-> m (Either OllamaError EmbeddingResp)
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO (Either OllamaError EmbeddingResp)
 -> m (Either OllamaError EmbeddingResp))
-> IO (Either OllamaError EmbeddingResp)
-> m (Either OllamaError EmbeddingResp)
forall a b. (a -> b) -> a -> b
$ Text
-> [Text]
-> Maybe Bool
-> Maybe Int
-> Maybe ModelOptions
-> Maybe OllamaConfig
-> IO (Either OllamaError EmbeddingResp)
embeddingOps Text
m [Text]
ip Maybe Bool
mbTruncate Maybe Int
mbKeepAlive Maybe ModelOptions
mbOptions Maybe OllamaConfig
mbCfg