diff --git a/CHANGELOG.md b/CHANGELOG.md new file mode 100644 --- /dev/null +++ b/CHANGELOG.md @@ -0,0 +1,6 @@ +# Changelog for langchain-hs-openai + +## 0.0.6.0 - 2026-10-02 + +- Initial release. Split from the monolithic `langchain-hs` package. +- Provides `Langchain.Provider.OpenAI` and `Langchain.Embeddings.OpenAI`. diff --git a/LICENSE b/LICENSE new file mode 100644 --- /dev/null +++ b/LICENSE @@ -0,0 +1,20 @@ +Copyright (c) 2025 Tushar Adhatrao + +Permission is hereby granted, free of charge, to any person obtaining +a copy of this software and associated documentation files (the +"Software"), to deal in the Software without restriction, including +without limitation the rights to use, copy, modify, merge, publish, +distribute, sublicense, and/or sell copies of the Software, and to +permit persons to whom the Software is furnished to do so, subject to +the following conditions: + +The above copyright notice and this permission notice shall be included +in all copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, +EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF +MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. +IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY +CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, +TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE +SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. diff --git a/README.md b/README.md new file mode 100644 --- /dev/null +++ b/README.md @@ -0,0 +1,32 @@ +# langchain-hs-openai + +OpenAI ChatModel provider and embeddings for [langchain-hs](https://hackage.haskell.org/package/langchain-hs). + +Supports OpenAI and any OpenAI-compatible endpoint (OpenRouter, Fireworks, +Together AI). Does not depend on Ollama or Gemini packages. + +## Installation + +```yaml +# stack.yaml extra-deps +- langchain-hs-openai-0.0.6.0 +``` + +```yaml +# package.yaml / .cabal +dependencies: + - langchain-hs-openai +``` + +## Usage + +```haskell +import Langchain.Provider.OpenAI +import Control.Monad.Except (runExceptT) + +main :: IO () +main = do + let model = newOpenAI "sk-..." "gpt-4o" + res <- runExceptT $ invoke model [userMessage "Hello"] Nothing + print res +``` diff --git a/langchain-hs-openai.cabal b/langchain-hs-openai.cabal new file mode 100644 --- /dev/null +++ b/langchain-hs-openai.cabal @@ -0,0 +1,104 @@ +cabal-version: 1.12 + +-- This file has been generated from package.yaml by hpack version 0.39.6. +-- +-- see: https://github.com/sol/hpack + +name: langchain-hs-openai +version: 0.0.6.0 +synopsis: OpenAI provider and embeddings for langchain-hs +description: OpenAI ChatModel provider, streaming support, structured output, and text embeddings integration for langchain-hs. Compatible with OpenAI and OpenAI-compatible endpoints (OpenRouter, Fireworks, Together AI). +category: AI +homepage: https://github.com/tusharad/langchain-hs#readme +bug-reports: https://github.com/tusharad/langchain-hs/issues +author: Tushar Adhatrao +maintainer: tusharadhatrao@gmail.com +copyright: 2025-2026 Tushar Adhatrao +license: MIT +license-file: LICENSE +build-type: Simple +tested-with: + GHC == 9.12.4 + , GHC == 9.10.3 + , GHC == 9.8.4 +extra-source-files: + README.md + CHANGELOG.md + +source-repository head + type: git + location: https://github.com/tusharad/langchain-hs + +library + exposed-modules: + Langchain.Provider.OpenAI + Langchain.Embeddings.OpenAI + other-modules: + Paths_langchain_hs_openai + hs-source-dirs: + src + ghc-options: -Wall -Wcompat -Widentities -Wincomplete-record-updates -Wincomplete-uni-patterns -Wmissing-export-lists -Wmissing-home-modules -Wpartial-fields -Wredundant-constraints + build-depends: + aeson >=2.0 && <3 + , async ==2.2.* + , base >=4.17 && <5 + , bytestring >=0.10 && <0.13 + , conduit ==1.3.* + , containers >=0.6 && <0.9 + , http-client ==0.7.* + , http-client-tls >=0.3 && <0.5 + , http-conduit ==2.* + , http-types >=0.11 && <0.13 + , langchain-hs-core ==0.0.6.* + , mtl >=2.2 && <2.4 + , openai >=2.2.1 && <3 + , servant >=0.19 && <0.21 + , servant-client >=0.19 && <0.21 + , servant-client-core >=0.19 && <0.21 + , servant-conduit >=0.15 && <0.17 + , servant-event-stream ==0.4.* + , text >=1.2 && <3 + , transformers >=0.5 && <0.7 + , vector >=0.12 && <0.14 + default-language: Haskell2010 + +test-suite langchain-hs-openai-test + type: exitcode-stdio-1.0 + main-is: Spec.hs + other-modules: + Test.Langchain.Provider.OpenAI + Test.Langchain.Provider.TestSseServer + Paths_langchain_hs_openai + hs-source-dirs: + test + ghc-options: -Wall -Wcompat -Widentities -Wincomplete-record-updates -Wincomplete-uni-patterns -Wmissing-export-lists -Wmissing-home-modules -Wpartial-fields -Wredundant-constraints -threaded -rtsopts -with-rtsopts=-N + build-depends: + aeson + , async ==2.2.* + , base >=4.17 && <5 + , bytestring + , conduit ==1.3.* + , containers + , http-client ==0.7.* + , http-client-tls >=0.3 && <0.5 + , http-conduit ==2.* + , http-types >=0.11 && <0.13 + , langchain-hs-core + , langchain-hs-openai + , mtl >=2.2 && <2.4 + , openai + , resourcet + , servant >=0.19 && <0.21 + , servant-client >=0.19 && <0.21 + , servant-client-core >=0.19 && <0.21 + , servant-conduit >=0.15 && <0.17 + , servant-event-stream ==0.4.* + , stm + , tasty + , tasty-hunit + , text + , transformers >=0.5 && <0.7 + , vector >=0.12 && <0.14 + , wai + , warp + default-language: Haskell2010 diff --git a/src/Langchain/Embeddings/OpenAI.hs b/src/Langchain/Embeddings/OpenAI.hs new file mode 100644 --- /dev/null +++ b/src/Langchain/Embeddings/OpenAI.hs @@ -0,0 +1,221 @@ +{-# LANGUAGE DeriveGeneric #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE RecordWildCards #-} +{-# LANGUAGE ScopedTypeVariables #-} + +{- | +Module : Langchain.Embeddings.OpenAI +Description : OpenAI integration for text embeddings in LangChain Haskell +Copyright : (c) 2025-2026 Tushar Adhatrao +License : MIT +Maintainer : Tushar Adhatrao +Stability : experimental + +OpenAI implementation of LangChain's embedding interface. +-} +module Langchain.Embeddings.OpenAI + ( OpenAIEmbeddings (..) + , defaultOpenAIEmbeddings + , textEmbedding3Small + , textEmbedding3Large + , textEmbeddingAda + , EncodingFormat (..) + ) where + +import Control.Exception (SomeException, try) +import Control.Monad.Except (throwError) +import Control.Monad.IO.Class (liftIO) +import Data.Aeson +import qualified Data.ByteString.Lazy as LBS +import Data.Maybe +import Data.Text (Text, unpack) +import qualified Data.Text as T +import Data.Text.Encoding (encodeUtf8) +import qualified Data.Text.Lazy as TL +import qualified Data.Vector as V +import GHC.Generics + +import Langchain.DocumentLoader.Core +import Langchain.Embeddings.Core +import Langchain.Error (llmError) +import Network.HTTP.Conduit +import Network.HTTP.Simple + ( getResponseBody + , getResponseStatus + , setRequestBodyJSON + , setRequestHeader + , setRequestMethod + ) +import Network.HTTP.Types.Status (statusCode) + +-- Internal types for serialization of OpenAI request. +data EncodingFormat = FloatFormat | Base64Format + deriving (Eq, Show, Generic) + +data EmbeddingsInput = TextInput Text | TextList [Text] + deriving (Show, Eq) + +data OpenAIEmbeddingsRequest = OpenAIEmbeddingsRequest + { inputReq :: EmbeddingsInput + , modelReq :: Text + , dimensionsReq :: Maybe Int + , encodingFormatReq :: Maybe EncodingFormat + } + deriving (Show, Eq, Generic) + +instance ToJSON EncodingFormat where + toJSON FloatFormat = String "float" + toJSON Base64Format = String "base64" + +instance ToJSON EmbeddingsInput where + toJSON (TextInput t) = String t + toJSON (TextList t) = Array (V.fromList $ map String t) + +instance ToJSON OpenAIEmbeddingsRequest where + toJSON OpenAIEmbeddingsRequest {..} = + object $ + [ "input" .= inputReq + , "model" .= modelReq + ] + ++ catMaybes + [ ("dimensions" .=) <$> dimensionsReq + , ("encoding_format" .=) <$> encodingFormatReq + ] + +-- Response +data EmbeddingsUsage = EmbeddingsUsage + { promptTokens :: Int + , totalTokens :: Int + } + deriving (Eq, Show, Generic) + +data EmbeddingsObject = EmbeddingsObject + { embeddings :: [Float] + , index :: Maybe Int + , objectType :: Text + } + deriving (Eq, Show, Generic) + +data OpenAIEmbeddingsResponse = OpenAIEmbeddingsResponse + { objectTypeResp :: Text + , dataList :: [EmbeddingsObject] + , responseModel :: Text + , usage :: Maybe EmbeddingsUsage + } + deriving (Eq, Show, Generic) + +instance FromJSON EmbeddingsUsage where + parseJSON = withObject "EmbeddingsUsage" $ \v -> + EmbeddingsUsage + <$> v .: "prompt_tokens" + <*> v .: "total_tokens" + +instance FromJSON EmbeddingsObject where + parseJSON = withObject "EmbeddingsObject" $ \v -> + EmbeddingsObject + <$> v .: "embedding" + <*> v .:? "index" + <*> v .: "object" + +instance FromJSON OpenAIEmbeddingsResponse where + parseJSON = withObject "OpenAIEmbeddingsResponse" $ \v -> + OpenAIEmbeddingsResponse + <$> v .: "object" + <*> v .: "data" + <*> v .: "model" + <*> v .:? "usage" + +-- | Embeddings type for OpenAI +data OpenAIEmbeddings = OpenAIEmbeddings + { apiKey :: Text + , baseUrl :: Maybe String + , model :: Text + , dimensions :: Maybe Int + , encodingFormat :: Maybe EncodingFormat + , timeout :: Maybe Int + } + deriving (Eq, Generic) + +instance Show OpenAIEmbeddings where + show OpenAIEmbeddings {..} = "OpenAIEmbeddings " <> "model " <> unpack model + +openAIEmbeddingsRequest :: + OpenAIEmbeddings -> [Text] -> IO (Either String OpenAIEmbeddingsResponse) +openAIEmbeddingsRequest OpenAIEmbeddings {..} txts = do + eReq <- try $ parseRequest $ fromMaybe "https://api.openai.com/v1" baseUrl <> "/embeddings" + case eReq of + Left (err :: SomeException) -> pure $ Left $ "Invalid URL: " ++ show err + Right request_ -> do + manager <- + newManager + tlsManagerSettings + { managerResponseTimeout = + responseTimeoutMicro (fromMaybe 60 timeout * 1000000) + } + let req = + setRequestMethod "POST" $ + setRequestHeader "Content-Type" ["application/json"] $ + setRequestHeader "Authorization" ["Bearer " <> encodeUtf8 apiKey] $ + setRequestBodyJSON + ( OpenAIEmbeddingsRequest + { inputReq = TextList txts + , modelReq = model + , dimensionsReq = dimensions + , encodingFormatReq = encodingFormat + } + ) + request_ + + eResponse <- try (httpLbs req manager) :: IO (Either SomeException (Response LBS.ByteString)) + case eResponse of + Left err -> pure $ Left $ "Network error: " ++ show err + Right response -> do + let status = statusCode $ getResponseStatus response + if status >= 200 && status < 300 + then case eitherDecode (getResponseBody response) of + Left err -> return $ Left $ "JSON parse error: " <> err + Right completionResponse -> return $ Right completionResponse + else + return $ + Left $ + "API error: " + <> show status + <> " " + <> show (getResponseBody response) + +instance Embeddings OpenAIEmbeddings where + embedDocuments openAIEmbeddings docs = do + eRes <- liftIO $ openAIEmbeddingsRequest openAIEmbeddings (map (TL.toStrict . pageContent) docs) + case eRes of + Left err -> throwError $ llmError (T.pack err) (Just "OpenAIEmbeddings") Nothing + Right (OpenAIEmbeddingsResponse {..}) -> pure $ map embeddings dataList + + embedQuery openAIEmbeddings query = do + eRes <- liftIO $ openAIEmbeddingsRequest openAIEmbeddings [query] + case eRes of + Left err -> throwError $ llmError (T.pack err) (Just "OpenAIEmbeddings") Nothing + Right (OpenAIEmbeddingsResponse {..}) -> + case listToMaybe dataList of + Nothing -> throwError $ llmError "Embeddings are empty" (Just "OpenAIEmbeddings") Nothing + Just x -> pure $ embeddings x + +textEmbedding3Small :: Text +textEmbedding3Small = "text-embedding-3-small" + +textEmbedding3Large :: Text +textEmbedding3Large = "text-embedding-3-large" + +textEmbeddingAda :: Text +textEmbeddingAda = "text-embedding-ada-002" + +defaultOpenAIEmbeddings :: OpenAIEmbeddings +defaultOpenAIEmbeddings = + OpenAIEmbeddings + { apiKey = "" + , baseUrl = pure "https://api.openai.com/v1" + , model = textEmbedding3Small + , dimensions = Nothing + , encodingFormat = Nothing + , timeout = Nothing + } diff --git a/src/Langchain/Provider/OpenAI.hs b/src/Langchain/Provider/OpenAI.hs new file mode 100644 --- /dev/null +++ b/src/Langchain/Provider/OpenAI.hs @@ -0,0 +1,622 @@ +{-# LANGUAGE AllowAmbiguousTypes #-} +{-# LANGUAGE ConstraintKinds #-} +{-# LANGUAGE DataKinds #-} +{-# LANGUAGE DeriveAnyClass #-} +{-# LANGUAGE DeriveGeneric #-} +{-# LANGUAGE DuplicateRecordFields #-} +{-# LANGUAGE FlexibleInstances #-} +{-# LANGUAGE LambdaCase #-} +{-# LANGUAGE MultiParamTypeClasses #-} +{-# LANGUAGE NamedFieldPuns #-} +{-# LANGUAGE OverloadedLists #-} +{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE RecordWildCards #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE TypeApplications #-} +{-# LANGUAGE TypeFamilies #-} +{-# LANGUAGE TypeOperators #-} + +{- | +Module : Langchain.Provider.OpenAI +Description : OpenAI provider implementing effect-polymorphic ChatModel +Copyright : (c) 2025-2026 Tushar Adhatrao +License : MIT +Maintainer : Tushar Adhatrao +Stability : experimental + +OpenAI and OpenAICompatible provider using the @openai@ Haskell package +for typed API calls. Multi-modal content and streaming support. +-} +module Langchain.Provider.OpenAI + ( OpenAI (..) + , OpenAIConfig (..) + , defaultOpenAIConfig + , OpenAIToolChoice (..) + , openAITools + , newOpenAI + , openAICompatible + , normalizeBaseUrl + , OpenAIOptions (..) + , withOpenAIJsonSchema + , withStructuredOutput + , ToSchema + ) where + +import Control.Applicative ((<|>)) +import Control.Concurrent.Async (AsyncCancelled (..)) +import Control.Exception (SomeException, fromException, throwIO, try) +import Control.Monad.Except (throwError) +import Control.Monad.IO.Class (liftIO) +import Control.Monad.Trans.Class (lift) +import qualified Data.Conduit.Combinators as C + +import Data.Aeson (Value (..), object, (.=)) +import qualified Data.Aeson as Aeson +import qualified Data.Aeson.KeyMap as KeyMap +import Data.Aeson.Types (Parser, parseMaybe) +import Data.Bifunctor (first) +import qualified Data.ByteString.Lazy as LBS +import Data.Conduit +import qualified Data.List as List +import qualified Data.List.NonEmpty as NonEmpty +import qualified Data.Map.Strict as Map +import Data.Maybe (fromMaybe) +import Data.Proxy (Proxy (..)) +import Data.Text (Text) +import qualified Data.Text as T +import qualified Data.Text.Encoding as TE +import qualified Data.Vector as V +import GHC.Generics (Generic) + +import qualified OpenAI.V1 as OAI +import qualified OpenAI.V1.Chat.Completions as CC +import qualified OpenAI.V1.Models as OM +import qualified OpenAI.V1.ToolCall as OTC +import qualified OpenAI.V1.Usage as OU + +import qualified Data.Aeson.KeyMap as KM +import Langchain.Cache.Core (CacheableChatModel (..)) +import Langchain.Error (LangchainError, llmError) +import Langchain.Model +import Langchain.OutputParser.Structured (StructuredOutput (..)) +import Langchain.Stream (StreamEvent (..), StreamM, TokenUsage (..), callbackSource) +import Langchain.Tool (Tool, toolToValue) +import Langchain.Tool.Binding (ToolBinder (..)) +import Network.HTTP.Client (newManager) +import Network.HTTP.Client.TLS (tlsManagerSettings) +import Servant.API (Header, JSON, ReqBody, (:>)) +import Servant.API.EventStream + ( FromServerEvent (fromServerEvent) + , PostServerSentEvents + , ServerEvent (eventData) + , jsonData + ) +import Servant.Client.Core.BaseUrl (parseBaseUrl) +import Servant.Client.Streaming (ClientM, client, mkClientEnv, withClientM) +import Servant.Conduit () + +-- | Configuration for OpenAI provider +data OpenAIConfig = OpenAIConfig + { configApiKey :: Text + , configModel :: Text + , configBaseUrl :: Maybe Text + , configTemperature :: Maybe Double + } + deriving (Eq, Show, Generic, Aeson.ToJSON, Aeson.FromJSON) + +defaultOpenAIConfig :: Text -> OpenAIConfig +defaultOpenAIConfig key = OpenAIConfig key "gpt-4o" Nothing (Just 0.7) + +-- | OpenAI ChatModel implementation +data OpenAI = OpenAI + { apiKey :: Text + , model :: Text + , baseUrl :: Text + -- ^ Base URL (e.g. @"https://api.openai.com"@). The @openai@ package automatically appends @/v1/chat/completions@. + , temperature :: Maybe Double + } + deriving (Eq, Show) + +-- | Controls how OpenAI chooses among request tool definitions. +data OpenAIToolChoice + = OpenAIToolAuto + | OpenAIToolNone + | OpenAIToolRequired + | OpenAIToolFunction Text + deriving (Eq, Show) + +-- | Build stream request options from langchain tools. +openAITools :: [Tool m] -> OpenAIToolChoice -> Value +openAITools tools choice = + object + [ "tools" .= map toolToValue tools + , "tool_choice" .= toolChoiceValue choice + ] + +toolChoiceValue :: OpenAIToolChoice -> Value +toolChoiceValue OpenAIToolAuto = String "auto" +toolChoiceValue OpenAIToolNone = String "none" +toolChoiceValue OpenAIToolRequired = String "required" +toolChoiceValue (OpenAIToolFunction name) = + object + [ "type" .= ("function" :: Text) + , "function" .= object ["name" .= name] + ] + +data OpenAIStreamEvent + = OpenAIChunk OpenAIStreamChunk + | OpenAIDone + +instance FromServerEvent OpenAIStreamEvent where + fromServerEvent event + | eventData event == "[DONE]" = Right OpenAIDone + | otherwise = OpenAIChunk <$> jsonData event + +data OpenAIStreamChunk = OpenAIStreamChunk + { streamChoices :: [OpenAIStreamChoice] + , streamUsage :: Maybe TokenUsage + } + +instance Aeson.FromJSON OpenAIStreamChunk where + parseJSON = Aeson.withObject "OpenAIStreamChunk" $ \obj -> + OpenAIStreamChunk + <$> obj Aeson..:? "choices" Aeson..!= [] + <*> (obj Aeson..:? "usage" >>= traverse parseOpenAIStreamUsage) + +parseOpenAIStreamUsage :: Value -> Parser TokenUsage +parseOpenAIStreamUsage = Aeson.withObject "OpenAIStreamUsage" $ \obj -> + TokenUsage + <$> obj Aeson..: "prompt_tokens" + <*> obj Aeson..: "completion_tokens" + <*> obj Aeson..: "total_tokens" + +data OpenAIStreamChoice = OpenAIStreamChoice + { streamChoiceIndex :: Int + , streamChoiceDelta :: OpenAIStreamDelta + } + +instance Aeson.FromJSON OpenAIStreamChoice where + parseJSON = Aeson.withObject "OpenAIStreamChoice" $ \obj -> + OpenAIStreamChoice + <$> obj Aeson..: "index" + <*> obj Aeson..: "delta" + +data OpenAIStreamDelta = OpenAIStreamDelta + { streamContent :: Maybe Text + , streamToolCalls :: [OpenAIStreamToolCall] + } + +instance Aeson.FromJSON OpenAIStreamDelta where + parseJSON = Aeson.withObject "OpenAIStreamDelta" $ \obj -> + OpenAIStreamDelta + <$> obj Aeson..:? "content" + <*> obj Aeson..:? "tool_calls" Aeson..!= [] + +data OpenAIStreamToolCall = OpenAIStreamToolCall + { streamToolCallIndex :: Int + , streamToolCallId :: Maybe Text + , streamToolCallName :: Maybe Text + , streamToolCallArguments :: Maybe Text + } + +instance Aeson.FromJSON OpenAIStreamToolCall where + parseJSON = Aeson.withObject "OpenAIStreamToolCall" $ \obj -> do + streamToolCallIndex <- obj Aeson..: "index" + streamToolCallId <- obj Aeson..:? "id" + streamFunction <- obj Aeson..:? "function" + let streamToolCallName = streamFunction >>= parseMaybe (Aeson..: "name") + streamToolCallArguments = streamFunction >>= parseMaybe (Aeson..: "arguments") + pure + OpenAIStreamToolCall + { streamToolCallIndex + , streamToolCallId + , streamToolCallName + , streamToolCallArguments + } + +data PartialToolCall = PartialToolCall + { partialToolCallId :: Maybe Text + , partialToolCallName :: Maybe Text + , partialToolCallArguments :: Text + } + +type OpenAIStreamApi = + "v1" + :> "chat" + :> "completions" + :> Header "Authorization" Text + :> ReqBody '[JSON] Value + :> PostServerSentEvents (ConduitT () OpenAIStreamEvent IO ()) + +openAIStreamClient :: + Maybe Text -> Value -> ClientM (ConduitT () OpenAIStreamEvent IO ()) +openAIStreamClient = client (Proxy :: Proxy OpenAIStreamApi) + +streamRequestBody :: CC.CreateChatCompletion -> Maybe OpenAIOptions -> Value +streamRequestBody request options = case Aeson.toJSON request of + Object fields -> + Object $ + KeyMap.insert "stream_options" (object ["include_usage" Aeson..= True]) $ + KeyMap.insert "stream" (Bool True) $ + KeyMap.union fields optionFields + value -> value + where + optionFields = case options of + Just (OpenAIOptions (Object fields)) -> fields + _ -> mempty + +-- | Create standard OpenAI provider instance +newOpenAI :: + Text -> + -- | Api key + Text -> + -- | Model name + OpenAI +newOpenAI key mName = + OpenAI + { apiKey = key + , model = mName + , baseUrl = "https://api.openai.com" + , temperature = Nothing + } + +-- | Create OpenAICompatible provider instance for OpenRouter/Fireworks/Together. +openAICompatible :: + Text -> + -- | Api key + Text -> + -- \| The @endpoint@ should be the __base URL__ only (e.g. + -- @"https://openrouter.ai/api"@), not the full chat completions path. + -- The @openai@ package appends @/v1/chat/completions@ automatically. + -- + + -- | Model name + Text -> + OpenAI +openAICompatible key mName endpoint = + OpenAI + { apiKey = key + , model = mName + , baseUrl = endpoint + , temperature = Nothing + } + +-- --------------------------------------------------------------------------- +-- Conversion: langchain-hs Message -> openai package Message +-- --------------------------------------------------------------------------- + +-- | Convert a langchain 'ContentBlock' to an openai 'CC.Content'. +contentBlockToOAI :: ContentBlock -> CC.Content +contentBlockToOAI (TextBlock t) = CC.Text {CC.text = t} +contentBlockToOAI (ImageBlock ImageContent {imageSource = ImageUrl url}) = + CC.Image_URL {CC.image_url = CC.ImageURL {CC.url = url, CC.detail = Nothing}} +contentBlockToOAI (ImageBlock ImageContent {imageSource = ImageBase64 (Just mime) b64}) = + CC.Image_URL + { CC.image_url = + CC.ImageURL + { CC.url = "data:" <> mime <> ";base64," <> b64 + , CC.detail = Nothing + } + } +contentBlockToOAI (ImageBlock ImageContent {imageSource = ImageBase64 Nothing b64}) = + CC.Image_URL + { CC.image_url = + CC.ImageURL + { CC.url = "data:application/octet-stream;base64," <> b64 + , CC.detail = Nothing + } + } +contentBlockToOAI (AudioBlock _mime _b64) = + -- Audio blocks are represented as text placeholders in the request + CC.Text {CC.text = "[Audio content]"} +contentBlockToOAI (DataBlock _) = + CC.Text {CC.text = "[Data block]"} + +-- | Convert a langchain 'Message' to an openai package 'CC.Message'. +toLangchainOAIMessage :: Message -> CC.Message (V.Vector CC.Content) +toLangchainOAIMessage msg = + let contents = V.fromList $ map contentBlockToOAI (NonEmpty.toList (messageContents msg)) + in case messageRole msg of + System -> + CC.System {CC.content = contents, CC.name = messageName msg} + User -> + CC.User {CC.content = contents, CC.name = messageName msg} + Assistant -> + CC.Assistant + { CC.assistant_content = Just contents + , CC.refusal = Nothing + , CC.name = messageName msg + , CC.assistant_audio = Nothing + , CC.tool_calls = V.fromList . map toOAIToolCall <$> messageToolCalls msg + } + Tool -> + CC.Tool + { CC.content = contents + , CC.tool_call_id = fromMaybe "" (messageToolId msg) + } + -- Developer and Function map to System for the openai package + Developer -> + CC.System {CC.content = contents, CC.name = messageName msg} + Function -> + CC.System {CC.content = contents, CC.name = messageName msg} + where + toOAIToolCall toolCall = + OTC.ToolCall_Function + { OTC.id = toolCallId toolCall + , OTC.function = OTC.Function {OTC.name = toolCallName toolCall, OTC.arguments = arguments toolCall} + } + + arguments = TE.decodeUtf8 . LBS.toStrict . Aeson.encode . toolCallArguments + +-- --------------------------------------------------------------------------- +-- Conversion: openai package response -> langchain-hs Message +-- --------------------------------------------------------------------------- + +-- | Convert an openai package 'CC.Choice' response message to a langchain 'Message'. +fromOAIMessage :: CC.Message Text -> Message +fromOAIMessage oaiMsg = case oaiMsg of + CC.Assistant {CC.assistant_content, CC.tool_calls = oaiToolCalls, CC.name = nm} -> + let contentText = fromMaybe "" assistant_content + baseMsg = (assistantMessage contentText) {messageName = nm} + tcList = case oaiToolCalls of + Nothing -> Nothing + Just tcs -> + Just $ + map + ( \(OTC.ToolCall_Function {OTC.id = tcId, OTC.function = fn}) -> + let argVal = case Aeson.decode (LBS.fromStrict (TE.encodeUtf8 (OTC.arguments fn))) of + Just v -> v + Nothing -> object [] + in ToolCall tcId "function" (OTC.name fn) argVal + ) + (V.toList tcs) + in baseMsg {messageToolCalls = tcList} + CC.System {CC.content = c} -> systemMessage c + CC.User {CC.content = c} -> userMessage c + CC.Tool {CC.content = c} -> + textMessage Tool c + +-- | Convert openai 'OU.Usage' to langchain 'TokenUsage'. +fromOAIUsage :: OU.Usage ctd ptd -> TokenUsage +fromOAIUsage u = + TokenUsage + { promptTokens = fromIntegral (OU.prompt_tokens u) + , completionTokens = fromIntegral (OU.completion_tokens u) + , totalTokens = fromIntegral (OU.total_tokens u) + } + +newtype OpenAIOptions = OpenAIOptions {unOptions :: Value} + +-- --------------------------------------------------------------------------- +-- ChatModel instance +-- --------------------------------------------------------------------------- + +instance ChatModel OpenAI where + type ModelConfig OpenAI = OpenAIOptions + invoke provider inputMsgs mbOptions = do + resp <- liftIO $ first asText <$> try createComplention + case resp of + Left err -> throwError $ llmError' err + Right (CC.ChatCompletionObject {CC.choices = choicesVec, CC.usage = oaiUsage}) -> do + case V.toList choicesVec of + [] -> throwError $ llmError "Empty choices array in OpenAI response" Nothing Nothing + (choice : _) -> do + let respMsg = fromOAIMessage $ CC.message choice + _usage = fromOAIUsage oaiUsage + pure respMsg {messageToolCalls = messageToolCalls respMsg} + where + createComplention = do + methods <- getMethods + let baseBody = reqBody provider inputMsgs + body = mergeOptions baseBody mbOptions + OAI.createChatCompletion methods body + getMethods = do + clientEnv <- OAI.getClientEnv (normalizeBaseUrl (baseUrl provider)) + pure $ OAI.makeMethods clientEnv (apiKey provider) Nothing Nothing + + stream provider inputMsgs options = do + yield $ LLMStart rId (model provider) inputMsgs + + (accumulated, toolCalls, usage) <- + callbackSource openAIEvents + .| receiveChunks "" Map.empty Nothing + + yield $ LLMEnd rId ((assistantMessage accumulated) {messageToolCalls = toolCalls}) usage + where + rId = "openai-stream-run" + receiveChunks accumulated toolCalls usage = + await >>= \case + Nothing -> finishStream + Just (Left err) -> throwError $ llmError' err + Just (Right OpenAIDone) -> finishStream + Just (Right (OpenAIChunk OpenAIStreamChunk {streamChoices, streamUsage})) -> do + let (texts, nextToolCalls) = handleChoice $ choice0 streamChoices + nextUsage = streamUsage <|> usage + mapM_ (\text -> yield $ LLMChunk rId text Nothing) texts + receiveChunks (accumulated <> mconcat texts) nextToolCalls nextUsage + where + choice0 = List.find $ (== 0) . streamChoiceIndex + + handleChoice Nothing = ([], toolCalls) + handleChoice (Just OpenAIStreamChoice {streamChoiceDelta = OpenAIStreamDelta {streamContent, streamToolCalls}}) = + let nextToolCalls = List.foldl' addToolCall toolCalls streamToolCalls + in (maybe [] pure streamContent, nextToolCalls) + + addToolCall + toolCalls' + OpenAIStreamToolCall + { streamToolCallIndex + , streamToolCallId + , streamToolCallName + , streamToolCallArguments + } = + Map.alter (Just . update) streamToolCallIndex toolCalls' + where + update curr = + let prevId = curr >>= partialToolCallId + prevName = curr >>= partialToolCallName + prevArgs = maybe "" partialToolCallArguments curr + nextArgs = fromMaybe "" streamToolCallArguments + in PartialToolCall + { partialToolCallId = streamToolCallId <|> prevId + , partialToolCallName = streamToolCallName <|> prevName + , partialToolCallArguments = prevArgs <> nextArgs + } + + finishStream = do + let finalizeToolCalls = traverse toToolCall $ Map.elems toolCalls + finalToolCalls <- lift finalizeToolCalls + mapM_ (yield . LLMChunk rId "" . Just) finalToolCalls + pure (accumulated, nonEmpty finalToolCalls, usage) + where + nonEmpty [] = Nothing + nonEmpty xs = Just xs + + toToolCall :: PartialToolCall -> StreamM ToolCall + toToolCall PartialToolCall {partialToolCallId, partialToolCallName, partialToolCallArguments} = do + toolCallId <- + fromMaybeOrThrow "OpenAI stream ended with a tool call missing an id" partialToolCallId + toolCallName <- + fromMaybeOrThrow "OpenAI stream ended with a tool call missing a function name" partialToolCallName + toolCallArguments <- + either + (throwLlmError . ("Invalid JSON arguments in OpenAI tool call: " <>) . T.pack) + pure + (decode partialToolCallArguments) + pure ToolCall {toolCallId, toolCallType = "function", toolCallName, toolCallArguments} + + throwLlmError = throwError . llmError' + fromMaybeOrThrow err = maybe (throwLlmError err) pure + decode = Aeson.eitherDecode . LBS.fromStrict . TE.encodeUtf8 + + -- \| Internal function to handle streaming events from OpenAI. + openAIEvents emit = do + result <- try $ do + manager <- newManager tlsManagerSettings + let baseUrl' = T.unpack $ normalizeBaseUrl $ baseUrl provider + clientEnv <- mkClientEnv manager <$> parseBaseUrl baseUrl' + let body = reqBody provider inputMsgs + bearerToken = Just $ "Bearer " <> apiKey provider + request = openAIStreamClient bearerToken $ streamRequestBody body options + + withClientM request clientEnv $ \case + Left err -> + emit $ Left $ T.pack $ show err + Right source -> + runConduit $ + source .| C.mapM_ (emit . Right) + + case result of + Left err + | Just AsyncCancelled <- fromException err -> throwIO err + | otherwise -> emit $ Left $ asText err + Right () -> pure () + +instance CacheableChatModel OpenAI where + cacheModelIdentity OpenAI {..} _ = + Aeson.object + [ "provider" .= ("openai" :: Text) + , "model" .= model + , "baseUrl" .= baseUrl + , "temperature" .= temperature + ] + +{- | Inject a JSON Schema into an 'OpenAIOptions' config to request structured output. +Pass the result of 'outputSchema' or any 'Value' representing a JSON Schema. + +>>> let opts = withOpenAIJsonSchema (outputSchema (Proxy @MyType)) defaultOpts +-} +withOpenAIJsonSchema :: Value -> OpenAIOptions -> OpenAIOptions +withOpenAIJsonSchema schema (OpenAIOptions (Object configFields)) = + let strictSchema = case schema of + Object schemaFields -> + Object (KM.insert "additionalProperties" (Bool False) schemaFields) + value -> value + in OpenAIOptions . Object $ + KM.insert + "response_format" + ( object + [ "type" .= ("json_schema" :: Text) + , "json_schema" + .= object + [ "name" .= ("structured_output" :: Text) + , "strict" .= True + , "schema" .= strictSchema + ] + ] + ) + configFields +withOpenAIJsonSchema _ opts = opts + +{- | Type alias for types that can produce a JSON schema. +Alias for 'StructuredOutput' from 'Langchain.OutputParser.Structured'. +-} +type ToSchema = StructuredOutput + +-- | Inject a JSON schema derived from 'StructuredOutput' into 'OpenAIOptions'. +withStructuredOutput :: + forall a. + (StructuredOutput a) => + OpenAIOptions -> + OpenAIOptions +withStructuredOutput = withOpenAIJsonSchema (outputSchema (Proxy @a)) + +-- | Convert a 'SomeException' to 'Text' for error reporting. +asText :: SomeException -> Text +asText ex = T.pack $ show (ex :: SomeException) + +-- | Construct a 'LangchainError' for LLM errors with optional details. +llmError' :: Text -> LangchainError +llmError' msg = llmError msg Nothing Nothing + +-- | Merge option fields (tools, tool_choice, response_format, etc.) into a 'CreateChatCompletion'. +mergeOptions :: CC.CreateChatCompletion -> Maybe OpenAIOptions -> CC.CreateChatCompletion +mergeOptions body Nothing = body +mergeOptions body (Just (OpenAIOptions opts)) = + case Aeson.toJSON body of + Object baseFields -> + let optFields = case opts of + Object fs -> fs + _ -> mempty + merged = KeyMap.union optFields baseFields + in case Aeson.fromJSON (Object merged) of + Aeson.Success merged' -> merged' + _ -> body -- fallback: ignore unparseable options + _ -> body + +-- | Construct the request body for OpenAI chat completion. +reqBody :: OpenAI -> [Message] -> CC.CreateChatCompletion +reqBody provider inputMsgs = + CC._CreateChatCompletion + { CC.messages = toVec inputMsgs + , CC.model = OM.Model $ model provider + , CC.temperature = temperature provider + } + where + toVec = V.fromList . map toLangchainOAIMessage + +{- | Normalize base URL to ensure compatibility with the @openai@ package. + @openai@ package will always append `/v1/chat/completions` at the end. +-} +normalizeBaseUrl :: Text -> Text +normalizeBaseUrl rawUrl = + let u0 = T.dropWhileEnd (== '/') rawUrl + u1 + | "/v1/chat/completions" `T.isSuffixOf` u0 = T.dropEnd (T.length "/v1/chat/completions") u0 + | "/chat/completions" `T.isSuffixOf` u0 = T.dropEnd (T.length "/chat/completions") u0 + | "/v1" `T.isSuffixOf` u0 = T.dropEnd (T.length "/v1") u0 + | otherwise = u0 + in T.dropWhileEnd (== '/') u1 + +-- | Bind tools to an OpenAI model by merging tool definitions into the options Value +instance ToolBinder OpenAI m where + bindToolsConfig tools mbOpts = + case tools of + [] -> mbOpts + _ -> + let toolsVal = openAITools tools OpenAIToolAuto + in Just $ case (mbOpts, toolsVal) of + (Nothing, v) -> OpenAIOptions v + (Just (OpenAIOptions (Object existing)), Object newFields) -> + OpenAIOptions . Object $ KeyMap.union newFields existing + (Just existing, _) -> existing diff --git a/test/Spec.hs b/test/Spec.hs new file mode 100644 --- /dev/null +++ b/test/Spec.hs @@ -0,0 +1,17 @@ +{-# LANGUAGE OverloadedStrings #-} + +module Main (main) where + +import qualified Test.Langchain.Provider.OpenAI as OpenAIProviderTest +import Test.Tasty + +main :: IO () +main = + defaultMain $ + testGroup + "langchain-hs-openai Test Suite" + [ testGroup + "Provider Tests" + [ OpenAIProviderTest.tests + ] + ] diff --git a/test/Test/Langchain/Provider/OpenAI.hs b/test/Test/Langchain/Provider/OpenAI.hs new file mode 100644 --- /dev/null +++ b/test/Test/Langchain/Provider/OpenAI.hs @@ -0,0 +1,416 @@ +{-# LANGUAGE OverloadedStrings #-} + +module Test.Langchain.Provider.OpenAI (tests) where + +import Control.Concurrent (newEmptyMVar, putMVar, takeMVar) +import Control.Concurrent.Async (async, poll, wait) +import Control.Concurrent.STM + ( atomically + , modifyTVar' + , newTVarIO + , readTVarIO + ) +import Control.Monad (forM, void) +import Control.Monad.Except (runExceptT) +import Control.Monad.IO.Class (liftIO) +import Control.Monad.Trans.Resource (runResourceT) +import Data.Aeson (Value) +import qualified Data.Aeson as Aeson +import qualified Data.Aeson.KeyMap as KeyMap +import qualified Data.ByteString.Lazy as LBS +import Data.Conduit (await, runConduit, (.|)) +import qualified Data.Conduit.Combinators as C +import Data.Maybe (fromMaybe, isJust, isNothing) +import qualified Data.Text as T +import qualified Data.Vector as V +import Network.HTTP.Types (status500) +import Network.Wai (Application, responseLBS) +import System.Environment (lookupEnv) +import System.Timeout (timeout) +import Test.Tasty +import Test.Tasty.HUnit + +import Langchain.Error (LangchainError) +import Langchain.Model +import Langchain.Stream (StreamEvent (..), TokenUsage (..), collectEvents) +import Langchain.Tool (Tool, createTool, toolToValue) +import qualified Langchain.Tool as CoreTool + +import Langchain.Provider.OpenAI +import Test.Langchain.Provider.TestSseServer + ( cancellationAwareSseServer + , capturingRawSseServer + , collectModelStream + , gatedSseServer + , rawSseServer + , sseFrame + , withTestApplication + ) + +withErrorProvider :: (OpenAI -> IO a) -> IO a +withErrorProvider action = + withTestApplication errorServer $ \url -> withOpenAIProvider url action + +withRawTestProvider :: [LBS.ByteString] -> (OpenAI -> IO a) -> IO a +withRawTestProvider frames action = + withTestApplication (rawSseServer frames) $ \url -> withOpenAIProvider url action + +withRequestCapturingProvider :: (Maybe Value -> IO ()) -> (OpenAI -> IO a) -> IO a +withRequestCapturingProvider captureRequest action = + withTestApplication (capturingRawSseServer (captureRequest . Aeson.decode) [sseFrame "[DONE]"]) $ \url -> + withOpenAIProvider url action + +withCancellationAwareProvider :: IO () -> (OpenAI -> IO a) -> IO a +withCancellationAwareProvider signalClientClosed action = + withTestApplication (cancellationAwareSseServer (sseFrame $ chunk "Hello") signalClientClosed) $ \url -> + withOpenAIProvider url action + +withGatedProvider :: IO () -> (OpenAI -> IO a) -> IO a +withGatedProvider waitForContinuation action = + withTestApplication + ( gatedSseServer + (sseFrame $ chunk "Hel") + waitForContinuation + [sseFrame (chunk "lo"), sseFrame "[DONE]"] + ) + $ \url -> withOpenAIProvider url action + +withOpenAIProvider :: T.Text -> (OpenAI -> IO a) -> IO a +withOpenAIProvider url action = + action $ (newOpenAI "test-key" "test-model") {baseUrl = url} + +errorServer :: Application +errorServer _request respond = respond $ responseLBS status500 [] "" + +collectRawStream :: [LBS.ByteString] -> IO (Either LangchainError [StreamEvent]) +collectRawStream frames = + withRawTestProvider frames $ \provider -> + collectModelStream provider [userMessage "Hello"] Nothing + +chunk :: LBS.ByteString -> LBS.ByteString +chunk content = + "{\"id\":\"chatcmpl-test\",\"object\":\"chat.completion.chunk\",\"created\":0,\"model\":\"test-model\",\"choices\":[{\"index\":0,\"delta\":{\"content\":\"" + <> content + <> "\"},\"finish_reason\":null}]}" + +emptyChoices :: LBS.ByteString +emptyChoices = + "{\"id\":\"chatcmpl-test\",\"object\":\"chat.completion.chunk\",\"created\":0,\"model\":\"test-model\",\"choices\":[]}" + +tests :: TestTree +tests = + testGroup + "Langchain.Provider.OpenAI" + [ testCase "newOpenAI initializes default provider" $ do + let p = newOpenAI "sk-test" "gpt-4o" + model p @?= "gpt-4o" + baseUrl p @?= "https://api.openai.com" + , testCase "openAICompatible initializes custom endpoint" $ do + let p = openAICompatible "sk-test" "custom-llm" "https://custom-ai.example.com" + model p @?= "custom-llm" + baseUrl p @?= "https://custom-ai.example.com" + , testCase "live OpenAI stream emits text and usage" $ do + mbApiKey <- lookupEnv "OPENAI_API_KEY" + case mbApiKey of + Nothing -> putStrLn " [SKIPPED] OPENAI_API_KEY is not set" + Just envApiKey -> do + envModel <- fromMaybe "gpt-4o-mini" <$> lookupEnv "OPENAI_STREAM_TEST_MODEL" + result <- + timeout 60000000 $ + runResourceT $ + runExceptT $ + collectEvents $ + stream + (newOpenAI (T.pack envApiKey) (T.pack envModel)) + [userMessage "Reply with exactly OK."] + Nothing + case result of + Nothing -> assertFailure "OpenAI stream timed out" + Just (Left err) -> assertFailure $ "Expected stream success, got: " ++ show err + Just (Right events) -> do + print events + case reverse events of + LLMEnd _ responseMessage (Just usage) : _ -> do + assertBool "Expected non-empty streamed text" $ not $ T.null $ extractMessageText responseMessage + assertBool "Expected positive total token usage" $ totalTokens usage > 0 + _ -> assertFailure $ "Expected LLMEnd with usage, got: " ++ show events + , testCase "live OpenAI stream invokes a tool and continues with its result" $ do + mbApiKey <- lookupEnv "OPENAI_API_KEY" + case mbApiKey of + Nothing -> putStrLn " [SKIPPED] OPENAI_API_KEY is not set" + Just envApiKey -> do + envModel <- fromMaybe "gpt-4o-mini" <$> lookupEnv "OPENAI_STREAM_TEST_MODEL" + let weatherTool :: Tool IO + weatherTool = + createTool + "get_weather" + "Returns the current weather for a city." + ( Aeson.object + [ "type" Aeson..= ("object" :: T.Text) + , "properties" + Aeson..= Aeson.object + [ "city" Aeson..= Aeson.object ["type" Aeson..= ("string" :: T.Text)] + ] + , "required" Aeson..= ["city" :: T.Text] + , "additionalProperties" Aeson..= False + ] + ) + (const $ pure $ Right "The weather in Paris is sunny and 22 C.") + provider = newOpenAI (T.pack envApiKey) (T.pack envModel) + runLive messages config = + timeout 60000000 $ + runResourceT $ + runExceptT $ + collectEvents $ + stream provider messages config + prompt = userMessage "Use get_weather to look up the weather in Paris, then answer using the tool result." + + firstResult <- + runLive + [prompt] + (Just $ OpenAIOptions $ openAITools [weatherTool] (OpenAIToolFunction "get_weather")) + firstEvents <- case firstResult of + Nothing -> assertFailure "OpenAI tool-call stream timed out" >> fail "unreachable" + Just (Left err) -> assertFailure ("Expected tool-call stream success, got: " ++ show err) >> fail "unreachable" + Just (Right events) -> pure events + (assistant, toolCalls) <- case reverse firstEvents of + LLMEnd _ responseMessage _ : _ -> case messageToolCalls responseMessage of + Just calls@[toolCall] + | toolCallName toolCall == "get_weather" -> pure (responseMessage, calls) + _ -> assertFailure ("Expected OpenAI tool call, got: " ++ show firstEvents) >> fail "unreachable" + _ -> assertFailure ("Expected tool-call stream end, got: " ++ show firstEvents) >> fail "unreachable" + toolResults <- forM toolCalls $ \toolCall -> do + output <- CoreTool.toolExecute weatherTool (toolCallArguments toolCall) + case output of + Left err -> assertFailure ("Tool execution failed: " ++ show err) >> fail "unreachable" + Right text -> + pure $ + (textMessage Tool text) + { messageName = Just (toolCallName toolCall) + , messageToolId = Just (toolCallId toolCall) + } + secondResult <- runLive ([prompt, assistant] <> toolResults) Nothing + case secondResult of + Nothing -> assertFailure "OpenAI tool-result stream timed out" + Just (Left err) -> assertFailure $ "Expected tool-result stream success, got: " ++ show err + Just (Right events) -> case reverse events of + LLMEnd _ responseMessage (Just usage) : _ -> do + assertBool "Expected final text after tool result" $ + not $ + T.null $ + extractMessageText responseMessage + assertBool "Expected positive total token usage" $ totalTokens usage > 0 + _ -> assertFailure $ "Expected LLMEnd with usage, got: " ++ show events + , testCase "normalizeBaseUrl strips endpoint paths for servant compatibility" $ do + normalizeBaseUrl "https://api.openai.com" @?= "https://api.openai.com" + normalizeBaseUrl "https://api.openai.com/" @?= "https://api.openai.com" + normalizeBaseUrl "https://api.openai.com/v1" @?= "https://api.openai.com" + normalizeBaseUrl "https://api.openai.com/v1/" @?= "https://api.openai.com" + normalizeBaseUrl "https://api.openai.com/v1/chat/completions" @?= "https://api.openai.com" + normalizeBaseUrl "https://openrouter.ai/api" @?= "https://openrouter.ai/api" + normalizeBaseUrl "https://openrouter.ai/api/v1" @?= "https://openrouter.ai/api" + normalizeBaseUrl "https://openrouter.ai/api/v1/chat/completions" @?= "https://openrouter.ai/api" + normalizeBaseUrl "http://localhost:11434/v1" @?= "http://localhost:11434" + , testCase "stream emits chunks and ends at [DONE]" $ do + result <- collectRawStream [sseFrame $ chunk "Hel", sseFrame $ chunk "lo", sseFrame "[DONE]"] + case result of + Left err -> assertFailure $ "Expected stream success, got: " ++ show err + Right events -> case events of + [ LLMStart {} + , LLMChunk _ "Hel" Nothing + , LLMChunk _ "lo" Nothing + , LLMEnd _ responseMessage Nothing + ] -> extractMessageText responseMessage @?= "Hello" + _ -> assertFailure $ "Unexpected stream events: " ++ show events + , testCase "stream delivers a chunk before the response completes" $ do + firstChunkReceived <- newEmptyMVar + continueResponse <- newEmptyMVar + receivedEvents <- newTVarIO [] + withGatedProvider (takeMVar continueResponse) $ \provider -> do + consumer <- + async + . runResourceT + . runExceptT + . runConduit + $ stream provider [userMessage "Hello"] Nothing + .| C.mapM_ + ( \event -> do + liftIO . atomically $ modifyTVar' receivedEvents (event :) + case event of + LLMChunk _ "Hel" _ -> liftIO $ putMVar firstChunkReceived () + _ -> pure () + ) + received <- timeout 500000 $ takeMVar firstChunkReceived + assertBool "expected first chunk before releasing the response" $ isJust received + stillStreaming <- poll consumer + assertBool "consumer should wait for the remaining response" $ isNothing stillStreaming + putMVar continueResponse () + result <- timeout 500000 $ wait consumer + case result of + Nothing -> assertFailure "stream did not finish after releasing the response" + Just (Left err) -> assertFailure $ "Expected stream success, got: " ++ show err + Just (Right ()) -> do + events <- reverse <$> readTVarIO receivedEvents + case reverse events of + LLMEnd _ responseMessage Nothing : _ -> + extractMessageText responseMessage @?= "Hello" + _ -> assertFailure $ "Expected a completed stream, got: " ++ show events + , testCase "stream finishes when the SSE connection closes" $ do + result <- collectRawStream [sseFrame $ chunk "Hello"] + case result of + Left err -> assertFailure $ "Expected stream success, got: " ++ show err + Right events -> case events of + [LLMStart {}, LLMChunk _ "Hello" Nothing, LLMEnd _ responseMessage Nothing] -> + extractMessageText responseMessage @?= "Hello" + _ -> assertFailure $ "Unexpected stream events: " ++ show events + , testCase "stream ignores chunks without choices" $ do + result <- collectRawStream [sseFrame emptyChoices, sseFrame "[DONE]"] + case result of + Left err -> assertFailure $ "Expected stream success, got: " ++ show err + Right events -> case events of + [LLMStart {}, LLMEnd _ responseMessage Nothing] -> + extractMessageText responseMessage @?= "" + _ -> assertFailure $ "Unexpected stream events: " ++ show events + , testCase "stream converts malformed SSE data to LangchainError" $ do + result <- collectRawStream [sseFrame "not JSON"] + case result of + Left _ -> pure () + Right events -> assertFailure $ "Expected stream failure, got: " ++ show events + , testCase "stream converts HTTP errors to LangchainError" $ do + result <- withErrorProvider $ \provider -> + runResourceT $ runExceptT $ collectEvents (stream provider [userMessage "Hello"] Nothing) + case result of + Left _ -> pure () + Right events -> assertFailure $ "Expected stream failure, got: " ++ show events + , testCase "stream handles SSE frames written in multiple pieces" $ do + let frame = sseFrame $ chunk "Hello" + splitPoint = LBS.length frame `div` 2 + fragments = [LBS.take splitPoint frame, LBS.drop splitPoint frame, "data: [DONE]\n\n"] + result <- collectRawStream fragments + case result of + Left err -> assertFailure $ "Expected stream success, got: " ++ show err + Right events -> case events of + [LLMStart {}, LLMChunk _ "Hello" Nothing, LLMEnd _ responseMessage Nothing] -> + extractMessageText responseMessage @?= "Hello" + _ -> assertFailure $ "Unexpected stream events: " ++ show events + , testCase "stream requests usage in stream options" $ do + requestBody <- newEmptyMVar + withRequestCapturingProvider (putMVar requestBody) $ \provider -> do + void . runResourceT . runExceptT $ collectEvents (stream provider [userMessage "Hello"] Nothing) + mbRequest <- takeMVar requestBody + case mbRequest of + Nothing -> assertFailure "Expected JSON request body" + Just (Aeson.Object fields) -> do + KeyMap.lookup "stream" fields @?= Just (Aeson.Bool True) + KeyMap.lookup "stream_options" fields @?= Just (Aeson.object ["include_usage" Aeson..= True]) + Just request -> assertFailure $ "Expected JSON object, got: " ++ show request + , testCase "stream sends tool definitions and tool choice" $ do + let weatherTool :: Tool IO + weatherTool = createTool "get_weather" "Gets the weather" (Aeson.object []) (const $ pure $ Right "sunny") + config = openAITools [weatherTool] (OpenAIToolFunction "get_weather") + requestBody <- newEmptyMVar + withRequestCapturingProvider (putMVar requestBody) $ \provider -> do + void . runResourceT . runExceptT $ + collectEvents (stream provider [userMessage "Hello"] (Just $ OpenAIOptions config)) + mbRequest <- takeMVar requestBody + case mbRequest of + Just (Aeson.Object fields) -> do + KeyMap.lookup "tools" fields @?= Just (Aeson.toJSON [toolToValue weatherTool]) + KeyMap.lookup "tool_choice" fields + @?= Just + ( Aeson.object + [ "type" Aeson..= ("function" :: T.Text) + , "function" Aeson..= Aeson.object ["name" Aeson..= ("get_weather" :: T.Text)] + ] + ) + Just request -> assertFailure $ "Expected JSON object, got: " ++ show request + Nothing -> assertFailure "Expected JSON request body" + , testCase "stream sends assistant tool calls before tool results" $ do + let toolCall = + ToolCall + "call_weather" + "function" + "get_weather" + (Aeson.object ["city" Aeson..= ("Paris" :: T.Text)]) + assistant = (assistantMessage "") {messageToolCalls = Just [toolCall]} + toolResult = (textMessage Tool "Sunny") {messageToolId = Just "call_weather"} + requestBody <- newEmptyMVar + withRequestCapturingProvider (putMVar requestBody) $ \provider -> do + void . runResourceT . runExceptT $ + collectEvents (stream provider [userMessage "Weather?", assistant, toolResult] Nothing) + mbRequest <- takeMVar requestBody + case mbRequest of + Just (Aeson.Object fields) -> case KeyMap.lookup "messages" fields of + Just (Aeson.Array messages) -> case V.toList messages of + [_, Aeson.Object assistantFields, Aeson.Object toolResultFields] -> do + KeyMap.lookup "tool_calls" assistantFields + @?= Just + ( Aeson.toJSON + [ Aeson.object + [ "id" Aeson..= ("call_weather" :: T.Text) + , "type" Aeson..= ("function" :: T.Text) + , "function" + Aeson..= Aeson.object + [ "name" Aeson..= ("get_weather" :: T.Text) + , "arguments" Aeson..= ("{\"city\":\"Paris\"}" :: T.Text) + ] + ] + ] + ) + KeyMap.lookup "tool_call_id" toolResultFields @?= Just (Aeson.String "call_weather") + messages' -> assertFailure $ "Expected three request messages, got: " ++ show messages' + request -> assertFailure $ "Expected messages array, got: " ++ show request + Just request -> assertFailure $ "Expected JSON object, got: " ++ show request + Nothing -> assertFailure "Expected JSON request body" + , testCase "stream accumulates text, fragmented tool calls, and usage" $ do + let frames = + [ sseFrame + "{\"id\":\"chatcmpl-test\",\"object\":\"chat.completion.chunk\",\"created\":0,\"model\":\"test-model\",\"choices\":[{\"index\":0,\"delta\":{\"content\":\"Checking weather...\"},\"finish_reason\":null}]}" + , sseFrame + "{\"id\":\"chatcmpl-test\",\"object\":\"chat.completion.chunk\",\"created\":0,\"model\":\"test-model\",\"choices\":[{\"index\":0,\"delta\":{\"tool_calls\":[{\"index\":0,\"id\":\"call_1\",\"type\":\"function\",\"function\":{\"name\":\"get_weather\",\"arguments\":\"{\\\"city\\\":\\\"\"}}]},\"finish_reason\":null}]}" + , sseFrame + "{\"id\":\"chatcmpl-test\",\"object\":\"chat.completion.chunk\",\"created\":0,\"model\":\"test-model\",\"choices\":[{\"index\":0,\"delta\":{\"tool_calls\":[{\"index\":0,\"function\":{\"arguments\":\"Paris\\\"}\"}}]},\"finish_reason\":\"tool_calls\"}]}" + , sseFrame + "{\"id\":\"chatcmpl-test\",\"object\":\"chat.completion.chunk\",\"created\":0,\"model\":\"test-model\",\"choices\":[],\"usage\":{\"prompt_tokens\":7,\"completion_tokens\":5,\"total_tokens\":12}}" + , sseFrame "[DONE]" + ] + expectedToolCall = + ToolCall + { toolCallId = "call_1" + , toolCallType = "function" + , toolCallName = "get_weather" + , toolCallArguments = Aeson.object ["city" Aeson..= ("Paris" :: T.Text)] + } + expectedUsage = TokenUsage 7 5 12 + result <- collectRawStream frames + case result of + Left err -> assertFailure $ "Expected stream success, got: " ++ show err + Right events -> case events of + [ LLMStart {} + , LLMChunk _ "Checking weather..." Nothing + , LLMChunk _ "" (Just toolCall) + , LLMEnd _ responseMessage (Just usage) + ] -> do + toolCall @?= expectedToolCall + extractMessageText responseMessage @?= "Checking weather..." + messageToolCalls responseMessage @?= Just [expectedToolCall] + usage @?= expectedUsage + _ -> assertFailure $ "Unexpected stream events: " ++ show events + , testCase "stream rejects invalid completed tool arguments" $ do + let frames = + [ sseFrame + "{\"id\":\"chatcmpl-test\",\"object\":\"chat.completion.chunk\",\"created\":0,\"model\":\"test-model\",\"choices\":[{\"index\":0,\"delta\":{\"tool_calls\":[{\"index\":0,\"id\":\"call_1\",\"type\":\"function\",\"function\":{\"name\":\"get_weather\",\"arguments\":\"not-json\"}}]},\"finish_reason\":\"tool_calls\"}]}" + , sseFrame "[DONE]" + ] + result <- collectRawStream frames + case result of + Left _ -> pure () + Right events -> assertFailure $ "Expected stream failure, got: " ++ show events + , testCase "stream closes the SSE connection when the consumer stops after a chunk" $ do + clientClosed <- newEmptyMVar + withCancellationAwareProvider (putMVar clientClosed ()) $ \provider -> do + void . runResourceT . runExceptT . runConduit $ + stream provider [userMessage "Hello"] Nothing .| (await >> await) + closed <- timeout 500000 $ takeMVar clientClosed + assertBool "expected the SSE connection to close" $ isJust closed + ] diff --git a/test/Test/Langchain/Provider/TestSseServer.hs b/test/Test/Langchain/Provider/TestSseServer.hs new file mode 100644 --- /dev/null +++ b/test/Test/Langchain/Provider/TestSseServer.hs @@ -0,0 +1,85 @@ +{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE TypeFamilies #-} + +module Test.Langchain.Provider.TestSseServer + ( withTestApplication + , rawSseServer + , capturingRawSseServer + , capturingRawSseRequestServer + , sseFrame + , gatedSseServer + , cancellationAwareSseServer + , collectModelStream + ) where + +import Control.Concurrent (threadDelay) +import Control.Exception (SomeException, catch) +import Control.Monad.Except (runExceptT) +import Control.Monad.Trans.Resource (runResourceT) +import qualified Data.ByteString.Builder as Builder +import qualified Data.ByteString.Lazy as LBS +import qualified Data.Text as T +import Network.HTTP.Types (hContentType, status200) +import Network.Wai (Application, Request, responseStream, strictRequestBody) +import Network.Wai.Handler.Warp (testWithApplication) + +import Langchain.Error (LangchainError) +import Langchain.Model (ChatModel (..), Message) +import Langchain.Stream (StreamEvent, collectEvents) + +withTestApplication :: Application -> (T.Text -> IO a) -> IO a +withTestApplication app action = + testWithApplication (pure app) $ \port -> + action $ "http://127.0.0.1:" <> T.pack (show port) + +rawSseServer :: [LBS.ByteString] -> Application +rawSseServer frames _request respond = + respond $ + responseStream status200 [(hContentType, "text/event-stream")] $ \write flush -> + mapM_ (\frame -> write (Builder.lazyByteString frame) >> flush) frames + +capturingRawSseServer :: (LBS.ByteString -> IO ()) -> [LBS.ByteString] -> Application +capturingRawSseServer captureRequest frames request respond = do + captureRequest =<< strictRequestBody request + rawSseServer frames request respond + +capturingRawSseRequestServer :: + (Request -> LBS.ByteString -> IO ()) -> [LBS.ByteString] -> Application +capturingRawSseRequestServer captureRequest frames request respond = do + body <- strictRequestBody request + captureRequest request body + rawSseServer frames request respond + +sseFrame :: LBS.ByteString -> LBS.ByteString +sseFrame payload = "data: " <> payload <> "\n\n" + +gatedSseServer :: LBS.ByteString -> IO () -> [LBS.ByteString] -> Application +gatedSseServer firstFrame waitForContinuation remainingFrames _request respond = + respond $ + responseStream status200 [(hContentType, "text/event-stream")] $ \write flush -> do + write $ Builder.lazyByteString firstFrame + flush + waitForContinuation + mapM_ (write . Builder.lazyByteString) remainingFrames + flush + +cancellationAwareSseServer :: LBS.ByteString -> IO () -> Application +cancellationAwareSseServer firstFrame signalClientClosed _request respond = + respond $ + responseStream status200 [(hContentType, "text/event-stream")] $ \write flush -> do + let keepAlive = do + write ": keepalive\n\n" + flush + threadDelay 1000 + keepAlive + onDisconnect :: SomeException -> IO () + onDisconnect _ = signalClientClosed + write $ Builder.lazyByteString firstFrame + flush + keepAlive `catch` onDisconnect + +collectModelStream :: + ChatModel model => + model -> [Message] -> Maybe (ModelConfig model) -> IO (Either LangchainError [StreamEvent]) +collectModelStream provider messages config = + runResourceT $ runExceptT $ collectEvents (stream provider messages config)