{-# LANGUAGE DuplicateRecordFields #-}
{-# LANGUAGE OverloadedStrings #-}

module Data.Ollama.Embeddings
  ( -- * Embedding API
    embedding
  , embeddingOps
  , EmbeddingOps (..)
  , EmbeddingResp (..)
  ) where

import Data.Aeson
import Data.Ollama.Common.Utils as CU
import Data.Text (Text)
import Data.Text qualified as T
import Network.HTTP.Client
import Control.Exception (try)
import Data.ByteString.Lazy.Char8 (ByteString)

-- TODO: Add Options parameter
data EmbeddingOps = EmbeddingOps
  { EmbeddingOps -> Text
model :: Text
  , EmbeddingOps -> Text
input :: Text
  , EmbeddingOps -> Maybe Bool
truncate :: Maybe Bool
  , EmbeddingOps -> Maybe Text
keepAlive :: Maybe Text
  }
  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)

data EmbeddingResp = EmbeddingResp
  { EmbeddingResp -> Text
model :: Text
  , EmbeddingResp -> [[Float]]
embedding_ :: [[Float]]
  }
  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 Text
keepAlive_) =
    [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 Text -> Pair
forall v. ToJSON v => Key -> v -> Pair
forall e kv v. (KeyValue e kv, ToJSON v) => Key -> v -> kv
.= Maybe Text
keepAlive_
      ]

-- TODO: Add Options parameter

-- | Embedding API
embeddingOps ::
  -- | Model
  Text ->
  -- | Input
  Text ->
  -- | Truncate
  Maybe Bool ->
  -- | Keep Alive
  Maybe Text ->
  IO (Either String EmbeddingResp)
embeddingOps :: Text
-> Text
-> Maybe Bool
-> Maybe Text
-> IO (Either String EmbeddingResp)
embeddingOps Text
modelName Text
input_ Maybe Bool
mTruncate Maybe Text
mKeepAlive = do
  let url :: Text
url = Text
defaultOllamaUrl
  Manager
manager <- ManagerSettings -> IO Manager
newManager ManagerSettings
defaultManagerSettings
  --einitialRequest <- parseRequest $ T.unpack (url <> "/api/embed")
  Either HttpException Request
eInitialRequest <-
    IO Request -> IO (Either HttpException Request)
forall e a. Exception e => IO a -> IO (Either e a)
try (IO Request -> IO (Either HttpException Request))
-> IO Request -> IO (Either HttpException Request)
forall a b. (a -> b) -> a -> b
$ String -> IO Request
forall (m :: * -> *). MonadThrow m => String -> m Request
parseRequest (String -> IO Request) -> String -> IO Request
forall a b. (a -> b) -> a -> b
$ Text -> String
T.unpack (Text
url Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
"/api/embed") :: IO (Either HttpException Request)
  case Either HttpException Request
eInitialRequest of
    Left HttpException
e -> do
      Either String EmbeddingResp -> IO (Either String EmbeddingResp)
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return (Either String EmbeddingResp -> IO (Either String EmbeddingResp))
-> Either String EmbeddingResp -> IO (Either String EmbeddingResp)
forall a b. (a -> b) -> a -> b
$ String -> Either String EmbeddingResp
forall a b. a -> Either a b
Left (String -> Either String EmbeddingResp)
-> String -> Either String EmbeddingResp
forall a b. (a -> b) -> a -> b
$ HttpException -> String
forall a. Show a => a -> String
show HttpException
e
    Right Request
initialRequest -> do
      let reqBody :: EmbeddingOps
reqBody =
            EmbeddingOps
              { $sel:model:EmbeddingOps :: Text
model = Text
modelName
              , $sel:input:EmbeddingOps :: Text
input = Text
input_
              , $sel:truncate:EmbeddingOps :: Maybe Bool
truncate = Maybe Bool
mTruncate
              , $sel:keepAlive:EmbeddingOps :: Maybe Text
keepAlive = Maybe Text
mKeepAlive
              }
          request :: Request
request =
            Request
initialRequest
              { method = "POST"
              , requestBody = RequestBodyLBS $ encode reqBody
              }
      Either HttpException (Response ByteString)
eResp <- IO (Response ByteString)
-> IO (Either HttpException (Response ByteString))
forall e a. Exception e => IO a -> IO (Either e a)
try (IO (Response ByteString)
 -> IO (Either HttpException (Response ByteString)))
-> IO (Response ByteString)
-> IO (Either HttpException (Response ByteString))
forall a b. (a -> b) -> a -> b
$ Request -> Manager -> IO (Response ByteString)
httpLbs Request
request Manager
manager :: IO (Either HttpException (Response ByteString))
      case Either HttpException (Response ByteString)
eResp of
        Left HttpException
err -> Either String EmbeddingResp -> IO (Either String EmbeddingResp)
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return (Either String EmbeddingResp -> IO (Either String EmbeddingResp))
-> Either String EmbeddingResp -> IO (Either String EmbeddingResp)
forall a b. (a -> b) -> a -> b
$ String -> Either String EmbeddingResp
forall a b. a -> Either a b
Left (HttpException -> String
forall a. Show a => a -> String
show HttpException
err)
        Right Response ByteString
resp -> 
          case ByteString -> Maybe EmbeddingResp
forall a. FromJSON a => ByteString -> Maybe a
decode (Response ByteString -> ByteString
forall body. Response body -> body
responseBody Response ByteString
resp) of
            Maybe EmbeddingResp
Nothing -> Either String EmbeddingResp -> IO (Either String EmbeddingResp)
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return (Either String EmbeddingResp -> IO (Either String EmbeddingResp))
-> Either String EmbeddingResp -> IO (Either String EmbeddingResp)
forall a b. (a -> b) -> a -> b
$ String -> Either String EmbeddingResp
forall a b. a -> Either a b
Left (String -> Either String EmbeddingResp)
-> String -> Either String EmbeddingResp
forall a b. (a -> b) -> a -> b
$ String
"Couldn't decode response: " String -> ShowS
forall a. Semigroup a => a -> a -> a
<> ByteString -> String
forall a. Show a => a -> String
show (Response ByteString -> ByteString
forall body. Response body -> body
responseBody Response ByteString
resp)
            Just EmbeddingResp
r -> Either String EmbeddingResp -> IO (Either String EmbeddingResp)
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return (Either String EmbeddingResp -> IO (Either String EmbeddingResp))
-> Either String EmbeddingResp -> IO (Either String EmbeddingResp)
forall a b. (a -> b) -> a -> b
$ EmbeddingResp -> Either String EmbeddingResp
forall a b. b -> Either a b
Right EmbeddingResp
r

-- Higher level binding that only takes important params

-- | Embedding API
embedding ::
  -- | Model
  Text ->
  -- | Input
  Text ->
  IO (Either String EmbeddingResp)
embedding :: Text -> Text -> IO (Either String EmbeddingResp)
embedding Text
modelName Text
input_ =
  Text
-> Text
-> Maybe Bool
-> Maybe Text
-> IO (Either String EmbeddingResp)
embeddingOps Text
modelName Text
input_ Maybe Bool
forall a. Maybe a
Nothing Maybe Text
forall a. Maybe a
Nothing