diff --git a/README.md b/README.md --- a/README.md +++ b/README.md @@ -5,9 +5,18 @@ ## Introduction This is an example of a typed-hole plugin for GHC that uses the [Ollama](https://ollama.com/) to host a local LLM to fill in holes in Haskell code. +Before using this plugin, make sure you have the Ollama CLI installed and the model you want to use is available. +You can install the Ollama CLI by following the instructions at [https://ollama.com/download](https://ollama.com/download), +and you can install the default model (gemma3:27b) by running `ollama pull gemma3:27b`. +Note that the speed and quality of the hole-fits generated by the plugin depends on +the model you use, and the default model requires a GPU to run efficiently. +For a smaller model, we suggest `gemma3:4b-it-qat`, or `deepcoder:1.5b`. + +This plugin is also availble on Hackage [https://hackage.haskell.org/package/ollama-holes-plugin](https://hackage.haskell.org/package/ollama-holes-plugin) + ## Example -Given +Given ```haskell {-# OPTIONS_GHC -fplugin=GHC.Plugin.OllamaHoles #-} @@ -73,3 +82,16 @@ ``` 5. Enjoy! If you want to change the underlying model, make sure to pass the model name via the plugin arguments (see example) + +## OpenAI and Gemini backends + +The plugin now supports using the OpenAI API and Gemini APIs to generate valid hole fits. +Simply set the backend flag `-fplugin-opt=GHC.Plugin.OllamaHoles:backend=openai`, +or `-fplugin-opt=GHC.Plugin.OllamaHoles:backend=gemini`, and make sure that +you have the `OPENAI_API_KEY` or `GEMINI_API_KEY` set in your environment. + +To use with any other OpenAI compatible api (like groq or OpenRouter), simply set +`-fplugin-opt=GHC.Plugin.OllamaHoles:backend=openai`, +and +`-fplugin-opt=GHC.Plugin.OllamaHoles:openai_base_url=https://api.groq.com/openai`, +`-fplugin-opt=GHC.Plugin.OllamaHoles:openai_key_name=GROQ_API_KEY`, diff --git a/ollama-holes-plugin.cabal b/ollama-holes-plugin.cabal --- a/ollama-holes-plugin.cabal +++ b/ollama-holes-plugin.cabal @@ -20,33 +20,50 @@ -- PVP summary: +-+------- breaking API changes -- | | +----- non-breaking API additions -- | | | +--- code changes with no API change -version: 0.1.0.0 +version: 0.1.1.0 -- A short (one-line) description of the package. -synopsis: A typed-hole plugin that uses LLMs via Ollama to generate valid hole-fits +synopsis: A typed-hole plugin that uses LLMs to generate valid hole-fits -- A longer description of the package. -description: This package provides a GHC plugin that uses LLMs via Ollama to generate valid hole-fits. - . +description: This package provides a GHC plugin that uses LLMs to generate valid hole-fits. + It supports multiple backends including Ollama, OpenAI, and Gemini. + The following flags are available: - . + To specify the model to use: - . + > -fplugin-opt=GHC.Plugin.OllamaHoles:model= - . + + To specify the backend to use (ollama, openai, or gemini): + + > -fplugin-opt=GHC.Plugin.OllamaHoles:backend= + + When using the openai backend, you can specify a custom base_url, e.g. + + > -fplugin-opt=GHC.Plugin.OllamaHoles:openai_base_url=api.groq.com/api + + You can also specify which key to use. + + > -fplugin-opt=GHC.Plugin.OllamaHoles:openai_key_name=GROQ_API_KEY + To specify how many fits to generate (passed to the model) - . + > -fplugin-opt=GHC.Plugin.OllamaHoles:n=5 - . + To enable debug output: - . + > -fplugin-opt=GHC.Plugin.OllamaHoles:debug=True - . - Before using this plugin, make sure you have the Ollama CLI installed and the model + + For the Ollama backend, make sure you have the Ollama CLI installed and the model you want to use is available. You can install the Ollama CLI by following the - instructions at [https://ollama.com/download](https://ollama.com/download), + instructions at https://ollama.com/download, and you can install the default model (gemma3:27b) by running `ollama pull gemma3:27b`. - . + + For the OpenAI backend, you'll need to set the OPENAI_API_KEY environment variable with your API key. + + For the Gemini backend, you'll need to set the GEMINI_API_KEY environment variable with your API key. + Note that the speed and quality of the hole-fits generated by the plugin depends on the model you use, and the default model requires a GPU to run efficiently. For a smaller model, we suggest `gemma3:4b-it-qat`, or `deepcoder:1.5b`. @@ -67,7 +84,7 @@ -- A copyright notice. -copyright: 2025- Matthias Pall Gissurarson +copyright: 2025 (c) Matthias Pall Gissurarson category: Development, Compiler Plugin tested-with: GHC == 9.6.* build-type: Simple @@ -96,7 +113,10 @@ exposed-modules: GHC.Plugin.OllamaHoles -- Modules included in this library but not exported. - -- other-modules: + other-modules: GHC.Plugin.OllamaHoles.Backend + , GHC.Plugin.OllamaHoles.Backend.Ollama + , GHC.Plugin.OllamaHoles.Backend.OpenAI + , GHC.Plugin.OllamaHoles.Backend.Gemini -- LANGUAGE extensions used by modules in this package. -- other-extensions: @@ -106,6 +126,9 @@ ghc ^>=9.6, ollama-haskell ^>= 0.1, text ^>= 2.1, + req ^>= 3.13, + modern-uri ^>= 0.3, + aeson ^>= 2.2, -- Directories containing source files. diff --git a/src/GHC/Plugin/OllamaHoles.hs b/src/GHC/Plugin/OllamaHoles.hs --- a/src/GHC/Plugin/OllamaHoles.hs +++ b/src/GHC/Plugin/OllamaHoles.hs @@ -13,17 +13,13 @@ import GHC.Tc.Types import GHC.Tc.Types.Constraint (Hole (..), ctLocEnv, ctLocSpan) import GHC.Tc.Utils.Monad (getGblEnv, newTcRef) -import Ollama (GenerateOps (..)) -import Ollama qualified --- | Options for the LLM model -genOps :: Ollama.GenerateOps -genOps = - Ollama.defaultGenerateOps - { modelName = "" - , prompt = "" - } +import GHC.Plugin.OllamaHoles.Backend +import GHC.Plugin.OllamaHoles.Backend.Ollama (ollamaBackend) +import GHC.Plugin.OllamaHoles.Backend.OpenAI (openAICompatibleBackend) +import GHC.Plugin.OllamaHoles.Backend.Gemini (geminiBackend) + promptTemplate :: Text promptTemplate = "You are a typed-hole plugin within GHC, the Glasgow Haskell Compiler.\n" @@ -39,6 +35,15 @@ <> "Feel free to include any other functions from the list of imports to generate more complicated expressions.\n" <> "Output a maximum of {numexpr} expresssions.\n" +getBackend :: Flags -> Backend +getBackend Flags {backend_name = "ollama"} = ollamaBackend +getBackend Flags {backend_name = "gemini"}= geminiBackend +getBackend Flags{backend_name = "openai", ..} = openAICompatibleBackend openai_base_url openai_key_name +getBackend Flags{..} = error $ "unknown backend: " <> T.unpack backend_name + +pluginName :: Text +pluginName = "Ollama Plugin" + -- | Ollama plugin for GHC plugin :: Plugin plugin = @@ -52,76 +57,94 @@ const HoleFitPlugin { candPlugin = \_ c -> return c -- Don't filter candidates - , fitPlugin = \hole fits -> do - let Flags{..} = parseFlags opts - dflags <- getDynFlags - gbl_env <- getGblEnv - let mod_name = moduleNameString $ moduleName $ tcg_mod gbl_env - imports = tcg_imports gbl_env - liftIO $ do - available_models <- Ollama.list - case available_models of - Nothing -> T.putStrLn "--- Ollama plugin: No models available.--" - Just (Ollama.Models models) -> do - unless (any ((== model_name) . Ollama.name) models) $ - error $ - "--- Ollama plugin: Model " - <> T.unpack model_name - <> " not found. " - <> "Use `ollama pull` to download the model, or specify another model using " - <> "`-fplugin-opt=GHC.Plugin.OllamaHoles:model=` ---" - when debug $ T.putStrLn "--- Ollama Plugin: Hole Found ---" - let mn = "Module: " <> mod_name - let lc = "Location: " <> showSDoc dflags (ppr $ ctLocSpan . hole_loc <$> th_hole hole) - let im = "Imports: " <> showSDoc dflags (ppr $ moduleEnvKeys $ imp_mods imports) - - case th_hole hole of - Just h -> do - let lcl_env = ctLocEnv (hole_loc h) - let hv = "Hole variable: _" <> occNameString (occName $ hole_occ h) - let ht = "Hole type: " <> showSDoc dflags (ppr $ hole_ty h) - let rc = "Relevant constraints: " <> showSDoc dflags (ppr $ th_relevant_cts hole) - let le = "Local environment (bindings): " <> showSDoc dflags (ppr $ tcl_rdr lcl_env) - let ge = "Global environment (bindings): " <> showSDoc dflags (ppr $ tcg_binds gbl_env) - let cf = "Candidate fits: " <> showSDoc dflags (ppr fits) - let prompt' = - replacePlaceholders - promptTemplate - [ ("{module}", mn) - , ("{location}", lc) - , ("{imports}", im) - , ("{hole_var}", hv) - , ("{hole_type}", ht) - , ("{relevant_constraints}", rc) - , ("{local_env}", le) - , ("{global_env}", ge) - , ("{candidate_fits}", cf) - , ("{numexpr}", show num_expr) - ] - res <- Ollama.generate genOps{prompt = prompt', modelName = model_name} - case res of - Right rsp -> do - let lns = (preProcess . T.lines) $ Ollama.response_ rsp - when debug $ do - T.putStrLn $ "--- Ollama Plugin: Prompt ---\n" <> prompt' - T.putStrLn $ "--- Ollama Plugin: Response ---\n" <> Ollama.response_ rsp - let fits' = map (RawHoleFit . text . T.unpack) lns - -- Return the generated fits - return fits' - Left err -> do - when debug $ - putStrLn $ - "Ollama plugin failed to generate a response.\n" <> err - -- Return the original fits without modification - return fits - Nothing -> return fits + , fitPlugin = fitPlugin opts } } } + where + fitPlugin opts hole fits = do + let flags@Flags{..} = parseFlags opts + dflags <- getDynFlags + gbl_env <- getGblEnv + let mod_name = moduleNameString $ moduleName $ tcg_mod gbl_env + imports = tcg_imports gbl_env + let backend = getBackend flags + liftIO $ do + available_models <- listModels backend + case available_models of + Nothing -> + error $ + "--- " <> T.unpack pluginName <> ": No models available, check your configuration ---" + Just models -> do + unless (model_name `elem` models) $ + error $ + "--- " + <> T.unpack pluginName + <> ": Model " + <> T.unpack model_name + <> " not found. " + <> ( if backend_name == "ollama" + then "Use `ollama pull` to download the model, or " + else "" + ) + <> "specify another model using " + <> "`-fplugin-opt=GHC.Plugin.OllamaHoles:model=` ---" + <> "--- Availble models: " + <> T.unpack (T.unlines models) + <> " ---" + when debug $ T.putStrLn $ "--- " <> pluginName <> ": Hole Found ---" + let mn = "Module: " <> mod_name + let lc = "Location: " <> showSDoc dflags (ppr $ ctLocSpan . hole_loc <$> th_hole hole) + let im = "Imports: " <> showSDoc dflags (ppr $ moduleEnvKeys $ imp_mods imports) + case th_hole hole of + Just h -> do + let lcl_env = ctLocEnv (hole_loc h) + let hv = "Hole variable: _" <> occNameString (occName $ hole_occ h) + let ht = "Hole type: " <> showSDoc dflags (ppr $ hole_ty h) + let rc = "Relevant constraints: " <> showSDoc dflags (ppr $ th_relevant_cts hole) + let le = "Local environment (bindings): " <> showSDoc dflags (ppr $ tcl_rdr lcl_env) + let ge = "Global environment (bindings): " <> showSDoc dflags (ppr $ tcg_binds gbl_env) + let cf = "Candidate fits: " <> showSDoc dflags (ppr fits) + let prompt' = + replacePlaceholders + promptTemplate + [ ("{module}", mn) + , ("{location}", lc) + , ("{imports}", im) + , ("{hole_var}", hv) + , ("{hole_type}", ht) + , ("{relevant_constraints}", rc) + , ("{local_env}", le) + , ("{global_env}", ge) + , ("{candidate_fits}", cf) + , ("{numexpr}", show num_expr) + ] + res <- generateFits backend prompt' model_name + case res of + Right rsp -> do + let lns = (preProcess . T.lines) rsp + when debug $ do + T.putStrLn $ "--- " <> pluginName <> ": Prompt ---\n" <> prompt' + T.putStrLn $ "--- " <> pluginName <> ": Response ---\n" <> rsp + let fits' = map (RawHoleFit . text . T.unpack) lns + -- Return the generated fits + return fits' + Left err -> do + when debug $ + putStrLn $ + T.unpack pluginName <> " failed to generate a response.\n" <> err + -- Return the original fits without modification + return fits + Nothing -> return fits + -- | Preprocess the response to remove empty lines, lines with only spaces, and code blocks preProcess :: [Text] -> [Text] preProcess [] = [] +-- \| Remove lines between and tags from e.g. deepseek +preProcess (ln : lns) + | T.isPrefixOf "" ln = + preProcess (tail $ dropWhile (not . T.isPrefixOf "") lns) preProcess (ln : lns) | should_drop = preProcess lns where should_drop :: Bool @@ -132,13 +155,16 @@ preProcess (ln : lns) = transform ln : preProcess lns where transform :: Text -> Text - transform = id + transform = T.strip -- | Command line options for the plugin data Flags = Flags { model_name :: Text + , backend_name :: Text , num_expr :: Int , debug :: Bool + , openai_base_url :: Text + , openai_key_name :: Text } -- | Default flags for the plugin @@ -146,8 +172,11 @@ defaultFlags = Flags { model_name = "gemma3:27b-it-qat" + , backend_name = "ollama" , num_expr = 5 , debug = False + , openai_base_url = "https://api.openai.com" + , openai_key_name = "OPENAI_API_KEY" } -- | Parse command line options @@ -160,6 +189,18 @@ | T.isPrefixOf "model=" (T.pack opt) = let model_name = T.drop (T.length "model=") (T.pack opt) in parseFlags' flags{model_name = model_name} opts + parseFlags' flags (opt : opts) + | T.isPrefixOf "backend=" (T.pack opt) = + let backend_name = T.drop (T.length "backend=") (T.pack opt) + in parseFlags' flags{backend_name = backend_name} opts + parseFlags' flags (opt : opts) + | T.isPrefixOf "openai_base_url=" (T.pack opt) = + let openai_base_url = T.drop (T.length "openai_base_url=") (T.pack opt) + in parseFlags' flags{openai_base_url = openai_base_url} opts + parseFlags' flags (opt : opts) + | T.isPrefixOf "openai_key_name=" (T.pack opt) = + let openai_key_name = T.drop (T.length "openai_key_name=") (T.pack opt) + in parseFlags' flags{openai_key_name = openai_key_name} opts parseFlags' flags (opt : opts) | T.isPrefixOf "debug=" (T.pack opt) = let debug = T.unpack $ T.drop (T.length "debug=") (T.pack opt) diff --git a/src/GHC/Plugin/OllamaHoles/Backend.hs b/src/GHC/Plugin/OllamaHoles/Backend.hs new file mode 100644 --- /dev/null +++ b/src/GHC/Plugin/OllamaHoles/Backend.hs @@ -0,0 +1,11 @@ +-- | The Backend data type +module GHC.Plugin.OllamaHoles.Backend (Backend(..)) where + +import Data.Text (Text) + +-- | Backend to use. +data Backend = Backend + { listModels :: IO (Maybe [Text]) + , generateFits :: Text -> Text -> IO (Either String Text) + } + diff --git a/src/GHC/Plugin/OllamaHoles/Backend/Gemini.hs b/src/GHC/Plugin/OllamaHoles/Backend/Gemini.hs new file mode 100644 --- /dev/null +++ b/src/GHC/Plugin/OllamaHoles/Backend/Gemini.hs @@ -0,0 +1,86 @@ +{-# 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{..} + where + apiEndpoint = https "generativelanguage.googleapis.com" /: "v1beta" + listModels = do + apiKey <- lookupEnv "GEMINI_API_KEY" + case apiKey of + Nothing -> return Nothing + Just key -> do + let url = apiEndpoint /: "models" + response <- runReq defaultHttpConfig $ req GET url NoReqBody jsonResponse ("key" =: key) + return $ Just $ parseGeminiModels (responseBody response) + + parseGeminiResponse :: Value -> Maybe Text + parseGeminiResponse = parseMaybe parseResponse + where + parseResponse :: Value -> Parser Text + parseResponse val = do + obj <- parseJSON val + candidates <- obj .: "candidates" + case candidates of + [] -> fail "No candidates in response" + (candidate : _) -> do + candidateObj <- parseJSON candidate + content <- candidateObj .: "content" + parts <- content .: "parts" + case parts of + [] -> fail "No parts in content" + (part : _) -> do + partObj <- parseJSON part + partObj .: "text" + + generateFits prompt modelName = do + apiKey <- lookupEnv "GEMINI_API_KEY" + case apiKey of + Nothing -> return $ Left "Gemini API key not found. Set the GEMINI_API_KEY environment variable." + Just key -> do + let requestBody = + object + [ "contents" .= [object ["parts" .= [object ["text" .= prompt]]]] + ] + + let url = apiEndpoint /: "models" /: (modelName <> ":generateContent") + params = "key" =: key + headers = header "Content-Type" "application/json" + + response <- runReq defaultHttpConfig $ req POST url (ReqBodyJson requestBody) jsonResponse (headers <> params) + + case parseGeminiResponse (responseBody response) of + Just content -> return $ Right content + Nothing -> + return $ Left $ "Failed to parse Gemini response: " <> show (responseBody response) + +-- | Parse the models from the endpoint +parseGeminiModels :: Value -> [Text] +parseGeminiModels value = + fromMaybe [] (parseMaybe parseModels value) + where + extractModelId :: Value -> Parser Text + extractModelId model = do + obj <- parseJSON model + T.drop (T.length "models/") <$> obj .: "name" + + parseModels :: Value -> Parser [Text] + parseModels val = do + obj <- parseJSON val + models <- obj .: "models" + mapM extractModelId models diff --git a/src/GHC/Plugin/OllamaHoles/Backend/Ollama.hs b/src/GHC/Plugin/OllamaHoles/Backend/Ollama.hs new file mode 100644 --- /dev/null +++ b/src/GHC/Plugin/OllamaHoles/Backend/Ollama.hs @@ -0,0 +1,27 @@ +{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE RecordWildCards #-} + +-- | The locally hosted ollama backend +module GHC.Plugin.OllamaHoles.Backend.Ollama (ollamaBackend) where + +import Ollama (GenerateOps (..)) +import Ollama qualified + +import GHC.Plugin.OllamaHoles.Backend + +-- | The locally hosted ollama backend +ollamaBackend :: Backend +ollamaBackend = Backend{..} + where + listModels = fmap getMs <$> Ollama.list + getMs (Ollama.Models models) = fmap Ollama.name models + generateFits prompt modelName = do + fmap Ollama.response_ <$> Ollama.generate genOps{prompt = prompt, modelName = modelName} + + -- \| Options for the LLM model + genOps :: Ollama.GenerateOps + genOps = + Ollama.defaultGenerateOps + { modelName = "" + , prompt = "" + } diff --git a/src/GHC/Plugin/OllamaHoles/Backend/OpenAI.hs b/src/GHC/Plugin/OllamaHoles/Backend/OpenAI.hs new file mode 100644 --- /dev/null +++ b/src/GHC/Plugin/OllamaHoles/Backend/OpenAI.hs @@ -0,0 +1,96 @@ +{-# 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 = openAICompatibleBackend "api.openai.com" "OPENAI_API_KEY" + +-- | Any OpenAI compatible backend +openAICompatibleBackend :: Text -> Text -> Backend +openAICompatibleBackend base_url key_name = Backend {..} + where + apiEndpoint = do uri <- mkURI base_url + case useHttpsURI uri of + Just (url, opts) -> return (url /: "v1", opts) + Nothing -> error $ "could not parse " <> T.unpack base_url <> " as a URI" + listModels = do + apiKey <- lookupEnv $ T.unpack key_name + case apiKey of + Nothing -> return Nothing + Just key -> do + (url, opts) <- apiEndpoint + let headers = header "Authorization" ("Bearer " <> encodeUtf8 (T.pack key)) + response <- runReq defaultHttpConfig $ req GET (url /: "models") NoReqBody bsResponse (headers <> opts) + return $ parseOpenAIModels <$> Aeson.decodeStrict (responseBody response) + + parseOpenAIResponse :: Value -> Maybe Text + parseOpenAIResponse = parseMaybe parseResponse + where + parseResponse :: Value -> Parser Text + parseResponse val = do + obj <- parseJSON val + choices <- obj .: "choices" + case choices of + [] -> fail "No choices in response" + (choice : _) -> do + choiceObj <- parseJSON choice + message <- choiceObj .: "message" + messageObj <- parseJSON message + messageObj .: "content" + + generateFits prompt modelName = do + apiKey <- lookupEnv $ T.unpack key_name + case apiKey of + Nothing -> return $ Left $ "API key not found. Set the " <> T.unpack key_name <> " environment variable." + Just key -> do + let requestBody = + object + [ "model" .= modelName + , "messages" .= [object ["role" .= ("user" :: Text), "content" .= prompt]] + ] + + (url, opts) <- apiEndpoint + let headers = + header "Content-Type" "application/json" + <> header "Authorization" ("Bearer " <> encodeUtf8 (T.pack key)) + + response <- runReq defaultHttpConfig $ + req POST (url /: "chat" /: "completions") (ReqBodyJson requestBody) bsResponse (headers <> opts) + case Aeson.decodeStrict (responseBody response) >>= parseOpenAIResponse of + Just content -> return $ Right content + Nothing -> + return $ Left $ "Failed to parse OpenAI response: " <> show (responseBody response) + +-- | Parse the models from the endpoint +parseOpenAIModels :: Value -> [Text] +parseOpenAIModels value = + fromMaybe [] (parseMaybe parseModels value) + where + extractModelId :: Value -> Parser Text + extractModelId model = do + obj <- parseJSON model + obj .: "id" + + parseModels :: Value -> Parser [Text] + parseModels val = do + obj <- parseJSON val + models <- obj .: "data" + mapM extractModelId models