Fix --host argument

assistant.py

@@ -10,6 +10,7 @@ from pygments.lexers import get_lexer_by_name
 from pygments.formatters import TerminalFormatter
 
 
+server = 'localhost:11434'
 model = 'llama3.1:8b-instruct-q8_0'
 temp = 0.2
 
@@ -87,8 +88,6 @@ def copy_string_to_clipboard(string):
     except:
         return
 
-# Point to the local server
-client = Client(host='localhost:11434')
 
 code_history = []
 
@@ -150,9 +149,15 @@ def parse_args():
     parser.add_argument('--model', '-m', nargs='?', const=True, default=False, help='Specify model')
     # Add the --temp (-t) argument
     parser.add_argument('--temp', '-t', nargs='?', const=True, default=False, help='Specify temperature')
+    # Add the --host (-h) argument
+    parser.add_argument('--host', nargs='?', const=True, default=False, help='Specify host of ollama server')
     # Parse the arguments
     return parser.parse_args()
 
+def set_host(host):
+    global server
+    server = host
+
 def arg_follow_up(args):
     sys.stdin = open('/dev/tty')
     if args.follow_up != True:
@@ -203,8 +208,15 @@ def handle_non_piped_input(args):
         else:
             chat(user_input)
 
+client = None
+
 def main():
     args = parse_args()
+    if args.host:
+        set_host(args.host)
+    # Point to the local server
+    global client
+    client = Client(host=server)
     if args.model:
         global model 
         model = args.model