Merge branch 'main' into pip
This commit is contained in:
Daniel Han
commit
bf63f79414
63 changed files with 2574 additions and 552 deletions
2
COPYING
2
COPYING
|
|
@ -661,4 +661,4 @@ For more information on this, and how to apply and follow the GNU AGPL, see
|
|||
<https://www.gnu.org/licenses/>.
|
||||
|
||||
Files under unsloth/*, tests/*, scripts/* are Apache 2.0 licensed.
|
||||
Files under studio/*, cli/* which is optional to install are AGPLv3 licensed.
|
||||
Files under studio/*, unsloth_cli/* which is optional to install are AGPLv3 licensed.
|
||||
2
LICENSE
2
LICENSE
|
|
@ -188,7 +188,7 @@
|
|||
|
||||
Copyright [2024-] [Unsloth AI. Inc team, Daniel Han-Chen & Michael Han-Chen]
|
||||
Files under unsloth/*, tests/*, scripts/* are Apache 2.0 licensed.
|
||||
Files under studio/*, cli/* which is optional to install are AGPLv3 licensed.
|
||||
Files under studio/*, unsloth_cli/* which is optional to install are AGPLv3 licensed.
|
||||
|
||||
Licensed under the Apache License, Version 2.0 (the "License");
|
||||
you may not use this file except in compliance with the License.
|
||||
|
|
|
|||
98
README.md
98
README.md
|
|
@ -16,7 +16,8 @@ Run and train AI models with a unified local interface.
|
|||
<a href="https://unsloth.ai/docs">Documentation</a> •
|
||||
<a href="https://discord.com/invite/unsloth">Discord</a>
|
||||
</p>
|
||||
<img alt="unsloth studio ui homepage" src="https://raw.githubusercontent.com/unslothai/unsloth/main/studio/frontend/public/studio%20github%20landscape%20colab%20display.png" style="max-width: 100%; margin-bottom: 0;">
|
||||
<a href="https://unsloth.ai/docs/new/studio">
|
||||
<img alt="unsloth studio ui homepage" src="https://raw.githubusercontent.com/unslothai/unsloth/main/studio/frontend/public/studio%20github%20landscape%20colab%20display.png" style="max-width: 100%; margin-bottom: 0;"></a>
|
||||
|
||||
Unsloth Studio lets you run and train models for text, [audio](https://unsloth.ai/docs/basics/text-to-speech-tts-fine-tuning), [embedding](https://unsloth.ai/docs/new/embedding-finetuning), [vision](https://unsloth.ai/docs/basics/vision-fine-tuning) and more. Available on Windows, Linux and macOS.
|
||||
## ⭐ Features
|
||||
|
|
@ -43,38 +44,100 @@ Unsloth can be used in two ways: through **[Unsloth Studio](https://unsloth.ai/d
|
|||
Unsloth Studio works on **Windows, Linux, WSL** and **macOS**.
|
||||
|
||||
* **CPU:** Supported for **chat inference only**
|
||||
* **NVIDIA GPUs:** Training works on RTX 30/40/50, Blackwell, DGX Spark, DGX Station and more
|
||||
* **NVIDIA:** Training works on RTX 30/40/50, Blackwell, DGX Spark, Station and more
|
||||
* **macOS:** Currently supports chat only; **MLX training** is coming very soon
|
||||
* **AMD:** Chat works. Train with [Unsloth Core](#unsloth-core-code-based). Studio support is coming soon.
|
||||
* **Coming soon:** Training support for Apple MLX, AMD, and Intel.
|
||||
* **Multi-GPU:** Available now, with a major upgrade on the way
|
||||
|
||||
#### Windows, MacOS Linux or WSL:
|
||||
```
|
||||
git clone https://github.com/unslothai/unsloth.git
|
||||
cd unsloth
|
||||
pip install -e .
|
||||
#### MacOS, Linux or WSL Setup (One time):
|
||||
```bash
|
||||
pip install -U pip uv
|
||||
uv venv unsloth_studio
|
||||
source unsloth_studio/bin/activate
|
||||
uv pip install unsloth --torch-backend=auto
|
||||
unsloth studio setup
|
||||
unsloth studio -H 0.0.0.0 -p 8888
|
||||
```
|
||||
Use our [Docker image](https://hub.docker.com/r/unsloth/unsloth) ```unsloth/unsloth``` container. Read our [Docker Guide](https://unsloth.ai/docs/get-started/install/docker).
|
||||
#### Git from source
|
||||
Then to launch every time:
|
||||
```bash
|
||||
source unsloth_studio/bin/activate
|
||||
unsloth studio -H 0.0.0.0 -p 8888
|
||||
```
|
||||
git clone https://github.com/unslothai/unsloth.git
|
||||
cd unsloth
|
||||
pip install -e .
|
||||
|
||||
#### Windows PowerShell (One time):
|
||||
```bash
|
||||
pip install -U pip uv
|
||||
uv venv unsloth_studio
|
||||
.\unsloth_studio\Scripts\activate
|
||||
uv pip install unsloth --torch-backend=auto
|
||||
unsloth studio setup
|
||||
unsloth studio -H 0.0.0.0 -p 8888
|
||||
```
|
||||
Then to launch every time:
|
||||
```bash
|
||||
.\unsloth_studio\Scripts\activate
|
||||
unsloth studio -H 0.0.0.0 -p 8888
|
||||
```
|
||||
|
||||
Use our [Docker image](https://hub.docker.com/r/unsloth/unsloth) ```unsloth/unsloth``` container. Read our [Docker Guide](https://unsloth.ai/docs/get-started/install/docker).
|
||||
|
||||
#### Nightly Installation - MacOS, Linux or WSL Setup (One time):
|
||||
```bash
|
||||
pip install -U pip uv
|
||||
git clone --filter=blob:none https://github.com/unslothai/unsloth.git unsloth_studio
|
||||
cd unsloth_studio
|
||||
uv venv
|
||||
source .venv/bin/activate
|
||||
uv pip install -e . --torch-backend=auto
|
||||
unsloth studio setup
|
||||
unsloth studio -H 0.0.0.0 -p 8888
|
||||
```
|
||||
Then to launch every time:
|
||||
```bash
|
||||
cd unsloth_studio
|
||||
source .venv/bin/activate
|
||||
unsloth studio -H 0.0.0.0 -p 8888
|
||||
```
|
||||
|
||||
#### Nightly Installation - Windows Powershell (One time):
|
||||
```bash
|
||||
pip install -U pip uv
|
||||
git clone --filter=blob:none https://github.com/unslothai/unsloth.git unsloth_studio
|
||||
cd unsloth_studio
|
||||
uv venv
|
||||
.\.venv\Scripts\activate
|
||||
uv pip install -e . --torch-backend=auto
|
||||
unsloth studio setup
|
||||
unsloth studio -H 0.0.0.0 -p 8888
|
||||
```
|
||||
Then to launch every time:
|
||||
```bash
|
||||
cd unsloth_studio
|
||||
.\.venv\Scripts\activate
|
||||
unsloth studio -H 0.0.0.0 -p 8888
|
||||
```
|
||||
|
||||
### Unsloth Core (code-based)
|
||||
#### Windows, Linux, WSL
|
||||
#### Linux, WSL
|
||||
```bash
|
||||
pip install unsloth
|
||||
pip install -U pip uv
|
||||
uv venv unsloth_env
|
||||
source unsloth_env/bin/activate
|
||||
uv pip install unsloth --torch-backend=auto
|
||||
```
|
||||
#### Windows Powershell
|
||||
```bash
|
||||
pip install -U pip uv
|
||||
uv venv unsloth_env
|
||||
.\unsloth_env\Scripts\activate
|
||||
uv pip install unsloth --torch-backend=auto
|
||||
```
|
||||
For Windows, `pip install unsloth` works only if you have Pytorch installed. Read our [Windows Guide](https://unsloth.ai/docs/get-started/install/windows-installation).
|
||||
You can use the same Docker image as Unsloth Studio.
|
||||
|
||||
#### AMD, Intel
|
||||
For RTX 50x, B200, 6000 GPUs: `pip install unsloth`. Read our guides for: [Blackwell](https://unsloth.ai/docs/blog/fine-tuning-llms-with-blackwell-rtx-50-series-and-unsloth) and [DGX Spark](https://unsloth.ai/docs/blog/fine-tuning-llms-with-nvidia-dgx-spark-and-unsloth). <br>
|
||||
For RTX 50x, B200, 6000 GPUs: `uv pip install unsloth --torch-backend=auto`. Read our guides for: [Blackwell](https://unsloth.ai/docs/blog/fine-tuning-llms-with-blackwell-rtx-50-series-and-unsloth) and [DGX Spark](https://unsloth.ai/docs/blog/fine-tuning-llms-with-nvidia-dgx-spark-and-unsloth). <br>
|
||||
To install Unsloth on **AMD** and **Intel** GPUs, follow our [AMD Guide](https://unsloth.ai/docs/get-started/install/amd) and [Intel Guide](https://unsloth.ai/docs/get-started/install/intel).
|
||||
|
||||
## ✨ Free Notebooks
|
||||
|
|
@ -132,6 +195,11 @@ You can cite the Unsloth repo as follows:
|
|||
```
|
||||
If you trained a model with 🦥Unsloth, you can use this cool sticker! <img src="https://raw.githubusercontent.com/unslothai/unsloth/main/images/made with unsloth.png" width="200" align="center" />
|
||||
|
||||
### License
|
||||
Unsloth uses a dual-licensing model of Apache 2.0 and AGPL-3.0. The core Unsloth package remains licensed under **[Apache 2.0](https://github.com/unslothai/unsloth?tab=Apache-2.0-1-ov-file)**, while certain optional components, such as the Unsloth Studio UI are licensed under the open-source license **[AGPL-3.0](https://github.com/unslothai/unsloth?tab=AGPL-3.0-2-ov-file)**.
|
||||
|
||||
This structure helps support ongoing Unsloth development while keeping the project open source and enabling the broader ecosystem to continue growing.
|
||||
|
||||
### Thank You to
|
||||
- The [llama.cpp library](https://github.com/ggml-org/llama.cpp) that lets users run and save models with Unsloth
|
||||
- The Hugging Face team and their libraries: [transformers](https://github.com/huggingface/transformers) and [TRL](https://github.com/huggingface/trl)
|
||||
|
|
|
|||
2
cli.py
2
cli.py
|
|
@ -1,7 +1,7 @@
|
|||
# SPDX-License-Identifier: AGPL-3.0-only
|
||||
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
|
||||
|
||||
from cli import app
|
||||
from unsloth_cli import app
|
||||
|
||||
if __name__ == "__main__":
|
||||
app()
|
||||
|
|
|
|||
|
|
@ -55,7 +55,7 @@ dependencies = [
|
|||
]
|
||||
|
||||
[project.scripts]
|
||||
unsloth = "cli:app"
|
||||
unsloth = "unsloth_cli:app"
|
||||
|
||||
[tool.setuptools.dynamic]
|
||||
version = {attr = "unsloth.models._utils.__version__"}
|
||||
|
|
@ -737,6 +737,6 @@ cu130-ampere-torch2100 = [
|
|||
]
|
||||
|
||||
[project.urls]
|
||||
homepage = "http://www.unsloth.ai"
|
||||
documentation = "https://github.com/unslothai/unsloth"
|
||||
homepage = "https://unsloth.ai"
|
||||
documentation = "https://unsloth.ai/docs"
|
||||
repository = "https://github.com/unslothai/unsloth"
|
||||
|
|
|
|||
|
|
@ -15,11 +15,14 @@
|
|||
"To install Unsloth Studio on your local device, follow [our guide](https://unsloth.ai/docs/new/unsloth-studio/install). Unsloth Studio is licensed [AGPL-3.0](https://github.com/unslothai/unsloth/blob/main/studio/LICENSE.AGPL-3.0).\n",
|
||||
"\n",
|
||||
"### Unsloth Studio\n",
|
||||
"\n",
|
||||
"Train and run open models with [**Unsloth Studio**](https://unsloth.ai/docs/new/unsloth-studio/start). Installation may take 10 mins.\n",
|
||||
"\n",
|
||||
"[Features](https://unsloth.ai/docs/new/unsloth-studio#features) • [Quickstart](https://unsloth.ai/docs/new/unsloth-studio/start) • [Data Recipes](https://unsloth.ai/docs/new/unsloth-studio/data-recipe) • [Studio Chat](https://unsloth.ai/docs/new/unsloth-studio/chat) • [Export](https://unsloth.ai/docs/new/unsloth-studio/export)"
|
||||
]
|
||||
"\n",
|
||||
"Train and run open models with [**Unsloth Studio**](https://unsloth.ai/docs/new/unsloth-studio/start). Currently, installation may take 30+ mins so use a newer GPU.\n",
|
||||
"\n",
|
||||
"\n",
|
||||
"We are actively working on making Unsloth Studio install on Colab T4 GPUs faster.\n",
|
||||
"\n",
|
||||
"[Features](https://unsloth.ai/docs/new/unsloth-studio#features) • [Quickstart](https://unsloth.ai/docs/new/unsloth-studio/start) • [Data Recipes](https://unsloth.ai/docs/new/unsloth-studio/data-recipe) • [Studio Chat](https://unsloth.ai/docs/new/unsloth-studio/chat) • [Export](https://unsloth.ai/docs/new/unsloth-studio/export)"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
|
|
|
|||
|
|
@ -9,6 +9,32 @@ Uses Colab's built-in proxy - no external tunneling needed!
|
|||
from pathlib import Path
|
||||
import sys
|
||||
|
||||
|
||||
def _bootstrap_studio_venv() -> None:
|
||||
"""Expose the Studio venv's site-packages to the current interpreter.
|
||||
|
||||
On Colab, notebook cells run outside the venv subshell. Instead of
|
||||
installing the full stack into system Python, we prepend the venv's
|
||||
site-packages so that packages like structlog, fastapi, etc. are
|
||||
importable from notebook cells and take priority over system copies.
|
||||
"""
|
||||
venv_lib = Path.home() / ".unsloth" / "studio" / ".venv" / "lib"
|
||||
if not venv_lib.exists():
|
||||
import warnings
|
||||
|
||||
warnings.warn(
|
||||
f"Studio venv not found at {venv_lib.parent} -- run 'unsloth studio setup' first",
|
||||
stacklevel = 2,
|
||||
)
|
||||
return
|
||||
for sp in venv_lib.glob("python*/site-packages"):
|
||||
sp_str = str(sp)
|
||||
if sp_str not in sys.path:
|
||||
sys.path.insert(0, sp_str)
|
||||
|
||||
|
||||
_bootstrap_studio_venv()
|
||||
|
||||
# Add backend to path early so local modules like loggers can be imported
|
||||
backend_path = str(Path(__file__).parent)
|
||||
if backend_path not in sys.path:
|
||||
|
|
|
|||
|
|
@ -10,6 +10,16 @@ like unsloth, transformers, or torch before the version activation
|
|||
code has a chance to run.
|
||||
"""
|
||||
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
# Ensure the backend directory is on sys.path so that bare "from utils.*"
|
||||
# imports used throughout the backend work when core is imported as a package
|
||||
# (e.g. from the CLI: "from studio.backend.core import ModelConfig").
|
||||
_backend_dir = str(Path(__file__).resolve().parent.parent)
|
||||
if _backend_dir not in sys.path:
|
||||
sys.path.insert(0, _backend_dir)
|
||||
|
||||
__all__ = [
|
||||
# Inference
|
||||
"InferenceBackend",
|
||||
|
|
|
|||
|
|
@ -40,59 +40,25 @@ def _activate_transformers_version(model_name: str) -> None:
|
|||
if backend_path not in sys.path:
|
||||
sys.path.insert(0, backend_path)
|
||||
|
||||
from utils.transformers_version import needs_transformers_5, _resolve_base_model
|
||||
from utils.transformers_version import (
|
||||
needs_transformers_5,
|
||||
_resolve_base_model,
|
||||
_ensure_venv_t5_exists,
|
||||
_VENV_T5_DIR,
|
||||
)
|
||||
|
||||
resolved = _resolve_base_model(model_name)
|
||||
if needs_transformers_5(resolved):
|
||||
venv_t5 = os.path.join(
|
||||
os.path.expanduser("~"), ".unsloth", "studio", ".venv_t5"
|
||||
)
|
||||
if os.path.isdir(venv_t5):
|
||||
sys.path.insert(0, venv_t5)
|
||||
logger.info("Activated transformers 5.x from %s", venv_t5)
|
||||
else:
|
||||
# Fallback: pip install at runtime (slower, ~10-15s)
|
||||
logger.warning(".venv_t5 not found at %s — installing at runtime", venv_t5)
|
||||
import subprocess as sp
|
||||
|
||||
os.makedirs(venv_t5, exist_ok = True)
|
||||
r1 = sp.run(
|
||||
[
|
||||
sys.executable,
|
||||
"-m",
|
||||
"pip",
|
||||
"install",
|
||||
"--target",
|
||||
venv_t5,
|
||||
"--no-deps",
|
||||
"transformers==5.3.0",
|
||||
],
|
||||
stdout = sp.PIPE,
|
||||
stderr = sp.STDOUT,
|
||||
if not _ensure_venv_t5_exists():
|
||||
raise RuntimeError(
|
||||
f"Cannot activate transformers 5.x: .venv_t5 missing at {_VENV_T5_DIR}"
|
||||
)
|
||||
r2 = sp.run(
|
||||
[
|
||||
sys.executable,
|
||||
"-m",
|
||||
"pip",
|
||||
"install",
|
||||
"--target",
|
||||
venv_t5,
|
||||
"--no-deps",
|
||||
"huggingface_hub==1.3.0",
|
||||
],
|
||||
stdout = sp.PIPE,
|
||||
stderr = sp.STDOUT,
|
||||
)
|
||||
if r1.returncode != 0 or r2.returncode != 0:
|
||||
raise RuntimeError(
|
||||
f"Failed to install transformers 5.x into {venv_t5}. "
|
||||
f"pip returncode: transformers={r1.returncode}, huggingface_hub={r2.returncode}"
|
||||
)
|
||||
sys.path.insert(0, venv_t5)
|
||||
if _VENV_T5_DIR not in sys.path:
|
||||
sys.path.insert(0, _VENV_T5_DIR)
|
||||
logger.info("Activated transformers 5.x from %s", _VENV_T5_DIR)
|
||||
# Propagate to child subprocesses (e.g. GGUF converter)
|
||||
_pp = os.environ.get("PYTHONPATH", "")
|
||||
os.environ["PYTHONPATH"] = venv_t5 + (os.pathsep + _pp if _pp else "")
|
||||
os.environ["PYTHONPATH"] = _VENV_T5_DIR + (os.pathsep + _pp if _pp else "")
|
||||
else:
|
||||
logger.info("Using default transformers (4.57.x) for %s", model_name)
|
||||
|
||||
|
|
|
|||
|
|
@ -1173,50 +1173,113 @@ class LlamaCppBackend:
|
|||
Handles formats like:
|
||||
<tool_call>{"name":"web_search","arguments":{"query":"..."}}</tool_call>
|
||||
<tool_call><function=web_search><parameter=query>...</parameter></function></tool_call>
|
||||
Closing </tool_call> tag is optional (models sometimes omit it).
|
||||
Closing tags (</tool_call>, </function>, </parameter>) are all optional
|
||||
since models frequently omit them.
|
||||
"""
|
||||
import re
|
||||
|
||||
tool_calls = []
|
||||
# Pattern 1: JSON inside <tool_call> tags (closing tag optional)
|
||||
for match in re.finditer(
|
||||
r"<tool_call>\s*(\{.*?\})\s*(?:</tool_call>)?", content, re.DOTALL
|
||||
):
|
||||
try:
|
||||
obj = json.loads(match.group(1))
|
||||
tc = {
|
||||
"id": f"call_{len(tool_calls)}",
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": obj.get("name", ""),
|
||||
"arguments": obj.get("arguments", {}),
|
||||
},
|
||||
}
|
||||
if isinstance(tc["function"]["arguments"], dict):
|
||||
tc["function"]["arguments"] = json.dumps(
|
||||
tc["function"]["arguments"]
|
||||
)
|
||||
tool_calls.append(tc)
|
||||
except (json.JSONDecodeError, ValueError):
|
||||
pass
|
||||
|
||||
# Pattern 1: JSON inside <tool_call> tags.
|
||||
# Use balanced-brace extraction that skips braces inside JSON strings.
|
||||
for m in re.finditer(r"<tool_call>\s*\{", content):
|
||||
brace_start = m.end() - 1 # position of the opening {
|
||||
depth, i = 0, brace_start
|
||||
in_string = False
|
||||
while i < len(content):
|
||||
ch = content[i]
|
||||
if in_string:
|
||||
if ch == "\\" and i + 1 < len(content):
|
||||
i += 2 # skip escaped character
|
||||
continue
|
||||
if ch == '"':
|
||||
in_string = False
|
||||
elif ch == '"':
|
||||
in_string = True
|
||||
elif ch == "{":
|
||||
depth += 1
|
||||
elif ch == "}":
|
||||
depth -= 1
|
||||
if depth == 0:
|
||||
break
|
||||
i += 1
|
||||
if depth == 0:
|
||||
json_str = content[brace_start : i + 1]
|
||||
try:
|
||||
obj = json.loads(json_str)
|
||||
tc = {
|
||||
"id": f"call_{len(tool_calls)}",
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": obj.get("name", ""),
|
||||
"arguments": obj.get("arguments", {}),
|
||||
},
|
||||
}
|
||||
if isinstance(tc["function"]["arguments"], dict):
|
||||
tc["function"]["arguments"] = json.dumps(
|
||||
tc["function"]["arguments"]
|
||||
)
|
||||
tool_calls.append(tc)
|
||||
except (json.JSONDecodeError, ValueError):
|
||||
pass
|
||||
|
||||
# Pattern 2: XML-style <function=name><parameter=key>value</parameter></function>
|
||||
# Closing </tool_call> optional
|
||||
# All closing tags optional -- models frequently omit </parameter>,
|
||||
# </function>, and/or </tool_call>.
|
||||
if not tool_calls:
|
||||
for match in re.finditer(
|
||||
r"<tool_call>\s*<function=(\w+)>(.*?)</function>\s*(?:</tool_call>)?",
|
||||
content,
|
||||
re.DOTALL,
|
||||
):
|
||||
func_name = match.group(1)
|
||||
params_text = match.group(2)
|
||||
# Step 1: Find all <function=name> positions and extract their bodies.
|
||||
# Body boundary: use only </tool_call> or next <function= as hard
|
||||
# boundaries. We avoid using </function> as a boundary because
|
||||
# code parameter values can contain that literal string.
|
||||
# After extracting, we trim a trailing </function> if present.
|
||||
func_starts = list(re.finditer(r"<function=(\w+)>\s*", content))
|
||||
for idx, fm in enumerate(func_starts):
|
||||
func_name = fm.group(1)
|
||||
body_start = fm.end()
|
||||
# Hard boundaries: next <function= tag or </tool_call>
|
||||
next_func = (
|
||||
func_starts[idx + 1].start()
|
||||
if idx + 1 < len(func_starts)
|
||||
else len(content)
|
||||
)
|
||||
end_tag = re.search(r"</tool_call>", content[body_start:])
|
||||
if end_tag:
|
||||
body_end = body_start + end_tag.start()
|
||||
else:
|
||||
body_end = len(content)
|
||||
body_end = min(body_end, next_func)
|
||||
body = content[body_start:body_end]
|
||||
# Trim trailing </function> if present (it's the real closing tag)
|
||||
body = re.sub(r"\s*</function>\s*$", "", body)
|
||||
|
||||
# Step 2: Extract parameters from body.
|
||||
# For single-parameter functions (the common case: code, command,
|
||||
# query), use body end as the only boundary to avoid false matches
|
||||
# on </parameter> inside code strings.
|
||||
arguments = {}
|
||||
for param_match in re.finditer(
|
||||
r"<parameter=(\w+)>\s*(.*?)\s*</parameter>",
|
||||
params_text,
|
||||
re.DOTALL,
|
||||
):
|
||||
arguments[param_match.group(1)] = param_match.group(2)
|
||||
param_starts = list(re.finditer(r"<parameter=(\w+)>\s*", body))
|
||||
if len(param_starts) == 1:
|
||||
# Single parameter: value is everything from after the tag
|
||||
# to end of body, trimming any trailing </parameter>.
|
||||
pm = param_starts[0]
|
||||
val = body[pm.end() :]
|
||||
val = re.sub(r"\s*</parameter>\s*$", "", val)
|
||||
arguments[pm.group(1)] = val.strip()
|
||||
else:
|
||||
for pidx, pm in enumerate(param_starts):
|
||||
param_name = pm.group(1)
|
||||
val_start = pm.end()
|
||||
# Value ends at next <parameter= or end of body
|
||||
next_param = (
|
||||
param_starts[pidx + 1].start()
|
||||
if pidx + 1 < len(param_starts)
|
||||
else len(body)
|
||||
)
|
||||
val = body[val_start:next_param]
|
||||
# Trim trailing </parameter> if present
|
||||
val = re.sub(r"\s*</parameter>\s*$", "", val)
|
||||
arguments[param_name] = val.strip()
|
||||
|
||||
tc = {
|
||||
"id": f"call_{len(tool_calls)}",
|
||||
"type": "function",
|
||||
|
|
@ -1273,10 +1336,11 @@ class LlamaCppBackend:
|
|||
) -> Generator[str, None, None]:
|
||||
"""Iterate over an httpx streaming response with cancel support.
|
||||
|
||||
Uses a short read timeout on the stream so that cancel_event is
|
||||
checked at least every 0.5s, even if the model is slow to produce
|
||||
the next token. Without this, iter_text() blocks until the next
|
||||
chunk arrives and cancellation can take many seconds on large models.
|
||||
Checks cancel_event between chunks and on ReadTimeout. The
|
||||
cancel watcher in _stream_with_retry also calls response.close()
|
||||
on cancel, which unblocks iter_text() once the response exists.
|
||||
During normal streaming llama-server sends tokens frequently,
|
||||
so the cancel check between chunks is the primary mechanism.
|
||||
"""
|
||||
text_iter = response.iter_text()
|
||||
while True:
|
||||
|
|
@ -1301,24 +1365,85 @@ class LlamaCppBackend:
|
|||
payload: dict,
|
||||
cancel_event: Optional[threading.Event] = None,
|
||||
):
|
||||
"""Open an httpx streaming POST, retrying on ReadTimeout.
|
||||
"""Open an httpx streaming POST with cancel support.
|
||||
|
||||
The short read timeout (0.5 s) that enables cancel-checking during
|
||||
streaming can also fire while waiting for the server to produce
|
||||
its first response bytes (e.g. a reasoning model thinking).
|
||||
This wrapper retries the connection until headers arrive or
|
||||
cancel_event is set.
|
||||
Sends the request once with a long read timeout (120 s) so
|
||||
prompt processing (prefill) can finish without triggering a
|
||||
retry storm. The previous 0.5 s timeout caused duplicate POST
|
||||
requests every half second, forcing llama-server to restart
|
||||
processing each time.
|
||||
|
||||
A background watcher thread provides cancel by closing the
|
||||
response when cancel_event is set. Limitation: httpx does not
|
||||
allow interrupting a blocked read from another thread before
|
||||
the response object exists, so cancel during the initial
|
||||
header wait (prefill phase) only takes effect once headers
|
||||
arrive. After that, response.close() unblocks reads promptly.
|
||||
In practice llama-server prefill is 1-5 s for typical prompts,
|
||||
during which cancel is deferred -- still much better than the
|
||||
old retry storm which made prefill slower.
|
||||
"""
|
||||
while True:
|
||||
if cancel_event is not None and cancel_event.is_set():
|
||||
raise GeneratorExit
|
||||
|
||||
# Background watcher: close the response if cancel is requested.
|
||||
# Only effective after response headers arrive (httpx limitation).
|
||||
_cancel_closed = threading.Event()
|
||||
_response_ref: list = [None]
|
||||
|
||||
def _cancel_watcher():
|
||||
while not _cancel_closed.is_set():
|
||||
if cancel_event.wait(timeout = 0.3):
|
||||
# Cancel requested. Keep polling until the response object
|
||||
# exists so we can close it, or until the main thread
|
||||
# finishes on its own (_cancel_closed is set in finally).
|
||||
while not _cancel_closed.is_set():
|
||||
r = _response_ref[0]
|
||||
if r is not None:
|
||||
try:
|
||||
r.close()
|
||||
return
|
||||
except Exception as e:
|
||||
logger.debug(
|
||||
f"Error closing response in cancel watcher: {e}"
|
||||
)
|
||||
# Response not created yet -- wait briefly and retry
|
||||
_cancel_closed.wait(timeout = 0.1)
|
||||
return
|
||||
|
||||
watcher = None
|
||||
if cancel_event is not None:
|
||||
watcher = threading.Thread(
|
||||
target = _cancel_watcher, daemon = True, name = "prefill-cancel"
|
||||
)
|
||||
watcher.start()
|
||||
|
||||
try:
|
||||
# Long read timeout so prefill (prompt processing) can finish
|
||||
# without triggering a retry storm. Cancel during both
|
||||
# prefill and streaming is handled by the watcher thread
|
||||
# which closes the response, unblocking any httpx read.
|
||||
prefill_timeout = httpx.Timeout(
|
||||
connect = 30,
|
||||
read = 120.0,
|
||||
write = 10,
|
||||
pool = 10,
|
||||
)
|
||||
with client.stream(
|
||||
"POST", url, json = payload, timeout = prefill_timeout
|
||||
) as response:
|
||||
_response_ref[0] = response
|
||||
if cancel_event is not None and cancel_event.is_set():
|
||||
raise GeneratorExit
|
||||
yield response
|
||||
return
|
||||
except (httpx.ReadError, httpx.RemoteProtocolError, httpx.CloseError):
|
||||
# Response was closed by the cancel watcher
|
||||
if cancel_event is not None and cancel_event.is_set():
|
||||
raise GeneratorExit
|
||||
try:
|
||||
with client.stream("POST", url, json = payload) as response:
|
||||
yield response
|
||||
return
|
||||
except httpx.ReadTimeout:
|
||||
# Server still thinking -- retry
|
||||
continue
|
||||
raise
|
||||
finally:
|
||||
_cancel_closed.set()
|
||||
|
||||
def generate_chat_completion(
|
||||
self,
|
||||
|
|
@ -1371,8 +1496,9 @@ class LlamaCppBackend:
|
|||
in_thinking = False
|
||||
|
||||
try:
|
||||
# Use a short read timeout so we can check cancel_event
|
||||
# frequently instead of blocking indefinitely on slow models.
|
||||
# _stream_with_retry uses a 120 s read timeout so prefill
|
||||
# can finish. Cancel during streaming is handled by the
|
||||
# watcher thread (closes the response on cancel_event).
|
||||
stream_timeout = httpx.Timeout(connect = 10, read = 0.5, write = 10, pool = 10)
|
||||
with httpx.Client(timeout = stream_timeout) as client:
|
||||
with self._stream_with_retry(
|
||||
|
|
@ -1468,7 +1594,10 @@ class LlamaCppBackend:
|
|||
stop: Optional[list[str]] = None,
|
||||
cancel_event: Optional[threading.Event] = None,
|
||||
enable_thinking: Optional[bool] = None,
|
||||
max_tool_iterations: int = 5,
|
||||
max_tool_iterations: int = 10,
|
||||
auto_heal_tool_calls: bool = True,
|
||||
tool_call_timeout: int = 300,
|
||||
session_id: Optional[str] = None,
|
||||
) -> Generator[dict, None, None]:
|
||||
"""
|
||||
Agentic loop: let the model call tools, execute them, and continue.
|
||||
|
|
@ -1533,16 +1662,45 @@ class LlamaCppBackend:
|
|||
tool_calls = message.get("tool_calls")
|
||||
|
||||
# Fallback: detect tool calls embedded as XML/text in content
|
||||
# Some models output <tool_call> XML instead of structured tool_calls
|
||||
# Some models output <tool_call> XML instead of structured tool_calls,
|
||||
# or bare <function=...> tags without <tool_call> wrapper.
|
||||
content_text = message.get("content", "") or ""
|
||||
if not tool_calls and "<tool_call>" in content_text:
|
||||
if (
|
||||
auto_heal_tool_calls
|
||||
and not tool_calls
|
||||
and ("<tool_call>" in content_text or "<function=" in content_text)
|
||||
):
|
||||
tool_calls = self._parse_tool_calls_from_text(content_text)
|
||||
if tool_calls:
|
||||
# Strip the tool call markup from content
|
||||
# Strip the tool call markup from content.
|
||||
# Use greedy match within <tool_call> blocks since they
|
||||
# can contain arbitrary content including code.
|
||||
import re
|
||||
|
||||
# Strip <tool_call>...</tool_call> blocks (greedy inside)
|
||||
content_text = re.sub(
|
||||
r"<tool_call>.*?(?:</tool_call>|$)",
|
||||
r"<tool_call>.*?</tool_call>",
|
||||
"",
|
||||
content_text,
|
||||
flags = re.DOTALL,
|
||||
)
|
||||
# Strip unterminated <tool_call>... to end
|
||||
content_text = re.sub(
|
||||
r"<tool_call>.*$",
|
||||
"",
|
||||
content_text,
|
||||
flags = re.DOTALL,
|
||||
)
|
||||
# Strip bare <function=...>...</function> blocks
|
||||
content_text = re.sub(
|
||||
r"<function=\w+>.*?</function>",
|
||||
"",
|
||||
content_text,
|
||||
flags = re.DOTALL,
|
||||
)
|
||||
# Strip unterminated bare <function=...> to end
|
||||
content_text = re.sub(
|
||||
r"<function=\w+>.*$",
|
||||
"",
|
||||
content_text,
|
||||
flags = re.DOTALL,
|
||||
|
|
@ -1569,7 +1727,10 @@ class LlamaCppBackend:
|
|||
try:
|
||||
arguments = json.loads(raw_args)
|
||||
except (json.JSONDecodeError, ValueError):
|
||||
arguments = {"query": raw_args}
|
||||
if auto_heal_tool_calls:
|
||||
arguments = {"query": raw_args}
|
||||
else:
|
||||
arguments = {"raw": raw_args}
|
||||
else:
|
||||
arguments = raw_args
|
||||
|
||||
|
|
@ -1596,10 +1757,33 @@ class LlamaCppBackend:
|
|||
status_text = f"Calling: {tool_name}"
|
||||
yield {"type": "status", "text": status_text}
|
||||
|
||||
# Emit tool_start so the frontend can record inputs
|
||||
yield {
|
||||
"type": "tool_start",
|
||||
"tool_name": tool_name,
|
||||
"tool_call_id": tc.get("id", ""),
|
||||
"arguments": arguments,
|
||||
}
|
||||
|
||||
# Execute the tool
|
||||
result = execute_tool(
|
||||
tool_name, arguments, cancel_event = cancel_event
|
||||
_effective_timeout = (
|
||||
None if tool_call_timeout >= 9999 else tool_call_timeout
|
||||
)
|
||||
result = execute_tool(
|
||||
tool_name,
|
||||
arguments,
|
||||
cancel_event = cancel_event,
|
||||
timeout = _effective_timeout,
|
||||
session_id = session_id,
|
||||
)
|
||||
|
||||
# Emit tool_end so the frontend can record outputs
|
||||
yield {
|
||||
"type": "tool_end",
|
||||
"tool_name": tool_name,
|
||||
"tool_call_id": tc.get("id", ""),
|
||||
"result": result,
|
||||
}
|
||||
|
||||
# Append tool result to conversation
|
||||
tool_msg = {
|
||||
|
|
@ -1651,7 +1835,30 @@ class LlamaCppBackend:
|
|||
if stop:
|
||||
stream_payload["stop"] = stop
|
||||
|
||||
import re as _re_final
|
||||
|
||||
# Closed blocks only -- safe to strip mid-stream without shrinking later.
|
||||
_TOOL_CLOSED_PATTERNS = [
|
||||
_re_final.compile(r"<tool_call>.*?</tool_call>", _re_final.DOTALL),
|
||||
_re_final.compile(r"<function=\w+>.*?</function>", _re_final.DOTALL),
|
||||
]
|
||||
# Open-ended patterns strip from an opening tag to end-of-string.
|
||||
# Only applied on the final flush to avoid non-monotonic shrinking.
|
||||
_TOOL_ALL_PATTERNS = _TOOL_CLOSED_PATTERNS + [
|
||||
_re_final.compile(r"<tool_call>.*$", _re_final.DOTALL),
|
||||
_re_final.compile(r"<function=\w+>.*$", _re_final.DOTALL),
|
||||
]
|
||||
|
||||
def _strip_tool_markup(text: str, *, final: bool = False) -> str:
|
||||
if not auto_heal_tool_calls:
|
||||
return text
|
||||
patterns = _TOOL_ALL_PATTERNS if final else _TOOL_CLOSED_PATTERNS
|
||||
for pat in patterns:
|
||||
text = pat.sub("", text)
|
||||
return text.strip() if final else text
|
||||
|
||||
cumulative = ""
|
||||
_last_emitted = ""
|
||||
in_thinking = False
|
||||
has_content_tokens = False
|
||||
reasoning_text = ""
|
||||
|
|
@ -1683,7 +1890,12 @@ class LlamaCppBackend:
|
|||
if in_thinking:
|
||||
if has_content_tokens:
|
||||
cumulative += "</think>"
|
||||
yield {"type": "content", "text": cumulative}
|
||||
yield {
|
||||
"type": "content",
|
||||
"text": _strip_tool_markup(
|
||||
cumulative, final = True
|
||||
),
|
||||
}
|
||||
else:
|
||||
cumulative = reasoning_text
|
||||
yield {"type": "content", "text": cumulative}
|
||||
|
|
@ -1713,7 +1925,11 @@ class LlamaCppBackend:
|
|||
cumulative += "</think>"
|
||||
in_thinking = False
|
||||
cumulative += token
|
||||
yield {"type": "content", "text": cumulative}
|
||||
cleaned = _strip_tool_markup(cumulative)
|
||||
# Only emit when cleaned text grows (monotonic).
|
||||
if len(cleaned) > len(_last_emitted):
|
||||
_last_emitted = cleaned
|
||||
yield {"type": "content", "text": cleaned}
|
||||
except json.JSONDecodeError:
|
||||
logger.debug(
|
||||
f"Skipping malformed SSE line: {line[:100]}"
|
||||
|
|
|
|||
|
|
@ -130,17 +130,17 @@ class InferenceOrchestrator:
|
|||
)
|
||||
if resp.status_code == 200:
|
||||
models = resp.json()
|
||||
# Top 8 GGUFs (frontend deduplicates against downloaded,
|
||||
# so we fetch extra to always fill 4 slots)
|
||||
# Top 40 GGUFs - frontend pages through them on-demand via
|
||||
# infinite scroll, so we send a deep pool.
|
||||
gguf_ids = [
|
||||
m["id"] for m in models if m.get("id", "").upper().endswith("-GGUF")
|
||||
][:8]
|
||||
# Top 8 non-GGUF hub models
|
||||
][:40]
|
||||
# Top 40 non-GGUF hub models
|
||||
hub_ids = [
|
||||
m["id"]
|
||||
for m in models
|
||||
if not m.get("id", "").upper().endswith("-GGUF")
|
||||
][:8]
|
||||
][:40]
|
||||
if gguf_ids:
|
||||
self._top_gguf_cache = gguf_ids
|
||||
logger.info("Top GGUF models: %s", gguf_ids)
|
||||
|
|
|
|||
|
|
@ -16,12 +16,42 @@ import sys
|
|||
import tempfile
|
||||
import threading
|
||||
|
||||
from loggers import get_logger
|
||||
from unsloth_zoo.rl_environments import check_signal_escape_patterns
|
||||
|
||||
logger = get_logger(__name__)
|
||||
|
||||
_EXEC_TIMEOUT = 300 # 5 minutes
|
||||
_MAX_OUTPUT_CHARS = 8000 # truncate long output
|
||||
_BASH_BLOCKED_WORDS = {"rm", "sudo", "dd", "chmod", "mkfs", "shutdown", "reboot"}
|
||||
|
||||
# Per-session working directories so each chat thread gets its own sandbox.
|
||||
# Falls back to a shared ~/studio_sandbox/ for API callers without a session_id.
|
||||
_workdirs: dict[str, str] = {}
|
||||
|
||||
|
||||
def _get_workdir(session_id: str | None = None) -> str:
|
||||
"""Return (and lazily create) a persistent working directory for tool execution."""
|
||||
global _workdirs
|
||||
key = session_id or "_default"
|
||||
if key not in _workdirs or not os.path.isdir(_workdirs[key]):
|
||||
home = os.path.expanduser("~")
|
||||
sandbox_root = os.path.join(home, "studio_sandbox")
|
||||
if session_id:
|
||||
# Sanitize: strip path separators and parent-dir references
|
||||
safe_id = os.path.basename(session_id.replace("..", ""))
|
||||
if not safe_id:
|
||||
safe_id = "_invalid"
|
||||
workdir = os.path.join(sandbox_root, safe_id)
|
||||
# Verify resolved path stays under sandbox root
|
||||
if not os.path.realpath(workdir).startswith(os.path.realpath(sandbox_root)):
|
||||
workdir = os.path.join(sandbox_root, "_invalid")
|
||||
else:
|
||||
workdir = sandbox_root
|
||||
os.makedirs(workdir, exist_ok = True)
|
||||
_workdirs[key] = workdir
|
||||
return _workdirs[key]
|
||||
|
||||
|
||||
WEB_SEARCH_TOOL = {
|
||||
"type": "function",
|
||||
|
|
@ -80,25 +110,47 @@ TERMINAL_TOOL = {
|
|||
ALL_TOOLS = [WEB_SEARCH_TOOL, PYTHON_TOOL, TERMINAL_TOOL]
|
||||
|
||||
|
||||
def execute_tool(name: str, arguments: dict, cancel_event = None) -> str:
|
||||
"""Execute a tool by name with the given arguments. Returns result as a string."""
|
||||
_TIMEOUT_UNSET = object()
|
||||
|
||||
|
||||
def execute_tool(
|
||||
name: str,
|
||||
arguments: dict,
|
||||
cancel_event = None,
|
||||
timeout: int | None = _TIMEOUT_UNSET,
|
||||
session_id: str | None = None,
|
||||
) -> str:
|
||||
"""Execute a tool by name with the given arguments. Returns result as a string.
|
||||
|
||||
``timeout``: int sets per-call limit in seconds, ``None`` means no limit,
|
||||
unset (default) uses ``_EXEC_TIMEOUT`` (300 s).
|
||||
``session_id``: optional thread/session ID for per-conversation sandbox isolation.
|
||||
"""
|
||||
logger.info(
|
||||
f"execute_tool: name={name}, session_id={session_id}, timeout={timeout}"
|
||||
)
|
||||
effective_timeout = _EXEC_TIMEOUT if timeout is _TIMEOUT_UNSET else timeout
|
||||
if name == "web_search":
|
||||
return _web_search(arguments.get("query", ""))
|
||||
return _web_search(arguments.get("query", ""), timeout = effective_timeout)
|
||||
if name == "python":
|
||||
return _python_exec(arguments.get("code", ""), cancel_event)
|
||||
return _python_exec(
|
||||
arguments.get("code", ""), cancel_event, effective_timeout, session_id
|
||||
)
|
||||
if name == "terminal":
|
||||
return _bash_exec(arguments.get("command", ""), cancel_event)
|
||||
return _bash_exec(
|
||||
arguments.get("command", ""), cancel_event, effective_timeout, session_id
|
||||
)
|
||||
return f"Unknown tool: {name}"
|
||||
|
||||
|
||||
def _web_search(query: str, max_results: int = 5) -> str:
|
||||
def _web_search(query: str, max_results: int = 5, timeout: int = _EXEC_TIMEOUT) -> str:
|
||||
"""Search the web using DuckDuckGo and return formatted results."""
|
||||
if not query.strip():
|
||||
return "No query provided."
|
||||
try:
|
||||
from ddgs import DDGS
|
||||
|
||||
results = DDGS().text(query, max_results = max_results)
|
||||
results = DDGS(timeout = timeout).text(query, max_results = max_results)
|
||||
if not results:
|
||||
return "No results found."
|
||||
parts = []
|
||||
|
|
@ -147,7 +199,12 @@ def _truncate(text: str, limit: int = _MAX_OUTPUT_CHARS) -> str:
|
|||
return text
|
||||
|
||||
|
||||
def _python_exec(code: str, cancel_event = None) -> str:
|
||||
def _python_exec(
|
||||
code: str,
|
||||
cancel_event = None,
|
||||
timeout: int = _EXEC_TIMEOUT,
|
||||
session_id: str | None = None,
|
||||
) -> str:
|
||||
"""Execute Python code in a subprocess sandbox."""
|
||||
if not code or not code.strip():
|
||||
return "No code provided."
|
||||
|
|
@ -158,8 +215,11 @@ def _python_exec(code: str, cancel_event = None) -> str:
|
|||
return error
|
||||
|
||||
tmp_path = None
|
||||
workdir = _get_workdir(session_id)
|
||||
try:
|
||||
fd, tmp_path = tempfile.mkstemp(suffix = ".py", prefix = "studio_exec_")
|
||||
fd, tmp_path = tempfile.mkstemp(
|
||||
suffix = ".py", prefix = "studio_exec_", dir = workdir
|
||||
)
|
||||
with os.fdopen(fd, "w") as f:
|
||||
f.write(code)
|
||||
|
||||
|
|
@ -168,7 +228,7 @@ def _python_exec(code: str, cancel_event = None) -> str:
|
|||
stdout = subprocess.PIPE,
|
||||
stderr = subprocess.STDOUT,
|
||||
text = True,
|
||||
cwd = tempfile.gettempdir(),
|
||||
cwd = workdir,
|
||||
)
|
||||
|
||||
# Spawn cancel watcher if we have a cancel event
|
||||
|
|
@ -179,11 +239,11 @@ def _python_exec(code: str, cancel_event = None) -> str:
|
|||
watcher.start()
|
||||
|
||||
try:
|
||||
output, _ = proc.communicate(timeout = _EXEC_TIMEOUT)
|
||||
output, _ = proc.communicate(timeout = timeout)
|
||||
except subprocess.TimeoutExpired:
|
||||
proc.kill()
|
||||
proc.communicate()
|
||||
return _truncate("Execution timed out after 5 minutes.")
|
||||
return _truncate(f"Execution timed out after {timeout} seconds.")
|
||||
|
||||
if cancel_event is not None and cancel_event.is_set():
|
||||
return "Execution cancelled."
|
||||
|
|
@ -203,7 +263,12 @@ def _python_exec(code: str, cancel_event = None) -> str:
|
|||
pass
|
||||
|
||||
|
||||
def _bash_exec(command: str, cancel_event = None) -> str:
|
||||
def _bash_exec(
|
||||
command: str,
|
||||
cancel_event = None,
|
||||
timeout: int = _EXEC_TIMEOUT,
|
||||
session_id: str | None = None,
|
||||
) -> str:
|
||||
"""Execute a bash command in a subprocess sandbox."""
|
||||
if not command or not command.strip():
|
||||
return "No command provided."
|
||||
|
|
@ -215,35 +280,35 @@ def _bash_exec(command: str, cancel_event = None) -> str:
|
|||
return f"Blocked command(s) for safety: {', '.join(sorted(blocked))}"
|
||||
|
||||
try:
|
||||
with tempfile.TemporaryDirectory() as tmpdir:
|
||||
proc = subprocess.Popen(
|
||||
["bash", "-c", command],
|
||||
stdout = subprocess.PIPE,
|
||||
stderr = subprocess.STDOUT,
|
||||
text = True,
|
||||
cwd = tmpdir,
|
||||
workdir = _get_workdir(session_id)
|
||||
proc = subprocess.Popen(
|
||||
["bash", "-c", command],
|
||||
stdout = subprocess.PIPE,
|
||||
stderr = subprocess.STDOUT,
|
||||
text = True,
|
||||
cwd = workdir,
|
||||
)
|
||||
|
||||
if cancel_event is not None:
|
||||
watcher = threading.Thread(
|
||||
target = _cancel_watcher, args = (proc, cancel_event), daemon = True
|
||||
)
|
||||
watcher.start()
|
||||
|
||||
if cancel_event is not None:
|
||||
watcher = threading.Thread(
|
||||
target = _cancel_watcher, args = (proc, cancel_event), daemon = True
|
||||
)
|
||||
watcher.start()
|
||||
try:
|
||||
output, _ = proc.communicate(timeout = timeout)
|
||||
except subprocess.TimeoutExpired:
|
||||
proc.kill()
|
||||
proc.communicate()
|
||||
return _truncate(f"Execution timed out after {timeout} seconds.")
|
||||
|
||||
try:
|
||||
output, _ = proc.communicate(timeout = _EXEC_TIMEOUT)
|
||||
except subprocess.TimeoutExpired:
|
||||
proc.kill()
|
||||
proc.communicate()
|
||||
return _truncate("Execution timed out after 5 minutes.")
|
||||
if cancel_event is not None and cancel_event.is_set():
|
||||
return "Execution cancelled."
|
||||
|
||||
if cancel_event is not None and cancel_event.is_set():
|
||||
return "Execution cancelled."
|
||||
|
||||
result = output or ""
|
||||
if proc.returncode != 0:
|
||||
result = f"Exit code {proc.returncode}:\n{result}"
|
||||
return _truncate(result) if result.strip() else "(no output)"
|
||||
result = output or ""
|
||||
if proc.returncode != 0:
|
||||
result = f"Exit code {proc.returncode}:\n{result}"
|
||||
return _truncate(result) if result.strip() else "(no output)"
|
||||
|
||||
except Exception as e:
|
||||
return f"Execution error: {e}"
|
||||
|
|
|
|||
|
|
@ -42,59 +42,25 @@ def _activate_transformers_version(model_name: str) -> None:
|
|||
if backend_path not in sys.path:
|
||||
sys.path.insert(0, backend_path)
|
||||
|
||||
from utils.transformers_version import needs_transformers_5, _resolve_base_model
|
||||
from utils.transformers_version import (
|
||||
needs_transformers_5,
|
||||
_resolve_base_model,
|
||||
_ensure_venv_t5_exists,
|
||||
_VENV_T5_DIR,
|
||||
)
|
||||
|
||||
resolved = _resolve_base_model(model_name)
|
||||
if needs_transformers_5(resolved):
|
||||
venv_t5 = os.path.join(
|
||||
os.path.expanduser("~"), ".unsloth", "studio", ".venv_t5"
|
||||
)
|
||||
if os.path.isdir(venv_t5):
|
||||
sys.path.insert(0, venv_t5)
|
||||
logger.info("Activated transformers 5.x from %s", venv_t5)
|
||||
else:
|
||||
# Fallback: pip install at runtime (slower, ~10-15s)
|
||||
logger.warning(".venv_t5 not found at %s — installing at runtime", venv_t5)
|
||||
import subprocess as sp
|
||||
|
||||
os.makedirs(venv_t5, exist_ok = True)
|
||||
r1 = sp.run(
|
||||
[
|
||||
sys.executable,
|
||||
"-m",
|
||||
"pip",
|
||||
"install",
|
||||
"--target",
|
||||
venv_t5,
|
||||
"--no-deps",
|
||||
"transformers==5.3.0",
|
||||
],
|
||||
stdout = sp.PIPE,
|
||||
stderr = sp.STDOUT,
|
||||
if not _ensure_venv_t5_exists():
|
||||
raise RuntimeError(
|
||||
f"Cannot activate transformers 5.x: .venv_t5 missing at {_VENV_T5_DIR}"
|
||||
)
|
||||
r2 = sp.run(
|
||||
[
|
||||
sys.executable,
|
||||
"-m",
|
||||
"pip",
|
||||
"install",
|
||||
"--target",
|
||||
venv_t5,
|
||||
"--no-deps",
|
||||
"huggingface_hub==1.3.0",
|
||||
],
|
||||
stdout = sp.PIPE,
|
||||
stderr = sp.STDOUT,
|
||||
)
|
||||
if r1.returncode != 0 or r2.returncode != 0:
|
||||
raise RuntimeError(
|
||||
f"Failed to install transformers 5.x into {venv_t5}. "
|
||||
f"pip returncode: transformers={r1.returncode}, huggingface_hub={r2.returncode}"
|
||||
)
|
||||
sys.path.insert(0, venv_t5)
|
||||
if _VENV_T5_DIR not in sys.path:
|
||||
sys.path.insert(0, _VENV_T5_DIR)
|
||||
logger.info("Activated transformers 5.x from %s", _VENV_T5_DIR)
|
||||
# Propagate to child subprocesses (e.g. GGUF converter)
|
||||
_pp = os.environ.get("PYTHONPATH", "")
|
||||
os.environ["PYTHONPATH"] = venv_t5 + (os.pathsep + _pp if _pp else "")
|
||||
os.environ["PYTHONPATH"] = _VENV_T5_DIR + (os.pathsep + _pp if _pp else "")
|
||||
else:
|
||||
logger.info("Using default transformers (4.57.x) for %s", model_name)
|
||||
|
||||
|
|
|
|||
|
|
@ -466,6 +466,7 @@ class UnslothTrainer:
|
|||
is_dataset_image: bool = False,
|
||||
is_dataset_audio: bool = False,
|
||||
trust_remote_code: bool = False,
|
||||
full_finetuning: bool = False,
|
||||
) -> bool:
|
||||
"""Load model for training (supports both text and vision models)"""
|
||||
self.load_in_4bit = load_in_4bit # Store for training_meta.json
|
||||
|
|
@ -612,6 +613,7 @@ class UnslothTrainer:
|
|||
dtype = None,
|
||||
auto_model = CsmForConditionalGeneration,
|
||||
load_in_4bit = False,
|
||||
full_finetuning = full_finetuning,
|
||||
token = hf_token,
|
||||
trust_remote_code = trust_remote_code,
|
||||
)
|
||||
|
|
@ -626,6 +628,7 @@ class UnslothTrainer:
|
|||
model_name = model_name,
|
||||
dtype = None,
|
||||
load_in_4bit = False,
|
||||
full_finetuning = full_finetuning,
|
||||
auto_model = WhisperForConditionalGeneration,
|
||||
whisper_language = "English",
|
||||
whisper_task = "transcribe",
|
||||
|
|
@ -646,6 +649,7 @@ class UnslothTrainer:
|
|||
max_seq_length = max_seq_length,
|
||||
dtype = None,
|
||||
load_in_4bit = load_in_4bit,
|
||||
full_finetuning = full_finetuning,
|
||||
token = hf_token,
|
||||
trust_remote_code = trust_remote_code,
|
||||
)
|
||||
|
|
@ -684,6 +688,7 @@ class UnslothTrainer:
|
|||
max_seq_length = max_seq_length,
|
||||
dtype = torch.float32, # Spark-TTS requires float32
|
||||
load_in_4bit = False,
|
||||
full_finetuning = full_finetuning,
|
||||
token = hf_token,
|
||||
trust_remote_code = trust_remote_code,
|
||||
)
|
||||
|
|
@ -697,6 +702,7 @@ class UnslothTrainer:
|
|||
model_name,
|
||||
max_seq_length = max_seq_length,
|
||||
load_in_4bit = False,
|
||||
full_finetuning = full_finetuning,
|
||||
token = hf_token,
|
||||
trust_remote_code = trust_remote_code,
|
||||
)
|
||||
|
|
@ -712,6 +718,7 @@ class UnslothTrainer:
|
|||
max_seq_length = max_seq_length,
|
||||
dtype = None,
|
||||
load_in_4bit = load_in_4bit,
|
||||
full_finetuning = full_finetuning,
|
||||
token = hf_token,
|
||||
trust_remote_code = trust_remote_code,
|
||||
)
|
||||
|
|
@ -724,6 +731,7 @@ class UnslothTrainer:
|
|||
max_seq_length = max_seq_length,
|
||||
dtype = None, # Auto-detect
|
||||
load_in_4bit = load_in_4bit,
|
||||
full_finetuning = full_finetuning,
|
||||
token = hf_token,
|
||||
trust_remote_code = trust_remote_code,
|
||||
)
|
||||
|
|
@ -755,6 +763,7 @@ class UnslothTrainer:
|
|||
max_seq_length = max_seq_length,
|
||||
dtype = None, # Auto-detect
|
||||
load_in_4bit = load_in_4bit,
|
||||
full_finetuning = full_finetuning,
|
||||
token = hf_token,
|
||||
trust_remote_code = trust_remote_code,
|
||||
)
|
||||
|
|
@ -779,13 +788,14 @@ class UnslothTrainer:
|
|||
self._source_code_retried = True
|
||||
logger.info(f"\n'could not get source code' — retrying once...\n")
|
||||
return self.load_model(
|
||||
model_name,
|
||||
max_seq_length,
|
||||
load_in_4bit,
|
||||
hf_token,
|
||||
is_dataset_image,
|
||||
is_dataset_audio,
|
||||
trust_remote_code,
|
||||
model_name = model_name,
|
||||
max_seq_length = max_seq_length,
|
||||
load_in_4bit = load_in_4bit,
|
||||
hf_token = hf_token,
|
||||
is_dataset_image = is_dataset_image,
|
||||
is_dataset_audio = is_dataset_audio,
|
||||
trust_remote_code = trust_remote_code,
|
||||
full_finetuning = full_finetuning,
|
||||
)
|
||||
error_msg = str(e)
|
||||
error_lower = error_msg.lower()
|
||||
|
|
|
|||
|
|
@ -36,59 +36,25 @@ def _activate_transformers_version(model_name: str) -> None:
|
|||
if backend_path not in sys.path:
|
||||
sys.path.insert(0, backend_path)
|
||||
|
||||
from utils.transformers_version import needs_transformers_5, _resolve_base_model
|
||||
from utils.transformers_version import (
|
||||
needs_transformers_5,
|
||||
_resolve_base_model,
|
||||
_ensure_venv_t5_exists,
|
||||
_VENV_T5_DIR,
|
||||
)
|
||||
|
||||
resolved = _resolve_base_model(model_name)
|
||||
if needs_transformers_5(resolved):
|
||||
venv_t5 = os.path.join(
|
||||
os.path.expanduser("~"), ".unsloth", "studio", ".venv_t5"
|
||||
)
|
||||
if os.path.isdir(venv_t5):
|
||||
sys.path.insert(0, venv_t5)
|
||||
logger.info("Activated transformers 5.x from %s", venv_t5)
|
||||
else:
|
||||
# Fallback: pip install at runtime (slower, ~10-15s)
|
||||
logger.warning(".venv_t5 not found at %s — installing at runtime", venv_t5)
|
||||
import subprocess as sp
|
||||
|
||||
os.makedirs(venv_t5, exist_ok = True)
|
||||
r1 = sp.run(
|
||||
[
|
||||
sys.executable,
|
||||
"-m",
|
||||
"pip",
|
||||
"install",
|
||||
"--target",
|
||||
venv_t5,
|
||||
"--no-deps",
|
||||
"transformers==5.3.0",
|
||||
],
|
||||
stdout = sp.PIPE,
|
||||
stderr = sp.STDOUT,
|
||||
if not _ensure_venv_t5_exists():
|
||||
raise RuntimeError(
|
||||
f"Cannot activate transformers 5.x: .venv_t5 missing at {_VENV_T5_DIR}"
|
||||
)
|
||||
r2 = sp.run(
|
||||
[
|
||||
sys.executable,
|
||||
"-m",
|
||||
"pip",
|
||||
"install",
|
||||
"--target",
|
||||
venv_t5,
|
||||
"--no-deps",
|
||||
"huggingface_hub==1.3.0",
|
||||
],
|
||||
stdout = sp.PIPE,
|
||||
stderr = sp.STDOUT,
|
||||
)
|
||||
if r1.returncode != 0 or r2.returncode != 0:
|
||||
raise RuntimeError(
|
||||
f"Failed to install transformers 5.x into {venv_t5}. "
|
||||
f"pip returncode: transformers={r1.returncode}, huggingface_hub={r2.returncode}"
|
||||
)
|
||||
sys.path.insert(0, venv_t5)
|
||||
if _VENV_T5_DIR not in sys.path:
|
||||
sys.path.insert(0, _VENV_T5_DIR)
|
||||
logger.info("Activated transformers 5.x from %s", _VENV_T5_DIR)
|
||||
# Propagate to child subprocesses (e.g. GGUF converter)
|
||||
_pp = os.environ.get("PYTHONPATH", "")
|
||||
os.environ["PYTHONPATH"] = venv_t5 + (os.pathsep + _pp if _pp else "")
|
||||
os.environ["PYTHONPATH"] = _VENV_T5_DIR + (os.pathsep + _pp if _pp else "")
|
||||
else:
|
||||
logger.info("Using default transformers (4.57.x) for %s", model_name)
|
||||
|
||||
|
|
@ -444,12 +410,16 @@ def run_training_process(
|
|||
_tqdm_thread = _th.Thread(target = _monitor_tqdm, daemon = True)
|
||||
_tqdm_thread.start()
|
||||
|
||||
training_type = config.get("training_type", "LoRA/QLoRA")
|
||||
use_lora = training_type == "LoRA/QLoRA"
|
||||
|
||||
# ── 4c. Load training model (uses VRAM — dataset already formatted) ──
|
||||
_send_status(event_queue, "Loading model...")
|
||||
success = trainer.load_model(
|
||||
model_name = model_name,
|
||||
max_seq_length = config["max_seq_length"],
|
||||
load_in_4bit = config["load_in_4bit"],
|
||||
full_finetuning = not use_lora,
|
||||
hf_token = hf_token,
|
||||
is_dataset_image = config.get("is_dataset_image", False),
|
||||
is_dataset_audio = config.get("is_dataset_audio", False),
|
||||
|
|
@ -473,8 +443,6 @@ def run_training_process(
|
|||
return
|
||||
|
||||
# ── 4d. Prepare model (LoRA or full finetuning) ──
|
||||
training_type = config.get("training_type", "LoRA/QLoRA")
|
||||
use_lora = training_type == "LoRA/QLoRA"
|
||||
if use_lora:
|
||||
_send_status(event_queue, "Configuring LoRA adapters...")
|
||||
success = trainer.prepare_model_for_training(
|
||||
|
|
|
|||
|
|
@ -318,6 +318,24 @@ class ChatCompletionRequest(BaseModel):
|
|||
None,
|
||||
description = "[x-unsloth] List of enabled tool names (e.g. ['web_search', 'python', 'terminal']). If None, all tools are enabled.",
|
||||
)
|
||||
auto_heal_tool_calls: Optional[bool] = Field(
|
||||
True,
|
||||
description = "[x-unsloth] Auto-detect and fix malformed tool calls from model output.",
|
||||
)
|
||||
max_tool_calls_per_message: Optional[int] = Field(
|
||||
10,
|
||||
ge = 0,
|
||||
description = "[x-unsloth] Maximum number of tool call iterations per message (0 = disabled, 9999 = unlimited).",
|
||||
)
|
||||
tool_call_timeout: Optional[int] = Field(
|
||||
300,
|
||||
ge = 1,
|
||||
description = "[x-unsloth] Timeout in seconds for each tool call execution (9999 = no limit).",
|
||||
)
|
||||
session_id: Optional[str] = Field(
|
||||
None,
|
||||
description = "[x-unsloth] Session/thread ID for scoping tool execution sandbox.",
|
||||
)
|
||||
|
||||
|
||||
# ── Streaming response chunks ────────────────────────────────────
|
||||
|
|
|
|||
|
|
@ -8,7 +8,7 @@ snac
|
|||
# TRL and related packages
|
||||
trl==0.23.1
|
||||
git+https://github.com/meta-pytorch/OpenEnv.git
|
||||
executorch>=1.0.1
|
||||
# executorch>=1.0.1 # 41.5 MB - no imports in unsloth/zoo/studio
|
||||
torch-c-dlpack-ext
|
||||
sentence_transformers==5.2.0
|
||||
transformers==4.57.1
|
||||
transformers==4.57.6
|
||||
|
|
|
|||
|
|
@ -4,7 +4,7 @@ tomli-w
|
|||
|
||||
# ExecuTorch dependencies
|
||||
ruamel.yaml
|
||||
coremltools
|
||||
# coremltools # 10.2 MB - Apple CoreML, no imports in unsloth/zoo/studio
|
||||
expecttest
|
||||
flatbuffers
|
||||
hydra-core
|
||||
|
|
@ -15,7 +15,7 @@ pytest<9.0
|
|||
pytest-json-report
|
||||
pytest-rerunfailures==15.1
|
||||
pytest-xdist
|
||||
# Also needed by sentence_transformers
|
||||
# Also needed by sentence_transformers (installed with --no-deps in extras-no-deps.txt)
|
||||
scikit-learn==1.7.1
|
||||
|
||||
# Additional extras
|
||||
|
|
@ -26,8 +26,8 @@ omegaconf
|
|||
einx
|
||||
pyloudnorm
|
||||
openai-whisper
|
||||
uroman
|
||||
MeCab
|
||||
# uroman # 4.0 MB - romanization, no imports found
|
||||
# MeCab # 19.9 MB - Japanese tokenizer, no imports found
|
||||
loguru
|
||||
flatten_dict
|
||||
ffmpy
|
||||
|
|
|
|||
|
|
@ -1,6 +1,6 @@
|
|||
# Single-env pins for unsloth + studio + data-designer
|
||||
# Keep compatible with unsloth transformers bounds.
|
||||
transformers==4.57.1
|
||||
transformers==4.57.6
|
||||
trl==0.23.1
|
||||
huggingface-hub==0.36.2
|
||||
|
||||
|
|
|
|||
|
|
@ -10,7 +10,7 @@ datasets==4.3.0
|
|||
pyjwt
|
||||
easydict
|
||||
addict
|
||||
gradio>=4.0.0
|
||||
# gradio>=4.0.0 # 148 MB - Studio uses React + FastAPI, not Gradio
|
||||
huggingface-hub==0.36.2
|
||||
structlog>=24.1.0
|
||||
diceware
|
||||
|
|
|
|||
|
|
@ -1028,7 +1028,7 @@ async def openai_chat_completions(
|
|||
if use_tools:
|
||||
from core.inference.tools import ALL_TOOLS
|
||||
|
||||
if payload.enabled_tools:
|
||||
if payload.enabled_tools is not None:
|
||||
tools_to_use = [
|
||||
t
|
||||
for t in ALL_TOOLS
|
||||
|
|
@ -1050,6 +1050,16 @@ async def openai_chat_completions(
|
|||
presence_penalty = payload.presence_penalty,
|
||||
cancel_event = cancel_event,
|
||||
enable_thinking = payload.enable_thinking,
|
||||
auto_heal_tool_calls = payload.auto_heal_tool_calls
|
||||
if payload.auto_heal_tool_calls is not None
|
||||
else True,
|
||||
max_tool_iterations = payload.max_tool_calls_per_message
|
||||
if payload.max_tool_calls_per_message is not None
|
||||
else 10,
|
||||
tool_call_timeout = payload.tool_call_timeout
|
||||
if payload.tool_call_timeout is not None
|
||||
else 300,
|
||||
session_id = payload.session_id,
|
||||
)
|
||||
|
||||
_tool_sentinel = object()
|
||||
|
|
@ -1093,6 +1103,10 @@ async def openai_chat_completions(
|
|||
yield f"data: {status_data}\n\n"
|
||||
continue
|
||||
|
||||
if event["type"] in ("tool_start", "tool_end"):
|
||||
yield f"data: {json.dumps(event)}\n\n"
|
||||
continue
|
||||
|
||||
# "content" type -- cumulative text
|
||||
cumulative = event.get("text", "")
|
||||
new_text = cumulative[len(prev_text) :]
|
||||
|
|
|
|||
|
|
@ -26,6 +26,7 @@ import json
|
|||
import structlog
|
||||
from loggers import get_logger
|
||||
import os
|
||||
import shutil
|
||||
import subprocess
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
|
@ -58,7 +59,7 @@ _tokenizer_class_cache: dict[str, bool] = {}
|
|||
|
||||
# Versions
|
||||
TRANSFORMERS_5_VERSION = "5.3.0"
|
||||
TRANSFORMERS_DEFAULT_VERSION = "4.57.1"
|
||||
TRANSFORMERS_DEFAULT_VERSION = "4.57.6"
|
||||
|
||||
# Pre-installed directory for transformers 5.x — created by setup.sh / setup.ps1
|
||||
_VENV_T5_DIR = str(Path.home() / ".unsloth" / "studio" / ".venv_t5")
|
||||
|
|
@ -216,15 +217,87 @@ def _purge_modules() -> int:
|
|||
return len(to_remove)
|
||||
|
||||
|
||||
def _ensure_venv_t5_exists() -> bool:
|
||||
"""Ensure .venv_t5/ exists. Install at runtime if missing."""
|
||||
if os.path.isdir(_VENV_T5_DIR) and os.listdir(_VENV_T5_DIR):
|
||||
return True
|
||||
_VENV_T5_PACKAGES = (
|
||||
f"transformers=={TRANSFORMERS_5_VERSION}",
|
||||
"huggingface_hub==1.7.1",
|
||||
"hf_xet==1.4.2",
|
||||
"tiktoken",
|
||||
)
|
||||
|
||||
logger.warning(".venv_t5 not found at %s — installing at runtime", _VENV_T5_DIR)
|
||||
os.makedirs(_VENV_T5_DIR, exist_ok = True)
|
||||
for pkg in (f"transformers=={TRANSFORMERS_5_VERSION}", "huggingface_hub==1.3.0"):
|
||||
cmd = [
|
||||
|
||||
def _venv_t5_is_valid() -> bool:
|
||||
"""Return True if .venv_t5/ has all required packages at the correct versions."""
|
||||
if not os.path.isdir(_VENV_T5_DIR) or not os.listdir(_VENV_T5_DIR):
|
||||
return False
|
||||
# Check that the key package directories exist AND match the required version
|
||||
for pkg_spec in _VENV_T5_PACKAGES:
|
||||
parts = pkg_spec.split("==")
|
||||
pkg_name = parts[0]
|
||||
pkg_version = parts[1] if len(parts) > 1 else None
|
||||
pkg_name_norm = pkg_name.replace("-", "_")
|
||||
# Check directory exists
|
||||
if not any(
|
||||
(Path(_VENV_T5_DIR) / d).is_dir()
|
||||
for d in (pkg_name_norm, pkg_name_norm.replace("_", "-"))
|
||||
):
|
||||
return False
|
||||
# For unpinned packages, existence is enough
|
||||
if pkg_version is None:
|
||||
continue
|
||||
# Check version via .dist-info metadata
|
||||
dist_info_found = False
|
||||
for di in Path(_VENV_T5_DIR).glob(f"{pkg_name_norm}-*.dist-info"):
|
||||
metadata = di / "METADATA"
|
||||
if not metadata.is_file():
|
||||
continue
|
||||
for line in metadata.read_text(errors = "replace").splitlines():
|
||||
if line.startswith("Version:"):
|
||||
installed_ver = line.split(":", 1)[1].strip()
|
||||
if installed_ver != pkg_version:
|
||||
logger.info(
|
||||
".venv_t5 has %s==%s but need %s",
|
||||
pkg_name,
|
||||
installed_ver,
|
||||
pkg_version,
|
||||
)
|
||||
return False
|
||||
dist_info_found = True
|
||||
break
|
||||
if dist_info_found:
|
||||
break
|
||||
if not dist_info_found:
|
||||
return False
|
||||
return True
|
||||
|
||||
|
||||
def _install_to_venv_t5(pkg: str) -> bool:
|
||||
"""Install a single package into .venv_t5/, preferring uv then pip."""
|
||||
# Try uv first (faster) if already on PATH -- do NOT install uv at runtime
|
||||
if shutil.which("uv"):
|
||||
result = subprocess.run(
|
||||
[
|
||||
"uv",
|
||||
"pip",
|
||||
"install",
|
||||
"--python",
|
||||
sys.executable,
|
||||
"--target",
|
||||
_VENV_T5_DIR,
|
||||
"--no-deps",
|
||||
"--upgrade",
|
||||
pkg,
|
||||
],
|
||||
stdout = subprocess.PIPE,
|
||||
stderr = subprocess.STDOUT,
|
||||
text = True,
|
||||
)
|
||||
if result.returncode == 0:
|
||||
return True
|
||||
logger.warning("uv install of %s failed, falling back to pip", pkg)
|
||||
|
||||
# Fallback to pip
|
||||
result = subprocess.run(
|
||||
[
|
||||
sys.executable,
|
||||
"-m",
|
||||
"pip",
|
||||
|
|
@ -232,13 +305,31 @@ def _ensure_venv_t5_exists() -> bool:
|
|||
"--target",
|
||||
_VENV_T5_DIR,
|
||||
"--no-deps",
|
||||
"--upgrade",
|
||||
pkg,
|
||||
]
|
||||
result = subprocess.run(
|
||||
cmd, stdout = subprocess.PIPE, stderr = subprocess.STDOUT, text = True
|
||||
)
|
||||
if result.returncode != 0:
|
||||
logger.error("pip install failed:\n%s", result.stdout)
|
||||
],
|
||||
stdout = subprocess.PIPE,
|
||||
stderr = subprocess.STDOUT,
|
||||
text = True,
|
||||
)
|
||||
if result.returncode != 0:
|
||||
logger.error("install failed:\n%s", result.stdout)
|
||||
return False
|
||||
return True
|
||||
|
||||
|
||||
def _ensure_venv_t5_exists() -> bool:
|
||||
"""Ensure .venv_t5/ exists with all required packages. Install if missing."""
|
||||
if _venv_t5_is_valid():
|
||||
return True
|
||||
|
||||
logger.warning(
|
||||
".venv_t5 not found or incomplete at %s -- installing at runtime", _VENV_T5_DIR
|
||||
)
|
||||
shutil.rmtree(_VENV_T5_DIR, ignore_errors = True)
|
||||
os.makedirs(_VENV_T5_DIR, exist_ok = True)
|
||||
for pkg in _VENV_T5_PACKAGES:
|
||||
if not _install_to_venv_t5(pkg):
|
||||
return False
|
||||
logger.info("Installed transformers 5.x to %s", _VENV_T5_DIR)
|
||||
return True
|
||||
|
|
|
|||
|
|
@ -5,7 +5,7 @@
|
|||
"": {
|
||||
"name": "unsloth-theme",
|
||||
"dependencies": {
|
||||
"@assistant-ui/react": "^0.12.17",
|
||||
"@assistant-ui/react": "^0.12.19",
|
||||
"@assistant-ui/react-markdown": "^0.12.3",
|
||||
"@assistant-ui/react-streamdown": "^0.1.2",
|
||||
"@base-ui/react": "^1.2.0",
|
||||
|
|
@ -43,7 +43,7 @@
|
|||
"framer-motion": "^11.18.2",
|
||||
"js-yaml": "^4.1.1",
|
||||
"katex": "^0.16.28",
|
||||
"lucide-react": "^0.575.0",
|
||||
"lucide-react": "^0.577.0",
|
||||
"mammoth": "^1.11.0",
|
||||
"motion": "^12.34.0",
|
||||
"next": "^16.1.6",
|
||||
|
|
@ -88,17 +88,17 @@
|
|||
|
||||
"@antfu/ni": ["@antfu/ni@25.0.0", "", { "dependencies": { "ansis": "^4.0.0", "fzf": "^0.5.2", "package-manager-detector": "^1.3.0", "tinyexec": "^1.0.1" }, "bin": { "na": "bin/na.mjs", "ni": "bin/ni.mjs", "nr": "bin/nr.mjs", "nci": "bin/nci.mjs", "nlx": "bin/nlx.mjs", "nun": "bin/nun.mjs", "nup": "bin/nup.mjs" } }, "sha512-9q/yCljni37pkMr4sPrI3G4jqdIk074+iukc5aFJl7kmDCCsiJrbZ6zKxnES1Gwg+i9RcDZwvktl23puGslmvA=="],
|
||||
|
||||
"@assistant-ui/core": ["@assistant-ui/core@0.1.5", "", { "dependencies": { "assistant-stream": "^0.3.5", "nanoid": "^5.1.6" }, "peerDependencies": { "@assistant-ui/store": "^0.2.2", "@assistant-ui/tap": "^0.5.2", "@types/react": "*", "assistant-cloud": "^0.1.21", "react": "^18 || ^19", "zustand": "^5.0.11" }, "optionalPeers": ["@types/react", "assistant-cloud", "react", "zustand"] }, "sha512-kLqFbRULZvE+hIwxGz705BW3QYhfwiVaVWoolfTGYkg+4xwah1PGuH0zqjXP5AMADtz+L69Lp+LX0xU9MQZ0DA=="],
|
||||
"@assistant-ui/core": ["@assistant-ui/core@0.1.7", "", { "dependencies": { "assistant-stream": "^0.3.6", "nanoid": "^5.1.6" }, "peerDependencies": { "@assistant-ui/store": "^0.2.3", "@assistant-ui/tap": "^0.5.3", "@types/react": "*", "assistant-cloud": "^0.1.22", "react": "^18 || ^19", "zustand": "^5.0.11" }, "optionalPeers": ["@types/react", "assistant-cloud", "react", "zustand"] }, "sha512-219T42ihVOicbJXZLWgD2CW5Bylg9Nk7geC331X4RfJxTDYlm2zIjViGlGaqfj6URXBp6kMulO2BTUrHGmAvdw=="],
|
||||
|
||||
"@assistant-ui/react": ["@assistant-ui/react@0.12.17", "", { "dependencies": { "@assistant-ui/core": "^0.1.5", "@assistant-ui/store": "^0.2.2", "@assistant-ui/tap": "^0.5.2", "@radix-ui/primitive": "^1.1.3", "@radix-ui/react-compose-refs": "^1.1.2", "@radix-ui/react-context": "^1.1.3", "@radix-ui/react-primitive": "^2.1.4", "@radix-ui/react-use-callback-ref": "^1.1.1", "@radix-ui/react-use-escape-keydown": "^1.1.1", "assistant-cloud": "^0.1.21", "assistant-stream": "^0.3.4", "nanoid": "^5.1.6", "radix-ui": "^1.4.3", "react-textarea-autosize": "^8.5.9", "zod": "^4.3.6", "zustand": "^5.0.11" }, "peerDependencies": { "@types/react": "*", "@types/react-dom": "*", "react": "^18 || ^19", "react-dom": "^18 || ^19" }, "optionalPeers": ["@types/react", "@types/react-dom"] }, "sha512-t4Z8LatD3LQrtURLaYPG47r4iG7UQgkdoi5YEv+EhzvYiG8I7kAyV4SbnFH6sXPrnleV4IpBHAd8Wc7ynkQtsw=="],
|
||||
"@assistant-ui/react": ["@assistant-ui/react@0.12.19", "", { "dependencies": { "@assistant-ui/core": "^0.1.7", "@assistant-ui/store": "^0.2.3", "@assistant-ui/tap": "^0.5.3", "@radix-ui/primitive": "^1.1.3", "@radix-ui/react-compose-refs": "^1.1.2", "@radix-ui/react-context": "^1.1.3", "@radix-ui/react-primitive": "^2.1.4", "@radix-ui/react-use-callback-ref": "^1.1.1", "@radix-ui/react-use-escape-keydown": "^1.1.1", "assistant-cloud": "^0.1.22", "assistant-stream": "^0.3.6", "nanoid": "^5.1.6", "radix-ui": "^1.4.3", "react-textarea-autosize": "^8.5.9", "zod": "^4.3.6", "zustand": "^5.0.11" }, "peerDependencies": { "@types/react": "*", "@types/react-dom": "*", "react": "^18 || ^19", "react-dom": "^18 || ^19" }, "optionalPeers": ["@types/react", "@types/react-dom"] }, "sha512-scAf0o8cwjuHT9Y44EFGXcE2y6BSmpeMvt0NxOn8+Y/HBlNttQMLNvrM0p2AjacXCUufagiafAnWybzBV3nKEQ=="],
|
||||
|
||||
"@assistant-ui/react-markdown": ["@assistant-ui/react-markdown@0.12.4", "", { "dependencies": { "@radix-ui/react-primitive": "^2.1.4", "@radix-ui/react-use-callback-ref": "^1.1.1", "classnames": "^2.5.1", "react-markdown": "^10.1.0" }, "peerDependencies": { "@assistant-ui/react": "^0.12.11", "@types/react": "*", "react": "^18 || ^19" }, "optionalPeers": ["@types/react"] }, "sha512-6TD9guiuLJxJoOwSjNHUYAVma2ctDCG9uypUqKHE0OUhDwTDD3NsMvTnQ0n0Lh8nnCEwVglOwKKlSEYpV7SnWA=="],
|
||||
|
||||
"@assistant-ui/react-streamdown": ["@assistant-ui/react-streamdown@0.1.3", "", { "dependencies": { "rehype-harden": "^1.1.7", "rehype-raw": "^7.0.0", "rehype-sanitize": "^6.0.0", "streamdown": "^2.1.0" }, "peerDependencies": { "@assistant-ui/react": "^0.12.11", "@streamdown/cjk": "^1.0.0", "@streamdown/code": "^1.0.0", "@streamdown/math": "^1.0.0", "@streamdown/mermaid": "^1.0.0", "@types/react": "*", "react": "^18 || ^19" }, "optionalPeers": ["@streamdown/cjk", "@streamdown/code", "@streamdown/math", "@streamdown/mermaid", "@types/react"] }, "sha512-n1UCjXQ3svmDtJBMJj/vXqz/BqAQBuy7myrXeymz2tD9l+ENQgqu2JY5ir3J19juJTe5lsi/P3+tOJ2C1jc/nw=="],
|
||||
|
||||
"@assistant-ui/store": ["@assistant-ui/store@0.2.2", "", { "dependencies": { "use-effect-event": "^2.0.3" }, "peerDependencies": { "@assistant-ui/tap": "^0.5.2", "@types/react": "*", "react": "^18 || ^19" }, "optionalPeers": ["@types/react"] }, "sha512-JzQseWFp3UmbByBSWQmiGi/bz5jbfru04hIgb2DJBpnnTyns8Zl+8wDPnwiYGF/6SA+IzTg5M0V1wf77rwU0dA=="],
|
||||
"@assistant-ui/store": ["@assistant-ui/store@0.2.3", "", { "dependencies": { "use-effect-event": "^2.0.3" }, "peerDependencies": { "@assistant-ui/tap": "^0.5.3", "@types/react": "*", "react": "^18 || ^19" }, "optionalPeers": ["@types/react"] }, "sha512-daStbgSQiX7+csqK6Cvo7A8p8UZkTCSMxBHxbhJvwrlVbp7BRJWTxq3U3rpTkSGIar23SXIyVRRfXU8VW7pswA=="],
|
||||
|
||||
"@assistant-ui/tap": ["@assistant-ui/tap@0.5.2", "", { "peerDependencies": { "@types/react": "*", "react": "^18 || ^19" }, "optionalPeers": ["@types/react", "react"] }, "sha512-w6gXhr+mF6cPG6ZCnkqV4kkOHzR+Fb+52S4T34PnrH0cs8l2Gqlwo/kB9BcB9fGmjwL7izdwubQ7t2VBhWpz/Q=="],
|
||||
"@assistant-ui/tap": ["@assistant-ui/tap@0.5.3", "", { "peerDependencies": { "@types/react": "*", "react": "^18 || ^19" }, "optionalPeers": ["@types/react", "react"] }, "sha512-wy06ksqF2LfFxe4JXy31Ns89N/be1Dy3c+mG363cFHFp3CbLkRu8CrCN2SQSgCkXt628E+D8QyzqdBcl9kD4NQ=="],
|
||||
|
||||
"@babel/code-frame": ["@babel/code-frame@7.29.0", "", { "dependencies": { "@babel/helper-validator-identifier": "^7.28.5", "js-tokens": "^4.0.0", "picocolors": "^1.1.1" } }, "sha512-9NhCeYjq9+3uxgdtp20LSiJXJvN0FeCtNGpJxuMFZ1Kv3cWUNb6DOhJwUvcVCzKGR66cw4njwM6hrJLqgOwbcw=="],
|
||||
|
||||
|
|
@ -846,7 +846,7 @@
|
|||
|
||||
"aria-hidden": ["aria-hidden@1.2.6", "", { "dependencies": { "tslib": "^2.0.0" } }, "sha512-ik3ZgC9dY/lYVVM++OISsaYDeg1tb0VtP5uL3ouh1koGOaUMDPpbFIei4JkFimWUFPn90sbMNMXQAIVOlnYKJA=="],
|
||||
|
||||
"assistant-cloud": ["assistant-cloud@0.1.21", "", { "dependencies": { "assistant-stream": "^0.3.4" } }, "sha512-KZ9ZsF1i1zMhozvD4m8TsmTdtufqULMaqgOoSLRyVtnhwvxkDufL87tSjv7epddZ4kbebe31biWSg7KIlgzvQA=="],
|
||||
"assistant-cloud": ["assistant-cloud@0.1.22", "", { "dependencies": { "assistant-stream": "^0.3.6" } }, "sha512-AEE9shV+oFrGDv/MRTRERctNKpIYS0n34UpAQXXICiOkSWD6QZnS1ljLqruFko7fJoT5CIWq8dNeJWdzQLTBLg=="],
|
||||
|
||||
"assistant-stream": ["assistant-stream@0.3.3", "", { "dependencies": { "@standard-schema/spec": "^1.1.0", "nanoid": "^5.1.6", "secure-json-parse": "^4.1.0" } }, "sha512-Ne/uTseMIiZx740dTbr/SWxONM8nYj4Z5BRmUfqQN+TNgtOCgWOlC/oTUQ+A7LIUHtmGbcoyZwDf8yd2RASnDA=="],
|
||||
|
||||
|
|
@ -1450,7 +1450,7 @@
|
|||
|
||||
"lru-cache": ["lru-cache@5.1.1", "", { "dependencies": { "yallist": "^3.0.2" } }, "sha512-KpNARQA3Iwv+jTA0utUVVbrh+Jlrr1Fv0e56GGzAFOXN7dk/FviaDW8LHmK52DlcH4WP2n6gI8vN1aesBFgo9w=="],
|
||||
|
||||
"lucide-react": ["lucide-react@0.575.0", "", { "peerDependencies": { "react": "^16.5.1 || ^17.0.0 || ^18.0.0 || ^19.0.0" } }, "sha512-VuXgKZrk0uiDlWjGGXmKV6MSk9Yy4l10qgVvzGn2AWBx1Ylt0iBexKOAoA6I7JO3m+M9oeovJd3yYENfkUbOeg=="],
|
||||
"lucide-react": ["lucide-react@0.577.0", "", { "peerDependencies": { "react": "^16.5.1 || ^17.0.0 || ^18.0.0 || ^19.0.0" } }, "sha512-4LjoFv2eEPwYDPg/CUdBJQSDfPyzXCRrVW1X7jrx/trgxnxkHFjnVZINbzvzxjN70dxychOfg+FTYwBiS3pQ5A=="],
|
||||
|
||||
"magic-string": ["magic-string@0.30.21", "", { "dependencies": { "@jridgewell/sourcemap-codec": "^1.5.5" } }, "sha512-vd2F4YUyEXKGcLHoq+TEyCjxueSeHnFxyyjNp80yg0XV4vUhnDer/lvvlqM/arB5bXQN5K2/3oinyCRyx8T2CQ=="],
|
||||
|
||||
|
|
@ -2082,9 +2082,9 @@
|
|||
|
||||
"zwitch": ["zwitch@2.0.4", "", {}, "sha512-bXE4cR/kVZhKZX/RjPEflHaKVhUVl85noU3v6b8apfQEc1x4A+zBxjZ4lN8LqGd6WZ3dl98pY4o717VFmoPp+A=="],
|
||||
|
||||
"@assistant-ui/core/assistant-stream": ["assistant-stream@0.3.5", "", { "dependencies": { "@standard-schema/spec": "^1.1.0", "nanoid": "^5.1.6", "secure-json-parse": "^4.1.0" } }, "sha512-OGxVClfpEOoSsJDraPoe+GYTwh9TJX1wxK3hT5Qs7gIOyD/MZbqwyWwabRO2KTnNU4w7usvmC/vneUzxSk4bBg=="],
|
||||
"@assistant-ui/core/assistant-stream": ["assistant-stream@0.3.6", "", { "dependencies": { "@standard-schema/spec": "^1.1.0", "nanoid": "^5.1.6", "secure-json-parse": "^4.1.0" } }, "sha512-NdtSRrQfWCDA/aqQ1xhobf/xnhuMZkhFAw9xzAt5iAoL3ouxVXOowSRN87OL4MYBQEvqtcjw9/CE6YcsXoBtuw=="],
|
||||
|
||||
"@assistant-ui/react/assistant-stream": ["assistant-stream@0.3.5", "", { "dependencies": { "@standard-schema/spec": "^1.1.0", "nanoid": "^5.1.6", "secure-json-parse": "^4.1.0" } }, "sha512-OGxVClfpEOoSsJDraPoe+GYTwh9TJX1wxK3hT5Qs7gIOyD/MZbqwyWwabRO2KTnNU4w7usvmC/vneUzxSk4bBg=="],
|
||||
"@assistant-ui/react/assistant-stream": ["assistant-stream@0.3.6", "", { "dependencies": { "@standard-schema/spec": "^1.1.0", "nanoid": "^5.1.6", "secure-json-parse": "^4.1.0" } }, "sha512-NdtSRrQfWCDA/aqQ1xhobf/xnhuMZkhFAw9xzAt5iAoL3ouxVXOowSRN87OL4MYBQEvqtcjw9/CE6YcsXoBtuw=="],
|
||||
|
||||
"@assistant-ui/react/zod": ["zod@4.3.6", "", {}, "sha512-rftlrkhHZOcjDwkGlnUtZZkvaPHCsDATp4pGpuOOMDaTdDDXF91wuVDJoWoPsKX/3YPQ5fHuF3STjcYyKr+Qhg=="],
|
||||
|
||||
|
|
@ -2282,7 +2282,7 @@
|
|||
|
||||
"ajv-formats/ajv": ["ajv@8.18.0", "", { "dependencies": { "fast-deep-equal": "^3.1.3", "fast-uri": "^3.0.1", "json-schema-traverse": "^1.0.0", "require-from-string": "^2.0.2" } }, "sha512-PlXPeEWMXMZ7sPYOHqmDyCJzcfNrUr3fGNKtezX14ykXOEIvyK81d+qydx89KY5O71FKMPaQ2vBfBFI5NHR63A=="],
|
||||
|
||||
"assistant-cloud/assistant-stream": ["assistant-stream@0.3.5", "", { "dependencies": { "@standard-schema/spec": "^1.1.0", "nanoid": "^5.1.6", "secure-json-parse": "^4.1.0" } }, "sha512-OGxVClfpEOoSsJDraPoe+GYTwh9TJX1wxK3hT5Qs7gIOyD/MZbqwyWwabRO2KTnNU4w7usvmC/vneUzxSk4bBg=="],
|
||||
"assistant-cloud/assistant-stream": ["assistant-stream@0.3.6", "", { "dependencies": { "@standard-schema/spec": "^1.1.0", "nanoid": "^5.1.6", "secure-json-parse": "^4.1.0" } }, "sha512-NdtSRrQfWCDA/aqQ1xhobf/xnhuMZkhFAw9xzAt5iAoL3ouxVXOowSRN87OL4MYBQEvqtcjw9/CE6YcsXoBtuw=="],
|
||||
|
||||
"chalk/ansi-styles": ["ansi-styles@4.3.0", "", { "dependencies": { "color-convert": "^2.0.1" } }, "sha512-zbB9rCJAT1rbjiVDb2hqKFHNYLxgtk8NURxZ3IZwD3F6NtxbXZQCnnSi1Lkx+IDohdPlFp222wVALIheZJQSEg=="],
|
||||
|
||||
|
|
|
|||
|
|
@ -13,7 +13,7 @@
|
|||
"biome:fix": "biome check . --write"
|
||||
},
|
||||
"dependencies": {
|
||||
"@assistant-ui/react": "^0.12.17",
|
||||
"@assistant-ui/react": "^0.12.19",
|
||||
"@assistant-ui/react-markdown": "^0.12.3",
|
||||
"@assistant-ui/react-streamdown": "^0.1.2",
|
||||
"@base-ui/react": "^1.2.0",
|
||||
|
|
@ -51,7 +51,7 @@
|
|||
"framer-motion": "^11.18.2",
|
||||
"js-yaml": "^4.1.1",
|
||||
"katex": "^0.16.28",
|
||||
"lucide-react": "^0.575.0",
|
||||
"lucide-react": "^0.577.0",
|
||||
"mammoth": "^1.11.0",
|
||||
"motion": "^12.34.0",
|
||||
"next": "^16.1.6",
|
||||
|
|
|
|||
67
studio/frontend/src/components/assistant-ui/badge.tsx
Normal file
67
studio/frontend/src/components/assistant-ui/badge.tsx
Normal file
|
|
@ -0,0 +1,67 @@
|
|||
"use client";
|
||||
|
||||
import type { ComponentProps } from "react";
|
||||
import { Slot } from "radix-ui";
|
||||
import { cva, type VariantProps } from "class-variance-authority";
|
||||
import { cn } from "@/lib/utils";
|
||||
|
||||
const badgeVariants = cva(
|
||||
"inline-flex items-center justify-center gap-1 rounded-md font-medium text-xs transition-colors [&_svg]:size-3 [&_svg]:shrink-0",
|
||||
{
|
||||
variants: {
|
||||
variant: {
|
||||
outline:
|
||||
"border border-input bg-transparent text-muted-foreground hover:bg-accent hover:text-accent-foreground",
|
||||
secondary:
|
||||
"bg-secondary text-secondary-foreground hover:bg-secondary/80",
|
||||
muted:
|
||||
"bg-muted text-muted-foreground hover:bg-muted/80 hover:text-foreground",
|
||||
ghost:
|
||||
"bg-transparent text-muted-foreground hover:bg-accent hover:text-accent-foreground",
|
||||
info: "bg-blue-100 text-blue-700 hover:bg-blue-100/80 dark:bg-blue-900/50 dark:text-blue-300",
|
||||
warning:
|
||||
"bg-amber-100 text-amber-700 hover:bg-amber-100/80 dark:bg-amber-900/50 dark:text-amber-300",
|
||||
success:
|
||||
"bg-emerald-100 text-emerald-700 hover:bg-emerald-100/80 dark:bg-emerald-900/50 dark:text-emerald-300",
|
||||
destructive:
|
||||
"bg-red-100 text-red-700 hover:bg-red-100/80 dark:bg-red-900/50 dark:text-red-300",
|
||||
},
|
||||
size: {
|
||||
sm: "px-1.5 py-0.5",
|
||||
default: "px-2 py-1",
|
||||
lg: "px-2.5 py-1.5 text-sm",
|
||||
},
|
||||
},
|
||||
defaultVariants: {
|
||||
variant: "outline",
|
||||
size: "default",
|
||||
},
|
||||
},
|
||||
);
|
||||
|
||||
export type BadgeProps = ComponentProps<"span"> &
|
||||
VariantProps<typeof badgeVariants> & {
|
||||
asChild?: boolean;
|
||||
};
|
||||
|
||||
function Badge({
|
||||
className,
|
||||
variant,
|
||||
size,
|
||||
asChild = false,
|
||||
...props
|
||||
}: BadgeProps) {
|
||||
const Comp = asChild ? Slot.Root : "span";
|
||||
|
||||
return (
|
||||
<Comp
|
||||
data-slot="badge"
|
||||
data-variant={variant}
|
||||
data-size={size}
|
||||
className={cn(badgeVariants({ variant, size }), className)}
|
||||
{...props}
|
||||
/>
|
||||
);
|
||||
}
|
||||
|
||||
export { Badge, badgeVariants };
|
||||
|
|
@ -10,7 +10,7 @@ import { HugeiconsIcon } from "@hugeicons/react";
|
|||
import { code } from "@streamdown/code";
|
||||
import { math } from "@streamdown/math";
|
||||
import { mermaid } from "@streamdown/mermaid";
|
||||
import { DownloadIcon } from "lucide-react";
|
||||
import { DownloadIcon, Maximize2Icon, Minimize2Icon } from "lucide-react";
|
||||
import { useEffect, useRef, useState } from "react";
|
||||
import { Block, type BlockProps, Streamdown } from "streamdown";
|
||||
import "katex/dist/katex.min.css";
|
||||
|
|
@ -84,6 +84,11 @@ function isSvgFence(codeFence: CodeFence): boolean {
|
|||
return false;
|
||||
}
|
||||
|
||||
function isHtmlFence(codeFence: CodeFence): boolean {
|
||||
const lang = codeFence.language?.toLowerCase() ?? "";
|
||||
return lang === "html" && !codeFence.source.trimStart().startsWith("<svg");
|
||||
}
|
||||
|
||||
const UNSAFE_SVG_RE = /<script[\s>]|on\w+\s*=|javascript:|<foreignObject[\s>]|<iframe[\s>]|<embed[\s>]|<object[\s>]/i;
|
||||
|
||||
function sanitizeSvg(source: string): string | null {
|
||||
|
|
@ -104,6 +109,96 @@ function SvgPreview({ source }: { source: string }) {
|
|||
);
|
||||
}
|
||||
|
||||
const HTML_PREVIEW_DEFAULT_HEIGHT = 400;
|
||||
const HTML_PREVIEW_MAX_HEIGHT = 800;
|
||||
|
||||
function HtmlPreview({ source }: { source: string }) {
|
||||
const iframeRef = useRef<HTMLIFrameElement>(null);
|
||||
const [height, setHeight] = useState(HTML_PREVIEW_DEFAULT_HEIGHT);
|
||||
const [enlarged, setEnlarged] = useState(false);
|
||||
|
||||
useEffect(() => {
|
||||
const handler = (e: MessageEvent) => {
|
||||
if (e.source !== iframeRef.current?.contentWindow) return;
|
||||
if (typeof e.data?.htmlPreviewHeight === "number") {
|
||||
setHeight(Math.min(Math.max(e.data.htmlPreviewHeight, 100), HTML_PREVIEW_MAX_HEIGHT));
|
||||
}
|
||||
};
|
||||
window.addEventListener("message", handler);
|
||||
return () => window.removeEventListener("message", handler);
|
||||
}, []);
|
||||
|
||||
useEffect(() => {
|
||||
if (!enlarged) return;
|
||||
const handler = (e: KeyboardEvent) => {
|
||||
if (e.key === "Escape") setEnlarged(false);
|
||||
};
|
||||
window.addEventListener("keydown", handler);
|
||||
return () => window.removeEventListener("keydown", handler);
|
||||
}, [enlarged]);
|
||||
|
||||
const resizeScript = `<script>new ResizeObserver(()=>{
|
||||
parent.postMessage({htmlPreviewHeight:document.documentElement.scrollHeight},"*");
|
||||
}).observe(document.documentElement);</script>`;
|
||||
|
||||
const srcDoc = source + resizeScript;
|
||||
|
||||
if (enlarged) {
|
||||
return (
|
||||
<>
|
||||
<div className="mt-2 overflow-hidden rounded-lg border border-border" style={{ height }}>
|
||||
{/* Placeholder keeps layout stable while overlay is shown */}
|
||||
</div>
|
||||
<div
|
||||
className="fixed inset-0 z-50 flex flex-col bg-background/80 backdrop-blur-sm"
|
||||
onClick={(e) => { if (e.target === e.currentTarget) setEnlarged(false); }}
|
||||
>
|
||||
<div className="flex items-center justify-end gap-2 px-4 py-2">
|
||||
<button
|
||||
type="button"
|
||||
className="flex items-center gap-1.5 rounded-md border border-border bg-background px-3 py-1.5 text-sm text-muted-foreground transition-colors hover:bg-muted hover:text-foreground"
|
||||
onClick={() => setEnlarged(false)}
|
||||
title="Exit fullscreen (Esc)"
|
||||
>
|
||||
<Minimize2Icon className="size-4" />
|
||||
Exit fullscreen
|
||||
</button>
|
||||
</div>
|
||||
<div className="mx-4 mb-4 flex-1 overflow-hidden rounded-lg border border-border bg-background">
|
||||
<iframe
|
||||
ref={iframeRef}
|
||||
srcDoc={srcDoc}
|
||||
sandbox="allow-scripts"
|
||||
style={{ width: "100%", height: "100%", border: "none", display: "block" }}
|
||||
title="HTML preview"
|
||||
/>
|
||||
</div>
|
||||
</div>
|
||||
</>
|
||||
);
|
||||
}
|
||||
|
||||
return (
|
||||
<div className="group/html-preview relative mt-2 overflow-hidden rounded-lg border border-border">
|
||||
<button
|
||||
type="button"
|
||||
className="absolute top-2 right-2 z-10 rounded-md border border-border bg-background/80 p-1.5 text-muted-foreground opacity-0 transition-all hover:bg-muted hover:text-foreground group-hover/html-preview:opacity-100 supports-[backdrop-filter]:backdrop-blur"
|
||||
onClick={() => setEnlarged(true)}
|
||||
title="Enlarge preview"
|
||||
>
|
||||
<Maximize2Icon className="size-4" />
|
||||
</button>
|
||||
<iframe
|
||||
ref={iframeRef}
|
||||
srcDoc={srcDoc}
|
||||
sandbox="allow-scripts"
|
||||
style={{ width: "100%", height, border: "none", display: "block" }}
|
||||
title="HTML preview"
|
||||
/>
|
||||
</div>
|
||||
);
|
||||
}
|
||||
|
||||
function downloadTextFile(filename: string, text: string): void {
|
||||
const blob = new Blob([text], { type: "text/plain;charset=utf-8" });
|
||||
const url = URL.createObjectURL(blob);
|
||||
|
|
@ -238,6 +333,14 @@ function StreamdownBlock(props: BlockProps) {
|
|||
);
|
||||
}
|
||||
|
||||
if (props.isIncomplete && codeFence && isHtmlFence(codeFence)) {
|
||||
return (
|
||||
<div className="my-4 flex h-48 items-center justify-center rounded-xl border border-border bg-muted/30 text-sm text-muted-foreground animate-pulse">
|
||||
Loading preview...
|
||||
</div>
|
||||
);
|
||||
}
|
||||
|
||||
if (mermaidSource) {
|
||||
return (
|
||||
<div className="relative isolate">
|
||||
|
|
@ -249,6 +352,7 @@ function StreamdownBlock(props: BlockProps) {
|
|||
|
||||
if (codeFence) {
|
||||
const svgSource = !props.isIncomplete && isSvgFence(codeFence) ? sanitizeSvg(codeFence.source) : null;
|
||||
const htmlSource = !props.isIncomplete && isHtmlFence(codeFence) ? codeFence.source : null;
|
||||
return (
|
||||
<>
|
||||
<div className="relative isolate">
|
||||
|
|
@ -260,6 +364,7 @@ function StreamdownBlock(props: BlockProps) {
|
|||
/>
|
||||
</div>
|
||||
{svgSource && <SvgPreview source={svgSource} />}
|
||||
{htmlSource && <HtmlPreview source={htmlSource} />}
|
||||
</>
|
||||
);
|
||||
}
|
||||
|
|
|
|||
|
|
@ -151,7 +151,7 @@ function ModelSelectorContent({
|
|||
<button
|
||||
type="button"
|
||||
onClick={onEject}
|
||||
className="flex w-full items-center justify-center gap-1.5 rounded-md px-2 py-1.5 text-xs text-muted-foreground transition-colors hover:bg-accent hover:text-foreground"
|
||||
className="flex w-full items-center justify-center gap-1.5 rounded-md px-2 py-1.5 text-xs text-destructive transition-colors hover:bg-destructive/10"
|
||||
title="Eject model"
|
||||
>
|
||||
<HugeiconsIcon icon={Logout01Icon} className="size-3.5" />
|
||||
|
|
|
|||
|
|
@ -35,7 +35,7 @@ import { checkVramFit, estimateLoadingVram } from "@/lib/vram";
|
|||
import { Search01Icon } from "@hugeicons/core-free-icons";
|
||||
import { HugeiconsIcon } from "@hugeicons/react";
|
||||
import { Trash2Icon } from "lucide-react";
|
||||
import { useCallback, useEffect, useMemo, useState, type ReactNode } from "react";
|
||||
import { useCallback, useEffect, useMemo, useRef, useState, type ReactNode } from "react";
|
||||
import { toast } from "sonner";
|
||||
import type {
|
||||
LoraModelOption,
|
||||
|
|
@ -455,21 +455,43 @@ export function HubModelPicker({
|
|||
const all = dedupe([...models.map((model) => model.id), value ?? ""])
|
||||
.filter((id) => !downloadedSet.has(id.toLowerCase()))
|
||||
.filter((id) => !chatOnly || isGgufRepo(id));
|
||||
// Cap at 4 GGUFs + 4 non-GGUFs so the list stays manageable
|
||||
// Sort: GGUFs first, then hub models
|
||||
const gguf: string[] = [];
|
||||
const hub: string[] = [];
|
||||
for (const id of all) {
|
||||
if (isGgufRepo(id) && gguf.length < 4) gguf.push(id);
|
||||
else if (!isGgufRepo(id) && hub.length < 4) hub.push(id);
|
||||
if (isGgufRepo(id)) gguf.push(id);
|
||||
else hub.push(id);
|
||||
}
|
||||
return [...gguf, ...hub];
|
||||
}, [models, value, downloadedSet, chatOnly]);
|
||||
|
||||
// Infinite scroll paging for the recommended section
|
||||
const [recommendedPage, setRecommendedPage] = useState(1);
|
||||
// Reset page when the underlying list changes
|
||||
useEffect(() => { setRecommendedPage(1); }, [models, chatOnly]);
|
||||
|
||||
const visibleRecommendedIds = useMemo(() => {
|
||||
const hubStartIndex = recommendedIds.findIndex((id) => !isGgufRepo(id));
|
||||
const allGguf = hubStartIndex === -1 ? recommendedIds : recommendedIds.slice(0, hubStartIndex);
|
||||
const allHub = hubStartIndex === -1 ? [] : recommendedIds.slice(hubStartIndex);
|
||||
// Interleave in chunks of 4: [4 gguf, 4 hub, 4 gguf, 4 hub, ...]
|
||||
const result: string[] = [];
|
||||
for (let p = 0; p < recommendedPage; p++) {
|
||||
result.push(...allGguf.slice(p * 4, (p + 1) * 4));
|
||||
result.push(...allHub.slice(p * 4, (p + 1) * 4));
|
||||
}
|
||||
return result;
|
||||
}, [recommendedIds, recommendedPage]);
|
||||
|
||||
const hasMoreRecommended = visibleRecommendedIds.length < recommendedIds.length;
|
||||
|
||||
// Fetch VRAM info for the full pool once (recommendedIds is stable across
|
||||
// page increments) so we don't re-fetch on every scroll.
|
||||
const { paramCountById: recommendedParamCountById } =
|
||||
useRecommendedModelVram(recommendedIds);
|
||||
|
||||
const showHfSection = debouncedQuery.trim().length > 0;
|
||||
const recommendedSet = useMemo(() => new Set(recommendedIds), [recommendedIds]);
|
||||
const recommendedSet = useMemo(() => new Set(visibleRecommendedIds), [visibleRecommendedIds]);
|
||||
|
||||
const hfIds = useMemo(() => {
|
||||
if (!showHfSection) return [];
|
||||
|
|
@ -519,7 +541,7 @@ export function HubModelPicker({
|
|||
string,
|
||||
{ est: number; status: VramFitStatus | null; detail: string | null }
|
||||
>();
|
||||
for (const id of recommendedIds) {
|
||||
for (const id of visibleRecommendedIds) {
|
||||
const totalParams = recommendedParamCountById.get(id);
|
||||
if (totalParams) {
|
||||
const est = estimateLoadingVram(totalParams, "qlora");
|
||||
|
|
@ -531,10 +553,36 @@ export function HubModelPicker({
|
|||
}
|
||||
}
|
||||
return map;
|
||||
}, [recommendedIds, recommendedParamCountById, gpu]);
|
||||
}, [visibleRecommendedIds, recommendedParamCountById, gpu]);
|
||||
|
||||
const { scrollRef, sentinelRef } = useInfiniteScroll(fetchMore, results.length);
|
||||
|
||||
// Sentinel + IntersectionObserver for recommended infinite scroll.
|
||||
// We disconnect after each fire so the observer doesn't loop while
|
||||
// React re-renders; the effect re-creates it on the next page.
|
||||
// Uses a callback ref for the sentinel so we detect mount/unmount reliably.
|
||||
const [recommendedSentinel, setRecommendedSentinel] = useState<HTMLDivElement | null>(null);
|
||||
const recommendedSentinelRef = useCallback((node: HTMLDivElement | null) => {
|
||||
setRecommendedSentinel(node);
|
||||
}, []);
|
||||
useEffect(() => {
|
||||
if (!recommendedSentinel || !hasMoreRecommended) return;
|
||||
const root = scrollRef.current;
|
||||
if (!root) return;
|
||||
const obs = new IntersectionObserver(
|
||||
([e]) => {
|
||||
if (e.isIntersecting) {
|
||||
obs.disconnect();
|
||||
setRecommendedPage((p) => p + 1);
|
||||
}
|
||||
},
|
||||
{ threshold: 0, root },
|
||||
);
|
||||
// Small delay so the browser finishes layout after the previous page render
|
||||
const timer = setTimeout(() => obs.observe(recommendedSentinel), 100);
|
||||
return () => { clearTimeout(timer); obs.disconnect(); };
|
||||
}, [recommendedSentinel, hasMoreRecommended, recommendedPage, scrollRef]);
|
||||
|
||||
/** Handle clicking a model row — GGUF repos expand, others load directly. */
|
||||
const handleModelClick = useCallback(
|
||||
(id: string) => {
|
||||
|
|
@ -622,12 +670,12 @@ export function HubModelPicker({
|
|||
{!showHfSection && cachedReady ? (
|
||||
<>
|
||||
<ListLabel>{"\uD83E\uDDA5"} Recommended</ListLabel>
|
||||
{recommendedIds.length === 0 ? (
|
||||
{visibleRecommendedIds.length === 0 ? (
|
||||
<div className="px-2.5 py-2 text-xs text-muted-foreground">
|
||||
No default models.
|
||||
</div>
|
||||
) : (
|
||||
recommendedIds.map((id) => {
|
||||
visibleRecommendedIds.map((id) => {
|
||||
const vram = recommendedVramMap.get(id);
|
||||
return (
|
||||
<div key={id}>
|
||||
|
|
@ -651,6 +699,14 @@ export function HubModelPicker({
|
|||
);
|
||||
})
|
||||
)}
|
||||
{hasMoreRecommended && (
|
||||
<>
|
||||
<div ref={recommendedSentinelRef} className="h-px" />
|
||||
<div className="flex items-center justify-center py-2">
|
||||
<Spinner className="size-3.5 text-muted-foreground" />
|
||||
</div>
|
||||
</>
|
||||
)}
|
||||
</>
|
||||
) : null}
|
||||
|
||||
|
|
|
|||
|
|
@ -19,10 +19,11 @@ import {
|
|||
useAuiState,
|
||||
useScrollLock,
|
||||
} from "@assistant-ui/react";
|
||||
import { copyToClipboard } from "@/lib/copy-to-clipboard";
|
||||
import { Idea01Icon } from "@hugeicons/core-free-icons";
|
||||
import { HugeiconsIcon } from "@hugeicons/react";
|
||||
import { type VariantProps, cva } from "class-variance-authority";
|
||||
import { ChevronDownIcon } from "lucide-react";
|
||||
import { ChevronDownIcon, CopyIcon, CheckIcon } from "lucide-react";
|
||||
import {
|
||||
type CSSProperties,
|
||||
type ComponentProps,
|
||||
|
|
@ -263,6 +264,45 @@ function ReasoningText({
|
|||
|
||||
const ReasoningImpl: ReasoningMessagePartComponent = () => <MarkdownText />;
|
||||
|
||||
const COPY_RESET_MS = 2000;
|
||||
|
||||
function ReasoningCopyButton({ startIndex, endIndex }: { startIndex: number; endIndex: number }) {
|
||||
const [copied, setCopied] = useState(false);
|
||||
const resetRef = useRef<ReturnType<typeof setTimeout> | null>(null);
|
||||
|
||||
const reasoningText = useAuiState(({ message }) => {
|
||||
return message.parts
|
||||
.slice(startIndex, endIndex + 1)
|
||||
.filter((p) => p.type === "reasoning")
|
||||
.map((p) => ("text" in p ? (p as { text: string }).text : ""))
|
||||
.join("\n");
|
||||
});
|
||||
|
||||
const handleCopy = useCallback(() => {
|
||||
if (copyToClipboard(reasoningText)) {
|
||||
setCopied(true);
|
||||
if (resetRef.current) clearTimeout(resetRef.current);
|
||||
resetRef.current = setTimeout(() => setCopied(false), COPY_RESET_MS);
|
||||
}
|
||||
}, [reasoningText]);
|
||||
|
||||
return (
|
||||
<button
|
||||
type="button"
|
||||
onClick={handleCopy}
|
||||
className="inline-flex items-center gap-1 rounded px-1.5 py-0.5 text-xs text-muted-foreground transition-colors hover:text-foreground hover:bg-muted"
|
||||
aria-label="Copy reasoning"
|
||||
>
|
||||
{copied ? (
|
||||
<CheckIcon className="size-3" />
|
||||
) : (
|
||||
<CopyIcon className="size-3" />
|
||||
)}
|
||||
{copied ? "Copied" : "Copy"}
|
||||
</button>
|
||||
);
|
||||
}
|
||||
|
||||
const ReasoningGroupImpl: ReasoningGroupComponent = ({
|
||||
children,
|
||||
startIndex,
|
||||
|
|
@ -328,10 +368,15 @@ const ReasoningGroupImpl: ReasoningGroupComponent = ({
|
|||
onOpenChange={handleOpenChange}
|
||||
variant={variant}
|
||||
>
|
||||
<ReasoningTrigger
|
||||
active={isReasoningStreaming}
|
||||
duration={duration || persistedDuration}
|
||||
/>
|
||||
<div className="flex items-center justify-between">
|
||||
<ReasoningTrigger
|
||||
active={isReasoningStreaming}
|
||||
duration={duration || persistedDuration}
|
||||
/>
|
||||
{isOpen && !isReasoningStreaming && (
|
||||
<ReasoningCopyButton startIndex={startIndex} endIndex={endIndex} />
|
||||
)}
|
||||
</div>
|
||||
<ReasoningContent
|
||||
aria-busy={isReasoningStreaming}
|
||||
streaming={isReasoningStreaming}
|
||||
|
|
|
|||
137
studio/frontend/src/components/assistant-ui/sources.tsx
Normal file
137
studio/frontend/src/components/assistant-ui/sources.tsx
Normal file
|
|
@ -0,0 +1,137 @@
|
|||
"use client";
|
||||
|
||||
import { memo, useState, type ComponentProps } from "react";
|
||||
import type { SourceMessagePartComponent } from "@assistant-ui/react";
|
||||
import { cn } from "@/lib/utils";
|
||||
import { Badge, badgeVariants, type BadgeProps } from "./badge";
|
||||
|
||||
const extractDomain = (url: string): string => {
|
||||
try {
|
||||
return new URL(url).hostname.replace(/^www\./, "");
|
||||
} catch {
|
||||
return url;
|
||||
}
|
||||
};
|
||||
|
||||
const getDomainInitial = (url: string): string => {
|
||||
const domain = extractDomain(url);
|
||||
return domain.charAt(0).toUpperCase();
|
||||
};
|
||||
|
||||
function SourceIcon({
|
||||
url,
|
||||
className,
|
||||
...props
|
||||
}: ComponentProps<"span"> & { url: string }) {
|
||||
const [hasError, setHasError] = useState(false);
|
||||
const domain = extractDomain(url);
|
||||
|
||||
if (hasError) {
|
||||
return (
|
||||
<span
|
||||
data-slot="source-icon-fallback"
|
||||
className={cn(
|
||||
"flex size-3 shrink-0 items-center justify-center rounded-sm bg-muted font-medium text-[10px]",
|
||||
className,
|
||||
)}
|
||||
{...props}
|
||||
>
|
||||
{getDomainInitial(url)}
|
||||
</span>
|
||||
);
|
||||
}
|
||||
|
||||
return (
|
||||
<img
|
||||
data-slot="source-icon"
|
||||
src={`https://www.google.com/s2/favicons?domain=${domain}&sz=32`}
|
||||
alt=""
|
||||
className={cn("size-3 shrink-0 rounded-sm", className)}
|
||||
onError={() => setHasError(true)}
|
||||
{...(props as ComponentProps<"img">)}
|
||||
/>
|
||||
);
|
||||
}
|
||||
|
||||
function SourceTitle({ className, ...props }: ComponentProps<"span">) {
|
||||
return (
|
||||
<span
|
||||
data-slot="source-title"
|
||||
className={cn("max-w-37.5 truncate", className)}
|
||||
{...props}
|
||||
/>
|
||||
);
|
||||
}
|
||||
|
||||
export type SourceProps = Omit<BadgeProps, "asChild"> &
|
||||
ComponentProps<"a"> & {
|
||||
asChild?: boolean;
|
||||
};
|
||||
|
||||
function Source({
|
||||
className,
|
||||
variant,
|
||||
size,
|
||||
asChild = false,
|
||||
target = "_blank",
|
||||
rel = "noopener noreferrer",
|
||||
...props
|
||||
}: SourceProps) {
|
||||
return (
|
||||
<Badge
|
||||
asChild
|
||||
variant={variant}
|
||||
size={size}
|
||||
className={cn(
|
||||
"cursor-pointer outline-none focus-visible:border-ring focus-visible:ring-[3px] focus-visible:ring-ring/50",
|
||||
className,
|
||||
)}
|
||||
>
|
||||
<a
|
||||
data-slot="source"
|
||||
target={target}
|
||||
rel={rel}
|
||||
{...(props as ComponentProps<"a">)}
|
||||
/>
|
||||
</Badge>
|
||||
);
|
||||
}
|
||||
|
||||
const SourcesImpl: SourceMessagePartComponent = ({
|
||||
url,
|
||||
title,
|
||||
sourceType,
|
||||
}) => {
|
||||
if (sourceType !== "url" || !url) return null;
|
||||
|
||||
const domain = extractDomain(url);
|
||||
const displayTitle = title || domain;
|
||||
|
||||
return (
|
||||
<span className="mr-1 mt-1 inline-block first:mt-2">
|
||||
<Source href={url}>
|
||||
<SourceIcon url={url} />
|
||||
<SourceTitle>{displayTitle}</SourceTitle>
|
||||
</Source>
|
||||
</span>
|
||||
);
|
||||
};
|
||||
|
||||
const Sources = memo(SourcesImpl) as unknown as SourceMessagePartComponent & {
|
||||
Root: typeof Source;
|
||||
Icon: typeof SourceIcon;
|
||||
Title: typeof SourceTitle;
|
||||
};
|
||||
|
||||
Sources.displayName = "Sources";
|
||||
Sources.Root = Source;
|
||||
Sources.Icon = SourceIcon;
|
||||
Sources.Title = SourceTitle;
|
||||
|
||||
export {
|
||||
Sources,
|
||||
Source,
|
||||
SourceIcon,
|
||||
SourceTitle,
|
||||
badgeVariants as sourceVariants,
|
||||
};
|
||||
|
|
@ -9,7 +9,12 @@ import {
|
|||
import { MessageTiming } from "@/components/assistant-ui/message-timing";
|
||||
import { MarkdownText } from "@/components/assistant-ui/markdown-text";
|
||||
import { Reasoning, ReasoningGroup } from "@/components/assistant-ui/reasoning";
|
||||
import { Sources } from "@/components/assistant-ui/sources";
|
||||
import { ToolFallback } from "@/components/assistant-ui/tool-fallback";
|
||||
import { ToolGroup } from "@/components/assistant-ui/tool-group";
|
||||
import { WebSearchToolUI } from "@/components/assistant-ui/tool-ui-web-search";
|
||||
import { PythonToolUI } from "@/components/assistant-ui/tool-ui-python";
|
||||
import { TerminalToolUI } from "@/components/assistant-ui/tool-ui-terminal";
|
||||
import { TooltipIconButton } from "@/components/assistant-ui/tooltip-icon-button";
|
||||
import { Button } from "@/components/ui/button";
|
||||
import { sentAudioNames } from "@/features/chat/api/chat-adapter";
|
||||
|
|
@ -52,7 +57,7 @@ import {
|
|||
TerminalIcon,
|
||||
XIcon,
|
||||
} from "lucide-react";
|
||||
import { type FC, useCallback, useRef, useState } from "react";
|
||||
import { type FC, useCallback, useEffect, useRef, useState } from "react";
|
||||
import { useChatRuntimeStore } from "@/features/chat/stores/chat-runtime-store";
|
||||
|
||||
export const Thread: FC<{ hideComposer?: boolean; hideWelcome?: boolean }> = ({
|
||||
|
|
@ -384,6 +389,20 @@ const CodeToolsToggle: FC = () => {
|
|||
|
||||
const ToolStatusDisplay: FC = () => {
|
||||
const toolStatus = useChatRuntimeStore((s) => s.toolStatus);
|
||||
const [elapsed, setElapsed] = useState(0);
|
||||
|
||||
useEffect(() => {
|
||||
if (!toolStatus) {
|
||||
setElapsed(0);
|
||||
return;
|
||||
}
|
||||
setElapsed(0);
|
||||
const interval = setInterval(() => {
|
||||
setElapsed((prev) => prev + 1);
|
||||
}, 1000);
|
||||
return () => clearInterval(interval);
|
||||
}, [toolStatus]);
|
||||
|
||||
if (!toolStatus) return null;
|
||||
const isRunning = toolStatus.startsWith("Running");
|
||||
const StatusIcon = isRunning ? TerminalIcon : GlobeIcon;
|
||||
|
|
@ -392,6 +411,7 @@ const ToolStatusDisplay: FC = () => {
|
|||
<div className="flex animate-pulse items-center gap-2 rounded-full border border-primary/20 bg-primary/5 px-3 py-1.5 text-xs text-primary">
|
||||
<StatusIcon className="size-3.5" />
|
||||
<span>{toolStatus}</span>
|
||||
<span className="tabular-nums opacity-60">{elapsed}s</span>
|
||||
</div>
|
||||
</div>
|
||||
);
|
||||
|
|
@ -485,7 +505,16 @@ const AssistantMessage: FC = () => {
|
|||
Text: MarkdownText,
|
||||
Reasoning: Reasoning,
|
||||
ReasoningGroup: ReasoningGroup,
|
||||
tools: { Fallback: ToolFallback },
|
||||
Source: Sources,
|
||||
ToolGroup: ToolGroup,
|
||||
tools: {
|
||||
by_name: {
|
||||
web_search: WebSearchToolUI,
|
||||
python: PythonToolUI,
|
||||
terminal: TerminalToolUI,
|
||||
},
|
||||
Fallback: ToolFallback,
|
||||
},
|
||||
}}
|
||||
/>
|
||||
<MessageError />
|
||||
|
|
|
|||
|
|
@ -77,7 +77,7 @@ function ToolFallbackRoot({
|
|||
open={isOpen}
|
||||
onOpenChange={handleOpenChange}
|
||||
className={cn(
|
||||
"aui-tool-fallback-root group/tool-fallback-root w-full rounded-lg border py-3",
|
||||
"aui-tool-fallback-root group/tool-fallback-root w-full corner-squircle rounded-lg border py-3",
|
||||
className,
|
||||
)}
|
||||
style={
|
||||
|
|
@ -104,18 +104,20 @@ const statusIconMap: Record<ToolStatus, ElementType> = {
|
|||
function ToolFallbackTrigger({
|
||||
toolName,
|
||||
status,
|
||||
icon: ToolIcon,
|
||||
className,
|
||||
...props
|
||||
}: ComponentProps<typeof CollapsibleTrigger> & {
|
||||
toolName: string;
|
||||
status?: ToolCallMessagePartStatus;
|
||||
icon?: ElementType;
|
||||
}) {
|
||||
const statusType = status?.type ?? "complete";
|
||||
const isRunning = statusType === "running";
|
||||
const isCancelled =
|
||||
status?.type === "incomplete" && status.reason === "cancelled";
|
||||
|
||||
const Icon = statusIconMap[statusType];
|
||||
const StatusIcon = statusIconMap[statusType];
|
||||
const label = isCancelled ? "Cancelled tool" : "Used tool";
|
||||
|
||||
return (
|
||||
|
|
@ -127,14 +129,30 @@ function ToolFallbackTrigger({
|
|||
)}
|
||||
{...props}
|
||||
>
|
||||
<Icon
|
||||
data-slot="tool-fallback-trigger-icon"
|
||||
className={cn(
|
||||
"aui-tool-fallback-trigger-icon size-4 shrink-0",
|
||||
isCancelled && "text-muted-foreground",
|
||||
isRunning && "animate-spin",
|
||||
)}
|
||||
/>
|
||||
{isRunning ? (
|
||||
<StatusIcon
|
||||
data-slot="tool-fallback-trigger-icon"
|
||||
className="aui-tool-fallback-trigger-icon size-4 shrink-0 animate-spin"
|
||||
/>
|
||||
) : (
|
||||
ToolIcon ? (
|
||||
<ToolIcon
|
||||
data-slot="tool-fallback-trigger-icon"
|
||||
className={cn(
|
||||
"aui-tool-fallback-trigger-icon size-4 shrink-0",
|
||||
isCancelled && "text-muted-foreground",
|
||||
)}
|
||||
/>
|
||||
) : (
|
||||
<StatusIcon
|
||||
data-slot="tool-fallback-trigger-icon"
|
||||
className={cn(
|
||||
"aui-tool-fallback-trigger-icon size-4 shrink-0",
|
||||
isCancelled && "text-muted-foreground",
|
||||
)}
|
||||
/>
|
||||
)
|
||||
)}
|
||||
<span
|
||||
data-slot="tool-fallback-trigger-label"
|
||||
className={cn(
|
||||
|
|
|
|||
230
studio/frontend/src/components/assistant-ui/tool-group.tsx
Normal file
230
studio/frontend/src/components/assistant-ui/tool-group.tsx
Normal file
|
|
@ -0,0 +1,230 @@
|
|||
"use client";
|
||||
|
||||
import {
|
||||
memo,
|
||||
useCallback,
|
||||
useRef,
|
||||
useState,
|
||||
type FC,
|
||||
type PropsWithChildren,
|
||||
} from "react";
|
||||
import { ChevronDownIcon, LoaderIcon } from "lucide-react";
|
||||
import { cva, type VariantProps } from "class-variance-authority";
|
||||
import { useScrollLock } from "@assistant-ui/react";
|
||||
import {
|
||||
Collapsible,
|
||||
CollapsibleContent,
|
||||
CollapsibleTrigger,
|
||||
} from "@/components/ui/collapsible";
|
||||
import { cn } from "@/lib/utils";
|
||||
|
||||
const ANIMATION_DURATION = 200;
|
||||
|
||||
const toolGroupVariants = cva("aui-tool-group-root group/tool-group w-full", {
|
||||
variants: {
|
||||
variant: {
|
||||
outline: "corner-squircle rounded-lg border py-3",
|
||||
ghost: "",
|
||||
muted: "corner-squircle rounded-lg border border-muted-foreground/30 bg-muted/30 py-3",
|
||||
},
|
||||
},
|
||||
defaultVariants: { variant: "outline" },
|
||||
});
|
||||
|
||||
export type ToolGroupRootProps = Omit<
|
||||
React.ComponentProps<typeof Collapsible>,
|
||||
"open" | "onOpenChange"
|
||||
> &
|
||||
VariantProps<typeof toolGroupVariants> & {
|
||||
open?: boolean;
|
||||
onOpenChange?: (open: boolean) => void;
|
||||
defaultOpen?: boolean;
|
||||
};
|
||||
|
||||
function ToolGroupRoot({
|
||||
className,
|
||||
variant,
|
||||
open: controlledOpen,
|
||||
onOpenChange: controlledOnOpenChange,
|
||||
defaultOpen = false,
|
||||
children,
|
||||
...props
|
||||
}: ToolGroupRootProps) {
|
||||
const collapsibleRef = useRef<HTMLDivElement>(null);
|
||||
const [uncontrolledOpen, setUncontrolledOpen] = useState(defaultOpen);
|
||||
const lockScroll = useScrollLock(collapsibleRef, ANIMATION_DURATION);
|
||||
|
||||
const isControlled = controlledOpen !== undefined;
|
||||
const isOpen = isControlled ? controlledOpen : uncontrolledOpen;
|
||||
|
||||
const handleOpenChange = useCallback(
|
||||
(open: boolean) => {
|
||||
if (!open) {
|
||||
lockScroll();
|
||||
}
|
||||
if (!isControlled) {
|
||||
setUncontrolledOpen(open);
|
||||
}
|
||||
controlledOnOpenChange?.(open);
|
||||
},
|
||||
[lockScroll, isControlled, controlledOnOpenChange],
|
||||
);
|
||||
|
||||
return (
|
||||
<Collapsible
|
||||
ref={collapsibleRef}
|
||||
data-slot="tool-group-root"
|
||||
data-variant={variant ?? "outline"}
|
||||
open={isOpen}
|
||||
onOpenChange={handleOpenChange}
|
||||
className={cn(
|
||||
toolGroupVariants({ variant }),
|
||||
"group/tool-group-root",
|
||||
className,
|
||||
)}
|
||||
style={
|
||||
{
|
||||
"--animation-duration": `${ANIMATION_DURATION}ms`,
|
||||
} as React.CSSProperties
|
||||
}
|
||||
{...props}
|
||||
>
|
||||
{children}
|
||||
</Collapsible>
|
||||
);
|
||||
}
|
||||
|
||||
function ToolGroupTrigger({
|
||||
count,
|
||||
active = false,
|
||||
className,
|
||||
...props
|
||||
}: React.ComponentProps<typeof CollapsibleTrigger> & {
|
||||
count: number;
|
||||
active?: boolean;
|
||||
}) {
|
||||
const label = `${count} tool ${count === 1 ? "call" : "calls"}`;
|
||||
|
||||
return (
|
||||
<CollapsibleTrigger
|
||||
data-slot="tool-group-trigger"
|
||||
className={cn(
|
||||
"aui-tool-group-trigger group/trigger flex items-center gap-2 text-sm transition-colors",
|
||||
"group-data-[variant=outline]/tool-group-root:w-full group-data-[variant=outline]/tool-group-root:px-4",
|
||||
"group-data-[variant=muted]/tool-group-root:w-full group-data-[variant=muted]/tool-group-root:px-4",
|
||||
className,
|
||||
)}
|
||||
{...props}
|
||||
>
|
||||
{active && (
|
||||
<LoaderIcon
|
||||
data-slot="tool-group-trigger-loader"
|
||||
className="aui-tool-group-trigger-loader size-4 shrink-0 animate-spin"
|
||||
/>
|
||||
)}
|
||||
<span
|
||||
data-slot="tool-group-trigger-label"
|
||||
className={cn(
|
||||
"aui-tool-group-trigger-label-wrapper relative inline-block text-left font-medium leading-none",
|
||||
"group-data-[variant=outline]/tool-group-root:grow",
|
||||
"group-data-[variant=muted]/tool-group-root:grow",
|
||||
)}
|
||||
>
|
||||
<span>{label}</span>
|
||||
{active && (
|
||||
<span
|
||||
aria-hidden
|
||||
data-slot="tool-group-trigger-shimmer"
|
||||
className="aui-tool-group-trigger-shimmer shimmer pointer-events-none absolute inset-0 motion-reduce:animate-none"
|
||||
>
|
||||
{label}
|
||||
</span>
|
||||
)}
|
||||
</span>
|
||||
<ChevronDownIcon
|
||||
data-slot="tool-group-trigger-chevron"
|
||||
className={cn(
|
||||
"aui-tool-group-trigger-chevron size-4 shrink-0",
|
||||
"transition-transform duration-(--animation-duration) ease-out",
|
||||
"group-data-[state=closed]/trigger:-rotate-90",
|
||||
"group-data-[state=open]/trigger:rotate-0",
|
||||
)}
|
||||
/>
|
||||
</CollapsibleTrigger>
|
||||
);
|
||||
}
|
||||
|
||||
function ToolGroupContent({
|
||||
className,
|
||||
children,
|
||||
...props
|
||||
}: React.ComponentProps<typeof CollapsibleContent>) {
|
||||
return (
|
||||
<CollapsibleContent
|
||||
data-slot="tool-group-content"
|
||||
className={cn(
|
||||
"aui-tool-group-content relative overflow-hidden text-sm outline-none",
|
||||
"group/collapsible-content ease-out",
|
||||
"data-[state=closed]:animate-collapsible-up",
|
||||
"data-[state=open]:animate-collapsible-down",
|
||||
"data-[state=closed]:fill-mode-forwards",
|
||||
"data-[state=closed]:pointer-events-none",
|
||||
"data-[state=open]:duration-(--animation-duration)",
|
||||
"data-[state=closed]:duration-(--animation-duration)",
|
||||
className,
|
||||
)}
|
||||
{...props}
|
||||
>
|
||||
<div
|
||||
className={cn(
|
||||
"mt-2 flex flex-col gap-2",
|
||||
"group-data-[variant=outline]/tool-group-root:mt-3 group-data-[variant=outline]/tool-group-root:border-t group-data-[variant=outline]/tool-group-root:px-4 group-data-[variant=outline]/tool-group-root:pt-3",
|
||||
"group-data-[variant=muted]/tool-group-root:mt-3 group-data-[variant=muted]/tool-group-root:border-t group-data-[variant=muted]/tool-group-root:px-4 group-data-[variant=muted]/tool-group-root:pt-3",
|
||||
)}
|
||||
>
|
||||
{children}
|
||||
</div>
|
||||
</CollapsibleContent>
|
||||
);
|
||||
}
|
||||
|
||||
type ToolGroupComponent = FC<
|
||||
PropsWithChildren<{ startIndex: number; endIndex: number }>
|
||||
> & {
|
||||
Root: typeof ToolGroupRoot;
|
||||
Trigger: typeof ToolGroupTrigger;
|
||||
Content: typeof ToolGroupContent;
|
||||
};
|
||||
|
||||
const ToolGroupImpl: FC<
|
||||
PropsWithChildren<{ startIndex: number; endIndex: number }>
|
||||
> = ({ children, startIndex, endIndex }) => {
|
||||
const toolCount = endIndex - startIndex + 1;
|
||||
|
||||
// Single tool call — render directly without wrapper
|
||||
if (toolCount <= 1) {
|
||||
return <>{children}</>;
|
||||
}
|
||||
|
||||
return (
|
||||
<ToolGroupRoot>
|
||||
<ToolGroupTrigger count={toolCount} />
|
||||
<ToolGroupContent>{children}</ToolGroupContent>
|
||||
</ToolGroupRoot>
|
||||
);
|
||||
};
|
||||
|
||||
const ToolGroup = memo(ToolGroupImpl) as unknown as ToolGroupComponent;
|
||||
|
||||
ToolGroup.displayName = "ToolGroup";
|
||||
ToolGroup.Root = ToolGroupRoot;
|
||||
ToolGroup.Trigger = ToolGroupTrigger;
|
||||
ToolGroup.Content = ToolGroupContent;
|
||||
|
||||
export {
|
||||
ToolGroup,
|
||||
ToolGroupRoot,
|
||||
ToolGroupTrigger,
|
||||
ToolGroupContent,
|
||||
toolGroupVariants,
|
||||
};
|
||||
136
studio/frontend/src/components/assistant-ui/tool-ui-python.tsx
Normal file
136
studio/frontend/src/components/assistant-ui/tool-ui-python.tsx
Normal file
|
|
@ -0,0 +1,136 @@
|
|||
// SPDX-License-Identifier: AGPL-3.0-only
|
||||
// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
|
||||
|
||||
"use client";
|
||||
|
||||
import { copyToClipboard } from "@/lib/copy-to-clipboard";
|
||||
import type { ToolCallMessagePartComponent } from "@assistant-ui/react";
|
||||
import { code as codePlugin } from "@streamdown/code";
|
||||
import { CheckIcon, CodeIcon, CopyIcon, LoaderIcon } from "lucide-react";
|
||||
import { memo, useCallback, useMemo, useRef, useState } from "react";
|
||||
import { Streamdown } from "streamdown";
|
||||
import {
|
||||
ToolFallbackContent,
|
||||
ToolFallbackRoot,
|
||||
ToolFallbackTrigger,
|
||||
} from "./tool-fallback";
|
||||
|
||||
const MAX_DISPLAY = 10_000;
|
||||
const COPY_RESET_MS = 2000;
|
||||
const SHIKI_THEME = ["github-light", "github-dark"] as const;
|
||||
|
||||
function truncate(text: string): string {
|
||||
return text.length <= MAX_DISPLAY
|
||||
? text
|
||||
: `${text.slice(0, MAX_DISPLAY)}\n... (truncated)`;
|
||||
}
|
||||
|
||||
function CopyBtn({ text }: { text: string }) {
|
||||
const [copied, setCopied] = useState(false);
|
||||
const timer = useRef<ReturnType<typeof setTimeout> | null>(null);
|
||||
const copy = useCallback(() => {
|
||||
if (copyToClipboard(text)) {
|
||||
setCopied(true);
|
||||
if (timer.current) {
|
||||
clearTimeout(timer.current);
|
||||
}
|
||||
timer.current = setTimeout(() => setCopied(false), COPY_RESET_MS);
|
||||
}
|
||||
}, [text]);
|
||||
|
||||
return (
|
||||
<button
|
||||
type="button"
|
||||
onClick={copy}
|
||||
className="inline-flex items-center gap-1 rounded px-1.5 py-0.5 text-xs text-muted-foreground transition-colors hover:bg-muted hover:text-foreground"
|
||||
aria-label="Copy to clipboard"
|
||||
>
|
||||
{copied ? (
|
||||
<CheckIcon className="size-3" />
|
||||
) : (
|
||||
<CopyIcon className="size-3" />
|
||||
)}
|
||||
{copied ? "Copied" : "Copy"}
|
||||
</button>
|
||||
);
|
||||
}
|
||||
|
||||
/** Render code with syntax highlighting via Streamdown + shiki. No extra borders — inherits parent container. */
|
||||
function HighlightedCode({ code: source, language }: { code: string; language: string }) {
|
||||
const markdown = useMemo(
|
||||
() => `\`\`\`${language}\n${truncate(source)}\n\`\`\``,
|
||||
[source, language],
|
||||
);
|
||||
return (
|
||||
<div className="max-h-48 overflow-auto text-xs [&_pre]:!m-0 [&_pre]:!bg-transparent [&_pre]:!p-0 [&_pre]:!text-xs [&_[data-streamdown=code-block]]:!my-0 [&_[data-streamdown=code-block]]:!p-0 [&_[data-streamdown=code-block]]:!border-0">
|
||||
<Streamdown
|
||||
mode="static"
|
||||
plugins={{ code: codePlugin }}
|
||||
controls={{ code: false }}
|
||||
shikiTheme={SHIKI_THEME}
|
||||
>
|
||||
{markdown}
|
||||
</Streamdown>
|
||||
</div>
|
||||
);
|
||||
}
|
||||
|
||||
const PythonToolUIImpl: ToolCallMessagePartComponent = ({
|
||||
args,
|
||||
result,
|
||||
status,
|
||||
}) => {
|
||||
const code = (args as { code?: string })?.code ?? "";
|
||||
const firstLine = code.split("\n")[0]?.slice(0, 60) ?? "";
|
||||
const isRunning = status?.type === "running";
|
||||
const output =
|
||||
typeof result === "string"
|
||||
? result
|
||||
: result
|
||||
? JSON.stringify(result, null, 2)
|
||||
: "";
|
||||
|
||||
return (
|
||||
<ToolFallbackRoot>
|
||||
<ToolFallbackTrigger
|
||||
toolName={firstLine ? `Python: ${firstLine}` : "Python"}
|
||||
status={status}
|
||||
icon={CodeIcon}
|
||||
/>
|
||||
<ToolFallbackContent>
|
||||
<div className="flex flex-col px-4">
|
||||
{/* Code + copy */}
|
||||
{code && (
|
||||
<div className="flex justify-end">
|
||||
<CopyBtn text={code} />
|
||||
</div>
|
||||
)}
|
||||
<HighlightedCode code={code} language="python" />
|
||||
|
||||
{/* Output */}
|
||||
{isRunning ? (
|
||||
<div className="mt-2 flex items-center gap-2 text-sm text-muted-foreground">
|
||||
<LoaderIcon className="size-3.5 animate-spin" />
|
||||
<span>Running…</span>
|
||||
</div>
|
||||
) : output ? (
|
||||
<div className="mt-2 border-t border-dashed pt-2">
|
||||
<div className="flex items-center justify-between">
|
||||
<span className="text-xs font-medium text-muted-foreground">output</span>
|
||||
<CopyBtn text={output} />
|
||||
</div>
|
||||
<pre className="mt-1 max-h-60 overflow-auto whitespace-pre-wrap break-words font-mono text-xs">
|
||||
{truncate(output)}
|
||||
</pre>
|
||||
</div>
|
||||
) : null}
|
||||
</div>
|
||||
</ToolFallbackContent>
|
||||
</ToolFallbackRoot>
|
||||
);
|
||||
};
|
||||
|
||||
export const PythonToolUI = memo(
|
||||
PythonToolUIImpl,
|
||||
) as unknown as ToolCallMessagePartComponent;
|
||||
PythonToolUI.displayName = "PythonToolUI";
|
||||
103
studio/frontend/src/components/assistant-ui/tool-ui-terminal.tsx
Normal file
103
studio/frontend/src/components/assistant-ui/tool-ui-terminal.tsx
Normal file
|
|
@ -0,0 +1,103 @@
|
|||
// SPDX-License-Identifier: AGPL-3.0-only
|
||||
// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
|
||||
|
||||
"use client";
|
||||
|
||||
import { copyToClipboard } from "@/lib/copy-to-clipboard";
|
||||
import type { ToolCallMessagePartComponent } from "@assistant-ui/react";
|
||||
import { CheckIcon, CopyIcon, LoaderIcon, TerminalIcon } from "lucide-react";
|
||||
import { memo, useCallback, useRef, useState } from "react";
|
||||
import {
|
||||
ToolFallbackContent,
|
||||
ToolFallbackRoot,
|
||||
ToolFallbackTrigger,
|
||||
} from "./tool-fallback";
|
||||
|
||||
const MAX_DISPLAY = 10_000;
|
||||
const COPY_RESET_MS = 2000;
|
||||
|
||||
function truncate(text: string): string {
|
||||
return text.length <= MAX_DISPLAY
|
||||
? text
|
||||
: `${text.slice(0, MAX_DISPLAY)}\n... (truncated)`;
|
||||
}
|
||||
|
||||
function CopyBtn({ text }: { text: string }) {
|
||||
const [copied, setCopied] = useState(false);
|
||||
const timer = useRef<ReturnType<typeof setTimeout> | null>(null);
|
||||
const copy = useCallback(() => {
|
||||
if (copyToClipboard(text)) {
|
||||
setCopied(true);
|
||||
if (timer.current) {
|
||||
clearTimeout(timer.current);
|
||||
}
|
||||
timer.current = setTimeout(() => setCopied(false), COPY_RESET_MS);
|
||||
}
|
||||
}, [text]);
|
||||
|
||||
return (
|
||||
<button
|
||||
type="button"
|
||||
onClick={copy}
|
||||
className="inline-flex items-center gap-1 rounded px-1.5 py-0.5 text-xs text-muted-foreground transition-colors hover:bg-muted hover:text-foreground"
|
||||
aria-label="Copy to clipboard"
|
||||
>
|
||||
{copied ? (
|
||||
<CheckIcon className="size-3" />
|
||||
) : (
|
||||
<CopyIcon className="size-3" />
|
||||
)}
|
||||
{copied ? "Copied" : "Copy"}
|
||||
</button>
|
||||
);
|
||||
}
|
||||
|
||||
const TerminalToolUIImpl: ToolCallMessagePartComponent = ({
|
||||
args,
|
||||
result,
|
||||
status,
|
||||
}) => {
|
||||
const command = (args as { command?: string })?.command ?? "";
|
||||
const isRunning = status?.type === "running";
|
||||
const output =
|
||||
typeof result === "string"
|
||||
? result
|
||||
: result
|
||||
? JSON.stringify(result, null, 2)
|
||||
: "";
|
||||
|
||||
return (
|
||||
<ToolFallbackRoot>
|
||||
<ToolFallbackTrigger
|
||||
toolName={command ? `$ ${command.slice(0, 60)}` : "Terminal"}
|
||||
status={status}
|
||||
icon={TerminalIcon}
|
||||
/>
|
||||
<ToolFallbackContent>
|
||||
<div className="flex flex-col px-4">
|
||||
{isRunning ? (
|
||||
<div className="flex items-center gap-2 text-sm text-muted-foreground">
|
||||
<LoaderIcon className="size-3.5 animate-spin" />
|
||||
<span>Running…</span>
|
||||
</div>
|
||||
) : output ? (
|
||||
<div>
|
||||
<div className="flex items-center justify-between">
|
||||
<span className="text-xs font-medium text-muted-foreground">output</span>
|
||||
<CopyBtn text={output} />
|
||||
</div>
|
||||
<pre className="mt-1 max-h-60 overflow-auto whitespace-pre-wrap break-words font-mono text-xs">
|
||||
{truncate(output)}
|
||||
</pre>
|
||||
</div>
|
||||
) : null}
|
||||
</div>
|
||||
</ToolFallbackContent>
|
||||
</ToolFallbackRoot>
|
||||
);
|
||||
};
|
||||
|
||||
export const TerminalToolUI = memo(
|
||||
TerminalToolUIImpl,
|
||||
) as unknown as ToolCallMessagePartComponent;
|
||||
TerminalToolUI.displayName = "TerminalToolUI";
|
||||
|
|
@ -0,0 +1,122 @@
|
|||
// SPDX-License-Identifier: AGPL-3.0-only
|
||||
// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
|
||||
|
||||
"use client";
|
||||
|
||||
import { type ToolCallMessagePartComponent, useAuiState } from "@assistant-ui/react";
|
||||
import { GlobeIcon, LoaderIcon } from "lucide-react";
|
||||
import { memo, useEffect, useState } from "react";
|
||||
import { Source, SourceIcon, SourceTitle } from "./sources";
|
||||
import {
|
||||
ToolFallbackContent,
|
||||
ToolFallbackRoot,
|
||||
ToolFallbackTrigger,
|
||||
} from "./tool-fallback";
|
||||
|
||||
interface ParsedSource {
|
||||
title: string;
|
||||
url: string;
|
||||
snippet: string;
|
||||
}
|
||||
|
||||
const RE_BLOCK_SEP = /\n---\n/;
|
||||
const RE_TITLE = /Title:\s*(.+)/;
|
||||
const RE_URL = /URL:\s*(.+)/;
|
||||
const RE_SNIPPET = /Snippet:\s*(.+)/s;
|
||||
|
||||
/** Parse the backend's "Title: ...\nURL: ...\nSnippet: ...\n---" format into structured sources. */
|
||||
function parseSearchResults(raw: string): ParsedSource[] {
|
||||
if (!raw) {
|
||||
return [];
|
||||
}
|
||||
const blocks = raw.split(RE_BLOCK_SEP).filter(Boolean);
|
||||
const sources: ParsedSource[] = [];
|
||||
for (const block of blocks) {
|
||||
const titleMatch = block.match(RE_TITLE);
|
||||
const urlMatch = block.match(RE_URL);
|
||||
const snippetMatch = block.match(RE_SNIPPET);
|
||||
if (titleMatch && urlMatch) {
|
||||
sources.push({
|
||||
title: titleMatch[1].trim(),
|
||||
url: urlMatch[1].trim(),
|
||||
snippet: snippetMatch?.[1]?.trim() ?? "",
|
||||
});
|
||||
}
|
||||
}
|
||||
return sources;
|
||||
}
|
||||
|
||||
const WebSearchToolUIImpl: ToolCallMessagePartComponent = ({
|
||||
args,
|
||||
result,
|
||||
status,
|
||||
}) => {
|
||||
const query = (args as { query?: string })?.query ?? "";
|
||||
const isRunning = status?.type === "running";
|
||||
const sources = result
|
||||
? parseSearchResults(
|
||||
typeof result === "string" ? result : JSON.stringify(result),
|
||||
)
|
||||
: [];
|
||||
|
||||
// Collapse when LLM starts generating text after the tool call
|
||||
const hasText = useAuiState(({ message }) =>
|
||||
message.content.some((p) => p.type === "text" && "text" in p && (p as { text: string }).text.length > 0),
|
||||
);
|
||||
const [open, setOpen] = useState(isRunning);
|
||||
useEffect(() => {
|
||||
if (isRunning) {
|
||||
setOpen(true);
|
||||
} else if (hasText) {
|
||||
setOpen(false);
|
||||
}
|
||||
}, [isRunning, hasText]);
|
||||
|
||||
return (
|
||||
<ToolFallbackRoot open={open} onOpenChange={setOpen}>
|
||||
<ToolFallbackTrigger
|
||||
toolName={query ? `Searched "${query}"` : "Web Search"}
|
||||
status={status}
|
||||
icon={GlobeIcon}
|
||||
/>
|
||||
<ToolFallbackContent>
|
||||
{isRunning ? (
|
||||
<div className="flex items-center gap-2 px-4 text-sm text-muted-foreground">
|
||||
<LoaderIcon className="size-3.5 animate-spin" />
|
||||
<span>Searching for “{query}”…</span>
|
||||
</div>
|
||||
) : sources.length > 0 ? (
|
||||
<div className="flex flex-col gap-1.5 px-4">
|
||||
{sources.map((source) => (
|
||||
<Source
|
||||
key={source.url}
|
||||
href={source.url}
|
||||
variant="outline"
|
||||
size="default"
|
||||
className="flex w-full max-w-full items-center gap-2 py-1.5"
|
||||
>
|
||||
<SourceIcon url={source.url} className="size-3.5" />
|
||||
<SourceTitle className="max-w-none flex-1 truncate">
|
||||
{source.title}
|
||||
</SourceTitle>
|
||||
</Source>
|
||||
))}
|
||||
</div>
|
||||
) : result ? (
|
||||
<div className="px-4">
|
||||
<pre className="max-h-40 overflow-auto whitespace-pre-wrap break-words rounded bg-muted/50 p-2 text-xs">
|
||||
{typeof result === "string"
|
||||
? result
|
||||
: JSON.stringify(result, null, 2)}
|
||||
</pre>
|
||||
</div>
|
||||
) : null}
|
||||
</ToolFallbackContent>
|
||||
</ToolFallbackRoot>
|
||||
);
|
||||
};
|
||||
|
||||
export const WebSearchToolUI = memo(
|
||||
WebSearchToolUIImpl,
|
||||
) as unknown as ToolCallMessagePartComponent;
|
||||
WebSearchToolUI.displayName = "WebSearchToolUI";
|
||||
|
|
@ -26,6 +26,28 @@ type RunMessage = RunMessages[number];
|
|||
/** Tracks which user messages were sent with an audio file (messageId → filename). */
|
||||
export const sentAudioNames = new Map<string, string>();
|
||||
|
||||
/** Parse "Title: ...\nURL: ...\nSnippet: ..." blocks into source content parts. */
|
||||
function parseSourcesFromResult(raw: string): { type: "source"; sourceType: "url"; id: string; url: string; title: string }[] {
|
||||
if (!raw) return [];
|
||||
const blocks = raw.split(/\n---\n/).filter(Boolean);
|
||||
const sources: { type: "source"; sourceType: "url"; id: string; url: string; title: string }[] = [];
|
||||
for (const block of blocks) {
|
||||
const titleMatch = block.match(/Title:\s*(.+)/);
|
||||
const urlMatch = block.match(/URL:\s*(.+)/);
|
||||
if (titleMatch && urlMatch) {
|
||||
const url = urlMatch[1].trim();
|
||||
sources.push({
|
||||
type: "source" as const,
|
||||
sourceType: "url" as const,
|
||||
id: url,
|
||||
url,
|
||||
title: titleMatch[1].trim(),
|
||||
});
|
||||
}
|
||||
}
|
||||
return sources;
|
||||
}
|
||||
|
||||
function estimateTokenCount(text: string): number | undefined {
|
||||
const trimmed = text.trim();
|
||||
if (!trimmed) {
|
||||
|
|
@ -40,6 +62,7 @@ function buildTiming(
|
|||
firstTokenTime?: number,
|
||||
totalStreamTime?: number,
|
||||
tokenCount?: number,
|
||||
toolCallCount = 0,
|
||||
): MessageTiming {
|
||||
return {
|
||||
streamStartTime,
|
||||
|
|
@ -53,7 +76,7 @@ function buildTiming(
|
|||
? tokenCount / (totalStreamTime / 1000)
|
||||
: undefined,
|
||||
totalChunks,
|
||||
toolCallCount: 0,
|
||||
toolCallCount,
|
||||
};
|
||||
}
|
||||
|
||||
|
|
@ -502,6 +525,9 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter {
|
|||
let cumulativeText = "";
|
||||
let reasoningStartAt: number | null = null;
|
||||
let reasoningDuration = 0;
|
||||
// Tool call content parts — accumulated and yielded cumulatively.
|
||||
// result is set directly on the tool-call part when tool_end arrives.
|
||||
const toolCallParts: { type: "tool-call"; toolCallId: string; toolName: string; args: Record<string, unknown>; result?: unknown }[] = [];
|
||||
|
||||
try {
|
||||
const { supportsReasoning, reasoningEnabled } = runtime;
|
||||
|
|
@ -528,6 +554,13 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter {
|
|||
...(toolsEnabled ? ["web_search"] : []),
|
||||
...(codeToolsEnabled ? ["python", "terminal"] : []),
|
||||
],
|
||||
auto_heal_tool_calls: useChatRuntimeStore.getState().autoHealToolCalls,
|
||||
max_tool_calls_per_message: useChatRuntimeStore.getState().maxToolCallsPerMessage,
|
||||
tool_call_timeout: (() => {
|
||||
const mins = useChatRuntimeStore.getState().toolCallTimeout;
|
||||
return mins >= 9999 ? 9999 : mins * 60;
|
||||
})(),
|
||||
session_id: unstable_threadId || undefined,
|
||||
}
|
||||
: {}),
|
||||
},
|
||||
|
|
@ -542,6 +575,39 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter {
|
|||
continue;
|
||||
}
|
||||
|
||||
// Emit tool-call content parts for assistant-ui.
|
||||
// On tool_start: add a new tool-call part (renders in "running" state).
|
||||
// On tool_end: set result on the existing part (transitions to "complete").
|
||||
const toolEvent = (chunk as unknown as { _toolEvent?: Record<string, unknown> })._toolEvent;
|
||||
if (toolEvent !== undefined) {
|
||||
if (toolEvent.type === "tool_start") {
|
||||
const id = (toolEvent.tool_call_id as string) || `${toolEvent.tool_name}_${Date.now()}`;
|
||||
toolCallParts.push({
|
||||
type: "tool-call" as const,
|
||||
toolCallId: id,
|
||||
toolName: toolEvent.tool_name as string,
|
||||
args: (toolEvent.arguments as Record<string, unknown>) ?? {},
|
||||
});
|
||||
} else if (toolEvent.type === "tool_end") {
|
||||
const id = (toolEvent.tool_call_id as string) ||
|
||||
toolCallParts[toolCallParts.length - 1]?.toolCallId || "";
|
||||
const idx = toolCallParts.findIndex((p) => p.toolCallId === id);
|
||||
if (idx !== -1) {
|
||||
toolCallParts[idx] = { ...toolCallParts[idx], result: toolEvent.result as string };
|
||||
}
|
||||
}
|
||||
// Yield cumulative state so tool UI updates (tools first, text after)
|
||||
const textParts = parseAssistantContent(cumulativeText);
|
||||
yield {
|
||||
content: [...toolCallParts, ...textParts],
|
||||
metadata: {
|
||||
timing: buildTiming(streamStartTime, totalChunks, firstTokenTime),
|
||||
custom: { reasoningDuration },
|
||||
},
|
||||
};
|
||||
continue;
|
||||
}
|
||||
|
||||
totalChunks += 1;
|
||||
const delta = chunk.choices?.[0]?.delta?.content;
|
||||
if (!delta) {
|
||||
|
|
@ -564,9 +630,9 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter {
|
|||
reasoningDuration = Math.round((Date.now() - reasoningStartAt) / 1000);
|
||||
}
|
||||
|
||||
if (parts.length > 0) {
|
||||
if (parts.length > 0 || toolCallParts.length > 0) {
|
||||
yield {
|
||||
content: parts,
|
||||
content: [...toolCallParts, ...parts],
|
||||
metadata: {
|
||||
timing: buildTiming(
|
||||
streamStartTime,
|
||||
|
|
@ -579,7 +645,19 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter {
|
|||
}
|
||||
}
|
||||
settleFirstTokenOk();
|
||||
|
||||
// Extract source parts from completed web_search tool calls
|
||||
const sourceParts = toolCallParts.flatMap((tc) => {
|
||||
if (tc.toolName !== "web_search" || !tc.result) return [];
|
||||
return parseSourcesFromResult(typeof tc.result === "string" ? tc.result : "");
|
||||
});
|
||||
|
||||
yield {
|
||||
content: [
|
||||
...toolCallParts,
|
||||
...parseAssistantContent(cumulativeText),
|
||||
...sourceParts,
|
||||
],
|
||||
metadata: {
|
||||
timing: buildTiming(
|
||||
streamStartTime,
|
||||
|
|
@ -587,6 +665,7 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter {
|
|||
firstTokenTime,
|
||||
Date.now() - streamStartTime,
|
||||
estimateTokenCount(cumulativeText),
|
||||
toolCallParts.length,
|
||||
),
|
||||
custom: { reasoningDuration },
|
||||
},
|
||||
|
|
|
|||
|
|
@ -234,6 +234,12 @@ export async function* streamChatCompletions(
|
|||
separatorIndex = buffer.search(/\r?\n\r?\n/);
|
||||
continue;
|
||||
}
|
||||
// Tool start/end events carry full input/output for the tool outputs panel
|
||||
if ("type" in parsed && (parsed.type === "tool_start" || parsed.type === "tool_end")) {
|
||||
yield { _toolEvent: parsed } as unknown as OpenAIChatChunk;
|
||||
separatorIndex = buffer.search(/\r?\n\r?\n/);
|
||||
continue;
|
||||
}
|
||||
yield parsed as OpenAIChatChunk;
|
||||
separatorIndex = buffer.search(/\r?\n\r?\n/);
|
||||
}
|
||||
|
|
|
|||
|
|
@ -426,6 +426,9 @@ export function ChatSettingsPanel({
|
|||
</Select>
|
||||
</div>
|
||||
)}
|
||||
<AutoHealToolCallsToggle />
|
||||
<MaxToolCallsSlider />
|
||||
<ToolCallTimeoutSlider />
|
||||
</div>
|
||||
</CollapsibleSection>
|
||||
|
||||
|
|
@ -436,6 +439,73 @@ export function ChatSettingsPanel({
|
|||
);
|
||||
}
|
||||
|
||||
function MaxToolCallsSlider() {
|
||||
const maxToolCalls = useChatRuntimeStore((s) => s.maxToolCallsPerMessage);
|
||||
const setMaxToolCalls = useChatRuntimeStore((s) => s.setMaxToolCallsPerMessage);
|
||||
|
||||
// Slider range 0-41; 41 maps to 9999 ("Max")
|
||||
const sliderValue = maxToolCalls >= 9999 ? 41 : Math.min(maxToolCalls, 40);
|
||||
|
||||
return (
|
||||
<ParamSlider
|
||||
label="Max Tool Calls Per Message"
|
||||
value={sliderValue}
|
||||
min={0}
|
||||
max={41}
|
||||
step={1}
|
||||
onChange={(v) => setMaxToolCalls(v >= 41 ? 9999 : v)}
|
||||
displayValue={sliderValue >= 41 ? "Max" : sliderValue === 0 ? "Off" : undefined}
|
||||
/>
|
||||
);
|
||||
}
|
||||
|
||||
function ToolCallTimeoutSlider() {
|
||||
const timeout = useChatRuntimeStore((s) => s.toolCallTimeout);
|
||||
const setTimeout_ = useChatRuntimeStore((s) => s.setToolCallTimeout);
|
||||
|
||||
// Slider 1-31; 31 maps to 9999 ("Max")
|
||||
const sliderValue = timeout >= 9999 ? 31 : Math.min(Math.max(timeout, 1), 30);
|
||||
|
||||
const displayValue =
|
||||
sliderValue >= 31
|
||||
? "Max"
|
||||
: sliderValue === 1
|
||||
? "1 minute"
|
||||
: `${sliderValue} minutes`;
|
||||
|
||||
return (
|
||||
<ParamSlider
|
||||
label="Max Tool Call Duration"
|
||||
value={sliderValue}
|
||||
min={1}
|
||||
max={31}
|
||||
step={1}
|
||||
onChange={(v) => setTimeout_(v >= 31 ? 9999 : v)}
|
||||
displayValue={displayValue}
|
||||
/>
|
||||
);
|
||||
}
|
||||
|
||||
function AutoHealToolCallsToggle() {
|
||||
const autoHealToolCalls = useChatRuntimeStore((s) => s.autoHealToolCalls);
|
||||
const setAutoHealToolCalls = useChatRuntimeStore((s) => s.setAutoHealToolCalls);
|
||||
|
||||
return (
|
||||
<div className="flex items-center justify-between gap-3">
|
||||
<div className="min-w-0">
|
||||
<div className="text-xs font-medium">Auto Heal Tool Calls 🦥</div>
|
||||
<div className="text-[11px] text-muted-foreground">
|
||||
Fix malformed tool calls from the model automatically.
|
||||
</div>
|
||||
</div>
|
||||
<Switch
|
||||
checked={autoHealToolCalls}
|
||||
onCheckedChange={setAutoHealToolCalls}
|
||||
/>
|
||||
</div>
|
||||
);
|
||||
}
|
||||
|
||||
function ChatTemplateSection({
|
||||
onReloadModel,
|
||||
}: {
|
||||
|
|
|
|||
|
|
@ -6,6 +6,7 @@ import { Spinner } from "@/components/ui/spinner";
|
|||
import { Button } from "@/components/ui/button";
|
||||
|
||||
type ModelLoadDescriptionProps = {
|
||||
title?: string | null;
|
||||
message?: string | null;
|
||||
progressPercent?: number | null;
|
||||
progressLabel?: string | null;
|
||||
|
|
@ -17,6 +18,7 @@ function clampProgress(value: number): number {
|
|||
}
|
||||
|
||||
export function ModelLoadDescription({
|
||||
title,
|
||||
message,
|
||||
progressPercent,
|
||||
progressLabel,
|
||||
|
|
@ -25,10 +27,14 @@ export function ModelLoadDescription({
|
|||
const hasProgress = typeof progressPercent === "number";
|
||||
|
||||
return (
|
||||
<div className="flex items-center gap-1.5">
|
||||
<div className="min-w-0 flex-1">
|
||||
<div className="relative flex min-h-12 w-full items-stretch gap-2">
|
||||
<div className="flex h-full shrink-0 items-center self-center">
|
||||
<Spinner className="size-4 text-foreground" />
|
||||
</div>
|
||||
<div className="min-w-0 flex-1 pr-5">
|
||||
{title ? <p className="text-foreground leading-5 font-semibold">{title}</p> : null}
|
||||
{hasProgress ? (
|
||||
<div className="w-[12.5rem] max-w-full">
|
||||
<div className="w-full pt-1">
|
||||
<div className="flex items-center justify-between text-[10px] font-medium tracking-[0.08em] text-muted-foreground/80">
|
||||
<span>{progressLabel}</span>
|
||||
<span>{Math.round(clampProgress(progressPercent))}%</span>
|
||||
|
|
@ -36,18 +42,19 @@ export function ModelLoadDescription({
|
|||
<Progress value={clampProgress(progressPercent)} className="h-1 bg-foreground/[0.08]" />
|
||||
</div>
|
||||
) : message ? (
|
||||
<p className="text-xs leading-relaxed text-muted-foreground">{message}</p>
|
||||
<p className="pt-1 text-xs leading-relaxed text-muted-foreground">{message}</p>
|
||||
) : null}
|
||||
</div>
|
||||
{onStop ? (
|
||||
<Button
|
||||
type="button"
|
||||
size="xs"
|
||||
variant="outline"
|
||||
className="h-5 shrink-0 px-2 text-[10px]"
|
||||
variant="ghost"
|
||||
aria-label="Stop model loading"
|
||||
className="h-auto self-stretch shrink-0 !rounded-none !border-0 bg-transparent px-1 text-[10px] text-muted-foreground hover:bg-transparent hover:text-destructive focus-visible:text-destructive"
|
||||
onClick={onStop}
|
||||
>
|
||||
Stop
|
||||
Cancel
|
||||
</Button>
|
||||
) : null}
|
||||
</div>
|
||||
|
|
|
|||
|
|
@ -3,7 +3,6 @@
|
|||
|
||||
import { createElement, useCallback, useRef, useState } from "react";
|
||||
import { toast } from "sonner";
|
||||
import { Spinner } from "@/components/ui/spinner";
|
||||
import { ModelLoadDescription } from "../components/model-load-status";
|
||||
import {
|
||||
getDownloadProgress,
|
||||
|
|
@ -34,12 +33,10 @@ type SelectedModelInput = {
|
|||
};
|
||||
|
||||
const MODEL_LOAD_TOAST_CLASSNAMES = {
|
||||
toast: "items-start gap-2.5 pr-8",
|
||||
content: "gap-0.5",
|
||||
toast: "items-start gap-2.5",
|
||||
content: "gap-0.5 flex-1 min-w-0",
|
||||
title: "leading-5",
|
||||
description: "mt-0",
|
||||
closeButton:
|
||||
"!left-auto !right-1.5 !top-1.5 !translate-x-0 !translate-y-0 !border-transparent !bg-transparent !shadow-none hover:!bg-transparent hover:opacity-70",
|
||||
description: "mt-0 w-full",
|
||||
} as const;
|
||||
|
||||
const LORA_SUFFIX_RE = /_(\d{9,})$/;
|
||||
|
|
@ -200,12 +197,14 @@ export function useChatModelRuntime() {
|
|||
|
||||
const renderLoadDescription = useCallback(
|
||||
(
|
||||
title: string,
|
||||
message: string,
|
||||
progressPercent?: number | null,
|
||||
progressLabel?: string | null,
|
||||
onStop?: () => void,
|
||||
) =>
|
||||
createElement(ModelLoadDescription, {
|
||||
title,
|
||||
message,
|
||||
progressPercent,
|
||||
progressLabel,
|
||||
|
|
@ -448,18 +447,19 @@ export function useChatModelRuntime() {
|
|||
}
|
||||
}
|
||||
|
||||
const toastTitle = isDownloaded ? "Starting model…" : "Downloading model…";
|
||||
const toastId = toast(
|
||||
isDownloaded ? "Starting model…" : "Downloading model…",
|
||||
null,
|
||||
{
|
||||
icon: createElement(Spinner, { className: "size-4" }),
|
||||
description: renderLoadDescription(
|
||||
toastTitle,
|
||||
loadingDescription,
|
||||
isDownloaded ? null : 0,
|
||||
isDownloaded ? null : "Preparing download",
|
||||
cancelLoading,
|
||||
),
|
||||
duration: Infinity,
|
||||
closeButton: true,
|
||||
closeButton: false,
|
||||
classNames: MODEL_LOAD_TOAST_CLASSNAMES,
|
||||
onDismiss: (dismissedToast) => {
|
||||
if (loadToastIdRef.current !== dismissedToast.id) {
|
||||
|
|
@ -505,18 +505,18 @@ export function useChatModelRuntime() {
|
|||
});
|
||||
if (loadToastDismissedRef.current) return;
|
||||
toast(
|
||||
"Downloading model…",
|
||||
null,
|
||||
{
|
||||
id: toastId,
|
||||
icon: createElement(Spinner, { className: "size-4" }),
|
||||
description: renderLoadDescription(
|
||||
"Downloading model…",
|
||||
loadingDescription,
|
||||
pct,
|
||||
progressLabel,
|
||||
cancelLoading,
|
||||
),
|
||||
duration: Infinity,
|
||||
closeButton: true,
|
||||
closeButton: false,
|
||||
classNames: MODEL_LOAD_TOAST_CLASSNAMES,
|
||||
onDismiss: (dismissedToast) => {
|
||||
if (loadToastIdRef.current !== dismissedToast.id) return;
|
||||
|
|
@ -542,17 +542,17 @@ export function useChatModelRuntime() {
|
|||
if (progressInterval) clearInterval(progressInterval);
|
||||
return;
|
||||
}
|
||||
toast("Starting model…", {
|
||||
toast(null, {
|
||||
id: toastId,
|
||||
icon: createElement(Spinner, { className: "size-4" }),
|
||||
description: renderLoadDescription(
|
||||
"Starting model…",
|
||||
"Download complete. Loading the model into memory.",
|
||||
100,
|
||||
"Download complete",
|
||||
cancelLoading,
|
||||
),
|
||||
duration: Infinity,
|
||||
closeButton: true,
|
||||
closeButton: false,
|
||||
classNames: MODEL_LOAD_TOAST_CLASSNAMES,
|
||||
onDismiss: (dismissedToast) => {
|
||||
if (loadToastIdRef.current !== dismissedToast.id) return;
|
||||
|
|
|
|||
|
|
@ -363,8 +363,8 @@ export function SharedComposer({
|
|||
// Side 1: load → generate → wait
|
||||
if (handle1 && model1?.id) {
|
||||
toast("Loading Model 1…", { id: toastId, description: name1, duration: Infinity });
|
||||
await ensureModelLoaded(model1);
|
||||
toast("Generating with Model 1…", { id: toastId, description: name1, duration: Infinity });
|
||||
const status1 = await ensureModelLoaded(model1);
|
||||
toast("Generating with Model 1…", { id: toastId, description: `${name1} (${status1})`, duration: Infinity });
|
||||
const done = handle1.waitForRunEnd();
|
||||
handle1.startRun();
|
||||
await done;
|
||||
|
|
@ -377,8 +377,8 @@ export function SharedComposer({
|
|||
if (needsLoad) {
|
||||
toast("Loading Model 2…", { id: toastId, description: name2, duration: Infinity });
|
||||
}
|
||||
await ensureModelLoaded(model2);
|
||||
toast("Generating with Model 2…", { id: toastId, description: name2, duration: Infinity });
|
||||
const status2 = await ensureModelLoaded(model2);
|
||||
toast("Generating with Model 2…", { id: toastId, description: `${name2} (${status2})`, duration: Infinity });
|
||||
const done = handle2.waitForRunEnd();
|
||||
handle2.startRun();
|
||||
await done;
|
||||
|
|
|
|||
|
|
@ -10,6 +10,9 @@ import {
|
|||
} from "../types/runtime";
|
||||
|
||||
const AUTO_TITLE_KEY = "unsloth_chat_auto_title";
|
||||
const AUTO_HEAL_TOOL_CALLS_KEY = "unsloth_auto_heal_tool_calls";
|
||||
const MAX_TOOL_CALLS_KEY = "unsloth_max_tool_calls_per_message";
|
||||
const TOOL_CALL_TIMEOUT_KEY = "unsloth_tool_call_timeout";
|
||||
|
||||
function canUseStorage(): boolean {
|
||||
return typeof window !== "undefined";
|
||||
|
|
@ -35,6 +38,27 @@ function saveBool(key: string, value: boolean): void {
|
|||
}
|
||||
}
|
||||
|
||||
function loadInt(key: string, fallback: number): number {
|
||||
if (!canUseStorage()) return fallback;
|
||||
try {
|
||||
const raw = localStorage.getItem(key);
|
||||
if (raw === null) return fallback;
|
||||
const parsed = parseInt(raw, 10);
|
||||
return Number.isNaN(parsed) ? fallback : parsed;
|
||||
} catch {
|
||||
return fallback;
|
||||
}
|
||||
}
|
||||
|
||||
function saveInt(key: string, value: number): void {
|
||||
if (!canUseStorage()) return;
|
||||
try {
|
||||
localStorage.setItem(key, String(value));
|
||||
} catch {
|
||||
// ignore
|
||||
}
|
||||
}
|
||||
|
||||
type ChatRuntimeStore = {
|
||||
params: InferenceParams;
|
||||
models: ChatModelSummary[];
|
||||
|
|
@ -51,6 +75,9 @@ type ChatRuntimeStore = {
|
|||
codeToolsEnabled: boolean;
|
||||
toolStatus: string | null;
|
||||
generatingStatus: string | null;
|
||||
autoHealToolCalls: boolean;
|
||||
maxToolCallsPerMessage: number;
|
||||
toolCallTimeout: number;
|
||||
kvCacheDtype: string | null;
|
||||
defaultChatTemplate: string | null;
|
||||
chatTemplateOverride: string | null;
|
||||
|
|
@ -73,6 +100,9 @@ type ChatRuntimeStore = {
|
|||
setCodeToolsEnabled: (enabled: boolean) => void;
|
||||
setToolStatus: (status: string | null) => void;
|
||||
setGeneratingStatus: (status: string | null) => void;
|
||||
setAutoHealToolCalls: (enabled: boolean) => void;
|
||||
setMaxToolCallsPerMessage: (value: number) => void;
|
||||
setToolCallTimeout: (value: number) => void;
|
||||
setKvCacheDtype: (dtype: string | null) => void;
|
||||
setChatTemplateOverride: (template: string | null) => void;
|
||||
setPendingAudio: (base64: string, name: string) => void;
|
||||
|
|
@ -95,6 +125,9 @@ export const useChatRuntimeStore = create<ChatRuntimeStore>((set) => ({
|
|||
codeToolsEnabled: false,
|
||||
toolStatus: null,
|
||||
generatingStatus: null,
|
||||
autoHealToolCalls: loadBool(AUTO_HEAL_TOOL_CALLS_KEY, true),
|
||||
maxToolCallsPerMessage: loadInt(MAX_TOOL_CALLS_KEY, 10),
|
||||
toolCallTimeout: loadInt(TOOL_CALL_TIMEOUT_KEY, 5),
|
||||
kvCacheDtype: null,
|
||||
defaultChatTemplate: null,
|
||||
chatTemplateOverride: null,
|
||||
|
|
@ -154,6 +187,21 @@ export const useChatRuntimeStore = create<ChatRuntimeStore>((set) => ({
|
|||
setCodeToolsEnabled: (codeToolsEnabled) => set({ codeToolsEnabled }),
|
||||
setToolStatus: (toolStatus) => set({ toolStatus }),
|
||||
setGeneratingStatus: (generatingStatus) => set({ generatingStatus }),
|
||||
setAutoHealToolCalls: (autoHealToolCalls) =>
|
||||
set(() => {
|
||||
saveBool(AUTO_HEAL_TOOL_CALLS_KEY, autoHealToolCalls);
|
||||
return { autoHealToolCalls };
|
||||
}),
|
||||
setMaxToolCallsPerMessage: (maxToolCallsPerMessage) =>
|
||||
set(() => {
|
||||
saveInt(MAX_TOOL_CALLS_KEY, maxToolCallsPerMessage);
|
||||
return { maxToolCallsPerMessage };
|
||||
}),
|
||||
setToolCallTimeout: (toolCallTimeout) =>
|
||||
set(() => {
|
||||
saveInt(TOOL_CALL_TIMEOUT_KEY, toolCallTimeout);
|
||||
return { toolCallTimeout };
|
||||
}),
|
||||
setKvCacheDtype: (kvCacheDtype) => set({ kvCacheDtype }),
|
||||
setChatTemplateOverride: (chatTemplateOverride) => set({ chatTemplateOverride }),
|
||||
setPendingAudio: (base64, name) =>
|
||||
|
|
|
|||
|
|
@ -156,6 +156,11 @@ export interface OpenAIChatCompletionsRequest {
|
|||
use_adapter?: boolean | string | null;
|
||||
enable_thinking?: boolean | null;
|
||||
enable_tools?: boolean | null;
|
||||
enabled_tools?: string[];
|
||||
auto_heal_tool_calls?: boolean;
|
||||
max_tool_calls_per_message?: number;
|
||||
tool_call_timeout?: number;
|
||||
session_id?: string;
|
||||
}
|
||||
|
||||
export interface OpenAIChatDelta {
|
||||
|
|
|
|||
|
|
@ -310,8 +310,8 @@ export function DatasetSection() {
|
|||
if (datasetSource !== "upload") return;
|
||||
if (!uploadedFile) return;
|
||||
if (selectedLocalDataset) return;
|
||||
// Don't clear if this is a direct file upload (not a recipe directory)
|
||||
if (isLikelyLocalDatasetRef(uploadedFile)) return;
|
||||
// Don't clear if this is a direct file upload (e.g. user uploaded a .jsonl/.csv)
|
||||
if (/\.(jsonl|json|csv|parquet|arrow)$/i.test(uploadedFile)) return;
|
||||
selectLocalDataset(null);
|
||||
}, [
|
||||
datasetSource,
|
||||
|
|
|
|||
|
|
@ -372,7 +372,7 @@ export function ModelSection() {
|
|||
|
||||
<div data-tour="studio-base-model" className="flex flex-col gap-2">
|
||||
<span className="flex items-center gap-1.5 text-xs font-medium text-muted-foreground">
|
||||
Base Model
|
||||
Hugging Face Model
|
||||
<Tooltip>
|
||||
<TooltipTrigger asChild={true}>
|
||||
<button
|
||||
|
|
|
|||
|
|
@ -9,6 +9,14 @@ import {
|
|||
CollapsibleTrigger,
|
||||
} from "@/components/ui/collapsible";
|
||||
import { Input } from "@/components/ui/input";
|
||||
import {
|
||||
Combobox,
|
||||
ComboboxContent,
|
||||
ComboboxEmpty,
|
||||
ComboboxInput,
|
||||
ComboboxItem,
|
||||
ComboboxList,
|
||||
} from "@/components/ui/combobox";
|
||||
import {
|
||||
Select,
|
||||
SelectContent,
|
||||
|
|
@ -37,7 +45,7 @@ import {
|
|||
Settings04Icon,
|
||||
} from "@hugeicons/core-free-icons";
|
||||
import { HugeiconsIcon } from "@hugeicons/react";
|
||||
import { type ReactElement, type ReactNode, useState } from "react";
|
||||
import { type ReactElement, type ReactNode, useEffect, useRef, useState } from "react";
|
||||
|
||||
function Row({
|
||||
label,
|
||||
|
|
@ -120,6 +128,25 @@ export function ParamsSection(): ReactElement {
|
|||
const showVisionLora = store.isVisionModel && store.isDatasetImage === true;
|
||||
const [loraOpen, setLoraOpen] = useState(false);
|
||||
const [hyperOpen, setHyperOpen] = useState(false);
|
||||
const [ctxInput, setCtxInput] = useState(String(store.contextLength));
|
||||
const ctxAnchorRef = useRef<HTMLDivElement>(null);
|
||||
const ctxItems = CONTEXT_LENGTHS.map(String);
|
||||
|
||||
// Keep input in sync when the store value changes externally
|
||||
// (e.g. model defaults being applied after model selection).
|
||||
useEffect(() => {
|
||||
setCtxInput(String(store.contextLength));
|
||||
}, [store.contextLength]);
|
||||
|
||||
const trySetContextLength = (input: string): number | null => {
|
||||
const n = Number(input);
|
||||
if (Number.isInteger(n) && n > 0) {
|
||||
store.setContextLength(n);
|
||||
return n;
|
||||
}
|
||||
return null;
|
||||
};
|
||||
|
||||
const { useEpochs, toggleUseEpochs } = useMaxStepsEpochsToggle({
|
||||
maxSteps: store.maxSteps,
|
||||
epochs: store.epochs,
|
||||
|
|
@ -259,21 +286,51 @@ export function ParamsSection(): ReactElement {
|
|||
</TooltipContent>
|
||||
</Tooltip>
|
||||
</span>
|
||||
<Select
|
||||
value={String(store.contextLength)}
|
||||
onValueChange={(v) => store.setContextLength(Number(v))}
|
||||
>
|
||||
<SelectTrigger className="w-full font-mono">
|
||||
<SelectValue />
|
||||
</SelectTrigger>
|
||||
<SelectContent>
|
||||
{CONTEXT_LENGTHS.map((len) => (
|
||||
<SelectItem key={len} value={String(len)} className="font-mono">
|
||||
{len.toLocaleString()}
|
||||
</SelectItem>
|
||||
))}
|
||||
</SelectContent>
|
||||
</Select>
|
||||
<div ref={ctxAnchorRef}>
|
||||
<Combobox
|
||||
items={ctxItems}
|
||||
filteredItems={ctxItems}
|
||||
filter={null}
|
||||
value={String(store.contextLength)}
|
||||
onValueChange={(v) => {
|
||||
if (v && trySetContextLength(v)) {
|
||||
setCtxInput(v);
|
||||
}
|
||||
}}
|
||||
onInputValueChange={setCtxInput}
|
||||
itemToStringValue={(id) => Number(id).toLocaleString()}
|
||||
autoHighlight={false}
|
||||
>
|
||||
<ComboboxInput
|
||||
placeholder={String(store.contextLength)}
|
||||
className="w-full font-mono"
|
||||
onBlur={() => {
|
||||
trySetContextLength(ctxInput);
|
||||
setCtxInput(String(store.contextLength));
|
||||
}}
|
||||
onKeyDown={(e) => {
|
||||
if (e.key !== "Enter") { return; }
|
||||
const n = trySetContextLength(ctxInput);
|
||||
if (n === null) { return; }
|
||||
if (!ctxItems.includes(ctxInput.trim())) {
|
||||
e.stopPropagation();
|
||||
e.preventDefault();
|
||||
}
|
||||
setCtxInput(String(n));
|
||||
}}
|
||||
/>
|
||||
<ComboboxContent anchor={ctxAnchorRef}>
|
||||
<ComboboxEmpty>Enter a custom value</ComboboxEmpty>
|
||||
<ComboboxList className="p-1">
|
||||
{(id: string) => (
|
||||
<ComboboxItem key={id} value={id} className="font-mono">
|
||||
{Number(id).toLocaleString()}
|
||||
</ComboboxItem>
|
||||
)}
|
||||
</ComboboxList>
|
||||
</ComboboxContent>
|
||||
</Combobox>
|
||||
</div>
|
||||
<p className="text-[10px] text-muted-foreground">
|
||||
Max sequence length for training samples
|
||||
</p>
|
||||
|
|
|
|||
|
|
@ -6,7 +6,7 @@ import { ReadMore, type TourStep } from "@/features/tour";
|
|||
export const studioBaseModelStep: TourStep = {
|
||||
id: "base-model",
|
||||
target: "studio-base-model",
|
||||
title: "Base model from Hugging Face",
|
||||
title: "Hugging Face Model",
|
||||
body: (
|
||||
<>
|
||||
Paste <span className="font-mono">org/model</span> or search. Pick a base
|
||||
|
|
|
|||
|
|
@ -157,7 +157,6 @@
|
|||
--shadow-2xl: 0px 0px 0px 0px hsl(0 0% 0% / 0);
|
||||
}
|
||||
|
||||
|
||||
@theme inline {
|
||||
--font-sans: "Inter Variable", ui-sans-serif, sans-serif, system-ui;
|
||||
--font-heading: "Hellix", "Space Grotesk Variable", ui-sans-serif, sans-serif;
|
||||
|
|
@ -325,8 +324,22 @@
|
|||
[data-streamdown="list-item"] {
|
||||
display: list-item;
|
||||
}
|
||||
}
|
||||
|
||||
/* Flatten code blocks: single border, language label, then code directly */
|
||||
[data-streamdown="code-block-body"] {
|
||||
border: none !important;
|
||||
border-radius: 0 !important;
|
||||
background: transparent !important;
|
||||
padding: 0 !important;
|
||||
}
|
||||
[data-streamdown="code-block"] {
|
||||
gap: 0;
|
||||
padding: 0.5rem;
|
||||
}
|
||||
[data-streamdown="code-block-header"] {
|
||||
padding-left: 0.75rem;
|
||||
}
|
||||
}
|
||||
|
||||
/* Minimal scrollbar — thumb only, no track */
|
||||
* {
|
||||
|
|
@ -373,4 +386,4 @@
|
|||
::view-transition-old(root), ::view-transition-new(root) {
|
||||
animation: none;
|
||||
mix-blend-mode: normal;
|
||||
}
|
||||
}
|
||||
|
|
@ -18,7 +18,7 @@
|
|||
|
||||
/* Linting */
|
||||
"strict": true,
|
||||
"noUnusedLocals": true,
|
||||
"noUnusedLocals": false,
|
||||
"noUnusedParameters": true,
|
||||
"erasableSyntaxOnly": true,
|
||||
"noFallthroughCasesInSwitch": true,
|
||||
|
|
|
|||
|
|
@ -140,15 +140,15 @@ def _bootstrap_uv() -> bool:
|
|||
global UV_NEEDS_SYSTEM
|
||||
if not shutil.which("uv"):
|
||||
return False
|
||||
# Probe: try a dry-run install without --system.
|
||||
# If uv can't find a venv it exits with code 2.
|
||||
# Probe: try a dry-run install targeting the current Python explicitly.
|
||||
# Without --python, uv can ignore the activated venv on some platforms.
|
||||
probe = subprocess.run(
|
||||
["uv", "pip", "install", "--dry-run", "pip"],
|
||||
["uv", "pip", "install", "--dry-run", "--python", sys.executable, "pip"],
|
||||
stdout = subprocess.PIPE,
|
||||
stderr = subprocess.STDOUT,
|
||||
)
|
||||
if probe.returncode != 0:
|
||||
# Retry with --system to confirm it works
|
||||
# Retry with --system (some envs need it when uv can't find a venv)
|
||||
probe_sys = subprocess.run(
|
||||
["uv", "pip", "install", "--dry-run", "--system", "pip"],
|
||||
stdout = subprocess.PIPE,
|
||||
|
|
@ -204,6 +204,10 @@ def _build_uv_cmd(args: tuple[str, ...]) -> list[str]:
|
|||
cmd = ["uv", "pip", "install"]
|
||||
if UV_NEEDS_SYSTEM:
|
||||
cmd.append("--system")
|
||||
# Always pass --python so uv targets the correct environment.
|
||||
# Without this, uv can ignore an activated venv and install into
|
||||
# the system Python (observed on Colab and similar environments).
|
||||
cmd.extend(["--python", sys.executable])
|
||||
cmd.extend(_translate_pip_args_for_uv(args))
|
||||
cmd.append("--torch-backend=auto")
|
||||
return cmd
|
||||
|
|
|
|||
376
studio/setup.ps1
376
studio/setup.ps1
|
|
@ -8,7 +8,7 @@
|
|||
Always installs Node.js if needed. When running from pip install:
|
||||
skips frontend build (already bundled). When running from git repo:
|
||||
full setup including frontend build.
|
||||
Requires an NVIDIA GPU -- CPU-only machines are not supported.
|
||||
Supports NVIDIA GPU (full training + inference) and CPU-only (GGUF chat mode).
|
||||
.NOTES
|
||||
Usage: powershell -ExecutionPolicy Bypass -File setup.ps1
|
||||
#>
|
||||
|
|
@ -107,11 +107,15 @@ function Find-Nvcc {
|
|||
# Returns e.g. "80" for A100 (8.0), "89" for RTX 4090 (8.9), etc.
|
||||
# Returns $null if detection fails.
|
||||
function Get-CudaComputeCapability {
|
||||
$nvSmi = Get-Command nvidia-smi -ErrorAction SilentlyContinue
|
||||
if (-not $nvSmi) { return $null }
|
||||
# Use the resolved absolute path ($NvidiaSmiExe) to survive Refresh-Environment
|
||||
$smiExe = if ($script:NvidiaSmiExe) { $script:NvidiaSmiExe } else {
|
||||
$cmd = Get-Command nvidia-smi -ErrorAction SilentlyContinue
|
||||
if ($cmd) { $cmd.Source } else { $null }
|
||||
}
|
||||
if (-not $smiExe) { return $null }
|
||||
|
||||
try {
|
||||
$raw = & nvidia-smi --query-gpu=compute_cap --format=csv,noheader 2>$null
|
||||
$raw = & $smiExe --query-gpu=compute_cap --format=csv,noheader 2>$null
|
||||
if ($LASTEXITCODE -ne 0 -or -not $raw) { return $null }
|
||||
|
||||
# nvidia-smi may return multiple GPUs; take the first one
|
||||
|
|
@ -168,14 +172,17 @@ function Get-NvccMaxArch {
|
|||
# https://download.pytorch.org/whl/<tag>. The tag must not exceed the driver's
|
||||
# capability: e.g. driver "CUDA Version: 12.9" → cu128 (not cu130).
|
||||
function Get-PytorchCudaTag {
|
||||
$nvSmi = Get-Command nvidia-smi -ErrorAction SilentlyContinue
|
||||
if (-not $nvSmi) { return "cu124" }
|
||||
$smiExe = if ($script:NvidiaSmiExe) { $script:NvidiaSmiExe } else {
|
||||
$cmd = Get-Command nvidia-smi -ErrorAction SilentlyContinue
|
||||
if ($cmd) { $cmd.Source } else { $null }
|
||||
}
|
||||
if (-not $smiExe) { return "cu124" }
|
||||
|
||||
try {
|
||||
# 2>&1 | Out-String merges stderr into stdout then converts to a single
|
||||
# string. Plain 2>$null doesn't fully suppress stderr in PS 5.1 —
|
||||
# string. Plain 2>$null doesn't fully suppress stderr in PS 5.1 --
|
||||
# ErrorRecord objects leak into $output and break the -match.
|
||||
$output = & nvidia-smi 2>&1 | Out-String
|
||||
$output = & $smiExe 2>&1 | Out-String
|
||||
if ($output -match 'CUDA Version:\s+(\d+)\.(\d+)') {
|
||||
$major = [int]$Matches[1]
|
||||
$minor = [int]$Matches[2]
|
||||
|
|
@ -251,23 +258,50 @@ Write-Host "+==============================================+" -ForegroundColor G
|
|||
# ==========================================================================
|
||||
|
||||
# ============================================
|
||||
# 1a. GPU requirement check
|
||||
# 1a. GPU detection
|
||||
# ============================================
|
||||
$HasNvidiaSmi = $false
|
||||
$NvidiaSmiExe = $null # Absolute path -- survives Refresh-Environment
|
||||
try {
|
||||
nvidia-smi 2>&1 | Out-Null
|
||||
if ($LASTEXITCODE -eq 0) { $HasNvidiaSmi = $true }
|
||||
$nvSmiCmd = Get-Command nvidia-smi -ErrorAction SilentlyContinue
|
||||
if ($nvSmiCmd) {
|
||||
& $nvSmiCmd.Source 2>&1 | Out-Null
|
||||
if ($LASTEXITCODE -eq 0) {
|
||||
$HasNvidiaSmi = $true
|
||||
$NvidiaSmiExe = $nvSmiCmd.Source
|
||||
}
|
||||
}
|
||||
} catch {}
|
||||
# Fallback: nvidia-smi may not be on PATH even though a GPU + driver exist.
|
||||
# Check the default install location and the Windows driver store.
|
||||
if (-not $HasNvidiaSmi) {
|
||||
$nvSmiDefaults = @(
|
||||
"$env:ProgramFiles\NVIDIA Corporation\NVSMI\nvidia-smi.exe",
|
||||
"$env:SystemRoot\System32\nvidia-smi.exe"
|
||||
)
|
||||
foreach ($p in $nvSmiDefaults) {
|
||||
if (Test-Path $p) {
|
||||
try {
|
||||
& $p 2>&1 | Out-Null
|
||||
if ($LASTEXITCODE -eq 0) {
|
||||
$HasNvidiaSmi = $true
|
||||
$NvidiaSmiExe = $p
|
||||
Write-Host " Found nvidia-smi at $(Split-Path $p -Parent)" -ForegroundColor Gray
|
||||
break
|
||||
}
|
||||
} catch {}
|
||||
}
|
||||
}
|
||||
}
|
||||
if (-not $HasNvidiaSmi) {
|
||||
Write-Host ""
|
||||
Write-Host "[ERROR] Unsloth Studio requires an NVIDIA GPU." -ForegroundColor Red
|
||||
Write-Host " CPU-only machines are not supported." -ForegroundColor Red
|
||||
Write-Host "[WARN] No NVIDIA GPU detected. Studio will run in chat-only (GGUF) mode." -ForegroundColor Yellow
|
||||
Write-Host " Training and GPU inference require an NVIDIA GPU with drivers installed." -ForegroundColor Yellow
|
||||
Write-Host " https://www.nvidia.com/Download/index.aspx" -ForegroundColor Yellow
|
||||
Write-Host ""
|
||||
Write-Host " If you have an NVIDIA GPU, ensure the driver is installed:" -ForegroundColor Yellow
|
||||
Write-Host " https://www.nvidia.com/Download/index.aspx" -ForegroundColor Yellow
|
||||
exit 1
|
||||
} else {
|
||||
Write-Host "[OK] NVIDIA GPU detected" -ForegroundColor Green
|
||||
}
|
||||
Write-Host "[OK] NVIDIA GPU detected" -ForegroundColor Green
|
||||
|
||||
# ============================================
|
||||
# 1a.5. Windows Long Paths (required for deep node_modules / Python paths)
|
||||
|
|
@ -341,6 +375,30 @@ if (-not $HasCmake) {
|
|||
$HasCmake = $null -ne (Get-Command cmake -ErrorAction SilentlyContinue)
|
||||
} catch { }
|
||||
}
|
||||
# winget may succeed but cmake isn't on PATH yet (MSI PATH changes need a
|
||||
# new shell). Try the default install location as a fallback.
|
||||
if (-not $HasCmake) {
|
||||
$cmakeDefaults = @(
|
||||
"$env:ProgramFiles\CMake\bin",
|
||||
"${env:ProgramFiles(x86)}\CMake\bin",
|
||||
"$env:LOCALAPPDATA\CMake\bin"
|
||||
)
|
||||
foreach ($d in $cmakeDefaults) {
|
||||
if (Test-Path (Join-Path $d "cmake.exe")) {
|
||||
$env:Path = "$d;$env:Path"
|
||||
# Persist to user PATH so Refresh-Environment does not drop it later
|
||||
$userPath = [Environment]::GetEnvironmentVariable('Path', 'User')
|
||||
if (-not $userPath -or $userPath -notlike "*$d*") {
|
||||
[Environment]::SetEnvironmentVariable('Path', "$d;$userPath", 'User')
|
||||
}
|
||||
$HasCmake = $null -ne (Get-Command cmake -ErrorAction SilentlyContinue)
|
||||
if ($HasCmake) {
|
||||
Write-Host " Found cmake at $d (added to PATH)" -ForegroundColor Gray
|
||||
break
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
if ($HasCmake) {
|
||||
Write-Host "[OK] CMake installed" -ForegroundColor Green
|
||||
} else {
|
||||
|
|
@ -389,6 +447,7 @@ if ($vsResult) {
|
|||
# ============================================
|
||||
# 1e. CUDA Toolkit (nvcc for llama.cpp build + env vars)
|
||||
# ============================================
|
||||
if ($HasNvidiaSmi) {
|
||||
# IMPORTANT: The CUDA Toolkit version must be <= the max CUDA version the
|
||||
# NVIDIA driver supports. nvidia-smi reports this as "CUDA Version: X.Y".
|
||||
# If we install a toolkit newer than the driver supports, llama-server will
|
||||
|
|
@ -397,7 +456,7 @@ if ($vsResult) {
|
|||
# -- Detect max CUDA version the driver supports --
|
||||
$DriverMaxCuda = $null
|
||||
try {
|
||||
$smiOut = nvidia-smi 2>&1 | Out-String
|
||||
$smiOut = & $NvidiaSmiExe 2>&1 | Out-String
|
||||
if ($smiOut -match "CUDA Version:\s+([\d]+)\.([\d]+)") {
|
||||
$DriverMaxCuda = "$($Matches[1]).$($Matches[2])"
|
||||
Write-Host " Driver supports up to CUDA $DriverMaxCuda" -ForegroundColor Gray
|
||||
|
|
@ -624,11 +683,24 @@ if ($VsInstallPath -and $CudaToolkitRoot) {
|
|||
Copy-Item "$cudaExtras\*" $vsCustomizations -Force -ErrorAction Stop
|
||||
Write-Host " [OK] CUDA VS integration files installed" -ForegroundColor Green
|
||||
} catch {
|
||||
Write-Host " [WARN] Could not copy CUDA VS integration files (may need admin)" -ForegroundColor Yellow
|
||||
Write-Host " Manual fix: copy contents of" -ForegroundColor Yellow
|
||||
Write-Host " $cudaExtras" -ForegroundColor Cyan
|
||||
Write-Host " into:" -ForegroundColor Yellow
|
||||
Write-Host " $vsCustomizations" -ForegroundColor Cyan
|
||||
# Direct copy failed (needs admin). Try elevated copy via Start-Process.
|
||||
try {
|
||||
$copyCmd = "Copy-Item '$cudaExtras\*' '$vsCustomizations' -Force"
|
||||
Start-Process powershell -ArgumentList "-NoProfile -Command $copyCmd" -Verb RunAs -Wait -ErrorAction Stop
|
||||
$hasTargetsRetry = Get-ChildItem $vsCustomizations -Filter "CUDA *.targets" -ErrorAction SilentlyContinue
|
||||
if ($hasTargetsRetry) {
|
||||
Write-Host " [OK] CUDA VS integration files installed (elevated)" -ForegroundColor Green
|
||||
} else {
|
||||
throw "Copy did not produce .targets files"
|
||||
}
|
||||
} catch {
|
||||
Write-Host " [WARN] Could not copy CUDA VS integration files" -ForegroundColor Yellow
|
||||
Write-Host " The llama.cpp build may fail with 'No CUDA toolset found'." -ForegroundColor Yellow
|
||||
Write-Host " Manual fix: copy contents of" -ForegroundColor Yellow
|
||||
Write-Host " $cudaExtras" -ForegroundColor Cyan
|
||||
Write-Host " into:" -ForegroundColor Yellow
|
||||
Write-Host " $vsCustomizations" -ForegroundColor Cyan
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
|
@ -643,6 +715,9 @@ Write-Host " CudaToolkitDir = $CudaToolkitRoot\" -ForegroundColor Gray
|
|||
if (-not $CudaArch) {
|
||||
Write-Host " [WARN] Could not detect compute capability -- cmake will use defaults" -ForegroundColor Yellow
|
||||
}
|
||||
} else {
|
||||
Write-Host "[SKIP] CUDA Toolkit -- no NVIDIA GPU detected" -ForegroundColor Yellow
|
||||
}
|
||||
|
||||
# ============================================
|
||||
# 1f. Node.js / npm (skip if pip-installed -- only needed for frontend build)
|
||||
|
|
@ -748,9 +823,40 @@ Write-Host ""
|
|||
# ==========================================================================
|
||||
# PHASE 2: Frontend build (skip if pip-installed -- already bundled)
|
||||
# ==========================================================================
|
||||
$DistDir = Join-Path $FrontendDir "dist"
|
||||
# Skip build if dist/ exists and no tracked input is newer than dist/.
|
||||
# Checks src/, public/, package.json, config files -- not just src/.
|
||||
$NeedFrontendBuild = $true
|
||||
if ($IsPipInstall) {
|
||||
$NeedFrontendBuild = $false
|
||||
Write-Host "[OK] Running from pip install - frontend already bundled, skipping build" -ForegroundColor Green
|
||||
} else {
|
||||
} elseif (Test-Path $DistDir) {
|
||||
$DistTime = (Get-Item $DistDir).LastWriteTime
|
||||
$NewerFile = $null
|
||||
# Check src/ and public/ recursively (probe paths directly, not via -Include)
|
||||
foreach ($subDir in @("src", "public")) {
|
||||
$subPath = Join-Path $FrontendDir $subDir
|
||||
if (Test-Path $subPath) {
|
||||
$NewerFile = Get-ChildItem -Path $subPath -Recurse -File -ErrorAction SilentlyContinue |
|
||||
Where-Object { $_.LastWriteTime -gt $DistTime } | Select-Object -First 1
|
||||
if ($NewerFile) { break }
|
||||
}
|
||||
}
|
||||
# Also check all top-level files (package.json, bun.lock, vite.config.ts, index.html, etc.)
|
||||
if (-not $NewerFile) {
|
||||
$NewerFile = Get-ChildItem -Path $FrontendDir -File -ErrorAction SilentlyContinue |
|
||||
Where-Object { $_.LastWriteTime -gt $DistTime } |
|
||||
Select-Object -First 1
|
||||
}
|
||||
if (-not $NewerFile) {
|
||||
$NeedFrontendBuild = $false
|
||||
Write-Host "[OK] Frontend already built and up to date -- skipping build" -ForegroundColor Green
|
||||
} else {
|
||||
Write-Host "[INFO] Frontend source changed since last build -- rebuilding..." -ForegroundColor Yellow
|
||||
}
|
||||
}
|
||||
$NeedFrontendBuild = $true
|
||||
if ($NeedFrontendBuild -and -not $IsPipInstall) {
|
||||
Write-Host ""
|
||||
Write-Host "Building frontend..." -ForegroundColor Cyan
|
||||
# npm writes warnings to stderr; lower ErrorActionPreference so PS doesn't
|
||||
|
|
@ -758,9 +864,6 @@ if ($IsPipInstall) {
|
|||
$prevEAP_npm = $ErrorActionPreference
|
||||
$ErrorActionPreference = "Continue"
|
||||
Push-Location $FrontendDir
|
||||
# Remove stale node_modules and package-lock.json to avoid version conflicts
|
||||
if (Test-Path "node_modules") { Remove-Item -Recurse -Force "node_modules" }
|
||||
if (Test-Path "package-lock.json") { Remove-Item -Force "package-lock.json" }
|
||||
npm install 2>&1 | Out-Null
|
||||
if ($LASTEXITCODE -ne 0) {
|
||||
Pop-Location
|
||||
|
|
@ -845,7 +948,35 @@ $ErrorActionPreference = "Continue"
|
|||
|
||||
$ActivateScript = Join-Path $VenvDir "Scripts\Activate.ps1"
|
||||
. $ActivateScript
|
||||
pip install --upgrade pip 2>&1 | Out-Null
|
||||
|
||||
# Try to use uv (much faster than pip), fall back to pip if unavailable
|
||||
$UseUv = $false
|
||||
if (Get-Command uv -ErrorAction SilentlyContinue) {
|
||||
$UseUv = $true
|
||||
} else {
|
||||
Write-Host " Installing uv package manager..." -ForegroundColor Cyan
|
||||
try {
|
||||
powershell -ExecutionPolicy ByPass -c "irm https://astral.sh/uv/install.ps1 | iex" 2>&1 | Out-Null
|
||||
Refresh-Environment
|
||||
# Re-activate venv since Refresh-Environment rebuilds PATH from
|
||||
# registry and drops the venv's Scripts directory
|
||||
. $ActivateScript
|
||||
if (Get-Command uv -ErrorAction SilentlyContinue) { $UseUv = $true }
|
||||
} catch { }
|
||||
}
|
||||
|
||||
# Helper: install a package, preferring uv with pip fallback
|
||||
function Fast-Install {
|
||||
param([Parameter(ValueFromRemainingArguments=$true)]$Args_)
|
||||
if ($UseUv) {
|
||||
$VenvPy = (Get-Command python).Source
|
||||
$result = & uv pip install --python $VenvPy @Args_ 2>&1
|
||||
if ($LASTEXITCODE -eq 0) { return }
|
||||
}
|
||||
& python -m pip install @Args_ 2>&1
|
||||
}
|
||||
|
||||
Fast-Install --upgrade pip | Out-Null
|
||||
|
||||
# if (-not $IsPipInstall) {
|
||||
# # Running from repo: copy requirements and do editable install
|
||||
|
|
@ -880,16 +1011,37 @@ $env:TORCHINDUCTOR_CACHE_DIR = $TorchCacheDir
|
|||
[Environment]::SetEnvironmentVariable('TORCHINDUCTOR_CACHE_DIR', $TorchCacheDir, 'User')
|
||||
Write-Host "[OK] TORCHINDUCTOR_CACHE_DIR set to $TorchCacheDir (avoids MAX_PATH issues)" -ForegroundColor Green
|
||||
|
||||
$CuTag = Get-PytorchCudaTag
|
||||
Write-Host " Installing PyTorch with CUDA support ($CuTag)..." -ForegroundColor Cyan
|
||||
pip install torch torchvision torchaudio --index-url "https://download.pytorch.org/whl/$CuTag" 2>&1 | Out-Null
|
||||
if ($HasNvidiaSmi) {
|
||||
$CuTag = Get-PytorchCudaTag
|
||||
Write-Host " Installing PyTorch with CUDA support ($CuTag)..." -ForegroundColor Cyan
|
||||
Write-Host " (This download is ~2.8 GB -- may take a few minutes)" -ForegroundColor Gray
|
||||
$output = Fast-Install torch torchvision torchaudio --index-url "https://download.pytorch.org/whl/$CuTag" | Out-String
|
||||
if ($LASTEXITCODE -ne 0) {
|
||||
Write-Host "[FAILED] PyTorch CUDA install failed (exit code $LASTEXITCODE)" -ForegroundColor Red
|
||||
Write-Host $output -ForegroundColor Red
|
||||
exit 1
|
||||
}
|
||||
|
||||
# Install Triton for Windows (enables torch.compile — without it training can hang)
|
||||
Write-Host " Installing Triton for Windows..." -ForegroundColor Cyan
|
||||
pip install "triton-windows<3.7" 2>&1 | Out-Null
|
||||
Write-Host "[OK] Triton for Windows installed (enables torch.compile)" -ForegroundColor Green
|
||||
# Install Triton for Windows (enables torch.compile -- without it training can hang)
|
||||
Write-Host " Installing Triton for Windows..." -ForegroundColor Cyan
|
||||
$output = Fast-Install "triton-windows<3.7" | Out-String
|
||||
if ($LASTEXITCODE -ne 0) {
|
||||
Write-Host "[WARN] Triton install failed -- torch.compile may not work" -ForegroundColor Yellow
|
||||
Write-Host $output -ForegroundColor Yellow
|
||||
} else {
|
||||
Write-Host "[OK] Triton for Windows installed (enables torch.compile)" -ForegroundColor Green
|
||||
}
|
||||
} else {
|
||||
Write-Host " Installing PyTorch (CPU-only)..." -ForegroundColor Cyan
|
||||
$output = Fast-Install torch torchvision torchaudio --index-url "https://download.pytorch.org/whl/cpu" | Out-String
|
||||
if ($LASTEXITCODE -ne 0) {
|
||||
Write-Host "[FAILED] PyTorch install failed (exit code $LASTEXITCODE)" -ForegroundColor Red
|
||||
Write-Host $output -ForegroundColor Red
|
||||
exit 1
|
||||
}
|
||||
}
|
||||
|
||||
# Ordered heavy dependency installation — shared cross-platform script
|
||||
# Ordered heavy dependency installation -- shared cross-platform script
|
||||
Write-Host " Running ordered dependency installation..." -ForegroundColor Cyan
|
||||
python "$PSScriptRoot\install_python_stack.py"
|
||||
# Restore ErrorActionPreference after pip/python work
|
||||
|
|
@ -898,7 +1050,7 @@ $ErrorActionPreference = $prevEAP
|
|||
# ── Pre-install transformers 5.x into .venv_t5/ ──
|
||||
# Models like GLM-4.7-Flash need transformers>=5.3.0. Instead of pip-installing
|
||||
# at runtime (slow, ~10-15s), we pre-install into a separate directory.
|
||||
# The training subprocess just prepends .venv_t5/ to sys.path — instant switch.
|
||||
# The training subprocess just prepends .venv_t5/ to sys.path -- instant switch.
|
||||
Write-Host ""
|
||||
Write-Host " Pre-installing transformers 5.x for newer model support..." -ForegroundColor Cyan
|
||||
$VenvT5Dir = Join-Path $env:USERPROFILE ".unsloth\studio\.venv_t5"
|
||||
|
|
@ -906,17 +1058,20 @@ if (Test-Path $VenvT5Dir) { Remove-Item -Recurse -Force $VenvT5Dir }
|
|||
New-Item -ItemType Directory -Path $VenvT5Dir -Force | Out-Null
|
||||
$prevEAP_t5 = $ErrorActionPreference
|
||||
$ErrorActionPreference = "Continue"
|
||||
pip install --target $VenvT5Dir --no-deps "transformers==5.3.0" 2>&1 | Out-Null
|
||||
if ($LASTEXITCODE -ne 0) {
|
||||
Write-Host "[FAIL] Could not install transformers 5.3.0 into .venv_t5/" -ForegroundColor Red
|
||||
$ErrorActionPreference = $prevEAP_t5
|
||||
exit 1
|
||||
foreach ($pkg in @("transformers==5.3.0", "huggingface_hub==1.7.1", "hf_xet==1.4.2")) {
|
||||
$output = Fast-Install --target $VenvT5Dir --no-deps $pkg | Out-String
|
||||
if ($LASTEXITCODE -ne 0) {
|
||||
Write-Host "[FAIL] Could not install $pkg into .venv_t5/" -ForegroundColor Red
|
||||
Write-Host $output -ForegroundColor Red
|
||||
$ErrorActionPreference = $prevEAP_t5
|
||||
exit 1
|
||||
}
|
||||
}
|
||||
pip install --target $VenvT5Dir --no-deps "huggingface_hub==1.3.0" 2>&1 | Out-Null
|
||||
# tiktoken is needed by Qwen-family tokenizers -- install with deps since
|
||||
# regex/requests may be missing on Windows
|
||||
$output = Fast-Install --target $VenvT5Dir tiktoken | Out-String
|
||||
if ($LASTEXITCODE -ne 0) {
|
||||
Write-Host "[FAIL] Could not install huggingface_hub 1.3.0 into .venv_t5/" -ForegroundColor Red
|
||||
$ErrorActionPreference = $prevEAP_t5
|
||||
exit 1
|
||||
Write-Host "[WARN] Could not install tiktoken into .venv_t5/ -- Qwen tokenizers may fail" -ForegroundColor Yellow
|
||||
}
|
||||
$ErrorActionPreference = $prevEAP_t5
|
||||
Write-Host "[OK] Transformers 5.x pre-installed to .venv_t5/" -ForegroundColor Green
|
||||
|
|
@ -982,12 +1137,46 @@ $LlamaCppDir = Join-Path $UnslothHome "llama.cpp"
|
|||
$BuildDir = Join-Path $LlamaCppDir "build"
|
||||
$LlamaServerBin = Join-Path $BuildDir "bin\Release\llama-server.exe"
|
||||
|
||||
$HasCmakeForBuild = $null -ne (Get-Command cmake -ErrorAction SilentlyContinue)
|
||||
|
||||
# Check if existing llama-server matches current GPU mode. A CUDA-built binary
|
||||
# on a now-CPU-only machine (or vice versa) needs to be rebuilt.
|
||||
$NeedRebuild = $false
|
||||
if (Test-Path $LlamaServerBin) {
|
||||
$CmakeCacheFile = Join-Path $BuildDir "CMakeCache.txt"
|
||||
if (Test-Path $CmakeCacheFile) {
|
||||
$cachedCuda = Select-String -Path $CmakeCacheFile -Pattern 'GGML_CUDA:BOOL=ON' -Quiet
|
||||
if ($HasNvidiaSmi -and -not $cachedCuda) {
|
||||
Write-Host " Existing llama-server is CPU-only but GPU is available -- rebuilding" -ForegroundColor Yellow
|
||||
$NeedRebuild = $true
|
||||
} elseif (-not $HasNvidiaSmi -and $cachedCuda) {
|
||||
Write-Host " Existing llama-server was built with CUDA but no GPU detected -- rebuilding" -ForegroundColor Yellow
|
||||
$NeedRebuild = $true
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if ((Test-Path $LlamaServerBin) -and -not $NeedRebuild) {
|
||||
Write-Host ""
|
||||
Write-Host "[OK] llama-server already exists at $LlamaServerBin" -ForegroundColor Green
|
||||
} elseif (-not $HasCmakeForBuild) {
|
||||
Write-Host ""
|
||||
if (-not $HasNvidiaSmi) {
|
||||
# CPU-only machines depend entirely on llama-server for GGUF chat -- cmake is required
|
||||
Write-Host "[ERROR] CMake is required to build llama-server for GGUF chat mode." -ForegroundColor Red
|
||||
Write-Host " Install CMake from https://cmake.org/download/ and re-run setup." -ForegroundColor Yellow
|
||||
exit 1
|
||||
}
|
||||
Write-Host "[SKIP] llama-server build -- cmake not available" -ForegroundColor Yellow
|
||||
Write-Host " GGUF inference and export will not be available." -ForegroundColor Yellow
|
||||
Write-Host " Install CMake from https://cmake.org/download/ and re-run setup." -ForegroundColor Yellow
|
||||
} else {
|
||||
Write-Host ""
|
||||
Write-Host "Building llama.cpp with CUDA support..." -ForegroundColor Cyan
|
||||
if ($HasNvidiaSmi) {
|
||||
Write-Host "Building llama.cpp with CUDA support..." -ForegroundColor Cyan
|
||||
} else {
|
||||
Write-Host "Building llama.cpp (CPU-only, no NVIDIA GPU detected)..." -ForegroundColor Cyan
|
||||
}
|
||||
Write-Host " This typically takes 5-10 minutes on first build." -ForegroundColor Gray
|
||||
Write-Host ""
|
||||
|
||||
|
|
@ -1007,17 +1196,19 @@ if (Test-Path $LlamaServerBin) {
|
|||
# Re-sanitize CUDA_PATH_V* vars — Refresh-Environment (called during
|
||||
# Node/Python installs above) may have repopulated conflicting versioned
|
||||
# vars from the Machine registry.
|
||||
$cudaPathVars2 = @([Environment]::GetEnvironmentVariables('Process').Keys | Where-Object { $_ -match '^CUDA_PATH_V' })
|
||||
foreach ($v2 in $cudaPathVars2) {
|
||||
[Environment]::SetEnvironmentVariable($v2, $null, 'Process')
|
||||
if ($HasNvidiaSmi -and $CudaToolkitRoot) {
|
||||
$cudaPathVars2 = @([Environment]::GetEnvironmentVariables('Process').Keys | Where-Object { $_ -match '^CUDA_PATH_V' })
|
||||
foreach ($v2 in $cudaPathVars2) {
|
||||
[Environment]::SetEnvironmentVariable($v2, $null, 'Process')
|
||||
}
|
||||
$tkDirName2 = Split-Path $CudaToolkitRoot -Leaf
|
||||
if ($tkDirName2 -match '^v(\d+)\.(\d+)') {
|
||||
[Environment]::SetEnvironmentVariable("CUDA_PATH_V$($Matches[1])_$($Matches[2])", $CudaToolkitRoot, 'Process')
|
||||
}
|
||||
# Also re-assert CUDA_PATH and CudaToolkitDir in case they were overwritten
|
||||
[Environment]::SetEnvironmentVariable('CUDA_PATH', $CudaToolkitRoot, 'Process')
|
||||
[Environment]::SetEnvironmentVariable('CudaToolkitDir', "$CudaToolkitRoot\", 'Process')
|
||||
}
|
||||
$tkDirName2 = Split-Path $CudaToolkitRoot -Leaf
|
||||
if ($tkDirName2 -match '^v(\d+)\.(\d+)') {
|
||||
[Environment]::SetEnvironmentVariable("CUDA_PATH_V$($Matches[1])_$($Matches[2])", $CudaToolkitRoot, 'Process')
|
||||
}
|
||||
# Also re-assert CUDA_PATH and CudaToolkitDir in case they were overwritten
|
||||
[Environment]::SetEnvironmentVariable('CUDA_PATH', $CudaToolkitRoot, 'Process')
|
||||
[Environment]::SetEnvironmentVariable('CudaToolkitDir', "$CudaToolkitRoot\", 'Process')
|
||||
|
||||
# -- Step A: Clone or pull llama.cpp --
|
||||
|
||||
|
|
@ -1037,7 +1228,14 @@ if (Test-Path $LlamaServerBin) {
|
|||
}
|
||||
}
|
||||
|
||||
# -- Step B: cmake configure (CUDA + Unsloth flags) --
|
||||
# -- Step B: cmake configure --
|
||||
# Clean stale CMake cache to prevent previous CUDA settings from leaking
|
||||
# into a CPU-only rebuild (or vice versa).
|
||||
$CmakeCacheFile = Join-Path $BuildDir "CMakeCache.txt"
|
||||
if (Test-Path $CmakeCacheFile) {
|
||||
Remove-Item -Recurse -Force $BuildDir
|
||||
}
|
||||
|
||||
if ($BuildOk) {
|
||||
Write-Host ""
|
||||
Write-Host "--- cmake configure ---" -ForegroundColor Cyan
|
||||
|
|
@ -1066,37 +1264,45 @@ if (Test-Path $LlamaServerBin) {
|
|||
$CmakeArgs += '-DLLAMA_CURL=OFF'
|
||||
}
|
||||
$CmakeArgs += '-DCMAKE_EXE_LINKER_FLAGS=/NODEFAULTLIB:LIBCMT'
|
||||
# CUDA flags (Unsloth-aligned)
|
||||
$CmakeArgs += '-DGGML_CUDA=ON'
|
||||
$CmakeArgs += "-DCUDAToolkit_ROOT=$CudaToolkitRoot"
|
||||
$CmakeArgs += "-DCUDA_TOOLKIT_ROOT_DIR=$CudaToolkitRoot"
|
||||
$CmakeArgs += "-DCMAKE_CUDA_COMPILER=$NvccPath"
|
||||
$CmakeArgs += '-DGGML_CUDA_FA_ALL_QUANTS=ON'
|
||||
$CmakeArgs += '-DGGML_CUDA_F16=OFF'
|
||||
$CmakeArgs += '-DGGML_CUDA_GRAPHS=OFF'
|
||||
$CmakeArgs += '-DGGML_CUDA_FORCE_CUBLAS=OFF'
|
||||
$CmakeArgs += '-DGGML_CUDA_PEER_MAX_BATCH_SIZE=8192'
|
||||
if ($CudaArch) {
|
||||
# Validate nvcc actually supports this architecture
|
||||
if (Test-NvccArchSupport -NvccExe $NvccPath -Arch $CudaArch) {
|
||||
$CmakeArgs += "-DCMAKE_CUDA_ARCHITECTURES=$CudaArch"
|
||||
} else {
|
||||
# GPU arch too new for this toolkit — fall back to highest supported.
|
||||
# PTX forward-compatibility will JIT-compile for the actual GPU at runtime.
|
||||
$maxArch = Get-NvccMaxArch -NvccExe $NvccPath
|
||||
if ($maxArch) {
|
||||
$CmakeArgs += "-DCMAKE_CUDA_ARCHITECTURES=$maxArch"
|
||||
Write-Host " [WARN] GPU is sm_$CudaArch but nvcc only supports up to sm_$maxArch" -ForegroundColor Yellow
|
||||
Write-Host " Building with sm_$maxArch (PTX will JIT for your GPU at runtime)" -ForegroundColor Yellow
|
||||
# CUDA flags -- only if GPU available, otherwise explicitly disable
|
||||
if ($HasNvidiaSmi -and $NvccPath) {
|
||||
$CmakeArgs += '-DGGML_CUDA=ON'
|
||||
$CmakeArgs += "-DCUDAToolkit_ROOT=$CudaToolkitRoot"
|
||||
$CmakeArgs += "-DCUDA_TOOLKIT_ROOT_DIR=$CudaToolkitRoot"
|
||||
$CmakeArgs += "-DCMAKE_CUDA_COMPILER=$NvccPath"
|
||||
if ($CudaArch) {
|
||||
# Validate nvcc actually supports this architecture
|
||||
if (Test-NvccArchSupport -NvccExe $NvccPath -Arch $CudaArch) {
|
||||
$CmakeArgs += "-DCMAKE_CUDA_ARCHITECTURES=$CudaArch"
|
||||
} else {
|
||||
# GPU arch too new for this toolkit -- fall back to highest supported.
|
||||
# PTX forward-compatibility will JIT-compile for the actual GPU at runtime.
|
||||
$maxArch = Get-NvccMaxArch -NvccExe $NvccPath
|
||||
if ($maxArch) {
|
||||
$CmakeArgs += "-DCMAKE_CUDA_ARCHITECTURES=$maxArch"
|
||||
Write-Host " [WARN] GPU is sm_$CudaArch but nvcc only supports up to sm_$maxArch" -ForegroundColor Yellow
|
||||
Write-Host " Building with sm_$maxArch (PTX will JIT for your GPU at runtime)" -ForegroundColor Yellow
|
||||
}
|
||||
# else: omit flag entirely, let cmake pick defaults
|
||||
}
|
||||
# else: omit flag entirely, let cmake pick defaults
|
||||
}
|
||||
} else {
|
||||
$CmakeArgs += '-DGGML_CUDA=OFF'
|
||||
}
|
||||
|
||||
cmake @CmakeArgs 2>&1 | Out-Null
|
||||
$cmakeOutput = cmake @CmakeArgs 2>&1 | Out-String
|
||||
if ($LASTEXITCODE -ne 0) {
|
||||
$BuildOk = $false
|
||||
$FailedStep = "cmake configure"
|
||||
Write-Host $cmakeOutput -ForegroundColor Red
|
||||
if ($cmakeOutput -match 'No CUDA toolset found|CUDA_TOOLKIT_ROOT_DIR|nvcc') {
|
||||
Write-Host ""
|
||||
Write-Host " Hint: CUDA VS integration may be missing. Try running as admin:" -ForegroundColor Yellow
|
||||
Write-Host " Copy contents of:" -ForegroundColor Yellow
|
||||
Write-Host " <CUDA_PATH>\extras\visual_studio_integration\MSBuildExtensions" -ForegroundColor Yellow
|
||||
Write-Host " into:" -ForegroundColor Yellow
|
||||
Write-Host " <VS_PATH>\MSBuild\Microsoft\VC\v170\BuildCustomizations" -ForegroundColor Yellow
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
|
|
@ -1110,10 +1316,11 @@ if (Test-Path $LlamaServerBin) {
|
|||
Write-Host " Parallel jobs: $NumCpu" -ForegroundColor Gray
|
||||
Write-Host ""
|
||||
|
||||
cmake --build $BuildDir --config Release --target llama-server -j $NumCpu 2>&1 | Out-Null
|
||||
$output = cmake --build $BuildDir --config Release --target llama-server -j $NumCpu 2>&1 | Out-String
|
||||
if ($LASTEXITCODE -ne 0) {
|
||||
$BuildOk = $false
|
||||
$FailedStep = "cmake build (llama-server)"
|
||||
Write-Host $output -ForegroundColor Red
|
||||
}
|
||||
}
|
||||
|
||||
|
|
@ -1121,9 +1328,10 @@ if (Test-Path $LlamaServerBin) {
|
|||
if ($BuildOk) {
|
||||
Write-Host ""
|
||||
Write-Host "--- cmake build (llama-quantize) ---" -ForegroundColor Cyan
|
||||
cmake --build $BuildDir --config Release --target llama-quantize -j $NumCpu 2>&1 | Out-Null
|
||||
$output = cmake --build $BuildDir --config Release --target llama-quantize -j $NumCpu 2>&1 | Out-String
|
||||
if ($LASTEXITCODE -ne 0) {
|
||||
Write-Host " [WARN] llama-quantize build failed (GGUF export may be unavailable)" -ForegroundColor Yellow
|
||||
Write-Host $output -ForegroundColor Yellow
|
||||
}
|
||||
}
|
||||
|
||||
|
|
|
|||
144
studio/setup.sh
144
studio/setup.sh
|
|
@ -41,13 +41,27 @@ if [[ "$keynames" == *$'\nCOLAB_'* ]]; then
|
|||
fi
|
||||
|
||||
# ── Detect whether frontend needs building ──
|
||||
# Only skip when BOTH conditions are true:
|
||||
# 1. We're inside site-packages (PyPI / pip install, not editable)
|
||||
# 2. dist/ already exists (pre-built in the wheel)
|
||||
# Otherwise always (re)build — handles upgrades, editable installs, and
|
||||
# pip-from-source where dist/ was never built.
|
||||
if [[ "$SCRIPT_DIR" == */site-packages/* ]] && [ -d "$SCRIPT_DIR/frontend/dist" ]; then
|
||||
echo "✅ Frontend pre-built (PyPI) — skipping Node/npm check."
|
||||
# Skip if dist/ exists AND no tracked input is newer than dist/.
|
||||
# Checks top-level config/entry files and src/, public/ recursively.
|
||||
# This handles: PyPI installs (dist/ bundled), repeat runs (no changes),
|
||||
# and upgrades/pulls (source newer than dist/ triggers rebuild).
|
||||
_NEED_FRONTEND_BUILD=true
|
||||
if [ -d "$SCRIPT_DIR/frontend/dist" ]; then
|
||||
# Check all top-level files (package.json, bun.lock, vite.config.ts, index.html, etc.)
|
||||
_changed=$(find "$SCRIPT_DIR/frontend" -maxdepth 1 -type f \
|
||||
-newer "$SCRIPT_DIR/frontend/dist" -print -quit 2>/dev/null)
|
||||
# Check src/ and public/ recursively (|| true guards against set -e when dirs are missing)
|
||||
if [ -z "$_changed" ]; then
|
||||
_changed=$(find "$SCRIPT_DIR/frontend/src" "$SCRIPT_DIR/frontend/public" \
|
||||
-type f -newer "$SCRIPT_DIR/frontend/dist" -print -quit 2>/dev/null) || true
|
||||
fi
|
||||
if [ -z "$_changed" ]; then
|
||||
_NEED_FRONTEND_BUILD=false
|
||||
fi
|
||||
fi
|
||||
_NEED_FRONTEND_BUILD=true
|
||||
if [ "$_NEED_FRONTEND_BUILD" = false ]; then
|
||||
echo "✅ Frontend already built and up to date -- skipping Node/npm check."
|
||||
else
|
||||
NEED_NODE=true
|
||||
if command -v node &>/dev/null && command -v npm &>/dev/null; then
|
||||
|
|
@ -146,12 +160,17 @@ run_quiet "npm run build" npm run build
|
|||
|
||||
_restore_gitignores
|
||||
trap - EXIT
|
||||
cd "$SCRIPT_DIR/backend/core/data_recipe/oxc-validator"
|
||||
run_quiet "npm install (oxc validator runtime)" npm install
|
||||
cd "$SCRIPT_DIR"
|
||||
echo "✅ Frontend built to frontend/dist"
|
||||
|
||||
fi # end frontend dist check
|
||||
fi # end frontend build check
|
||||
|
||||
# ── oxc-validator runtime (needs npm -- skip if not available) ──
|
||||
if [ -d "$SCRIPT_DIR/backend/core/data_recipe/oxc-validator" ] && command -v npm &>/dev/null; then
|
||||
cd "$SCRIPT_DIR/backend/core/data_recipe/oxc-validator"
|
||||
run_quiet "npm install (oxc validator runtime)" npm install
|
||||
cd "$SCRIPT_DIR"
|
||||
fi
|
||||
|
||||
# ── 6. Python venv + deps ──
|
||||
|
||||
|
|
@ -223,50 +242,76 @@ install_python_stack() {
|
|||
python "$SCRIPT_DIR/install_python_stack.py"
|
||||
}
|
||||
|
||||
if [ "$IS_COLAB" = true ]; then
|
||||
# Colab: install packages directly without venv
|
||||
install_python_stack
|
||||
else
|
||||
# Local: create venv under ~/.unsloth/studio/ (shared location, not in repo)
|
||||
STUDIO_HOME="$HOME/.unsloth/studio"
|
||||
VENV_DIR="$STUDIO_HOME/.venv"
|
||||
VENV_T5_DIR="$STUDIO_HOME/.venv_t5"
|
||||
mkdir -p "$STUDIO_HOME"
|
||||
# Create venv under ~/.unsloth/studio/ (shared location, not in repo).
|
||||
# All platforms (including Colab) use the same isolated venv so that
|
||||
# studio dependencies are never installed into the system Python.
|
||||
STUDIO_HOME="$HOME/.unsloth/studio"
|
||||
VENV_DIR="$STUDIO_HOME/.venv"
|
||||
VENV_T5_DIR="$STUDIO_HOME/.venv_t5"
|
||||
mkdir -p "$STUDIO_HOME"
|
||||
|
||||
# Clean up legacy in-repo venvs if they exist
|
||||
[ -d "$REPO_ROOT/.venv" ] && rm -rf "$REPO_ROOT/.venv"
|
||||
[ -d "$REPO_ROOT/.venv_overlay" ] && rm -rf "$REPO_ROOT/.venv_overlay"
|
||||
[ -d "$REPO_ROOT/.venv_t5" ] && rm -rf "$REPO_ROOT/.venv_t5"
|
||||
# Clean up legacy in-repo venvs if they exist
|
||||
[ -d "$REPO_ROOT/.venv" ] && rm -rf "$REPO_ROOT/.venv"
|
||||
[ -d "$REPO_ROOT/.venv_overlay" ] && rm -rf "$REPO_ROOT/.venv_overlay"
|
||||
[ -d "$REPO_ROOT/.venv_t5" ] && rm -rf "$REPO_ROOT/.venv_t5"
|
||||
|
||||
rm -rf "$VENV_DIR"
|
||||
rm -rf "$VENV_T5_DIR"
|
||||
"$BEST_PY" -m venv "$VENV_DIR"
|
||||
rm -rf "$VENV_DIR"
|
||||
rm -rf "$VENV_T5_DIR"
|
||||
# Try creating venv with pip; fall back to --without-pip + bootstrap
|
||||
# (some environments like Colab have broken ensurepip)
|
||||
if ! "$BEST_PY" -m venv "$VENV_DIR" 2>/dev/null; then
|
||||
"$BEST_PY" -m venv --without-pip "$VENV_DIR"
|
||||
source "$VENV_DIR/bin/activate"
|
||||
cd "$SCRIPT_DIR"
|
||||
install_python_stack
|
||||
curl -sS https://bootstrap.pypa.io/get-pip.py | python > /dev/null
|
||||
else
|
||||
source "$VENV_DIR/bin/activate"
|
||||
fi
|
||||
|
||||
# ── 6b. Pre-install transformers 5.x into .venv_t5/ ──
|
||||
# Models like GLM-4.7-Flash need transformers>=5.3.0. Instead of pip-installing
|
||||
# at runtime (slow, ~10-15s), we pre-install into a separate directory.
|
||||
# The training subprocess just prepends .venv_t5/ to sys.path — instant switch.
|
||||
echo ""
|
||||
echo " Pre-installing transformers 5.x for newer model support..."
|
||||
mkdir -p "$VENV_T5_DIR"
|
||||
run_quiet "pip install transformers 5.x" pip install --target "$VENV_T5_DIR" --no-deps "transformers==5.3.0"
|
||||
run_quiet "pip install huggingface_hub for t5" pip install --target "$VENV_T5_DIR" --no-deps "huggingface_hub==1.3.0"
|
||||
echo "✅ Transformers 5.x pre-installed to $VENV_T5_DIR/"
|
||||
# ── Ensure uv is available (much faster than pip) ──
|
||||
USE_UV=false
|
||||
if command -v uv &>/dev/null; then
|
||||
USE_UV=true
|
||||
elif curl -LsSf https://astral.sh/uv/install.sh | sh > /dev/null 2>&1; then
|
||||
export PATH="$HOME/.local/bin:$PATH"
|
||||
command -v uv &>/dev/null && USE_UV=true
|
||||
fi
|
||||
|
||||
# ── 7. WSL: pre-install GGUF build dependencies ──
|
||||
# On WSL, sudo requires a password and can't be entered during GGUF export
|
||||
# (runs in a non-interactive subprocess). Install build deps here instead.
|
||||
if grep -qi microsoft /proc/version 2>/dev/null; then
|
||||
echo ""
|
||||
echo "⚠️ WSL detected — installing build dependencies for GGUF export..."
|
||||
echo " You may be prompted for your password."
|
||||
sudo apt-get update -y
|
||||
sudo apt-get install -y build-essential cmake curl git libcurl4-openssl-dev
|
||||
echo "✅ GGUF build dependencies installed"
|
||||
# Helper: install a package, preferring uv with pip fallback
|
||||
fast_install() {
|
||||
if [ "$USE_UV" = true ]; then
|
||||
uv pip install --python "$(command -v python)" "$@" && return 0
|
||||
fi
|
||||
python -m pip install "$@"
|
||||
}
|
||||
|
||||
cd "$SCRIPT_DIR"
|
||||
install_python_stack
|
||||
|
||||
# ── 6b. Pre-install transformers 5.x into .venv_t5/ ──
|
||||
# Models like GLM-4.7-Flash need transformers>=5.3.0. Instead of pip-installing
|
||||
# at runtime (slow, ~10-15s), we pre-install into a separate directory.
|
||||
# The training subprocess just prepends .venv_t5/ to sys.path -- instant switch.
|
||||
echo ""
|
||||
echo " Pre-installing transformers 5.x for newer model support..."
|
||||
mkdir -p "$VENV_T5_DIR"
|
||||
run_quiet "install transformers 5.x" fast_install --target "$VENV_T5_DIR" --no-deps "transformers==5.3.0"
|
||||
run_quiet "install huggingface_hub for t5" fast_install --target "$VENV_T5_DIR" --no-deps "huggingface_hub==1.7.1"
|
||||
run_quiet "install hf_xet for t5" fast_install --target "$VENV_T5_DIR" --no-deps "hf_xet==1.4.2"
|
||||
# tiktoken is needed by Qwen-family tokenizers. Install with deps since
|
||||
# regex/requests may be missing on Windows.
|
||||
run_quiet "install tiktoken for t5" fast_install --target "$VENV_T5_DIR" "tiktoken"
|
||||
echo "✅ Transformers 5.x pre-installed to $VENV_T5_DIR/"
|
||||
|
||||
# ── 7. WSL: pre-install GGUF build dependencies ──
|
||||
# On WSL, sudo requires a password and can't be entered during GGUF export
|
||||
# (runs in a non-interactive subprocess). Install build deps here instead.
|
||||
if grep -qi microsoft /proc/version 2>/dev/null; then
|
||||
echo ""
|
||||
echo "⚠️ WSL detected -- installing build dependencies for GGUF export..."
|
||||
echo " You may be prompted for your password."
|
||||
sudo apt-get update -y
|
||||
sudo apt-get install -y build-essential cmake curl git libcurl4-openssl-dev
|
||||
echo "✅ GGUF build dependencies installed"
|
||||
fi
|
||||
|
||||
# ── 8. Build llama.cpp binaries for GGUF inference + export ──
|
||||
|
|
@ -405,6 +450,9 @@ if [ "$IS_COLAB" = true ]; then
|
|||
echo "╠══════════════════════════════════════╣"
|
||||
echo "║ Unsloth Studio is ready to start ║"
|
||||
echo "║ in your Colab notebook! ║"
|
||||
echo "║ ║"
|
||||
echo "║ from colab import start ║"
|
||||
echo "║ start() ║"
|
||||
echo "╚══════════════════════════════════════╝"
|
||||
else
|
||||
echo "╔══════════════════════════════════════╗"
|
||||
|
|
|
|||
|
|
@ -12,7 +12,7 @@
|
|||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
__version__ = "2026.3.5"
|
||||
__version__ = "2026.3.6"
|
||||
|
||||
__all__ = [
|
||||
"SUPPORTS_BFLOAT16",
|
||||
|
|
@ -242,8 +242,12 @@ def prefer_flex_attn_if_supported(model_class, config):
|
|||
# decode q_len=1, causing ValueError. Needs transformers update.
|
||||
# Gemma3N: timm vision wrappers (eg Gemma3nVisionConfig) do not
|
||||
# support flex_attention.
|
||||
# NemotronH: hybrid Mamba-2 + Transformer model that does not
|
||||
# support flex_attention (raises NotImplementedError from transformers).
|
||||
model_type = getattr(config, "model_type", "") if config else ""
|
||||
if model_type in ("gpt_oss", "mllama") or str(model_type).startswith("gemma3n"):
|
||||
if model_type in ("gpt_oss", "mllama", "nemotron_h") or str(
|
||||
model_type
|
||||
).startswith("gemma3n"):
|
||||
return None
|
||||
if config is not None:
|
||||
setattr(config, "_attn_implementation", "flex_attention")
|
||||
|
|
|
|||
|
|
@ -3,11 +3,11 @@
|
|||
|
||||
import typer
|
||||
|
||||
from cli.commands.train import train
|
||||
from cli.commands.inference import inference
|
||||
from cli.commands.export import export, list_checkpoints
|
||||
from cli.commands.ui import ui
|
||||
from cli.commands.studio import studio_app
|
||||
from unsloth_cli.commands.train import train
|
||||
from unsloth_cli.commands.inference import inference
|
||||
from unsloth_cli.commands.export import export, list_checkpoints
|
||||
from unsloth_cli.commands.ui import ui
|
||||
from unsloth_cli.commands.studio import studio_app
|
||||
|
||||
app = typer.Typer(
|
||||
help = "Command-line interface for Unsloth training, inference, and export.",
|
||||
|
|
@ -14,7 +14,7 @@ studio_app = typer.Typer(help = "Unsloth Studio commands.")
|
|||
|
||||
STUDIO_HOME = Path.home() / ".unsloth" / "studio"
|
||||
|
||||
# __file__ is cli/commands/studio.py — two parents up is the package root
|
||||
# __file__ is unsloth_cli/commands/studio.py -- two parents up is the package root
|
||||
# (either site-packages or the repo root for editable installs).
|
||||
_PACKAGE_ROOT = Path(__file__).resolve().parent.parent.parent
|
||||
|
||||
|
|
@ -33,7 +33,7 @@ def _find_run_py() -> Optional[Path]:
|
|||
|
||||
No CWD dependency — works from any directory.
|
||||
Since studio/ is now a proper package (has __init__.py), it lives in
|
||||
site-packages after pip install, right next to cli/.
|
||||
site-packages after pip install, right next to unsloth_cli/.
|
||||
"""
|
||||
# 1. Relative to __file__ (site-packages or editable repo root)
|
||||
run_py = _PACKAGE_ROOT / "studio" / "backend" / "run.py"
|
||||
|
|
@ -7,8 +7,8 @@ from typing import Optional
|
|||
|
||||
import typer
|
||||
|
||||
from cli.config import Config, load_config
|
||||
from cli.options import add_options_from_config
|
||||
from unsloth_cli.config import Config, load_config
|
||||
from unsloth_cli.options import add_options_from_config
|
||||
|
||||
|
||||
@add_options_from_config(Config)
|
||||
|
|
@ -25,7 +25,11 @@ def ui(
|
|||
),
|
||||
):
|
||||
"""Launch the Unsloth web UI backend server (alias for 'unsloth studio')."""
|
||||
from cli.commands.studio import _studio_venv_python, _find_run_py, STUDIO_HOME
|
||||
from unsloth_cli.commands.studio import (
|
||||
_studio_venv_python,
|
||||
_find_run_py,
|
||||
STUDIO_HOME,
|
||||
)
|
||||
|
||||
# Re-execute in studio venv if available and not already inside it
|
||||
studio_venv_dir = STUDIO_HOME / ".venv"
|
||||
Loading…
Add table
Add a link
Reference in a new issue