diff --git a/CHANGELOG.md b/CHANGELOG.md --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,5 +1,10 @@ # Revision history for ollama-haskell +## 0.1.2.0 -- 2024-11-20 + +* Added hostUrl and responseTimeOut options in generate function. +* Added hostUrl and responseTimeOut options in chat function. + ## 0.1.1.3 -- 2024-11-08 * Increase response timeout to 15 minutes diff --git a/README.md b/README.md --- a/README.md +++ b/README.md @@ -38,7 +38,7 @@ 3. **Garbage collection**: Haskell handles memory management automatically, freeing you from worries about manual memory deallocation. ``` -You can find practical examples demonstrating how to use the library in the `src/OllamaExamples.hs` file. +You can find practical examples demonstrating how to use the library in the `examples/OllamaExamples.hs` file. ## Prerequisite diff --git a/ollama-haskell.cabal b/ollama-haskell.cabal --- a/ollama-haskell.cabal +++ b/ollama-haskell.cabal @@ -5,7 +5,7 @@ -- see: https://github.com/sol/hpack name: ollama-haskell -version: 0.1.1.3 +version: 0.1.2.0 synopsis: Haskell bindings for ollama. description: Please see the README on GitHub at category: Web @@ -41,7 +41,6 @@ Data.Ollama.Push Data.Ollama.Show Ollama - OllamaExamples other-modules: Paths_ollama_haskell hs-source-dirs: diff --git a/src/Data/Ollama/Chat.hs b/src/Data/Ollama/Chat.hs --- a/src/Data/Ollama/Chat.hs +++ b/src/Data/Ollama/Chat.hs @@ -14,11 +14,12 @@ , ChatResponse (..) ) where +import Control.Exception (try) import Data.Aeson import Data.ByteString.Char8 qualified as BS import Data.ByteString.Lazy.Char8 qualified as BSL import Data.List.NonEmpty as NonEmpty -import Data.Maybe (isNothing) +import Data.Maybe (fromMaybe, isNothing) import Data.Ollama.Common.Utils as CU import Data.Text (Text) import Data.Text qualified as T @@ -72,6 +73,10 @@ -- ^ Optional streaming functions where the first handles each chunk of the response, and the second flushes the stream. , keepAlive :: Maybe Text -- ^ Optional text to specify keep-alive behavior. + , hostUrl :: Maybe Text + -- ^ Override default Ollama host url. Default url = "http://127.0.0.1:11434" + , responseTimeOut :: Maybe Int + -- ^ Override default response timeout in minutes. Default = 15 minutes } instance Show ChatOps where @@ -123,7 +128,7 @@ deriving (Show, Eq) instance ToJSON ChatOps where - toJSON (ChatOps model_ messages_ tools_ format_ stream_ keepAlive_) = + toJSON (ChatOps model_ messages_ tools_ format_ stream_ keepAlive_ _ _) = object [ "model" .= model_ , "messages" .= messages_ @@ -165,6 +170,8 @@ , format = Nothing , stream = Nothing , keepAlive = Nothing + , hostUrl = Nothing + , responseTimeOut = Nothing } {- | @@ -183,45 +190,57 @@ -} chat :: ChatOps -> IO (Either String ChatResponse) chat cOps = do - let url = CU.host defaultOllama + let url = fromMaybe defaultOllamaUrl (hostUrl cOps) + responseTimeout = fromMaybe 15 (responseTimeOut cOps) manager <- newManager defaultManagerSettings -- Setting response timeout to 5 minutes, since llm takes time - { managerResponseTimeout = responseTimeoutMicro (15 * 60 * 1000000) + { managerResponseTimeout = responseTimeoutMicro (responseTimeout * 60 * 1000000) } - initialRequest <- parseRequest $ T.unpack (url <> "/api/chat") - let reqBody = cOps - request = - initialRequest - { method = "POST" - , requestBody = RequestBodyLBS $ encode reqBody - } - withResponse request manager $ \response -> do - let streamResponse sendChunk flush = do - bs <- brRead $ responseBody response - if BS.null bs - then putStrLn "" >> pure (Left "") - else do - let eRes = eitherDecode (BSL.fromStrict bs) :: Either String ChatResponse - case eRes of - Left e -> pure (Left e) - Right r -> do - _ <- sendChunk r - _ <- flush - if done r then pure (Left "") else streamResponse sendChunk flush - let genResponse op = do - bs <- brRead $ responseBody response - if BS.null bs - then do - let eRes = eitherDecode (BSL.fromStrict op) :: Either String ChatResponse - case eRes of - Left e -> pure (Left e) - Right r -> pure (Right r) - else genResponse (op <> bs) - case stream cOps of - Nothing -> genResponse "" - Just (sendChunk, flush) -> streamResponse sendChunk flush + eInitialRequest <- try $ parseRequest $ T.unpack (url <> "/api/chat") :: IO (Either HttpException Request) + case eInitialRequest of + Left e -> return $ Left $ "Failed to parse host url: " <> show e + Right initialRequest -> do + let reqBody = cOps + request = + initialRequest + { method = "POST" + , requestBody = RequestBodyLBS $ encode reqBody + } + eRes <- + try (withResponse request manager $ handleRequest cOps) :: + IO (Either HttpException (Either String ChatResponse)) + case eRes of + Left e -> return $ Left $ "HTTP error occured: " <> show e + Right r -> return r +handleRequest :: ChatOps -> Response BodyReader -> IO (Either String ChatResponse) +handleRequest cOps response = do + let streamResponse sendChunk flush = do + bs <- brRead $ responseBody response + if BS.null bs + then putStrLn "" >> pure (Left "") + else do + let eRes = eitherDecode (BSL.fromStrict bs) :: Either String ChatResponse + case eRes of + Left e -> pure (Left e) + Right r -> do + _ <- sendChunk r + _ <- flush + if done r then pure (Left "") else streamResponse sendChunk flush + let genResponse op = do + bs <- brRead $ responseBody response + if BS.null bs + then do + let eRes = eitherDecode (BSL.fromStrict op) :: Either String ChatResponse + case eRes of + Left e -> pure (Left e) + Right r -> pure (Right r) + else genResponse (op <> bs) + case stream cOps of + Nothing -> genResponse "" + Just (sendChunk, flush) -> streamResponse sendChunk flush + {- | chatJson is a higher level function that takes ChatOps (similar to chat) and also takes a Haskell type (that has To and From JSON instance) and returns the response in provided type. @@ -258,8 +277,10 @@ chatJson :: (FromJSON jsonResult, ToJSON jsonResult) => ChatOps -> - jsonResult -> -- ^ Haskell type that you want your result in - Maybe Int -> -- ^ Max retries + -- | Haskell type that you want your result in + jsonResult -> + -- | Max retries + Maybe Int -> IO (Either String jsonResult) chatJson cOps@ChatOps {..} jsonStructure mMaxRetries = do let lastMessage = NonEmpty.last messages @@ -288,8 +309,8 @@ Nothing -> return $ Left "Something went wrong" Just res -> do case decode (BSL.fromStrict . T.encodeUtf8 $ content res) of - Nothing -> do - case mMaxRetries of - Nothing -> return $ Left "Decoding Failed :(" - Just n -> if n < 1 then return $ Left "Decoding Failed :(" else chatJson cOps jsonStructure (Just (n - 1)) + Nothing -> do + case mMaxRetries of + Nothing -> return $ Left "Decoding Failed :(" + Just n -> if n < 1 then return $ Left "Decoding Failed :(" else chatJson cOps jsonStructure (Just (n - 1)) Just resultInType -> return $ Right resultInType diff --git a/src/Data/Ollama/Common/Utils.hs b/src/Data/Ollama/Common/Utils.hs --- a/src/Data/Ollama/Common/Utils.hs +++ b/src/Data/Ollama/Common/Utils.hs @@ -1,6 +1,6 @@ {-# LANGUAGE OverloadedStrings #-} -module Data.Ollama.Common.Utils (defaultOllama, OllamaClient (..), encodeImage) where +module Data.Ollama.Common.Utils (defaultOllamaUrl, OllamaClient (..), encodeImage) where import Control.Exception (IOException, try) import Data.ByteString qualified as BS @@ -12,8 +12,8 @@ import System.Directory import System.FilePath -defaultOllama :: OllamaClient -defaultOllama = OllamaClient "http://127.0.0.1:11434" +defaultOllamaUrl :: Text +defaultOllamaUrl = "http://127.0.0.1:11434" supportedExtensions :: [String] supportedExtensions = [".jpg", ".jpeg", ".png"] @@ -33,7 +33,7 @@ {- | encodeImage is a utility function that takes an image file path (jpg, jpeg, png) and - returns the image data in Base64 encoded format. Since GenerateOps' images field + returns the image data in Base64 encoded format. Since GenerateOps' images field expects image data in base64. It is helper function that we are providing out of the box. -} encodeImage :: FilePath -> IO (Maybe Text) diff --git a/src/Data/Ollama/Copy.hs b/src/Data/Ollama/Copy.hs --- a/src/Data/Ollama/Copy.hs +++ b/src/Data/Ollama/Copy.hs @@ -36,7 +36,7 @@ source_ destination_ = do - let url = CU.host CU.defaultOllama + let url = CU.defaultOllamaUrl manager <- newManager defaultManagerSettings initialRequest <- parseRequest $ T.unpack (url <> "/api/copy") let reqBody = diff --git a/src/Data/Ollama/Create.hs b/src/Data/Ollama/Create.hs --- a/src/Data/Ollama/Create.hs +++ b/src/Data/Ollama/Create.hs @@ -69,7 +69,7 @@ stream_ path_ = do - let url = CU.host defaultOllama + let url = defaultOllamaUrl manager <- newManager defaultManagerSettings initialRequest <- parseRequest $ T.unpack (url <> "/api/create") let reqBody = diff --git a/src/Data/Ollama/Delete.hs b/src/Data/Ollama/Delete.hs --- a/src/Data/Ollama/Delete.hs +++ b/src/Data/Ollama/Delete.hs @@ -28,7 +28,7 @@ IO () deleteModel modelName = do - let url = CU.host CU.defaultOllama + let url = CU.defaultOllamaUrl manager <- newManager defaultManagerSettings initialRequest <- parseRequest $ T.unpack (url <> "/api/delete") let reqBody = diff --git a/src/Data/Ollama/Embeddings.hs b/src/Data/Ollama/Embeddings.hs --- a/src/Data/Ollama/Embeddings.hs +++ b/src/Data/Ollama/Embeddings.hs @@ -54,7 +54,7 @@ Maybe Text -> IO (Maybe EmbeddingResp) embeddingOps modelName input_ mTruncate mKeepAlive = do - let url = CU.host defaultOllama + let url = defaultOllamaUrl manager <- newManager defaultManagerSettings initialRequest <- parseRequest $ T.unpack (url <> "/api/embed") let reqBody = diff --git a/src/Data/Ollama/Generate.hs b/src/Data/Ollama/Generate.hs --- a/src/Data/Ollama/Generate.hs +++ b/src/Data/Ollama/Generate.hs @@ -11,6 +11,7 @@ , GenerateResponse (..) ) where +import Control.Exception (try) import Data.Aeson import Data.ByteString.Char8 qualified as BS import Data.ByteString.Lazy.Char8 qualified as BSL @@ -66,6 +67,10 @@ -- ^ An optional flag to return the raw response. , keepAlive :: Maybe Text -- ^ Optional text to specify keep-alive behavior. + , hostUrl :: Maybe Text + -- ^ Override default Ollama host url. Default url = "http://127.0.0.1:11434" + , responseTimeOut :: Maybe Int + -- ^ Override default response timeout in minutes. Default = 15 minutes } instance Show GenerateOps where @@ -146,6 +151,8 @@ stream raw keepAlive + _ -- Host url + _ -- Response timeout ) = object [ "model" .= model @@ -197,6 +204,8 @@ , stream = Nothing , raw = Nothing , keepAlive = Nothing + , hostUrl = Nothing + , responseTimeOut = Nothing } {- | @@ -239,42 +248,59 @@ -} generate :: GenerateOps -> IO (Either String GenerateResponse) generate genOps = do - let url = CU.host defaultOllama + let url = fromMaybe defaultOllamaUrl (hostUrl genOps) + responseTimeout = fromMaybe 15 (responseTimeOut genOps) manager <- newManager -- Setting response timeout to 5 minutes, since llm takes time - defaultManagerSettings {managerResponseTimeout = responseTimeoutMicro (15 * 60 * 1000000)} - initialRequest <- parseRequest $ T.unpack (url <> "/api/generate") - let reqBody = genOps - request = - initialRequest - { method = "POST" - , requestBody = RequestBodyLBS $ encode reqBody - } - withResponse request manager $ \response -> do - let streamResponse sendChunk flush = do - bs <- brRead $ responseBody response - if BS.null bs - then putStrLn "" >> pure (Left "") - else do - let eRes = eitherDecode (BSL.fromStrict bs) :: Either String GenerateResponse - case eRes of - Left e -> pure (Left e) - Right r -> do - _ <- sendChunk r - _ <- flush - if done r then pure (Left "") else streamResponse sendChunk flush - let genResponse op = do - bs <- brRead $ responseBody response - if bs == "" - then do - let eRes0 = eitherDecode (BSL.fromStrict op) :: Either String GenerateResponse - case eRes0 of - Left e -> pure (Left e) - Right r -> pure (Right r) - else genResponse (op <> bs) - case stream genOps of - Nothing -> genResponse "" - Just (sendChunk, flush) -> streamResponse sendChunk flush + defaultManagerSettings + { managerResponseTimeout = responseTimeoutMicro (responseTimeout * 60 * 1000000) + } + eInitialRequest <- + try $ parseRequest $ T.unpack (url <> "/api/generate") :: IO (Either HttpException Request) + case eInitialRequest of + Left e -> do + return $ Left $ show e + Right initialRequest -> do + let reqBody = genOps + request = + initialRequest + { method = "POST" + , requestBody = RequestBodyLBS $ encode reqBody + } + eRes <- + try (withResponse request manager $ handleRequest genOps) :: + IO (Either HttpException (Either String GenerateResponse)) + case eRes of + Left e -> do + return $ Left $ "HTTP error occured: " <> show e + Right r -> return r + +handleRequest :: GenerateOps -> Response BodyReader -> IO (Either String GenerateResponse) +handleRequest genOps response = do + let streamResponse sendChunk flush = do + bs <- brRead $ responseBody response + if BS.null bs + then putStrLn "" >> pure (Left "") + else do + let eRes = eitherDecode (BSL.fromStrict bs) :: Either String GenerateResponse + case eRes of + Left e -> pure (Left e) + Right r -> do + _ <- sendChunk r + _ <- flush + if done r then pure (Left "") else streamResponse sendChunk flush + let genResponse op = do + bs <- brRead $ responseBody response + if bs == "" + then do + let eRes0 = eitherDecode (BSL.fromStrict op) :: Either String GenerateResponse + case eRes0 of + Left e -> pure (Left e) + Right r -> pure (Right r) + else genResponse (op <> bs) + case stream genOps of + Nothing -> genResponse "" + Just (sendChunk, flush) -> streamResponse sendChunk flush {- | generateJson is a higher level function that takes generateOps (similar to generate) and also takes diff --git a/src/Data/Ollama/List.hs b/src/Data/Ollama/List.hs --- a/src/Data/Ollama/List.hs +++ b/src/Data/Ollama/List.hs @@ -46,7 +46,7 @@ -- | List all models from local list :: IO (Maybe Models) list = do - let url = CU.host defaultOllama + let url = defaultOllamaUrl manager <- newManager defaultManagerSettings request <- parseRequest $ T.unpack (url <> "/api/tags") response <- httpLbs request manager diff --git a/src/Data/Ollama/Ps.hs b/src/Data/Ollama/Ps.hs --- a/src/Data/Ollama/Ps.hs +++ b/src/Data/Ollama/Ps.hs @@ -48,7 +48,7 @@ -- | List running models ps :: IO (Maybe RunningModels) ps = do - let url = CU.host defaultOllama + let url = defaultOllamaUrl manager <- newManager defaultManagerSettings request <- parseRequest $ T.unpack (url <> "/api/ps") response <- httpLbs request manager diff --git a/src/Data/Ollama/Pull.hs b/src/Data/Ollama/Pull.hs --- a/src/Data/Ollama/Pull.hs +++ b/src/Data/Ollama/Pull.hs @@ -66,7 +66,7 @@ Maybe Bool -> IO () pullOps modelName mInsecure mStream = do - let url = CU.host defaultOllama + let url = defaultOllamaUrl manager <- newManager defaultManagerSettings initialRequest <- parseRequest $ T.unpack (url <> "/api/pull") let reqBody = diff --git a/src/Data/Ollama/Push.hs b/src/Data/Ollama/Push.hs --- a/src/Data/Ollama/Push.hs +++ b/src/Data/Ollama/Push.hs @@ -44,7 +44,7 @@ Maybe Bool -> IO () pushOps modelName mInsecure mStream = do - let url = CU.host defaultOllama + let url = defaultOllamaUrl manager <- newManager defaultManagerSettings initialRequest <- parseRequest $ T.unpack (url <> "/api/push") let reqBody = diff --git a/src/Data/Ollama/Show.hs b/src/Data/Ollama/Show.hs --- a/src/Data/Ollama/Show.hs +++ b/src/Data/Ollama/Show.hs @@ -141,7 +141,7 @@ modelName verbose_ = do - let url = CU.host CU.defaultOllama + let url = CU.defaultOllamaUrl manager <- newManager defaultManagerSettings initialRequest <- parseRequest $ T.unpack (url <> "/api/show") let reqBody = diff --git a/src/OllamaExamples.hs b/src/OllamaExamples.hs deleted file mode 100644 --- a/src/OllamaExamples.hs +++ /dev/null @@ -1,202 +0,0 @@ -{-# 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 - { sortedList :: [String] - , wasListAlreadSorted :: Bool - } - deriving (Show, Eq, Generic, FromJSON, ToJSON) - -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. - void $ - generate - 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. - eRes <- - generate - defaultGenerateOps - { modelName = "llama3.2" - , prompt = "What is 2+2?" - } - case eRes of - Left e -> putStrLn e - Right Ollama.GenerateResponse {..} -> T.putStrLn 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 = Ollama.Message User "What is functional programming?" Nothing - defaultMsg = Ollama.Message User "" Nothing - void $ - chat - 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. - eRes1 <- - chat - defaultChatOps - { Chat.chatModelName = "llama3.2" - , Chat.messages = msg :| [] - } - case eRes1 of - Left e -> putStrLn e - Right r -> do - let mMessage = Ollama.message r - case mMessage of - Nothing -> putStrLn "Something went wrong" - Just res -> T.putStrLn $ Ollama.content 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. - res <- Ollama.ps - print res - - -- Example 6: Simple Embedding - -- This demonstrates how to request embeddings for a given text using a specific model. - void $ Ollama.embedding "llama3.1" "What is 5+2?" - - -- Example 7: Embedding with Options - -- This example uses the `embeddingOps` function, allowing for additional configuration like options and streaming. - void $ Ollama.embeddingOps "llama3.1" "What is 5+2?" Nothing 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 - { sortedList = ["sorted List here"] - , wasListAlreadSorted = False - } - eRes2 <- - generateJson - defaultGenerateOps - { modelName = "llama3.2" - , prompt = "Sort given list: [14, 12 , 13, 67]. Also tell whether list was already sorted or not." - } - expectedJsonStrucutre - (Just 2) - case eRes2 of - Left e -> putStrLn e - Right r -> print ("JSON response: " :: String, 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 = - Ollama.Message - User - "Sort given list: [4, 2 , 3, 67]. Also tell whether list was already sorted or not." - Nothing - eRes3 <- - chatJson - defaultChatOps - { Chat.chatModelName = "llama3.2" - , Chat.messages = msg0 :| [] - } - expectedJsonStrucutre - (Just 2) - print eRes3 - - -- Example 10: Chat with Image - -- This example demonstrates chatting with example using an image. - mImg <- encodeImage "/home/user/sample.png" - void $ - generate - defaultGenerateOps - { modelName = "llama3.2-vision" - , prompt = "Describe the given image" - , images = (\x -> Just [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) --} diff --git a/test/Main.hs b/test/Main.hs --- a/test/Main.hs +++ b/test/Main.hs @@ -1,75 +1,123 @@ {-# LANGUAGE OverloadedStrings #-} -module Main where -import Ollama qualified -import Ollama (defaultGenerateOps, GenerateOps(..), defaultChatOps, Role(..)) +module Main (main) where + +import Data.Either +import Data.List.NonEmpty hiding (length) +import Data.Maybe import Data.Ollama.Chat qualified as Chat import Data.Text.IO qualified as T +import Ollama (GenerateOps (..), Role (..), defaultChatOps, defaultGenerateOps) +import Ollama qualified +import System.IO.Silently (capture) import Test.Tasty import Test.Tasty.HUnit -import System.IO.Silently (capture) -import Data.Either -import Data.List.NonEmpty hiding (length) -import Data.Maybe tests :: TestTree -tests = testGroup "Tests" [ - generateTest - , chatTest - , psTest - , showTest - ] +tests = + testGroup + "Tests" + [ generateTest + , chatTest + , psTest + , showTest + ] generateTest :: TestTree -generateTest = testGroup "Generate tests" [ - testCase "generate stream" $ do - output <- capture $ Ollama.generate defaultGenerateOps { modelName = "smollm:360m" - , prompt = "what is 4 + 2?" - , stream = Just (T.putStr . Ollama.response_, pure ()) } - assertBool "Checking if generate function is printing anything" (length output > 0) +generateTest = + testGroup + "Generate tests" + [ testCase "generate stream" $ do + output <- + capture $ + Ollama.generate + defaultGenerateOps + { modelName = "llama3.2" + , prompt = "what is 4 + 2?" + , stream = Just (T.putStr . Ollama.response_, pure ()) + } + assertBool "Checking if generate function is printing anything" (length output > 0) , testCase "Generate non-stream" $ do - eRes <- Ollama.generate defaultGenerateOps { modelName = "smollm:360m" - , prompt = "what is 4 + 2?" } - assertBool "Checking if generate function returns a valid value" (isRight eRes) + eRes <- + Ollama.generate + defaultGenerateOps + { modelName = "llama3.2" + , prompt = "what is 4 + 2?" + } + assertBool "Checking if generate function returns a valid value" (isRight eRes) + , testCase "Generate with invalid host" $ do + eRes <- Ollama.generate defaultGenerateOps { + modelName = "llama3.2" + , prompt = "what is 23 + 9?" + , hostUrl = pure "http://some-site" + , responseTimeOut = pure 2 + } + print ("got response" :: String ,eRes) + assertBool "Expecting Left" (isLeft eRes) , testCase "Generate with invalid model" $ do - eRes <- Ollama.generate defaultGenerateOps { modelName = "invalid-model" } - assertBool "Expecting generation to fail with invalid model" (isLeft eRes) - ] + eRes <- Ollama.generate defaultGenerateOps {modelName = "invalid-model"} + assertBool "Expecting generation to fail with invalid model" (isLeft eRes) + ] chatTest :: TestTree -chatTest = testGroup "Chat tests" [ - testCase "chat stream" $ do - let msg = Ollama.Message User "What is 29 + 3?" Nothing - defaultMsg = Ollama.Message User "" Nothing - output <- capture $ Ollama.chat defaultChatOps { - Chat.chatModelName = "smollm:360m" - , Chat.messages = msg :| [] - , Chat.stream = Just (T.putStr . Chat.content . fromMaybe defaultMsg . Chat.message, pure ()) - } - assertBool "Checking if chat function is printing anything" (length output > 0) - +chatTest = + testGroup + "Chat tests" + [ testCase "chat stream" $ do + let msg = Ollama.Message User "What is 29 + 3?" Nothing + defaultMsg = Ollama.Message User "" Nothing + output <- + capture $ + Ollama.chat + defaultChatOps + { Chat.chatModelName = "llama3.2" + , Chat.messages = msg :| [] + , Chat.stream = Just (T.putStr . Chat.content . fromMaybe defaultMsg . Chat.message, pure ()) + } + assertBool "Checking if chat function is printing anything" (length output > 0) , testCase "Chat non-stream" $ do - let msg = Ollama.Message User "What is 29 + 3?" Nothing - eRes <- Ollama.chat defaultChatOps { - Chat.chatModelName = "smollm:360m" - , Chat.messages = msg :| [] - } - assertBool "Checking if chat function returns a valid value" (isRight eRes) - ] + let msg = Ollama.Message User "What is 29 + 3?" Nothing + eRes <- + Ollama.chat + defaultChatOps + { Chat.chatModelName = "llama3.2" + , Chat.messages = msg :| [] + } + assertBool "Checking if chat function returns a valid value" (isRight eRes) + , testCase "Chat invalid host url" $ do + let msg = Ollama.Message User "What is 29 + 3?" Nothing + eRes <- + Ollama.chat + defaultChatOps + { Chat.chatModelName = "llama3.2" + , Chat.messages = msg :| [] + , Chat.hostUrl = pure "some random value" + , Chat.responseTimeOut = pure 2 + } + assertBool "It should return Left" (isLeft eRes) + ] psTest :: TestTree -psTest = testGroup "PS test" [ - testCase "check ps" $ do - mRes <- Ollama.ps - assertBool "Check if ps returns anything" (isJust mRes) - ] +psTest = + testGroup + "PS test" + [ testCase "check ps" $ do + mRes <- Ollama.ps + assertBool "Check if ps returns anything" (isJust mRes) + ] showTest :: TestTree -showTest = testGroup "Show test" [ - testCase "check show" $ do - mRes <- Ollama.showModel "smollm:360m" - assertBool "Check if model exists or not" (isJust mRes) - ] +showTest = + testGroup + "Show test" + [ testCase "check show" $ do + mRes <- Ollama.showModel "llama3.2" + assertBool "Check if model exists or not" (isJust mRes) + ] main :: IO () -main = defaultMain tests +main = do + mRes <- Ollama.list + case mRes of + Nothing -> pure () -- Ollama is likely not running. Not running tests. + Just _ -> defaultMain tests