{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE RecordWildCards #-}

-- | The OpenAI backend
module GHC.Plugin.OllamaHoles.Backend.OpenAI (openAIBackend, openAICompatibleBackend) where

import Network.HTTP.Req
import System.Environment (lookupEnv)

import Data.Aeson (FromJSON (..), Value, object, parseJSON, (.:), (.=))
import Data.Aeson.Types (Parser, parseMaybe)
import qualified Data.Aeson as Aeson

import Data.Text (Text)
import Data.Text qualified as T
import Data.Text.Encoding (encodeUtf8)
import GHC.Plugin.OllamaHoles.Backend
import Data.Maybe (fromMaybe)
import Text.URI (mkURI)


-- | The OpenAI backend
openAIBackend :: Backend
openAIBackend :: Backend
openAIBackend = Text -> Text -> Backend
openAICompatibleBackend Text
"api.openai.com" Text
"OPENAI_API_KEY"

-- | Any OpenAI compatible backend
openAICompatibleBackend :: Text -> Text -> Backend
openAICompatibleBackend :: Text -> Text -> Backend
openAICompatibleBackend Text
base_url Text
key_name = Backend {IO (Maybe [Text])
Text -> Text -> IO (Either String Text)
forall {p} {p}.
(ToJSON p, ToJSON p) =>
p -> p -> IO (Either String Text)
listModels :: IO (Maybe [Text])
generateFits :: forall {p} {p}.
(ToJSON p, ToJSON p) =>
p -> p -> IO (Either String Text)
listModels :: IO (Maybe [Text])
generateFits :: Text -> Text -> IO (Either String Text)
..}
  where
    apiEndpoint :: IO (Url 'Https, Option scheme)
apiEndpoint = do URI
uri <- Text -> IO URI
forall (m :: * -> *). MonadThrow m => Text -> m URI
mkURI Text
base_url
                     case URI -> Maybe (Url 'Https, Option scheme)
forall (scheme :: Scheme). URI -> Maybe (Url 'Https, Option scheme)
useHttpsURI URI
uri of
                      Just (Url 'Https
url, Option scheme
opts) -> (Url 'Https, Option scheme) -> IO (Url 'Https, Option scheme)
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return (Url 'Https
url Url 'Https -> Text -> Url 'Https
forall (scheme :: Scheme). Url scheme -> Text -> Url scheme
/: Text
"v1", Option scheme
opts)
                      Maybe (Url 'Https, Option scheme)
Nothing -> String -> IO (Url 'Https, Option scheme)
forall a. HasCallStack => String -> a
error (String -> IO (Url 'Https, Option scheme))
-> String -> IO (Url 'Https, Option scheme)
forall a b. (a -> b) -> a -> b
$ String
"could not parse " String -> String -> String
forall a. Semigroup a => a -> a -> a
<> Text -> String
T.unpack Text
base_url String -> String -> String
forall a. Semigroup a => a -> a -> a
<> String
" as a URI"
    listModels :: IO (Maybe [Text])
listModels = do
        Maybe String
apiKey <- String -> IO (Maybe String)
lookupEnv (String -> IO (Maybe String)) -> String -> IO (Maybe String)
forall a b. (a -> b) -> a -> b
$ Text -> String
T.unpack Text
key_name
        case Maybe String
apiKey of
            Maybe String
Nothing -> Maybe [Text] -> IO (Maybe [Text])
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return Maybe [Text]
forall a. Maybe a
Nothing
            Just String
key -> do
                (Url 'Https
url, Option 'Https
opts) <-  IO (Url 'Https, Option 'Https)
forall {scheme :: Scheme}. IO (Url 'Https, Option scheme)
apiEndpoint
                let headers :: Option scheme
headers = ByteString -> ByteString -> Option scheme
forall (scheme :: Scheme).
ByteString -> ByteString -> Option scheme
header ByteString
"Authorization" (ByteString
"Bearer " ByteString -> ByteString -> ByteString
forall a. Semigroup a => a -> a -> a
<> Text -> ByteString
encodeUtf8 (String -> Text
T.pack String
key))
                BsResponse
response <- HttpConfig -> Req BsResponse -> IO BsResponse
forall (m :: * -> *) a. MonadIO m => HttpConfig -> Req a -> m a
runReq HttpConfig
defaultHttpConfig (Req BsResponse -> IO BsResponse)
-> Req BsResponse -> IO BsResponse
forall a b. (a -> b) -> a -> b
$ GET
-> Url 'Https
-> NoReqBody
-> Proxy BsResponse
-> Option 'Https
-> Req BsResponse
forall (m :: * -> *) method body response (scheme :: Scheme).
(MonadHttp m, HttpMethod method, HttpBody body,
 HttpResponse response,
 HttpBodyAllowed (AllowsBody method) (ProvidesBody body)) =>
method
-> Url scheme
-> body
-> Proxy response
-> Option scheme
-> m response
req GET
GET (Url 'Https
url Url 'Https -> Text -> Url 'Https
forall (scheme :: Scheme). Url scheme -> Text -> Url scheme
/: Text
"models") NoReqBody
NoReqBody Proxy BsResponse
bsResponse (Option 'Https
forall {scheme :: Scheme}. Option scheme
headers Option 'Https -> Option 'Https -> Option 'Https
forall a. Semigroup a => a -> a -> a
<> Option 'Https
opts)
                Maybe [Text] -> IO (Maybe [Text])
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return (Maybe [Text] -> IO (Maybe [Text]))
-> Maybe [Text] -> IO (Maybe [Text])
forall a b. (a -> b) -> a -> b
$ Value -> [Text]
parseOpenAIModels (Value -> [Text]) -> Maybe Value -> Maybe [Text]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ByteString -> Maybe Value
forall a. FromJSON a => ByteString -> Maybe a
Aeson.decodeStrict (BsResponse -> HttpResponseBody BsResponse
forall response.
HttpResponse response =>
response -> HttpResponseBody response
responseBody BsResponse
response)

    parseOpenAIResponse :: Value -> Maybe Text
    parseOpenAIResponse :: Value -> Maybe Text
parseOpenAIResponse = (Value -> Parser Text) -> Value -> Maybe Text
forall a b. (a -> Parser b) -> a -> Maybe b
parseMaybe Value -> Parser Text
parseResponse
      where
        parseResponse :: Value -> Parser Text
        parseResponse :: Value -> Parser Text
parseResponse Value
val = do
            Object
obj <- Value -> Parser Object
forall a. FromJSON a => Value -> Parser a
parseJSON Value
val
            [Value]
choices <- Object
obj Object -> Key -> Parser [Value]
forall a. FromJSON a => Object -> Key -> Parser a
.: Key
"choices"
            case [Value]
choices of
                [] -> String -> Parser Text
forall a. String -> Parser a
forall (m :: * -> *) a. MonadFail m => String -> m a
fail String
"No choices in response"
                (Value
choice : [Value]
_) -> do
                    Object
choiceObj <- Value -> Parser Object
forall a. FromJSON a => Value -> Parser a
parseJSON Value
choice
                    Value
message <- Object
choiceObj Object -> Key -> Parser Value
forall a. FromJSON a => Object -> Key -> Parser a
.: Key
"message"
                    Object
messageObj <- Value -> Parser Object
forall a. FromJSON a => Value -> Parser a
parseJSON Value
message
                    Object
messageObj Object -> Key -> Parser Text
forall a. FromJSON a => Object -> Key -> Parser a
.: Key
"content"

    generateFits :: p -> p -> IO (Either String Text)
generateFits p
prompt p
modelName = do
        Maybe String
apiKey <- String -> IO (Maybe String)
lookupEnv (String -> IO (Maybe String)) -> String -> IO (Maybe String)
forall a b. (a -> b) -> a -> b
$ Text -> String
T.unpack Text
key_name
        case Maybe String
apiKey of
            Maybe String
Nothing -> Either String Text -> IO (Either String Text)
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return (Either String Text -> IO (Either String Text))
-> Either String Text -> IO (Either String Text)
forall a b. (a -> b) -> a -> b
$ String -> Either String Text
forall a b. a -> Either a b
Left (String -> Either String Text) -> String -> Either String Text
forall a b. (a -> b) -> a -> b
$ String
"API key not found. Set the " String -> String -> String
forall a. Semigroup a => a -> a -> a
<> Text -> String
T.unpack Text
key_name String -> String -> String
forall a. Semigroup a => a -> a -> a
<> String
" environment variable."
            Just String
key -> do
                let requestBody :: Value
requestBody =
                        [Pair] -> Value
object
                            [ Key
"model" Key -> p -> Pair
forall v. ToJSON v => Key -> v -> Pair
forall e kv v. (KeyValue e kv, ToJSON v) => Key -> v -> kv
.= p
modelName
                            , Key
"messages" Key -> [Value] -> Pair
forall v. ToJSON v => Key -> v -> Pair
forall e kv v. (KeyValue e kv, ToJSON v) => Key -> v -> kv
.= [[Pair] -> Value
object [Key
"role" Key -> Text -> Pair
forall v. ToJSON v => Key -> v -> Pair
forall e kv v. (KeyValue e kv, ToJSON v) => Key -> v -> kv
.= (Text
"user" :: Text), Key
"content" Key -> p -> Pair
forall v. ToJSON v => Key -> v -> Pair
forall e kv v. (KeyValue e kv, ToJSON v) => Key -> v -> kv
.= p
prompt]]
                            ]

                (Url 'Https
url, Option 'Https
opts) <- IO (Url 'Https, Option 'Https)
forall {scheme :: Scheme}. IO (Url 'Https, Option scheme)
apiEndpoint
                let headers :: Option scheme
headers =
                        ByteString -> ByteString -> Option scheme
forall (scheme :: Scheme).
ByteString -> ByteString -> Option scheme
header ByteString
"Content-Type" ByteString
"application/json"
                            Option scheme -> Option scheme -> Option scheme
forall a. Semigroup a => a -> a -> a
<> ByteString -> ByteString -> Option scheme
forall (scheme :: Scheme).
ByteString -> ByteString -> Option scheme
header ByteString
"Authorization" (ByteString
"Bearer " ByteString -> ByteString -> ByteString
forall a. Semigroup a => a -> a -> a
<> Text -> ByteString
encodeUtf8 (String -> Text
T.pack String
key))

                BsResponse
response <- HttpConfig -> Req BsResponse -> IO BsResponse
forall (m :: * -> *) a. MonadIO m => HttpConfig -> Req a -> m a
runReq HttpConfig
defaultHttpConfig (Req BsResponse -> IO BsResponse)
-> Req BsResponse -> IO BsResponse
forall a b. (a -> b) -> a -> b
$
                              POST
-> Url 'Https
-> ReqBodyJson Value
-> Proxy BsResponse
-> Option 'Https
-> Req BsResponse
forall (m :: * -> *) method body response (scheme :: Scheme).
(MonadHttp m, HttpMethod method, HttpBody body,
 HttpResponse response,
 HttpBodyAllowed (AllowsBody method) (ProvidesBody body)) =>
method
-> Url scheme
-> body
-> Proxy response
-> Option scheme
-> m response
req POST
POST (Url 'Https
url Url 'Https -> Text -> Url 'Https
forall (scheme :: Scheme). Url scheme -> Text -> Url scheme
/: Text
"chat" Url 'Https -> Text -> Url 'Https
forall (scheme :: Scheme). Url scheme -> Text -> Url scheme
/: Text
"completions") (Value -> ReqBodyJson Value
forall a. a -> ReqBodyJson a
ReqBodyJson Value
requestBody) Proxy BsResponse
bsResponse (Option 'Https
forall {scheme :: Scheme}. Option scheme
headers Option 'Https -> Option 'Https -> Option 'Https
forall a. Semigroup a => a -> a -> a
<> Option 'Https
opts)
                case ByteString -> Maybe Value
forall a. FromJSON a => ByteString -> Maybe a
Aeson.decodeStrict (BsResponse -> HttpResponseBody BsResponse
forall response.
HttpResponse response =>
response -> HttpResponseBody response
responseBody BsResponse
response) Maybe Value -> (Value -> Maybe Text) -> Maybe Text
forall a b. Maybe a -> (a -> Maybe b) -> Maybe b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= Value -> Maybe Text
parseOpenAIResponse of
                    Just Text
content -> Either String Text -> IO (Either String Text)
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return (Either String Text -> IO (Either String Text))
-> Either String Text -> IO (Either String Text)
forall a b. (a -> b) -> a -> b
$ Text -> Either String Text
forall a b. b -> Either a b
Right Text
content
                    Maybe Text
Nothing ->
                        Either String Text -> IO (Either String Text)
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return (Either String Text -> IO (Either String Text))
-> Either String Text -> IO (Either String Text)
forall a b. (a -> b) -> a -> b
$ String -> Either String Text
forall a b. a -> Either a b
Left (String -> Either String Text) -> String -> Either String Text
forall a b. (a -> b) -> a -> b
$ String
"Failed to parse OpenAI response: " String -> String -> String
forall a. Semigroup a => a -> a -> a
<> ByteString -> String
forall a. Show a => a -> String
show (BsResponse -> HttpResponseBody BsResponse
forall response.
HttpResponse response =>
response -> HttpResponseBody response
responseBody BsResponse
response)

-- | Parse the models from the endpoint
parseOpenAIModels :: Value -> [Text]
parseOpenAIModels :: Value -> [Text]
parseOpenAIModels Value
value =
    [Text] -> Maybe [Text] -> [Text]
forall a. a -> Maybe a -> a
fromMaybe [] ((Value -> Parser [Text]) -> Value -> Maybe [Text]
forall a b. (a -> Parser b) -> a -> Maybe b
parseMaybe Value -> Parser [Text]
parseModels Value
value)
  where
    extractModelId :: Value -> Parser Text
    extractModelId :: Value -> Parser Text
extractModelId Value
model = do
        Object
obj <- Value -> Parser Object
forall a. FromJSON a => Value -> Parser a
parseJSON Value
model
        Object
obj Object -> Key -> Parser Text
forall a. FromJSON a => Object -> Key -> Parser a
.: Key
"id"

    parseModels :: Value -> Parser [Text]
    parseModels :: Value -> Parser [Text]
parseModels Value
val = do
        Object
obj <- Value -> Parser Object
forall a. FromJSON a => Value -> Parser a
parseJSON Value
val
        [Value]
models <- Object
obj Object -> Key -> Parser [Value]
forall a. FromJSON a => Object -> Key -> Parser a
.: Key
"data"
        (Value -> Parser Text) -> [Value] -> Parser [Text]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM Value -> Parser Text
extractModelId [Value]
models