| """Local provider: YCode's own from-scratch model (``ycode-lm``), no API key. |
| |
| A small model trained on a laptop can answer coding questions and write |
| short functions, but it has not learned the tool-calling protocol, so it |
| runs in answer-only mode: it never reads, edits or runs anything. |
| """ |
|
|
| from __future__ import annotations |
|
|
| from pathlib import Path |
|
|
| from ycode.errors import ConfigError |
| from ycode.llm.base import ( |
| AssistantMessage, |
| LLMProvider, |
| LLMResponse, |
| Message, |
| StopReason, |
| StreamCallbacks, |
| ToolSpec, |
| Usage, |
| UserMessage, |
| ) |
|
|
| _NOTE_PREFIX = "[Note:" |
|
|
|
|
| class LocalProvider(LLMProvider): |
| name = "local" |
| supports_tools = False |
|
|
| def __init__(self, model_dir: Path, *, max_new_tokens: int = 400, temperature: float = 0.7, |
| device: str = "auto", lm=None) -> None: |
| super().__init__(str(model_dir)) |
| self.model_dir = Path(model_dir) |
| self.max_new_tokens = max_new_tokens |
| self.temperature = temperature |
| self.lm = lm if lm is not None else self._load(device) |
| from ycode.lm.assist import Assistant |
|
|
| |
| big = getattr(self.lm, "num_params", 0) > 500e6 |
| self.assistant = Assistant(self.lm, max_new_tokens=max_new_tokens, temperature=temperature, |
| candidates=2 if big else 4) |
|
|
| def _load(self, device: str): |
| try: |
| from ycode.lm.generate import LocalLM |
| except ImportError: |
| raise ConfigError( |
| "The local model needs PyTorch.", |
| hint='Install it with: pip install -e ".[local]"', |
| ) from None |
| try: |
| return LocalLM(self.model_dir, device=device) |
| except FileNotFoundError: |
| from ycode.config import bundled_models |
|
|
| available = ", ".join(m.removeprefix("ycode-lm-") for m in bundled_models()) |
| raise ConfigError( |
| f"No model found at {self.model_dir}.", |
| hint=(f"Bundled models available: {available} (e.g. `ycode --local v4`)." if available else |
| "YCode's trained models live in the repository's models/ folder: install from a clone " |
| "(git clone https://github.com/yazhiniyanMS/Code-agent && pip install -e .), " |
| "or set YCODE_MODELS_DIR to a folder containing them."), |
| ) from None |
| except (ValueError, RuntimeError) as exc: |
| raise ConfigError(f"Could not load the local model at {self.model_dir}: {exc}") from None |
|
|
| @property |
| def display_name(self) -> str: |
| params = getattr(self.lm, "num_params", 0) |
| version = getattr(self.lm, "version", None) |
| size = f"{params / 1e6:.1f}M params, " if params else "" |
| name = f"YCode-LM v{version}" if isinstance(version, int) else "YCode-LM" |
| return f"{name} ({size}local, {self.model_dir})" |
|
|
| @staticmethod |
| def _turns(messages: list[Message]) -> list[tuple[str, str]]: |
| turns: list[tuple[str, str]] = [] |
| for msg in messages: |
| if isinstance(msg, UserMessage): |
| text = msg.text |
| if text.startswith(_NOTE_PREFIX) and "\n\n" in text: |
| text = text.split("\n\n", 1)[1] |
| turns.append(("user", text)) |
| elif isinstance(msg, AssistantMessage) and msg.text: |
| turns.append(("assistant", msg.text)) |
| return turns |
|
|
| def complete(self, *, system: str, messages: list[Message], tools: list[ToolSpec], |
| callbacks: StreamCallbacks | None = None) -> LLMResponse: |
| callbacks = callbacks or StreamCallbacks() |
| turns = self._turns(messages)[-6:] |
| if not turns or turns[-1][0] != "user": |
| turns.append(("user", "Continue.")) |
| answer = self.assistant.answer(turns, on_text=callbacks.text) |
| if not answer.strip(): |
| answer = "(The local model did not produce an answer. Try rephrasing the question.)" |
| callbacks.text(answer) |
| return LLMResponse( |
| AssistantMessage(text=answer, provider=self.name), |
| StopReason.END_TURN, |
| Usage(output_tokens=len(answer) // 4), |
| ) |
|
|
