1  # Ollama integration for NovaChat
   2  from ollama import Client
   3  import logging
   4  
   5  def list_models() -> list:
   6      try:
   7          client = Client()
   8          return [model['name'] for model in client.list()['models']]
   9      except Exception as e:
  10          logging.error(f"Model listing failed: {str(e)}")
  11          return []
  12  
  13  
  14  def pull_model(model_name: str) -> None:
  15      try:
  16          client = Client()
  17          client.pull(model_name, stream=False)
  18      except Exception as e:
  19          logging.error(f"Model pull failed: {str(e)}")
  20          raise
  21      response = requests.post(
  22          'http://localhost:11434/api/pull',
  23          json={'name': model_name},
  24          stream=True
  25      )
  26      
  27      for line in response.iter_lines():
  28          status = json.loads(line)
  29          if status.get('error'):
  30              raise Exception(status['error'])
  31  
  32  
  33  def query_ollama(prompt, model="cogito", temperature=0.7, max_tokens=150):
  34      """Query local Ollama server using its REST API"""
  35      try:
  36          response = requests.post(
  37              'http://localhost:11434/api/generate',
  38              json={
  39                  'model': model,
  40                  'prompt': prompt,
  41                  'temperature': temperature,
  42                  'max_tokens': max_tokens,
  43                  'stream': False
  44              },
  45              timeout=30
  46          )
  47          response.raise_for_status()
  48          return response.json().get('response', '')
  49      except Exception as e:
  50          return f"Ollama Error: {str(e)}"