Update assistant.py

assistant.py

@@ -1,7 +1,7 @@
 from command_classifier import classify_command
 from executor import execute_command
 
-def get_assistant_response(pipe, command: str, execute: bool) -> str:
+def get_assistant_response(pipe, command: str, execute: bool, max_new_tokens: int = 1024) -> str:
     """
     Process the user command using the language model pipeline and optionally execute a system command.
 
@@ -9,6 +9,7 @@ def get_assistant_response(pipe, command: str, execute: bool) -> str:
         pipe: The Hugging Face model pipeline
         command (str): The user's text input
         execute (bool): Whether or not to execute a real command on the machine
+        max_new_tokens (int): Maximum tokens to generate (longer = more detailed)
 
     Returns:
         str: Response from the assistant (and optionally OS output)
@@ -21,11 +22,17 @@ def get_assistant_response(pipe, command: str, execute: bool) -> str:
         if exec_output:
             return f"🛠️ Executed:\n{exec_output}"
 
-    # Format prompt for Qwen2.5-VL instruction-tuned model
+    # Format prompt for instruction-tuned models like Gemma or Qwen
     prompt = f"<|user|>\n{command}\n<|assistant|>\n"
 
     try:
-        output = pipe(prompt)[0]["generated_text"]
+        output = pipe(
+            prompt,
+            max_new_tokens=max_new_tokens,
+            do_sample=True,
+            temperature=0.7,
+            top_p=0.9
+        )[0]["generated_text"]
 
         # Cleanly extract assistant response
         if "<|assistant|>" in output: