forked from ThePrimeagen/99
-
Notifications
You must be signed in to change notification settings - Fork 0
Expand file tree
/
Copy pathinit.lua
More file actions
526 lines (473 loc) · 14 KB
/
init.lua
File metadata and controls
526 lines (473 loc) · 14 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
local Logger = require("99.logger.logger")
local Level = require("99.logger.level")
local ops = require("99.ops")
local Languages = require("99.language")
local Window = require("99.window")
local get_id = require("99.id")
local RequestContext = require("99.request-context")
local geo = require("99.geo")
local Range = geo.Range
local Point = geo.Point
local Extensions = require("99.extensions")
local Agents = require("99.extensions.agents")
local Providers = require("99.providers")
local time = require("99.time")
---@param path_or_rule string | _99.Agents.Rule
---@return _99.Agents.Rule | string
local function expand(path_or_rule)
if type(path_or_rule) == "string" then
return vim.fn.expand(path_or_rule)
end
return {
name = path_or_rule.name,
path = vim.fn.expand(path_or_rule.path),
}
end
--- @param opts _99.ops.Opts?
--- @return _99.ops.Opts
local function process_opts(opts)
opts = opts or {}
for i, rule in ipairs(opts.additional_rules or {}) do
local r = expand(rule)
assert(
type(r) ~= "string",
"broken configuration. additional_rules must never be a string"
)
opts.additional_rules[i] = r
end
return opts
end
--- @alias _99.Cleanup fun(): nil
--- @class _99.RequestEntry
--- @field id number
--- @field operation string
--- @field status "running" | "success" | "failed" | "cancelled"
--- @field filename string
--- @field lnum number
--- @field col number
--- @field started_at number
--- @class _99.ActiveRequest
--- @field clean_up _99.Cleanup
--- @field request_id number
--- @class _99.StateProps
--- @field model string
--- @field md_files string[]
--- @field prompts _99.Prompts
--- @field ai_stdout_rows number
--- @field languages string[]
--- @field display_errors boolean
--- @field auto_add_skills boolean
--- @field provider_override _99.Providers.BaseProvider?
--- @field __active_requests table<number, _99.ActiveRequest>
--- @field __view_log_idx number
--- @field __request_history _99.RequestEntry[]
--- @field __request_by_id table<number, _99.RequestEntry>
--- @return _99.StateProps
local function create_99_state()
return {
model = "opencode/claude-sonnet-4-5",
md_files = {},
prompts = require("99.prompt-settings"),
ai_stdout_rows = 3,
languages = { "lua", "go", "java", "elixir", "cpp", "ruby" },
display_errors = false,
provider_override = nil,
auto_add_skills = false,
__active_requests = {},
__view_log_idx = 1,
__request_history = {},
__request_by_id = {},
}
end
--- @class _99.Completion
--- @field source "cmp" | nil
--- @field custom_rules string[]
--- @class _99.Options
--- @field logger _99.Logger.Options?
--- @field model string?
--- @field md_files string[]?
--- @field provider _99.Providers.BaseProvider?
--- @field debug_log_prefix string?
--- @field display_errors? boolean
--- @field auto_add_skills? boolean
--- @field completion _99.Completion?
--- unanswered question -- will i need to queue messages one at a time or
--- just send them all... So to prepare ill be sending around this state object
--- @class _99.State
--- @field completion _99.Completion
--- @field model string
--- @field md_files string[]
--- @field prompts _99.Prompts
--- @field ai_stdout_rows number
--- @field languages string[]
--- @field display_errors boolean
--- @field provider_override _99.Providers.BaseProvider?
--- @field auto_add_skills boolean
--- @field rules _99.Agents.Rules
--- @field __active_requests table<number, _99.ActiveRequest>
--- @field __view_log_idx number
--- @field __request_history _99.RequestEntry[]
--- @field __request_by_id table<number, _99.RequestEntry>
local _99_State = {}
_99_State.__index = _99_State
--- @return _99.State
function _99_State.new()
local props = create_99_state()
---@diagnostic disable-next-line: return-type-mismatch
return setmetatable(props, _99_State)
end
--- TODO: This is something to understand. I bet that this is going to need
--- a lot of performance tuning. I am just reading every file, and this could
--- take a decent amount of time if there are lots of rules.
---
--- Simple perfs:
--- 1. read 4096 bytes at a tiem instead of whole file and parse out lines
--- 2. don't show the docs
--- 3. do the operation once at setup instead of every time.
--- likely not needed to do this all the time.
function _99_State:refresh_rules()
self.rules = Agents.rules(self)
Extensions.refresh(self)
end
--- @param context _99.RequestContext
--- @return _99.RequestEntry
function _99_State:track_request(context)
local point = context.range and context.range.start or Point:from_cursor()
local entry = {
id = context.xid,
operation = context.operation or "request",
status = "running",
filename = context.full_path,
lnum = point.row,
col = point.col,
started_at = time.now(),
}
table.insert(self.__request_history, entry)
self.__request_by_id[entry.id] = entry
return entry
end
--- @param id number
--- @param status "success" | "failed" | "cancelled"
function _99_State:finish_request(id, status)
local entry = self.__request_by_id[id]
if entry then
entry.status = status
end
end
--- @param id number
function _99_State:remove_request(id)
for i, entry in ipairs(self.__request_history) do
if entry.id == id then
table.remove(self.__request_history, i)
break
end
end
self.__request_by_id[id] = nil
end
--- @return number
function _99_State:previous_request_count()
local count = 0
for _, entry in ipairs(self.__request_history) do
if entry.status ~= "running" then
count = count + 1
end
end
return count
end
function _99_State:clear_previous_requests()
local keep = {}
for _, entry in ipairs(self.__request_history) do
if entry.status == "running" then
table.insert(keep, entry)
else
self.__request_by_id[entry.id] = nil
end
end
self.__request_history = keep
end
local _active_request_id = 0
---@param clean_up _99.Cleanup
---@param request_id number
---@return number
function _99_State:add_active_request(clean_up, request_id)
_active_request_id = _active_request_id + 1
Logger:debug("adding active request", "id", _active_request_id)
self.__active_requests[_active_request_id] = {
clean_up = clean_up,
request_id = request_id,
}
return _active_request_id
end
function _99_State:active_request_count()
local count = 0
for _ in pairs(self.__active_requests) do
count = count + 1
end
return count
end
---@param id number
function _99_State:remove_active_request(id)
local logger = Logger:set_id(id)
local r = self.__active_requests[id]
logger:assert(r, "there is no active request for id. implementation broken")
logger:debug("removing active request")
self.__active_requests[id] = nil
end
local _99_state = _99_State.new()
--- @class _99
local _99 = {
DEBUG = Level.DEBUG,
INFO = Level.INFO,
WARN = Level.WARN,
ERROR = Level.ERROR,
FATAL = Level.FATAL,
}
--- you can only set those marks after the visual selection is removed
local function set_selection_marks()
vim.api.nvim_feedkeys(
vim.api.nvim_replace_termcodes("<Esc>", true, false, true),
"x",
false
)
end
--- @param cb fun(ok: boolean, o: _99.ops.Opts?): nil
--- @param context _99.RequestContext
--- @param opts _99.ops.Opts
--- @return fun(ok: boolean, response: string): nil
local function wrap_window_capture(cb, context, opts)
--- @param ok boolean
--- @param response string
return function(ok, response)
context.logger:debug("capture_prompt", "success", ok, "response", response)
if not ok then
return cb(false)
end
local rules_and_names = Agents.by_name(_99_state.rules, response)
opts.additional_rules = opts.additional_rules or {}
for _, r in ipairs(rules_and_names.rules) do
table.insert(opts.additional_rules, r)
end
opts.additional_prompt = response
cb(true, opts)
end
end
--- @param operation_name string
--- @return _99.RequestContext
local function get_context(operation_name)
_99_state:refresh_rules()
local trace_id = get_id()
local context = RequestContext.from_current_buffer(_99_state, trace_id)
context.operation = operation_name
context.logger:debug("99 Request", "method", operation_name)
return context
end
function _99.info()
local info = {}
_99_state:refresh_rules()
table.insert(
info,
string.format("Previous Requests: %d", _99_state:previous_request_count())
)
table.insert(
info,
string.format("custom rules(%d):", #(_99_state.rules.custom or {}))
)
for _, rule in ipairs(_99_state.rules.custom or {}) do
table.insert(info, string.format("* %s", rule.name))
end
Window.display_centered_message(info)
end
--- @param path string
function _99:rule_from_path(path)
_ = self
path = expand(path) --[[ @as string]]
return Agents.get_rule_by_path(_99_state.rules, path)
end
--- @param opts? _99.ops.Opts
function _99.fill_in_function_prompt(opts)
opts = process_opts(opts)
local context = get_context("fill-in-function-with-prompt")
context.logger:debug("start")
Window.capture_input({
cb = wrap_window_capture(function(ok, o)
if not ok then
return
end
assert(o ~= nil, "if ok, then opts must exist")
ops.fill_in_function(context, o)
end, context, opts),
on_load = function()
Extensions.setup_buffer(_99_state)
end,
rules = _99_state.rules,
})
end
--- @param opts? _99.ops.Opts
function _99.fill_in_function(opts)
opts = process_opts(opts)
ops.fill_in_function(get_context("fill_in_function"), opts)
end
--- @param opts _99.ops.Opts
function _99.visual_prompt(opts)
opts = process_opts(opts)
local context = get_context("over-range-with-prompt")
context.logger:debug("start")
Window.capture_input({
cb = wrap_window_capture(function(ok, o)
if not ok then
return
end
assert(o ~= nil, "if ok, then opts must exist")
_99.visual(context, o)
end, context, opts),
on_load = function()
Extensions.setup_buffer(_99_state)
end,
rules = _99_state.rules,
})
end
--- @param context _99.RequestContext?
--- @param opts _99.ops.Opts?
function _99.visual(context, opts)
opts = process_opts(opts)
--- TODO: Talk to teej about this.
--- Visual selection marks are only set in place post visual selection.
--- that means for this function to work i must escape out of visual mode
--- which i dislike very much. because maybe you dont want this
set_selection_marks()
context = context or get_context("over-range")
local range = Range.from_visual_selection()
ops.over_range(context, range, opts)
end
--- View all the logs that are currently cached. Cached log count is determined
--- by _99.Logger.Options that are passed in.
function _99.view_logs()
_99_state.__view_log_idx = 1
local logs = Logger.logs()
if #logs == 0 then
print("no logs to display")
return
end
Window.display_full_screen_message(logs[1])
end
function _99.prev_request_logs()
local logs = Logger.logs()
if #logs == 0 then
print("no logs to display")
return
end
_99_state.__view_log_idx = math.min(_99_state.__view_log_idx + 1, #logs)
Window.display_full_screen_message(logs[_99_state.__view_log_idx])
end
function _99.next_request_logs()
local logs = Logger.logs()
if #logs == 0 then
print("no logs to display")
return
end
_99_state.__view_log_idx = math.max(_99_state.__view_log_idx - 1, 1)
Window.display_full_screen_message(logs[_99_state.__view_log_idx])
end
function _99.stop_all_requests()
for _, active in pairs(_99_state.__active_requests) do
_99_state:remove_request(active.request_id)
active.clean_up()
end
_99_state.__active_requests = {}
end
function _99.previous_requests_to_qfix()
local items = {}
for _, entry in ipairs(_99_state.__request_history) do
table.insert(items, {
filename = entry.filename,
lnum = entry.lnum,
col = entry.col,
text = string.format("[%s] %s", entry.status, entry.operation),
})
end
vim.fn.setqflist({}, "r", { title = "99 Requests", items = items })
vim.cmd("copen")
end
function _99.clear_previous_requests()
_99_state:clear_previous_requests()
end
--- if you touch this function you will be fired
--- @return _99.State
function _99.__get_state()
return _99_state
end
--- @param opts _99.Options?
function _99.setup(opts)
opts = opts or {}
_99_state = _99_State.new()
_99_state.provider_override = opts.provider
_99_state.completion = opts.completion
or {
source = nil,
custom_rules = {},
}
_99_state.completion.custom_rules = _99_state.completion.custom_rules or {}
_99_state.auto_add_skills = opts.auto_add_skills or false
local crules = _99_state.completion.custom_rules
for i, rule in ipairs(crules) do
local str = expand(rule)
assert(type(str) == "string", "rule path must be a string")
crules[i] = str
end
vim.api.nvim_create_autocmd("VimLeavePre", {
callback = function()
_99.stop_all_requests()
end,
})
Logger:configure(opts.logger)
if opts.model then
assert(type(opts.model) == "string", "opts.model is not a string")
_99_state.model = opts.model
else
local provider = opts.provider or Providers.OpenCodeProvider
if provider._get_default_model then
_99_state.model = provider._get_default_model()
end
end
if opts.md_files then
assert(type(opts.md_files) == "table", "opts.md_files is not a table")
for _, md in ipairs(opts.md_files) do
_99.add_md_file(md)
end
end
_99_state.display_errors = opts.display_errors or false
_99_state:refresh_rules()
Languages.initialize(_99_state)
Extensions.init(_99_state)
end
--- @param md string
--- @return _99
function _99.add_md_file(md)
table.insert(_99_state.md_files, md)
return _99
end
--- @param md string
--- @return _99
function _99.rm_md_file(md)
for i, name in ipairs(_99_state.md_files) do
if name == md then
table.remove(_99_state.md_files, i)
break
end
end
return _99
end
--- @param model string
--- @return _99
function _99.set_model(model)
_99_state.model = model
return _99
end
function _99.__debug()
Logger:configure({
path = nil,
level = Level.DEBUG,
})
end
_99.Providers = Providers
return _99