Spaces:
Running on Zero
Running on Zero
File size: 4,444 Bytes
9936912 9e637cd 9936912 9e637cd 9936912 9e637cd 9936912 6be46a5 9e637cd 6be46a5 9e637cd 6be46a5 9e637cd 6be46a5 9e637cd 6be46a5 9936912 9e637cd 6be46a5 9e637cd 6be46a5 9e637cd 9936912 9e637cd 9936912 9e637cd 9936912 6be46a5 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 | #!/usr/bin/env python3
"""Interactive terminal client for ControlAI."""
from __future__ import annotations
import argparse
import sys
import time
from pathlib import Path
PROJECT_ROOT = Path(__file__).resolve().parent
if str(PROJECT_ROOT) not in sys.path:
sys.path.insert(0, str(PROJECT_ROOT))
from controlai_agent.agent import ControlAgent
from controlai_agent.engine import DEFAULT_ADAPTER, DEFAULT_MODEL, LocalEngine
DIM, RESET, BOLD = "\033[2m", "\033[0m", "\033[1m"
def converse(agent: ControlAgent, question: str, history: list[dict], verbose: bool) -> str:
started = time.time()
first_text: float | None = None
answer = ""
in_thought = False
for event in agent.stream(question, history):
kind = event["type"]
if kind == "thinking":
if not in_thought:
print(f"{DIM}thinking… ", end="", flush=True)
in_thought = True
if verbose:
print(f"{DIM}{event['text']}{RESET}", end="", flush=True)
elif kind == "text":
if in_thought:
print(RESET, flush=True)
in_thought = False
if first_text is None:
first_text = time.time() - started
print(event["text"], end="", flush=True)
elif kind == "tool_start":
args = str(event["arguments"])
print(f"\n{DIM}→ {event['tool']}({args[:100]}{'…' if len(args) > 100 else ''}){RESET}", flush=True)
elif kind == "tool_end":
print(f"{DIM} {event['status']}{RESET}", flush=True)
elif kind == "plot":
print(f"\n{DIM}[plot saved: outputs/plots/{Path(event['url']).name}]{RESET}", flush=True)
elif kind == "done":
answer = event["answer"]
if event["sources"]:
print(f"\n\n{DIM}Sources: {'; '.join(dict.fromkeys(event['sources']))}{RESET}")
stats = event["stats"]
print(
f"\n{DIM}{time.time() - started:.1f}s total · first token {first_text or 0:.2f}s · "
f"{stats['cached_tokens']}/{stats['prompt_tokens']} prompt tokens cached · "
f"{stats['decode_tps']} tok/s{RESET}"
)
return answer
def main() -> int:
parser = argparse.ArgumentParser(description="ControlAI -- offline control engineering agent")
parser.add_argument("prompt", nargs="*", help="ask one question and exit")
parser.add_argument("--model", default=DEFAULT_MODEL, help="MLX model id or local path")
parser.add_argument("--adapter", default=DEFAULT_ADAPTER, help="optional LoRA adapter path")
parser.add_argument(
"--thinking",
choices=("off", "auto", "on"),
default=None,
help="reasoning before answering (default: auto -- on for conceptual questions only)",
)
parser.add_argument("--think-budget", type=int, default=None, help="max tokens spent reasoning")
parser.add_argument("--verbose", action="store_true", help="print reasoning and full tool output")
args = parser.parse_args()
print(f"Loading {args.model}…")
engine = LocalEngine(model_id=args.model, adapter_path=args.adapter)
kwargs = {}
if args.thinking:
kwargs["thinking"] = args.thinking
if args.think_budget:
kwargs["think_budget"] = args.think_budget
agent = ControlAgent(engine=engine, **kwargs)
print(f"Ready in {engine.load_seconds:.1f}s.\n")
if args.prompt:
converse(agent, " ".join(args.prompt), [], args.verbose)
return 0
print(f"{BOLD}ControlAI{RESET} — control engineering assistant. Ctrl-D or 'exit' to quit.")
history: list[dict] = []
while True:
try:
question = input(f"\n{BOLD}you ›{RESET} ").strip()
except (KeyboardInterrupt, EOFError):
print("\nBye.")
return 0
if not question:
continue
if question.lower() in ("exit", "quit", "q"):
return 0
if question.lower() in ("clear", "reset"):
history.clear()
agent.engine.reset_cache()
print(f"{DIM}history cleared{RESET}")
continue
print()
answer = converse(agent, question, history, args.verbose)
history.append({"role": "user", "content": question})
history.append({"role": "assistant", "content": answer})
if __name__ == "__main__":
raise SystemExit(main())
|