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

-- | The Gemini backend
module GHC.Plugin.OllamaHoles.Backend.Gemini (geminiBackend) where

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

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

import Data.Text (Text)
import qualified Data.Text as T
import GHC.Plugin.OllamaHoles.Backend
import Data.Maybe

-- | The Gemini backend
geminiBackend :: Backend
geminiBackend :: Backend
geminiBackend = Backend{IO (Maybe [Text])
Text -> Text -> IO (Either String Text)
forall {p}. ToJSON p => p -> Text -> IO (Either String Text)
listModels :: IO (Maybe [Text])
generateFits :: forall {p}. ToJSON p => p -> Text -> IO (Either String Text)
listModels :: IO (Maybe [Text])
generateFits :: Text -> Text -> IO (Either String Text)
..}
  where
    apiEndpoint :: Url 'Https
apiEndpoint = Text -> Url 'Https
https Text
"generativelanguage.googleapis.com"  Url 'Https -> Text -> Url 'Https
forall (scheme :: Scheme). Url scheme -> Text -> Url scheme
/: Text
"v1beta"
    listModels :: IO (Maybe [Text])
listModels = do
        Maybe String
apiKey <- String -> IO (Maybe String)
lookupEnv String
"GEMINI_API_KEY"
        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
                let url :: Url 'Https
url = Url 'Https
apiEndpoint Url 'Https -> Text -> Url 'Https
forall (scheme :: Scheme). Url scheme -> Text -> Url scheme
/: Text
"models"
                JsonResponse Value
response <- HttpConfig -> Req (JsonResponse Value) -> IO (JsonResponse Value)
forall (m :: * -> *) a. MonadIO m => HttpConfig -> Req a -> m a
runReq HttpConfig
defaultHttpConfig (Req (JsonResponse Value) -> IO (JsonResponse Value))
-> Req (JsonResponse Value) -> IO (JsonResponse Value)
forall a b. (a -> b) -> a -> b
$ GET
-> Url 'Https
-> NoReqBody
-> Proxy (JsonResponse Value)
-> Option 'Https
-> Req (JsonResponse Value)
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 NoReqBody
NoReqBody Proxy (JsonResponse Value)
forall a. Proxy (JsonResponse a)
jsonResponse (Text
"key" Text -> String -> Option 'Https
forall param a.
(QueryParam param, ToHttpApiData a) =>
Text -> a -> param
=: String
key)
                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
$ [Text] -> Maybe [Text]
forall a. a -> Maybe a
Just ([Text] -> Maybe [Text]) -> [Text] -> Maybe [Text]
forall a b. (a -> b) -> a -> b
$ Value -> [Text]
parseGeminiModels (JsonResponse Value -> HttpResponseBody (JsonResponse Value)
forall response.
HttpResponse response =>
response -> HttpResponseBody response
responseBody JsonResponse Value
response)

    parseGeminiResponse :: Value -> Maybe Text
    parseGeminiResponse :: Value -> Maybe Text
parseGeminiResponse = (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]
candidates <- Object
obj Object -> Key -> Parser [Value]
forall a. FromJSON a => Object -> Key -> Parser a
.: Key
"candidates"
            case [Value]
candidates of
                [] -> String -> Parser Text
forall a. String -> Parser a
forall (m :: * -> *) a. MonadFail m => String -> m a
fail String
"No candidates in response"
                (Value
candidate : [Value]
_) -> do
                    Object
candidateObj <- Value -> Parser Object
forall a. FromJSON a => Value -> Parser a
parseJSON Value
candidate
                    Object
content <- Object
candidateObj Object -> Key -> Parser Object
forall a. FromJSON a => Object -> Key -> Parser a
.: Key
"content"
                    [Value]
parts <- Object
content Object -> Key -> Parser [Value]
forall a. FromJSON a => Object -> Key -> Parser a
.: Key
"parts"
                    case [Value]
parts of
                        [] -> String -> Parser Text
forall a. String -> Parser a
forall (m :: * -> *) a. MonadFail m => String -> m a
fail String
"No parts in content"
                        (Value
part : [Value]
_) -> do
                            Object
partObj <- Value -> Parser Object
forall a. FromJSON a => Value -> Parser a
parseJSON Value
part
                            Object
partObj Object -> Key -> Parser Text
forall a. FromJSON a => Object -> Key -> Parser a
.: Key
"text"

    generateFits :: p -> Text -> IO (Either String Text)
generateFits p
prompt Text
modelName = do
        Maybe String
apiKey <- String -> IO (Maybe String)
lookupEnv String
"GEMINI_API_KEY"
        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
"Gemini API key not found. Set the GEMINI_API_KEY environment variable."
            Just String
key -> do
                let requestBody :: Value
requestBody =
                        [Pair] -> Value
object
                            [ Key
"contents" 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
"parts" 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
"text" Key -> p -> Pair
forall v. ToJSON v => Key -> v -> Pair
forall e kv v. (KeyValue e kv, ToJSON v) => Key -> v -> kv
.= p
prompt]]]]
                            ]

                let url :: Url 'Https
url = Url 'Https
apiEndpoint Url 'Https -> Text -> Url 'Https
forall (scheme :: Scheme). Url scheme -> Text -> Url scheme
/: Text
"models" Url 'Https -> Text -> Url 'Https
forall (scheme :: Scheme). Url scheme -> Text -> Url scheme
/: (Text
modelName Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<>  Text
":generateContent")
                    params :: Option 'Https
params = Text
"key" Text -> String -> Option 'Https
forall param a.
(QueryParam param, ToHttpApiData a) =>
Text -> a -> param
=: String
key
                    headers :: Option scheme
headers = ByteString -> ByteString -> Option scheme
forall (scheme :: Scheme).
ByteString -> ByteString -> Option scheme
header ByteString
"Content-Type" ByteString
"application/json"

                JsonResponse Value
response <- HttpConfig -> Req (JsonResponse Value) -> IO (JsonResponse Value)
forall (m :: * -> *) a. MonadIO m => HttpConfig -> Req a -> m a
runReq HttpConfig
defaultHttpConfig (Req (JsonResponse Value) -> IO (JsonResponse Value))
-> Req (JsonResponse Value) -> IO (JsonResponse Value)
forall a b. (a -> b) -> a -> b
$ POST
-> Url 'Https
-> ReqBodyJson Value
-> Proxy (JsonResponse Value)
-> Option 'Https
-> Req (JsonResponse Value)
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 (Value -> ReqBodyJson Value
forall a. a -> ReqBodyJson a
ReqBodyJson Value
requestBody) Proxy (JsonResponse Value)
forall a. Proxy (JsonResponse a)
jsonResponse (Option 'Https
forall {scheme :: Scheme}. Option scheme
headers Option 'Https -> Option 'Https -> Option 'Https
forall a. Semigroup a => a -> a -> a
<> Option 'Https
params)

                case Value -> Maybe Text
parseGeminiResponse (JsonResponse Value -> HttpResponseBody (JsonResponse Value)
forall response.
HttpResponse response =>
response -> HttpResponseBody response
responseBody JsonResponse Value
response) 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 Gemini response: " String -> String -> String
forall a. Semigroup a => a -> a -> a
<> Value -> String
forall a. Show a => a -> String
show (JsonResponse Value -> HttpResponseBody (JsonResponse Value)
forall response.
HttpResponse response =>
response -> HttpResponseBody response
responseBody JsonResponse Value
response)

-- | Parse the models from the endpoint
parseGeminiModels :: Value -> [Text]
parseGeminiModels :: Value -> [Text]
parseGeminiModels 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
        Int -> Text -> Text
T.drop (Text -> Int
T.length Text
"models/") (Text -> Text) -> Parser Text -> Parser Text
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Object
obj Object -> Key -> Parser Text
forall a. FromJSON a => Object -> Key -> Parser a
.: Key
"name"

    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
"models"
        (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