Spaces:
Paused
Paused
File size: 6,541 Bytes
daea45b | 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 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 | """smolcode agent engine — backed by the Rust smolcode_core agent loop."""
from __future__ import annotations
import asyncio
import os
import tempfile
from collections.abc import Callable
from dataclasses import dataclass
from .config import Preset, load_preset
from .rust_session import RustRunResult, RustSession, rust_available
from .sandbox import Workspace
from .trace_collector import TraceCollector
# Legacy prompt kept for docs; Rust agent uses prompts.rs system prompts.
SYSTEM_PROMPT = """You are smolcode, a precise coding assistant running on a small local model."""
@dataclass
class Step:
number: int
kind: str
detail: str
total_tokens: int | None = None
class SmallCodeAgent:
"""Agent facade: uses the Rust engine when smolcode_core is installed."""
def __init__(
self,
preset: Preset | None = None,
model: str | None = None,
max_steps: int = 12,
*,
system_prompt: str | None = None,
registry_builder: Callable | None = None,
workspace: Workspace | None = None,
name: str = "smolcode",
agent: str = "build",
profile: str = "full",
yolo: bool = False,
workspace_dir: str | None = None,
approval_handler=None,
rust_session: RustSession | None = None,
) -> None:
self.preset = preset or load_preset()
self.model = model or self.preset.default_model
self.max_steps = max_steps
self._system_prompt = system_prompt # unused by Rust; kept for API compat
self._registry_builder = registry_builder
self.hit_max_steps = False
self.errored = False
ws_path = workspace_dir or os.environ.get("SMALLCODE_WORKSPACE")
if workspace is not None:
ws_path = str(workspace.root)
elif ws_path is None:
ws_path = tempfile.mkdtemp(prefix="smallcode-")
self._owns_workspace = True
else:
self._owns_workspace = False
self.workspace = workspace or Workspace(root=ws_path)
profile_name = profile
if registry_builder is not None:
profile_name = "web"
if not rust_available():
raise RuntimeError(
"smolcode_core required; install with maturin in smolcode-cli/crates/smolcode-py"
)
if rust_session is not None:
self._rust = rust_session
else:
self._rust = RustSession(
workspace=ws_path,
agent=agent,
yolo=yolo,
model=self.model,
base_url=self.preset.base_url,
api_key=self.preset.api_key,
profile=profile_name,
approval_handler=approval_handler,
)
self.trace_collector = self._rust.trace_collector
if registry_builder is not None:
self._register_web_tools()
def _register_web_tools(self) -> None:
from .tools import check_app_impl
ws = self.workspace
collector = self.trace_collector
def check_app(args: dict) -> dict:
return check_app_impl(ws, collector, args)
self._rust.register_tool("check_app", check_app)
async def run(self, task: str, *, think: str | None = None, yolo: bool | None = None) -> tuple[str, list[Step]]:
self.hit_max_steps = False
self.errored = False
result: RustRunResult = await self._rust.run(task, think=think, yolo=yolo)
self.hit_max_steps = result.hit_max_steps
self.errored = result.errored
steps = self._steps_from_trace()
return result.final, steps
async def run_live_turn(
self,
task: str,
*,
think: str | None = None,
yolo: bool | None = None,
poll_interval: float = 0.35,
):
"""Async generator yielding LiveFrame snapshots during a Rust agent turn."""
from .live_run import LiveFrame
self.hit_max_steps = False
self.errored = False
self.trace_collector.events.clear()
self._rust.clear_cancel()
self._rust._session.start_turn(task, think=think, yolo=yolo)
final_text = ""
done = False
interrupted = False
while not done:
if self._rust.cancelled:
interrupted = True
done = True
break
ev = await asyncio.to_thread(self._rust._session.poll_event)
if ev is None:
yield LiveFrame(
events=self.trace_collector.snapshot(),
files=self.files(),
)
await asyncio.sleep(poll_interval)
continue
kind = ev.get("kind")
if kind == "approval":
approved = True
if self._rust.approval_handler is not None:
approved = await self._rust.approval_handler(ev.get("desc", ""))
self._rust._session.approve(approved)
continue
self._rust._ingest_event(ev)
if kind == "final":
final_text = ev.get("text", "")
if kind == "done":
done = True
yield LiveFrame(
events=self.trace_collector.snapshot(),
files=self.files(),
raw_event=ev,
)
if interrupted:
final_text = final_text or "interrupted"
self.errored = True
if final_text and not interrupted:
self._rust._session.record_turn(task, final_text)
steps = self._steps_from_trace()
yield LiveFrame(
steps=steps,
events=self.trace_collector.snapshot(),
files=self.files(),
done=True,
result=(final_text, steps),
)
def _steps_from_trace(self) -> list[Step]:
out: list[Step] = []
for i, ev in enumerate(self.trace_collector.events):
out.append(Step(number=i, kind=ev.kind, detail=ev.detail))
return out
def current_steps(self) -> list[Step]:
return self._steps_from_trace()
def raw_history(self) -> list:
return self.current_steps()
def files(self) -> dict[str, str]:
return self._rust.files()
@property
def rust_session(self) -> RustSession:
return self._rust
def cleanup(self) -> None:
if getattr(self, "_owns_workspace", False):
self.workspace.cleanup()
|