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

module Data.Ollama.Create
  ( -- * Create Model API
    createModel
  , createModelOps
  ) where

import Control.Monad (unless)
import Data.Aeson
import Data.ByteString.Lazy.Char8 qualified as BSL
import Data.Ollama.Common.Utils as CU
import Data.Text (Text)
import Data.Text qualified as T
import Data.Text.IO qualified as T
import Network.HTTP.Client

-- TODO: Add Options parameter
-- TODO: Add Context parameter
data CreateModelOps = CreateModelOps
  { CreateModelOps -> Text
name :: Text
  , CreateModelOps -> Maybe Text
modelFile :: Maybe Text
  , CreateModelOps -> Maybe Bool
stream :: Maybe Bool
  , CreateModelOps -> Maybe FilePath
path :: Maybe FilePath
  }
  deriving (Int -> CreateModelOps -> ShowS
[CreateModelOps] -> ShowS
CreateModelOps -> FilePath
(Int -> CreateModelOps -> ShowS)
-> (CreateModelOps -> FilePath)
-> ([CreateModelOps] -> ShowS)
-> Show CreateModelOps
forall a.
(Int -> a -> ShowS) -> (a -> FilePath) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> CreateModelOps -> ShowS
showsPrec :: Int -> CreateModelOps -> ShowS
$cshow :: CreateModelOps -> FilePath
show :: CreateModelOps -> FilePath
$cshowList :: [CreateModelOps] -> ShowS
showList :: [CreateModelOps] -> ShowS
Show, CreateModelOps -> CreateModelOps -> Bool
(CreateModelOps -> CreateModelOps -> Bool)
-> (CreateModelOps -> CreateModelOps -> Bool) -> Eq CreateModelOps
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: CreateModelOps -> CreateModelOps -> Bool
== :: CreateModelOps -> CreateModelOps -> Bool
$c/= :: CreateModelOps -> CreateModelOps -> Bool
/= :: CreateModelOps -> CreateModelOps -> Bool
Eq)

-- TODO: Add Context Param
newtype CreateModelResp = CreateModelResp {CreateModelResp -> Text
status :: Text}
  deriving (Int -> CreateModelResp -> ShowS
[CreateModelResp] -> ShowS
CreateModelResp -> FilePath
(Int -> CreateModelResp -> ShowS)
-> (CreateModelResp -> FilePath)
-> ([CreateModelResp] -> ShowS)
-> Show CreateModelResp
forall a.
(Int -> a -> ShowS) -> (a -> FilePath) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> CreateModelResp -> ShowS
showsPrec :: Int -> CreateModelResp -> ShowS
$cshow :: CreateModelResp -> FilePath
show :: CreateModelResp -> FilePath
$cshowList :: [CreateModelResp] -> ShowS
showList :: [CreateModelResp] -> ShowS
Show, CreateModelResp -> CreateModelResp -> Bool
(CreateModelResp -> CreateModelResp -> Bool)
-> (CreateModelResp -> CreateModelResp -> Bool)
-> Eq CreateModelResp
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: CreateModelResp -> CreateModelResp -> Bool
== :: CreateModelResp -> CreateModelResp -> Bool
$c/= :: CreateModelResp -> CreateModelResp -> Bool
/= :: CreateModelResp -> CreateModelResp -> Bool
Eq)

instance ToJSON CreateModelOps where
  toJSON :: CreateModelOps -> Value
toJSON
    ( CreateModelOps
        Text
name
        Maybe Text
modelFile
        Maybe Bool
stream
        Maybe FilePath
path
      ) =
      [Pair] -> Value
object
        [ Key
"name" Key -> Text -> Pair
forall v. ToJSON v => Key -> v -> Pair
forall e kv v. (KeyValue e kv, ToJSON v) => Key -> v -> kv
.= Text
name
        , Key
"modelfile" 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
modelFile
        , Key
"stream" 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
stream
        , Key
"path" Key -> Maybe FilePath -> Pair
forall v. ToJSON v => Key -> v -> Pair
forall e kv v. (KeyValue e kv, ToJSON v) => Key -> v -> kv
.= Maybe FilePath
path
        ]

instance FromJSON CreateModelResp where
  parseJSON :: Value -> Parser CreateModelResp
parseJSON = FilePath
-> (Object -> Parser CreateModelResp)
-> Value
-> Parser CreateModelResp
forall a. FilePath -> (Object -> Parser a) -> Value -> Parser a
withObject FilePath
"CreateModelResp" ((Object -> Parser CreateModelResp)
 -> Value -> Parser CreateModelResp)
-> (Object -> Parser CreateModelResp)
-> Value
-> Parser CreateModelResp
forall a b. (a -> b) -> a -> b
$ \Object
v ->
    Text -> CreateModelResp
CreateModelResp
      (Text -> CreateModelResp) -> Parser Text -> Parser CreateModelResp
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
"status"

{- | Create a new model either from ModelFile or Path
Please note, if you specify both ModelFile and Path, ModelFile will be used.
-}
createModelOps ::
  -- | Model Name
  Text ->
  -- | Model File
  Maybe Text ->
  -- | Stream
  Maybe Bool ->
  -- | Path
  Maybe FilePath ->
  IO ()
createModelOps :: Text -> Maybe Text -> Maybe Bool -> Maybe FilePath -> IO ()
createModelOps
  Text
modelName
  Maybe Text
modelFile
  Maybe Bool
stream
  Maybe FilePath
path =
    do
      let url :: Text
url = OllamaClient -> Text
CU.host OllamaClient
defaultOllama
      Manager
manager <- ManagerSettings -> IO Manager
newManager ManagerSettings
defaultManagerSettings
      Request
initialRequest <- FilePath -> IO Request
forall (m :: * -> *). MonadThrow m => FilePath -> m Request
parseRequest (FilePath -> IO Request) -> FilePath -> IO Request
forall a b. (a -> b) -> a -> b
$ Text -> FilePath
T.unpack (Text
url Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
"/api/create")
      let reqBody :: CreateModelOps
reqBody =
            CreateModelOps
              { $sel:name:CreateModelOps :: Text
name = Text
modelName
              , $sel:modelFile:CreateModelOps :: Maybe Text
modelFile = Maybe Text
modelFile
              , $sel:stream:CreateModelOps :: Maybe Bool
stream = Maybe Bool
stream
              , $sel:path:CreateModelOps :: Maybe FilePath
path = Maybe FilePath
path
              }
          request :: Request
request =
            Request
initialRequest
              { method = "POST"
              , requestBody = RequestBodyLBS $ encode reqBody
              }
      Request -> Manager -> (Response BodyReader -> IO ()) -> IO ()
forall a.
Request -> Manager -> (Response BodyReader -> IO a) -> IO a
withResponse Request
request Manager
manager ((Response BodyReader -> IO ()) -> IO ())
-> (Response BodyReader -> IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ \Response BodyReader
response -> do
        let go :: IO ()
go = do
              Method
bs <- BodyReader -> BodyReader
brRead (BodyReader -> BodyReader) -> BodyReader -> BodyReader
forall a b. (a -> b) -> a -> b
$ Response BodyReader -> BodyReader
forall body. Response body -> body
responseBody Response BodyReader
response
              let eRes :: Either FilePath CreateModelResp
eRes =
                    ByteString -> Either FilePath CreateModelResp
forall a. FromJSON a => ByteString -> Either FilePath a
eitherDecode (Method -> ByteString
BSL.fromStrict Method
bs) ::
                      Either String CreateModelResp
              case Either FilePath CreateModelResp
eRes of
                Left FilePath
err -> do
                  FilePath -> IO ()
putStrLn (FilePath -> IO ()) -> FilePath -> IO ()
forall a b. (a -> b) -> a -> b
$ FilePath
"Error: " FilePath -> ShowS
forall a. Semigroup a => a -> a -> a
<> FilePath
err
                Right CreateModelResp
res -> do
                  Bool -> IO () -> IO ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
unless
                    (CreateModelResp -> Text
status CreateModelResp
res Text -> Text -> Bool
forall a. Eq a => a -> a -> Bool
/= Text
"success")
                    ( do
                        Text -> IO ()
T.putStr (Text -> IO ()) -> Text -> IO ()
forall a b. (a -> b) -> a -> b
$ CreateModelResp -> Text
status CreateModelResp
res
                        IO ()
go
                    )
        IO ()
go

{- | Create a new model
| Please note, if you specify both ModelFile and Path, ModelFile will be used.
-}
createModel ::
  -- | Model Name
  Text ->
  -- | Model File
  Maybe Text ->
  -- | Path
  Maybe FilePath ->
  IO ()
createModel :: Text -> Maybe Text -> Maybe FilePath -> IO ()
createModel Text
modelName Maybe Text
modelFile =
  Text -> Maybe Text -> Maybe Bool -> Maybe FilePath -> IO ()
createModelOps
    Text
modelName
    Maybe Text
modelFile
    Maybe Bool
forall a. Maybe a
Nothing