Simple Ollama Client

I made this simple Ollama client.

ollamaclient.py

import argparse
import commands
import ollama
import sys

parser = argparse.ArgumentParser(
                    prog='ollama-client',
                    description='Simple client for ollama',
                    epilog='made by Michael Norris')

parser.add_argument('--model', default="llama3:latest")
parser.add_argument('--system', default='')

args = parser.parse_args()
model = args.model

class Session:
  def __init__(self, stream=True):
    self.messages = []
    self.stream = stream

  def add_message(self, content, role="user"):
    self.messages.append({'role': role, 'content': content})

session = Session()

if args.system:
  session.add_message(args.system, role="system")

while True:
  inp = input("> ")
  if inp[0] == '/':
    commands.do_command_or_print_error(inp)
    continue

  session.add_message(inp, role="user")
  stream = ollama.chat(model=model, messages = session.messages, stream=session.stream)
  accumulated_stream = ""
  for chunk in stream:
    print(chunk['message']['content'], end='')
    sys.stdout.flush()
    accumulated_stream += chunk['message']['content']
  print()
  session.add_message(accumulated_stream, role="assistant")

And commands.py

def do_command_or_print_error(command):
  if command in commands:
    commands[command].do()
  else:
    Help().do()

class Command:
  def do(self):
    pass

class Help(Command):
  def do(self):
    print('\n'.join(commands.keys()))

commands = {
  '/help': Help()
}

Then just run it with ` $ python3 ollamaclient.py`