Feature/add cmps

lua/docscribe/config.lua

@@ -7,16 +7,28 @@ local prompt_templates = require("docscribe.prompt_templates")
 local config = {
     ui = {
         highlight = {
-            style = "signature",        -- "signature" | "full" | "none" function highlight
-            timeout = 2000,             -- Time (ms) before highlight fades
-            bg = "#545454",             -- Highlight background color
+            style = "signature", -- "signature" | "full" | "none" function highlight
+            timeout = 2000,      -- Time (ms) before highlight fades
+            bg = "#545454",      -- Highlight background color
         },
     },
     llm = {
-        provider = "ollama",            -- Backend used for LLM (e.g., ollama, openai)
-        model = "llama3.2",             -- Default model used for docs
+        provider = "ollama", -- Backend used for LLM (e.g., ollama, google, groq)
+        provider_opts = {
+            ollama = {
+                model = "llama3.2", -- Default model used for docs
+            },
+            google = {
+                model = "gemini-1.5-flash",
+                api_key = os.getenv("GOOGLE_API_KEY"), -- API key for Google's Gemini
+            },
+            groq = {
+                model = "llama-3.1-8b-instant",      -- Default model used for docs
+                api_key = os.getenv("GROQ_API_KEY"), -- API key for Groq
+            },
+        },
     },
-    prompt_templates = {                -- Set of prompt templates for each programming language
+    prompt_templates = { -- Set of prompt templates for each programming language
         default = prompt_templates.default_template,
         h = prompt_templates.c_template,
         c = prompt_templates.c_template,

lua/docscribe/llm/init.lua

@@ -2,7 +2,12 @@
 --- Generates documentation using LLM providers.
 
 local config = require("docscribe.config")
-local ollama = require("docscribe.llm.providers.ollama")
+
+local providers = {
+    ollama = require("docscribe.llm.providers.ollama"),
+    google = require("docscribe.llm.providers.google"),
+    groq = require("docscribe.llm.providers.groq"),
+}
 
 local M = {}
 
@@ -16,12 +21,16 @@ function M.generate_docs(function_code, callback)
     local prompt_template = prompt_templates[lang] or prompt_templates.default
     local prompt = prompt_template:gsub("{{code}}", function_code)
 
-    local llm = config.get_config("llm")
+    local llm_config = config.get_config("llm")
+    local provider_name = llm_config.provider
+    local provider_opts = llm_config.provider_opts[provider_name]
+
+    local provider = providers[provider_name]
 
-    if llm.provider == "ollama" then
-        ollama.generate_response(prompt, llm.model, callback)
+    if provider then
+        provider.generate_response(prompt, callback, provider_opts)
     else
-        callback(nil, 'Invalid LLM runner "' .. llm.provider .. '"')
+        callback(nil, 'Invalid LLM provider "' .. provider_name .. '"')
     end
 end
 

lua/docscribe/llm/providers/google.lua

@@ -0,0 +1,88 @@
+--- @module "docscribe.llm.providers.google"
+--- Interface to the Google LLM provider.
+
+local Job = require("plenary.job")
+
+local M = {}
+
+--- Handles the exit of the curl job for the Google provider.
+--- @param job table The job object.
+--- @param code number The exit code.
+--- @param callback fun(docs: string|nil, err_msg: string|nil) The callback function.
+local function handle_job_exit(job, code, callback)
+    if code ~= 0 then
+        local error_message = table.concat(job:stderr_result(), "")
+        callback(nil, "Error generating Google response: " .. error_message)
+        return
+    end
+
+    local result = table.concat(job:result(), "")
+    local ok, data = pcall(vim.json.decode, result)
+
+    if not ok then
+        callback(nil, "Failed to decode Google API response: " .. tostring(data))
+        return
+    end
+
+    if data.error then
+        callback(nil, "Google API error: " .. data.error.message)
+        return
+    end
+
+    if
+        data.candidates
+        and data.candidates[1]
+        and data.candidates[1].content
+        and data.candidates[1].content.parts
+        and data.candidates[1].content.parts[1]
+    then
+        callback(data.candidates[1].content.parts[1].text)
+    else
+        callback(nil, "Could not extract text from Google API response")
+    end
+end
+
+--- Generates a response using the Google LLM provider.
+--- @param prompt string The prompt to send to Google.
+--- @param callback fun(docs: string|nil, err_msg: string|nil) Callback with the response or an error.
+--- @param opts table The options for the Google provider. Must contain `model` and `api_key`.
+function M.generate_response(prompt, callback, opts)
+    if not opts.api_key or opts.api_key == "" then
+        callback(nil, "Google API key is not configured. Set the GOOGLE_API_KEY environment variable.")
+        return
+    end
+
+    local url = "https://generativelanguage.googleapis.com/v1beta/models/" .. opts.model .. ":generateContent"
+
+    local body = {
+        contents = {
+            {
+                parts = {
+                    {
+                        text = prompt,
+                    },
+                },
+            },
+        },
+    }
+
+    Job:new({
+        command = "curl",
+        args = {
+            "-X",
+            "POST",
+            "-H",
+            "Content-Type: application/json",
+            "-H",
+            "x-goog-api-key: " .. opts.api_key,
+            url,
+            "-d",
+            vim.json.encode(body),
+        },
+        on_exit = function(job, code)
+            handle_job_exit(job, code, callback)
+        end,
+    }):start()
+end
+
+return M

lua/docscribe/llm/providers/groq.lua

@@ -0,0 +1,81 @@
+--- @module "docscribe.llm.providers.groq"
+--- Interface to the Groq LLM provider.
+
+local Job = require("plenary.job")
+
+local M = {}
+
+--- Handles the exit of the curl job for the Groq provider.
+--- @param job table The job object.
+--- @param code number The exit code.
+--- @param callback fun(docs: string|nil, err_msg: string|nil) The callback function.
+local function handle_job_exit(job, code, callback)
+    if code ~= 0 then
+        local error_message = table.concat(job:stderr_result(), "")
+        callback(nil, "Error generating Groq response: " .. error_message)
+        return
+    end
+
+    local result = table.concat(job:result(), "")
+    local ok, data = pcall(vim.json.decode, result)
+
+    if not ok then
+        callback(nil, "Failed to decode Groq API response: " .. tostring(data))
+        return
+    end
+
+    if data.error then
+        callback(nil, "Groq API error: " .. data.error.message)
+        return
+    end
+
+    if data.choices and data.choices[1] and data.choices[1].message and data.choices[1].message.content then
+        callback(data.choices[1].message.content)
+    else
+        callback(nil, "Could not extract text from Groq API response")
+    end
+end
+
+--- Generates a response using the Groq LLM provider.
+--- @param prompt string The prompt to send to Groq.
+--- @param callback fun(docs: string|nil, err_msg: string|nil) Callback with the response or an error.
+--- @param opts table The options for the Groq provider. Must contain `model` and `api_key`.
+function M.generate_response(prompt, callback, opts)
+    if not opts.api_key or opts.api_key == "" then
+        callback(nil, "Groq API key is not configured. Set the GROQ_API_KEY environment variable.")
+        return
+    end
+
+    local url = "https://api.groq.com/openai/v1/chat/completions"
+
+    local body = {
+        model = opts.model,
+        messages = {
+            {
+                role = "user",
+                content = prompt,
+            },
+        },
+    }
+
+    Job:new({
+        command = "curl",
+        args = {
+            "-X",
+            "POST",
+            "-H",
+            "Content-Type: application/json",
+            "-H",
+            "Authorization: Bearer " .. opts.api_key,
+            url,
+            "-d",
+            vim.json.encode(body),
+        },
+        on_exit = function(job, code)
+            handle_job_exit(job, code, callback)
+        end,
+    }):start()
+end
+
+return M
+

lua/docscribe/llm/providers/ollama.lua

@@ -5,23 +5,31 @@ local Job = require("plenary.job")
 
 local M = {}
 
+--- Handles the exit of the Ollama job.
+--- @param job table The job object.
+--- @param code number The exit code.
+--- @param callback fun(docs: string|nil, err_msg: string|nil) The callback function.
+local function handle_job_exit(job, code, callback)
+    if code ~= 0 then
+        callback(nil, "Error generating Ollama response")
+        return
+    end
+
+    local docs = table.concat(job:result(), "\n")
+    callback(docs)
+end
+
 --- Generates a response using the Ollama LLM provider.
 --- @param prompt string The prompt to send to Ollama.
---- @param model string The Ollama model to use.
 --- @param callback fun(docs: string|nil, err_msg: string|nil) Callback with the response or an error.
-function M.generate_response(prompt, model, callback)
-    --- @diagnostic disable-next-line: missing-fields
+--- @param opts table The Ollama model to use.
+function M.generate_response(prompt, callback, opts)
     Job:new({
         command = "ollama",
-        args = { "run", model },
+        args = { "run", opts.model },
         writer = prompt,
         on_exit = function(job, code)
-            if code == 0 then
-                local docs = table.concat(job:result(), "\n")
-                callback(docs)
-            else
-                callback(nil, "Error generating Ollama response")
-            end
+            handle_job_exit(job, code, callback)
         end,
     }):start()
 end

lua/docscribe/prompt_templates/c.lua

@@ -1,21 +1,23 @@
+--- @module "docscribe.prompt_templates.c"
+
 return [[
-    You are a documentation assistant.
+You are an expert C programmer specializing in writing documentation that follows the Doxygen standard.
 
-    Generate a Doxygen-style comment block for the C function below. The output must:
+Your response must contain ONLY the Doxygen comment block.
+Do NOT include the original function code in your response.
+Do NOT wrap the documentation in markdown code fences (e.g., ```).
 
-    1. Start with `/**` and end with `*/`.
-    2. Begin with a `@brief` tag summarizing the function.
-    3. Follow `@brief` with any additional tags, in this order:
-       - `@param` for each parameter (name, type, purpose).
-       - `@return` to describe the return value.
-       - `@example` for one or more usage examples.
-    4. Document all function parameters using `@param`, describing their name, type, and purpose.
-    5. Include a `@return` tag describing the return value.
-    6. If needed, include an `@example` showing how to use the function.
+Generate a Doxygen-style comment block for the following C function.
 
-    Do NOT include the function code or any Markdown formatting.
+Follow these rules:
+1.  Use `/** ... */` for the comment block.
+2.  The main description should be a single, brief paragraph.
+3.  Use `@param` for parameters and `@return` for the return value.
+4.  If the function's logic is complex, add a `@note` section to explain important details or non-obvious usage. Do not use `@example`.
+5.  Add a blank line between the main description and the tag sections.
 
-    ```c
-    {{code}}
-    ```
+Code:
+```c
+{{code}}
+```
 ]]

lua/docscribe/prompt_templates/cpp.lua

@@ -1,23 +1,23 @@
+--- @module "docscribe.prompt_templates.cpp"
+
 return [[
-    You are a documentation assistant.
+You are an expert C++ programmer specializing in writing documentation that follows the Doxygen standard.
 
-    Generate a Doxygen-style comment block for the C++ function below. The output must:
+Your response must contain ONLY the Doxygen comment block.
+Do NOT include the original function code in your response.
+Do NOT wrap the documentation in markdown code fences (e.g., ```).
 
-    1. Start with `/**` and end with `*/`.
-    2. Begin with a `@brief` tag summarizing the function.
-    3. Follow `@brief` with any additional tags, in this order:
-       - `@param` for each parameter (name, type, purpose).
-       - `@return` to describe the return value.
-       - `@throws` for each possible exception.
-       - `@example` for one or more usage examples.
-    4. Document all function parameters using `@param`, describing their name, type, and purpose.
-    5. Include a `@return` tag describing the return value.
-    6. If the function may throw exceptions, use `@throws` for each possible exception.
-    7. Include an `@example` block demonstrating typical usage.
+Generate a Doxygen-style comment block for the following C++ function.
 
-    Do NOT include the function code or any Markdown formatting.
+Follow these rules:
+1.  Use `/** ... */` for the comment block.
+2.  The main description should be a single, brief paragraph.
+3.  Use `@param` for parameters and `@return` for the return value.
+4.  If the function's logic is complex, add a `@note` section to explain important details or non-obvious usage. Do not use `@example`.
+5.  Add a blank line between the main description and the tag sections.
 
-    ```cpp
-    {{code}}
-    ```
+Code:
+```cpp
+{{code}}
+```
 ]]

lua/docscribe/prompt_templates/default.lua

@@ -1,19 +1,23 @@
+--- @module "docscribe.prompt_templates.default"
+
 return [[
-    You are a documentation assistant.
+You are an expert programmer tasked with writing clear, concise, and accurate documentation for a given function.
 
-    Generate a **proper documentation comment block** for the function below. The output must:
+Your response must contain ONLY the documentation block.
+Do NOT include the original function code in your response.
+Do NOT wrap the documentation in markdown code fences (e.g., ```).
 
-    1. Include a **clear and concise function description**.
-    2. Document **each parameter** with the appropriate tags, describing the name, type, and purpose.
-    3. Include a **return value** description with the return type and a brief explanation.
-    4. If the function throws any exceptions, include an exception tag.
-    5. Include an example block showing typical usage.
+Generate documentation for the following code.
 
-    **Do NOT wrap the output in backticks, triple backticks, or any Markdown formatting.**
-    Do **not** include the function code in your output.
-    Just return the unwrapped documentation block as plain text.
+Follow these rules:
+1.  The documentation should be a block comment.
+2.  The main description should be a single, brief paragraph explaining the function's purpose. Do not describe the parameters or return value in the description.
+3.  Use tags (e.g., @param, @return) to describe parameters and return values.
+4.  Only include a usage example if the function's logic is complex or has important edge cases. The example must be meaningful and demonstrate a non-obvious use case. Do not show a trivial call.
+5.  Add a blank line between the main description and the tag sections.
 
-    ```
-    {{code}}
-    ```
+Code:
+```
+{{code}}
+```
 ]]