feat: add temp flag

assistant.py

@@ -6,7 +6,8 @@ import pyperclip
 import sys
 import argparse
 
-model = 'mistral-nemo'
+model = 'llama3.1:8b-instruct-q8_0'
+temp = 0.2
 
 pattern = r'```[a-z]*\n[\s\S]*?\n```'
 line_pattern = r'`[a-z]*[\s\S]*?`'
@@ -49,6 +50,7 @@ def chat(message, stream=True):
     history.append({"role": "user", "content": message})
     completion = client.chat(
             model=model,
+            options={"temperature":temp},
             messages=history,
             stream=stream
             )
@@ -56,7 +58,9 @@ def chat(message, stream=True):
     for chunk in completion:
         if stream:
             print(chunk['message']['content'], end='', flush=True)
-        result += chunk['message']['content']
+            result += chunk['message']['content']
+    if not stream:
+        result = completion['message']['content']
     if stream:
         print()
     history.append({"role": 'assistant', 'content': result})
@@ -73,6 +77,8 @@ def parse_args():
     parser.add_argument('--shell', '-s', nargs='?', const=True, default=False, help='output a shell command that does as described')
     # Add the --model (-m) argument
     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')
     # Parse the arguments
     return parser.parse_args()
 
@@ -85,7 +91,7 @@ def arg_follow_up(args):
     return second_input
 
 def arg_shell(args):
-    query = 'Form a shell command based on the following description. Only output a working shell command and then <|eot_id|>.\nDescription: '
+    query = 'Form a shell command based on the following description. Only output a working shell command .\nDescription: '
     if args.shell != True:
         query += args.shell
     else:
@@ -131,6 +137,9 @@ def main():
     if args.model:
         global model 
         model = args.model
+    if args.temp:
+        global temp
+        temp = float(args.temp)
     if not sys.stdin.isatty():
         handle_piped_input(args)
     else: