MLConnectorsSpec.hs

{-# LANGUAGE OverloadedStrings #-}

module Test.MLConnectorsSpec (spec) where

import Data.ByteString.Lazy.Char8 qualified as LBS
import Database.Bloodhound.OpenSearch3.Types
import TestsUtils.Import
import Prelude

-- | Create response: @{connector_id}@.
sampleCreateResponse :: LBS.ByteString
sampleCreateResponse = "{\"connector_id\":\"a1eMb4kBJ1eYAeTMAljY\"}"

-- | Get response: the bare stored connector document. Note @credential@
-- is never returned and @version@ comes back as a JSON string.
sampleGetResponse :: LBS.ByteString
sampleGetResponse =
  "{\
  \  \"name\": \"OpenAI Chat Connector\",\
  \  \"description\": \"The connector to public OpenAI model service\",\
  \  \"version\": \"1\",\
  \  \"protocol\": \"http\",\
  \  \"parameters\": {\"endpoint\": \"api.openai.com\", \"model\": \"gpt-3.5-turbo\"},\
  \  \"actions\": [\
  \    {\
  \      \"action_type\": \"predict\",\
  \      \"method\": \"POST\",\
  \      \"url\": \"https://api.openai.com/v1/chat/completions\",\
  \      \"headers\": {\"Authorization\": \"Bearer XXX\"},\
  \      \"request_body\": \"{ \\\"model\\\": \\\"gpt-3.5-turbo\\\" }\"\
  \    }\
  \  ]\
  \}"

spec :: Spec
spec = describe "ML Commons connector APIs" $ do
  describe "ConnectorProtocol" $ do
    it "encodes the documented protocols" $ do
      encode ConnectorProtocolHttp `shouldBe` "\"http\""
      encode ConnectorProtocolAwsSigv4 `shouldBe` "\"aws_sigv4\""
    it "decodes the documented protocols" $ do
      decode "\"http\"" `shouldBe` Just ConnectorProtocolHttp
      decode "\"aws_sigv4\"" `shouldBe` Just ConnectorProtocolAwsSigv4
    it "round-trips an unknown protocol through Other" $ do
      let other = ConnectorProtocolOther "custom"
      decode (encode other) `shouldBe` Just other

  describe "CreateConnectorResponse" $
    it "decodes the doc sample" $
      decode sampleCreateResponse
        `shouldBe` Just (CreateConnectorResponse {createConnectorResponseId = Just "a1eMb4kBJ1eYAeTMAljY"})

  describe "ConnectorInfo" $ do
    it "decodes the doc get sample (no credential; version is a string)" $ do
      let decoded = decode sampleGetResponse :: Maybe ConnectorInfo
      case decoded of
        Just info -> do
          connectorInfoName info `shouldBe` Just "OpenAI Chat Connector"
          connectorInfoVersion info `shouldBe` Just "1"
          connectorInfoProtocol info `shouldBe` Just ConnectorProtocolHttp
          connectorInfoActions info `shouldSatisfy` isJust
        Nothing -> expectationFailure "expected ConnectorInfo decode"
    it "round-trips a connector with an action" $ do
      let info =
            ConnectorInfo
              { connectorInfoName = Just "c",
                connectorInfoDescription = Nothing,
                connectorInfoVersion = Just "1",
                connectorInfoProtocol = Just ConnectorProtocolHttp,
                connectorInfoParameters = Nothing,
                connectorInfoActions =
                  Just
                    [ MLConnectorAction
                        { mlConnectorActionType = "predict",
                          mlConnectorActionMethod = Just "POST",
                          mlConnectorActionUrl = Just "https://x",
                          mlConnectorActionHeaders = Nothing,
                          mlConnectorActionRequestBody = Nothing,
                          mlConnectorActionPreProcessFunction = Nothing,
                          mlConnectorActionPostProcessFunction = Nothing
                        }
                    ],
                connectorInfoBackendRoles = Nothing,
                connectorInfoAccessMode = Nothing
              }
      decode (encode info) `shouldBe` Just info