diff --git a/.cargo/config.toml b/.cargo/config.toml new file mode 100644 index 0000000..9f2e555 --- /dev/null +++ b/.cargo/config.toml @@ -0,0 +1,17 @@ +# Build flags shared by every developer and by CI. +# +# The assistant's quantized inference kernels (Candle) are only compiled when +# the target CPU has AVX2/FMA; without them the copilot is 3-4x slower. All +# x86-64 machines since 2013 (Intel Haswell, AMD Excavator/Zen) support the +# x86-64-v3 level. Apple Silicon (aarch64) always has NEON and needs no flag. +[target.x86_64-unknown-linux-gnu] +rustflags = ["-C", "target-cpu=x86-64-v3"] + +[target.x86_64-pc-windows-msvc] +rustflags = ["-C", "target-cpu=x86-64-v3"] + +[target.x86_64-pc-windows-gnu] +rustflags = ["-C", "target-cpu=x86-64-v3"] + +[target.x86_64-apple-darwin] +rustflags = ["-C", "target-cpu=x86-64-v3"] diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index ae59494..fd9c496 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -21,7 +21,13 @@ jobs: with: python-version: "3.12" - name: Install IfcOpenShell + IfcTester (the IFC reader bridge) - run: python -m pip install --upgrade pip && python -m pip install ifcopenshell "ifctester" --no-deps && python -m pip install lxml xmlschema + # ifctester is installed without its optional reporting extras (odfpy has no wheel); + # everything the IDS engine itself needs is listed explicitly. + run: | + python -m pip install --upgrade pip + python -m pip install ifcopenshell typing_extensions numpy lxml xmlschema + python -m pip install --no-deps ifctester + python -c "import ifcopenshell, ifctester; print('ifcopenshell', ifcopenshell.version)" - name: Install Tauri system dependencies (Linux) run: | sudo apt-get update diff --git a/.github/workflows/release.yml b/.github/workflows/release.yml index 06a6062..41e295d 100644 --- a/.github/workflows/release.yml +++ b/.github/workflows/release.yml @@ -7,6 +7,12 @@ name: Release on: push: tags: ["v*"] + workflow_dispatch: + inputs: + version: + description: "Version to release (e.g. 1.0.0). The tag v is created on the selected ref if it does not exist." + required: true + type: string permissions: contents: write @@ -76,8 +82,19 @@ jobs: merge-multiple: true - name: Checksums run: cd dist && sha256sum * > SHA256SUMS.txt && cat SHA256SUMS.txt + - name: Resolve tag name + id: tag + shell: bash + run: | + if [ "${{ github.event_name }}" = "workflow_dispatch" ]; then + echo "name=v${{ inputs.version }}" >> "$GITHUB_OUTPUT" + else + echo "name=${{ github.ref_name }}" >> "$GITHUB_OUTPUT" + fi - uses: softprops/action-gh-release@v2 with: - name: VeriModel ${{ github.ref_name }} + tag_name: ${{ steps.tag.outputs.name }} + target_commitish: ${{ github.sha }} + name: VeriModel ${{ steps.tag.outputs.name }} body_path: docs/RELEASE_NOTES.md files: dist/* diff --git a/README.md b/README.md index 28864a8..69f4a59 100644 --- a/README.md +++ b/README.md @@ -50,6 +50,7 @@ Not a bot that clicks. Not a flow builder. A quality gate a model must cross wit | Need | Why | Install | |---|---|---| | Rust 1.80+ | engine, CLI, desktop shell | | +| x86-64 CPU with AVX2 (2013+) or Apple Silicon | the assistant's inference kernels (`.cargo/config.toml` targets x86-64-v3) | — | | Python 3.9+ with `ifcopenshell` and `ifctester` | the IFC reader bridge (parsing, geometry, IDS) | `pip install ifcopenshell ifctester` | | Tauri v2 system deps (Linux only) | desktop app | `libwebkit2gtk-4.1-dev libayatana-appindicator3-dev librsvg2-dev patchelf` | | `cargo tauri` CLI (only to build the desktop app) | bundling | `cargo install tauri-cli --version "^2" --locked` | @@ -330,7 +331,7 @@ Standards interpretations made where the specification leaves room, also listed - Base quantities are trusted as declared; VeriModel does not recompute them when they disagree with the geometry (it reports the source of every value so you can). - BCF snapshots (PNG) are not generated. - The UI has been exercised in a browser with mock data and through the Tauri commands' tests; a full packaged-app UI test harness is not included. -- The assistant is a 0.5B-parameter model: good at finding, explaining and drafting from the retrieved context, not a substitute for reading the report. It runs on CPU at a few tokens per second per core; larger GGUF models can be dropped into the same folder. +- The assistant is a 0.5B-parameter model: good at finding, explaining and drafting from the retrieved context, not a substitute for reading the report. Measured on a 4-core CPU with the default build flags (AVX2 kernels): model load 1.5 s, prompt processing 33–38 tokens/s, generation 12–15 tokens/s, so a question takes roughly 30–50 s, almost all of it reading the prompt (Candle processes prompts one token row at a time; caching the fixed prefix is the planned fix). RAM: about 1.0 GB resident, 1.4 GB peak. Builds without AVX2 (`-C target-cpu`, see `.cargo/config.toml`) are 3–4× slower. Larger GGUF models can be dropped into the same folder. ## License diff --git a/crates/assistant/prompts/ask.txt b/crates/assistant/prompts/ask.txt new file mode 100644 index 0000000..3102194 --- /dev/null +++ b/crates/assistant/prompts/ask.txt @@ -0,0 +1,7 @@ +[system] +Task: answer the user's question about VeriModel, its rules, its configuration or the current report. +JSON schema: +{"answer": string} +Instructions: +- "answer": 1 to 6 plain sentences based only on the context. Name rule ids, property sets and keys exactly as written there. +- If the context does not contain the answer, say that you do not know and suggest where to look. diff --git a/crates/assistant/prompts/config_help.txt b/crates/assistant/prompts/config_help.txt new file mode 100644 index 0000000..7f3b413 --- /dev/null +++ b/crates/assistant/prompts/config_help.txt @@ -0,0 +1,18 @@ +[system] +Task: help the user write VeriModel project.yaml configuration. +JSON schema: +{"yaml": string, "explanation": string} +Instructions: +- "yaml": a project.yaml snippet with only the keys needed. Use only documented keys. Top-level keys: project, requirements, tolerances, qto, clash, output, checks, naming, classification. +- A required property goes under requirements.required_psets as an entry with pset, applies_to (IFC types), properties. +- Standard properties live in the buildingSMART common property set of the type: walls (IfcWall) use Pset_WallCommon, doors (IfcDoor) Pset_DoorCommon, slabs (IfcSlab) Pset_SlabCommon, columns (IfcColumn) Pset_ColumnCommon, windows (IfcWindow) Pset_WindowCommon. +- Naming rules go under naming.patterns (IFC type: regex). +- "explanation": one or two sentences. +[user] +Question: columns must have LoadBearing +[assistant] +{"yaml": "requirements:\n required_psets:\n - pset: Pset_ColumnCommon\n applies_to: [IfcColumn]\n properties: [LoadBearing]\n", "explanation": "Every IfcColumn must have Pset_ColumnCommon with the property LoadBearing; a missing value is reported as an error."} +[user] +Question: door names must look like D-001 +[assistant] +{"yaml": "naming:\n patterns:\n IfcDoor: \"^D-[0-9]{3}$\"\n", "explanation": "Door names must match the regular expression ^D-[0-9]{3}$."} diff --git a/crates/assistant/prompts/draft_rules.txt b/crates/assistant/prompts/draft_rules.txt new file mode 100644 index 0000000..99f5413 --- /dev/null +++ b/crates/assistant/prompts/draft_rules.txt @@ -0,0 +1,16 @@ +[system] +Task: turn requirement text from project documentation into a VeriModel project.yaml fragment. +JSON schema: +{"yaml": string, "explanation": string} +Instructions: +- "yaml" may only use these sections: + requirements.required_psets: list of {pset, applies_to: [IFC types], properties: [names], severity: error|warning|info} + naming.patterns: map of IFC type to regex + classification: {required: true/false, applies_to: [IFC types], code_pattern: regex} +- Only include requirements that the text states. Do not invent property names. +- "explanation": which sentences of the text each part comes from. +[user] +Requirements text: +All doors shall carry Pset_DoorCommon with FireRating and AcousticRating. Door names shall follow D-### (three digits). +[assistant] +{"yaml": "requirements:\n required_psets:\n - pset: Pset_DoorCommon\n applies_to: [IfcDoor]\n properties: [FireRating, AcousticRating]\nnaming:\n patterns:\n IfcDoor: \"^D-[0-9]{3}$\"\n", "explanation": "The first sentence gives the required door properties; the second gives the door naming pattern."} diff --git a/crates/assistant/prompts/explain.txt b/crates/assistant/prompts/explain.txt new file mode 100644 index 0000000..a16914a --- /dev/null +++ b/crates/assistant/prompts/explain.txt @@ -0,0 +1,19 @@ +[system] +Task: explain one finding of the VeriModel report to a BIM coordinator. +JSON schema: +{"text": string} +Instructions: +- "text": 2 to 5 plain sentences. Say what is wrong, on which element (its name and GlobalId), why it matters, and how to fix it in the authoring tool. +- Use only the finding, the rule and the context. Do not invent values. +[user] +Finding: +- rule: NAM-001 (severity WARNING) +- element: IfcDoor "Door 7" (GlobalId 3vB2YO$MX4xv5uCqZZG05x), storey Ground Floor +- message: Name "Door 7" does not match the naming pattern ^D-[0-9]{3}$ +- expected: ^D-[0-9]{3}$ +- actual: Door 7 +- fix hint: Rename the element to follow the project naming convention. + +Rule: NAM-001 [naming] Element names follow the project naming pattern. Status FAIL, 4 failures out of 30 elements. +[assistant] +{"text": "The door \"Door 7\" (GlobalId 3vB2YO$MX4xv5uCqZZG05x) on Ground Floor is named \"Door 7\", but the project requires door names like D-001. Consistent names let schedules and issue reports refer to the right element. Rename it in the authoring tool to the D-### pattern, for example D-007, and export the IFC again."} diff --git a/crates/assistant/prompts/filter.txt b/crates/assistant/prompts/filter.txt new file mode 100644 index 0000000..e0120f7 --- /dev/null +++ b/crates/assistant/prompts/filter.txt @@ -0,0 +1,42 @@ +[system] +Task: turn the user's request into a filter for the findings table. +JSON schema: +{"severity": [string], "rule_id": string or null, "storey": string or null, "ifc_type": string or null, "search": string or null, "explanation": string} +Instructions: +- Use only values listed under "Allowed values" and copy them exactly. +- Use null (or [] for severity) for every field the request does not mention. +- Set "rule_id" only when the request names a rule id; otherwise null. +- "errors" means severity "ERROR", "warnings" means "WARNING". +- Walls, doors, slabs... map to the matching IFC type (IfcWall, IfcDoor, IfcSlab...). +- "search" is free text to match in the finding message; usually null. +- "explanation": one short sentence describing the filter. +[user] +Allowed values: +rule_id: GEO-002, NAM-001, CLS-001 +storey: Ground Floor, Roof +ifc_type: IfcDoor, IfcWindow, IfcColumn +severity: ERROR, WARNING, INFO + +Request: warnings for windows on the roof +[assistant] +{"severity": ["WARNING"], "rule_id": null, "storey": "Roof", "ifc_type": "IfcWindow", "search": null, "explanation": "Warnings on windows (IfcWindow) on storey Roof."} +[user] +Allowed values: +rule_id: GEO-002, NAM-001, CLS-001 +storey: Ground Floor, Roof +ifc_type: IfcDoor, IfcWindow, IfcColumn +severity: ERROR, WARNING, INFO + +Request: muéstrame los fallos de NAM-001 +[assistant] +{"severity": [], "rule_id": "NAM-001", "storey": null, "ifc_type": null, "search": null, "explanation": "Hallazgos de la regla NAM-001."} +[user] +Allowed values: +rule_id: GEO-002, NAM-001, CLS-001 +storey: Ground Floor, Roof +ifc_type: IfcDoor, IfcWindow, IfcColumn +severity: ERROR, WARNING, INFO + +Request: errors for columns on the ground floor +[assistant] +{"severity": ["ERROR"], "rule_id": null, "storey": "Ground Floor", "ifc_type": "IfcColumn", "search": null, "explanation": "Errors on columns (IfcColumn) on storey Ground Floor."} diff --git a/crates/assistant/prompts/system.txt b/crates/assistant/prompts/system.txt new file mode 100644 index 0000000..bd29b2a --- /dev/null +++ b/crates/assistant/prompts/system.txt @@ -0,0 +1,7 @@ +You are the VeriModel assistant, a copilot inside a tool that checks IFC building models. +Rules you always follow: +- You only read. You never change the model, the configuration or the report. +- You never decide whether anything passes or fails. The VeriModel rule engine does that; repeat its results only as they appear in the context. +- Use only facts from the context and from the user's message. If the answer is not there, say that you do not know. +- Answer in the user's language. +- Answer ONLY with one JSON object matching the schema below. No text before or after it, no markdown. diff --git a/crates/assistant/src/backend/candle.rs b/crates/assistant/src/backend/candle.rs index 4ffae12..f3e437a 100644 --- a/crates/assistant/src/backend/candle.rs +++ b/crates/assistant/src/backend/candle.rs @@ -1,40 +1,448 @@ //! Candle backend: quantized GGUF inference in pure Rust (CPU; Metal/CUDA via -//! candle features later). STUB: to be implemented. +//! candle features later). //! //! Contract: //! * `CandleBackend::load(gguf, tokenizer_json)` reads the GGUF header to pick -//! the architecture (`general.architecture`: "qwen2" | "llama" | ...) and -//! builds the matching `candle_transformers::models::quantized_*` model. -//! * `generate` tokenises the ChatML prompt, runs greedy decoding (or sampling -//! when `temperature > 0`) up to `max_tokens`, stops on EOS or any `stop` +//! the architecture (`general.architecture`: "qwen2" | "llama") and builds +//! the matching `candle_transformers::models::quantized_*` model. +//! * `generate` takes a ChatML prompt (see [`super::chatml`]); for the Llama +//! family it is re-rendered with the Llama 3 chat template. It tokenises the +//! prompt, runs greedy decoding (or sampling when `temperature > 0`) with a +//! mild repeat penalty, up to `max_tokens`, stops on EOS or any `stop` //! string, and returns the decoded text without the stop sequence. -//! * Must never panic on a bad file: return `AssistantError::Load`. +//! * The KV cache is cleared at the start of every call: each answer is +//! computed from position 0, so the same prompt always gives the same text. +//! * Never panics on a bad file: returns `AssistantError::Load`. +//! +//! Performance: candle picks its AVX2/FMA quantized kernels at compile time +//! (`cfg(target_feature = "avx2")`). A default x86-64 build falls back to +//! scalar code and is several times slower; release builds should use +//! `-C target-cpu=x86-64-v3` (or `native` for local runs). use std::path::Path; -use super::{Backend, GenOptions}; +use candle_core::quantized::gguf_file; +use candle_core::{Device, Tensor}; +use candle_transformers::generation::LogitsProcessor; +use candle_transformers::models::{quantized_llama, quantized_qwen2}; +use tokenizers::Tokenizer; + +use super::{Backend, GenOptions, Message}; use crate::error::{AssistantError, AssistantResult}; +/// Repeat penalty applied to the last [`REPEAT_LAST_N`] generated tokens. +const REPEAT_PENALTY: f32 = 1.1; +const REPEAT_LAST_N: usize = 64; +/// Seed for sampling (only used when `temperature > 0`). +const SEED: u64 = 299_792_458; +/// Token strings treated as end-of-sequence when the tokenizer knows them. +const EOS_TOKENS: &[&str] = &[ + "<|im_end|>", + "<|endoftext|>", + "<|eot_id|>", + "<|end_of_text|>", +]; + +enum Weights { + Qwen2(quantized_qwen2::ModelWeights), + Llama(quantized_llama::ModelWeights), +} + +impl Weights { + fn forward(&mut self, x: &Tensor, pos: usize) -> candle_core::Result { + match self { + Weights::Qwen2(m) => m.forward(x, pos), + Weights::Llama(m) => m.forward(x, pos), + } + } + + fn clear_kv_cache(&mut self) { + match self { + Weights::Qwen2(m) => m.clear_kv_cache(), + Weights::Llama(m) => m.clear_kv_cache(), + } + } +} + +/// Timing of the last [`Backend::generate`] call (for diagnostics and benchmarks). +#[derive(Debug, Clone, Copy, PartialEq, Default)] +pub struct GenStats { + pub prompt_tokens: usize, + pub generated_tokens: usize, + /// Time to process the prompt (first forward pass), in seconds. + pub prefill_secs: f64, + /// Time spent generating the remaining tokens, in seconds. + pub decode_secs: f64, +} + +impl GenStats { + /// Prompt tokens processed per second. + pub fn prefill_tokens_per_sec(&self) -> f64 { + if self.prefill_secs > 0.0 { + self.prompt_tokens as f64 / self.prefill_secs + } else { + 0.0 + } + } + + /// Generated tokens per second, excluding the prefill. + pub fn decode_tokens_per_sec(&self) -> f64 { + if self.decode_secs > 0.0 { + self.generated_tokens.saturating_sub(1) as f64 / self.decode_secs + } else { + 0.0 + } + } +} + +/// A loaded GGUF model plus its tokenizer. pub struct CandleBackend { name: String, + family: String, + weights: Weights, + tokenizer: Tokenizer, + eos_ids: Vec, + context_length: usize, + device: Device, + last_stats: Option, } impl CandleBackend { + /// Load a GGUF model and its `tokenizer.json`. Never panics on a bad file. pub fn load(gguf: &Path, tokenizer: &Path) -> AssistantResult { - let _ = (gguf, tokenizer); - Err(AssistantError::Load( - "candle backend not implemented yet".into(), - )) + let gguf = gguf.to_path_buf(); + let tokenizer = tokenizer.to_path_buf(); + match std::panic::catch_unwind(move || Self::load_inner(&gguf, &tokenizer)) { + Ok(r) => r, + Err(_) => Err(AssistantError::Load( + "the model file is malformed (loader panicked)".into(), + )), + } + } + + fn load_inner(gguf: &Path, tokenizer_path: &Path) -> AssistantResult { + let mut file = std::fs::File::open(gguf) + .map_err(|e| AssistantError::Load(format!("cannot open {}: {e}", gguf.display())))?; + let content = gguf_file::Content::read(&mut file).map_err(|e| { + AssistantError::Load(format!("{} is not a valid GGUF file: {e}", gguf.display())) + })?; + let arch = content + .metadata + .get("general.architecture") + .and_then(|v| v.to_string().ok()) + .cloned() + .ok_or_else(|| { + AssistantError::Load("the GGUF file does not declare general.architecture".into()) + })?; + let name = content + .metadata + .get("general.name") + .and_then(|v| v.to_string().ok()) + .cloned() + .unwrap_or_else(|| arch.clone()); + let context_length = content + .metadata + .get(&format!("{arch}.context_length")) + .and_then(|v| v.to_u32().ok()) + .map(|v| v as usize) + .unwrap_or(4096); + let gguf_eos = content + .metadata + .get("tokenizer.ggml.eos_token_id") + .and_then(|v| v.to_u32().ok()); + let device = Device::Cpu; + let weights = match arch.as_str() { + "qwen2" => Weights::Qwen2( + quantized_qwen2::ModelWeights::from_gguf(content, &mut file, &device).map_err( + |e| AssistantError::Load(format!("cannot build the qwen2 model: {e}")), + )?, + ), + "llama" => Weights::Llama( + quantized_llama::ModelWeights::from_gguf(content, &mut file, &device).map_err( + |e| AssistantError::Load(format!("cannot build the llama model: {e}")), + )?, + ), + other => { + return Err(AssistantError::Load(format!( + "unsupported model architecture '{other}' (supported: qwen2, llama)" + ))) + } + }; + let tokenizer = Tokenizer::from_file(tokenizer_path).map_err(|e| { + AssistantError::Load(format!( + "cannot read tokenizer {}: {e}", + tokenizer_path.display() + )) + })?; + let mut eos_ids: Vec = EOS_TOKENS + .iter() + .filter_map(|t| tokenizer.token_to_id(t)) + .collect(); + if let Some(id) = gguf_eos { + if !eos_ids.contains(&id) { + eos_ids.push(id); + } + } + Ok(CandleBackend { + name: format!("candle:{name}"), + family: arch, + weights, + tokenizer, + eos_ids, + context_length, + device, + last_stats: None, + }) + } + + /// Model family from the GGUF header ("qwen2" | "llama"). + pub fn family(&self) -> &str { + &self.family + } + + /// Token counts and timings of the last generation. + pub fn last_stats(&self) -> Option { + self.last_stats } + + /// Render the ChatML prompt in the template this model family expects. + fn render_prompt(&self, prompt: &str) -> String { + if self.family == "llama" { + if let Some(messages) = parse_chatml(prompt) { + // Keep any text already written in the open assistant turn (a primed answer prefix). + let primed = prompt + .rsplit_once("<|im_start|>assistant\n") + .map(|(_, t)| t) + .filter(|t| !t.contains("<|im_end|>")) + .unwrap_or(""); + return format!("{}{primed}", llama3_chat(&messages)); + } + } + prompt.to_string() + } + + fn generate_inner(&mut self, prompt: &str, opts: &GenOptions) -> AssistantResult { + let gen_err = |e: candle_core::Error| AssistantError::Generation(e.to_string()); + let prompt = self.render_prompt(prompt); + let enc = self + .tokenizer + .encode(prompt.as_str(), false) + .map_err(|e| AssistantError::Generation(format!("tokenizer: {e}")))?; + let prompt_ids: Vec = enc.get_ids().to_vec(); + if prompt_ids.is_empty() { + return Ok(String::new()); + } + if prompt_ids.len() + opts.max_tokens > self.context_length { + return Err(AssistantError::Generation(format!( + "the prompt is too long ({} tokens, model context {})", + prompt_ids.len(), + self.context_length + ))); + } + let mut stops: Vec = opts.stop.clone(); + if self.family == "llama" { + stops.push("<|eot_id|>".into()); + } + let mut sampler = if opts.temperature > 0.0 { + Some(LogitsProcessor::new(SEED, Some(opts.temperature), None)) + } else { + None + }; + + self.weights.clear_kv_cache(); + self.last_stats = None; + let mut stats = GenStats { + prompt_tokens: prompt_ids.len(), + ..GenStats::default() + }; + let started = std::time::Instant::now(); + let mut generated: Vec = Vec::new(); + let mut text = String::new(); + let mut pos = 0usize; + let mut input = prompt_ids; + while generated.len() < opts.max_tokens { + let x = Tensor::new(input.as_slice(), &self.device) + .and_then(|t| t.unsqueeze(0)) + .map_err(gen_err)?; + let logits = self.weights.forward(&x, pos).map_err(gen_err)?; + if pos == 0 { + stats.prefill_secs = started.elapsed().as_secs_f64(); + } + pos += input.len(); + let logits = logits + .squeeze(0) + .and_then(|l| l.to_dtype(candle_core::DType::F32)) + .map_err(gen_err)?; + let logits = if generated.is_empty() { + logits + } else { + let start = generated.len().saturating_sub(REPEAT_LAST_N); + candle_transformers::utils::apply_repeat_penalty( + &logits, + REPEAT_PENALTY, + &generated[start..], + ) + .map_err(gen_err)? + }; + let next = match sampler.as_mut() { + Some(s) => s.sample(&logits).map_err(gen_err)?, + None => logits + .argmax(candle_core::D::Minus1) + .and_then(|t| t.to_scalar::()) + .map_err(gen_err)?, + }; + stats.generated_tokens += 1; + stats.decode_secs = started.elapsed().as_secs_f64() - stats.prefill_secs; + self.last_stats = Some(stats); + if self.eos_ids.contains(&next) { + break; + } + generated.push(next); + text = self + .tokenizer + .decode(&generated, false) + .map_err(|e| AssistantError::Generation(format!("tokenizer: {e}")))?; + if let Some(cut) = first_stop(&text, &stops) { + text.truncate(cut); + return Ok(text); + } + input = vec![next]; + } + Ok(text) + } +} + +/// Byte offset of the earliest stop string in `text`, if any. +fn first_stop(text: &str, stops: &[String]) -> Option { + stops + .iter() + .filter(|s| !s.is_empty()) + .filter_map(|s| text.find(s.as_str())) + .min() } impl Backend for CandleBackend { fn name(&self) -> &str { &self.name } - fn generate(&mut self, _prompt: &str, _opts: &GenOptions) -> AssistantResult { - Err(AssistantError::Generation( - "candle backend not implemented yet".into(), - )) + + fn generate(&mut self, prompt: &str, opts: &GenOptions) -> AssistantResult { + let r = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| { + self.generate_inner(prompt, opts) + })); + match r { + Ok(r) => r, + Err(_) => { + self.weights.clear_kv_cache(); + Err(AssistantError::Generation( + "inference failed unexpectedly".into(), + )) + } + } + } +} + +/// Parse a prompt produced by [`super::chatml`] back into messages. The final +/// open `assistant` turn is dropped. Returns `None` if the text is not ChatML. +pub fn parse_chatml(prompt: &str) -> Option> { + let mut out = Vec::new(); + let mut rest = prompt; + while let Some(start) = rest.find("<|im_start|>") { + rest = &rest[start + "<|im_start|>".len()..]; + let nl = rest.find('\n')?; + let role = match rest[..nl].trim() { + "system" => "system", + "user" => "user", + "assistant" => "assistant", + _ => return None, + }; + rest = &rest[nl + 1..]; + match rest.find("<|im_end|>") { + Some(end) => { + out.push(Message { + role, + content: rest[..end].to_string(), + }); + rest = &rest[end + "<|im_end|>".len()..]; + } + None => break, // the open assistant turn + } + } + if out.is_empty() { + None + } else { + Some(out) + } +} + +/// Llama 3.x chat template, ending with an open assistant header. +pub fn llama3_chat(messages: &[Message]) -> String { + let mut s = String::from("<|begin_of_text|>"); + for m in messages { + s.push_str("<|start_header_id|>"); + s.push_str(m.role); + s.push_str("<|end_header_id|>\n\n"); + s.push_str(m.content.trim()); + s.push_str("<|eot_id|>"); + } + s.push_str("<|start_header_id|>assistant<|end_header_id|>\n\n"); + s +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn chatml_roundtrip_to_llama3() { + let msgs = vec![ + Message { + role: "system", + content: "sys".into(), + }, + Message { + role: "user", + content: "hi\nthere".into(), + }, + ]; + let p = super::super::chatml(&msgs); + let parsed = parse_chatml(&p).unwrap(); + assert_eq!(parsed, msgs); + let l = llama3_chat(&parsed); + assert_eq!( + l, + "<|begin_of_text|><|start_header_id|>system<|end_header_id|>\n\nsys<|eot_id|><|start_header_id|>user<|end_header_id|>\n\nhi\nthere<|eot_id|><|start_header_id|>assistant<|end_header_id|>\n\n" + ); + assert!(parse_chatml("plain text").is_none()); + } + + #[test] + fn bad_files_are_load_errors_not_panics() { + let dir = tempfile::tempdir().unwrap(); + let g = dir.path().join("model.gguf"); + let t = dir.path().join("tokenizer.json"); + std::fs::write(&g, b"GGUF garbage").unwrap(); + std::fs::write(&t, b"{}").unwrap(); + assert!(matches!( + CandleBackend::load(&g, &t), + Err(AssistantError::Load(_)) + )); + let missing = dir.path().join("nope.gguf"); + assert!(matches!( + CandleBackend::load(&missing, &t), + Err(AssistantError::Load(_)) + )); + } + + #[test] + fn stop_strings() { + assert_eq!( + first_stop("abc<|im_end|>x", &["<|im_end|>".into()]), + Some(3) + ); + assert_eq!( + first_stop("abc", &["<|im_end|>".into(), String::new()]), + None + ); } } diff --git a/crates/assistant/src/convert.rs b/crates/assistant/src/convert.rs index 1e90de3..74a222e 100644 --- a/crates/assistant/src/convert.rs +++ b/crates/assistant/src/convert.rs @@ -4,7 +4,7 @@ //! requirements, naming conventions, ...) into [`Knowledge`] chunks the //! assistant can retrieve. Pure text extraction, no model involved. //! -//! STUB: to be implemented. Contract: +//! Contract: //! * `convert_paths(paths)` accepts files and directories (recursed) with //! extensions .md .markdown .txt .yaml .yml .json .docx .pdf; others are //! skipped and listed in `ConvertReport.skipped`. @@ -16,13 +16,22 @@ //! * Every chunk gets `source` = relative path, `keywords` = top 8 distinct //! capitalised or underscored tokens (identifiers like Pset_WallCommon). //! * Deterministic order (sorted paths). +//! +//! Notes: "capitalised or underscored tokens" means identifier-like tokens: +//! containing `_`, an inner capital (`IfcWall`, `BIM`) or a capital plus a +//! digit (`ID-001`); ordinary capitalised words ("The") are not keywords. +//! Hidden files and directories (`.git`, ...) are ignored when recursing. +//! A file that cannot be read or parsed (e.g. a broken PDF) is listed in +//! `skipped` with the reason; it never fails the whole conversion. +use std::collections::HashMap; +use std::io::Read; use std::path::{Path, PathBuf}; use serde::{Deserialize, Serialize}; use crate::error::AssistantResult; -use crate::knowledge::Knowledge; +use crate::knowledge::{chunk_markdown, split_long, Chunk, Knowledge}; #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, Default)] pub struct ConvertReport { @@ -35,17 +44,572 @@ pub const SUPPORTED_EXTENSIONS: &[&str] = &[ "md", "markdown", "txt", "yaml", "yml", "json", "docx", "pdf", ]; +/// Maximum characters per chunk (same as markdown chunking). +const MAX_CHUNK_CHARS: usize = 1800; +/// Number of keywords per chunk. +const KEYWORDS: usize = 8; + /// Convert files/directories into knowledge. pub fn convert_paths(paths: &[PathBuf]) -> AssistantResult { - let _ = paths; - Err(crate::error::AssistantError::Convert( - "not implemented yet".into(), + let mut report = ConvertReport::default(); + // (absolute-ish path used for ordering, source label) + let mut files: Vec<(PathBuf, String)> = Vec::new(); + for p in paths { + if p.is_dir() { + collect_dir(p, p, &mut files); + } else if p.is_file() { + let label = p + .file_name() + .map(|n| n.to_string_lossy().into_owned()) + .unwrap_or_else(|| p.display().to_string()); + files.push((p.clone(), label)); + } else { + report.skipped.push(format!("{}: not found", p.display())); + } + } + files.sort_by(|a, b| a.0.cmp(&b.0)); + files.dedup_by(|a, b| a.0 == b.0); + + for (path, source) in files { + let ext = path + .extension() + .map(|e| e.to_string_lossy().to_lowercase()) + .unwrap_or_default(); + if !SUPPORTED_EXTENSIONS.contains(&ext.as_str()) { + report.skipped.push(format!( + "{source}: unsupported file type{}", + if ext.is_empty() { + String::new() + } else { + format!(" .{ext}") + } + )); + continue; + } + match convert_file(&path, &source, &ext) { + Ok(mut chunks) => { + let stem = path + .file_stem() + .map(|s| s.to_string_lossy().into_owned()) + .unwrap_or_else(|| source.clone()); + for c in &mut chunks { + if c.title.trim().is_empty() { + c.title = stem.clone(); + } + if c.keywords.is_empty() { + c.keywords = keywords(&format!("{}\n{}", c.title, c.text), KEYWORDS); + } + } + report.knowledge.chunks.extend(chunks); + report.files_read.push(source); + } + Err(reason) => report.skipped.push(format!("{source}: {reason}")), + } + } + Ok(report) +} + +/// Recurse into `dir`, labelling files relative to `root` with `/` separators. +fn collect_dir(root: &Path, dir: &Path, out: &mut Vec<(PathBuf, String)>) { + let Ok(entries) = std::fs::read_dir(dir) else { + return; + }; + let mut entries: Vec = entries.filter_map(|e| e.ok().map(|e| e.path())).collect(); + entries.sort(); + for p in entries { + if p.file_name() + .map(|n| n.to_string_lossy().starts_with('.')) + .unwrap_or(false) + { + continue; + } + if p.is_dir() { + collect_dir(root, &p, out); + } else if p.is_file() { + let rel = p.strip_prefix(root).unwrap_or(&p); + let label = rel + .components() + .map(|c| c.as_os_str().to_string_lossy().into_owned()) + .collect::>() + .join("/"); + out.push((p, label)); + } + } +} + +fn convert_file(path: &Path, source: &str, ext: &str) -> Result, String> { + match ext { + "md" | "markdown" | "txt" => Ok(chunk_markdown(source, &read_text(path)?)), + "yaml" | "yml" => convert_yaml(source, &read_text(path)?), + "json" => convert_json(source, &read_text(path)?), + "docx" => Ok(chunk_markdown(source, &docx_to_markdown(path)?)), + "pdf" => Ok(chunk_markdown(source, &pdf_to_text(path)?)), + other => Err(format!("unsupported file type .{other}")), + } +} + +fn read_text(path: &Path) -> Result { + let bytes = std::fs::read(path).map_err(|e| format!("cannot read: {e}"))?; + Ok(String::from_utf8_lossy(&bytes).into_owned()) +} + +/// A Knowledge file (top-level `chunks:`) is merged; any other YAML is chunked by top-level key. +fn convert_yaml(source: &str, text: &str) -> Result, String> { + let value: serde_yaml::Value = + serde_yaml::from_str(text).map_err(|e| format!("invalid YAML: {e}"))?; + if let serde_yaml::Value::Mapping(m) = &value { + if m.contains_key("chunks") { + let k = + Knowledge::from_yaml(text).map_err(|e| format!("invalid knowledge file: {e}"))?; + return Ok(k + .chunks + .into_iter() + .map(|mut c| { + if c.source.trim().is_empty() { + c.source = source.to_string(); + } + c + }) + .collect()); + } + let mut out = Vec::new(); + for (k, v) in m { + let key = match k { + serde_yaml::Value::String(s) => s.clone(), + other => serde_yaml::to_string(other) + .unwrap_or_default() + .trim() + .to_string(), + }; + let mut one = serde_yaml::Mapping::new(); + one.insert(k.clone(), v.clone()); + let body = serde_yaml::to_string(&one).unwrap_or_default(); + out.extend(pieces(source, &key, &body)); + } + return Ok(out); + } + Ok(pieces(source, "", text)) +} + +/// JSON is pretty-printed; an object is chunked by top-level key. +fn convert_json(source: &str, text: &str) -> Result, String> { + let value: serde_json::Value = + serde_json::from_str(text).map_err(|e| format!("invalid JSON: {e}"))?; + if let serde_json::Value::Object(m) = &value { + let mut out = Vec::new(); + for (k, v) in m { + let mut one = serde_json::Map::new(); + one.insert(k.clone(), v.clone()); + let body = + serde_json::to_string_pretty(&serde_json::Value::Object(one)).unwrap_or_default(); + out.extend(pieces(source, k, &body)); + } + return Ok(out); + } + Ok(pieces( + source, + "", + &serde_json::to_string_pretty(&value).unwrap_or_default(), )) } +fn pieces(source: &str, title: &str, text: &str) -> Vec { + let t = text.trim(); + if t.is_empty() { + return vec![]; + } + split_long(t, MAX_CHUNK_CHARS) + .into_iter() + .map(|text| Chunk { + source: source.to_string(), + title: title.to_string(), + text, + keywords: vec![], + }) + .collect() +} + +/// Extract `word/document.xml` paragraphs as markdown (Title/Heading1..3 → #..####). +fn docx_to_markdown(path: &Path) -> Result { + let file = std::fs::File::open(path).map_err(|e| format!("cannot read: {e}"))?; + let mut zip = + zip::ZipArchive::new(file).map_err(|e| format!("not a valid .docx (zip): {e}"))?; + let mut xml = String::new(); + zip.by_name("word/document.xml") + .map_err(|e| format!("not a valid .docx (word/document.xml): {e}"))? + .read_to_string(&mut xml) + .map_err(|e| format!("cannot read word/document.xml: {e}"))?; + docx_xml_to_markdown(&xml) +} + +fn docx_xml_to_markdown(xml: &str) -> Result { + use quick_xml::events::Event; + let mut reader = quick_xml::Reader::from_str(xml); + let mut paragraphs: Vec = Vec::new(); + let mut para = String::new(); + let mut style: Option = None; + let mut in_para = false; + let mut in_text = false; + loop { + match reader + .read_event() + .map_err(|e| format!("invalid word/document.xml: {e}"))? + { + Event::Start(e) => match e.name().as_ref() { + b"w:p" => { + in_para = true; + para.clear(); + style = None; + } + b"w:t" => in_text = true, + _ => {} + }, + Event::Empty(e) => match e.name().as_ref() { + b"w:pStyle" if in_para => { + style = e + .attributes() + .flatten() + .find(|a| a.key.as_ref() == b"w:val") + .and_then(|a| a.unescape_value().ok()) + .map(|v| v.into_owned()); + } + b"w:tab" if in_para => para.push('\t'), + b"w:br" | b"w:cr" if in_para => para.push('\n'), + _ => {} + }, + Event::Text(t) if in_text => { + para.push_str( + &t.unescape() + .map_err(|e| format!("invalid text in word/document.xml: {e}"))?, + ); + } + Event::End(e) => match e.name().as_ref() { + b"w:t" => in_text = false, + b"w:p" => { + in_para = false; + let text = para.trim(); + if !text.is_empty() { + let prefix = match style + .as_deref() + .map(|s| s.to_lowercase().replace([' ', '-', '_'], "")) + .as_deref() + { + Some("title") => "# ", + Some("heading1") => "## ", + Some("heading2") => "### ", + Some("heading3") => "#### ", + _ => "", + }; + let text = if prefix.is_empty() { + text.to_string() + } else { + text.replace(['\n', '\t'], " ") + }; + paragraphs.push(format!("{prefix}{text}")); + } + } + _ => {} + }, + Event::Eof => break, + _ => {} + } + } + Ok(paragraphs.join("\n\n")) +} + +/// Extract PDF text; paragraphs are separated by blank lines. +fn pdf_to_text(path: &Path) -> Result { + let p = path.to_path_buf(); + let text = std::panic::catch_unwind(move || pdf_extract::extract_text(&p)) + .map_err(|_| "the PDF could not be parsed".to_string())? + .map_err(|e| format!("no text could be extracted from the PDF: {e}"))?; + let paragraphs: Vec = text + .replace('\r', "") + .split("\n\n") + .map(|p| { + p.lines() + .map(str::trim) + .filter(|l| !l.is_empty()) + .collect::>() + .join(" ") + }) + .filter(|p| !p.is_empty()) + .collect(); + if paragraphs.is_empty() { + return Err("the PDF contains no extractable text (scanned image?)".into()); + } + Ok(paragraphs.join("\n\n")) +} + +/// Identifier-like token: contains `_`, an inner capital, or a capital and a digit. +fn is_identifier(t: &str) -> bool { + if t.chars().count() < 2 || !t.chars().any(|c| c.is_alphabetic()) { + return false; + } + let has_inner_upper = t.chars().skip(1).any(|c| c.is_uppercase()); + let starts_upper = t.chars().next().map(|c| c.is_uppercase()).unwrap_or(false); + t.contains('_') || has_inner_upper || (starts_upper && t.chars().any(|c| c.is_ascii_digit())) +} + +/// Top `n` distinct identifier-like tokens by frequency (ties: first occurrence). +pub fn keywords(text: &str, n: usize) -> Vec { + let mut counts: HashMap = HashMap::new(); + for (i, tok) in text + .split(|c: char| !(c.is_alphanumeric() || c == '_' || c == '-')) + .map(|t| t.trim_matches(|c| c == '-' || c == '_')) + .filter(|t| is_identifier(t)) + .enumerate() + { + let e = counts.entry(tok.to_string()).or_insert((0, i)); + e.0 += 1; + } + let mut v: Vec<(String, (usize, usize))> = counts.into_iter().collect(); + v.sort_by(|a, b| b.1 .0.cmp(&a.1 .0).then(a.1 .1.cmp(&b.1 .1))); + v.into_iter().take(n).map(|(t, _)| t).collect() +} + /// Convert and write `knowledge.yaml` next to the outputs. pub fn convert_to_file(paths: &[PathBuf], out: &Path) -> AssistantResult { let r = convert_paths(paths)?; std::fs::write(out, r.knowledge.to_yaml())?; Ok(r) } + +#[cfg(test)] +mod tests { + use super::*; + use std::io::Write; + + fn write_docx(path: &Path) { + let xml = r#" + +Property requirements +All walls shall carry Pset_WallCommon.FireRating & IsExternal. +Doors use Pset_DoorCommon; Pset_WallCommon applies to IfcWall only. +"#; + let f = std::fs::File::create(path).unwrap(); + let mut z = zip::ZipWriter::new(f); + let opts = zip::write::SimpleFileOptions::default(); + z.start_file("[Content_Types].xml", opts).unwrap(); + z.write_all(b"").unwrap(); + z.start_file("word/document.xml", opts).unwrap(); + z.write_all(xml.as_bytes()).unwrap(); + z.finish().unwrap(); + } + + /// A minimal uncompressed one-page PDF with two text paragraphs. + fn write_pdf(path: &Path) { + let content = "BT /F1 12 Tf 72 720 Td (BIM Execution Plan) Tj ET\nBT /F1 12 Tf 72 600 Td (Walls require Pset_WallCommon) Tj ET\n"; + let objects = [ + "<< /Type /Catalog /Pages 2 0 R >>".to_string(), + "<< /Type /Pages /Kids [3 0 R] /Count 1 >>".to_string(), + "<< /Type /Page /Parent 2 0 R /MediaBox [0 0 612 792] /Resources << /Font << /F1 5 0 R >> >> /Contents 4 0 R >>".to_string(), + format!("<< /Length {} >>\nstream\n{}endstream", content.len(), content), + "<< /Type /Font /Subtype /Type1 /BaseFont /Helvetica >>".to_string(), + ]; + let mut pdf = String::from("%PDF-1.4\n"); + let mut offsets = Vec::new(); + for (i, o) in objects.iter().enumerate() { + offsets.push(pdf.len()); + pdf.push_str(&format!("{} 0 obj\n{}\nendobj\n", i + 1, o)); + } + let xref = pdf.len(); + pdf.push_str(&format!( + "xref\n0 {}\n0000000000 65535 f \n", + objects.len() + 1 + )); + for o in offsets { + pdf.push_str(&format!("{o:010} 00000 n \n")); + } + pdf.push_str(&format!( + "trailer\n<< /Size {} /Root 1 0 R >>\nstartxref\n{}\n%%EOF\n", + objects.len() + 1, + xref + )); + std::fs::write(path, pdf).unwrap(); + } + + fn fixture() -> tempfile::TempDir { + let dir = tempfile::tempdir().unwrap(); + let d = dir.path(); + std::fs::write( + d.join("a_guide.md"), + "# Guide\nintro text\n## Naming\nWalls are named W-001 using IfcWall names.\n", + ) + .unwrap(); + std::fs::write( + d.join("b_notes.txt"), + "Plain notes about the BIM model.\n\nSecond paragraph mentions LOD_300.", + ) + .unwrap(); + std::fs::write( + d.join("c_data.json"), + r#"{"project": {"name": "Demo"}, "codes": ["Ss_25_10", "Ss_25_20"]}"#, + ) + .unwrap(); + std::fs::write( + d.join("d_knowledge.yaml"), + "chunks:\n - source: ''\n title: Existing chunk\n text: Some imported text about Pset_SlabCommon.\n keywords: [slab]\n", + ) + .unwrap(); + std::fs::write( + d.join("e_config.yml"), + "naming:\n patterns:\n IfcWall: '^W-'\nqto:\n currency: EUR\n", + ) + .unwrap(); + write_docx(&d.join("f_eir.docx")); + std::fs::write(d.join("g_image.xyz"), "binary").unwrap(); + std::fs::create_dir(d.join("sub")).unwrap(); + std::fs::write( + d.join("sub").join("h_more.md"), + "## Deep\nNested IfcSlab note.\n", + ) + .unwrap(); + std::fs::create_dir(d.join(".git")).unwrap(); + std::fs::write(d.join(".git").join("x.md"), "ignored").unwrap(); + dir + } + + #[test] + fn converts_mixed_directory() { + let dir = fixture(); + let r = convert_paths(&[dir.path().to_path_buf()]).unwrap(); + assert_eq!( + r.files_read, + vec![ + "a_guide.md", + "b_notes.txt", + "c_data.json", + "d_knowledge.yaml", + "e_config.yml", + "f_eir.docx", + "sub/h_more.md" + ] + ); + assert_eq!(r.skipped, vec!["g_image.xyz: unsupported file type .xyz"]); + + let by_source = |s: &str| { + r.knowledge + .chunks + .iter() + .filter(|c| c.source == s) + .collect::>() + }; + let md = by_source("a_guide.md"); + assert_eq!(md.len(), 2); + assert_eq!(md[1].title, "Guide › Naming"); + assert!( + md[1].keywords.contains(&"W-001".to_string()), + "{:?}", + md[1].keywords + ); + assert!(md[1].keywords.contains(&"IfcWall".to_string())); + assert!(!md[1].keywords.contains(&"Walls".to_string())); + + let txt = by_source("b_notes.txt"); + assert_eq!(txt.len(), 1); + assert_eq!(txt[0].title, "b_notes"); + assert_eq!(txt[0].keywords, vec!["b_notes", "BIM", "LOD_300"]); + + let json = by_source("c_data.json"); + assert_eq!( + json.iter().map(|c| c.title.as_str()).collect::>(), + vec!["project", "codes"] + ); + assert!(json[1].keywords.contains(&"Ss_25_10".to_string())); + + let k = by_source("d_knowledge.yaml"); + assert_eq!(k.len(), 1); + assert_eq!(k[0].title, "Existing chunk"); + assert_eq!(k[0].keywords, vec!["slab"]); + + let y = by_source("e_config.yml"); + assert_eq!( + y.iter().map(|c| c.title.as_str()).collect::>(), + vec!["naming", "qto"] + ); + assert!(y[0].text.contains("IfcWall")); + + let docx = by_source("f_eir.docx"); + assert_eq!(docx.len(), 1); + assert_eq!(docx[0].title, "Property requirements"); + assert!( + docx[0] + .text + .contains("All walls shall carry Pset_WallCommon.FireRating & IsExternal."), + "{}", + docx[0].text + ); + assert!(docx[0].text.contains("\n\nDoors use Pset_DoorCommon")); + assert_eq!(docx[0].keywords[0], "Pset_WallCommon"); + + assert_eq!(by_source("sub/h_more.md")[0].title, "Deep"); + assert_eq!(r.knowledge.chunks.len(), 2 + 1 + 2 + 1 + 2 + 1 + 1); + + // Deterministic: same output twice; a file reached twice is read once. + assert_eq!(convert_paths(&[dir.path().to_path_buf()]).unwrap(), r); + let again = convert_paths(&[dir.path().join("sub"), dir.path().to_path_buf()]).unwrap(); + assert_eq!(again.knowledge.chunks.len(), r.knowledge.chunks.len()); + } + + #[test] + fn missing_and_broken_files_are_skipped() { + let dir = tempfile::tempdir().unwrap(); + let bad_json = dir.path().join("bad.json"); + let bad_docx = dir.path().join("bad.docx"); + let bad_pdf = dir.path().join("bad.pdf"); + std::fs::write(&bad_json, "{ nope").unwrap(); + std::fs::write(&bad_docx, "not a zip").unwrap(); + std::fs::write(&bad_pdf, "%PDF-1.4 garbage").unwrap(); + let r = + convert_paths(&[bad_json, bad_docx, bad_pdf, dir.path().join("missing.md")]).unwrap(); + assert!(r.files_read.is_empty()); + assert_eq!(r.skipped.len(), 4, "{:?}", r.skipped); + assert!(r.skipped[0].ends_with("missing.md: not found")); + assert!(r.skipped[1].starts_with("bad.docx: not a valid .docx")); + assert!(r.skipped[2].starts_with("bad.json: invalid JSON")); + assert!(r.skipped[3].starts_with("bad.pdf: ")); + } + + #[test] + fn converts_minimal_pdf() { + let dir = tempfile::tempdir().unwrap(); + let p = dir.path().join("bep.pdf"); + write_pdf(&p); + let r = convert_paths(std::slice::from_ref(&p)).unwrap(); + assert_eq!(r.files_read, vec!["bep.pdf"], "{:?}", r.skipped); + let text: String = r + .knowledge + .chunks + .iter() + .map(|c| c.text.clone()) + .collect::>() + .join("\n"); + assert!(text.contains("BIM Execution Plan"), "{text}"); + assert!(text.contains("Pset_WallCommon"), "{text}"); + assert!(r.knowledge.chunks[0] + .keywords + .contains(&"Pset_WallCommon".to_string())); + } + + #[test] + fn convert_to_file_writes_loadable_knowledge() { + let dir = fixture(); + let out = dir.path().join("out.yaml"); + let r = convert_to_file(&[dir.path().join("a_guide.md")], &out).unwrap(); + let k = Knowledge::from_yaml(&std::fs::read_to_string(&out).unwrap()).unwrap(); + assert_eq!(k, r.knowledge); + assert!(!k.search("naming walls", 1).is_empty()); + } + + #[test] + fn keyword_rules() { + assert_eq!( + keywords( + "The Wall uses IfcWall and IfcWall, Pset_WallCommon, ID-001, plain words, X", + 8 + ), + vec!["IfcWall", "Pset_WallCommon", "ID-001"] + ); + } +} diff --git a/crates/assistant/src/knowledge.rs b/crates/assistant/src/knowledge.rs index 96659a6..8171496 100644 --- a/crates/assistant/src/knowledge.rs +++ b/crates/assistant/src/knowledge.rs @@ -87,7 +87,7 @@ pub fn chunk_markdown(source: &str, text: &str) -> Vec { }; for line in text.lines() { let hashes = line.chars().take_while(|c| *c == '#').count(); - if hashes >= 1 && hashes <= 4 && line.chars().nth(hashes) == Some(' ') { + if (1..=4).contains(&hashes) && line.chars().nth(hashes) == Some(' ') { flush(&mut buf, &path, &mut out); path.retain(|(lvl, _)| *lvl < hashes); path.push((hashes, line[hashes..].trim().to_string())); diff --git a/crates/assistant/src/model.rs b/crates/assistant/src/model.rs index 5f2ee2a..d03f1b9 100644 --- a/crates/assistant/src/model.rs +++ b/crates/assistant/src/model.rs @@ -4,11 +4,13 @@ //! Assistant panel (an explicit action, shown with size and license); it is //! stored under `/assistant//` and verified by SHA-256. +use std::io::{Read, Write}; use std::path::{Path, PathBuf}; use serde::{Deserialize, Serialize}; +use sha2::{Digest, Sha256}; -use crate::error::AssistantResult; +use crate::error::{AssistantError, AssistantResult}; /// A downloadable model. #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] @@ -34,12 +36,12 @@ pub static DEFAULT_MODEL: std::sync::LazyLock = std::sync::LazyLock:: name: "Qwen2.5 0.5B Instruct (Q4_K_M)".into(), family: "qwen2".into(), gguf_url: "https://huggingface.co/Qwen/Qwen2.5-0.5B-Instruct-GGUF/resolve/main/qwen2.5-0.5b-instruct-q4_k_m.gguf".into(), - gguf_sha256: String::new(), // filled in by the implementer after verifying the download + gguf_sha256: "74a4da8c9fdbcd15bd1f6d01d621410d31c6fc00986f5eb687824e7b93d7a9db".into(), gguf_size_bytes: 491_400_032, tokenizer_url: "https://huggingface.co/Qwen/Qwen2.5-0.5B-Instruct/resolve/main/tokenizer.json".into(), - tokenizer_sha256: String::new(), + tokenizer_sha256: "c0382117ea329cdf097041132f6d735924b697924d6f6fc3945713e96ce87539".into(), license: "Apache-2.0".into(), - approx_ram_mb: 700, + approx_ram_mb: 1400, } }); @@ -99,18 +101,197 @@ pub fn status_of(app_data: &Path, spec: &ModelSpec) -> ModelStatus { /// Download progress callback: (bytes_done, bytes_total_if_known). pub type Progress<'a> = &'a mut dyn FnMut(u64, Option); +/// Progress is reported roughly every this many bytes. +const PROGRESS_STEP: u64 = 1 << 20; + /// Download `spec` into `app_data`, verify SHA-256 (when the spec declares one), -/// and return the status. STUB: to be implemented with `ureq` streaming to a -/// `.part` file renamed on success; never leaves a half file as installed. +/// and return the status. +/// +/// Each file is streamed to `.part` and renamed only after the checksum +/// matches, so a failed or interrupted download never looks installed. The +/// tokenizer (small) is fetched first, then the GGUF weights. Progress is +/// cumulative over both files. A file already in place with the right +/// checksum is not downloaded again. pub fn download( app_data: &Path, spec: &ModelSpec, progress: Progress<'_>, ) -> AssistantResult { - let _ = (app_data, spec, progress); - Err(crate::error::AssistantError::Download( - "not implemented yet".into(), - )) + let dir = model_dir(app_data, spec); + std::fs::create_dir_all(&dir) + .map_err(|e| AssistantError::Download(format!("cannot create {}: {e}", dir.display())))?; + let files = [ + ( + &spec.tokenizer_url, + &spec.tokenizer_sha256, + "tokenizer.json", + ), + (&spec.gguf_url, &spec.gguf_sha256, "model.gguf"), + ]; + let mut done: u64 = 0; + let mut total: Option = if spec.gguf_size_bytes > 0 { + Some(spec.gguf_size_bytes) + } else { + None + }; + for (i, (url, sha, name)) in files.iter().enumerate() { + let target = dir.join(name); + if target.is_file() + && !sha.is_empty() + && sha256_file(&target) + .map(|h| h.eq_ignore_ascii_case(sha)) + .unwrap_or(false) + { + let len = std::fs::metadata(&target).map(|m| m.len()).unwrap_or(0); + if i == 0 { + total = total.map(|t| t + len); + } + done += len; + progress(done, total); + continue; + } + let resp = agent_for(url) + .get(url) + .call() + .map_err(|e| AssistantError::Download(format!("{url}: {e}")))?; + let len: Option = resp.header("Content-Length").and_then(|v| v.parse().ok()); + if i == 0 { + // tokenizer first: total = tokenizer + declared gguf size + total = match (total, len) { + (Some(t), Some(l)) => Some(t + l), + _ => None, + }; + } else if let (Some(l), None) = (len, total) { + total = Some(done + l); + } + let part = dir.join(format!("{name}.part")); + let result = stream_to_file(resp.into_reader(), &part, &mut done, total, progress) + .and_then(|hash| { + if !sha.is_empty() && !hash.eq_ignore_ascii_case(sha) { + return Err(AssistantError::Checksum(format!( + "{name}: expected {sha}, got {hash}" + ))); + } + std::fs::rename(&part, &target).map_err(|e| { + AssistantError::Download(format!("cannot move {name} into place: {e}")) + }) + }); + if let Err(e) = result { + let _ = std::fs::remove_file(&part); + return Err(e); + } + } + progress(done, Some(done)); + Ok(status_of(app_data, spec)) +} + +/// HTTP agent; honours HTTPS_PROXY / HTTP_PROXY / ALL_PROXY except for loopback URLs. +fn agent_for(url: &str) -> ureq::Agent { + let loopback = ["://127.0.0.1", "://localhost", "://[::1]"] + .iter() + .any(|h| url.contains(h)); + ureq::AgentBuilder::new() + .timeout_connect(std::time::Duration::from_secs(30)) + .timeout_read(std::time::Duration::from_secs(120)) + .try_proxy_from_env(!loopback) + .build() +} + +/// Copy `reader` into `path`, reporting progress, and return the SHA-256 (hex). +fn stream_to_file( + mut reader: impl Read, + path: &Path, + done: &mut u64, + total: Option, + progress: Progress<'_>, +) -> AssistantResult { + let mut out = std::fs::File::create(path) + .map_err(|e| AssistantError::Download(format!("cannot create {}: {e}", path.display())))?; + let mut hasher = Sha256::new(); + let mut buf = vec![0u8; 64 * 1024]; + let mut last_report = *done; + loop { + let n = match reader.read(&mut buf) { + Ok(0) => break, + Ok(n) => n, + Err(e) if e.kind() == std::io::ErrorKind::Interrupted => continue, + Err(e) => { + return Err(AssistantError::Download(format!( + "connection interrupted: {e}" + ))) + } + }; + out.write_all(&buf[..n]).map_err(|e| { + AssistantError::Download(format!("cannot write {}: {e}", path.display())) + })?; + hasher.update(&buf[..n]); + *done += n as u64; + if *done - last_report >= PROGRESS_STEP { + last_report = *done; + progress(*done, total); + } + } + out.flush() + .and_then(|_| out.sync_all()) + .map_err(|e| AssistantError::Download(format!("cannot write {}: {e}", path.display())))?; + Ok(hex::encode(hasher.finalize())) +} + +/// SHA-256 of a file (hex, lower case). +pub fn sha256_file(path: &Path) -> AssistantResult { + let mut f = std::fs::File::open(path)?; + let mut hasher = Sha256::new(); + let mut buf = vec![0u8; 256 * 1024]; + loop { + let n = f.read(&mut buf)?; + if n == 0 { + break; + } + hasher.update(&buf[..n]); + } + Ok(hex::encode(hasher.finalize())) +} + +/// Install a model from local files (for users who downloaded the files +/// manually, and for tests). Each file is copied to `.part`, verified +/// against the spec's SHA-256 (when declared) and renamed into place. +pub fn install_from_files( + app_data: &Path, + spec: &ModelSpec, + gguf_path: &Path, + tokenizer_path: &Path, +) -> AssistantResult { + let dir = model_dir(app_data, spec); + std::fs::create_dir_all(&dir)?; + for (src, sha, name) in [ + (tokenizer_path, &spec.tokenizer_sha256, "tokenizer.json"), + (gguf_path, &spec.gguf_sha256, "model.gguf"), + ] { + if !src.is_file() { + return Err(AssistantError::Other(format!( + "{} does not exist", + src.display() + ))); + } + let part = dir.join(format!("{name}.part")); + let result = std::fs::copy(src, &part) + .map_err(|e| AssistantError::Other(format!("cannot copy {}: {e}", src.display()))) + .and_then(|_| sha256_file(&part)) + .and_then(|hash| { + if !sha.is_empty() && !hash.eq_ignore_ascii_case(sha) { + return Err(AssistantError::Checksum(format!( + "{}: expected {sha}, got {hash}", + src.display() + ))); + } + std::fs::rename(&part, dir.join(name)).map_err(AssistantError::from) + }); + if let Err(e) = result { + let _ = std::fs::remove_file(&part); + return Err(e); + } + } + Ok(status_of(app_data, spec)) } /// Remove an installed model. @@ -121,3 +302,164 @@ pub fn remove(app_data: &Path, spec: &ModelSpec) -> AssistantResult<()> { } Ok(()) } + +#[cfg(test)] +mod tests { + use super::*; + + fn test_spec(gguf_sha: &str, tok_sha: &str) -> ModelSpec { + ModelSpec { + id: "test-model".into(), + gguf_sha256: gguf_sha.into(), + tokenizer_sha256: tok_sha.into(), + gguf_size_bytes: 0, + ..DEFAULT_MODEL.clone() + } + } + + fn sha(bytes: &[u8]) -> String { + hex::encode(Sha256::digest(bytes)) + } + + #[test] + fn default_model_declares_checksums() { + assert_eq!(DEFAULT_MODEL.gguf_sha256.len(), 64); + assert_eq!(DEFAULT_MODEL.tokenizer_sha256.len(), 64); + } + + #[test] + fn status_and_install_from_files() { + let app = tempfile::tempdir().unwrap(); + let src = tempfile::tempdir().unwrap(); + let g = src.path().join("x.gguf"); + let t = src.path().join("tok.json"); + std::fs::write(&g, b"weights").unwrap(); + std::fs::write(&t, b"{\"tok\":1}").unwrap(); + let spec = test_spec(&sha(b"weights"), &sha(b"{\"tok\":1}")); + + let before = status_of(app.path(), &spec); + assert!(!before.installed); + assert_eq!( + before.model_dir, + app.path().join("assistant").join("test-model") + ); + + let after = install_from_files(app.path(), &spec, &g, &t).unwrap(); + assert!(after.installed); + assert_eq!(std::fs::read(&after.gguf_path).unwrap(), b"weights"); + assert!(!after.model_dir.join("model.gguf.part").exists()); + assert_eq!(status_of(app.path(), &spec), after); + + remove(app.path(), &spec).unwrap(); + assert!(!status_of(app.path(), &spec).installed); + } + + #[test] + fn install_rejects_checksum_mismatch_and_leaves_nothing() { + let app = tempfile::tempdir().unwrap(); + let src = tempfile::tempdir().unwrap(); + let g = src.path().join("x.gguf"); + let t = src.path().join("tok.json"); + std::fs::write(&g, b"weights").unwrap(); + std::fs::write(&t, b"tok").unwrap(); + let spec = test_spec(&"0".repeat(64), ""); + let err = install_from_files(app.path(), &spec, &g, &t).unwrap_err(); + assert!(matches!(err, AssistantError::Checksum(_)), "{err:?}"); + let st = status_of(app.path(), &spec); + assert!(!st.installed); + assert!(!st.gguf_path.exists()); + assert!(!st.model_dir.join("model.gguf.part").exists()); + } + + /// Serve fixed bodies over plain HTTP on localhost (one request per connection). + fn serve(bodies: Vec<(&'static str, Vec)>) -> String { + use std::io::BufRead; + let listener = std::net::TcpListener::bind("127.0.0.1:0").unwrap(); + let addr = listener.local_addr().unwrap(); + std::thread::spawn(move || { + for stream in listener.incoming() { + let Ok(mut stream) = stream else { break }; + let mut reader = std::io::BufReader::new(stream.try_clone().unwrap()); + let mut line = String::new(); + reader.read_line(&mut line).unwrap(); + let path = line.split_whitespace().nth(1).unwrap_or("/").to_string(); + loop { + let mut h = String::new(); + if reader.read_line(&mut h).unwrap_or(0) == 0 || h == "\r\n" { + break; + } + } + match bodies.iter().find(|(p, _)| *p == path) { + Some((_, body)) => { + let _ = write!( + stream, + "HTTP/1.1 200 OK\r\nContent-Length: {}\r\nConnection: close\r\n\r\n", + body.len() + ); + let _ = stream.write_all(body); + } + None => { + let _ = write!(stream, "HTTP/1.1 404 Not Found\r\nContent-Length: 0\r\nConnection: close\r\n\r\n"); + } + } + } + }); + format!("http://{addr}") + } + + #[test] + fn download_streams_verifies_and_installs() { + let gguf: Vec = (0..3_000_000u32).map(|i| (i % 251) as u8).collect(); + let tok = b"{\"model\":{}}".to_vec(); + let base = serve(vec![("/m.gguf", gguf.clone()), ("/t.json", tok.clone())]); + let spec = ModelSpec { + gguf_url: format!("{base}/m.gguf"), + tokenizer_url: format!("{base}/t.json"), + gguf_size_bytes: gguf.len() as u64, + ..test_spec(&sha(&gguf), &sha(&tok)) + }; + let app = tempfile::tempdir().unwrap(); + let mut calls = Vec::new(); + let st = download(app.path(), &spec, &mut |d, t| calls.push((d, t))).unwrap(); + assert!(st.installed); + assert_eq!(std::fs::read(&st.gguf_path).unwrap(), gguf); + assert!(calls.len() >= 3, "{calls:?}"); + let full = (gguf.len() + tok.len()) as u64; + assert_eq!(*calls.last().unwrap(), (full, Some(full))); + assert!(calls.windows(2).all(|w| w[0].0 <= w[1].0)); + } + + #[test] + fn download_checksum_mismatch_is_not_installed() { + let base = serve(vec![ + ("/m.gguf", b"bad weights".to_vec()), + ("/t.json", b"tok".to_vec()), + ]); + let spec = ModelSpec { + gguf_url: format!("{base}/m.gguf"), + tokenizer_url: format!("{base}/t.json"), + ..test_spec(&"a".repeat(64), "") + }; + let app = tempfile::tempdir().unwrap(); + let err = download(app.path(), &spec, &mut |_, _| {}).unwrap_err(); + assert!(matches!(err, AssistantError::Checksum(_)), "{err:?}"); + let st = status_of(app.path(), &spec); + assert!(!st.installed); + assert!(!st.model_dir.join("model.gguf.part").exists()); + assert!(!st.gguf_path.exists()); + } + + #[test] + fn download_http_error_is_download_error() { + let base = serve(vec![]); + let spec = ModelSpec { + gguf_url: format!("{base}/missing.gguf"), + tokenizer_url: format!("{base}/missing.json"), + ..test_spec("", "") + }; + let app = tempfile::tempdir().unwrap(); + let err = download(app.path(), &spec, &mut |_, _| {}).unwrap_err(); + assert!(matches!(err, AssistantError::Download(_)), "{err:?}"); + assert!(!status_of(app.path(), &spec).installed); + } +} diff --git a/crates/assistant/src/tasks.rs b/crates/assistant/src/tasks.rs index 13315f3..a128f28 100644 --- a/crates/assistant/src/tasks.rs +++ b/crates/assistant/src/tasks.rs @@ -1,7 +1,7 @@ //! Tasks: what the assistant can be asked, how it is prompted, and how its //! answer is validated before anyone sees it. //! -//! STUB: `run` to be implemented. Contract: +//! Contract: //! * Build messages: a system prompt stating the role and the JSON schema of the //! expected answer, then up to `k` retrieved chunks (BM25 over the knowledge), //! then the user content. Never include anything the user did not provide or @@ -22,14 +22,36 @@ //! the same way. //! * `Task::Ask`: free question → `Answer::Text` with the titles of the chunks //! used as `sources`. +//! +//! Prompt templates live in `crates/assistant/prompts/*.txt` (data, not code). +//! A template starts with a `[system]` section (task instructions + schema) +//! and may continue with `[user]` / `[assistant]` example turns (few-shot). +//! For `ConfigHelp` / `DraftRules`, YAML that still does not validate after +//! the retry is returned as `ConfigSnippet { valid: false, .. }` (flagged, +//! never applied); only output that is not JSON at all becomes plain text. +//! +//! Implementation notes (all deterministic): +//! * The prompt ends with the opening of the answer object (`{"answer":`), +//! which keeps a 0.5B model in JSON; the prefix is re-attached before parsing. +//! * Invalid YAML (and a non-matching `ConfigHelp`/`DraftRules` answer) also +//! triggers the single retry, with the validation error in the message. +//! * Literal newlines inside JSON strings are repaired; for every task except +//! `Filter`, trailing junk after complete leading fields is ignored +//! (`{"answer": "...", "id": 001}`). Filters are never salvaged. +//! * `Filter` retrieves `RETRIEVE_K / 2` chunks (the allowed values are its +//! real context); a missing `Ifc` prefix and plurals ("errors") are +//! canonicalised to an existing value; `dropped` entries read `field: value`. +//! * `ConfigHelp`/`DraftRules` fragments (e.g. a bare `required_psets:`) are +//! wrapped into their `project.yaml` section before validation and returned +//! wrapped, so the snippet can be pasted as is. use serde::{Deserialize, Serialize}; use verimodel_core::model::{Finding, RuleResult}; -use crate::backend::Backend; +use crate::backend::{chatml, Backend, GenOptions, Message}; use crate::error::AssistantResult; -use crate::knowledge::Knowledge; +use crate::knowledge::{Chunk, Knowledge}; /// Allowed values for filter answers: taken from the current report. #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, Default)] @@ -69,7 +91,7 @@ pub enum Task { context: ReportContext, }, Explain { - finding: Finding, + finding: Box, rule: Option, }, ConfigHelp { @@ -109,14 +131,1224 @@ pub enum Answer { /// Number of retrieved chunks put in the prompt. pub const RETRIEVE_K: usize = 4; +/// Maximum characters of one retrieved chunk put in the prompt (keeps prefill fast). +const CHUNK_PROMPT_CHARS: usize = 1200; +/// Maximum characters of user-provided long text (requirements, current YAML). +const LONG_TEXT_CHARS: usize = 4000; + +const PROMPT_SYSTEM: &str = include_str!("../prompts/system.txt"); +const PROMPT_FILTER: &str = include_str!("../prompts/filter.txt"); +const PROMPT_EXPLAIN: &str = include_str!("../prompts/explain.txt"); +const PROMPT_CONFIG_HELP: &str = include_str!("../prompts/config_help.txt"); +const PROMPT_DRAFT_RULES: &str = include_str!("../prompts/draft_rules.txt"); +const PROMPT_ASK: &str = include_str!("../prompts/ask.txt"); + +const RETRY_MESSAGE: &str = "Format error in your reply ({error}). Write your answer again as ONLY one JSON object with this schema: {schema}"; + /// Run a task against a backend with the given knowledge. pub fn run( backend: &mut dyn Backend, knowledge: &Knowledge, task: &Task, ) -> AssistantResult { - let _ = (backend, knowledge, task); - Err(crate::error::AssistantError::Generation( - "tasks not implemented yet".into(), - )) + // A filter is grounded in the allowed values; extra documentation only adds noise and latency. + let k = if matches!(task, Task::Filter { .. }) { + RETRIEVE_K / 2 + } else { + RETRIEVE_K + }; + let hits = knowledge.search(&retrieval_query(task), k); + let chunks: Vec<&Chunk> = hits.iter().map(|(c, _)| *c).collect(); + let sources = source_titles(&chunks); + let mut messages = build_messages(task, &chunks); + let opts = GenOptions::default(); + + // The answer is primed with the opening of the JSON object ({"answer":), + // which keeps small models in the required format. + let prefix = json_prefix(task); + let mut raw = with_prefix( + &prefix, + &backend.generate(&format!("{}{prefix}", chatml(&messages)), &opts)?, + ); + let mut parsed = parse_answer(task, &raw); + if let Err(error) = &parsed { + log::debug!("assistant answer rejected ({error}), retrying once"); + messages.push(Message { + role: "assistant", + content: raw.trim().to_string(), + }); + let retry = RETRY_MESSAGE + .replace("{error}", error) + .replace("{schema}", &schema_line(task)); + messages.push(Message { + role: "user", + content: retry, + }); + raw = with_prefix( + &prefix, + &backend.generate(&format!("{}{prefix}", chatml(&messages)), &opts)?, + ); + parsed = parse_answer(task, &raw); + } + Ok(match parsed { + Ok(mut answer) => { + if let Answer::Text { sources: s, .. } = &mut answer { + *s = sources; + } + answer + } + // The YAML still does not validate after the retry: show it, flagged. + Err(_) if matches!(task, Task::ConfigHelp { .. } | Task::DraftRules { .. }) => { + match parse_snippet(&raw) { + Some(answer) => answer, + None => Answer::Text { + text: plain_text(&prefix, &raw), + sources, + structured: false, + }, + } + } + Err(_) => Answer::Text { + text: plain_text(&prefix, &raw), + sources, + structured: false, + }, + }) +} + +// ---------------------------------------------------------------- prompts + +fn retrieval_query(task: &Task) -> String { + match task { + Task::Filter { question, .. } => question.clone(), + Task::Explain { finding, rule } => { + let mut q = format!( + "{} {} {}", + finding.rule_id, + finding.message, + finding.fix_hint.clone().unwrap_or_default() + ); + if let Some(r) = rule { + q.push(' '); + q.push_str(&r.description); + } + q + } + Task::ConfigHelp { question, .. } => format!("project.yaml {question}"), + Task::DraftRules { requirements_text } => format!( + "required_psets naming patterns classification {}", + truncate_chars(requirements_text, 600) + ), + Task::Ask { question } => question.clone(), + } +} + +fn source_titles(chunks: &[&Chunk]) -> Vec { + let mut out: Vec = Vec::new(); + for c in chunks { + let t = if c.title.trim().is_empty() { + c.source.clone() + } else { + c.title.clone() + }; + if !out.contains(&t) { + out.push(t); + } + } + out +} + +/// The one-line JSON schema of a task (the line after "JSON schema:" in its template). +fn schema_line(task: &Task) -> String { + let (system, _) = parse_template(template(task)); + let mut lines = system.lines(); + lines.by_ref().find(|l| l.trim() == "JSON schema:"); + lines.next().unwrap_or_default().trim().to_string() +} + +/// `{"":`, appended after the assistant header of the prompt. +fn json_prefix(task: &Task) -> String { + let schema = schema_line(task); + match schema.split('"').nth(1) { + Some(key) if !key.is_empty() => format!("{{\"{key}\":"), + _ => "{".into(), + } +} + +/// Re-attach the primed prefix unless the model restarted the object itself. +fn with_prefix(prefix: &str, output: &str) -> String { + let t = output.trim_start(); + let unfenced = t + .strip_prefix("```json") + .or_else(|| t.strip_prefix("```")) + .unwrap_or(t); + if unfenced.trim_start().starts_with('{') { + output.to_string() + } else { + format!("{prefix}{output}") + } +} + +/// Unparseable output shown as plain text, without the primed JSON prefix. +fn plain_text(prefix: &str, raw: &str) -> String { + match raw.strip_prefix(prefix) { + Some(rest) => rest + .trim() + .trim_start_matches('"') + .trim_end_matches('}') + .trim_end() + .trim_end_matches('"') + .trim() + .to_string(), + None => raw.trim().to_string(), + } +} + +fn template(task: &Task) -> &'static str { + match task { + Task::Filter { .. } => PROMPT_FILTER, + Task::Explain { .. } => PROMPT_EXPLAIN, + Task::ConfigHelp { .. } => PROMPT_CONFIG_HELP, + Task::DraftRules { .. } => PROMPT_DRAFT_RULES, + Task::Ask { .. } => PROMPT_ASK, + } +} + +/// Split a template into its `[system]` text and example turns. +fn parse_template(text: &str) -> (String, Vec) { + let mut system = String::new(); + let mut turns: Vec = Vec::new(); + let mut current: Option<&'static str> = None; + let mut buf = String::new(); + let flush = |role: Option<&'static str>, + buf: &mut String, + system: &mut String, + turns: &mut Vec| { + let content = buf.trim().to_string(); + buf.clear(); + match role { + Some("system") => *system = content, + Some(r) if !content.is_empty() => turns.push(Message { role: r, content }), + _ => {} + } + }; + for line in text.lines() { + let role = match line.trim() { + "[system]" => Some("system"), + "[user]" => Some("user"), + "[assistant]" => Some("assistant"), + _ => None, + }; + if let Some(r) = role { + flush(current, &mut buf, &mut system, &mut turns); + current = Some(r); + } else { + buf.push_str(line); + buf.push('\n'); + } + } + flush(current, &mut buf, &mut system, &mut turns); + (system, turns) +} + +fn context_block(chunks: &[&Chunk]) -> String { + if chunks.is_empty() { + return "Context: (nothing relevant found)\n\n".into(); + } + let mut s = String::from("Context:\n"); + for (i, c) in chunks.iter().enumerate() { + s.push_str(&format!( + "[{}] {} (source: {})\n{}\n\n", + i + 1, + c.title, + c.source, + truncate_chars(c.text.trim(), CHUNK_PROMPT_CHARS) + )); + } + s +} + +fn user_content(task: &Task) -> String { + match task { + Task::Filter { question, context } => format!( + "Allowed values:\nrule_id: {}\nstorey: {}\nifc_type: {}\nseverity: {}\n\nRequest: {}", + context.rule_ids.join(", "), + context.storeys.join(", "), + context.ifc_types.join(", "), + context.severities.join(", "), + question.trim() + ), + Task::Explain { finding, rule } => { + let mut s = format!( + "Finding:\n- rule: {} (severity {})\n", + finding.rule_id, + finding.severity.as_str() + ); + let mut element = String::from("- element: "); + element.push_str(finding.ifc_type.as_deref().unwrap_or("model-level finding")); + if let Some(n) = &finding.name { + element.push_str(&format!(" \"{n}\"")); + } + if let Some(g) = &finding.global_id { + element.push_str(&format!(" (GlobalId {g})")); + } + if let Some(st) = &finding.storey { + element.push_str(&format!(", storey {st}")); + } + s.push_str(&element); + s.push('\n'); + s.push_str(&format!("- message: {}\n", finding.message)); + for (label, v) in [ + ("expected", &finding.expected), + ("actual", &finding.actual), + ("fix hint", &finding.fix_hint), + ] { + if let Some(v) = v { + s.push_str(&format!("- {label}: {v}\n")); + } + } + if let Some(r) = rule { + s.push_str(&format!( + "\nRule: {} [{}] {}. Status {:?}, {} failures out of {} elements.", + r.rule_id, r.category, r.description, r.status, r.failures, r.elements_checked + )); + } + s + } + Task::ConfigHelp { + question, + current_yaml, + } => { + let mut s = String::new(); + if let Some(y) = current_yaml.as_deref().filter(|y| !y.trim().is_empty()) { + s.push_str(&format!( + "Current project.yaml:\n{}\n\n", + truncate_chars(y.trim(), LONG_TEXT_CHARS) + )); + } + s.push_str(&format!("Question: {}", question.trim())); + s + } + Task::DraftRules { requirements_text } => format!( + "Requirements text:\n{}", + truncate_chars(requirements_text.trim(), LONG_TEXT_CHARS) + ), + Task::Ask { question } => format!("Question: {}", question.trim()), + } +} + +/// System prompt + schema, few-shot examples, then context + user content. +fn build_messages(task: &Task, chunks: &[&Chunk]) -> Vec { + let (task_system, examples) = parse_template(template(task)); + let mut messages = vec![Message { + role: "system", + content: format!("{}\n\n{}", PROMPT_SYSTEM.trim(), task_system), + }]; + messages.extend(examples); + messages.push(Message { + role: "user", + content: format!("{}{}", context_block(chunks), user_content(task)), + }); + messages +} + +fn truncate_chars(s: &str, max: usize) -> String { + match s.char_indices().nth(max) { + Some((i, _)) => format!("{}…", &s[..i]), + None => s.to_string(), + } +} + +// ---------------------------------------------------------------- parsing + +/// The first balanced `{...}` in `text` (markdown fences ignored, strings respected). +pub fn extract_json_object(text: &str) -> Option<&str> { + let start = text.find('{')?; + let bytes = text.as_bytes(); + let (mut depth, mut in_str, mut esc) = (0usize, false, false); + for (i, &b) in bytes.iter().enumerate().skip(start) { + if in_str { + match b { + _ if esc => esc = false, + b'\\' => esc = true, + b'"' => in_str = false, + _ => {} + } + continue; + } + match b { + b'"' => in_str = true, + b'{' => depth += 1, + b'}' => { + depth -= 1; + if depth == 0 { + return Some(&text[start..=i]); + } + } + _ => {} + } + } + None +} + +/// Escape raw control characters inside JSON strings (small models often put +/// literal newlines in a "yaml" value). +fn repair_json(s: &str) -> String { + let mut out = String::with_capacity(s.len()); + let (mut in_str, mut esc) = (false, false); + for ch in s.chars() { + if in_str { + if esc { + esc = false; + out.push(ch); + continue; + } + match ch { + '\\' => { + esc = true; + out.push(ch); + } + '"' => { + in_str = false; + out.push(ch); + } + '\n' => out.push_str("\\n"), + '\r' => {} + '\t' => out.push_str("\\t"), + c if (c as u32) < 0x20 => {} + c => out.push(c), + } + } else { + if ch == '"' { + in_str = true; + } + out.push(ch); + } + } + out +} + +/// Parse the first JSON object of a model output. +pub fn parse_json_object( + output: &str, +) -> Result, String> { + let obj = + extract_json_object(output).ok_or_else(|| "no complete JSON object found".to_string())?; + let value: serde_json::Value = serde_json::from_str(obj) + .or_else(|_| serde_json::from_str(&repair_json(obj))) + .map_err(|e| format!("invalid JSON: {e}"))?; + match value { + serde_json::Value::Object(m) => Ok(m), + _ => Err("the answer is not a JSON object".into()), + } +} + +/// Leading complete `"key": value` pairs of a malformed object, e.g. +/// `{"answer": "text", "id": 001}` → `{"answer": "text"}`. Used only for +/// tasks whose essential field comes first (never for filters, where a lost +/// field would silently widen the result). +fn salvage_leading_pairs(output: &str) -> Option> { + let start = output.find('{')?; + let repaired = repair_json(&output[start + 1..]); + let mut rest = repaired.as_str(); + let mut map = serde_json::Map::new(); + loop { + rest = rest.trim_start(); + if !rest.starts_with('"') { + break; + } + let mut keys = serde_json::Deserializer::from_str(rest).into_iter::(); + let Some(Ok(key)) = keys.next() else { break }; + rest = rest[keys.byte_offset()..].trim_start(); + let Some(after_colon) = rest.strip_prefix(':') else { + break; + }; + let mut values = + serde_json::Deserializer::from_str(after_colon).into_iter::(); + let Some(Ok(value)) = values.next() else { + break; + }; + rest = &after_colon[values.byte_offset()..]; + // A value is complete only if a separator follows it. + let next = rest.trim_start(); + if !(next.starts_with(',') || next.starts_with('}')) { + break; + } + map.insert(key, value); + match next.strip_prefix(',') { + Some(r) => rest = r, + None => break, + } + } + (!map.is_empty()).then_some(map) +} + +/// A string field; JSON null, missing, empty and "null"/"none" are `None`. +fn opt_str(m: &serde_json::Map, key: &str) -> Option { + match m.get(key)? { + serde_json::Value::String(s) => { + let t = s.trim(); + if t.is_empty() + || ["null", "none", "any", "all", "n/a"].contains(&t.to_lowercase().as_str()) + { + None + } else { + Some(t.to_string()) + } + } + serde_json::Value::Number(n) => Some(n.to_string()), + _ => None, + } +} + +fn first_str(m: &serde_json::Map, keys: &[&str]) -> Option { + keys.iter().find_map(|k| opt_str(m, k)) +} + +fn parse_answer(task: &Task, output: &str) -> Result { + let m = match parse_json_object(output) { + Ok(m) => m, + Err(e) if matches!(task, Task::Filter { .. }) => return Err(e), + Err(e) => salvage_leading_pairs(output).ok_or(e)?, + }; + match task { + Task::Filter { context, .. } => Ok(Answer::Filter { + filter: validate_filter(&m, context)?, + }), + Task::Explain { .. } => { + let text = first_str(&m, &["text", "explanation", "answer"]) + .ok_or_else(|| "missing \"text\"".to_string())?; + Ok(Answer::Explanation { text }) + } + Task::Ask { .. } => { + let text = first_str(&m, &["answer", "text"]) + .ok_or_else(|| "missing \"answer\"".to_string())?; + Ok(Answer::Text { + text, + sources: vec![], + structured: true, + }) + } + Task::ConfigHelp { .. } | Task::DraftRules { .. } => { + let answer = snippet_from(&m)?; + match &answer { + Answer::ConfigSnippet { + valid: false, + validation_error, + .. + } => Err(format!( + "the YAML is not a valid VeriModel configuration: {}", + validation_error.clone().unwrap_or_default() + )), + _ => Ok(answer), + } + } + } +} + +// ---------------------------------------------------------------- filter + +/// Canonical value from `allowed` (case-insensitive; plural "errors" and a +/// missing "Ifc" prefix are tolerated). +fn canonical(value: &str, allowed: &[String]) -> Option { + let v = value.trim().to_lowercase(); + let find = |cand: &str| { + allowed + .iter() + .find(|a| a.trim().to_lowercase() == cand) + .cloned() + }; + find(&v) + .or_else(|| v.strip_suffix('s').and_then(find)) + .or_else(|| { + if v.starts_with("ifc") { + None + } else { + find(&format!("ifc{v}")) + } + }) +} + +fn validate_filter( + m: &serde_json::Map, + ctx: &ReportContext, +) -> Result { + let known = [ + "severity", + "rule_id", + "storey", + "ifc_type", + "search", + "explanation", + ]; + if !m.keys().any(|k| known.contains(&k.as_str())) { + return Err(format!("expected the keys {}", known.join(", "))); + } + let mut f = FilterSpec { + explanation: opt_str(m, "explanation").unwrap_or_default(), + search: opt_str(m, "search"), + ..FilterSpec::default() + }; + let severities: Vec = match m.get("severity") { + Some(serde_json::Value::Array(a)) => a + .iter() + .filter_map(|v| v.as_str().map(str::to_string)) + .collect(), + Some(serde_json::Value::String(s)) => s.split(',').map(str::to_string).collect(), + _ => vec![], + }; + for s in severities + .iter() + .map(|s| s.trim()) + .filter(|s| !s.is_empty()) + { + match canonical(s, &ctx.severities) { + Some(c) if !f.severity.contains(&c) => f.severity.push(c), + Some(_) => {} + None => f.dropped.push(format!("severity: {s}")), + } + } + let mut field = |key: &str, allowed: &[String]| -> Option { + let v = opt_str(m, key)?; + let c = canonical(&v, allowed); + if c.is_none() { + f.dropped.push(format!("{key}: {v}")); + } + c + }; + let rule_id = field("rule_id", &ctx.rule_ids); + let storey = field("storey", &ctx.storeys); + let ifc_type = field("ifc_type", &ctx.ifc_types); + f.rule_id = rule_id; + f.storey = storey; + f.ifc_type = ifc_type; + Ok(f) +} + +// ---------------------------------------------------------------- config snippets + +/// Top-level keys of `project.yaml`. +const CONFIG_SECTIONS: &[&str] = &[ + "project", + "requirements", + "tolerances", + "qto", + "clash", + "output", + "checks", + "naming", + "classification", +]; + +/// Section that owns a key the model may have written without its parent. +fn section_of(key: &str) -> Option<&'static str> { + Some(match key { + "required_psets" | "ids_file" | "rule_pack" => "requirements", + "patterns" => "naming", + "required" | "applies_to" | "allowed_codes" | "allowed_codes_file" | "code_pattern" => { + "classification" + } + "rules" | "categories" | "unique_name_within" => "checks", + "clash_hard_mm" | "clash_clearance_mm" | "far_from_origin_m" | "geometry" => "tolerances", + "group_by" | "unit_rates_file" | "currency" | "sample_count" => "qto", + _ => return None, + }) +} + +fn strip_fences(s: &str) -> String { + let t = s.trim(); + if !t.starts_with("```") { + return t.to_string(); + } + let body = t.split_once('\n').map(|(_, b)| b).unwrap_or(""); + body.trim_end() + .trim_end_matches("```") + .trim_end() + .to_string() +} + +/// Wrap a fragment (e.g. a bare `required_psets:` list) into full +/// `project.yaml` sections, then validate with `Config::from_yaml`. +/// Returns (normalised YAML, validation error). +pub fn normalize_and_validate_yaml(yaml: &str) -> (String, Option) { + let text = strip_fences(yaml); + if text.trim().is_empty() { + return (text, Some("the YAML is empty".into())); + } + let value: serde_yaml::Value = match serde_yaml::from_str(&text) { + Ok(v) => v, + Err(e) => return (text, Some(format!("YAML syntax error: {e}"))), + }; + let mut text = text; + if let serde_yaml::Value::Mapping(map) = &value { + let keys: Vec = map + .keys() + .filter_map(|k| k.as_str().map(str::to_string)) + .collect(); + let needs_wrap = keys + .iter() + .any(|k| !CONFIG_SECTIONS.contains(&k.as_str()) && section_of(k).is_some()); + let all_known = keys + .iter() + .all(|k| CONFIG_SECTIONS.contains(&k.as_str()) || section_of(k).is_some()); + if needs_wrap && all_known { + let mut out = serde_yaml::Mapping::new(); + for (k, v) in map { + let key = k.as_str().unwrap_or_default(); + let (section, inner): (&str, Option<&str>) = if CONFIG_SECTIONS.contains(&key) { + (key, None) + } else { + (section_of(key).unwrap_or(key), Some(key)) + }; + let skey = serde_yaml::Value::String(section.to_string()); + match inner { + None => match (out.get_mut(&skey), v) { + ( + Some(serde_yaml::Value::Mapping(existing)), + serde_yaml::Value::Mapping(new), + ) => { + for (nk, nv) in new { + existing.insert(nk.clone(), nv.clone()); + } + } + _ => { + out.insert(skey, v.clone()); + } + }, + Some(inner) => { + let entry = out.entry(skey).or_insert_with(|| { + serde_yaml::Value::Mapping(serde_yaml::Mapping::new()) + }); + if let serde_yaml::Value::Mapping(m) = entry { + m.insert(serde_yaml::Value::String(inner.to_string()), v.clone()); + } + } + } + } + if let Ok(s) = serde_yaml::to_string(&serde_yaml::Value::Mapping(out)) { + text = s; + } + } + } + let err = verimodel_core::config::Config::from_yaml(&text) + .err() + .map(|e| e.to_string()); + (text.trim_end().to_string() + "\n", err) +} + +fn snippet_from(m: &serde_json::Map) -> Result { + let yaml = match m.get("yaml") { + Some(serde_json::Value::String(s)) => s.clone(), + // A model may emit the YAML as a nested object: render it. + Some(v @ serde_json::Value::Object(_)) => { + serde_yaml::to_string(v).map_err(|e| e.to_string())? + } + _ => return Err("missing \"yaml\"".into()), + }; + let explanation = opt_str(m, "explanation").unwrap_or_default(); + let (yaml, validation_error) = normalize_and_validate_yaml(&yaml); + Ok(Answer::ConfigSnippet { + yaml, + valid: validation_error.is_none(), + validation_error, + explanation, + }) +} + +fn parse_snippet(output: &str) -> Option { + parse_json_object(output) + .ok() + .and_then(|m| snippet_from(&m).ok()) +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::backend::MockBackend; + use verimodel_core::model::{RuleSource, RuleStatus, Severity}; + + fn ctx() -> ReportContext { + ReportContext { + rule_ids: vec!["ID-001".into(), "SP-001".into()], + storeys: vec!["Level 1".into(), "Level 2".into()], + ifc_types: vec!["IfcWall".into(), "IfcSlab".into()], + severities: vec!["ERROR".into(), "WARNING".into(), "INFO".into()], + } + } + + fn finding() -> Finding { + Finding { + rule_id: "ID-001".into(), + severity: Severity::Error, + global_id: Some("0abcDEF123ghiJKL456mno".into()), + step_id: Some(42), + ifc_type: Some("IfcWall".into()), + name: Some("Basic Wall 200".into()), + storey: Some("Level 1".into()), + message: "GlobalId is duplicated".into(), + expected: Some("unique GlobalId".into()), + actual: Some("used by 2 elements".into()), + fix_hint: Some("Regenerate the GlobalId".into()), + source: RuleSource::RulePack { + pack: "generic".into(), + }, + } + } + + fn rule() -> RuleResult { + RuleResult { + rule_id: "ID-001".into(), + category: "identity".into(), + description: "GlobalId is unique".into(), + severity: Severity::Error, + status: RuleStatus::Fail, + elements_checked: 10, + failures: 2, + skipped_reason: None, + source: RuleSource::RulePack { + pack: "generic".into(), + }, + } + } + + fn knowledge() -> Knowledge { + Knowledge { + chunks: vec![ + Chunk { source: "docs/x.md".into(), title: "Duplicate GlobalId".into(), text: "Rule ID-001 checks that every GlobalId is unique.".into(), keywords: vec!["id-001".into()] }, + Chunk { source: "docs/y.md".into(), title: "Required psets".into(), text: "requirements.required_psets lists pset, applies_to and properties. FireRating walls Pset_WallCommon.".into(), keywords: vec![] }, + ], + } + } + + fn run_mock(answers: &[&str], task: &Task) -> (Answer, MockBackend) { + let mut b = MockBackend::new(answers.iter().copied()); + let a = run(&mut b, &knowledge(), task).unwrap(); + (a, b) + } + + fn filter_task(q: &str) -> Task { + Task::Filter { + question: q.into(), + context: ctx(), + } + } + + #[test] + fn prompt_contains_retrieved_chunk_and_schema() { + let (_, b) = run_mock( + &[r#"{"answer":"It checks GlobalId uniqueness."}"#], + &Task::Ask { + question: "What does ID-001 check?".into(), + }, + ); + let p = &b.prompts[0]; + assert!(p.contains("Duplicate GlobalId"), "{p}"); + assert!(p.contains(r#"{"answer": string}"#), "{p}"); + assert!( + p.contains("you never decide") || p.contains("never decide whether anything passes"), + "{p}" + ); + assert!(p.starts_with("<|im_start|>system\n")); + assert!(p.ends_with("<|im_start|>assistant\n{\"answer\":"), "{p}"); + } + + #[test] + fn ask_valid_answer_has_sources() { + let (a, _) = run_mock( + &[r#"{"answer":"It checks GlobalId uniqueness."}"#], + &Task::Ask { + question: "What does ID-001 check?".into(), + }, + ); + assert_eq!( + a, + Answer::Text { + text: "It checks GlobalId uniqueness.".into(), + sources: vec!["Duplicate GlobalId".into()], + structured: true + } + ); + } + + #[test] + fn ask_invalid_twice_falls_back_to_plain_text() { + let (a, b) = run_mock( + &["I think it checks ids", "still no json"], + &Task::Ask { + question: "What does ID-001 check?".into(), + }, + ); + assert_eq!( + a, + Answer::Text { + text: "still no json".into(), + sources: vec!["Duplicate GlobalId".into()], + structured: false + } + ); + assert_eq!(b.prompts.len(), 2); + assert!(b.prompts[1].contains("Format error in your reply (no complete JSON object found)")); + assert!( + b.prompts[1].ends_with( + "schema: {\"answer\": string}<|im_end|>\n<|im_start|>assistant\n{\"answer\":" + ), + "{}", + b.prompts[1] + ); + assert!(b.prompts[1].contains("I think it checks ids")); + } + + #[test] + fn primed_prefix_is_reattached() { + // The model continues after `{"answer":` without repeating it. + let (a, _) = run_mock( + &[r#" "GlobalIds must be unique."}"#], + &Task::Ask { + question: "What does ID-001 check?".into(), + }, + ); + assert!( + matches!(a, Answer::Text { ref text, structured: true, .. } if text == "GlobalIds must be unique."), + "{a:?}" + ); + // Truncated twice: plain text without the JSON scaffolding. + let (a, _) = run_mock( + &[r#" "GlobalIds must"#, r#" "GlobalIds must be"#], + &Task::Ask { + question: "What does ID-001 check?".into(), + }, + ); + assert!( + matches!(a, Answer::Text { ref text, structured: false, .. } if text == "GlobalIds must be"), + "{a:?}" + ); + } + + #[test] + fn filter_valid_with_canonicalisation() { + let (a, b) = run_mock( + &["```json\n{\"severity\": [\"errors\"], \"rule_id\": null, \"storey\": \"level 1\", \"ifc_type\": \"wall\", \"search\": \"\", \"explanation\": \"Errors on walls.\"}\n```"], + &filter_task("show me the errors on Level 1 for walls"), + ); + let Answer::Filter { filter } = a else { + panic!("{a:?}") + }; + assert_eq!(filter.severity, vec!["ERROR"]); + assert_eq!(filter.storey.as_deref(), Some("Level 1")); + assert_eq!(filter.ifc_type.as_deref(), Some("IfcWall")); + assert_eq!(filter.rule_id, None); + assert_eq!(filter.search, None); + assert!(filter.dropped.is_empty()); + assert!(b.prompts[0].contains("rule_id: ID-001, SP-001")); + } + + #[test] + fn filter_invalid_json_then_valid_on_retry() { + let (a, b) = run_mock( + &[ + r#"{"severity": ["ERROR"], "storey": "Level 1""#, + r#"{"severity": ["ERROR"], "storey": "Level 1"}"#, + ], + &filter_task("errors on level 1"), + ); + let Answer::Filter { filter } = a else { + panic!("{a:?}") + }; + assert_eq!(filter.severity, vec!["ERROR"]); + assert_eq!(filter.storey.as_deref(), Some("Level 1")); + assert_eq!(b.prompts.len(), 2); + } + + #[test] + fn filter_hallucinated_values_are_dropped() { + let (a, _) = run_mock( + &[ + r#"{"severity": ["ERROR", "CRITICAL"], "rule_id": "XX-999", "storey": "Level 9", "ifc_type": "IfcSlab", "explanation": "x"}"#, + ], + &filter_task("critical slab problems"), + ); + let Answer::Filter { filter } = a else { + panic!("{a:?}") + }; + assert_eq!(filter.rule_id, None); + assert_eq!(filter.storey, None); + assert_eq!(filter.ifc_type.as_deref(), Some("IfcSlab")); + assert_eq!(filter.severity, vec!["ERROR"]); + assert_eq!( + filter.dropped, + vec!["severity: CRITICAL", "rule_id: XX-999", "storey: Level 9"] + ); + } + + #[test] + fn explain_valid_and_prompt_has_facts() { + let task = Task::Explain { + finding: Box::new(finding()), + rule: Some(rule()), + }; + let (a, b) = run_mock( + &[ + r#"{"text": "Basic Wall 200 (0abcDEF123ghiJKL456mno) shares its GlobalId. Regenerate it."}"#, + ], + &task, + ); + assert_eq!( + a, + Answer::Explanation { + text: "Basic Wall 200 (0abcDEF123ghiJKL456mno) shares its GlobalId. Regenerate it." + .into() + } + ); + let p = &b.prompts[0]; + assert!( + p.contains( + "IfcWall \"Basic Wall 200\" (GlobalId 0abcDEF123ghiJKL456mno), storey Level 1" + ), + "{p}" + ); + assert!(p.contains("Rule: ID-001 [identity] GlobalId is unique. Status Fail, 2 failures out of 10 elements."), "{p}"); + assert!(p.contains("Duplicate GlobalId")); + } + + #[test] + fn explain_invalid_then_valid_on_retry() { + let task = Task::Explain { + finding: Box::new(finding()), + rule: None, + }; + let (a, _) = run_mock( + &[ + "The wall has a duplicated id.", + "{\"text\": \"Line one.\nLine two.\"}", + ], + &task, + ); + assert_eq!( + a, + Answer::Explanation { + text: "Line one.\nLine two.".into() + } + ); + } + + #[test] + fn config_help_valid_snippet() { + let task = Task::ConfigHelp { + question: "require FireRating on all walls".into(), + current_yaml: Some("project:\n name: Demo\n".into()), + }; + let yaml = "requirements:\n required_psets:\n - pset: Pset_WallCommon\n applies_to: [IfcWall]\n properties: [FireRating]\n"; + let out = + serde_json::json!({"yaml": yaml, "explanation": "Walls need FireRating."}).to_string(); + let (a, b) = run_mock(&[out.as_str()], &task); + assert_eq!( + a, + Answer::ConfigSnippet { + yaml: yaml.into(), + valid: true, + validation_error: None, + explanation: "Walls need FireRating.".into() + } + ); + assert!(b.prompts[0].contains("Current project.yaml:\nproject:\n name: Demo")); + assert!(b.prompts[0].contains(r#"{"yaml": string, "explanation": string}"#)); + } + + #[test] + fn config_help_invalid_yaml_is_flagged_after_retry() { + let task = Task::ConfigHelp { + question: "require FireRating on all walls".into(), + current_yaml: None, + }; + let bad = serde_json::json!({"yaml": "requirements:\n required_psets:\n - name: Pset_WallCommon\n", "explanation": "x"}).to_string(); + let (a, b) = run_mock(&[bad.as_str(), bad.as_str()], &task); + let Answer::ConfigSnippet { + valid, + validation_error, + yaml, + .. + } = a + else { + panic!("{a:?}") + }; + assert!(!valid); + assert!(validation_error.unwrap().contains("unknown field")); + assert!(yaml.contains("Pset_WallCommon")); + assert_eq!(b.prompts.len(), 2); + assert!(b.prompts[1].contains("not a valid VeriModel configuration")); + } + + #[test] + fn config_help_invalid_then_valid_on_retry() { + let task = Task::ConfigHelp { + question: "door names D-001".into(), + current_yaml: None, + }; + let bad = + serde_json::json!({"yaml": "naming:\n patterns:\n IfcDoor: \"^D-[0-9{3}$\"\n"}) + .to_string(); + let good = + serde_json::json!({"yaml": "naming:\n patterns:\n IfcDoor: \"^D-[0-9]{3}$\"\n"}) + .to_string(); + let (a, _) = run_mock(&[bad.as_str(), good.as_str()], &task); + assert!( + matches!(a, Answer::ConfigSnippet { valid: true, .. }), + "{a:?}" + ); + } + + #[test] + fn draft_rules_fragment_is_wrapped_and_validated() { + let task = Task::DraftRules { + requirements_text: + "All walls shall carry Pset_WallCommon.FireRating. Walls are named W-###.".into(), + }; + // Literal newlines inside the JSON string: repaired by the parser. + let out = "{\"yaml\": \"required_psets:\n - pset: Pset_WallCommon\n applies_to: [IfcWall]\n properties: [FireRating]\npatterns:\n IfcWall: '^W-[0-9]{3}$'\n\", \"explanation\": \"From the two sentences.\"}"; + let (a, b) = run_mock(&[out], &task); + let Answer::ConfigSnippet { + yaml, + valid, + validation_error, + .. + } = a + else { + panic!("{a:?}") + }; + assert!(valid, "{validation_error:?}\n{yaml}"); + assert!( + yaml.starts_with("requirements:\n required_psets:"), + "{yaml}" + ); + assert!(yaml.contains("naming:\n patterns:"), "{yaml}"); + assert!(b.prompts[0].contains("Requirements text:\nAll walls shall carry")); + } + + #[test] + fn draft_rules_invalid_twice_flagged_not_applied() { + let task = Task::DraftRules { + requirements_text: "Walls need a fire rating.".into(), + }; + let bad = serde_json::json!({"yaml": "required_psets:\n - pset: Pset_WallCommon\n severity: fatal\n"}).to_string(); + let (a, _) = run_mock(&[bad.as_str(), bad.as_str()], &task); + let Answer::ConfigSnippet { + valid, + validation_error, + .. + } = a + else { + panic!("{a:?}") + }; + assert!(!valid); + assert!(validation_error + .unwrap() + .contains("severity 'fatal' is invalid")); + } + + #[test] + fn non_json_config_answer_falls_back_to_text() { + let task = Task::ConfigHelp { + question: "x".into(), + current_yaml: None, + }; + let (a, _) = run_mock(&["no idea", "still no idea"], &task); + assert!( + matches!( + a, + Answer::Text { + structured: false, + .. + } + ), + "{a:?}" + ); + } + + #[test] + fn backend_errors_propagate() { + let mut b = MockBackend::new(Vec::::new()); + assert!(run( + &mut b, + &knowledge(), + &Task::Ask { + question: "x".into() + } + ) + .is_err()); + } + + #[test] + fn salvage_keeps_complete_leading_pairs() { + let m = salvage_leading_pairs( + r#"{"answer": "Every element must have a unique GlobalId.", "id": 001}"#, + ) + .unwrap(); + assert_eq!(m.len(), 1); + assert_eq!(m["answer"], "Every element must have a unique GlobalId."); + assert!(salvage_leading_pairs(r#"{"answer": "trunc"#).is_none()); + let (a, b) = run_mock( + &[r#" "Unique GlobalIds.", "id": 001}"#], + &Task::Ask { + question: "What does ID-001 check?".into(), + }, + ); + assert!( + matches!(a, Answer::Text { ref text, structured: true, .. } if text == "Unique GlobalIds."), + "{a:?}" + ); + assert_eq!(b.prompts.len(), 1); + // Filters are never salvaged: a lost field would widen the filter. + let (_, b) = run_mock( + &[ + r#"{"severity": ["ERROR"], "storey": 01}"#, + r#"{"severity": ["ERROR"]}"#, + ], + &filter_task("errors"), + ); + assert_eq!(b.prompts.len(), 2); + } + + #[test] + fn json_extraction() { + assert_eq!( + extract_json_object("text {\"a\": \"}{\", \"b\": {\"c\": 1}} tail {}"), + Some("{\"a\": \"}{\", \"b\": {\"c\": 1}}") + ); + assert_eq!(extract_json_object("{\"a\": 1"), None); + assert!(parse_json_object("[1,2]").is_err()); + } + + #[test] + fn templates_parse() { + for t in [ + PROMPT_FILTER, + PROMPT_EXPLAIN, + PROMPT_CONFIG_HELP, + PROMPT_DRAFT_RULES, + PROMPT_ASK, + ] { + let (system, turns) = parse_template(t); + assert!(system.contains("JSON schema:"), "{system}"); + assert!(turns.len() % 2 == 0); + for pair in turns.chunks(2) { + assert_eq!(pair[0].role, "user"); + assert_eq!(pair[1].role, "assistant"); + assert!( + parse_json_object(&pair[1].content).is_ok(), + "{}", + pair[1].content + ); + } + } + // Few-shot config examples must themselves be valid configurations. + for t in [PROMPT_CONFIG_HELP, PROMPT_DRAFT_RULES] { + for turn in parse_template(t).1.iter().filter(|m| m.role == "assistant") { + let a = parse_snippet(&turn.content).unwrap(); + assert!( + matches!(a, Answer::ConfigSnippet { valid: true, .. }), + "{a:?}" + ); + } + } + } } diff --git a/crates/assistant/tests/real_model.rs b/crates/assistant/tests/real_model.rs new file mode 100644 index 0000000..7457e67 --- /dev/null +++ b/crates/assistant/tests/real_model.rs @@ -0,0 +1,278 @@ +//! Smoke test against the real default model (Qwen2.5-0.5B-Instruct Q4_K_M). +//! +//! Ignored by default. Run with a folder holding `model.gguf` (or the original +//! `qwen2.5-0.5b-instruct-q4_k_m.gguf`) and `tokenizer.json`: +//! +//! ```sh +//! VERIMODEL_ASSISTANT_MODEL_DIR=/path/to/models \ +//! cargo test -p verimodel-assistant --release --test real_model -- --ignored --nocapture +//! ``` +//! +//! Use a release build: debug inference is far too slow. +#![cfg(feature = "candle")] + +use std::path::{Path, PathBuf}; +use std::time::Instant; + +use verimodel_assistant::backend::candle::CandleBackend; +use verimodel_assistant::backend::{Backend, GenOptions}; +use verimodel_assistant::model::{sha256_file, DEFAULT_MODEL}; +use verimodel_assistant::tasks::{self, Answer, ReportContext, Task}; +use verimodel_assistant::AssistantResult; +use verimodel_assistant::Knowledge; +use verimodel_core::model::{Finding, RuleResult, RuleSource, RuleStatus, Severity}; + +fn find_model_files(dir: &Path) -> Option<(PathBuf, PathBuf)> { + let tok = dir.join("tokenizer.json"); + let gguf = [ + dir.join("model.gguf"), + dir.join(format!("{}.gguf", DEFAULT_MODEL.id)), + ] + .into_iter() + .find(|p| p.is_file())?; + tok.is_file().then_some((gguf, tok)) +} + +/// Resident memory of this process in MB: (current, peak). +fn rss_mb() -> (u64, u64) { + let status = std::fs::read_to_string("/proc/self/status").unwrap_or_default(); + let field = |name: &str| { + status + .lines() + .find(|l| l.starts_with(name)) + .and_then(|l| l.split_whitespace().nth(1)) + .and_then(|v| v.parse::().ok()) + .unwrap_or(0) + / 1024 + }; + (field("VmRSS:"), field("VmHWM:")) +} + +/// Prints every raw model output (and its timing) before the task parses it. +struct Logged<'a>(&'a mut CandleBackend); + +impl Backend for Logged<'_> { + fn name(&self) -> &str { + self.0.name() + } + fn generate(&mut self, prompt: &str, opts: &GenOptions) -> AssistantResult { + let out = self.0.generate(prompt, opts)?; + let s = self.0.last_stats().unwrap_or_default(); + println!( + " raw output ({} prompt tokens, prefill {:.2} s = {:.0} tok/s; {} generated, decode {:.2} s = {:.1} tok/s): {}", + s.prompt_tokens, + s.prefill_secs, + s.prefill_tokens_per_sec(), + s.generated_tokens, + s.decode_secs, + s.decode_tokens_per_sec(), + out.trim() + ); + Ok(out) + } +} + +fn run(backend: &mut CandleBackend, knowledge: &Knowledge, label: &str, task: &Task) -> Answer { + println!("\n=== {label}"); + let t = Instant::now(); + let answer = tasks::run(&mut Logged(backend), knowledge, task).expect("generation failed"); + println!(" total {:.1} s", t.elapsed().as_secs_f64()); + println!("{}", serde_json::to_string_pretty(&answer).unwrap()); + answer +} + +#[test] +#[ignore = "needs VERIMODEL_ASSISTANT_MODEL_DIR and a release build"] +fn real_model_smoke() { + let Some(dir) = std::env::var_os("VERIMODEL_ASSISTANT_MODEL_DIR") else { + panic!("set VERIMODEL_ASSISTANT_MODEL_DIR to a folder with model.gguf and tokenizer.json"); + }; + let (gguf, tok) = find_model_files(Path::new(&dir)) + .expect("model.gguf / tokenizer.json not found in VERIMODEL_ASSISTANT_MODEL_DIR"); + + // Verify the files against the catalogue checksums (install_from_files does the same + // after copying; it is unit-tested in model.rs, so the 0.5 GB copy is skipped here). + let t = Instant::now(); + assert_eq!( + sha256_file(&gguf).unwrap(), + DEFAULT_MODEL.gguf_sha256, + "model.gguf checksum" + ); + assert_eq!( + sha256_file(&tok).unwrap(), + DEFAULT_MODEL.tokenizer_sha256, + "tokenizer.json checksum" + ); + println!("SHA-256 verification: {:.1} s", t.elapsed().as_secs_f64()); + + let (rss0, _) = rss_mb(); + let t = Instant::now(); + let mut backend = CandleBackend::load(&gguf, &tok).expect("load"); + let (rss1, _) = rss_mb(); + println!( + "load: {:.2} s, family {}, RSS {} MB -> {} MB", + t.elapsed().as_secs_f64(), + backend.family(), + rss0, + rss1 + ); + assert_eq!(backend.family(), "qwen2"); + + let knowledge = Knowledge::builtin(); + let mut failures: Vec = Vec::new(); + + // (a) Ask + let a = run( + &mut backend, + &knowledge, + "Ask: What does rule ID-001 check?", + &Task::Ask { + question: "What does rule ID-001 check?".into(), + }, + ); + match &a { + Answer::Text { + text, structured, .. + } => { + if !text.to_lowercase().contains("globalid") { + failures.push(format!("(a) answer does not mention GlobalId: {text}")); + } + if !structured { + failures.push("(a) answer was not valid JSON (plain-text fallback)".into()); + } + } + other => failures.push(format!("(a) unexpected answer kind: {other:?}")), + } + + // (b) Filter + let context = ReportContext { + rule_ids: vec!["ID-001".into(), "SP-001".into()], + storeys: vec!["Level 1".into(), "Level 2".into()], + ifc_types: vec!["IfcWall".into(), "IfcSlab".into()], + severities: vec!["ERROR".into(), "WARNING".into(), "INFO".into()], + }; + let b = run( + &mut backend, + &knowledge, + "Filter: show me the errors on Level 1 for walls", + &Task::Filter { + question: "show me the errors on Level 1 for walls".into(), + context, + }, + ); + match &b { + Answer::Filter { filter } => { + if filter.severity != vec!["ERROR".to_string()] + || filter.storey.as_deref() != Some("Level 1") + || filter.ifc_type.as_deref() != Some("IfcWall") + { + failures.push(format!("(b) wrong filter: {filter:?}")); + } + } + other => failures.push(format!("(b) unexpected answer kind: {other:?}")), + } + + // (c) Explain + let finding = Finding { + rule_id: "SP-001".into(), + severity: Severity::Error, + global_id: Some("1hOSvn6df7F8_7GcBWlRGQ".into()), + step_id: Some(1234), + ifc_type: Some("IfcWall".into()), + name: Some("Basic Wall:Exterior 300".into()), + storey: None, + message: "Element is not contained in any spatial structure element".into(), + expected: Some("contained in a storey, space, site or facility part".into()), + actual: Some("no spatial container".into()), + fix_hint: Some( + "Assign the element to a level / storey in the authoring tool before exporting.".into(), + ), + source: RuleSource::RulePack { + pack: "generic".into(), + }, + }; + let rule = RuleResult { + rule_id: "SP-001".into(), + category: "spatial".into(), + description: "Every element must be contained in a spatial structure element (storey, space, site, facility part)".into(), + severity: Severity::Error, + status: RuleStatus::Fail, + elements_checked: 412, + failures: 3, + skipped_reason: None, + source: RuleSource::RulePack { pack: "generic".into() }, + }; + let c = run( + &mut backend, + &knowledge, + "Explain: SP-001 finding", + &Task::Explain { + finding: Box::new(finding), + rule: Some(rule), + }, + ); + match &c { + Answer::Explanation { text } | Answer::Text { text, .. } => { + let t = text.to_lowercase(); + if text.trim().is_empty() + || !(t.contains("basic wall") || t.contains("1hosvn6df7f8_7gcbwlrgq")) + { + failures.push(format!( + "(c) explanation does not mention the element: {text}" + )); + } + if matches!(c, Answer::Text { .. }) { + failures.push("(c) explanation was not valid JSON (plain-text fallback)".into()); + } + } + other => failures.push(format!("(c) unexpected answer kind: {other:?}")), + } + + // (d) ConfigHelp + let d = run( + &mut backend, + &knowledge, + "ConfigHelp: require FireRating on all walls", + &Task::ConfigHelp { + question: "require FireRating on all walls".into(), + current_yaml: None, + }, + ); + match &d { + Answer::ConfigSnippet { + yaml, + valid, + validation_error, + .. + } => { + if !(yaml.contains("Pset_WallCommon") && yaml.contains("FireRating") && *valid) { + failures.push(format!( + "(d) wrong snippet (valid {valid}, error {validation_error:?}):\n{yaml}" + )); + } + } + other => failures.push(format!("(d) unexpected answer kind: {other:?}")), + } + + // Determinism: greedy decoding gives the same answer twice. + let d2 = tasks::run( + &mut backend, + &knowledge, + &Task::ConfigHelp { + question: "require FireRating on all walls".into(), + current_yaml: None, + }, + ) + .unwrap(); + if d2 != d { + failures.push("(d) not deterministic".into()); + } + + let (rss, peak) = rss_mb(); + println!("\nRSS after the run: {rss} MB (peak {peak} MB)"); + assert!( + failures.is_empty(), + "failed checks:\n{}", + failures.join("\n") + ); +} diff --git a/docs/RELEASE_NOTES.md b/docs/RELEASE_NOTES.md index 2be712e..48679af 100644 --- a/docs/RELEASE_NOTES.md +++ b/docs/RELEASE_NOTES.md @@ -16,6 +16,7 @@ First release of VeriModel: a local-first IFC quality gate, quantity takeoff and ### Requirements - A Python 3.9+ with `ifcopenshell` and `ifctester` (`pip install ifcopenshell ifctester`). Run `verimodel doctor` to check. +- For the optional assistant: an x86-64 CPU with AVX2 (2013 or later) or Apple Silicon, about 1.5 GB of free RAM, and a one-time 0.5 GB model download (Qwen2.5-0.5B-Instruct, Apache-2.0). A question takes roughly 30–50 s on a 4-core CPU. ### Assets - `VeriModel_1.0.0_x64-setup.exe` / `.msi` (Windows), `.dmg` (macOS), `.AppImage` / `.deb` (Linux)