57 lines
1.6 KiB
Python
57 lines
1.6 KiB
Python
import ollama
|
||
import subprocess
|
||
|
||
|
||
def shell_command(command: str) -> str:
|
||
"""
|
||
Run a shell command.
|
||
|
||
Args:
|
||
command: the command to run
|
||
|
||
Returns:
|
||
stdout: the stdout
|
||
stderr: the stderr
|
||
"""
|
||
exec = subprocess.run(command, encoding="utf-8", shell=True, capture_output=True)
|
||
return f"stdout:\n{exec.stdout}\n\nstderr:\n{exec.stderr}"
|
||
|
||
|
||
tools = [shell_command] # ollama-python turns this into JSON
|
||
|
||
messages = [{"role": "user", "content": input("prompt > ")}]
|
||
|
||
while True:
|
||
stream = ollama.chat(model="gpt-oss", messages=messages,
|
||
tools=tools, stream=True)
|
||
|
||
ran_tool = False
|
||
saw_content = False
|
||
for chunk in stream:
|
||
msg = chunk.message
|
||
|
||
# 1️⃣ the model is asking us to run something
|
||
if msg.tool_calls:
|
||
for call in msg.tool_calls:
|
||
if call["function"]["name"] == "shell_command":
|
||
cmd = call["function"]["arguments"]["command"]
|
||
output = shell_command(cmd)
|
||
messages.append({
|
||
"role": "tool",
|
||
"name": "shell_command",
|
||
"content": output
|
||
})
|
||
ran_tool = True
|
||
|
||
if msg.thinking:
|
||
print(f"\033[90m{msg.thinking}", end="", flush=True)
|
||
|
||
if msg.content:
|
||
if not saw_content:
|
||
print("")
|
||
saw_content = True
|
||
print(f"\033[97m{msg.content}", end="", flush=True)
|
||
|
||
# loop again if we just satisfied a tool call
|
||
if not ran_tool:
|
||
break
|