From 56885984b2efaad9bf10b8f801539c60539edcad Mon Sep 17 00:00:00 2001
From: Maximilian Dietrich <m.dietrich16@web.de>
Date: Thu, 14 Mar 2024 11:51:44 +0100
Subject: [PATCH] fix(ollama): Fix parameters passthrough

Also add keep_alive, prompt, system prompt and template parameters.
---
 lua/cmp_ai/backends/ollama.lua | 9 ++++++---
 1 file changed, 6 insertions(+), 3 deletions(-)

diff --git a/lua/cmp_ai/backends/ollama.lua b/lua/cmp_ai/backends/ollama.lua
index 604115d..60f282e 100644
--- a/lua/cmp_ai/backends/ollama.lua
+++ b/lua/cmp_ai/backends/ollama.lua
@@ -2,11 +2,11 @@ local requests = require('cmp_ai.requests')
 
 Ollama = requests:new(nil)
 
-function Ollama:new(o, params)
+function Ollama:new(o)
   o = o or {}
   setmetatable(o, self)
   self.__index = self
-  self.params = vim.tbl_deep_extend('keep', params or {}, {
+  self.params = vim.tbl_deep_extend('keep', o or {}, {
     base_url = 'http://127.0.0.1:11434/api/generate',
     model = 'codellama:7b-code',
     options = {
@@ -20,7 +20,10 @@ end
 function Ollama:complete(lines_before, lines_after, cb)
   local data = {
     model = self.params.model,
-    prompt = '<PRE> ' .. lines_before .. ' <SUF>' .. lines_after .. ' <MID>',
+    prompt = self.params.prompt and self.params.prompt(lines_before, lines_after) or '<PRE> ' .. lines_before .. ' <SUF>' .. lines_after .. ' <MID>',
+    keep_alive = self.params.keep_alive,
+    template = self.params.template,
+    system = self.params.system,
     stream = false,
     options = self.params.options,
   }