diff --git a/lua/gemini/api.lua b/lua/gemini/api.lua
index c3f6233..1e157a6 100644
--- a/lua/gemini/api.lua
+++ b/lua/gemini/api.lua
@@ -1,4 +1,4 @@
--- plugin/gemini/api.lua
+-- lua/gemini/api.lua
 
 local M = {}
 
diff --git a/lua/gemini/init.lua b/lua/gemini/init.lua
index 6a43329..e3c50c6 100644
--- a/lua/gemini/init.lua
+++ b/lua/gemini/init.lua
@@ -1,6 +1,7 @@
--- plugin/gemini/init.lua
-print(vim.o.runtimepath)
-require("gemini.api")
+-- lua/gemini/init.lua
+
+local api = require("gemini.api")
+local M = {}
 
 local function gemini_query(prompt)
 	local response = api.get_response(prompt)
@@ -23,36 +24,32 @@ local function gemini_query(prompt)
 	end
 end
 
-vim.api.nvim_create_user_command("Gemini", function(opts)
-	local prompt = opts.args
-	if prompt == "" then
-		vim.notify("Please provide a prompt for Gemini.", vim.log.levels.WARN)
-		return
-	end
-	gemini_query(prompt)
-end, {
-	desc = "Query Google AI",
-	nargs = "+", -- Require at least one argument
-	complete = "shellcmd", -- Optional: Enable completion
-})
+-- Make gemini_query available in M so it can be used by the setup function
+M.gemini_query = gemini_query
 
--- Keymapping example (optional)
-vim.keymap.set("n", "<leader>g", function()
-	vim.ui.input({ prompt = "Gemini Query: " }, function(input)
-		if input then
-			gemini_query(input)
+function M.setup()
+	-- Create the user command
+	vim.api.nvim_create_user_command("Gemini", function(opts)
+		local prompt = opts.args
+		if prompt == "" then
+			vim.notify("Please provide a prompt for Gemini.", vim.log.levels.WARN)
+			return
 		end
-	end)
-end, { desc = "Query Google AI (via Input)" })
+		M.gemini_query(prompt)  -- Use M.gemini_query instead of gemini_query
+	end, {
+		desc = "Query Google AI",
+		nargs = "+",
+		complete = "shellcmd",
+	})
 
--- Add a setup function (even if it's empty)
-local function setup()
-	print("gemini.setup() called") -- For debugging
-	require("gemini.api")
-	print("api module loaded:", api) -- For debugging
-	-- Optional: Add any initialization logic here
+	-- Set up keymapping
+	vim.keymap.set("n", "<leader>g", function()
+		vim.ui.input({ prompt = "Gemini Query: " }, function(input)
+			if input then
+				M.gemini_query(input)  -- Use M.gemini_query instead of gemini_query
+			end
+		end)
+	end, { desc = "Query Google AI (via Input)" })
 end
 
-return {
-	setup = setup,
-}
+return M