{-# LANGUAGE OverloadedStrings #-}

module Data.Ollama.List
  ( -- * List Models API
    list
  , Models (..)
  , ModelInfo (..)
  )
where

import Data.Aeson
import Data.Ollama.Common.Types as CT
import Data.Ollama.Common.Utils as CU
import Data.Text (Text)
import Data.Text qualified as T
import Data.Time
import GHC.Int (Int64)
import Network.HTTP.Client
import Network.HTTP.Types.Status (statusCode)

newtype Models = Models [ModelInfo]
  deriving (Models -> Models -> Bool
(Models -> Models -> Bool)
-> (Models -> Models -> Bool) -> Eq Models
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: Models -> Models -> Bool
== :: Models -> Models -> Bool
$c/= :: Models -> Models -> Bool
/= :: Models -> Models -> Bool
Eq, Int -> Models -> ShowS
[Models] -> ShowS
Models -> String
(Int -> Models -> ShowS)
-> (Models -> String) -> ([Models] -> ShowS) -> Show Models
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> Models -> ShowS
showsPrec :: Int -> Models -> ShowS
$cshow :: Models -> String
show :: Models -> String
$cshowList :: [Models] -> ShowS
showList :: [Models] -> ShowS
Show)

data ModelInfo = ModelInfo
  { ModelInfo -> Text
name :: Text
  , ModelInfo -> UTCTime
modifiedAt :: UTCTime
  , ModelInfo -> Int64
size :: Int64
  , ModelInfo -> Text
digest :: Text
  , ModelInfo -> ModelDetails
details :: ModelDetails
  }
  deriving (ModelInfo -> ModelInfo -> Bool
(ModelInfo -> ModelInfo -> Bool)
-> (ModelInfo -> ModelInfo -> Bool) -> Eq ModelInfo
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: ModelInfo -> ModelInfo -> Bool
== :: ModelInfo -> ModelInfo -> Bool
$c/= :: ModelInfo -> ModelInfo -> Bool
/= :: ModelInfo -> ModelInfo -> Bool
Eq, Int -> ModelInfo -> ShowS
[ModelInfo] -> ShowS
ModelInfo -> String
(Int -> ModelInfo -> ShowS)
-> (ModelInfo -> String)
-> ([ModelInfo] -> ShowS)
-> Show ModelInfo
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> ModelInfo -> ShowS
showsPrec :: Int -> ModelInfo -> ShowS
$cshow :: ModelInfo -> String
show :: ModelInfo -> String
$cshowList :: [ModelInfo] -> ShowS
showList :: [ModelInfo] -> ShowS
Show)

-- Instances
instance FromJSON Models where
  parseJSON :: Value -> Parser Models
parseJSON = String -> (Object -> Parser Models) -> Value -> Parser Models
forall a. String -> (Object -> Parser a) -> Value -> Parser a
withObject String
"Models" ((Object -> Parser Models) -> Value -> Parser Models)
-> (Object -> Parser Models) -> Value -> Parser Models
forall a b. (a -> b) -> a -> b
$ \Object
v -> [ModelInfo] -> Models
Models ([ModelInfo] -> Models) -> Parser [ModelInfo] -> Parser Models
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Object
v Object -> Key -> Parser [ModelInfo]
forall a. FromJSON a => Object -> Key -> Parser a
.: Key
"models"

instance FromJSON ModelInfo where
  parseJSON :: Value -> Parser ModelInfo
parseJSON = String -> (Object -> Parser ModelInfo) -> Value -> Parser ModelInfo
forall a. String -> (Object -> Parser a) -> Value -> Parser a
withObject String
"ModelInfo" ((Object -> Parser ModelInfo) -> Value -> Parser ModelInfo)
-> (Object -> Parser ModelInfo) -> Value -> Parser ModelInfo
forall a b. (a -> b) -> a -> b
$ \Object
v ->
    Text -> UTCTime -> Int64 -> Text -> ModelDetails -> ModelInfo
ModelInfo
      (Text -> UTCTime -> Int64 -> Text -> ModelDetails -> ModelInfo)
-> Parser Text
-> Parser (UTCTime -> Int64 -> Text -> ModelDetails -> ModelInfo)
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
"name"
      Parser (UTCTime -> Int64 -> Text -> ModelDetails -> ModelInfo)
-> Parser UTCTime
-> Parser (Int64 -> Text -> ModelDetails -> ModelInfo)
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 UTCTime
forall a. FromJSON a => Object -> Key -> Parser a
.: Key
"modified_at"
      Parser (Int64 -> Text -> ModelDetails -> ModelInfo)
-> Parser Int64 -> Parser (Text -> ModelDetails -> ModelInfo)
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 Int64
forall a. FromJSON a => Object -> Key -> Parser a
.: Key
"size"
      Parser (Text -> ModelDetails -> ModelInfo)
-> Parser Text -> Parser (ModelDetails -> ModelInfo)
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 Text
forall a. FromJSON a => Object -> Key -> Parser a
.: Key
"digest"
      Parser (ModelDetails -> ModelInfo)
-> Parser ModelDetails -> Parser ModelInfo
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 ModelDetails
forall a. FromJSON a => Object -> Key -> Parser a
.: Key
"details"

-- | List all models from local
list :: IO (Maybe Models)
list :: IO (Maybe Models)
list = do
  let url :: Text
url = OllamaClient -> Text
CU.host OllamaClient
defaultOllama
  Manager
manager <- ManagerSettings -> IO Manager
newManager ManagerSettings
defaultManagerSettings
  Request
request <- 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/tags")
  Response ByteString
response <- Request -> Manager -> IO (Response ByteString)
httpLbs Request
request Manager
manager
  if Status -> Int
statusCode (Response ByteString -> Status
forall body. Response body -> Status
responseStatus Response ByteString
response) Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= Int
200
    then Maybe Models -> IO (Maybe Models)
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Maybe Models
forall a. Maybe a
Nothing
    else do
      let res :: Maybe Models
res = ByteString -> Maybe Models
forall a. FromJSON a => ByteString -> Maybe a
decode (Response ByteString -> ByteString
forall body. Response body -> body
responseBody Response ByteString
response) :: Maybe Models
      case Maybe Models
res of
        Maybe Models
Nothing -> Maybe Models -> IO (Maybe Models)
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Maybe Models
forall a. Maybe a
Nothing
        Just Models
l -> Maybe Models -> IO (Maybe Models)
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Maybe Models -> IO (Maybe Models))
-> Maybe Models -> IO (Maybe Models)
forall a b. (a -> b) -> a -> b
$ Models -> Maybe Models
forall a. a -> Maybe a
Just Models
l