369 lines
11 KiB
Haskell
369 lines
11 KiB
Haskell
{-# LANGUAGE OverloadedStrings #-}
|
|
{-# LANGUAGE DeriveGeneric #-}
|
|
{-# LANGUAGE ScopedTypeVariables #-}
|
|
|
|
-- UncloseAI Haskell Library
|
|
-- OpenAI-compatible API client with streaming support
|
|
-- Compatible with vLLM, Ollama, and OpenAI-compatible endpoints
|
|
|
|
import Network.HTTP.Simple
|
|
import Network.HTTP.Client (responseBody)
|
|
import Data.Aeson
|
|
import Data.Text (Text)
|
|
import qualified Data.Text as T
|
|
import qualified Data.Text.IO as TIO
|
|
import qualified Data.Text.Encoding as TE
|
|
import qualified Data.ByteString.Lazy as BL
|
|
import qualified Data.ByteString as BS
|
|
import GHC.Generics
|
|
import Control.Exception
|
|
import Control.Monad (unless)
|
|
import Control.Monad.IO.Class (liftIO)
|
|
import System.IO
|
|
import System.Environment (lookupEnv)
|
|
import Data.Maybe (fromMaybe, isJust)
|
|
import Data.Conduit
|
|
import qualified Data.Conduit.List as CL
|
|
import qualified Data.Conduit.Combinators as CC
|
|
|
|
-- Model info type
|
|
data ModelInfo = ModelInfo
|
|
{ modelId :: Text
|
|
, modelEndpoint :: String
|
|
, modelMaxTokens :: Int
|
|
} deriving (Show)
|
|
|
|
-- UncloseAI Client type
|
|
data UncloseAIClient = UncloseAIClient
|
|
{ clientModels :: [ModelInfo]
|
|
, clientTtsEndpoints :: [String]
|
|
, clientTimeout :: Int
|
|
} deriving (Show)
|
|
|
|
-- Message types
|
|
data ChatMessage = ChatMessage
|
|
{ role :: Text
|
|
, content :: Text
|
|
} deriving (Generic, Show)
|
|
|
|
instance ToJSON ChatMessage
|
|
|
|
data ChatRequest = ChatRequest
|
|
{ model :: Text
|
|
, messages :: [ChatMessage]
|
|
, max_tokens :: Int
|
|
, stream :: Maybe Bool
|
|
} deriving (Generic, Show)
|
|
|
|
instance ToJSON ChatRequest where
|
|
toJSON (ChatRequest m msgs mt s) = object $
|
|
[ "model" .= m
|
|
, "messages" .= msgs
|
|
, "max_tokens" .= mt
|
|
] ++ case s of
|
|
Just True -> ["stream" .= True]
|
|
_ -> []
|
|
|
|
data ChatResponse = ChatResponse
|
|
{ choices :: [Choice]
|
|
} deriving (Generic, Show)
|
|
|
|
data Choice = Choice
|
|
{ message :: ResponseMessage
|
|
} deriving (Generic, Show)
|
|
|
|
data ResponseMessage = ResponseMessage
|
|
{ respContent :: Text
|
|
} deriving (Generic, Show)
|
|
|
|
instance FromJSON ChatResponse
|
|
instance FromJSON Choice
|
|
instance FromJSON ResponseMessage where
|
|
parseJSON = withObject "ResponseMessage" $ \v ->
|
|
ResponseMessage <$> v .: "content"
|
|
|
|
data TTSRequest = TTSRequest
|
|
{ tts_model :: Text
|
|
, voice :: Text
|
|
, input :: Text
|
|
} deriving (Show)
|
|
|
|
instance ToJSON TTSRequest where
|
|
toJSON (TTSRequest m v i) = object
|
|
[ "model" .= m
|
|
, "voice" .= v
|
|
, "input" .= i
|
|
]
|
|
|
|
-- Models discovery types
|
|
data ModelData = ModelData
|
|
{ mdId :: Text
|
|
, mdMaxModelLen :: Maybe Int
|
|
} deriving (Generic, Show)
|
|
|
|
instance FromJSON ModelData where
|
|
parseJSON = withObject "ModelData" $ \v ->
|
|
ModelData
|
|
<$> v .: "id"
|
|
<*> v .:? "max_model_len"
|
|
|
|
data ModelsResponse = ModelsResponse
|
|
{ modelsData :: [ModelData]
|
|
} deriving (Generic, Show)
|
|
|
|
instance FromJSON ModelsResponse where
|
|
parseJSON = withObject "ModelsResponse" $ \v ->
|
|
ModelsResponse <$> v .: "data"
|
|
|
|
-- Initialize client with auto-discovery
|
|
initClient :: Int -> IO UncloseAIClient
|
|
initClient timeout = do
|
|
putStrLn "Initializing UncloseAI client..."
|
|
|
|
-- Discover chat/code models
|
|
models <- discoverModelsLoop 1 []
|
|
|
|
-- Discover TTS endpoints
|
|
ttsEndpoints <- discoverTtsLoop 1 []
|
|
|
|
putStrLn $ "Discovered " ++ show (length models) ++ " models, " ++
|
|
show (length ttsEndpoints) ++ " TTS endpoints\n"
|
|
|
|
return $ UncloseAIClient
|
|
{ clientModels = models
|
|
, clientTtsEndpoints = ttsEndpoints
|
|
, clientTimeout = timeout
|
|
}
|
|
|
|
-- Model discovery
|
|
discoverModelsFromEndpoint :: String -> IO [ModelInfo]
|
|
discoverModelsFromEndpoint ep = do
|
|
putStrLn $ "Endpoint: " ++ ep
|
|
|
|
result <- try $ do
|
|
request <- parseRequest $ "GET " ++ ep ++ "/models"
|
|
response <- httpLBS request
|
|
return $ getResponseBody response
|
|
|
|
case result of
|
|
Left (e :: SomeException) -> return []
|
|
Right body ->
|
|
case decode body :: Maybe ModelsResponse of
|
|
Nothing -> return []
|
|
Just modelsResp -> do
|
|
let modelsList = modelsData modelsResp
|
|
-- Filter out modelperm-* entries
|
|
filtered = filter (\md -> not $ T.isPrefixOf "modelperm-" (mdId md)) modelsList
|
|
mapM (\md -> do
|
|
let maxToks = fromMaybe 8192 (mdMaxModelLen md)
|
|
return $ ModelInfo (mdId md) ep maxToks
|
|
) filtered
|
|
|
|
discoverModelsLoop :: Int -> [ModelInfo] -> IO [ModelInfo]
|
|
discoverModelsLoop i acc | i > 9999 = return $ reverse acc
|
|
discoverModelsLoop i acc = do
|
|
maybeEndpoint <- lookupEnv $ "MODEL_ENDPOINT_" ++ show i
|
|
case maybeEndpoint of
|
|
Nothing -> return $ reverse acc
|
|
Just ep -> do
|
|
newModels <- discoverModelsFromEndpoint ep
|
|
discoverModelsLoop (i + 1) (reverse newModels ++ acc)
|
|
|
|
discoverTtsLoop :: Int -> [String] -> IO [String]
|
|
discoverTtsLoop i acc | i > 9999 = return $ reverse acc
|
|
discoverTtsLoop i acc = do
|
|
maybeEndpoint <- lookupEnv $ "TTS_ENDPOINT_" ++ show i
|
|
case maybeEndpoint of
|
|
Nothing -> return $ reverse acc
|
|
Just ep -> do
|
|
putStrLn $ "Discovering TTS from: " ++ ep
|
|
discoverTtsLoop (i + 1) (ep : acc)
|
|
|
|
-- Streaming response types
|
|
data StreamDelta = StreamDelta
|
|
{ deltaContent :: Maybe Text
|
|
} deriving (Generic, Show)
|
|
|
|
instance FromJSON StreamDelta where
|
|
parseJSON = withObject "StreamDelta" $ \v ->
|
|
StreamDelta <$> v .:? "content"
|
|
|
|
data StreamChoice = StreamChoice
|
|
{ delta :: StreamDelta
|
|
} deriving (Generic, Show)
|
|
|
|
instance FromJSON StreamChoice
|
|
|
|
data StreamChunk = StreamChunk
|
|
{ streamChoices :: [StreamChoice]
|
|
} deriving (Generic, Show)
|
|
|
|
instance FromJSON StreamChunk where
|
|
parseJSON = withObject "StreamChunk" $ \v ->
|
|
StreamChunk <$> v .: "choices"
|
|
|
|
-- Non-streaming chat completion
|
|
chat :: UncloseAIClient -> [ChatMessage] -> Maybe Int -> Maybe Int -> Maybe Double -> IO (Either String Text)
|
|
chat client msgs maybeModelIdx maybeMaxToks maybeTemp = do
|
|
let modelIdx = fromMaybe 0 maybeModelIdx
|
|
maxToks = fromMaybe 100 maybeMaxToks
|
|
temp = fromMaybe 0.7 maybeTemp
|
|
models = clientModels client
|
|
|
|
if modelIdx >= length models
|
|
then return $ Left "Invalid model index"
|
|
else do
|
|
let modelInfo = models !! modelIdx
|
|
let req = ChatRequest
|
|
{ model = modelId modelInfo
|
|
, messages = msgs
|
|
, max_tokens = maxToks
|
|
, stream = Nothing
|
|
}
|
|
|
|
result <- try $ do
|
|
request <- parseRequest $ "POST " ++ modelEndpoint modelInfo ++ "/chat/completions"
|
|
let request' = setRequestBodyJSON req request
|
|
response <- httpLBS request'
|
|
return $ getResponseBody response
|
|
|
|
case result of
|
|
Right body ->
|
|
case decode body :: Maybe ChatResponse of
|
|
Just resp ->
|
|
case choices resp of
|
|
(c:_) -> return $ Right $ respContent $ message c
|
|
[] -> return $ Left "No response choices"
|
|
Nothing -> return $ Left "Could not parse response"
|
|
Left (e :: SomeException) -> return $ Left $ show e
|
|
|
|
-- Streaming chat completion - yields content via IO action
|
|
chatStream :: UncloseAIClient -> [ChatMessage] -> Maybe Int -> Maybe Int -> Maybe Double -> IO (Either String ())
|
|
chatStream client msgs maybeModelIdx maybeMaxToks maybeTemp = do
|
|
let modelIdx = fromMaybe 0 maybeModelIdx
|
|
maxToks = fromMaybe 500 maybeMaxToks
|
|
temp = fromMaybe 0.7 maybeTemp
|
|
models = clientModels client
|
|
|
|
if modelIdx >= length models
|
|
then return $ Left "Invalid model index"
|
|
else do
|
|
let modelInfo = models !! modelIdx
|
|
let req = ChatRequest
|
|
{ model = modelId modelInfo
|
|
, messages = msgs
|
|
, max_tokens = maxToks
|
|
, stream = Just True
|
|
}
|
|
|
|
result <- try $ do
|
|
request <- parseRequest $ "POST " ++ modelEndpoint modelInfo ++ "/chat/completions"
|
|
let request' = setRequestBodyJSON req request
|
|
withResponse request' $ \response ->
|
|
runConduit $ getResponseBody response
|
|
.| CC.linesUnboundedAscii
|
|
.| CL.mapM_ processSSELine
|
|
|
|
case result of
|
|
Right () -> return $ Right ()
|
|
Left (e :: SomeException) -> return $ Left $ show e
|
|
|
|
-- Text-to-speech generation
|
|
tts :: UncloseAIClient -> Text -> Maybe Text -> Maybe String -> IO (Either String String)
|
|
tts client text maybeVoice maybeOutputFile = do
|
|
let voice = fromMaybe "alloy" maybeVoice
|
|
outputFile = fromMaybe "/tmp/speech.mp3" maybeOutputFile
|
|
ttsEndpoints = clientTtsEndpoints client
|
|
|
|
if null ttsEndpoints
|
|
then return $ Left "No TTS endpoints available"
|
|
else do
|
|
let endpoint = head ttsEndpoints
|
|
let req = TTSRequest
|
|
{ tts_model = "tts-1"
|
|
, voice = voice
|
|
, input = text
|
|
}
|
|
|
|
result <- try $ do
|
|
request <- parseRequest $ "POST " ++ endpoint ++ "/audio/speech"
|
|
let request' = setRequestBodyJSON req request
|
|
response <- httpLBS request'
|
|
let body = getResponseBody response
|
|
BL.writeFile outputFile body
|
|
return outputFile
|
|
|
|
case result of
|
|
Right file -> return $ Right file
|
|
Left (e :: SomeException) -> return $ Left $ show e
|
|
|
|
-- Process SSE line
|
|
processSSELine :: BS.ByteString -> IO ()
|
|
processSSELine line
|
|
| BS.isPrefixOf "data: " line = do
|
|
let dataStr = BS.drop 6 line
|
|
unless (dataStr == "[DONE]") $ do
|
|
case decode (BL.fromStrict dataStr) :: Maybe StreamChunk of
|
|
Just chunk ->
|
|
case streamChoices chunk of
|
|
(c:_) ->
|
|
case deltaContent (delta c) of
|
|
Just content -> TIO.putStr content >> hFlush stdout
|
|
Nothing -> return ()
|
|
[] -> return ()
|
|
Nothing -> return ()
|
|
| otherwise = return ()
|
|
|
|
-- Demo program showing library usage
|
|
main :: IO ()
|
|
main = do
|
|
hSetBuffering stdout NoBuffering
|
|
putStrLn "=== UncloseAI Haskell Client (with Streaming) ===\n"
|
|
|
|
-- Initialize client
|
|
client <- initClient 30
|
|
|
|
if null (clientModels client)
|
|
then do
|
|
putStrLn "ERROR: No models discovered"
|
|
else do
|
|
let models = clientModels client
|
|
let firstModel = head models
|
|
|
|
-- Non-streaming chat example
|
|
putStrLn "=== Non-Streaming Chat ==="
|
|
putStrLn $ "Model: " ++ T.unpack (modelId firstModel)
|
|
|
|
let messages = [ChatMessage "user" "Explain quantum computing in one sentence"]
|
|
result <- chat client messages Nothing Nothing Nothing
|
|
case result of
|
|
Right response -> putStrLn $ "Response: " ++ T.unpack response ++ "\n"
|
|
Left err -> putStrLn $ "Error: " ++ err ++ "\n"
|
|
|
|
-- Streaming chat example
|
|
let modelIdx = if length models >= 2 then 1 else 0
|
|
let streamModel = models !! modelIdx
|
|
|
|
putStrLn "=== Streaming Chat ==="
|
|
putStrLn $ "Model: " ++ T.unpack (modelId streamModel)
|
|
putStr "Response: "
|
|
|
|
let streamMessages = [ChatMessage "user" "Write a hello world program in Haskell"]
|
|
streamResult <- chatStream client streamMessages (Just modelIdx) Nothing Nothing
|
|
case streamResult of
|
|
Right () -> putStrLn "\n"
|
|
Left err -> putStrLn $ "\nError: " ++ err ++ "\n"
|
|
|
|
-- TTS example
|
|
if not (null (clientTtsEndpoints client))
|
|
then do
|
|
putStrLn "=== TTS Speech Generation ==="
|
|
putStrLn "Model: tts-1"
|
|
|
|
ttsResult <- tts client "Hello from UncloseAI Haskell client!" Nothing (Just "/tmp/speech.mp3")
|
|
case ttsResult of
|
|
Right file -> putStrLn $ "Audio saved to " ++ file
|
|
Left err -> putStrLn $ "TTS failed: " ++ err
|
|
else return ()
|
|
|
|
putStrLn "\n=== Examples Complete ==="
|