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

module OllamaExamples (main) where

import Control.Monad (void)
import Data.Aeson
import Data.List.NonEmpty (NonEmpty ((:|)))
import Data.Maybe (fromMaybe)
import Data.Ollama.Chat (chatJson)
import Data.Ollama.Chat qualified as Chat
import Data.Ollama.Common.Utils (encodeImage)
import Data.Ollama.Generate (generateJson)
import Data.Text.IO qualified as T
import GHC.Generics
import Ollama (GenerateOps (..), Role (..), chat, defaultChatOps, defaultGenerateOps, generate)
import Ollama qualified

data Example = Example
  { Example -> [String]
sortedList :: [String]
  , Example -> Bool
wasListAlreadSorted :: Bool
  }
  deriving (Int -> Example -> ShowS
[Example] -> ShowS
Example -> String
(Int -> Example -> ShowS)
-> (Example -> String) -> ([Example] -> ShowS) -> Show Example
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> Example -> ShowS
showsPrec :: Int -> Example -> ShowS
$cshow :: Example -> String
show :: Example -> String
$cshowList :: [Example] -> ShowS
showList :: [Example] -> ShowS
Show, Example -> Example -> Bool
(Example -> Example -> Bool)
-> (Example -> Example -> Bool) -> Eq Example
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: Example -> Example -> Bool
== :: Example -> Example -> Bool
$c/= :: Example -> Example -> Bool
/= :: Example -> Example -> Bool
Eq, (forall x. Example -> Rep Example x)
-> (forall x. Rep Example x -> Example) -> Generic Example
forall x. Rep Example x -> Example
forall x. Example -> Rep Example x
forall a.
(forall x. a -> Rep a x) -> (forall x. Rep a x -> a) -> Generic a
$cfrom :: forall x. Example -> Rep Example x
from :: forall x. Example -> Rep Example x
$cto :: forall x. Rep Example x -> Example
to :: forall x. Rep Example x -> Example
Generic, Maybe Example
Value -> Parser [Example]
Value -> Parser Example
(Value -> Parser Example)
-> (Value -> Parser [Example]) -> Maybe Example -> FromJSON Example
forall a.
(Value -> Parser a)
-> (Value -> Parser [a]) -> Maybe a -> FromJSON a
$cparseJSON :: Value -> Parser Example
parseJSON :: Value -> Parser Example
$cparseJSONList :: Value -> Parser [Example]
parseJSONList :: Value -> Parser [Example]
$comittedField :: Maybe Example
omittedField :: Maybe Example
FromJSON, [Example] -> Value
[Example] -> Encoding
Example -> Bool
Example -> Value
Example -> Encoding
(Example -> Value)
-> (Example -> Encoding)
-> ([Example] -> Value)
-> ([Example] -> Encoding)
-> (Example -> Bool)
-> ToJSON Example
forall a.
(a -> Value)
-> (a -> Encoding)
-> ([a] -> Value)
-> ([a] -> Encoding)
-> (a -> Bool)
-> ToJSON a
$ctoJSON :: Example -> Value
toJSON :: Example -> Value
$ctoEncoding :: Example -> Encoding
toEncoding :: Example -> Encoding
$ctoJSONList :: [Example] -> Value
toJSONList :: [Example] -> Value
$ctoEncodingList :: [Example] -> Encoding
toEncodingList :: [Example] -> Encoding
$comitField :: Example -> Bool
omitField :: Example -> Bool
ToJSON)

main :: IO ()
main :: IO ()
main = do
  -- Example 1: Streamed Text Generation
  -- This example demonstrates how to generate text using a model and stream the output directly to the console.
  -- The `stream` option enables processing of each chunk of the response as it arrives.
  IO (Either String GenerateResponse) -> IO ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (IO (Either String GenerateResponse) -> IO ())
-> IO (Either String GenerateResponse) -> IO ()
forall a b. (a -> b) -> a -> b
$
    GenerateOps -> IO (Either String GenerateResponse)
generate
      GenerateOps
defaultGenerateOps
        { modelName = "llama3.2"
        , prompt = "what is functional programming?"
        , stream = Just (T.putStr . Ollama.response_, pure ())
        }

  -- Example 2: Non-streamed Text Generation
  -- This example shows how to generate text and handle the complete response.
  -- The result is either an error message or the generated text.
  Either String GenerateResponse
eRes <-
    GenerateOps -> IO (Either String GenerateResponse)
generate
      GenerateOps
defaultGenerateOps
        { modelName = "llama3.2"
        , prompt = "What is 2+2?"
        }
  case Either String GenerateResponse
eRes of
    Left String
e -> String -> IO ()
putStrLn String
e
    Right Ollama.GenerateResponse {Bool
Maybe Int64
Text
UTCTime
$sel:response_:GenerateResponse :: GenerateResponse -> Text
model :: Text
createdAt :: UTCTime
response_ :: Text
done :: Bool
totalDuration :: Maybe Int64
loadDuration :: Maybe Int64
promptEvalCount :: Maybe Int64
promptEvalDuration :: Maybe Int64
evalCount :: Maybe Int64
evalDuration :: Maybe Int64
$sel:model:GenerateResponse :: GenerateResponse -> Text
$sel:createdAt:GenerateResponse :: GenerateResponse -> UTCTime
$sel:done:GenerateResponse :: GenerateResponse -> Bool
$sel:totalDuration:GenerateResponse :: GenerateResponse -> Maybe Int64
$sel:loadDuration:GenerateResponse :: GenerateResponse -> Maybe Int64
$sel:promptEvalCount:GenerateResponse :: GenerateResponse -> Maybe Int64
$sel:promptEvalDuration:GenerateResponse :: GenerateResponse -> Maybe Int64
$sel:evalCount:GenerateResponse :: GenerateResponse -> Maybe Int64
$sel:evalDuration:GenerateResponse :: GenerateResponse -> Maybe Int64
..} -> Text -> IO ()
T.putStrLn Text
response_

  -- Example 3: Chat with Streaming
  -- This example demonstrates setting up a chat session with streaming enabled.
  -- As messages are received, they are printed to the console.
  let msg :: Message
msg = Role -> Text -> Maybe [Text] -> Message
Ollama.Message Role
User Text
"What is functional programming?" Maybe [Text]
forall a. Maybe a
Nothing
      defaultMsg :: Message
defaultMsg = Role -> Text -> Maybe [Text] -> Message
Ollama.Message Role
User Text
"" Maybe [Text]
forall a. Maybe a
Nothing
  IO (Either String ChatResponse) -> IO ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (IO (Either String ChatResponse) -> IO ())
-> IO (Either String ChatResponse) -> IO ()
forall a b. (a -> b) -> a -> b
$
    ChatOps -> IO (Either String ChatResponse)
chat
      ChatOps
defaultChatOps
        { Chat.chatModelName = "llama3.2"
        , Chat.messages = msg :| []
        , Chat.stream =
            Just (T.putStr . Chat.content . fromMaybe defaultMsg . Chat.message, pure ())
        }

  -- Example 4: Non-streamed Chat
  -- Here, we handle a complete chat response, checking for potential errors.
  Either String ChatResponse
eRes1 <-
    ChatOps -> IO (Either String ChatResponse)
chat
      ChatOps
defaultChatOps
        { Chat.chatModelName = "llama3.2"
        , Chat.messages = msg :| []
        }
  case Either String ChatResponse
eRes1 of
    Left String
e -> String -> IO ()
putStrLn String
e
    Right ChatResponse
r -> do
      let mMessage :: Maybe Message
mMessage = ChatResponse -> Maybe Message
Ollama.message ChatResponse
r
      case Maybe Message
mMessage of
        Maybe Message
Nothing -> String -> IO ()
putStrLn String
"Something went wrong"
        Just Message
res -> Text -> IO ()
T.putStrLn (Text -> IO ()) -> Text -> IO ()
forall a b. (a -> b) -> a -> b
$ Message -> Text
Ollama.content Message
res

  -- Example 5: Check Model Status (ps)
  -- This example checks the status of models using the `ps` function.
  -- It outputs the status or details of the available models.
  Maybe RunningModels
res <- IO (Maybe RunningModels)
Ollama.ps
  Maybe RunningModels -> IO ()
forall a. Show a => a -> IO ()
print Maybe RunningModels
res

  -- Example 6: Simple Embedding
  -- This demonstrates how to request embeddings for a given text using a specific model.
  IO (Maybe EmbeddingResp) -> IO ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (IO (Maybe EmbeddingResp) -> IO ())
-> IO (Maybe EmbeddingResp) -> IO ()
forall a b. (a -> b) -> a -> b
$ Text -> Text -> IO (Maybe EmbeddingResp)
Ollama.embedding Text
"llama3.1" Text
"What is 5+2?"

  -- Example 7: Embedding with Options
  -- This example uses the `embeddingOps` function, allowing for additional configuration like options and streaming.
  IO (Maybe EmbeddingResp) -> IO ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (IO (Maybe EmbeddingResp) -> IO ())
-> IO (Maybe EmbeddingResp) -> IO ()
forall a b. (a -> b) -> a -> b
$ Text
-> Text -> Maybe Bool -> Maybe Text -> IO (Maybe EmbeddingResp)
Ollama.embeddingOps Text
"llama3.1" Text
"What is 5+2?" Maybe Bool
forall a. Maybe a
Nothing Maybe Text
forall a. Maybe a
Nothing

  -- Example 8: Stream Text Generation with JSON Body
  -- It is a higher level version of generate, here with genOps, you can also provide a Haskell type.
  -- You will get the response from LLM in this Haskell type.
  let expectedJsonStrucutre :: Example
expectedJsonStrucutre =
        Example
          { sortedList :: [String]
sortedList = [String
"sorted List here"]
          , wasListAlreadSorted :: Bool
wasListAlreadSorted = Bool
False
          }
  Either String Example
eRes2 <-
    GenerateOps -> Example -> Maybe Int -> IO (Either String Example)
forall jsonResult.
(ToJSON jsonResult, FromJSON jsonResult) =>
GenerateOps
-> jsonResult -> Maybe Int -> IO (Either String jsonResult)
generateJson
      GenerateOps
defaultGenerateOps
        { modelName = "llama3.2"
        , prompt = "Sort given list: [14, 12 , 13, 67]. Also tell whether list was already sorted or not."
        }
      Example
expectedJsonStrucutre
      (Int -> Maybe Int
forall a. a -> Maybe a
Just Int
2)
  case Either String Example
eRes2 of
    Left String
e -> String -> IO ()
putStrLn String
e
    Right Example
r -> (String, Example) -> IO ()
forall a. Show a => a -> IO ()
print (String
"JSON response: " :: String, Example
r)
  -- ("JSON response: ",Example {sortedList = ["1","2","3","4"], wasListAlreadSorted = False})

  -- Example 9: Chat with JSON Body
  -- This example demonstrates setting up a chat session but you receive the response in
  -- given haskell type.
  let msg0 :: Message
msg0 =
        Role -> Text -> Maybe [Text] -> Message
Ollama.Message
          Role
User
          Text
"Sort given list: [4, 2 , 3, 67]. Also tell whether list was already sorted or not."
          Maybe [Text]
forall a. Maybe a
Nothing
  Either String Example
eRes3 <-
    ChatOps -> Example -> Maybe Int -> IO (Either String Example)
forall jsonResult.
(FromJSON jsonResult, ToJSON jsonResult) =>
ChatOps -> jsonResult -> Maybe Int -> IO (Either String jsonResult)
chatJson
      ChatOps
defaultChatOps
        { Chat.chatModelName = "llama3.2"
        , Chat.messages = msg0 :| []
        }
      Example
expectedJsonStrucutre
      (Int -> Maybe Int
forall a. a -> Maybe a
Just Int
2)
  Either String Example -> IO ()
forall a. Show a => a -> IO ()
print Either String Example
eRes3

  -- Example 10: Chat with Image
  -- This example demonstrates chatting with example using an image.
  Maybe Text
mImg <- String -> IO (Maybe Text)
encodeImage String
"/home/user/sample.png"
  IO (Either String GenerateResponse) -> IO ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (IO (Either String GenerateResponse) -> IO ())
-> IO (Either String GenerateResponse) -> IO ()
forall a b. (a -> b) -> a -> b
$
    GenerateOps -> IO (Either String GenerateResponse)
generate
      GenerateOps
defaultGenerateOps
        { modelName = "llama3.2-vision"
        , prompt = "Describe the given image"
        , images = (\Text
x -> [Text] -> Maybe [Text]
forall a. a -> Maybe a
Just [Text
x]) =<< mImg
        , stream = Just (T.putStr . Ollama.response_, pure ())
        }

{-
Scotty example:
{-# LANGUAGE OverloadedStrings #-}

module Main where

import Web.Scotty
import Control.Monad.IO.Class (liftIO)
import Data.Aeson (FromJSON, ToJSON)
import Data.Text (Text)
import Data.Text qualified as T
import Database.SQLite.Simple
import Ollama (GenerateOps(..), defaultGenerateOps, generate)
import Data.Maybe (fromRight)

data PromptInput = PromptInput
  { conversation_id :: Int
  , prompt :: Text
  } deriving (Show, Generic)

instance FromJSON PromptInput
instance ToJSON PromptInput

main :: IO ()
main = do
  conn <- open "chat.db"
  execute_ conn "CREATE TABLE IF NOT EXISTS conversation (convo_id INTEGER PRIMARY KEY, convo_title TEXT)"
  execute_ conn "CREATE TABLE IF NOT EXISTS chats (chat_id INTEGER PRIMARY KEY, convo_id INTEGER, role TEXT, message TEXT, FOREIGN KEY(convo_id) REFERENCES conversation(convo_id))"

  scotty 3000 $ do
    post "/chat" $ do
      p <- jsonData :: ActionM PromptInput
      let cId = conversation_id p
      let trimmedP = T.dropEnd 3 $ T.drop 3 $ prompt p
      newConvoId <- case cId of
        -1 -> do
          liftIO $ execute conn "INSERT INTO conversation (convo_title) VALUES (?)" (Only ("latest title" :: String))
          [Only convoId] <- liftIO $ query_ conn "SELECT last_insert_rowid()" :: ActionM [Only Int]
          pure convoId
        _ -> pure cId

      liftIO $ execute conn "INSERT INTO chats (convo_id, role, message) VALUES (?, 'user', ?)" (newConvoId, trimmedP)

      stream $ \sendChunk flush -> do
        eRes <- generate defaultGenerateOps
                { modelName = "llama3.2"
                , prompt = prompt p
                , stream = Just (sendChunk . T.pack, flush)
                }
        case eRes of
            Left e -> return ()
            Right r -> do
                let res = response_ r
                liftIO $ execute conn "INSERT INTO chats (convo_id, role, message) VALUES (?, 'ai', ?)" (newConvoId, res)
-}