@@ -330,7 +330,7 @@ def prompt(
330330 raise GenerativeAIPromptError (str (exc )) from exc
331331
332332
333- class OllamaGenerativeAIModelType (GenerativeAIModelType ):
333+ class OllamaGenerativeAIModelType (BaseOpenAIGenerativeAIModelType ):
334334 type = "ollama"
335335
336336 def get_host (self , workspace = None , settings_override = None ):
@@ -339,46 +339,28 @@ def get_host(self, workspace=None, settings_override=None):
339339 or settings .BASEROW_OLLAMA_HOST
340340 )
341341
342+ def get_api_key (self , workspace = None , settings_override = None ):
343+ return "ollama"
344+
345+ def get_organization (self , workspace = None , settings_override = None ):
346+ return None
347+
348+ def get_base_url (self , workspace = None , settings_override = None ):
349+ host = self .get_host (workspace , settings_override )
350+ return f"{ host } /v1"
351+
342352 def get_enabled_models (self , workspace = None , settings_override = None ):
343353 workspace_models = self .get_workspace_setting (
344354 workspace , "models" , settings_override
345355 )
346356 return workspace_models or settings .BASEROW_OLLAMA_MODELS
347357
348358 def is_enabled (self , workspace = None , settings_override = None ):
349- ollama_host = self .get_host (workspace , settings_override )
350- return bool (ollama_host ) and bool (
359+ host = self .get_host (workspace , settings_override )
360+ return bool (host ) and bool (
351361 self .get_enabled_models (workspace , settings_override )
352362 )
353363
354- def get_client (self , workspace = None , settings_override = None ):
355- from ollama import Client as OllamaClient
356-
357- ollama_host = self .get_host (workspace , settings_override )
358- return OllamaClient (host = ollama_host )
359-
360- def prompt (
361- self , model , prompt , workspace = None , temperature = None , settings_override = None
362- ):
363- from ollama import RequestError as OllamaRequestError
364- from ollama import ResponseError as OllamaResponseError
365-
366- client = self .get_client (workspace , settings_override )
367- options = {}
368- if temperature :
369- # Because some LLMs can have a temperature of 2, this is the maximum by
370- # default. We're changing it to a maximum of 1 because Ollama only
371- # accepts 1.
372- options ["temperature" ] = min (temperature , 1 )
373- try :
374- response = client .generate (
375- model = model , prompt = prompt , stream = False , options = options
376- )
377- except (OllamaRequestError , OllamaResponseError ) as exc :
378- raise GenerativeAIPromptError (str (exc )) from exc
379-
380- return response ["response" ]
381-
382364 def get_settings_serializer (self ):
383365 from baserow .api .generative_ai .serializers import OllamaSettingsSerializer
384366
0 commit comments