@@ -310,34 +310,34 @@ def load(self):
310310
311311 if self .model .startswith ("ollama/" ):
312312 model_name = self .model .replace ("ollama/" , "" )
313+ api_base = getattr (self , 'api_base' , None ) or "http://localhost:11434"
314+ names = []
313315 try :
314316 # List out all downloaded ollama models. Will fail if ollama isn't installed
315- result = subprocess .run (
316- ["ollama" , "list" ], capture_output = True , text = True , check = True
317- )
317+ response = requests .get (f"{ api_base } /api/tags" )
318+ if response .ok :
319+ data = response .json ()
320+ names = [
321+ model ['name' ].replace (":latest" , "" )
322+ for model in data ['models' ]
323+ if 'name' in model and model ['name' ]
324+ ]
325+
318326 except Exception as e :
319327 print (str (e ))
320328 self .interpreter .display_message (
321329 f"> Ollama not found\n \n Please download Ollama from [ollama.com](https://ollama.com/) to use `{ model_name } `.\n "
322330 )
323331 exit ()
324332
325- lines = result .stdout .split ("\n " )
326- names = [
327- line .split ()[0 ].replace (":latest" , "" )
328- for line in lines [1 :]
329- if line .strip ()
330- ] # Extract names, trim out ":latest", skip header
331-
333+ # Download model if not already installed
332334 if model_name not in names :
333335 self .interpreter .display_message (f"\n Downloading { model_name } ...\n " )
334- subprocess . run ([ "ollama " , "pull" , model_name ], check = True )
336+ requests . post ( f" { api_base } /api/pull " , json = { "name" : model_name } )
335337
336338 # Get context window if not set
337339 if self .context_window == None :
338- response = requests .post (
339- "http://localhost:11434/api/show" , json = {"name" : model_name }
340- )
340+ response = requests .post (f"{ api_base } /api/show" , json = {"name" : model_name })
341341 model_info = response .json ().get ("model_info" , {})
342342 context_length = None
343343 for key in model_info :
0 commit comments