{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE RecordWildCards #-}

-- | The locally hosted ollama backend
module GHC.Plugin.OllamaHoles.Backend.Ollama (ollamaBackend) where

import Ollama (GenerateOps (..))
import Ollama qualified

import GHC.Plugin.OllamaHoles.Backend

-- | The locally hosted ollama backend
ollamaBackend :: Backend
ollamaBackend :: Backend
ollamaBackend = Backend{IO (Maybe [Text])
Text -> Text -> IO (Either String Text)
listModels :: IO (Maybe [Text])
generateFits :: Text -> Text -> IO (Either String Text)
listModels :: IO (Maybe [Text])
generateFits :: Text -> Text -> IO (Either String Text)
..}
  where
    listModels :: IO (Maybe [Text])
listModels = (Models -> [Text]) -> Maybe Models -> Maybe [Text]
forall a b. (a -> b) -> Maybe a -> Maybe b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap Models -> [Text]
getMs (Maybe Models -> Maybe [Text])
-> IO (Maybe Models) -> IO (Maybe [Text])
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> IO (Maybe Models)
Ollama.list
    getMs :: Models -> [Text]
getMs (Ollama.Models [ModelInfo]
models) = (ModelInfo -> Text) -> [ModelInfo] -> [Text]
forall a b. (a -> b) -> [a] -> [b]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap ModelInfo -> Text
Ollama.name [ModelInfo]
models
    generateFits :: Text -> Text -> IO (Either String Text)
generateFits Text
prompt Text
modelName = do
        (GenerateResponse -> Text)
-> Either String GenerateResponse -> Either String Text
forall a b. (a -> b) -> Either String a -> Either String b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap GenerateResponse -> Text
Ollama.response_ (Either String GenerateResponse -> Either String Text)
-> IO (Either String GenerateResponse) -> IO (Either String Text)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> GenerateOps -> IO (Either String GenerateResponse)
Ollama.generate GenerateOps
genOps{prompt = prompt, modelName = modelName}

    -- \| Options for the LLM model
    genOps :: Ollama.GenerateOps
    genOps :: GenerateOps
genOps =
        GenerateOps
Ollama.defaultGenerateOps
            { modelName = ""
            , prompt = ""
            }