diff --git a/CHANGELOG.md b/CHANGELOG.md --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,5 +1,9 @@ # Revision history for OllamaHoles +## 0.1.6.0 -- 2024-05-02 + +* Feature: Add model-options parameter + ## 0.1.5.3 -- 2024-05-02 * CI: Add GitHub Actions workflow for building and testing diff --git a/README.md b/README.md --- a/README.md +++ b/README.md @@ -17,6 +17,28 @@ This plugin is also availble on Hackage [https://hackage.haskell.org/package/ollama-holes-plugin](https://hackage.haskell.org/package/ollama-holes-plugin) +## Installation + +1. Install [Ollama](https://ollama.com/download) +2. Install the `qwen3` model (or any other model you prefer) using the following command: + +```bash + +ollama pull qwen3 +``` +3. Clone this repository and navigate to the directory, and build the project using: + +```bash +cabal build +``` +4. Run the example using: + +```bash +cabal build Test +``` + +5. Enjoy! If you want to change the underlying model, make sure to pass the model name via the plugin arguments (see example) + ## Example Given @@ -140,27 +162,17 @@ ... ``` -## Installation - -1. Install [Ollama](https://ollama.com/download) -2. Install the `qwen3` model (or any other model you prefer) using the following command: +## Model Options -```bash +Using -ollama pull qwen3 ``` -3. Clone this repository and navigate to the directory, and build the project using: - -```bash -cabal build +-fplugin-opt=GHC.Plugin.OllamaHoles:model-options={\"num_ctxt\": 32000, \"temperature\": 1.0} ``` -4. Run the example using: -```bash -cabal build Test -``` +You can pass further custom options to the model, e.g. here we increase the context length +and set the temperature. -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 diff --git a/ollama-holes-plugin.cabal b/ollama-holes-plugin.cabal --- a/ollama-holes-plugin.cabal +++ b/ollama-holes-plugin.cabal @@ -20,7 +20,7 @@ -- PVP summary: +-+------- breaking API changes -- | | +----- non-breaking API additions -- | | | +--- code changes with no API change -version: 0.1.5.3 +version: 0.1.6.0 -- A short (one-line) description of the package. synopsis: A typed-hole plugin that uses LLMs to generate valid hole-fits @@ -28,13 +28,13 @@ -- A longer description of the package. 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 include documentation in the LLM's context (not recommended for small models): > -fplugin-opt=GHC.Plugin.OllamaHoles:include-docs @@ -59,14 +59,22 @@ > -fplugin-opt=GHC.Plugin.OllamaHoles:debug + For custom model options: + + > -fplugin-opt=GHC.Plugin.OllamaHoles:model-options= + + Where json is an object with your parameters, e.g.: + + > -fplugin-opt=GHC.Plugin.OllamaHoles:model-options={"num_ctxt": 32000, "temperature": 1.0} + 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, and you can install the default model (qwen3) by running `ollama pull qwen3`. - For the OpenAI backend, you'll need to set the OPENAI_API_KEY environment variable with your API key. + 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. + 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. 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 @@ -50,6 +50,8 @@ import qualified GHC.Tc.Utils.TcType as GHC (tyCoFVsOfType, mkPhiTy) import qualified GHC.Tc.Solver as GHC (simplifyTop, simplifyInfer, captureTopConstraints, InferMode(..)) import qualified GHC.Tc.Solver.Monad as GHC (zonkTcType, runTcSEarlyAbort) +import Data.Aeson (Value) +import qualified Data.Aeson as Aeson -- | Prompt used to prompt the LLM promptTemplate :: Text @@ -107,15 +109,15 @@ imports = tcg_imports gbl_env let backend = getBackend flags available_models <- liftIO $ listModels backend + liftIO $ when debug $ T.putStrLn $ "Running " <> pluginName <> " with flags:" + liftIO $ when debug $ print flags + case available_models of Nothing -> - error $ - "--- " <> T.unpack pluginName <> ": No models available, check your configuration ---" + error $ T.unpack pluginName <> ": No models available, check your configuration" Just models -> do unless (model_name `elem` models) $ - error $ - "--- " - <> T.unpack pluginName + error $ T.unpack pluginName <> ": Model " <> T.unpack model_name <> " not found. " @@ -124,11 +126,10 @@ else "" ) <> "specify another model using " - <> "`-fplugin-opt=GHC.Plugin.OllamaHoles:model=` ---" - <> "--- Availble models: " + <> "`-fplugin-opt=GHC.Plugin.OllamaHoles:model=`\n" + <> "Availble models: \n" <> T.unpack (T.unlines models) - <> " ---" - liftIO $ when debug $ T.putStrLn $ "--- " <> pluginName <> ": Hole Found ---" + liftIO $ when debug $ T.putStrLn $ pluginName <> ": Hole Found" let mn = "Module: " <> mod_name let lc = "Location: " <> showSDoc dflags (ppr $ ctLocSpan . hole_loc <$> th_hole hole) #if __GLASGOW_HASKELL__ >= 912 @@ -161,12 +162,12 @@ , ("{scope}", scope) , ("{docs}", docs) ] - liftIO $ when debug $ do T.putStrLn $ "--- " <> pluginName <> ": Prompt ---\n" <> prompt' - res <- liftIO $ generateFits backend prompt' model_name + liftIO $ when debug $ do T.putStrLn $ pluginName <> " Prompt:\n```\n" <> prompt' <> "\n```" + res <- liftIO $ generateFits backend prompt' model_name model_options case res of Right rsp -> do let lns = (preProcess . T.lines) rsp - liftIO $ when debug $ do T.putStrLn $ "--- " <> pluginName <> ": Response ---\n" <> rsp + liftIO $ when debug $ do T.putStrLn $ pluginName <> " Response:\n```\n" <> rsp <> "\n```" verified <- filterM (verifyHoleFit debug hole) lns let fits' = map (RawHoleFit . text . T.unpack) verified -- Return the generated fits @@ -203,8 +204,7 @@ case parsed of Left err_msg-> do liftIO $ when debug $ do - putStrLn "--- Error when validating: ---" - T.putStrLn fit + T.putStrLn $ "Error during validation of " <> fit putStrLn err_msg return False Right p_e -> do @@ -276,7 +276,8 @@ , include_docs :: Bool , openai_base_url :: Text , openai_key_name :: Text - } + , model_options :: Maybe Value + } deriving (Show) -- | Default flags for the plugin defaultFlags :: Flags @@ -289,6 +290,7 @@ , include_docs = False , openai_base_url = "https://api.openai.com" , openai_key_name = "OPENAI_API_KEY" + , model_options = Nothing } -- | Produce the documentation of all the HolefitCandidates. @@ -358,6 +360,13 @@ | T.isPrefixOf "n=" (T.pack opt) = let num_expr = T.unpack $ T.drop (T.length "n=") (T.pack opt) in parseFlags' flags{num_expr = read num_expr} opts + parseFlags' flags (opt : opts) + | T.isPrefixOf "model-options=" (T.pack opt) = + let model_options = T.drop (T.length "model-options=") (T.pack opt) + m_opts = Aeson.eitherDecodeStrictText model_options + in case m_opts of + Right o -> parseFlags' flags{model_options = Just o } opts + Left err -> error $ "Failed to parse model-options: " <> err parseFlags' flags _ = flags -- | Helper function to replace placeholders in a template string diff --git a/src/GHC/Plugin/OllamaHoles/Backend.hs b/src/GHC/Plugin/OllamaHoles/Backend.hs --- a/src/GHC/Plugin/OllamaHoles/Backend.hs +++ b/src/GHC/Plugin/OllamaHoles/Backend.hs @@ -2,10 +2,11 @@ module GHC.Plugin.OllamaHoles.Backend (Backend(..)) where import Data.Text (Text) +import Data.Aeson (Value) -- | Backend to use. data Backend = Backend { listModels :: IO (Maybe [Text]) - , generateFits :: Text -> Text -> IO (Either String Text) + , generateFits :: Text -> Text -> Maybe Value -> IO (Either String Text) } diff --git a/src/GHC/Plugin/OllamaHoles/Backend/Gemini.hs b/src/GHC/Plugin/OllamaHoles/Backend/Gemini.hs --- a/src/GHC/Plugin/OllamaHoles/Backend/Gemini.hs +++ b/src/GHC/Plugin/OllamaHoles/Backend/Gemini.hs @@ -7,7 +7,7 @@ import Network.HTTP.Req import System.Environment (lookupEnv) -import Data.Aeson (FromJSON (..), Value, object, parseJSON, (.:), (.=)) +import Data.Aeson (FromJSON (..), Value (..), object, parseJSON, (.:), (.=)) import Data.Aeson.Types (Parser, parseMaybe) import Data.Text (Text) @@ -48,16 +48,18 @@ partObj <- parseJSON part partObj .: "text" - generateFits prompt modelName = do + generateFits prompt modelName options = 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 = + let base_req = object [ "contents" .= [object ["parts" .= [object ["text" .= prompt]]]] ] - + requestBody = case options of + Just (Object opts) | Object base <- base_req -> Object (base <> opts) + _ -> base_req let url = apiEndpoint /: "models" /: (modelName <> ":generateContent") params = "key" =: key headers = header "Content-Type" "application/json" diff --git a/src/GHC/Plugin/OllamaHoles/Backend/Ollama.hs b/src/GHC/Plugin/OllamaHoles/Backend/Ollama.hs --- a/src/GHC/Plugin/OllamaHoles/Backend/Ollama.hs +++ b/src/GHC/Plugin/OllamaHoles/Backend/Ollama.hs @@ -15,13 +15,8 @@ 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} + generateFits prompt modelName options = do + fmap Ollama.response_ <$> Ollama.generate Ollama.defaultGenerateOps{prompt = prompt, + modelName = modelName, + options = options} - -- \| 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 --- a/src/GHC/Plugin/OllamaHoles/Backend/OpenAI.hs +++ b/src/GHC/Plugin/OllamaHoles/Backend/OpenAI.hs @@ -7,7 +7,7 @@ import Network.HTTP.Req import System.Environment (lookupEnv) -import Data.Aeson (FromJSON (..), Value, object, parseJSON, (.:), (.=)) +import Data.Aeson (FromJSON (..), Value (Object), object, parseJSON, (.:), (.=)) import Data.Aeson.Types (Parser, parseMaybe) import qualified Data.Aeson as Aeson @@ -56,17 +56,19 @@ messageObj <- parseJSON message messageObj .: "content" - generateFits prompt modelName = do + generateFits prompt modelName options = 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 + let base_req = + object [ "model" .= modelName , "messages" .= [object ["role" .= ("user" :: Text), "content" .= prompt]] ] - + requestBody = case options of + Just (Object opts) | Object base <- base_req -> Object (base <> opts) + _ -> base_req (url, opts) <- apiEndpoint let headers = header "Content-Type" "application/json"