DeepSeek.hs

module LLM.Providers.DeepSeek
  ( deepSeekGateway,
    deepSeekGatewayWith,
    deepSeekProvider,
    deepSeekProviderWith,
    deepSeekBuildBodyPairs,
  )
where

import Data.Aeson
  ( KeyValue ((.=)),
    Value,
    encode,
    object,
  )
import Data.Aeson.Types (Pair)
import Data.Text (Text)
import Data.Text.Lazy qualified as TL
import Data.Text.Lazy.Encoding qualified as TLE
import LLM.Core.LLMProvider (LLMProvider (..), toGateway)
import LLM.Core.ProviderUtils (handleStreamResponse, lenientConfig)
import LLM.Core.Types
  ( commonImageMediaTypes,
    ChatRequest
      ( reqConversation,
        reqMaxTokens,
        reqModel,
        reqTemperature,
        reqThinking,
        reqTools
      ),
    LLMGateway,
    MessageEncodeOptions (..),
    ThinkingMode (..),
    deepSeekMessageEncodeOptions,
    pattern UserTurn,
  )
import LLM.Providers.OpenAI
  ( authHeader,
    buildMessages,
    encodeToolDef,
    parseOpenAIObjectResponse,
    parseOpenAIResponse,
    parseOpenAIStream,
  )
import Network.HTTP.Req
  ( Option,
    POST (POST),
    ReqBodyJson (ReqBodyJson),
    Url,
    https,
    bsResponse,
    req,
    reqBr,
    responseBody,
    responseStatusCode,
    runReq,
    (/:),
  )

-- | Create a LLMGateway for the DeepSeek provider. Takes the API key as a parameter.
deepSeekGateway :: Text -> LLMGateway
deepSeekGateway apiKey = toGateway $ deepSeekProvider apiKey

-- | Create a LLMGateway for the DeepSeek provider with a custom base URL. Takes the API key and base URL as parameters.
deepSeekGatewayWith :: Url scheme -> Option scheme -> Text -> LLMGateway
deepSeekGatewayWith baseUrl baseOpts apiKey = toGateway (deepSeekProviderWith baseUrl baseOpts apiKey)

-- | DeepSeek provider at api.deepseek.com.
deepSeekProvider :: Text -> LLMProvider
deepSeekProvider = deepSeekProviderWith (https "api.deepseek.com") mempty

-- | DeepSeek-compatible provider with a custom base URL.
deepSeekProviderWith :: Url scheme -> Option scheme -> Text -> LLMProvider
deepSeekProviderWith baseUrl baseOpts apiKey =
  LLMProvider
    { providerName = "deepseek",
      providerImageMediaTypes = commonImageMediaTypes,
      providerImageUrlSupport = True,
      buildBody = deepSeekBuildBody,
      sendRequest = sendRequest,
      openStream = \body k ->
        runReq lenientConfig $ do
          let url = baseUrl /: "v1" /: "chat" /: "completions"
              opts = baseOpts <> authHeader apiKey
          reqBr POST url (ReqBodyJson body) opts $ \resp ->
            handleStreamResponse resp k,
      parseStream = const parseOpenAIStream,
      parseResponse = pure . parseOpenAIResponse,
      buildObjectBody = \r schema ->
        let schemaText = TL.toStrict . TLE.decodeUtf8 $ encode schema
            instruction =
              "Respond with a raw JSON object matching this JSON schema. "
                <> "No markdown, no code fences, no explanation:\n"
                <> schemaText
            r' = r {reqConversation = r.reqConversation <> [UserTurn instruction]}
         in object $
              deepSeekBuildBodyPairs False r'
                <> ["response_format" .= object ["type" .= ("json_object" :: Text)]],
      sendObjectRequest = sendRequest,
      parseObjectResponse = pure . parseOpenAIObjectResponse
    }
  where
    sendRequest body =
      runReq lenientConfig $ do
        let url = baseUrl /: "v1" /: "chat" /: "completions"
            opts = baseOpts <> authHeader apiKey
        resp <- req POST url (ReqBodyJson body) bsResponse opts
        pure (responseStatusCode resp, responseBody resp)

-- | OpenAI-compatible body; DeepSeek uses @max_tokens@ (not @max_completion_tokens@).
deepSeekBuildBody :: Bool -> ChatRequest -> Value
deepSeekBuildBody stream r = object $ deepSeekBuildBodyPairs stream r

deepSeekBuildBodyPairs :: Bool -> ChatRequest -> [Pair]
deepSeekBuildBodyPairs stream r =
  [ "model" .= r.reqModel,
    "max_tokens" .= r.reqMaxTokens,
    "messages" .= buildMessages encodeOpts r
  ]
    ++ thinkingPairs r
    ++ ["temperature" .= t | Just t <- [r.reqTemperature]]
    ++ ["tools" .= map encodeToolDef r.reqTools | not (null r.reqTools)]
    ++ ["stream" .= True | stream]
    ++ ["stream_options" .= object ["include_usage" .= True] | stream]
  where
    encodeOpts =
      deepSeekMessageEncodeOptions
        { meoRequireReasoning = maybe False (.tmEnabled) r.reqThinking
        }

thinkingPairs :: ChatRequest -> [Pair]
thinkingPairs r =
  case r.reqThinking of
    Just tm
      | not tm.tmEnabled ->
          ["thinking" .= object ["type" .= ("disabled" :: Text)]]
    Just tm ->
      ("thinking" .= object ["type" .= ("enabled" :: Text)]) : ["reasoning_effort" .= e | Just e <- [tm.tmEffort]]
    Nothing ->
      ["thinking" .= object ["type" .= ("disabled" :: Text)]]