diff --git a/common/arg.cpp b/common/arg.cpp index ac0e9b71c..aad251529 100644 --- a/common/arg.cpp +++ b/common/arg.cpp @@ -489,12 +489,15 @@ void common_models_handler_apply(common_models_handler & handler, common_params task.opts = opts; tasks.push_back(task); } + + bool had_spec_url = false; if (!params.speculative.draft.mparams.url.empty()) { common_download_task task; task.url = params.speculative.draft.mparams.url; task.local_path = params.speculative.draft.mparams.path; task.opts = opts; tasks.push_back(task); + had_spec_url = true; } // handle hf_plan tasks @@ -514,6 +517,18 @@ void common_models_handler_apply(common_models_handler & handler, common_params }); } }; + + // handle plan_spec (e.g. --spec-draft-hf) + if (!plan_spec.model_files.empty() && !had_spec_url) { + add_tasks(plan_spec.model_files, plan_spec.primary, params.speculative.draft.mparams); + had_spec_url = true; + } + + // handle vocoder plan (e.g. --hf-repo-v) + if (!plan_voc.model_files.empty()) { + add_tasks(plan_voc.model_files, plan_voc.primary, params.vocoder.model); + } + if (!plan.model_files.empty()) { add_tasks(plan.model_files, plan.primary, params.model); } @@ -522,7 +537,7 @@ void common_models_handler_apply(common_models_handler & handler, common_params params.mmproj.path = hf_cache::finalize_file(plan.mmproj); }); } - if (!plan.mtp.local_path.empty()) { + if (!plan.mtp.local_path.empty() && !had_spec_url) { tasks.emplace_back(plan.mtp, opts, [&]() { // only fall back to the discovered MTP head when no draft was explicitly provided if (params.speculative.draft.mparams.empty()) { @@ -541,16 +556,6 @@ void common_models_handler_apply(common_models_handler & handler, common_params }); } - // handle plan_spec (e.g. --spec-draft-hf) - if (!plan_spec.model_files.empty()) { - add_tasks(plan_spec.model_files, plan_spec.primary, params.speculative.draft.mparams); - } - - // handle vocoder plan (e.g. --hf-repo-v) - if (!plan_voc.model_files.empty()) { - add_tasks(plan_voc.model_files, plan_voc.primary, params.vocoder.model); - } - // run all tasks in parallel if (!params.offline) { // if duplicated files are found, only download once (but still call on_done for each task) @@ -563,6 +568,7 @@ void common_models_handler_apply(common_models_handler & handler, common_params } std::vector unique_tasks_vec; for (auto & pair : unique_tasks) { + LOG_DBG("download task: %s -> %s\n", pair.second->url.c_str(), pair.second->local_path.c_str()); unique_tasks_vec.push_back(*pair.second); } common_download_run_tasks(unique_tasks_vec); @@ -3037,7 +3043,7 @@ common_params_context common_params_parser_init(common_params & params, llama_ex {"--tools"}, "TOOL1,TOOL2,...", "experimental: whether to enable built-in tools for AI agents - do not enable in untrusted environments (default: no tools)\n" "specify \"all\" to enable all tools\n" - "available tools: read_file, file_glob_search, grep_search, exec_shell_command, write_file, edit_file, apply_diff, get_datetime", + "available tools: read_file, file_glob_search, grep_search, exec_shell_command, write_file, edit_file, get_datetime", [](common_params & params, const std::string & value) { params.server_tools = parse_csv_row(value); } diff --git a/ggml/include/ggml-et.h b/ggml/include/ggml-et.h new file mode 100644 index 000000000..8b78f39aa --- /dev/null +++ b/ggml/include/ggml-et.h @@ -0,0 +1,28 @@ +#pragma once + +#include "ggml.h" +#include "ggml-backend.h" + +#ifdef __cplusplus +extern "C" { +#endif + +#define GGML_ET_NAME "ET" + +// backend API +GGML_BACKEND_API ggml_guid_t ggml_backend_et_guid(void); +GGML_BACKEND_API ggml_backend_t ggml_backend_et_init(size_t devidx); + +GGML_BACKEND_API bool ggml_backend_is_et(ggml_backend_t backend); +GGML_BACKEND_API int ggml_backend_et_get_device_count(void); +GGML_BACKEND_API void ggml_backend_et_get_device_description(int devidx, char * description, size_t description_size); +GGML_BACKEND_API void ggml_backend_et_get_device_memory(int devidx, size_t * free, size_t * total); + +GGML_BACKEND_API ggml_backend_buffer_type_t ggml_backend_et_buffer_type(size_t dev_num); +GGML_BACKEND_API ggml_backend_buffer_type_t ggml_backend_et_host_buffer_type(void); + +GGML_BACKEND_API ggml_backend_reg_t ggml_backend_et_reg(void); + +#ifdef __cplusplus +} +#endif diff --git a/ggml/src/ggml-backend-reg.cpp b/ggml/src/ggml-backend-reg.cpp index 516f7ab23..3b80d8c22 100644 --- a/ggml/src/ggml-backend-reg.cpp +++ b/ggml/src/ggml-backend-reg.cpp @@ -86,6 +86,10 @@ #include "ggml-openvino.h" #endif +#ifdef GGML_USE_ET +#include "ggml-et.h" +#endif + namespace fs = std::filesystem; static std::string path_str(const fs::path & path) { @@ -161,6 +165,9 @@ struct ggml_backend_registry { #ifdef GGML_USE_OPENVINO register_backend(ggml_backend_openvino_reg()); #endif +#ifdef GGML_USE_ET + register_backend(ggml_backend_et_reg()); +#endif #ifdef GGML_USE_CPU register_backend(ggml_backend_cpu_reg()); #endif diff --git a/ggml/src/ggml-cpu/arch/arm/quants.c b/ggml/src/ggml-cpu/arch/arm/quants.c index 636d7be12..b988abf99 100644 --- a/ggml/src/ggml-cpu/arch/arm/quants.c +++ b/ggml/src/ggml-cpu/arch/arm/quants.c @@ -263,13 +263,13 @@ void ggml_vec_dot_q2_0_q8_0(int n, float * GGML_RESTRICT s, size_t bs, const voi const uint8x16_t raw16 = vcombine_u8(raw, raw); // First 16 elements: replicate bytes 0-3, shift, mask, subtract 1 - uint8x16_t bytes0 = vqtbl1q_u8(raw16, idx_lo); + uint8x16_t bytes0 = ggml_vqtbl1q_u8(raw16, idx_lo); int8x16_t qv0 = vsubq_s8( vreinterpretq_s8_u8(vandq_u8(vshlq_u8(bytes0, shifts), mask2)), one); // Second 16 elements: replicate bytes 4-7, shift, mask, subtract 1 - uint8x16_t bytes1 = vqtbl1q_u8(raw16, idx_hi); + uint8x16_t bytes1 = ggml_vqtbl1q_u8(raw16, idx_hi); int8x16_t qv1 = vsubq_s8( vreinterpretq_s8_u8(vandq_u8(vshlq_u8(bytes1, shifts), mask2)), one); diff --git a/src/llama-graph.cpp b/src/llama-graph.cpp index 4a1fe0f46..6d215a957 100644 --- a/src/llama-graph.cpp +++ b/src/llama-graph.cpp @@ -646,7 +646,7 @@ static void dsv4_set_kq_mask( return; } - GGML_ASSERT(dst->type == GGML_TYPE_F32); + GGML_ASSERT(dst->type == GGML_TYPE_F32 || dst->type == GGML_TYPE_F16); GGML_ASSERT(n_stream > 0); GGML_ASSERT(n_tokens%n_stream == 0); GGML_ASSERT(dst->ne[0] == plan.n_kv); @@ -656,13 +656,27 @@ static void dsv4_set_kq_mask( GGML_ASSERT((int64_t) plan.n_visible.size() == (int64_t) n_tokens); GGML_ASSERT(ggml_backend_buffer_is_host(dst->buffer)); - float * data = (float *) dst->data; + if (dst->type == GGML_TYPE_F32) { + float * data = (float *) dst->data; - for (int64_t i = 0; i < (int64_t) n_tokens; ++i) { - const int32_t n_visible = plan.n_visible[i]; + for (int64_t i = 0; i < (int64_t) n_tokens; ++i) { + const int32_t n_visible = plan.n_visible[i]; - for (int64_t j = 0; j < dst->ne[0]; ++j) { - data[i*dst->ne[0] + j] = j < n_visible ? 0.0f : -INFINITY; + for (int64_t j = 0; j < dst->ne[0]; ++j) { + data[i*dst->ne[0] + j] = j < n_visible ? 0.0f : -INFINITY; + } + } + } else if (dst->type == GGML_TYPE_F16) { + ggml_fp16_t * data = (ggml_fp16_t *) dst->data; + const ggml_fp16_t fp16_ninf = llama_cast(-INFINITY); + const ggml_fp16_t fp16_zero = llama_cast(0.0f); + + for (int64_t i = 0; i < (int64_t) n_tokens; ++i) { + const int32_t n_visible = plan.n_visible[i]; + + for (int64_t j = 0; j < dst->ne[0]; ++j) { + data[i*dst->ne[0] + j] = j < n_visible ? fp16_zero : fp16_ninf; + } } } } @@ -679,8 +693,7 @@ static ggml_tensor * dsv4_build_raw_kq_mask( GGML_ASSERT(n_stream > 0); GGML_ASSERT(n_tokens%n_stream == 0); - const bool use_fattn = cparams.flash_attn && (!cparams.kv_unified || n_stream == 1); - const auto type = use_fattn ? GGML_TYPE_F16 : GGML_TYPE_F32; + const auto type = cparams.flash_attn ? GGML_TYPE_F16 : GGML_TYPE_F32; ggml_tensor * res = ggml_new_tensor_4d(ctx, type, n_kv, n_tokens/n_stream, 1, n_stream); ggml_set_input(res); @@ -815,6 +828,7 @@ static void dsv4_build_comp_inputs( llm_graph_input_dsv4::comp_input & inp, const llama_kv_cache_dsv4_context::comp_plan & plan, const char * name, + const llama_cparams & cparams, int64_t n_stream) { inp.state_pos = dsv4_build_input_1d(ctx, GGML_TYPE_I32, plan.state_pos.size(), std::string("dsv4_") + name + "_state_pos"); inp.state_persist_src_idxs = dsv4_build_input_1d(ctx, GGML_TYPE_I32, plan.state_persist_src_idxs.size(), std::string("dsv4_") + name + "_state_persist_src_idxs"); @@ -829,7 +843,7 @@ static void dsv4_build_comp_inputs( GGML_ASSERT(n_stream > 0); GGML_ASSERT(n_tokens%n_stream == 0); - inp.kq_mask = ggml_new_tensor_4d(ctx, GGML_TYPE_F32, plan.n_kv, n_tokens/n_stream, 1, n_stream); + inp.kq_mask = ggml_new_tensor_4d(ctx, cparams.flash_attn && strcmp(name, "lid") != 0 ? GGML_TYPE_F16 : GGML_TYPE_F32, plan.n_kv, n_tokens/n_stream, 1, n_stream); ggml_set_input(inp.kq_mask); ggml_set_name(inp.kq_mask, (std::string("dsv4_") + name + "_kq_mask").c_str()); } @@ -3077,9 +3091,9 @@ llm_graph_input_dsv4 * llm_graph_context::build_inp_dsv4() const { inp_raw->self_k_rot = raw_ctx->build_input_k_rot(ctx0); auto inp = std::make_unique(cparams, std::move(inp_raw), mctx_cur); - dsv4_build_comp_inputs(ctx0, inp->inp_csa, mctx_cur->get_csa_plan(ubatch), "csa", n_stream); - dsv4_build_comp_inputs(ctx0, inp->inp_hca, mctx_cur->get_hca_plan(ubatch), "hca", n_stream); - dsv4_build_comp_inputs(ctx0, inp->inp_lid, mctx_cur->get_lid_plan(ubatch), "lid", n_stream); + dsv4_build_comp_inputs(ctx0, inp->inp_csa, mctx_cur->get_csa_plan(ubatch), "csa", cparams, n_stream); + dsv4_build_comp_inputs(ctx0, inp->inp_hca, mctx_cur->get_hca_plan(ubatch), "hca", cparams, n_stream); + dsv4_build_comp_inputs(ctx0, inp->inp_lid, mctx_cur->get_lid_plan(ubatch), "lid", cparams, n_stream); inp->inp_csa.k_rot = mctx_cur->get_csa()->build_input_k_rot(ctx0); inp->inp_hca.k_rot = mctx_cur->get_hca()->build_input_k_rot(ctx0); inp->inp_lid.k_rot = mctx_cur->get_lid()->build_input_k_rot(ctx0); diff --git a/src/llama-model.cpp b/src/llama-model.cpp index de9b7ee89..c074bd891 100644 --- a/src/llama-model.cpp +++ b/src/llama-model.cpp @@ -474,38 +474,38 @@ struct ggml_backend_meta_split_state llama_meta_device_get_split_state(const str const llama_hparams & hparams = ud->model->hparams; const std::string tensor_name = tensor->name; - const std::regex pattern_q_weight ("blk\\.\\d*\\.attn_q.weight"); - const std::regex pattern_kv_weight ("blk\\.\\d*\\.attn_(k|v).weight"); - const std::regex pattern_qkv_weight ("blk\\.\\d*\\.attn_qkv.weight"); - const std::regex pattern_q_bias ("blk\\.\\d*\\.attn_q\\.bias"); - const std::regex pattern_kv_bias ("blk\\.\\d*\\.attn_(k|v)\\.bias"); - const std::regex pattern_qkv_bias ("blk\\.\\d*\\.attn_qkv.bias"); - const std::regex pattern_qk_norm ("blk\\.\\d*\\.attn_(q|k)_norm\\.weight"); - const std::regex pattern_kv_cache ("cache_(k|v)_l\\d*"); - const std::regex pattern_attn_sinks ("blk\\.\\d*\\.attn_sinks.weight"); - const std::regex pattern_attn_out_weight ("blk\\.\\d*\\.attn_output.weight"); - const std::regex pattern_attn_out_bias ("blk\\.\\d*\\.attn_output.bias"); - const std::regex pattern_attn_gate_weight("blk\\.\\d*\\.attn_gate.weight"); + static const std::regex pattern_q_weight ("blk\\.\\d*\\.attn_q.weight"); + static const std::regex pattern_kv_weight ("blk\\.\\d*\\.attn_(k|v).weight"); + static const std::regex pattern_qkv_weight ("blk\\.\\d*\\.attn_qkv.weight"); + static const std::regex pattern_q_bias ("blk\\.\\d*\\.attn_q\\.bias"); + static const std::regex pattern_kv_bias ("blk\\.\\d*\\.attn_(k|v)\\.bias"); + static const std::regex pattern_qkv_bias ("blk\\.\\d*\\.attn_qkv.bias"); + static const std::regex pattern_qk_norm ("blk\\.\\d*\\.attn_(q|k)_norm\\.weight"); + static const std::regex pattern_kv_cache ("cache_(k|v)_l\\d*"); + static const std::regex pattern_attn_sinks ("blk\\.\\d*\\.attn_sinks.weight"); + static const std::regex pattern_attn_out_weight ("blk\\.\\d*\\.attn_output.weight"); + static const std::regex pattern_attn_out_bias ("blk\\.\\d*\\.attn_output.bias"); + static const std::regex pattern_attn_gate_weight("blk\\.\\d*\\.attn_gate.weight"); - const std::regex pattern_ssm_dt ("blk\\.\\d*\\.ssm_dt.bias"); - const std::regex pattern_ssm_a ("blk\\.\\d*\\.ssm_a"); - const std::regex pattern_ssm_alpha ("blk\\.\\d*\\.ssm_alpha.weight"); - const std::regex pattern_ssm_beta ("blk\\.\\d*\\.ssm_beta.weight"); - const std::regex pattern_ssm_beta_alpha ("blk\\.\\d*\\.ssm_ba.weight"); - const std::regex pattern_r_cache ("cache_r_l\\d*"); - const std::regex pattern_s_cache ("cache_s_l\\d*"); - const std::regex pattern_ssm_conv1d ("blk\\.\\d*\\.ssm_conv1d.weight"); - const std::regex pattern_ssm_out_weight ("blk\\.\\d*\\.ssm_out.weight"); + static const std::regex pattern_ssm_dt ("blk\\.\\d*\\.ssm_dt.bias"); + static const std::regex pattern_ssm_a ("blk\\.\\d*\\.ssm_a"); + static const std::regex pattern_ssm_alpha ("blk\\.\\d*\\.ssm_alpha.weight"); + static const std::regex pattern_ssm_beta ("blk\\.\\d*\\.ssm_beta.weight"); + static const std::regex pattern_ssm_beta_alpha ("blk\\.\\d*\\.ssm_ba.weight"); + static const std::regex pattern_r_cache ("cache_r_l\\d*"); + static const std::regex pattern_s_cache ("cache_s_l\\d*"); + static const std::regex pattern_ssm_conv1d ("blk\\.\\d*\\.ssm_conv1d.weight"); + static const std::regex pattern_ssm_out_weight ("blk\\.\\d*\\.ssm_out.weight"); - const std::regex pattern_ffn_up_gate_weight("blk\\.\\d*\\.ffn_(up|gate)(_exps)?.weight"); - const std::regex pattern_ffn_up_gate_bias ("blk\\.\\d*\\.ffn_(up|gate)(_exps)?.bias"); - const std::regex pattern_ffn_gate_up_weight("blk\\.\\d*\\.ffn_gate_up(_exps)?.weight"); - const std::regex pattern_ffn_down_weight ("blk\\.\\d*\\.ffn_down(_exps)?.weight"); - const std::regex pattern_ffn_down_bias ("blk\\.\\d*\\.ffn_down.bias"); - const std::regex pattern_ffn_down_exps_bias("blk\\.\\d*\\.ffn_down_exps.bias"); + static const std::regex pattern_ffn_up_gate_weight("blk\\.\\d*\\.ffn_(up|gate)(_exps)?.weight"); + static const std::regex pattern_ffn_up_gate_bias ("blk\\.\\d*\\.ffn_(up|gate)(_exps)?.bias"); + static const std::regex pattern_ffn_gate_up_weight("blk\\.\\d*\\.ffn_gate_up(_exps)?.weight"); + static const std::regex pattern_ffn_down_weight ("blk\\.\\d*\\.ffn_down(_exps)?.weight"); + static const std::regex pattern_ffn_down_bias ("blk\\.\\d*\\.ffn_down.bias"); + static const std::regex pattern_ffn_down_exps_bias("blk\\.\\d*\\.ffn_down_exps.bias"); - const std::regex pattern_output_weight("output\\.weight"); - const std::regex pattern_output_bias ("output\\.bias"); + static const std::regex pattern_output_weight("output\\.weight"); + static const std::regex pattern_output_bias ("output\\.bias"); struct tensor_config { ggml_backend_meta_split_axis axis; diff --git a/src/models/deepseek4.cpp b/src/models/deepseek4.cpp index 3fb5bff1b..07aa477e1 100644 --- a/src/models/deepseek4.cpp +++ b/src/models/deepseek4.cpp @@ -184,32 +184,6 @@ static ggml_tensor * dsv4_with_zero_dep(ggml_context * ctx, ggml_tensor * t, ggm return ggml_add(ctx, t, zero); } -// Raw SWA K is stored once, but compressed K/masks can carry a stream axis. -// Repeat raw K at graph build time before concatenating raw and compressed K. -static ggml_tensor * dsv4_repeat_streams(ggml_context * ctx, ggml_tensor * t, int64_t n_stream) { - if (t->ne[3] == n_stream) { - return t; - } - - GGML_ASSERT(t->ne[3] == 1); - return ggml_repeat_4d(ctx, t, t->ne[0], t->ne[1], t->ne[2], n_stream); -} - -static ggml_tensor * dsv4_build_kq_zero_bias( - ggml_context * ctx, - const llama_cparams & cparams, - ggml_tensor * kq_mask, - int64_t n_head) { - if (!cparams.kv_unified || !cparams.flash_attn || kq_mask->ne[3] == 1) { - return nullptr; - } - - // Keep multi-stream unified DSV4 on the explicit attention path. - ggml_tensor * res = ggml_new_tensor_4d(ctx, GGML_TYPE_F32, - kq_mask->ne[0], kq_mask->ne[1], n_head, kq_mask->ne[3]); - return ggml_fill(ctx, res, 0.0f); -} - static constexpr int64_t DSV4_CSA_RATIO = 4; static constexpr int64_t DSV4_HCA_RATIO = 128; @@ -624,7 +598,7 @@ ggml_tensor * llama_model_deepseek4::graph::build_top_k_mask( ggml_tensor * top_k_3d = ggml_view_4d(ctx0, top_k, top_k->ne[0], top_k->ne[1], top_k->ne[3], 1, top_k->nb[1], top_k->nb[2], top_k->ne[3]*top_k->nb[3], 0); - ggml_tensor * zeros = ggml_new_tensor_4d(ctx0, GGML_TYPE_F32, 1, top_k_3d->ne[0], top_k_3d->ne[1], top_k_3d->ne[2]); + ggml_tensor * zeros = ggml_new_tensor_4d(ctx0, cparams.flash_attn ? GGML_TYPE_F16 : GGML_TYPE_F32, 1, top_k_3d->ne[0], top_k_3d->ne[1], top_k_3d->ne[2]); zeros = ggml_fill(ctx0, zeros, 0.0f); ggml_tensor * kq_mask_top_k = ggml_set_rows(ctx0, kq_mask_all, zeros, top_k_3d); @@ -681,26 +655,16 @@ ggml_tensor * llama_model_deepseek4::graph::build_csa_lid_attention( csa_k->nb[1], csa_k->nb[2], csa_k->nb[3], 0); cb(csa_k, "csa_comp_k", il); - raw_k = dsv4_repeat_streams(ctx0, raw_k, csa_k->ne[3]); - ggml_tensor * k_all = ggml_concat(ctx0, raw_k, csa_k, 2); cb(k_all, "csa_k_all", il); ggml_tensor * raw_mask = inp_attn->get_kq_mask(); ggml_tensor * csa_mask = build_top_k_mask(inp_csa.kq_mask, top_k, "csa_top_k_mask", il); - const bool use_fattn = cparams.flash_attn && (!cparams.kv_unified || csa_mask->ne[3] == 1); - if (use_fattn && csa_mask->type != GGML_TYPE_F16) { - csa_mask = ggml_cast(ctx0, csa_mask, GGML_TYPE_F16); - } - if (raw_mask->type != csa_mask->type) { - raw_mask = ggml_cast(ctx0, raw_mask, csa_mask->type); - } ggml_tensor * kq_mask = ggml_concat(ctx0, raw_mask, csa_mask, 0); cb(kq_mask, "csa_lid_kq_mask", il); - ggml_tensor * kq_b = dsv4_build_kq_zero_bias(ctx0, cparams, kq_mask, q->ne[1]); - ggml_tensor * out = build_attn_mha(q, k_all, k_all, kq_b, kq_mask, sinks, nullptr, kq_scale, il); + ggml_tensor * out = build_attn_mha(q, k_all, k_all, nullptr, kq_mask, sinks, nullptr, kq_scale, il); if (k_rot) { out = llama_mul_mat_hadamard(ctx0, out, k_rot); } @@ -746,26 +710,16 @@ ggml_tensor * llama_model_deepseek4::graph::build_hca_attention( hca_k->nb[1], hca_k->nb[2], hca_k->nb[3], 0); cb(hca_k, "hca_comp_k", il); - raw_k = dsv4_repeat_streams(ctx0, raw_k, hca_k->ne[3]); - ggml_tensor * k_all = ggml_concat(ctx0, raw_k, hca_k, 2); cb(k_all, "hca_k_all", il); ggml_tensor * raw_mask = inp_attn->get_kq_mask(); ggml_tensor * hca_mask = inp_hca.kq_mask; - const bool use_fattn = cparams.flash_attn && (!cparams.kv_unified || hca_mask->ne[3] == 1); - if (use_fattn && hca_mask->type != GGML_TYPE_F16) { - hca_mask = ggml_cast(ctx0, hca_mask, GGML_TYPE_F16); - } - if (raw_mask->type != hca_mask->type) { - raw_mask = ggml_cast(ctx0, raw_mask, hca_mask->type); - } ggml_tensor * kq_mask = ggml_concat(ctx0, raw_mask, hca_mask, 0); cb(kq_mask, "hca_kq_mask", il); - ggml_tensor * kq_b = dsv4_build_kq_zero_bias(ctx0, cparams, kq_mask, q->ne[1]); - ggml_tensor * out = build_attn_mha(q, k_all, k_all, kq_b, kq_mask, sinks, nullptr, kq_scale, il); + ggml_tensor * out = build_attn_mha(q, k_all, k_all, nullptr, kq_mask, sinks, nullptr, kq_scale, il); if (k_rot) { out = llama_mul_mat_hadamard(ctx0, out, k_rot); } @@ -800,10 +754,8 @@ ggml_tensor * llama_model_deepseek4::graph::build_raw_attention( ggml_tensor * kq_mask = inp_attn->get_kq_mask(); ggml_tensor * k = mctx_cur->get_k(ctx0, il); - k = dsv4_repeat_streams(ctx0, k, kq_mask->ne[3]); - ggml_tensor * kq_b = dsv4_build_kq_zero_bias(ctx0, cparams, kq_mask, q->ne[1]); - ggml_tensor * out = build_attn_mha(q, k, k, kq_b, kq_mask, sinks, nullptr, kq_scale, il); + ggml_tensor * out = build_attn_mha(q, k, k, nullptr, kq_mask, sinks, nullptr, kq_scale, il); if (k_rot) { out = llama_mul_mat_hadamard(ctx0, out, k_rot); } diff --git a/tools/cli/cli-context.cpp b/tools/cli/cli-context.cpp index 74d60eb19..0de8f6902 100644 --- a/tools/cli/cli-context.cpp +++ b/tools/cli/cli-context.cpp @@ -153,9 +153,19 @@ bool cli_context::init() { if (use_external_server) { spinner.reset(); - if (!list_and_ask_models()) { + try { + if (!list_and_ask_models()) { + return false; + } + } catch (const json::parse_error & e) { + ui::show_error(e.what()); + ui::show_message("This might be caused by an incorrect server-base endpoint URL"); + return false; + } catch (const std::exception & e) { + ui::show_error(e.what()); return false; } + // restore the spinner for the next step spinner.emplace("Waiting for server..."); } diff --git a/tools/mtmd/clip-graph.h b/tools/mtmd/clip-graph.h index c84b32880..a95de20a3 100644 --- a/tools/mtmd/clip-graph.h +++ b/tools/mtmd/clip-graph.h @@ -20,8 +20,8 @@ struct clip_graph { const clip_hparams & hparams; projector_type proj_type; - // we only support single image per batch - const clip_image_f32 & img; + const clip_image_f32 & img; // for backward compat + const clip_image_f32_batch * img_batch = nullptr; const int patch_size; const int n_patches_x; @@ -63,6 +63,12 @@ struct clip_graph { // void cb(ggml_tensor * cur0, const char * name, int il) const; + const clip_image_f32 & get_img(size_t idx) const { + GGML_ASSERT(img_batch); + GGML_ASSERT(idx < img_batch->entries.size()); + return img_batch->entries[idx]; + } + // siglip2 naflex ggml_tensor * resize_position_embeddings(uint32_t interpolation_mode = DEFAULT_INTERPOLATION_MODE); diff --git a/tools/mtmd/clip-model.h b/tools/mtmd/clip-model.h index 46be39a64..6d4336c40 100644 --- a/tools/mtmd/clip-model.h +++ b/tools/mtmd/clip-model.h @@ -69,6 +69,7 @@ struct clip_hparams { std::vector image_res_candidates; int32_t preproc_min_tiles = 0; int32_t preproc_max_tiles = 0; + int32_t preproc_tile_size = 0; // local tile size (deepseek-ocr) resize_algo image_resize_algo_rf = RESIZE_ALGO_BICUBIC; resize_algo image_resize_algo_ov = RESIZE_ALGO_BILINEAR; pad_style image_pad_rf = PAD_CEIL; // padding style for the refined image (e.g. llava-1.6) diff --git a/tools/mtmd/clip.cpp b/tools/mtmd/clip.cpp index 606c03a6d..1ca5f7e20 100644 --- a/tools/mtmd/clip.cpp +++ b/tools/mtmd/clip.cpp @@ -1072,6 +1072,8 @@ static std::unique_ptr clip_get_graph_builder(clip_ctx * ctx, const GGML_ABORT("missing cgraph builder"); } + builder->img_batch = &imgs; + // TODO [QWEN_VIDEO]: improve this in the future builder->n_batch = imgs.entries.size(); @@ -1656,7 +1658,16 @@ struct clip_model_loader { get_u32(KEY_SAM_N_HEAD, hparams.sam_n_head, true); get_u32(KEY_SAM_N_EMBD, hparams.sam_n_embd, true); get_u32(KEY_ATTN_WINDOW_SIZE, hparams.attn_window_size, true); + hparams.preproc_min_tiles = 2; + if (model.proj_type == PROJECTOR_TYPE_DEEPSEEKOCR) { + hparams.preproc_max_tiles = 9; + hparams.preproc_tile_size = 640; + // the CLIP/ViT body runs its layernorms at 1e-5 (the SAM stage uses 1e-6) + hparams.eps = 1e-5f; + } if (model.proj_type == PROJECTOR_TYPE_DEEPSEEKOCR2) { + hparams.preproc_max_tiles = 6; + hparams.preproc_tile_size = 768; // qwen2 encoder is GQA, requires KEY_N_HEAD_KV get_u32(string_format(KEY_N_HEAD_KV, "vision"), hparams.n_head_kv); } @@ -3332,6 +3343,9 @@ int clip_n_output_tokens_x(const clip_ctx * ctx, const clip_image_f32 * img) { return (img->nx() / params.patch_size) / 2; case PROJECTOR_TYPE_STEP3VL: return img->nx() / (params.patch_size * params.n_merge); + case PROJECTOR_TYPE_DEEPSEEKOCR: + case PROJECTOR_TYPE_DEEPSEEKOCR2: + return (img->nx() / params.patch_size) / 4; default: break; } @@ -3541,10 +3555,17 @@ int clip_n_output_tokens(const clip_ctx * ctx, const clip_image_f32 * img) { // E.g., 64x64 -> 16x16 patches n_patches /= 16; - // build_global_local_features adds image newlines and view separator - // Formula: h*(w+1) + 1 where h = w = sqrt(n_patches) - int h = static_cast(std::sqrt(static_cast(n_patches))); - n_patches = h * (h + 1) + 1; + if (img->add_viewsep) { + // global view: one image-newline per token-row + trailing view separator + const int h = static_cast(std::sqrt(static_cast(n_patches))); + n_patches = h * (h + 1) + 1; + } else if (img->ny() >= img->nx() && img->ny() % img->nx() == 0) { + // tile row: one image-newline per token-row + const int grid_w = img->ny() / img->nx(); + const int tile_patches = img->nx() / (patch_size * 4); // patches per tile side (SAM divides by 4) + const int h = tile_patches; + n_patches = (tile_patches * grid_w + 1) * h; + } } break; case PROJECTOR_TYPE_HUNYUANVL: { @@ -4184,7 +4205,10 @@ bool clip_image_batch_encode(clip_ctx * ctx, int n_threads, const clip_image_f32 case PROJECTOR_TYPE_DEEPSEEKOCR: case PROJECTOR_TYPE_DEEPSEEKOCR2: { - GGML_ASSERT(pos_w == pos_h); + GGML_ASSERT( + (pos_w == pos_h) // overview image + || (pos_h >= pos_w && pos_h % pos_w == 0) // tile images + ); const int window = hparams.attn_window_size; const int pos = pos_w; diff --git a/tools/mtmd/models/deepseekocr.cpp b/tools/mtmd/models/deepseekocr.cpp index c3c22d0a4..b9fea3538 100644 --- a/tools/mtmd/models/deepseekocr.cpp +++ b/tools/mtmd/models/deepseekocr.cpp @@ -96,6 +96,8 @@ ggml_tensor * clip_graph_deepseekocr::build_sam(ggml_tensor * inp_raw) { const int n_heads = hparams.sam_n_head; const int d_heads = n_embd / n_heads; const int window = hparams.attn_window_size; + // SAM stage runs its layernorms at 1e-6 + const float sam_eps = 1e-6f; ggml_tensor * inpL; @@ -134,7 +136,7 @@ ggml_tensor * clip_graph_deepseekocr::build_sam(ggml_tensor * inp_raw) { ggml_tensor * shortcut = cur; // layernorm1 - cur = build_norm(cur, layer.ln_1_w, layer.ln_1_b, NORM_TYPE_NORMAL, eps, il); + cur = build_norm(cur, layer.ln_1_w, layer.ln_1_b, NORM_TYPE_NORMAL, sam_eps, il); const int64_t w0 = cur->ne[1]; const int64_t h0 = cur->ne[2]; @@ -214,7 +216,7 @@ ggml_tensor * clip_graph_deepseekocr::build_sam(ggml_tensor * inp_raw) { ggml_tensor * inpFF = cur; // layernorm2 - cur = build_norm(inpFF, layer.ln_2_w, layer.ln_2_b, NORM_TYPE_NORMAL, eps, il); + cur = build_norm(inpFF, layer.ln_2_w, layer.ln_2_b, NORM_TYPE_NORMAL, sam_eps, il); // ffn cur = build_ffn(cur, layer.ff_up_w, layer.ff_up_b, nullptr, nullptr, layer.ff_down_w, layer.ff_down_b, @@ -229,12 +231,12 @@ ggml_tensor * clip_graph_deepseekocr::build_sam(ggml_tensor * inp_raw) { cur = ggml_conv_2d(ctx0, model.neck_0_w, cur, 1, 1, 0, 0, 1, 1); cur = ggml_cont(ctx0, ggml_permute(ctx0, cur, 1, 2, 0, 3)); - cur = build_norm(cur, model.neck_1_w, model.neck_1_b, NORM_TYPE_NORMAL, hparams.eps, -1); + cur = build_norm(cur, model.neck_1_w, model.neck_1_b, NORM_TYPE_NORMAL, sam_eps, -1); cur = ggml_cont(ctx0, ggml_permute(ctx0, cur, 2, 0, 1, 3)); cur = ggml_conv_2d(ctx0, model.neck_2_w, cur, 1, 1, 1, 1, 1, 1); cur = ggml_cont(ctx0, ggml_permute(ctx0, cur, 1, 2, 0, 3)); - cur = build_norm(cur, model.neck_3_w, model.neck_3_b, NORM_TYPE_NORMAL, hparams.eps, -1); + cur = build_norm(cur, model.neck_3_w, model.neck_3_b, NORM_TYPE_NORMAL, sam_eps, -1); cur = ggml_cont(ctx0, ggml_permute(ctx0, cur, 2, 0, 1, 3)); cur = ggml_conv_2d(ctx0, model.net_2, cur, 2, 2, 1, 1, 1, 1); @@ -248,8 +250,40 @@ ggml_tensor * clip_graph_deepseekocr::build_sam(ggml_tensor * inp_raw) { ggml_cgraph * clip_graph_deepseekocr::build() { // patch embedding ggml_tensor * inp_raw = build_inp_raw(); + + bool is_overview = img.add_viewsep; + int n_tiles_per_row = 0; + + // note: we expect either a batch of rows or a batch of overviews, but not a mix of both + + if (!is_overview) { + // handle the case where we have a batch of rows + // sanity check + for (auto & entry : img_batch->entries) { + if (entry.add_viewsep) { + throw std::runtime_error("DeepSeek-OCR: mixed overview and non-overview images in batch"); + } + if (entry.nx() != img.nx() || entry.ny() != img.ny()) { + throw std::runtime_error("DeepSeek-OCR: mixed image sizes in batch"); + } + } + + GGML_ASSERT(img.ny() >= img.nx()); + GGML_ASSERT(img.ny() % img.nx() == 0); + n_tiles_per_row = img.ny() / img.nx(); + + // input shape: [tile_size, tile_size * n_tiles_per_row, 3] + // we want to reshape it to [tile_size, tile_size, 3, n_tiles_per_row] + inp_raw = ggml_reshape_4d(ctx0, inp_raw, img.nx(), img.nx(), n_tiles_per_row, 3); + inp_raw = ggml_cont(ctx0, ggml_permute(ctx0, inp_raw, 0, 1, 3, 2)); + } + ggml_tensor * sam_out = build_sam(inp_raw); + if (!is_overview) { + n_batch = n_tiles_per_row; + } + const int clip_n_patches = sam_out->ne[0] * sam_out->ne[1]; ggml_tensor * clip_out; @@ -257,7 +291,9 @@ ggml_cgraph * clip_graph_deepseekocr::build() { { ggml_tensor * inp; - inp = ggml_reshape_2d(ctx0, sam_out, clip_n_patches, sam_out->ne[2]); + // sam_out: [patch_h, patch_w, n_embd, n_batch] + // -> [n_embd, clip_n_patches, n_batch] + inp = ggml_reshape_3d(ctx0, sam_out, clip_n_patches, sam_out->ne[2], sam_out->ne[3]); inp = ggml_cont(ctx0, ggml_permute(ctx0, inp, 1, 0, 2, 3)); ggml_tensor * new_pos_embd = model.position_embeddings; @@ -281,8 +317,11 @@ ggml_cgraph * clip_graph_deepseekocr::build() { n_pos = tgt_size * tgt_size + 1; } - // add CLS token - inp = ggml_concat(ctx0, model.class_embedding, inp, 1); + // add CLS token per batch item + // inp: [n_embd, clip_n_patches, n_batch] + // class_embedding: [n_embd] -> [n_embd, 1, n_batch] + ggml_tensor * cls_embd = ggml_repeat_4d(ctx0, model.class_embedding, n_embd, 1, n_batch, 1); + inp = ggml_concat(ctx0, cls_embd, inp, 1); // for selecting learned pos embd, used by ViT ggml_tensor * positions = ggml_cast(ctx0, ggml_arange(ctx0, 0, n_pos, 1), GGML_TYPE_I32); @@ -294,25 +333,56 @@ ggml_cgraph * clip_graph_deepseekocr::build() { clip_out = cur; } + // sam_out: [patch_h, patch_w, n_embd, n_batch] + // -> [n_embd, clip_n_patches, n_batch] sam_out = ggml_cont(ctx0, ggml_permute(ctx0, sam_out, 1, 2, 0, 3)); - sam_out = ggml_reshape_2d(ctx0, sam_out, sam_out->ne[0], clip_n_patches); - clip_out = ggml_view_2d(ctx0, clip_out, n_embd, clip_n_patches, clip_out->nb[1], clip_out->nb[1]); + sam_out = ggml_reshape_3d(ctx0, sam_out, sam_out->ne[0], clip_n_patches, n_batch); + + // clip_out: [n_embd, n_pos, n_batch] where n_pos = clip_n_patches + 1 (CLS) + // strip CLS token: skip first position, view only the patch tokens + clip_out = ggml_view_3d(ctx0, clip_out, n_embd, clip_n_patches, n_batch, + clip_out->nb[1], clip_out->nb[2], clip_out->nb[1]); ggml_tensor * cur; cur = ggml_concat(ctx0, clip_out, sam_out, 0); cur = ggml_mul_mat(ctx0, model.mm_fc_w, cur); cur = ggml_add(ctx0, cur, model.mm_fc_b); - const auto h = static_cast(std::sqrt(static_cast(cur->ne[1]))); - const auto w = h; - const auto n_dim = cur->ne[0]; + if (is_overview) { + // global view: weave one newline per row + trailing view separator + const auto h = static_cast(std::sqrt(static_cast(cur->ne[1]))); + const auto w = h; + const auto n_dim = cur->ne[0]; - ggml_tensor * imgnl; + ggml_tensor * imgnl = ggml_repeat_4d(ctx0, model.image_newline, n_dim, 1, h, 1); + cur = ggml_reshape_3d(ctx0, cur, n_dim, w, h); + cur = ggml_reshape_2d(ctx0, ggml_concat(ctx0, cur, imgnl, 1), n_dim, (w + 1) * h); + cur = ggml_concat(ctx0, cur, model.view_seperator, 1); // (n_dim, h*(w+1) + 1) + } else { + // tile row: interleave tiles within each row, add newline per row + const int grid_x = static_cast(std::sqrt(static_cast(clip_n_patches))); + const int grid_y = grid_x; + const auto n_dim = cur->ne[0]; - imgnl = ggml_repeat_4d(ctx0, model.image_newline, n_dim, 1, h, 1); - cur = ggml_reshape_3d(ctx0, cur, n_dim, w, h); - cur = ggml_reshape_2d(ctx0, ggml_concat(ctx0, cur, imgnl, 1), n_dim, (w + 1) * h); - cur = ggml_concat(ctx0, cur, model.view_seperator, 1); // (n_dim, h*(w+1) + 1) + // (n_dim, clip_n_patches, n_batch) -> (n_dim, grid_x, grid_y, n_batch) + cur = ggml_reshape_4d(ctx0, cur, n_dim, grid_x, grid_y, n_batch); + + // tiles: re-order from A.row0 A.row1 B.row0 B.row1 ... + // to A.row0 B.row0 A.row1 B.row1 ... + // then add nl: A.row0 B.row0 [nl] A.row1 B.row1 [nl] ... + // interleave tiles: (n_dim, grid_x, grid_y, n_batch) -> (n_dim, grid_x, n_batch, grid_y) + cur = ggml_cont(ctx0, ggml_permute(ctx0, cur, 0, 1, 3, 2)); + + // merge: (n_dim, grid_x, n_batch, grid_y) -> (n_dim, grid_x*n_batch, grid_y, 1) + cur = ggml_reshape_4d(ctx0, cur, n_dim, grid_x * n_batch, grid_y, 1); + + // append newline per row: (n_dim, grid_x*n_batch+1, grid_y, 1) + ggml_tensor * imgnl = ggml_repeat_4d(ctx0, model.image_newline, n_dim, 1, grid_y, 1); + cur = ggml_concat(ctx0, cur, imgnl, 1); + + // flatten: (n_dim, (grid_x*n_batch+1)*grid_y) + cur = ggml_reshape_2d(ctx0, cur, n_dim, (grid_x * n_batch + 1) * grid_y); + } cb(cur, "dsocr_output", -1); diff --git a/tools/mtmd/models/models.h b/tools/mtmd/models/models.h index 12d5e6949..5f1493fa6 100644 --- a/tools/mtmd/models/models.h +++ b/tools/mtmd/models/models.h @@ -127,6 +127,7 @@ struct clip_graph_deepseekocr : clip_graph { clip_graph_deepseekocr(clip_ctx * ctx, const clip_image_f32 & img) : clip_graph(ctx, img) {} ggml_cgraph * build() override; ggml_tensor * build_sam(ggml_tensor * inp); // build the SAM model + // bool support_batch() const override { return true; } // TODO: support batch for DeepSeek-OCR v1 }; struct clip_graph_deepseekocr2 : clip_graph_deepseekocr { diff --git a/tools/mtmd/mtmd-image.cpp b/tools/mtmd/mtmd-image.cpp index 01d9b4517..36cd463b2 100644 --- a/tools/mtmd/mtmd-image.cpp +++ b/tools/mtmd/mtmd-image.cpp @@ -1107,44 +1107,7 @@ mtmd_image_preproc_out mtmd_image_preprocessor_internvl::preprocess(const clip_i // mtmd_image_preprocessor_deepseekocr // -mtmd_image_preproc_out mtmd_image_preprocessor_deepseekocr::preprocess(const clip_image_u8 & img) { - static constexpr int native_resolutions[] = { 1024 /* base */, 1280 /* large */ }; - // TODO: support 512 (tiny) and 640 (small) once we have eval data for them - - const int64_t orig_area = static_cast(img.get_size().area()); - - size_t mode_i = 0; - int64_t min_diff = std::numeric_limits::max(); - for (size_t i = 0; i < std::size(native_resolutions); i++) { - const int64_t r = native_resolutions[i]; - const int64_t diff = std::abs(orig_area - r * r); - if (diff < min_diff) { - mode_i = i; - min_diff = diff; - } - } - const int image_size = native_resolutions[mode_i]; - - // Aspect-preserving fit-and-pad. Pillow bicubic + PAD_NEAREST for - // byte-parity with the upstream deepseek-ai/DeepSeek-OCR HF preprocessor. - clip_image_u8 padded; - img_tool::resize(img, padded, {image_size, image_size}, RESIZE_ALGO_BICUBIC_PILLOW, - PAD_NEAREST, hparams.image_pad_color); - mtmd_image_preproc_out output; - output.append_overview(hparams, padded, true); - output.grid_x = 0; - output.grid_y = 0; - // TODO @ngxson : support slicing for DeepSeek-OCR, to do in another PR - return output; -} - -// -// mtmd_image_preprocessor_deepseekocr2 -// - -// candidate tile grids (cols, rows) with min_tiles <= cols*rows <= max_tiles -// sorted by tile count -std::vector mtmd_image_preprocessor_deepseekocr2::get_target_ratios() { +std::vector mtmd_image_preprocessor_deepseekocr::get_target_ratios() const { std::vector ratios; for (int n = min_tiles; n <= max_tiles; n++) { for (int w = 1; w <= n; w++) { @@ -1171,13 +1134,11 @@ std::vector mtmd_image_preprocessor_deepseekocr2::get_target_ra return ratios; } -// pick the grid whose aspect ratio is closest to the image -// on a tie, prefer the larger grid when the image fits -clip_image_size mtmd_image_preprocessor_deepseekocr2::find_closest_aspect_ratio( +clip_image_size mtmd_image_preprocessor_deepseekocr::find_closest_aspect_ratio( float aspect_ratio, const std::vector & target_ratios, int width, - int height) { + int height) const { float best_ratio_diff = std::numeric_limits::max(); clip_image_size best_ratio = { 1, 1 }; const float area = static_cast(width * height); @@ -1198,37 +1159,69 @@ clip_image_size mtmd_image_preprocessor_deepseekocr2::find_closest_aspect_ratio( return best_ratio; } -mtmd_image_preproc_out mtmd_image_preprocessor_deepseekocr2::preprocess(const clip_image_u8 & img) { - // emit 768x768 local tiles when the image is larger than a tile in either - // dimension, then always a 1024x1024 global view. order: [tiles..., global]. - +mtmd_image_preproc_out mtmd_image_preprocessor_deepseekocr::preprocess(const clip_image_u8 & img) { mtmd_image_preproc_out output; + int grid_w = 0; + int grid_h = 0; const auto img_size = img.get_size(); + + // global view: aspect-preserving fit-and-pad to base_size + clip_image_u8 padded; + img_tool::resize(img, padded, + { base_size, base_size }, + RESIZE_ALGO_BICUBIC_PILLOW, + PAD_NEAREST, + hparams.image_pad_color); + output.append_overview(hparams, padded, true); + output.overview.add_viewsep = true; + + // if this condition doesn't hold, the output is overview only, no tiles if (img_size.width > tile_size || img_size.height > tile_size) { const float aspect_ratio = static_cast(img_size.width) / img_size.height; const auto target_ratios = get_target_ratios(); - const clip_image_size grid = find_closest_aspect_ratio(aspect_ratio, target_ratios, img_size.width, img_size.height); + const clip_image_size grid = + find_closest_aspect_ratio(aspect_ratio, target_ratios, img_size.width, img_size.height); + grid_w = grid.width; + grid_h = grid.height; - // stretch onto the grid (no aspect preserve), then crop tiles row-major. clip_image_u8 refined; - img_tool::resize(img, refined, { tile_size * grid.width, tile_size * grid.height }, - RESIZE_ALGO_BICUBIC_PILLOW, PAD_NONE); + img_tool::resize(img, refined, { tile_size * grid_w, tile_size * grid_h }, RESIZE_ALGO_BICUBIC_PILLOW, + PAD_NONE); - for (int row = 0; row < grid.height; row++) { - for (int col = 0; col < grid.width; col++) { - clip_image_u8 tile; - img_tool::crop(refined, tile, col * tile_size, row * tile_size, tile_size, tile_size); - output.append(hparams, tile, true); + for (int row = 0; row < grid_h; row++) { + if (fuse_row) { + // concat all tiles in this row into a single image, along the H axis + // output image size: w = tile_size, h = tile_size * grid_w + // this is to ensure the whole row is always processed together + clip_image_u8 row_img; + row_img.set_size({tile_size, tile_size * grid_w}, false); + for (int col = 0; col < grid_w; col++) { + for (int py = 0; py < tile_size; py++) { + for (int px = 0; px < tile_size; px++) { + row_img.set_pixel(px, col * tile_size + py, + refined.get_pixel(col * tile_size + px, row * tile_size + py)); + } + } + } + output.append(hparams, row_img, true); + } else { + for (int col = 0; col < grid_w; col++) { + clip_image_u8 tile; + img_tool::crop(refined, tile, col * tile_size, row * tile_size, tile_size, tile_size); + output.append(hparams, tile, true); + } } } + if (fuse_row) { + grid_w = 1; // each fused row is one image; a single output column + } } - // global view: aspect-preserving fit-and-pad to base_size. - clip_image_u8 padded; - img_tool::resize(img, padded, { base_size, base_size }, RESIZE_ALGO_BICUBIC_PILLOW, - PAD_NEAREST, hparams.image_pad_color); - output.append_overview(hparams, padded, true); - output.overview.add_viewsep = true; + LOG_DBG("%s: grid size: %d x %d (%d tiles) + global view\n", __func__, grid_w, grid_h, grid_w * grid_h); + LOG_DBG("%s: overview size: %d x %d\n", __func__, padded.get_size().width, padded.get_size().height); + + output.grid_x = grid_w; + output.grid_y = grid_h; return output; } diff --git a/tools/mtmd/mtmd-image.h b/tools/mtmd/mtmd-image.h index f458e39e7..115cba51e 100644 --- a/tools/mtmd/mtmd-image.h +++ b/tools/mtmd/mtmd-image.h @@ -160,29 +160,29 @@ struct mtmd_image_preprocessor_internvl : mtmd_image_preprocessor_llava_uhd { mtmd_image_preproc_out preprocess(const clip_image_u8 & img) override; }; +// DeepSeek-OCR (v1/v2) global view + optional local tile grid struct mtmd_image_preprocessor_deepseekocr : mtmd_image_preprocessor { - mtmd_image_preprocessor_deepseekocr(const clip_ctx * ctx) : mtmd_image_preprocessor(ctx) {} - mtmd_image_preproc_out preprocess(const clip_image_u8 & img) override; -}; - -// DeepSeek-OCR-2: a 1024x1024 global view, plus InternVL-style 768x768 local -// tiles when the image is larger than a tile in either dimension. -struct mtmd_image_preprocessor_deepseekocr2 : mtmd_image_preprocessor { - static constexpr int base_size = 1024; // global view - static constexpr int tile_size = 768; // local tile - static constexpr int min_tiles = 2; - static constexpr int max_tiles = 6; - - mtmd_image_preprocessor_deepseekocr2(const clip_ctx * ctx) : mtmd_image_preprocessor(ctx) {} + mtmd_image_preprocessor_deepseekocr(const clip_ctx * ctx) + : mtmd_image_preprocessor(ctx), + fuse_row(clip_get_projector_type(ctx) == PROJECTOR_TYPE_DEEPSEEKOCR), + base_size(hparams.image_size), + tile_size(hparams.preproc_tile_size), + min_tiles(hparams.preproc_min_tiles), + max_tiles(hparams.preproc_max_tiles) {} mtmd_image_preproc_out preprocess(const clip_image_u8 & img) override; private: - static std::vector get_target_ratios(); - static clip_image_size find_closest_aspect_ratio( - float aspect_ratio, - const std::vector & target_ratios, - int width, - int height); + bool fuse_row; // v1 fuses a tile-row into one image; v2 keeps tiles separate + int base_size; // global view + int tile_size; // each tile + int min_tiles; + int max_tiles; + + std::vector get_target_ratios() const; + clip_image_size find_closest_aspect_ratio( + float aspect_ratio, + const std::vector & target_ratios, + int width, int height) const; }; // custom image preprocessing for Step3VL diff --git a/tools/mtmd/mtmd.cpp b/tools/mtmd/mtmd.cpp index 5aa902be6..24b82db21 100644 --- a/tools/mtmd/mtmd.cpp +++ b/tools/mtmd/mtmd.cpp @@ -618,15 +618,10 @@ struct mtmd_context { image_preproc = std::make_unique(ctx_v); } break; case PROJECTOR_TYPE_DEEPSEEKOCR: - { - img_end = "\n"; // prevent empty batch on llama-server - image_preproc = std::make_unique(ctx_v); - ov_img_first = false; - } break; case PROJECTOR_TYPE_DEEPSEEKOCR2: { img_end = "\n"; // prevent empty batch on llama-server - image_preproc = std::make_unique(ctx_v); + image_preproc = std::make_unique(ctx_v); ov_img_first = false; } break; case PROJECTOR_TYPE_HUNYUANVL: @@ -1132,6 +1127,7 @@ struct mtmd_tokenizer { // add slices (or tiles) if (!chunks.empty()) { + LOG_DBG("%s: adding %d slices (%d rows x %d cols)\n", __func__, (int)chunks.size(), n_row, n_col); GGML_ASSERT((int)chunks.size() == n_row * n_col); add_text(ctx->tok_slices_start); for (int y = 0; y < n_row; y++) { @@ -1174,7 +1170,6 @@ struct mtmd_tokenizer { cur.entries.emplace_back(std::move(ov_chunk)); add_text(ctx->tok_ov_img_end); } - } else { if (preproc_out.entries.size() == 0) { diff --git a/tools/server/server-http.cpp b/tools/server/server-http.cpp index bb88dda21..87eee5fd4 100644 --- a/tools/server/server-http.cpp +++ b/tools/server/server-http.cpp @@ -175,6 +175,15 @@ bool server_http_context::init(const common_params & params) { // Middlewares // + // Frontend paths - all embedded UI assets + static const std::unordered_set frontend_paths = []() { + std::unordered_set paths { "/" }; + for (const llama_ui_asset & a : llama_ui_get_assets()) { + paths.insert("/" + a.name); + } + return paths; + }(); + // Public endpoints - API routes plus all embedded UI assets static const std::unordered_set get_public_endpoints = []() { std::unordered_set endpoints { @@ -182,11 +191,8 @@ bool server_http_context::init(const common_params & params) { "/v1/health", "/models", "/v1/models", - "/", }; - for (const llama_ui_asset & a : llama_ui_get_assets()) { - endpoints.insert("/" + a.name); - } + endpoints.insert(frontend_paths.begin(), frontend_paths.end()); return endpoints; }(); @@ -239,18 +245,9 @@ bool server_http_context::init(const common_params & params) { auto middleware_server_state = [this](const httplib::Request & req, httplib::Response & res) { if (!is_ready.load()) { -#if defined(LLAMA_UI_HAS_ASSETS) - if (const auto tmp = string_split(req.path, '.'); - req.path == "/" || (!tmp.empty() && tmp.back() == "html")) { - if (const llama_ui_asset * a = llama_ui_find_asset("loading.html")) { - res.status = 503; - res.set_content(reinterpret_cast(a->data), a->size, "text/html; charset=utf-8"); - return false; - } + if (frontend_paths.count(req.path)) { + return true; // frontend asset, allow it to load and show "loading" } -#else - (void)req; -#endif // no endpoints are allowed to be accessed when the server is not ready // this is to prevent any data races or inconsistent states res.status = 503; diff --git a/tools/server/server-schema.cpp b/tools/server/server-schema.cpp index 5713cc831..89026eb4e 100644 --- a/tools/server/server-schema.cpp +++ b/tools/server/server-schema.cpp @@ -568,10 +568,16 @@ static void handle_with_catch(const char * name, std::function func) { } } +// treat a null value as absent so clients can send null to request the server default +static bool has_value(const json & data, const char * n) { + auto it = data.find(n); + return it != data.end() && !it->is_null(); +} + template void field_num::eval(field_eval_context & ctx, const json & data) { for (const auto & n : name) { - if (data.contains(n)) { + if (has_value(data, n)) { handle_with_catch(n, [&]() { if (custom_handler) { custom_handler(ctx, data); @@ -593,7 +599,7 @@ void field_num::eval(field_eval_context & ctx, const json & data) { void field_str::eval(field_eval_context & ctx, const json & data) { GGML_ASSERT(custom_handler); for (const auto & n : name) { - if (data.contains(n)) { + if (has_value(data, n)) { handle_with_catch(n, [&]() { custom_handler(ctx, data); }); @@ -604,7 +610,7 @@ void field_str::eval(field_eval_context & ctx, const json & data) { void field_bool::eval(field_eval_context & ctx, const json & data) { for (const auto & n : name) { - if (data.contains(n)) { + if (has_value(data, n)) { handle_with_catch(n, [&]() { if (custom_handler) { custom_handler(ctx, data); @@ -620,7 +626,7 @@ void field_bool::eval(field_eval_context & ctx, const json & data) { void field_json::eval(field_eval_context & ctx, const json & data) { GGML_ASSERT(custom_handler); for (const auto & n : name) { - if (data.contains(n)) { + if (has_value(data, n)) { handle_with_catch(n, [&]() { custom_handler(ctx, data); }); diff --git a/tools/server/server-tools.cpp b/tools/server/server-tools.cpp index 790ed85a0..3f162a13e 100644 --- a/tools/server/server-tools.cpp +++ b/tools/server/server-tools.cpp @@ -19,91 +19,7 @@ namespace fs = std::filesystem; // internal helpers // -static std::vector to_cstr_vec(const std::vector & v) { - std::vector r; - r.reserve(v.size() + 1); - for (const auto & s : v) { - r.push_back(const_cast(s.c_str())); - } - r.push_back(nullptr); - return r; -} - -struct run_proc_result { - std::string output; - int exit_code = -1; - bool timed_out = false; -}; - -static run_proc_result run_process( - const std::vector & args, - size_t max_output, - int timeout_secs) { - run_proc_result res; - - subprocess_s proc; - auto argv = to_cstr_vec(args); - - int options = subprocess_option_no_window - | subprocess_option_combined_stdout_stderr - | subprocess_option_inherit_environment - | subprocess_option_search_user_path; - - if (subprocess_create(argv.data(), options, &proc) != 0) { - res.output = "failed to spawn process"; - return res; - } - - std::atomic done{false}; - std::atomic timed_out{false}; - - std::thread timeout_thread([&]() { - auto deadline = std::chrono::steady_clock::now() + std::chrono::seconds(timeout_secs); - while (!done.load()) { - if (std::chrono::steady_clock::now() >= deadline) { - timed_out.store(true); - subprocess_terminate(&proc); - return; - } - std::this_thread::sleep_for(std::chrono::milliseconds(100)); - } - }); - - FILE * f = subprocess_stdout(&proc); - std::string output; - bool truncated = false; - if (f) { - char buf[4096]; - while (fgets(buf, sizeof(buf), f) != nullptr) { - if (!truncated) { - size_t len = strlen(buf); - if (output.size() + len <= max_output) { - output.append(buf, len); - } else { - output.append(buf, max_output - output.size()); - truncated = true; - } - } - } - } - - done.store(true); - if (timeout_thread.joinable()) { - timeout_thread.join(); - } - - subprocess_join(&proc, &res.exit_code); - subprocess_destroy(&proc); - - res.output = output; - res.timed_out = timed_out.load(); - if (truncated) { - res.output += "\n[output truncated]"; - } - return res; -} - -json server_tool::to_json() { +json server_tool::to_json() const { return { {"display_name", display_name}, {"tool", name}, @@ -115,6 +31,228 @@ json server_tool::to_json() { }; } +static constexpr size_t SERVER_TOOL_GIT_LS_FILES_MAX_OUTPUT = 8 * 1024 * 1024; // 8 MB +static constexpr int SERVER_TOOL_GIT_LS_FILES_TIMEOUT = 15; // seconds + +class tools_io { +public: + struct exec_result { + std::string output; + int exit_code = -1; + bool timed_out = false; + }; + + virtual ~tools_io() = default; + + virtual bool is_directory(const std::string & path) const = 0; + virtual bool is_regular_file(const std::string & path) const = 0; + virtual bool file_size(const std::string & path, uintmax_t & out_size) const = 0; + virtual bool read_file(const std::string & path, std::string & out) const = 0; + virtual bool write_file(const std::string & path, const std::string & content) const = 0; + // paths relative to `base`, '/'-separated; sets `err` if `base` isn't a directory + virtual std::vector list_files(const std::string & base, std::string & err) const = 0; + virtual exec_result run(const std::vector & args, size_t max_output, int timeout_secs) const = 0; +}; + +class tools_io_basic : public tools_io { +public: + bool is_directory(const std::string & path) const override { + std::error_code ec; + return fs::is_directory(path, ec) && !ec; + } + + bool is_regular_file(const std::string & path) const override { + std::error_code ec; + return fs::is_regular_file(path, ec) && !ec; + } + + bool file_size(const std::string & path, uintmax_t & out_size) const override { + std::error_code ec; + out_size = fs::file_size(path, ec); + return !ec; + } + + bool read_file(const std::string & path, std::string & out) const override { + std::ifstream f(path, std::ios::binary); + if (!f) return false; + std::ostringstream ss; + ss << f.rdbuf(); + out = ss.str(); + return true; + } + + bool write_file(const std::string & path, const std::string & content) const override { + std::error_code ec; + fs::path fpath(path); + if (fpath.has_parent_path()) { + fs::create_directories(fpath.parent_path(), ec); + if (ec) return false; + } + std::ofstream f(path, std::ios::binary); + if (!f) return false; + f << content; + return (bool) f; + } + + std::vector list_files(const std::string & base, std::string & err) const override { + err.clear(); + if (!is_directory(base)) { + err = "path does not exist or is not a directory: " + base; + return {}; + } + + auto res = run( + {"git", "-C", base, "ls-files", "--cached", "--others", "--exclude-standard"}, + SERVER_TOOL_GIT_LS_FILES_MAX_OUTPUT, SERVER_TOOL_GIT_LS_FILES_TIMEOUT); + + if (res.exit_code == 0 && !res.timed_out) { + std::vector result; + std::istringstream iss(res.output); + std::string line; + while (std::getline(iss, line)) { + if (!line.empty() && line.back() == '\r') line.pop_back(); + if (line.empty()) continue; + std::replace(line.begin(), line.end(), '\\', '/'); + if (is_regular_file((fs::path(base) / line).string())) { + result.push_back(line); + } + } + return result; + } + + return list_files_fallback(base); + } + + exec_result run(const std::vector & args, size_t max_output, int timeout_secs) const override { + exec_result res; + + subprocess_s proc; + auto argv = to_cstr_vec(args); + + int options = subprocess_option_no_window + | subprocess_option_combined_stdout_stderr + | subprocess_option_inherit_environment + | subprocess_option_search_user_path; + + if (subprocess_create(argv.data(), options, &proc) != 0) { + res.output = "failed to spawn process"; + return res; + } + + std::atomic done{false}; + std::atomic timed_out{false}; + + std::thread timeout_thread([&]() { + auto deadline = std::chrono::steady_clock::now() + std::chrono::seconds(timeout_secs); + while (!done.load()) { + if (std::chrono::steady_clock::now() >= deadline) { + timed_out.store(true); + subprocess_terminate(&proc); + return; + } + std::this_thread::sleep_for(std::chrono::milliseconds(100)); + } + }); + + FILE * f = subprocess_stdout(&proc); + std::string output; + bool truncated = false; + if (f) { + char buf[4096]; + while (fgets(buf, sizeof(buf), f) != nullptr) { + if (!truncated) { + size_t len = strlen(buf); + if (output.size() + len <= max_output) { + output.append(buf, len); + } else { + output.append(buf, max_output - output.size()); + truncated = true; + } + } + } + } + + done.store(true); + if (timeout_thread.joinable()) { + timeout_thread.join(); + } + + subprocess_join(&proc, &res.exit_code); + subprocess_destroy(&proc); + + res.output = output; + res.timed_out = timed_out.load(); + if (truncated) { + res.output += "\n[output truncated]"; + } + return res; + } + +private: + static std::vector to_cstr_vec(const std::vector & v) { + std::vector r; + r.reserve(v.size() + 1); + for (const auto & s : v) { + r.push_back(const_cast(s.c_str())); + } + r.push_back(nullptr); + return r; + } + + static const std::unordered_set & junk_dir_names() { + static const std::unordered_set names = { + ".git", ".svn", ".hg", "node_modules", "__pycache__", + ".venv", "venv", "dist", "build", "target", ".cache", ".idea", ".vscode", + }; + return names; + } + + std::vector list_files_fallback(const std::string & base) const { + std::vector result; + std::error_code ec; + + std::vector> stack; + stack.emplace_back(fs::path(base), fs::path()); + + while (!stack.empty()) { + auto [dir, rel_dir] = stack.back(); + stack.pop_back(); + + for (const auto & entry : fs::directory_iterator(dir, fs::directory_options::skip_permission_denied, ec)) { + if (ec) break; + std::string fname = entry.path().filename().string(); + std::error_code tec; + if (entry.is_directory(tec)) { + if (junk_dir_names().count(fname) > 0) continue; + stack.emplace_back(entry.path(), rel_dir / fname); + } else if (entry.is_regular_file(tec)) { + std::string rel = (rel_dir / fname).string(); + std::replace(rel.begin(), rel.end(), '\\', '/'); + result.push_back(rel); + } + } + } + + return result; + } +}; + +static std::unique_ptr make_tools_io(const json & params) { + GGML_UNUSED(params); // TODO in follow-up PR + return std::make_unique(); +} + +// no '/' in pattern -> match basename at any depth; else match full relative path +static bool path_glob_match(const std::string & pattern, const std::string & rel_path) { + if (pattern.find('/') == std::string::npos) { + return glob_match(pattern, fs::path(rel_path).filename().string()); + } + if (pattern == "**" || pattern.rfind("**/", 0) == 0 || pattern.rfind('/', 0) == 0) { + return glob_match(pattern, rel_path); + } + return glob_match("**/" + pattern, rel_path); +} + // // read_file: read a file with optional line range and line-number prefix // @@ -128,7 +266,7 @@ struct server_tool_read_file : server_tool { permission_write = false; } - json get_definition() override { + json get_definition() const override { return { {"type", "function"}, {"function", { @@ -149,16 +287,17 @@ struct server_tool_read_file : server_tool { }; } - json invoke(json params) override { + json invoke(json params) const override { std::string path = params.at("path").get(); int start_line = json_value(params, "start_line", 1); int end_line = json_value(params, "end_line", -1); // -1 = no limit bool append_loc = json_value(params, "append_loc", false); - std::error_code ec; - uintmax_t file_size = fs::file_size(path, ec); - if (ec) { - return {{"error", "cannot stat file: " + ec.message()}}; + auto io = make_tools_io(params); + + uintmax_t file_size = 0; + if (!io->file_size(path, file_size)) { + return {{"error", "cannot stat file: " + path}}; } if (file_size > SERVER_TOOL_READ_FILE_MAX_SIZE && end_line == -1) { return {{"error", string_format( @@ -166,11 +305,12 @@ struct server_tool_read_file : server_tool { (size_t)file_size, SERVER_TOOL_READ_FILE_MAX_SIZE)}}; } - std::ifstream f(path); - if (!f) { + std::string content; + if (!io->read_file(path, content)) { return {{"error", "failed to open file: " + path}}; } + std::istringstream f(content); std::string result; std::string line; int lineno = 0; @@ -211,17 +351,23 @@ struct server_tool_file_glob_search : server_tool { permission_write = false; } - json get_definition() override { + json get_definition() const override { return { {"type", "function"}, {"function", { {"name", name}, - {"description", "Recursively search for files matching a glob pattern under a directory."}, + {"description", + "Recursively search for files matching a glob pattern under a directory. " + "Automatically skips files ignored by .gitignore (when the directory is inside a git repo) " + "and common junk directories (.git, node_modules, build, dist, etc.) otherwise. " + "A pattern with no '/' (e.g. \"*.cpp\") matches the file's basename at any depth. " + "A pattern containing '/' matches the full relative path; unless already anchored with " + "\"**/\" or a leading '/', it is automatically prefixed with \"**/\"."}, {"parameters", { {"type", "object"}, {"properties", { {"path", {{"type", "string"}, {"description", "Base directory to search in"}}}, - {"include", {{"type", "string"}, {"description", "Glob pattern for files to include (e.g. \"**/*.cpp\"). Default: **"}}}, + {"include", {{"type", "string"}, {"description", "Glob pattern for files to include (e.g. \"*.cpp\" or \"src/**/*.cpp\"). Default: **"}}}, {"exclude", {{"type", "string"}, {"description", "Glob pattern for files to exclude"}}}, }}, {"required", json::array({"path"})}, @@ -230,33 +376,39 @@ struct server_tool_file_glob_search : server_tool { }; } - json invoke(json params) override { + json invoke(json params) const override { std::string base = params.at("path").get(); std::string include = json_value(params, "include", std::string("**")); std::string exclude = json_value(params, "exclude", std::string("")); - std::ostringstream output_text; - size_t count = 0; - - std::error_code ec; - for (const auto & entry : fs::recursive_directory_iterator(base, - fs::directory_options::skip_permission_denied, ec)) { - if (!entry.is_regular_file()) continue; - - std::string rel = fs::relative(entry.path(), base, ec).string(); - if (ec) continue; - std::replace(rel.begin(), rel.end(), '\\', '/'); - - if (!glob_match(include, rel)) continue; - if (!exclude.empty() && glob_match(exclude, rel)) continue; - - output_text << entry.path().string() << "\n"; - if (++count >= SERVER_TOOL_FILE_SEARCH_MAX_RESULTS) { - break; - } + auto io = make_tools_io(params); + std::string err; + auto files = io->list_files(base, err); + if (!err.empty()) { + return {{"error", err}}; } - output_text << "\n---\nTotal matches: " << count << "\n"; + std::vector matches; + for (const auto & rel : files) { + if (!path_glob_match(include, rel)) continue; + if (!exclude.empty() && path_glob_match(exclude, rel)) continue; + matches.push_back(rel); + } + + size_t total = matches.size(); + size_t shown = std::min(total, SERVER_TOOL_FILE_SEARCH_MAX_RESULTS); + + std::ostringstream output_text; + for (size_t i = 0; i < shown; i++) { + output_text << matches[i] << "\n"; + } + + output_text << "\n---\nTotal matches: " << total << "\n"; + if (total > shown) { + output_text << string_format( + "[%zu results limit reached (%zu total matches). Refine the glob pattern to narrow the search.]\n", + shown, total); + } return {{"plain_text_response", output_text.str()}}; } @@ -275,20 +427,29 @@ struct server_tool_grep_search : server_tool { permission_write = false; } - json get_definition() override { + json get_definition() const override { return { {"type", "function"}, {"function", { {"name", name}, - {"description", "Search for a regex pattern in files under a path. Returns matching lines."}, + {"description", + "Search for a pattern in files under a path. Returns matching lines with file paths " + "(and, unless searching a single file, paths relative to the given directory). " + "Automatically skips files ignored by .gitignore (when the directory is inside a git repo) " + "and common junk directories (.git, node_modules, build, dist, etc.) otherwise. " + "include/exclude: a pattern with no '/' matches the basename at any depth; a pattern " + "containing '/' matches the full relative path (auto-anchored with \"**/\" unless already anchored)."}, {"parameters", { {"type", "object"}, {"properties", { {"path", {{"type", "string"}, {"description", "File or directory to search in"}}}, - {"pattern", {{"type", "string"}, {"description", "Regular expression pattern to search for"}}}, + {"pattern", {{"type", "string"}, {"description", "Pattern to search for (regular expression unless literal is true)"}}}, {"include", {{"type", "string"}, {"description", "Glob pattern to filter files (default: **)"}}}, {"exclude", {{"type", "string"}, {"description", "Glob pattern to exclude files"}}}, {"return_line_numbers", {{"type", "boolean"}, {"description", "If true, include line numbers in results"}}}, + {"literal", {{"type", "boolean"}, {"description", "Treat pattern as a literal string instead of a regular expression (default: false)"}}}, + {"ignore_case", {{"type", "boolean"}, {"description", "Case-insensitive search (default: false)"}}}, + {"context_lines", {{"type", "integer"}, {"description", "Number of lines of context to show before and after each match (default: 0)"}}}, }}, {"required", json::array({"path", "pattern"})}, }}, @@ -296,64 +457,109 @@ struct server_tool_grep_search : server_tool { }; } - json invoke(json params) override { - std::string path = params.at("path").get(); - std::string pat_str = params.at("pattern").get(); - std::string include = json_value(params, "include", std::string("**")); - std::string exclude = json_value(params, "exclude", std::string("")); - bool show_lineno = json_value(params, "return_line_numbers", false); + json invoke(json params) const override { + std::string path = params.at("path").get(); + std::string pat_str = params.at("pattern").get(); + std::string include = json_value(params, "include", std::string("**")); + std::string exclude = json_value(params, "exclude", std::string("")); + bool show_lineno = json_value(params, "return_line_numbers", false); + bool literal = json_value(params, "literal", false); + bool ignore_case = json_value(params, "ignore_case", false); + int ctx_lines = std::max(0, json_value(params, "context_lines", 0)); + + std::string pattern_src = pat_str; + if (literal) { + static const std::string specials = "\\^$.|?*+()[]{}"; + std::string escaped; + escaped.reserve(pat_str.size() * 2); + for (char c : pat_str) { + if (specials.find(c) != std::string::npos) escaped += '\\'; + escaped += c; + } + pattern_src = escaped; + } std::regex pattern; try { - pattern = std::regex(pat_str); + auto flags = std::regex::ECMAScript; + if (ignore_case) flags |= std::regex::icase; + pattern = std::regex(pattern_src, flags); } catch (const std::regex_error & e) { return {{"error", std::string("invalid regex: ") + e.what()}}; } - std::ostringstream output_text; - size_t total = 0; + auto io = make_tools_io(params); - auto search_file = [&](const fs::path & fpath) { - std::ifstream f(fpath); - if (!f) return; - std::string line; - int lineno = 0; - while (std::getline(f, line) && total < SERVER_TOOL_GREP_SEARCH_MAX_RESULTS) { - lineno++; - if (std::regex_search(line, pattern)) { - output_text << fpath.string() << ":"; - if (show_lineno) { - output_text << lineno << ":"; - } - output_text << line << "\n"; - total++; - } + // collect (absolute_path, display_path) pairs to search + std::vector> files; + + if (io->is_regular_file(path)) { + files.emplace_back(path, path); + } else if (io->is_directory(path)) { + std::string err; + auto candidates = io->list_files(path, err); + if (!err.empty()) { + return {{"error", err}}; } - }; - - std::error_code ec; - if (fs::is_regular_file(path, ec)) { - search_file(path); - } else if (fs::is_directory(path, ec)) { - for (const auto & entry : fs::recursive_directory_iterator(path, - fs::directory_options::skip_permission_denied, ec)) { - if (!entry.is_regular_file()) continue; - if (total >= SERVER_TOOL_GREP_SEARCH_MAX_RESULTS) break; - - std::string rel = fs::relative(entry.path(), path, ec).string(); - if (ec) continue; - std::replace(rel.begin(), rel.end(), '\\', '/'); - - if (!glob_match(include, rel)) continue; - if (!exclude.empty() && glob_match(exclude, rel)) continue; - - search_file(entry.path()); + for (const auto & rel : candidates) { + if (!path_glob_match(include, rel)) continue; + if (!exclude.empty() && path_glob_match(exclude, rel)) continue; + files.emplace_back((fs::path(path) / rel).string(), rel); } } else { return {{"error", "path does not exist: " + path}}; } - output_text << "\n\n---\nTotal matches: " << total << "\n"; + std::ostringstream output_text; + size_t total = 0; + bool limit_reached = false; + bool show_num = show_lineno || ctx_lines > 0; + + for (const auto & file_entry : files) { + if (limit_reached) break; + const std::string & fpath = file_entry.first; + const std::string & display_path = file_entry.second; + + std::string content; + if (!io->read_file(fpath, content)) continue; + std::vector lines; + { + std::istringstream f(content); + std::string line; + while (std::getline(f, line)) lines.push_back(line); + } + + for (size_t i = 0; i < lines.size(); i++) { + if (total >= SERVER_TOOL_GREP_SEARCH_MAX_RESULTS) { + limit_reached = true; + break; + } + if (!std::regex_search(lines[i], pattern)) continue; + + long ctx_start = ctx_lines > 0 ? std::max(0, (long) i - ctx_lines) : (long) i; + long ctx_end = ctx_lines > 0 ? std::min((long) lines.size() - 1, (long) i + ctx_lines) : (long) i; + + for (long j = ctx_start; j <= ctx_end; j++) { + bool is_match = (j == (long) i); + output_text << display_path << (is_match ? ':' : '-'); + if (show_num) { + output_text << (j + 1) << (is_match ? ':' : '-'); + } + output_text << lines[j] << "\n"; + } + if (ctx_lines > 0) { + output_text << "--\n"; + } + total++; + } + } + + output_text << "\n---\nTotal matches: " << total << "\n"; + if (limit_reached) { + output_text << string_format( + "[%zu matches limit reached. Narrow the path/pattern/include to see more.]\n", + SERVER_TOOL_GREP_SEARCH_MAX_RESULTS); + } return {{"plain_text_response", output_text.str()}}; } @@ -373,7 +579,7 @@ struct server_tool_exec_shell_command : server_tool { permission_write = true; } - json get_definition() override { + json get_definition() const override { return { {"type", "function"}, {"function", { @@ -392,7 +598,7 @@ struct server_tool_exec_shell_command : server_tool { }; } - json invoke(json params) override { + json invoke(json params) const override { std::string command = params.at("command").get(); int timeout = json_value(params, "timeout", 10); size_t max_output = (size_t) json_value(params, "max_output_size", (int) SERVER_TOOL_EXEC_SHELL_COMMAND_MAX_OUTPUT_SIZE); @@ -406,7 +612,8 @@ struct server_tool_exec_shell_command : server_tool { std::vector args = {"sh", "-c", command}; #endif - auto res = run_process(args, max_output, timeout); + auto io = make_tools_io(params); + auto res = io->run(args, max_output, timeout); std::string text_output = res.output; text_output += string_format("\n[exit code: %d]", res.exit_code); @@ -429,7 +636,7 @@ struct server_tool_write_file : server_tool { permission_write = true; } - json get_definition() override { + json get_definition() const override { return { {"type", "function"}, {"function", { @@ -447,25 +654,12 @@ struct server_tool_write_file : server_tool { }; } - json invoke(json params) override { + json invoke(json params) const override { std::string path = params.at("path").get(); std::string content = params.at("content").get(); - std::error_code ec; - fs::path fpath(path); - if (fpath.has_parent_path()) { - fs::create_directories(fpath.parent_path(), ec); - if (ec) { - return {{"error", "failed to create directories: " + ec.message()}}; - } - } - - std::ofstream f(path, std::ios::binary); - if (!f) { - return {{"error", "failed to open file for writing: " + path}}; - } - f << content; - if (!f) { + auto io = make_tools_io(params); + if (!io->write_file(path, content)) { return {{"error", "failed to write file: " + path}}; } @@ -474,7 +668,7 @@ struct server_tool_write_file : server_tool { }; // -// edit_file: edit file content via line-based changes +// edit_file: exact text replacement, one or more edits per call // struct server_tool_edit_file : server_tool { @@ -484,218 +678,322 @@ struct server_tool_edit_file : server_tool { permission_write = true; } - json get_definition() override { + json get_definition() const override { return { {"type", "function"}, {"function", { {"name", name}, {"description", - "Edit a file by applying a list of line-based changes. " - "Each change targets a 1-based inclusive line range and has a mode: " - "\"replace\" (replace lines with content), " - "\"delete\" (remove lines, content must be empty string), " - "\"append\" (insert content after line_end). " - "Set line_start to -1 to target the end of file (line_end is ignored in that case). " - "Changes must not overlap. They are applied in reverse line order automatically."}, + "Edit a file using exact text replacement. Each edits[].old_text must be unique in the file " + "and is matched against the original content, not incrementally. Merge nearby changes into " + "one edit instead of overlapping edits. Use write_file to replace the whole file."}, {"parameters", { {"type", "object"}, {"properties", { - {"path", {{"type", "string"}, {"description", "Path to the file to edit"}}}, - {"changes", { + {"path", {{"type", "string"}, {"description", "Path to the file to edit"}}}, + {"edits", { {"type", "array"}, - {"description", "List of changes to apply"}, + {"description", "One or more exact text replacements to apply"}, {"items", { {"type", "object"}, {"properties", { - {"mode", {{"type", "string"}, {"description", "\"replace\", \"delete\", or \"append\""}}}, - {"line_start", {{"type", "integer"}, {"description", "First line of the range (1-based); use -1 for end of file"}}}, - {"line_end", {{"type", "integer"}, {"description", "Last line of the range (1-based, inclusive); ignored when line_start is -1"}}}, - {"content", {{"type", "string"}, {"description", "Content to insert; must be empty string for delete mode"}}}, + {"old_text", {{"type", "string"}, {"description", "Exact text to find; must be unique in the file and must not overlap with other edits"}}}, + {"new_text", {{"type", "string"}, {"description", "Text to replace old_text with"}}}, }}, - {"required", json::array({"mode", "line_start", "line_end", "content"})}, + {"required", json::array({"old_text", "new_text"})}, }}, }}, }}, - {"required", json::array({"path", "changes"})}, + {"required", json::array({"path", "edits"})}, }}, }}, }; } - json invoke(json params) override { + json invoke(json params) const override { std::string path = params.at("path").get(); - const json & changes = params.at("changes"); + const json & edits_json = params.at("edits"); - if (!changes.is_array()) { - return {{"error", "\"changes\" must be an array"}}; + if (!edits_json.is_array() || edits_json.empty()) { + return {{"error", "\"edits\" must be a non-empty array"}}; } - // read file into lines - std::ifstream fin(path); - if (!fin) { + struct edit_req { + std::string old_text; + std::string new_text; + }; + std::vector edits; + edits.reserve(edits_json.size()); + for (const auto & e : edits_json) { + edit_req er; + er.old_text = e.at("old_text").get(); + er.new_text = e.at("new_text").get(); + if (er.old_text.empty()) { + return {{"error", string_format("edits[%zu].old_text must not be empty", edits.size())}}; + } + edits.push_back(std::move(er)); + } + + auto io = make_tools_io(params); + std::string original_content; + if (!io->read_file(path, original_content)) { return {{"error", "failed to open file: " + path}}; } - std::vector lines; - { - std::string line; - while (std::getline(fin, line)) { - lines.push_back(line); - } - } - fin.close(); - // validate and collect changes, then sort descending by line_start - struct change_entry { - std::string mode; - int line_start; // 1-based - int line_end; // 1-based inclusive - std::string content; - }; - std::vector entries; - entries.reserve(changes.size()); - - for (const auto & ch : changes) { - change_entry e; - e.mode = ch.at("mode").get(); - e.line_start = ch.at("line_start").get(); - e.line_end = ch.at("line_end").get(); - e.content = ch.at("content").get(); - - if (e.mode != "replace" && e.mode != "delete" && e.mode != "append") { - return {{"error", "invalid mode \"" + e.mode + "\"; must be replace, delete, or append"}}; + // does any old_text need fuzzy matching (no exact match found)? + bool any_fuzzy = false; + for (size_t i = 0; i < edits.size(); i++) { + if (original_content.find(edits[i].old_text) != std::string::npos) continue; + std::string fuzzy_content = normalize_for_fuzzy_match(original_content); + std::string fuzzy_old = normalize_for_fuzzy_match(edits[i].old_text); + if (fuzzy_content.find(fuzzy_old) == std::string::npos) { + return {{"error", string_format( + "could not find edits[%zu].old_text in %s, it must match the file's current content exactly", + i, path.c_str())}}; } - if (e.mode == "delete" && !e.content.empty()) { - return {{"error", "content must be empty string for delete mode"}}; - } - int n = (int) lines.size(); - if (e.line_start == -1) { - // -1 targets end of file -> valid for append only; line_end is ignored - if (e.mode != "append") { - return {{"error", "line_start -1 (end of file) is only valid for append mode"}}; - } - // append at end of file: insert position is the current line count - e.line_start = n; - e.line_end = n; - } else { - if (e.line_start < 1 || e.line_end < e.line_start) { - return {{"error", string_format("invalid line range [%d, %d]", e.line_start, e.line_end)}}; - } - if (e.line_end > n) { - return {{"error", string_format("line_end %d exceeds file length %d", e.line_end, n)}}; - } - } - entries.push_back(std::move(e)); + any_fuzzy = true; } - // sort descending so earlier-indexed changes don't shift later ones - std::sort(entries.begin(), entries.end(), [](const change_entry & a, const change_entry & b) { - return a.line_start > b.line_start; + std::string base_content = any_fuzzy ? normalize_for_fuzzy_match(original_content) : original_content; + + // uniqueness check always uses fuzzy-normalized text, so a whitespace-only duplicate still counts + std::vector matched; + matched.reserve(edits.size()); + for (size_t i = 0; i < edits.size(); i++) { + std::string needle = any_fuzzy ? normalize_for_fuzzy_match(edits[i].old_text) : edits[i].old_text; + size_t occurrences = count_occurrences( + normalize_for_fuzzy_match(original_content), + normalize_for_fuzzy_match(edits[i].old_text)); + if (occurrences > 1) { + return {{"error", string_format( + "found %zu occurrences of edits[%zu].old_text in %s, it must be unique", + occurrences, i, path.c_str())}}; + } + size_t idx = base_content.find(needle); + matched.push_back({i, idx, needle.size(), edits[i].new_text}); + } + + std::sort(matched.begin(), matched.end(), [](const matched_edit & a, const matched_edit & b) { + return a.match_index < b.match_index; }); - - // apply changes (0-based indices internally) - for (const auto & e : entries) { - int idx_start = e.line_start - 1; // 0-based - int idx_end = e.line_end - 1; // 0-based inclusive - - // split content into lines (preserve trailing newline awareness) - std::vector new_lines; - if (!e.content.empty()) { - std::istringstream ss(e.content); - std::string ln; - while (std::getline(ss, ln)) { - new_lines.push_back(ln); - } - // if content ends with \n, getline consumed it — no extra empty line needed - // if content does NOT end with \n, last line is still captured correctly - } - - if (e.mode == "replace") { - // erase [idx_start, idx_end] and insert new_lines - lines.erase(lines.begin() + idx_start, lines.begin() + idx_end + 1); - lines.insert(lines.begin() + idx_start, new_lines.begin(), new_lines.end()); - } else if (e.mode == "delete") { - lines.erase(lines.begin() + idx_start, lines.begin() + idx_end + 1); - } else { // append - // insert after idx_end; idx_end + 1 == lines.size() for end-of-file append - lines.insert(lines.begin() + (idx_end + 1), new_lines.begin(), new_lines.end()); + for (size_t i = 1; i < matched.size(); i++) { + if (matched[i - 1].match_index + matched[i - 1].match_length > matched[i].match_index) { + return {{"error", string_format( + "edits[%zu] and edits[%zu] overlap in %s; merge them into one edit or target disjoint regions", + matched[i - 1].edit_index, matched[i].edit_index, path.c_str())}}; } } - // write file back - std::ofstream fout(path, std::ios::binary); - if (!fout) { - return {{"error", "failed to open file for writing: " + path}}; + std::string new_content = any_fuzzy + ? apply_replacements_preserving_unchanged_lines(original_content, base_content, matched) + : apply_replacements(base_content, matched, 0); + + if (new_content == original_content) { + return {{"error", "no changes made: the replacement(s) produced identical content"}}; } - for (size_t i = 0; i < lines.size(); i++) { - fout << lines[i]; - if (i + 1 < lines.size()) { - fout << "\n"; - } - } - if (!lines.empty()) { - fout << "\n"; - } - if (!fout) { + + if (!io->write_file(path, new_content)) { return {{"error", "failed to write file: " + path}}; } - return {{"result", "file edited successfully"}, {"path", path}, {"lines", (int) lines.size()}}; - } -}; - -// -// apply_diff: apply a unified diff via git apply -// - -struct server_tool_apply_diff : server_tool { - server_tool_apply_diff() { - name = "apply_diff"; - display_name = "Apply diff"; - permission_write = true; + return {{"result", "file edited successfully"}, {"path", path}, {"edits_applied", (int) matched.size()}}; } - json get_definition() override { - return { - {"type", "function"}, - {"function", { - {"name", name}, - {"description", "Apply a unified diff to edit one or more files using git apply. Use this instead of edit_file when the changes are complex."}, - {"parameters", { - {"type", "object"}, - {"properties", { - {"diff", {{"type", "string"}, {"description", "Unified diff content in git diff format"}}}, - }}, - {"required", json::array({"diff"})}, - }}, - }}, - }; - } +private: + // strip trailing whitespace, normalize smart quotes/dashes/spaces to ASCII + static std::string normalize_line_for_fuzzy_match(const std::string & line) { + size_t end = line.size(); + while (end > 0 && (line[end - 1] == ' ' || line[end - 1] == '\t' || line[end - 1] == '\r')) { + end--; + } + std::string s = line.substr(0, end); - json invoke(json params) override { - std::string diff = params.at("diff").get(); - - // write diff to a temporary file - static std::atomic counter{0}; - std::string tmp_path = (fs::temp_directory_path() / - ("llama_patch_" + std::to_string(++counter) + ".patch")).string(); - - { - std::ofstream f(tmp_path, std::ios::binary); - if (!f) { - return {{"error", "failed to create temp patch file"}}; + auto replace_all = [](std::string & str, const std::string & from, const std::string & to) { + if (from.empty()) return; + size_t pos = 0; + while ((pos = str.find(from, pos)) != std::string::npos) { + str.replace(pos, from.size(), to); + pos += to.size(); + } + }; + + // smart single quotes -> ' + for (unsigned char b : {0x98, 0x99, 0x9A, 0x9B}) { + replace_all(s, std::string("\xE2\x80") + (char) b, "'"); + } + // smart double quotes -> " + for (unsigned char b : {0x9C, 0x9D, 0x9E, 0x9F}) { + replace_all(s, std::string("\xE2\x80") + (char) b, "\""); + } + // various dashes -> - + for (unsigned char b = 0x90; b <= 0x95; b++) { + replace_all(s, std::string("\xE2\x80") + (char) b, "-"); + } + replace_all(s, "\xE2\x88\x92", "-"); // minus sign + // special spaces -> ' ' + replace_all(s, "\xC2\xA0", " "); // no-break space + for (unsigned char b = 0x82; b <= 0x8A; b++) { + replace_all(s, std::string("\xE2\x80") + (char) b, " "); + } + replace_all(s, "\xE2\x80\xAF", " "); // narrow no-break space + replace_all(s, "\xE2\x81\x9F", " "); // medium mathematical space + replace_all(s, "\xE3\x80\x80", " "); // ideographic space + + return s; + } + + // applies the per-line transform above to every line; preserves line count/positions + static std::string normalize_for_fuzzy_match(const std::string & content) { + std::string result; + result.reserve(content.size()); + size_t start = 0; + while (true) { + size_t nl = content.find('\n', start); + bool is_last = nl == std::string::npos; + std::string line = is_last ? content.substr(start) : content.substr(start, nl - start); + result += normalize_line_for_fuzzy_match(line); + if (is_last) break; + result += '\n'; + start = nl + 1; + } + return result; + } + + // lines with trailing '\n' kept, so untouched ones can be reconstructed verbatim + static std::vector split_lines_with_endings(const std::string & content) { + std::vector lines; + size_t start = 0; + while (start < content.size()) { + size_t nl = content.find('\n', start); + if (nl == std::string::npos) { + lines.push_back(content.substr(start)); + break; + } + lines.push_back(content.substr(start, nl - start + 1)); + start = nl + 1; + } + return lines; + } + + struct line_span { + size_t start; + size_t end; + }; + + static std::vector get_line_spans(const std::string & content) { + std::vector spans; + size_t offset = 0; + for (const auto & line : split_lines_with_endings(content)) { + spans.push_back({offset, offset + line.size()}); + offset += line.size(); + } + return spans; + } + + // count non-overlapping occurrences of `needle` in `content` + static size_t count_occurrences(const std::string & content, const std::string & needle) { + if (needle.empty()) return 0; + size_t count = 0, pos = 0; + while ((pos = content.find(needle, pos)) != std::string::npos) { + count++; + pos += needle.size(); + } + return count; + } + + struct matched_edit { + size_t edit_index; + size_t match_index; // offset into the "base content" (see below) + size_t match_length; + std::string new_text; + }; + + // replacements must be sorted ascending by match_index and non-overlapping + static std::string apply_replacements( + const std::string & content, + const std::vector & replacements, + size_t offset) { + std::string result = content; + for (auto it = replacements.rbegin(); it != replacements.rend(); ++it) { + size_t local_index = it->match_index - offset; + result = result.substr(0, local_index) + it->new_text + result.substr(local_index + it->match_length); + } + return result; + } + + // widen a replacement's byte range to the line(s) of `lines` it touches + static bool get_replacement_line_range( + const std::vector & lines, + size_t match_index, size_t match_length, + size_t & out_start_line, size_t & out_end_line /* exclusive */) { + size_t replacement_start = match_index; + size_t replacement_end = match_index + match_length; + + size_t start_line = (size_t) -1; + for (size_t i = 0; i < lines.size(); i++) { + if (replacement_start >= lines[i].start && replacement_start < lines[i].end) { + start_line = i; + break; + } + } + if (start_line == (size_t) -1) return false; + + size_t end_line = start_line; + while (end_line < lines.size() && lines[end_line].end < replacement_end) { + end_line++; + } + if (end_line >= lines.size()) return false; + + out_start_line = start_line; + out_end_line = end_line + 1; + return true; + } + + // like apply_replacements, but untouched lines come from `original_content` + static std::string apply_replacements_preserving_unchanged_lines( + const std::string & original_content, + const std::string & base_content, + const std::vector & replacements /* ascending, non-overlapping */) { + auto original_lines = split_lines_with_endings(original_content); + auto base_lines = get_line_spans(base_content); + + struct group { + size_t start_line; + size_t end_line; // exclusive + std::vector reps; + }; + std::vector groups; + + for (const auto & rep : replacements) { + size_t start_line = 0, end_line = 0; + get_replacement_line_range(base_lines, rep.match_index, rep.match_length, start_line, end_line); + if (!groups.empty() && start_line < groups.back().end_line) { + groups.back().end_line = std::max(groups.back().end_line, end_line); + groups.back().reps.push_back(rep); + } else { + groups.push_back({start_line, end_line, {rep}}); } - f << diff; } - auto res = run_process({"git", "apply", tmp_path}, 4096, 10); + size_t original_line_index = 0; + std::string result; + for (auto & g : groups) { + for (size_t i = original_line_index; i < g.start_line; i++) { + result += original_lines[i]; + } - std::error_code ec; - fs::remove(tmp_path, ec); + size_t group_start_offset = base_lines[g.start_line].start; + size_t group_end_offset = base_lines[g.end_line - 1].end; + std::string slice = base_content.substr(group_start_offset, group_end_offset - group_start_offset); + result += apply_replacements(slice, g.reps, group_start_offset); - if (res.exit_code != 0) { - return {{"error", "git apply failed (exit " + std::to_string(res.exit_code) + "): " + res.output}}; + original_line_index = g.end_line; } - return {{"result", "patch applied successfully"}}; + for (size_t i = original_line_index; i < original_lines.size(); i++) { + result += original_lines[i]; + } + + return result; } }; @@ -710,7 +1008,7 @@ struct server_tool_get_datetime : server_tool { permission_write = false; } - json get_definition() override { + json get_definition() const override { return { {"type", "function"}, {"function", { @@ -720,7 +1018,7 @@ struct server_tool_get_datetime : server_tool { }; } - json invoke(json) override { + json invoke(json) const override { auto now = std::chrono::system_clock::now(); auto time = std::chrono::system_clock::to_time_t(now); @@ -740,7 +1038,6 @@ static std::vector> build_tools() { tools.push_back(std::make_unique()); tools.push_back(std::make_unique()); tools.push_back(std::make_unique()); - tools.push_back(std::make_unique()); tools.push_back(std::make_unique()); return tools; } diff --git a/tools/server/server-tools.h b/tools/server/server-tools.h index 444ef5f80..cc147379b 100644 --- a/tools/server/server-tools.h +++ b/tools/server/server-tools.h @@ -9,10 +9,10 @@ struct server_tool { bool permission_write = false; virtual ~server_tool() = default; - virtual json get_definition() = 0; - virtual json invoke(json params) = 0; + virtual json get_definition() const = 0; + virtual json invoke(json params) const = 0; - json to_json(); + json to_json() const; }; struct server_tools { diff --git a/tools/server/tests/unit/test_tools_builtin.py b/tools/server/tests/unit/test_tools_builtin.py new file mode 100755 index 000000000..d4fd5dc9b --- /dev/null +++ b/tools/server/tests/unit/test_tools_builtin.py @@ -0,0 +1,125 @@ +import os + +import pytest +from utils import * + +server: ServerProcess + +# project root, used as the search directory for grep_search/file_glob_search +PROJECT_ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "..", "..", "..")) + +# marker for the grep_search test to find in this file +GREP_MARKER = "llama_cpp_test_tools_builtin_marker_grep_search" + + +@pytest.fixture(autouse=True) +def create_server(): + global server + server = ServerPreset.router() + server.server_tools = "all" + + +def call_tool(name: str, params: dict) -> dict: + res = server.make_request("POST", "/tools", data={"tool": name, "params": params}) + assert res.status_code == 200, res.body + assert "error" not in res.body, res.body + return res.body + + +def call_tool_expect_error(name: str, params: dict) -> str: + res = server.make_request("POST", "/tools", data={"tool": name, "params": params}) + assert res.status_code == 200, res.body + assert "error" in res.body, res.body + return res.body["error"] + + +def test_tools_builtin_grep_search(): + global server + server.start() + + res = call_tool("grep_search", { + "path": PROJECT_ROOT, + "pattern": GREP_MARKER, + "include": "test_tools_builtin.py", # bare pattern -> matches basename at any depth + }) + text = res["plain_text_response"] + assert "test_tools_builtin.py" in text + assert GREP_MARKER in text + assert "Total matches: 1" in text + + +def test_tools_builtin_read_file(): + global server + server.start() + + this_file = os.path.join(PROJECT_ROOT, "tools", "server", "tests", "unit", "test_tools_builtin.py") + res = call_tool("read_file", {"path": this_file}) + text = res["plain_text_response"] + assert GREP_MARKER in text + assert "def test_tools_builtin_read_file" in text + + +def test_tools_builtin_write_then_edit_file(): + global server + server.start() + + log_path = os.path.join(PROJECT_ROOT, "test.log") + try: + write_res = call_tool("write_file", {"path": log_path, "content": "line1\nline2\nline3\n"}) + assert write_res["result"] == "file written successfully" + + read_before = call_tool("read_file", {"path": log_path}) + assert read_before["plain_text_response"] == "line1\nline2\nline3\n" + + edit_res = call_tool("edit_file", { + "path": log_path, + "edits": [ + {"old_text": "line2", "new_text": "line2-edited"}, + {"old_text": "line3\n", "new_text": "line3\nline4\n"}, + ], + }) + assert edit_res["result"] == "file edited successfully" + assert edit_res["edits_applied"] == 2 + + read_after = call_tool("read_file", {"path": log_path}) + assert read_after["plain_text_response"] == "line1\nline2-edited\nline3\nline4\n" + finally: + if os.path.exists(log_path): + os.remove(log_path) + + +def test_tools_builtin_edit_file_rejects_non_unique_old_text(): + global server + server.start() + + log_path = os.path.join(PROJECT_ROOT, "test.log") + try: + call_tool("write_file", {"path": log_path, "content": "dup\ndup\n"}) + err = call_tool_expect_error("edit_file", { + "path": log_path, + "edits": [{"old_text": "dup", "new_text": "changed"}], + }) + assert "unique" in err + finally: + if os.path.exists(log_path): + os.remove(log_path) + + +def test_tools_builtin_edit_file_rejects_overlapping_edits(): + global server + server.start() + + log_path = os.path.join(PROJECT_ROOT, "test.log") + try: + call_tool("write_file", {"path": log_path, "content": "line1\nline2\n"}) + err = call_tool_expect_error("edit_file", { + "path": log_path, + "edits": [ + {"old_text": "line1\nline2", "new_text": "a"}, + {"old_text": "line2", "new_text": "b"}, + ], + }) + assert "overlap" in err + finally: + if os.path.exists(log_path): + os.remove(log_path) diff --git a/tools/server/tests/utils.py b/tools/server/tests/utils.py index 67d7d20db..8c0de384f 100644 --- a/tools/server/tests/utils.py +++ b/tools/server/tests/utils.py @@ -113,6 +113,7 @@ class ServerProcess: ui_mcp_proxy: bool = False backend_sampling: bool = False gcp_compat: bool = False + server_tools: str | None = None # session variables process: subprocess.Popen | None = None @@ -256,6 +257,8 @@ class ServerProcess: server_args.append("--no-cache-idle-slots") if self.ui_mcp_proxy: server_args.append("--ui-mcp-proxy") + if self.server_tools: + server_args.extend(["--tools", self.server_tools]) if self.backend_sampling: server_args.append("--backend_sampling") if self.gcp_compat: diff --git a/tools/ui/embed.cpp b/tools/ui/embed.cpp index cdbb64232..914d51fa1 100644 --- a/tools/ui/embed.cpp +++ b/tools/ui/embed.cpp @@ -187,7 +187,6 @@ int main(int argc, char ** argv) { struct required_check { const char * label; match_fn match; bool found; }; required_check checks[] = { { "index.html", exact("index.html"), false }, - { "loading.html", exact("loading.html"), false }, { "manifest.webmanifest", exact("manifest.webmanifest"), false }, { "sw.js", exact("sw.js"), false }, { "build.json", exact("build.json"), false }, diff --git a/tools/ui/src/lib/components/app/chat/ChatScreen/ChatScreenServerError.svelte b/tools/ui/src/lib/components/app/chat/ChatScreen/ChatScreenServerError.svelte index 2a998dbeb..5345f19a5 100644 --- a/tools/ui/src/lib/components/app/chat/ChatScreen/ChatScreenServerError.svelte +++ b/tools/ui/src/lib/components/app/chat/ChatScreen/ChatScreenServerError.svelte @@ -1,10 +1,11 @@ {#if hasError} @@ -12,23 +13,31 @@ class="pointer-events-auto mx-auto mb-4 max-w-[48rem] px-1" use:fadeInView={{ y: 10, duration: 250 }} > - - + + {#if isLoadingModel} + + {:else} + + {/if} - Server unavailable + {isLoadingModel ? 'Loading model' : 'Server unavailable'} - + {#if !isLoadingModel} + + {/if} - {serverError()} + {#if !isLoadingModel} + {serverError()} + {/if} {/if} diff --git a/tools/ui/src/lib/components/ui/tooltip/tooltip-content.svelte b/tools/ui/src/lib/components/ui/tooltip/tooltip-content.svelte index 5b0c76818..0b173ee7c 100644 --- a/tools/ui/src/lib/components/ui/tooltip/tooltip-content.svelte +++ b/tools/ui/src/lib/components/ui/tooltip/tooltip-content.svelte @@ -5,7 +5,7 @@ let { ref = $bindable(null), class: className, - sideOffset = 0, + sideOffset = 4, side = 'top', children, arrowClasses, diff --git a/tools/ui/src/lib/constants/pwa.ts b/tools/ui/src/lib/constants/pwa.ts index b9171d4bb..343bcaf3c 100644 --- a/tools/ui/src/lib/constants/pwa.ts +++ b/tools/ui/src/lib/constants/pwa.ts @@ -258,12 +258,6 @@ export const GLOB_PATTERNS: string[] = [ '**/*.{js,css,html,ico,svg,png,webp,woff,woff2,json,webmanifest}' ]; -// loading.html is the model loading page served by llama-server itself. -// The SvelteKit PWA manifest transform strips the html extension from every -// precache entry to match clean URLs, but loading.html is a plain static asset -// with no clean URL, so static servers answer 404 and the SW install fails. -export const GLOB_IGNORES: string[] = ['**/loading.html']; - export const SW_CONFIG = { CHECK_INTERVAL_MS: 60000, UPDATE_FETCH_OPTIONS: { @@ -317,7 +311,6 @@ export const SVELTEKIT_PWA_OPTIONS: SvelteKitPWAOptions = { // Uses '**/' because SvelteKit outputs files under _app/immutable/ // subdirectories. globPatterns: GLOB_PATTERNS, - globIgnores: GLOB_IGNORES, maximumFileSizeToCacheInBytes: CACHE_SETTINGS.MAX_FILE_SIZE_BYTES, // Prevent @vite-pwa/sveltekit from auto-adding a NavigationRoute by diff --git a/tools/ui/src/lib/stores/conversations.svelte.ts b/tools/ui/src/lib/stores/conversations.svelte.ts index b29a900fe..47d2f060d 100644 --- a/tools/ui/src/lib/stores/conversations.svelte.ts +++ b/tools/ui/src/lib/stores/conversations.svelte.ts @@ -1115,21 +1115,18 @@ class ConversationsStore { } /** - * Downloads a conversation as JSON file. + * Downloads a single conversation as a JSONL file, serializing the full message tree. * @param convId - The conversation ID to download */ async downloadConversation(convId: string): Promise { - let conversation: DatabaseConversation | null; - let messages: DatabaseMessage[]; + const conversation = + this.activeConversation?.id === convId + ? this.activeConversation + : await DatabaseService.getConversation(convId); - if (this.activeConversation?.id === convId) { - conversation = this.activeConversation; - messages = this.activeMessages; - } else { - conversation = await DatabaseService.getConversation(convId); - if (!conversation) return; - messages = await DatabaseService.getConversationMessages(convId); - } + if (!conversation) return; + + const messages = await DatabaseService.getConversationMessages(convId); this.downloadConversationFile({ conv: conversation, messages }); } diff --git a/tools/ui/src/lib/stores/models.svelte.ts b/tools/ui/src/lib/stores/models.svelte.ts index 11316a193..4c7c7cdfe 100644 --- a/tools/ui/src/lib/stores/models.svelte.ts +++ b/tools/ui/src/lib/stores/models.svelte.ts @@ -145,6 +145,10 @@ class ModelsStore { */ getModelModalities(modelId: string): ModelModalities | null { + if (!isRouterMode() && serverStore.props?.modalities) { + return this.buildModalities(serverStore.props.modalities); + } + const model = this.models.find((m) => m.model === modelId || m.id === modelId); if (model?.modalities) { return model.modalities; @@ -629,7 +633,12 @@ class ModelsStore { } findModelByName(modelName: string): ModelOption | null { - return this.models.find((model) => model.model === modelName) ?? null; + return ( + this.models.find( + (model) => + model.model === modelName || model.id === modelName || model.aliases?.includes(modelName) + ) ?? null + ); } findModelById(modelId: string): ModelOption | null { diff --git a/tools/ui/src/lib/stores/server.svelte.ts b/tools/ui/src/lib/stores/server.svelte.ts index d9a9f855a..66ab41119 100644 --- a/tools/ui/src/lib/stores/server.svelte.ts +++ b/tools/ui/src/lib/stores/server.svelte.ts @@ -1,5 +1,8 @@ import { PropsService } from '$lib/services/props.service'; import { ServerRole } from '$lib/enums'; +import { ApiError } from '$lib/utils/api-fetch'; + +const LOADING_RETRY_INTERVAL_MS = 1000; /** * serverStore - Server connection state, configuration, and role detection @@ -29,8 +32,10 @@ class ServerStore { props = $state(null); loading = $state(false); error = $state(null); + status = $state(null); role = $state(null); private fetchPromise: Promise | null = null; + private retryTimer: ReturnType | null = null; /** * @@ -70,23 +75,43 @@ class ServerStore { * */ - async fetch(): Promise { + /** + * @param background - Set by the automatic "still loading" poll. Skips the + * `loading` flag flip so the UI doesn't bounce between the full loading + * splash and the chat screen every retry tick. + */ + async fetch({ background = false }: { background?: boolean } = {}): Promise { if (this.fetchPromise) return this.fetchPromise; - this.loading = true; - this.error = null; + this.clearRetryTimer(); + if (!background) { + this.loading = true; + } + // Don't clear an existing "still loading" error before a retry - + // doing so would unmount/remount the error banner every second. + if (this.status !== 503) { + this.error = null; + } const fetchPromise = (async () => { try { const props = await PropsService.fetch(); this.props = props; this.error = null; + this.status = null; this.detectRole(props); } catch (error: unknown) { this.error = error instanceof Error ? error.message : String(error); + this.status = error instanceof ApiError ? error.status : null; console.error('Error fetching server properties:', error); + + if (this.status === 503) { + this.scheduleRetry(); + } } finally { - this.loading = false; + if (!background) { + this.loading = false; + } this.fetchPromise = null; } })(); @@ -96,13 +121,30 @@ class ServerStore { } clear(): void { + this.clearRetryTimer(); this.props = null; this.error = null; + this.status = null; this.loading = false; this.role = null; this.fetchPromise = null; } + private scheduleRetry(): void { + if (this.retryTimer) return; + this.retryTimer = setTimeout(() => { + this.retryTimer = null; + this.fetch({ background: true }); + }, LOADING_RETRY_INTERVAL_MS); + } + + private clearRetryTimer(): void { + if (this.retryTimer) { + clearTimeout(this.retryTimer); + this.retryTimer = null; + } + } + /** * * @@ -125,6 +167,7 @@ export const serverStore = new ServerStore(); export const serverProps = () => serverStore.props; export const serverLoading = () => serverStore.loading; export const serverError = () => serverStore.error; +export const serverStatus = () => serverStore.status; export const serverRole = () => serverStore.role; export const defaultParams = () => serverStore.defaultParams; export const contextSize = () => serverStore.contextSize; diff --git a/tools/ui/src/lib/utils/api-fetch.ts b/tools/ui/src/lib/utils/api-fetch.ts index 82a9383dd..e9d906258 100644 --- a/tools/ui/src/lib/utils/api-fetch.ts +++ b/tools/ui/src/lib/utils/api-fetch.ts @@ -12,6 +12,21 @@ import { ERROR_MESSAGES, HTTP_CODE_TO_STRING } from '$lib/constants/error'; * - Base path resolution */ +/** + * Error thrown when an API request fails, carrying the HTTP status code + * so callers can distinguish e.g. a 503 "still loading" response from a + * genuine failure. + */ +export class ApiError extends Error { + status: number; + + constructor(message: string, status: number) { + super(message); + this.name = 'ApiError'; + this.status = status; + } +} + export interface ApiFetchOptions extends Omit { /** * Use auth-only headers (no Content-Type). @@ -67,7 +82,7 @@ export async function apiFetch(path: string, options: ApiFetchOptions = {}): if (!response.ok) { const errorMessage = await parseErrorMessage(response); - throw new Error(errorMessage); + throw new ApiError(errorMessage, response.status); } return response.json() as Promise; @@ -119,7 +134,7 @@ export async function apiFetchWithParams( if (!response.ok) { const errorMessage = await parseErrorMessage(response); - throw new Error(errorMessage); + throw new ApiError(errorMessage, response.status); } return response.json() as Promise; diff --git a/tools/ui/static/loading.html b/tools/ui/static/loading.html deleted file mode 100644 index c3fd19a0f..000000000 --- a/tools/ui/static/loading.html +++ /dev/null @@ -1,12 +0,0 @@ - - - - - - -
- The model is loading. Please wait.
- The user interface will appear soon. -
- - diff --git a/tools/ui/tests/unit/pwa.spec.ts b/tools/ui/tests/unit/pwa.spec.ts index 17da27986..64b9630d2 100644 --- a/tools/ui/tests/unit/pwa.spec.ts +++ b/tools/ui/tests/unit/pwa.spec.ts @@ -189,9 +189,5 @@ describe('PWA Build Output', () => { expect(existsSync(resolve(DIST_DIR, 'pwa-192x192.png'))).toBeTruthy(); expect(existsSync(resolve(DIST_DIR, 'pwa-512x512.png'))).toBeTruthy(); }); - - it('has loading.html fallback page', () => { - expect(existsSync(resolve(DIST_DIR, 'loading.html'))).toBeTruthy(); - }); }); });