Extract duplicated code (Whisper loading, audio recording, transcription, VAD processing) into reusable sttlib/ package. Rewrite all 3 scripts as thin wrappers. Add 24 unit tests with mocked hardware. Fix GPU fallback bug in assistant.py and args.system assignment bug.
139 lines
5.1 KiB
Python
139 lines
5.1 KiB
Python
import argparse
|
|
import subprocess
|
|
import json
|
|
import ollama
|
|
from sttlib import load_whisper_model, record_until_enter, transcribe
|
|
|
|
# --- Configuration ---
|
|
OLLAMA_MODEL = "qwen2.5-coder:7b"
|
|
CONFIRM_COMMANDS = True # Set to False to run commands instantly
|
|
|
|
|
|
# --- Terminal Tool ---
|
|
|
|
def run_terminal_command(command: str):
|
|
"""
|
|
Executes a bash command in the Linux terminal.
|
|
Used for file management, system info, and terminal tasks.
|
|
"""
|
|
if CONFIRM_COMMANDS:
|
|
print(f"\n{'='*40}")
|
|
print(f"\u26a0\ufe0f AI SUGGESTED: \033[1;32m{command}\033[0m")
|
|
choice = input(" Confirm? [Y/n] or provide feedback: ").strip()
|
|
|
|
if choice.lower() == 'n':
|
|
return "USER_REJECTION: The user did not approve this command. Please suggest an alternative."
|
|
elif choice and choice.lower() != 'y':
|
|
return f"USER_FEEDBACK: The user rejected the command with this reason: '{choice}'. Please adjust."
|
|
print(f"{'='*40}\n")
|
|
|
|
# Safety Guardrail
|
|
blacklist = ["rm -rf /", "mkfs", "dd if="]
|
|
if any(forbidden in command for forbidden in blacklist):
|
|
return "Error: Command blocked for security reasons."
|
|
|
|
try:
|
|
result = subprocess.run(command, shell=True,
|
|
capture_output=True, text=True, timeout=20)
|
|
output = f"STDOUT: {result.stdout}\nSTDERR: {result.stderr}"
|
|
return output if output.strip() else "Success (No output)."
|
|
except Exception as e:
|
|
return f"Execution Error: {str(e)}"
|
|
|
|
|
|
def main():
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument("--model", default=OLLAMA_MODEL)
|
|
parser.add_argument("--model-size", default="medium",
|
|
help="Whisper model size")
|
|
args, _ = parser.parse_known_args()
|
|
|
|
whisper_model = load_whisper_model(args.model_size)
|
|
|
|
# Initial System Prompt
|
|
messages = [{
|
|
'role': 'system',
|
|
'content': (
|
|
'You are a Linux expert assistant. When asked for a system task, '
|
|
'use the "run_terminal_command" tool. If the user rejects a command, '
|
|
'analyze their feedback and suggest a corrected alternative.'
|
|
)
|
|
}]
|
|
|
|
print(f"--- Assistant Active (Model: {args.model}) ---")
|
|
|
|
while True:
|
|
try:
|
|
# 1. Voice Capture
|
|
audio_data = record_until_enter()
|
|
user_text = transcribe(whisper_model, audio_data.flatten())
|
|
if not user_text:
|
|
continue
|
|
|
|
print(f"\nYOU: {user_text}")
|
|
messages.append({'role': 'user', 'content': user_text})
|
|
|
|
# 2. AI Interaction Loop (Supports up to 3 retries if rejected)
|
|
for attempt in range(3):
|
|
response = ollama.chat(
|
|
model=args.model,
|
|
messages=messages,
|
|
tools=[run_terminal_command],
|
|
options={'temperature': 0}
|
|
)
|
|
|
|
tool_calls = response.message.tool_calls
|
|
|
|
# Fallback Repair: Catch raw JSON output
|
|
if not tool_calls and '"run_terminal_command"' in response.message.content:
|
|
try:
|
|
c = response.message.content
|
|
start, end = c.find('{'), c.rfind('}') + 1
|
|
raw_json = json.loads(c[start:end])
|
|
tool_calls = [{'function': {
|
|
'name': 'run_terminal_command',
|
|
'arguments': raw_json.get('arguments', raw_json)
|
|
}}]
|
|
except:
|
|
pass
|
|
|
|
# 3. Execution Logic
|
|
if tool_calls:
|
|
call = tool_calls[0]
|
|
# Normalize arguments format
|
|
f_args = call.function.arguments if hasattr(
|
|
call, 'function') else call['function']['arguments']
|
|
|
|
result = run_terminal_command(f_args['command'])
|
|
|
|
# Update History
|
|
messages.append(response.message)
|
|
messages.append({'role': 'tool', 'content': result})
|
|
|
|
# If REJECTED, the loop continues and the AI sees the feedback
|
|
if "USER_REJECTION" in result or "USER_FEEDBACK" in result:
|
|
print("[RETHINKING] AI is adjusting the command...")
|
|
continue
|
|
else:
|
|
# Success: Let AI explain the result
|
|
final_res = ollama.chat(
|
|
model=args.model, messages=messages)
|
|
print(f"AI: {final_res.message.content}")
|
|
messages.append(final_res.message)
|
|
break
|
|
else:
|
|
# Normal Chat
|
|
print(f"AI: {response.message.content}")
|
|
messages.append(response.message)
|
|
break
|
|
|
|
except KeyboardInterrupt:
|
|
print("\nExiting...")
|
|
break
|
|
except Exception as e:
|
|
print(f"System Error: {e}")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|