@@ -5,13 +5,12 @@ local log = require("codecompanion.utils.log")
55
66local CONSTANTS = {
77 TIMEOUT = 3000 , -- 3 seconds
8- POLL_INTERVAL = 10 ,
98}
109
11- --- Whether there are already some requests running.
12- local running = false
10+ --- @type table<string , boolean>
11+ local _running = {}
1312
14- M = {}
13+ local M = {}
1514
1615--- @type table<string , table<string , { formatted_name : string ?, opts : { can_reason : boolean , has_vision : boolean , can_use_tools : boolean } } >>
1716local _cached_models = {}
@@ -20,85 +19,104 @@ local _cached_models = {}
2019
2120--- @param url string
2221--- @param opts ? OllamaGetModelsOpts
22+ --- @return table | string | nil
2323local function get_cached_models (url , opts )
24- assert (_cached_models [url ] ~= nil , " Model info is not available in the cache." )
2524 local models = _cached_models [url ]
25+ if not models or vim .tbl_isempty (models ) then
26+ return opts and opts .last and " " or {}
27+ end
2628 if opts and opts .last then
2729 return vim .tbl_keys (models )[1 ]
28- else
29- return models
3030 end
31+ return models
32+ end
33+
34+ --- Build auth headers from adapter env vars
35+ --- @param adapter CodeCompanion.HTTPAdapter
36+ --- @return table
37+ local function build_headers (adapter )
38+ local headers = adapter_utils .set_env_vars (adapter , adapter .headers ) or {}
39+
40+ if adapter .env_replaced .api_key then
41+ local prefix = adapter .env_replaced .authorization or " Bearer"
42+ headers [" Authorization" ] = prefix .. " " .. adapter .env_replaced .api_key
43+ end
44+
45+ return headers
46+ end
47+
48+ --- Parse capabilities from a model info response
49+ --- @param output table The curl response
50+ --- @return { can_reason : boolean , can_use_tools : boolean , has_vision : boolean }
51+ local function parse_capabilities (output )
52+ local opts = {}
53+ if output .status ~= 200 then
54+ return opts
55+ end
56+
57+ local ok , json = pcall (vim .json .decode , output .body , { array = true , object = true })
58+ if not ok then
59+ return opts
60+ end
61+
62+ local capabilities = json .capabilities or {}
63+ opts .can_reason = vim .list_contains (capabilities , " thinking" )
64+ opts .can_use_tools = vim .list_contains (capabilities , " tools" )
65+ opts .has_vision = vim .list_contains (capabilities , " vision" )
66+
67+ return opts
3168end
3269
3370--- Fetch model list and model info.
34- --- Aborts if there's another fetch job running.
35- --- Returns the number of models if the fetches are fired.
71+ --- Aborts if there's another fetch job running for this URL.
3672--- @param adapter CodeCompanion.HTTPAdapter Ollama adapter with env var replaced.
37- --- @param opts OllamaGetModelsOpts
38- local function fetch_async (adapter , opts )
73+ local function fetch_models (adapter )
3974 assert (adapter ~= nil )
40- if running then
75+ local url = adapter .env_replaced .url
76+
77+ if _running [url ] then
4178 return
4279 end
43- local url = adapter .env_replaced .url
44- running = true
80+ _running [url ] = true
4581 _cached_models [url ] = _cached_models [url ] or {}
4682
47- local headers = adapter_utils .set_env_vars (adapter , adapter .headers ) or {}
48-
49- local auth_header = " Bearer "
50- if adapter .env_replaced .authorization then
51- auth_header = adapter .env_replaced .authorization .. " "
52- end
53- if adapter .env_replaced .api_key then
54- headers [" Authorization" ] = auth_header .. adapter .env_replaced .api_key
55- end
83+ local headers = build_headers (adapter )
5684
5785 pcall (function ()
58- local job = Curl .get (url .. " /api/tags" , {
86+ Curl .get (url .. " /api/tags" , {
5987 headers = headers ,
6088 insecure = config .adapters .http .opts .allow_insecure ,
6189 proxy = config .adapters .http .opts .proxy ,
6290 timeout = CONSTANTS .TIMEOUT ,
6391 callback = function (response )
6492 if response .status ~= 200 then
93+ _running [url ] = false
6594 return log :error (" Could not get the Ollama models from " .. url .. " /api/tags.\n Error: %s" , response )
6695 end
6796
6897 local ok , json = pcall (vim .json .decode , response .body )
6998 if not ok then
99+ _running [url ] = false
70100 return log :error (" Could not parse the response from " .. url .. " /api/tags" )
71101 end
72102
73- -- A container for pending requests.
74- -- New jobs are added on creation and removed on completion.
75- local jobs = {}
103+ local pending = {}
76104
77105 for _ , model_obj in ipairs (json .models ) do
78- jobs [model_obj .name ] = Curl .post (url .. " /api/show" , {
106+ pending [model_obj .name ] = Curl .post (url .. " /api/show" , {
107+ body = vim .json .encode ({ model = model_obj .name }),
79108 headers = headers ,
80109 insecure = config .adapters .http .opts .allow_insecure ,
81110 proxy = config .adapters .http .opts .proxy ,
82- body = vim .json .encode ({ model = model_obj .name }),
83111 timeout = CONSTANTS .TIMEOUT ,
84112 callback = function (output )
85- _cached_models [url ][model_obj .name ] = { formatted_name = model_obj .name , opts = {} }
86- if output .status == 200 then
87- local ok , model_info_json = pcall (vim .json .decode , output .body , { array = true , object = true })
88- if ok then
89- _cached_models [url ][model_obj .name ].opts .can_reason =
90- vim .list_contains (model_info_json .capabilities or {}, " thinking" )
91- _cached_models [url ][model_obj .name ].opts .has_vision =
92- vim .list_contains (model_info_json .capabilities or {}, " vision" )
93- _cached_models [url ][model_obj .name ].opts .can_use_tools =
94- vim .list_contains (model_info_json .capabilities or {}, " tools" )
95- end
96- end
97- jobs [model_obj .name ] = nil
98- if vim .tbl_isempty (jobs ) then
99- -- when the last curl request job is removed,
100- -- mark the current `fetch_async` job as finished
101- running = false
113+ _cached_models [url ][model_obj .name ] = {
114+ formatted_name = model_obj .name ,
115+ opts = parse_capabilities (output ),
116+ }
117+ pending [model_obj .name ] = nil
118+ if vim .tbl_isempty (pending ) then
119+ _running [url ] = false
102120 end
103121 end ,
104122 })
@@ -107,8 +125,7 @@ local function fetch_async(adapter, opts)
107125 })
108126 if adapter .opts .cache_adapter == false then
109127 vim .wait (CONSTANTS .TIMEOUT , function ()
110- local models = _cached_models [url ]
111- return models ~= nil and not vim .tbl_isempty (models ) and not running
128+ return not _running [url ]
112129 end )
113130 end
114131 end )
@@ -128,14 +145,13 @@ function M.choices(self, opts)
128145 adapter_utils .get_env_vars (adapter , { timeout = config .adapters .opts .cmd_timeout })
129146 local url = adapter .env_replaced .url
130147 local is_uninitialised = _cached_models [url ] == nil
131-
132148 local should_block = (adapter .opts .cache_adapter == false ) or is_uninitialised or not opts .async
133149
134- fetch_async (adapter , { async = not should_block }) -- should_block means NO async
150+ fetch_models (adapter )
135151
136- if should_block and running then
152+ if should_block and _running [ url ] then
137153 vim .wait (CONSTANTS .TIMEOUT , function ()
138- return not running
154+ return not _running [ url ]
139155 end )
140156 end
141157 return get_cached_models (url , opts )
@@ -150,11 +166,11 @@ function M.check_thinking_capability(self, model)
150166 if type (model ) == " function" then
151167 model = model (self )
152168 end
153- local _choices = self .schema .model .choices
154- if type (_choices ) == " function" then
155- _choices = _choices (self )
169+ local choices = self .schema .model .choices
170+ if type (choices ) == " function" then
171+ choices = choices (self )
156172 end
157- if _choices and _choices [model ] and _choices [model ].opts and _choices [model ].opts .can_reason then
173+ if choices and choices [model ] and choices [model ].opts and choices [model ].opts .can_reason then
158174 return true
159175 end
160176 return false
0 commit comments