Create.hs

{-# LANGUAGE DeriveGeneric #-}
{-# LANGUAGE DuplicateRecordFields #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE RecordWildCards #-}

{- |
Module      : Data.Ollama.Create
Copyright   : (c) 2025 Tushar Adhatrao
License     : MIT
Maintainer  : Tushar Adhatrao <tusharadhatrao@gmail.com>
Stability   : experimental
Description : Functionality for creating new models in the Ollama client.

This module provides functions to create new models in the Ollama API. Models can be created from:

- Another existing model (using 'from' parameter)
- A safetensors directory (using 'files' parameter)
- A GGUF file (using 'files' parameter)

The create operation is performed via a POST request to the @\/api\/create@ endpoint, with streaming
support for progress updates. The module supports model quantization, custom templates, system prompts,
parameters, and LORA adapters.

Example creating from existing model:

>>> let ops = defaultCreateOps { modelName = "mario", fromModel = Just "llama3.2", systemPrompt = Just "You are Mario from Super Mario Bros." }
>>> createModel ops Nothing
Creating model...
Success

Example quantizing a model:

>>> let ops = defaultCreateOps { modelName = "llama3.2:quantized", fromModel = Just "llama3.2:3b-instruct-fp16", quantizeType = Just Q4_K_M }
>>> createModel ops Nothing
Quantizing model...
Success
-}
module Data.Ollama.Create
  ( -- * Create Model API
    createModel
  , createModelM

    -- * Configuration Types
  , CreateOps (..)
  , defaultCreateOps

    -- * Response Types
  , CreateResp (..)

    -- * Quantization Types
  , QuantizationType (..)
  ) where

import Control.Monad (void)
import Control.Monad.IO.Class (MonadIO (liftIO))
import Data.Aeson
import Data.Map.Strict (Map)
import Data.Maybe (catMaybes)
import Data.Ollama.Common.Config (OllamaConfig)
import Data.Ollama.Common.Types (HasDone (getDone), Message, ModelOptions)
import Data.Ollama.Common.Utils as CU
import Data.Text (Text)
import GHC.Generics (Generic)
import GHC.Int (Int64)

-- | Quantization types supported by Ollama.
data QuantizationType
  = -- | Recommended quantization type
    Q4_K_M
  | -- | Alternative quantization type
    Q4_K_S
  | -- | Recommended quantization type
    Q8_0
  deriving (Show, Eq, Generic)

instance ToJSON QuantizationType where
  toJSON Q4_K_M = String "q4_K_M"
  toJSON Q4_K_S = String "q4_K_S"
  toJSON Q8_0 = String "q8_0"

instance FromJSON QuantizationType where
  parseJSON = withText "QuantizationType" $ \t ->
    case t of
      "q4_K_M" -> pure Q4_K_M
      "q4_K_S" -> pure Q4_K_S
      "q8_0" -> pure Q8_0
      _ -> fail $ "Invalid QuantizationType: " <> show t

-- | Configuration for creating a new model.
data CreateOps = CreateOps
  { modelName :: !Text
  -- ^ The name of the model to create.
  , fromModel :: !(Maybe Text)
  -- ^ Optional name of an existing model to create the new model from.
  , files :: !(Maybe (Map Text Text))
  -- ^ Optional dictionary of file names to SHA256 digests of blobs to create the model from.
  , adapters :: !(Maybe (Map Text Text))
  -- ^ Optional dictionary of file names to SHA256 digests of blobs for LORA adapters.
  , template :: !(Maybe Text)
  -- ^ Optional prompt template for the model.
  , license :: !(Maybe [Text])
  -- ^ Optional list of strings containing the license(s) for the model.
  , systemPrompt :: !(Maybe Text)
  -- ^ Optional system prompt for the model.
  , parameters :: !(Maybe ModelOptions)
  -- ^ Optional parameters for the model.
  , messages :: !(Maybe [Message])
  -- ^ Optional list of message objects used to create a conversation.
  , stream :: !(Maybe Bool)
  -- ^ Optional flag to enable streaming progress updates.
  , quantizeType :: !(Maybe QuantizationType)
  -- ^ Optional quantization type for quantizing a non-quantized model.
  }
  deriving (Show, Eq, Generic)

-- | Default configuration for creating a model.
defaultCreateOps :: Text -> CreateOps
defaultCreateOps name =
  CreateOps
    { modelName = name
    , fromModel = Nothing
    , files = Nothing
    , adapters = Nothing
    , template = Nothing
    , license = Nothing
    , systemPrompt = Nothing
    , parameters = Nothing
    , messages = Nothing
    , stream = Nothing
    , quantizeType = Nothing
    }

-- | Response type for model creation operations.
data CreateResp = CreateResp
  { status :: !Text
  -- ^ The status of the create operation (e.g., "success", "reading model metadata").
  , digest :: !(Maybe Text)
  -- ^ Optional digest (hash) of the model layer being processed.
  , total :: !(Maybe Int64)
  -- ^ Optional total size in bytes for quantization operations.
  , completed :: !(Maybe Int64)
  -- ^ Optional number of bytes completed for quantization operations.
  }
  deriving (Show, Eq, Generic)

instance HasDone CreateResp where
  getDone CreateResp {..} = status == "success"

instance ToJSON CreateOps where
  toJSON CreateOps {..} =
    object $
      catMaybes
        [ Just ("model" .= modelName)
        , ("from" .=) <$> fromModel
        , ("files" .=) <$> files
        , ("adapters" .=) <$> adapters
        , ("template" .=) <$> template
        , ("license" .=) <$> license
        , ("system" .=) <$> systemPrompt
        , ("parameters" .=) <$> parameters
        , ("messages" .=) <$> messages
        , ("stream" .=) <$> stream
        , ("quantize" .=) <$> quantizeType
        ]

instance FromJSON CreateResp where
  parseJSON = withObject "CreateResp" $ \v ->
    CreateResp
      <$> v .: "status"
      <*> v .:? "digest"
      <*> v .:? "total"
      <*> v .:? "completed"

{- | Creates a new model according to the provided configuration.

Sends a POST request to the @\/api\/create@ endpoint to create a model. The model can be created from:

- Another existing model (specify 'fromModel')
- A safetensors directory (specify 'files' with file name to SHA256 digest mappings)
- A GGUF file (specify 'files' with file name to SHA256 digest mapping)

Supports quantization, custom templates, system prompts, parameters, and LORA adapters.
Prints progress messages to the console during creation.
-}
createModel ::
  -- | Model creation configuration
  CreateOps ->
  -- | Optional 'OllamaConfig' (defaults to 'defaultOllamaConfig' if 'Nothing')
  Maybe OllamaConfig ->
  IO ()
createModel createOps mbConfig =
  void $
    withOllamaRequest
      "/api/create"
      "POST"
      (Just createOps)
      mbConfig
      (commonStreamHandler (onToken, pure ()))
  where
    onToken :: CreateResp -> IO ()
    onToken CreateResp {..} = case (total, completed) of
      (Just t, Just c) -> putStrLn $ "Progress: " <> show c <> "/" <> show t <> " bytes"
      _ -> putStrLn $ "Status: " <> show status

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

Lifts the 'createModel' function into a 'MonadIO' context, allowing it to be used in monadic computations.
-}
createModelM ::
  MonadIO m =>
  CreateOps ->
  Maybe OllamaConfig ->
  m ()
createModelM createOps mbCfg = liftIO $ createModel createOps mbCfg