diff --git a/Llama.hs b/Llama.hs --- a/Llama.hs +++ b/Llama.hs @@ -1,15 +1,18 @@ -{-# LANGUAGE OverloadedStrings, DeriveGeneric #-} +{-# LANGUAGE OverloadedStrings, DeriveGeneric, DuplicateRecordFields #-} module Llama where -import Control.Monad.IO.Class (liftIO) +import Conduit import Data.Aeson import Data.Text (Text) import GHC.Generics import Network.HTTP.Conduit +import Network.HTTP.Simple hiding (httpLbs) import Network.HTTP.Types.Status import System.IO (hPutStrLn, stderr) +import Llama.Streaming + data Role = System | User | CustomRole Text deriving Show instance ToJSON Role where toJSON System = "system" @@ -25,7 +28,7 @@ toJSON m = object [ "role" .= role m - , "content" .= content m + , "content" .= Llama.content m ] newtype LlamaApplyTemplateRequest = LlamaApplyTemplateRequest @@ -33,24 +36,30 @@ } deriving (Show, Generic) instance ToJSON LlamaApplyTemplateRequest -data Health = HealthOk | HealthNok deriving (Show) +newtype LlamaApplyTemplateResponse = LlamaApplyTemplateResponse + { prompt :: Text + } deriving (Show, Generic) ---data LlamaApplyTemplateResponse = LlamaApplyTemplateResponse --- { prompt :: Text --- } deriving (Show) +instance FromJSON LlamaApplyTemplateResponse where + parseJSON = withObject "LlamaApplyTemplateResponse" $ \v -> LlamaApplyTemplateResponse + <$> v .: "prompt" +data Health = HealthOk | HealthNok deriving (Show) + -- Llama request and response -newtype LlamaRequest = LlamaRequest +data LlamaRequest = LlamaRequest { prompt :: Text + , stream :: Bool } deriving (Show) instance FromJSON LlamaRequest where parseJSON = withObject "LlamaRequest" $ \v -> LlamaRequest <$> v .: "prompt" + <*> v .: "stream" instance ToJSON LlamaRequest where - toJSON (LlamaRequest p) = - object ["prompt" .= p] + toJSON (LlamaRequest p s) = + object ["prompt" .= p, "stream" .= s] newtype LlamaResponse = LlamaResponse { generatedText :: Text @@ -72,7 +81,7 @@ } response <- httpLbs req manager case decode (responseBody response) of - Just (LlamaRequest text) -> return (Just text) + Just (LlamaApplyTemplateResponse text) -> return (Just text) Nothing -> do liftIO $ hPutStrLn stderr "Failed to decode Llama response" return Nothing @@ -81,7 +90,7 @@ sendToLlama :: URL -> Manager -> Text -> IO (Maybe Text) sendToLlama url manager input = do let request = parseRequest_ $ url ++ "/completion" - body = encode (LlamaRequest input) + body = encode (LlamaRequest input False) req = request { method = "POST" , requestBody = RequestBodyLBS body , requestHeaders = [("Content-Type", "application/json")] @@ -93,6 +102,16 @@ liftIO $ hPutStrLn stderr "Failed to decode Llama response" return Nothing +sendToLlamaStreaming :: (MonadThrow m, MonadResource m) => URL -> Manager -> Text -> IO (ConduitT () LlamaStreamingResponse m ()) +sendToLlamaStreaming url manager input = do + let request = setRequestManager manager $ parseRequest_ $ url ++ "/completion" + body = encode (LlamaRequest input True) + req = request { method = "POST" + , requestBody = RequestBodyLBS body + , requestHeaders = [("Content-Type", "application/json")] + } + pure $ httpSource req getResponseBody .| eventConduit + llama :: URL -> Text -> IO (Maybe Text) llama url input = do manager <- liftIO $ newManager tlsManagerSettings { managerResponseTimeout = responseTimeoutNone } @@ -105,6 +124,14 @@ case res of Just text -> sendToLlama url manager text _ -> pure Nothing + +llamaTemplatedStreaming :: (MonadThrow m, MonadResource m) => URL -> LlamaApplyTemplateRequest -> IO (ConduitT () LlamaStreamingResponse m ()) +llamaTemplatedStreaming url input = do + manager <- liftIO $ newManager tlsManagerSettings { managerResponseTimeout = responseTimeoutNone } + res <- applyTemplate url manager input + case res of + Just text -> sendToLlamaStreaming url manager text + _ -> pure $ yieldMany [] health :: URL -> IO Health health url = do diff --git a/Llama/Streaming.hs b/Llama/Streaming.hs new file mode 100644 --- /dev/null +++ b/Llama/Streaming.hs @@ -0,0 +1,100 @@ +{- +Copyright Kadena LLC (c) 2019 + +All rights reserved. + +Redistribution and use in source and binary forms, with or without +modification, are permitted provided that the following conditions are met: + + * Redistributions of source code must retain the above copyright + notice, this list of conditions and the following disclaimer. + + * Redistributions in binary form must reproduce the above + copyright notice, this list of conditions and the following + disclaimer in the documentation and/or other materials provided + with the distribution. + + * Neither the name of Author name here nor the names of other + contributors may be used to endorse or promote products derived + from this software without specific prior written permission. + +THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS +"AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT +LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR +A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT +OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, +SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT +LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, +DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY +THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT +(INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. +-} + +{-# LANGUAGE OverloadedStrings, DeriveGeneric #-} + +module Llama.Streaming + ( eventConduit + , LlamaStreamingResponse(..) + ) +where + +import Control.Applicative +import Control.Monad.Catch +import Data.Aeson +import Data.Attoparsec.ByteString +import Data.Attoparsec.ByteString.Char8 as AC8 +import Data.ByteString as B +import Data.Conduit +import Data.Conduit.Attoparsec +import Data.Text +import GHC.Generics + +data ServerEvent = ServerEvent + { eventName :: Maybe ByteString + , eventId :: Maybe ByteString + , eventData :: [ByteString] + } | CommentEvent + { eventComment :: ByteString + } | RetryEvent + { eventRetry :: Int + } | CloseEvent + deriving (Show) + +data LlamaStreamingResponse = LlamaStreamingResponse + { index :: Int + , content :: Text + , stop :: Bool + , tokens_predicted :: Int + , tokens_evaluated :: Int + } deriving (Show, Generic) +instance FromJSON LlamaStreamingResponse + +eventConduit :: (MonadThrow m) => ConduitT ByteString LlamaStreamingResponse m () +eventConduit = mapOutputMaybe parseEventJSON $ conduitParser event + +parseEventJSON :: (PositionRange, ServerEvent) -> Maybe LlamaStreamingResponse +parseEventJSON (_, ServerEvent _ _ [d]) = decodeStrict d +parseEventJSON _ = Nothing + +-- |text/event-stream parser +event :: Parser ServerEvent +event = (sevent <|> comment <|> retry) <* eol + +sevent :: Parser ServerEvent +sevent = ServerEvent + <$> optional (string "event" *> char ':' *> chars <* eol) + <*> optional (string "id" *> char ':' *> chars <* eol) + <*> many (string "data" *> char ':' *> chars <* eol) + +comment :: Parser ServerEvent +comment = CommentEvent <$> (char ':' *> chars <* eol) + +retry :: Parser ServerEvent +retry = RetryEvent <$> (string "retry:" *> decimal <* eol) + +chars :: Parser ByteString +chars = AC8.takeTill (== '\n') + +eol :: Parser Char +eol = char '\n' diff --git a/Main.hs b/Main.hs new file mode 100644 --- /dev/null +++ b/Main.hs @@ -0,0 +1,49 @@ +{-# LANGUAGE DataKinds #-} +{-# LANGUAGE DeriveGeneric #-} +{-# LANGUAGE FlexibleInstances #-} +{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE StandaloneDeriving #-} +{-# LANGUAGE TypeOperators #-} + +module Main where + +import Conduit +import Data.Conduit.Lazy +import qualified Data.Text.IO as T +import Options.Generic +import System.Exit +import System.IO + +import Llama +import Llama.Streaming as LS + +data Options w = Options + { systemPrompt :: w ::: Text "The system prompt to use" "You are a helpful assistant." + , url :: w ::: String "llama-server URL" "http://localhost:8080" + , streaming :: w ::: Bool "use to stream output from the LLM" + } deriving (Generic) + +instance ParseRecord (Options Wrapped) where + parseRecord = parseRecordWithModifiers lispCaseModifiers +deriving instance Show (Options Unwrapped) + +main :: IO () +main = do + opts <- unwrapRecord "A command line interface for llama-server" + input <- T.getContents + let request = [ LlamaMessage System $ systemPrompt opts + , LlamaMessage User input + ] + case streaming opts of + False -> do + response <- llamaTemplated (url opts) (LlamaApplyTemplateRequest request) + case response of + Nothing -> T.hPutStrLn stderr "Got no response from the server." >> exitFailure + Just r -> T.putStrLn r + True -> do + hSetBuffering stdout NoBuffering + conduit <- llamaTemplatedStreaming (url opts) (LlamaApplyTemplateRequest request) + runResourceT $ do + list <- lazyConsume conduit + liftIO $ mapM_ (T.putStr . LS.content) $ list + T.putStrLn "" diff --git a/llama-cpp-haskell.cabal b/llama-cpp-haskell.cabal --- a/llama-cpp-haskell.cabal +++ b/llama-cpp-haskell.cabal @@ -1,7 +1,7 @@ cabal-version: 2.2 name: llama-cpp-haskell -version: 0.1.0.2 -synopsis: Haskell bindings for the llama.cpp llama-server +version: 0.1.1 +synopsis: Haskell bindings for the llama.cpp llama-server and a simple CLI -- description: license: AGPL-3.0-only license-file: LICENSE @@ -10,6 +10,7 @@ -- copyright: category: Text build-type: Simple +tested-with: GHC == 9.12.2 -- extra-source-files: Source-repository head @@ -19,19 +20,29 @@ Source-repository this type: git location: https://github.com/l29ah/llama-cpp-haskell.git - tag: 0.1.0.2 + tag: 0.1.1 common stuff ghc-options: -Wall - -library - import: stuff - exposed-modules: Llama - -- other-modules: - -- other-extensions: + default-language: Haskell2010 + other-modules: Llama.Streaming build-depends: base >= 4 && < 5 + , conduit ^>= 1.3.5 + , conduit-extra ^>= 1.3.7 + , exceptions ^>= 0.10.9 , http-conduit ^>= 2.3 , aeson ^>= 2.2 , text ^>= 2.1 , http-types ^>= 0.12 + , bytestring ^>= 0.12.2.0 + , attoparsec ^>= 0.14.4 + +library + import: stuff + exposed-modules: Llama default-language: Haskell2010 +executable llamacall + import: stuff + main-is: Main.hs + other-modules: Llama + build-depends: optparse-generic ^>= 1.5.2