{-# LANGUAGE OverloadedStrings #-}

module Data.Ollama.Ps
  ( ps
  , RunningModels (..)
  , RunningModel (..)
  ) 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)

-- Types for Ps API
newtype RunningModels = RunningModels [RunningModel]
  deriving (RunningModels -> RunningModels -> Bool
(RunningModels -> RunningModels -> Bool)
-> (RunningModels -> RunningModels -> Bool) -> Eq RunningModels
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: RunningModels -> RunningModels -> Bool
== :: RunningModels -> RunningModels -> Bool
$c/= :: RunningModels -> RunningModels -> Bool
/= :: RunningModels -> RunningModels -> Bool
Eq, Int -> RunningModels -> ShowS
[RunningModels] -> ShowS
RunningModels -> String
(Int -> RunningModels -> ShowS)
-> (RunningModels -> String)
-> ([RunningModels] -> ShowS)
-> Show RunningModels
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> RunningModels -> ShowS
showsPrec :: Int -> RunningModels -> ShowS
$cshow :: RunningModels -> String
show :: RunningModels -> String
$cshowList :: [RunningModels] -> ShowS
showList :: [RunningModels] -> ShowS
Show)

data RunningModel = RunningModel
  { RunningModel -> Text
name_ :: Text
  , RunningModel -> Text
modelName :: Text
  , RunningModel -> Int64
size_ :: Int64
  , RunningModel -> Text
modelDigest :: Text
  , RunningModel -> ModelDetails
modelDetails :: ModelDetails
  , RunningModel -> UTCTime
expiresAt :: UTCTime
  , RunningModel -> Int64
sizeVRam :: Int64
  }
  deriving (RunningModel -> RunningModel -> Bool
(RunningModel -> RunningModel -> Bool)
-> (RunningModel -> RunningModel -> Bool) -> Eq RunningModel
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: RunningModel -> RunningModel -> Bool
== :: RunningModel -> RunningModel -> Bool
$c/= :: RunningModel -> RunningModel -> Bool
/= :: RunningModel -> RunningModel -> Bool
Eq, Int -> RunningModel -> ShowS
[RunningModel] -> ShowS
RunningModel -> String
(Int -> RunningModel -> ShowS)
-> (RunningModel -> String)
-> ([RunningModel] -> ShowS)
-> Show RunningModel
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> RunningModel -> ShowS
showsPrec :: Int -> RunningModel -> ShowS
$cshow :: RunningModel -> String
show :: RunningModel -> String
$cshowList :: [RunningModel] -> ShowS
showList :: [RunningModel] -> ShowS
Show)

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

instance FromJSON RunningModel where
  parseJSON :: Value -> Parser RunningModel
parseJSON = String
-> (Object -> Parser RunningModel) -> Value -> Parser RunningModel
forall a. String -> (Object -> Parser a) -> Value -> Parser a
withObject String
"RunningModel" ((Object -> Parser RunningModel) -> Value -> Parser RunningModel)
-> (Object -> Parser RunningModel) -> Value -> Parser RunningModel
forall a b. (a -> b) -> a -> b
$ \Object
v ->
    Text
-> Text
-> Int64
-> Text
-> ModelDetails
-> UTCTime
-> Int64
-> RunningModel
RunningModel
      (Text
 -> Text
 -> Int64
 -> Text
 -> ModelDetails
 -> UTCTime
 -> Int64
 -> RunningModel)
-> Parser Text
-> Parser
     (Text
      -> Int64
      -> Text
      -> ModelDetails
      -> UTCTime
      -> Int64
      -> RunningModel)
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
  (Text
   -> Int64
   -> Text
   -> ModelDetails
   -> UTCTime
   -> Int64
   -> RunningModel)
-> Parser Text
-> Parser
     (Int64 -> Text -> ModelDetails -> UTCTime -> Int64 -> RunningModel)
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
"model"
      Parser
  (Int64 -> Text -> ModelDetails -> UTCTime -> Int64 -> RunningModel)
-> Parser Int64
-> Parser
     (Text -> ModelDetails -> UTCTime -> Int64 -> RunningModel)
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 -> UTCTime -> Int64 -> RunningModel)
-> Parser Text
-> Parser (ModelDetails -> UTCTime -> Int64 -> RunningModel)
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 -> UTCTime -> Int64 -> RunningModel)
-> Parser ModelDetails -> Parser (UTCTime -> Int64 -> RunningModel)
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"
      Parser (UTCTime -> Int64 -> RunningModel)
-> Parser UTCTime -> Parser (Int64 -> RunningModel)
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
"expires_at"
      Parser (Int64 -> RunningModel)
-> Parser Int64 -> Parser RunningModel
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_vram"

-- | List running models
ps :: IO (Maybe RunningModels)
ps :: IO (Maybe RunningModels)
ps = 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/ps")
  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 RunningModels -> IO (Maybe RunningModels)
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Maybe RunningModels
forall a. Maybe a
Nothing
    else do
      let res :: Maybe RunningModels
res = ByteString -> Maybe RunningModels
forall a. FromJSON a => ByteString -> Maybe a
decode (Response ByteString -> ByteString
forall body. Response body -> body
responseBody Response ByteString
response) :: Maybe RunningModels
      case Maybe RunningModels
res of
        Maybe RunningModels
Nothing -> Maybe RunningModels -> IO (Maybe RunningModels)
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Maybe RunningModels
forall a. Maybe a
Nothing
        Just RunningModels
l -> Maybe RunningModels -> IO (Maybe RunningModels)
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Maybe RunningModels -> IO (Maybe RunningModels))
-> Maybe RunningModels -> IO (Maybe RunningModels)
forall a b. (a -> b) -> a -> b
$ RunningModels -> Maybe RunningModels
forall a. a -> Maybe a
Just RunningModels
l