@lzhzzzzwill/cofos 1.0.1 → 1.0.4

package/runtime/src/retrieval/text_search.py

@@ -0,0 +1,60 @@
+"""Shared lexical tokenization for BM25 indexing and querying.
+
+The tokenizer preserves chemical notation (·OH, O₂·⁻, SO₄·⁻, etc.) so that
+chemically distinct species stay distinct in the BM25 index.  The previous
+\\b\\w+\\b regex stripped radical dots, charge superscripts, and Unicode
+subscripts, collapsing ·OH → oh (same as OH⁻) and O₂·⁻ → o2 (same as O₂).
+
+IMPORTANT: after changing this tokenizer you MUST rebuild the BM25 index:
+    python scripts/build_retrieval.py
+"""
+
+import re
+import unicodedata
+from typing import List
+
+
+STOPWORDS = {
+    "a", "an", "and", "are", "as", "at", "be", "by", "does", "for",
+    "from", "how", "in", "is", "it", "of", "on", "or", "the", "to",
+    "under", "what", "which", "with",
+}
+
+# Punctuation characters that are stripped from token *boundaries* only.
+# Internal occurrences (e.g. hyphens in COF names, dots in radicals) are kept.
+_BOUNDARY_STRIP = ''.join(
+    set('.,;:!?()[]{}"\'“”‘’«»、。，；：！？）】…›‹«»„‚')
+)
+
+
+def tokenize_search_text(text: str) -> List[str]:
+    """Tokenize text preserving chemical notation (·OH, O₂·⁻, SO₄·⁻, etc.).
+
+    Splits on whitespace and strips boundary punctuation only.  Internal
+    special characters (radical dots, charge marks, hyphens, subscripts) are
+    kept so that chemically distinct species produce distinct tokens.
+    """
+    if not text:
+        return []
+
+    # NFKC normalisation keeps sub/superscripts readable while canonicalising
+    # fullwidth punctuation and ligatures.
+    text = unicodedata.normalize('NFKC', text.lower())
+    # Normalise common chemical-notation variants to a single form.
+    text = text.replace('−', '-')   # minus sign
+    text = text.replace('⋅', '·')  # dot operator → middle dot
+    text = text.replace('•', '·')  # bullet → middle dot
+
+    tokens: List[str] = []
+    for raw_token in text.split():
+        token = raw_token.strip(_BOUNDARY_STRIP)
+        if not token:
+            continue
+        if token in STOPWORDS:
+            continue
+        # Drop single-char punctuation-only tokens.
+        if len(token) == 1 and not token.isalnum():
+            continue
+        tokens.append(token)
+
+    return tokens

package/scripts/chat.py

@@ -1,85 +1,503 @@
 #!/usr/bin/env python3
-"""Persistent chat: loads an LLM once, reads prompts from stdin, streams to stdout."""
+"""Persistent COFOS RAG chat backend.
+
+Loads the merged model from Hugging Face (Willlzh/COFOS) and auto-downloads the
+KG + BM25 retrieval indexes from Willlzh/COFOS_data/runtime when they are not
+present locally.  Serves the Node CLI through stdin/stdout protocol markers.
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+import logging
+import os
 import sys
-from threading import Thread
+import time
+from pathlib import Path
+
+import re
 
-from transformers import (
-    AutoModelForCausalLM,
-    AutoTokenizer,
-    TextIteratorStreamer,
+HISTORY_FILE = os.path.join(os.path.expanduser("~"), ".cofos_history.jsonl")
+HISTORY_MAX_TURNS = 20
+
+DEFAULT_MODEL_PATH = os.environ.get(
+    "COFOS_MODEL",
+    "Willlzh/COFOS",
 )
+DEFAULT_CONFIG = os.environ.get("COFOS_CONFIG", "config/config.yaml")
+CACHE_ROOT = Path(os.environ.get("COFOS_CACHE_DIR", Path.home() / ".cache" / "cofos")).expanduser()
+
+
+def package_root() -> Path:
+    return Path(__file__).resolve().parents[1]
+
+
+def find_runtime_root() -> Path:
+    bundled = package_root() / "runtime"
+    if (bundled / "src" / "inference" / "student_adapter_inference.py").exists():
+        return bundled
+
+    current = Path(__file__).resolve()
+    for parent in [current.parent, *current.parents]:
+        if (parent / "src" / "inference" / "student_adapter_inference.py").exists():
+            return parent
+    cwd = Path.cwd().resolve()
+    if (cwd / "src" / "inference" / "student_adapter_inference.py").exists():
+        return cwd
+    raise FileNotFoundError(
+        "Could not find COFOS runtime files. The npm package should include "
+        "runtime/src and runtime/config."
+    )
+
+
+def default_config_path(runtime_root: Path) -> Path:
+    configured = Path(DEFAULT_CONFIG).expanduser()
+    if configured.is_absolute():
+        return configured
+    candidate = runtime_root / configured
+    if candidate.exists():
+        return candidate
+    return (runtime_root / "config" / "config.yaml").resolve()
+
+
+def runtime_cache_root() -> Path:
+    CACHE_ROOT.mkdir(parents=True, exist_ok=True)
+    return CACHE_ROOT.resolve()
+
+
+def configure_cache_paths(config: dict, pdf_dir: str | None = None) -> dict:
+    cache = runtime_cache_root()
+    updated = dict(config)
+    paths = dict(updated.get("paths", {}))
+    raw_pdf_dir = Path(pdf_dir).expanduser().resolve() if pdf_dir else cache / "raw_pdfs"
+    paths.update({
+        "raw_pdfs": str(raw_pdf_dir),
+        "raw_json": str(cache / "raw_json"),
+        "raw_wos": str(cache / "raw_wos" / "wos_abstracts.xlsx"),
+        "processed_dir": str(cache / "data" / "processed"),
+        "sft_dir": str(cache / "data" / "sft"),
+        "kg_dir": str(cache / "data" / "kg"),
+        "eval_dir": str(cache / "data" / "eval"),
+        "retrieval_store": str(cache / "retrieval_store" / "embedding"),
+        "bm25_store": str(cache / "retrieval_store" / "bm25"),
+        "checkpoints_dir": str(cache / "checkpoints"),
+        "outputs_dir": str(cache / "outputs"),
+        "logs_dir": str(cache / "outputs" / "logs"),
+    })
+    updated["paths"] = paths
+    return updated
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(description="Run persistent COFOS 9B + RAG chat backend.")
+    parser.add_argument("--config", default=DEFAULT_CONFIG)
+    parser.add_argument("--merged-model-path", "--model", dest="model_path", default=DEFAULT_MODEL_PATH)
+    parser.add_argument("--top-k", type=int, default=int(os.environ.get("COFOS_TOP_K", "5")))
+    parser.add_argument("--max-new-tokens", type=int, default=None)
+    parser.add_argument("--pdf-dir", default=os.environ.get("COFOS_PDF_DIR"))
+    parser.add_argument("--rebuild-pdf-index", action="store_true")
+    parser.add_argument("--no-rag", action="store_true")
+    return parser.parse_args()
+
+
+_HF_REPO_RE = re.compile(r"^[A-Za-z0-9][A-Za-z0-9_\-\.]*/[A-Za-z0-9][A-Za-z0-9_\-\.]*$")
+_LOCAL_PREFIXES = {
+    "config",
+    "data",
+    "retrieval_store",
+    "checkpoints",
+    "scripts",
+    "src",
+    "fwdemo",
+    ".",
+    "..",
+    "~",
+}
+
+
+def is_hf_repo_id(value: str) -> bool:
+    if not _HF_REPO_RE.match(value):
+        return False
+    first = value.split("/", 1)[0]
+    if first in _LOCAL_PREFIXES:
+        return False
+    if Path(value).suffix:
+        return False
+    return True
+
+
+def resolve_project_path(project_root: Path, value: str) -> Path | str:
+    if is_hf_repo_id(value):
+        return value
+    path = Path(value).expanduser()
+    if path.is_absolute():
+        return path
+    return (project_root / path).resolve()
 
-MODEL_ID = "Willlzh/COFOS"
+
+
+def _project_path(project_root: Path, value: str) -> Path:
+    path = Path(value).expanduser()
+    if path.is_absolute():
+        return path
+    return (project_root / path).resolve()
+
+
+def _with_pdf_dir(config: dict, pdf_dir: str | None) -> dict:
+    if not pdf_dir:
+        return config
+    updated = dict(config)
+    paths = dict(updated.get("paths", {}))
+    paths["raw_pdfs"] = str(Path(pdf_dir).expanduser().resolve())
+    updated["paths"] = paths
+    return updated
+
+
+def _iter_jsonl(path: Path):
+    if not path.exists():
+        return
+    with path.open("r", encoding="utf-8") as handle:
+        for line in handle:
+            line = line.strip()
+            if line:
+                yield json.loads(line)
+
+
+def _doc_key(meta: dict, text: str = "") -> tuple[str, str, str]:
+    source_type = str(meta.get("source_type") or meta.get("source") or "")
+    doc_id = str(meta.get("doc_id") or meta.get("chunk_id") or "")
+    if doc_id:
+        return source_type, doc_id, ""
+    return source_type, "", text[:256]
+
+
+def _pdf_manifest_path(project_root: Path, config: dict) -> Path:
+    return _project_path(project_root, config["paths"]["processed_dir"]) / "pdf_ingest_manifest.json"
+
+
+def _current_pdf_manifest(project_root: Path, config: dict) -> dict:
+    pdf_dir = _project_path(project_root, config["paths"]["raw_pdfs"])
+    pdf_files = sorted(pdf_dir.glob("*.pdf")) if pdf_dir.exists() else []
+    files = []
+    for pdf_path in pdf_files:
+        stat = pdf_path.stat()
+        files.append({
+            "name": pdf_path.name,
+            "path": str(pdf_path.resolve()),
+            "size": stat.st_size,
+            "mtime_ns": stat.st_mtime_ns,
+        })
+    return {"pdf_dir": str(pdf_dir), "files": files}
+
+
+def _load_pdf_manifest(project_root: Path, config: dict) -> dict | None:
+    path = _pdf_manifest_path(project_root, config)
+    if not path.exists():
+        return None
+    try:
+        return json.loads(path.read_text(encoding="utf-8"))
+    except (OSError, json.JSONDecodeError):
+        return None
+
+
+def _save_pdf_manifest(project_root: Path, config: dict, manifest: dict) -> None:
+    path = _pdf_manifest_path(project_root, config)
+    path.parent.mkdir(parents=True, exist_ok=True)
+    path.write_text(json.dumps(manifest, ensure_ascii=False, indent=2) + "\n", encoding="utf-8")
+
+
+def _pdf_inputs_changed(project_root: Path, config: dict) -> bool:
+    chunks_path = _project_path(project_root, config["paths"]["processed_dir"]) / "pdf_chunks.jsonl"
+    return (not chunks_path.exists()) or _load_pdf_manifest(project_root, config) != _current_pdf_manifest(project_root, config)
+
+
+def _load_existing_bm25_rows(project_root: Path, config: dict) -> list[tuple[str, dict]]:
+    bm25_dir = _project_path(project_root, config["paths"]["bm25_store"])
+    metadata = list(_iter_jsonl(bm25_dir / "bm25_metadata.jsonl") or [])
+    documents = list(_iter_jsonl(bm25_dir / "bm25_documents.jsonl") or [])
+    rows = []
+    for index, meta in enumerate(metadata):
+        doc_text = str(documents[index].get("text", "")) if index < len(documents) else ""
+        text = str(meta.get("text") or doc_text)
+        if text.strip():
+            merged_meta = dict(meta)
+            merged_meta["text"] = text
+            rows.append((text, merged_meta))
+    return rows
+
+
+def _load_pdf_rows(project_root: Path, config: dict) -> list[tuple[str, dict]]:
+    chunks_path = _project_path(project_root, config["paths"]["processed_dir"]) / "pdf_chunks.jsonl"
+    rows = []
+    for record in _iter_jsonl(chunks_path) or []:
+        text = str(record.get("text", ""))
+        if not text.strip() or not record.get("is_ros_related", False):
+            continue
+        rows.append((text, {
+            "doc_id": record.get("chunk_id", ""),
+            "source_type": "pdf",
+            "source": record.get("source_type", "pdf_fulltext"),
+            "title": record.get("title", ""),
+            "doi": record.get("doi", ""),
+            "page": record.get("page", 0),
+            "is_ros_related": record.get("is_ros_related", False),
+            "text": text,
+        }))
+    return rows
+
+
+def parse_pdfs_if_needed(project_root: Path, config: dict, force: bool = False) -> bool:
+    from data_processing.parse_pdf_text import PDFParser
+
+    pdf_dir = _project_path(project_root, config["paths"]["raw_pdfs"])
+    pdf_files = sorted(pdf_dir.glob("*.pdf")) if pdf_dir.exists() else []
+    if not pdf_files:
+        changed = _pdf_inputs_changed(project_root, config)
+        if changed:
+            chunks_path = _project_path(project_root, config["paths"]["processed_dir"]) / "pdf_chunks.jsonl"
+            chunks_path.parent.mkdir(parents=True, exist_ok=True)
+            chunks_path.write_text("", encoding="utf-8")
+            _save_pdf_manifest(project_root, config, _current_pdf_manifest(project_root, config))
+            print(f"[chat] No PDFs found in {pdf_dir}; cleared PDF chunks.", file=sys.stderr, flush=True)
+            return True
+        return False
+    if not force and not _pdf_inputs_changed(project_root, config):
+        print("[chat] PDF inputs unchanged; skipping PDF parsing.", file=sys.stderr, flush=True)
+        return False
+    print(f"[chat] Parsing {len(pdf_files)} PDFs from {pdf_dir} ...", file=sys.stderr, flush=True)
+    PDFParser(config).run()
+    _save_pdf_manifest(project_root, config, _current_pdf_manifest(project_root, config))
+    return True
+
+
+def build_augmented_bm25_if_needed(project_root: Path, config: dict, force: bool = False) -> None:
+    from retrieval.build_bm25_index import BM25IndexBuilder
+
+    pdf_rows = _load_pdf_rows(project_root, config)
+    existing_rows = _load_existing_bm25_rows(project_root, config)
+    pdf_keys = {_doc_key(meta, text) for text, meta in pdf_rows}
+    existing_keys = {_doc_key(meta, text) for text, meta in existing_rows}
+    if pdf_rows and pdf_keys.issubset(existing_keys) and not force:
+        print("[chat] PDF chunks already present in BM25; skipping BM25 rebuild.", file=sys.stderr, flush=True)
+        return
+    if not pdf_rows and not any(str(meta.get("source_type") or meta.get("source") or "") == "pdf" for _, meta in existing_rows):
+        return
+
+    merged = {}
+    for text, meta in existing_rows:
+        if str(meta.get("source_type") or meta.get("source") or "") == "pdf":
+            continue
+        merged[_doc_key(meta, text)] = (text, meta)
+    for text, meta in pdf_rows:
+        merged[_doc_key(meta, text)] = (text, meta)
+    documents = [row[0] for row in merged.values()]
+    metadata = [row[1] for row in merged.values()]
+    print(f"[chat] Rebuilding BM25 with {len(documents)} docs ({len(pdf_rows)} PDF chunks).", file=sys.stderr, flush=True)
+    builder = BM25IndexBuilder(config)
+    builder.build_index(documents, metadata)
+    builder.save_index()
+
+
+def ensure_runtime_data(project_root: Path, config: dict) -> None:
+    """Download runtime KG + BM25 data from Hugging Face if not present locally.
+
+    ``Willlzh/COFOS_data`` hosts ``runtime/kg/`` and ``runtime/bm25/`` with the
+    pre-built retrieval indexes.  This function mirrors them into the local
+    ``data/kg`` and ``retrieval_store/bm25`` directories, skipping files that
+    already exist.
+    """
+    import shutil
+    from pathlib import Path as _Path
+
+    kg_dir = _Path(config["paths"]["kg_dir"])
+    bm25_dir = _Path(config["paths"]["bm25_store"])
+
+    kg_done = (kg_dir / "cofros_kg_triples.jsonl").exists()
+    bm25_done = (bm25_dir / "bm25_index.pkl").exists()
+    if kg_done and bm25_done:
+        return
+
+    print("[chat] Runtime data not found locally; downloading from Willlzh/COFOS_data ...",
+          file=sys.stderr, flush=True)
+
+    try:
+        from huggingface_hub import hf_hub_download
+    except ImportError:
+        print("[chat] huggingface_hub not installed. Run: pip install huggingface_hub",
+              file=sys.stderr, flush=True)
+        return
+
+    repo = "Willlzh/COFOS_data"
+    file_map = {
+        "runtime/kg/cofros_kg_triples.jsonl": kg_dir,
+        "runtime/kg/cofros_kg_nodes.csv": kg_dir,
+        "runtime/kg/cofros_kg_edges.csv": kg_dir,
+        "runtime/bm25/bm25_index.pkl": bm25_dir,
+        "runtime/bm25/bm25_metadata.jsonl": bm25_dir,
+        "runtime/bm25/bm25_documents.jsonl": bm25_dir,
+        "runtime/bm25/bm25_info.json": bm25_dir,
+    }
+
+    for remote_path, target_dir in file_map.items():
+        target_dir.mkdir(parents=True, exist_ok=True)
+        local_name = remote_path.rsplit("/", 1)[-1]
+        local_path = target_dir / local_name
+        if local_path.exists():
+            continue
+        cached = hf_hub_download(
+            repo_id=repo,
+            repo_type="dataset",
+            filename=remote_path,
+        )
+        shutil.copy2(cached, local_path)
+        print(f"[chat] Downloaded {remote_path} -> {local_path}", file=sys.stderr, flush=True)
+
+    print("[chat] Runtime data ready.", file=sys.stderr, flush=True)
+
+
+def _load_history() -> list:
+    """Load recent conversation turns from the persistent history file."""
+    if not os.path.exists(HISTORY_FILE):
+        return []
+    turns = []
+    try:
+        with open(HISTORY_FILE, "r", encoding="utf-8") as fh:
+            for line in fh:
+                line = line.strip()
+                if not line:
+                    continue
+                try:
+                    turns.append(json.loads(line))
+                except json.JSONDecodeError:
+                    continue
+    except OSError:
+        return []
+    # Keep only the last N turns
+    max_items = HISTORY_MAX_TURNS * 2
+    return turns[-max_items:] if len(turns) > max_items else turns
+
+
+def _save_turn(role: str, content: str) -> None:
+    """Append a single turn to the persistent history file."""
+    try:
+        with open(HISTORY_FILE, "a", encoding="utf-8") as fh:
+            json.dump({"role": role, "content": content}, fh, ensure_ascii=False)
+            fh.write("\n")
+    except OSError:
+        pass  # best-effort
 
 
 def main() -> None:
-    model_id = sys.argv[1] if len(sys.argv) > 1 else MODEL_ID
+    args = parse_args()
+    runtime_root = find_runtime_root()
+    os.chdir(runtime_root)
+    src_path = runtime_root / "src"
+    scripts_path = runtime_root / "scripts"
+    sys.path.insert(0, str(src_path))
+    if scripts_path.exists():
+        sys.path.insert(0, str(scripts_path))
+
+    import yaml
+    from inference.student_adapter_inference import StudentAdapterInference
+
+    logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s")
+
+    config_path = default_config_path(runtime_root) if args.config == DEFAULT_CONFIG else resolve_project_path(runtime_root, args.config)
+    model_path = resolve_project_path(runtime_root, args.model_path)
 
-    print(f"[chat] Loading {model_id} …", file=sys.stderr, flush=True)
-    print(f"[chat] First run may take a while while the COFOS model is downloaded or read from cache.", file=sys.stderr, flush=True)
+    print(f"[chat] Runtime root: {runtime_root}", file=sys.stderr, flush=True)
+    print(f"[chat] Cache root: {runtime_cache_root()}", file=sys.stderr, flush=True)
+    print(f"[chat] Config: {config_path}", file=sys.stderr, flush=True)
+    print(f"[chat] Model: {model_path}", file=sys.stderr, flush=True)
+    print(f"[chat] RAG: {'disabled' if args.no_rag else 'enabled'}", file=sys.stderr, flush=True)
 
-    tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True)
-    if tokenizer.pad_token_id is None:
-        tokenizer.pad_token_id = tokenizer.eos_token_id
+    t0 = time.time()
+    with open(config_path, "r", encoding="utf-8") as handle:
+        config = yaml.safe_load(handle) or {}
+    config = configure_cache_paths(config, args.pdf_dir)
 
-    model = AutoModelForCausalLM.from_pretrained(
-        model_id, dtype="auto", device_map="auto", trust_remote_code=True, low_cpu_mem_usage=True,
+    if not args.no_rag:
+        ensure_runtime_data(runtime_root, config)
+        pdf_changed = parse_pdfs_if_needed(runtime_root, config, force=args.rebuild_pdf_index)
+        build_augmented_bm25_if_needed(runtime_root, config, force=args.rebuild_pdf_index or pdf_changed)
+
+    print("[chat] Loading model (this may take 1-2 min on first run) ...", file=sys.stderr, flush=True)
+
+    runner = StudentAdapterInference(
+        config=config,
+        model_path=str(model_path),
+        adapter_path=None,
+        use_rag=not args.no_rag,
+        max_new_tokens=args.max_new_tokens,
     )
-    model.eval()
 
-    print("[chat] Ready.", file=sys.stderr, flush=True)
+    runner.load_model()
+
+    elapsed = time.time() - t0
+    print(f"[chat] Ready. (loaded in {elapsed:.1f}s)", file=sys.stderr, flush=True)
     print("===READY===", flush=True)
 
-    messages: list[dict] = []
+    top_k = args.top_k
+    history = _load_history()
+    if history:
+        print(f"[chat] Restored {len(history) // 2} previous turns from history.",
+              file=sys.stderr, flush=True)
 
     for raw in sys.stdin:
-        line = raw.strip("\r\n")
-        if not line or line == "/exit":
+        question = raw.strip("\r\n")
+        if not question or question == "/exit":
             break
-        if line.startswith("/model "):
-            mid = line[len("/model "):].strip()
-            tokenizer = AutoTokenizer.from_pretrained(mid, trust_remote_code=True)
-            if tokenizer.pad_token_id is None:
-                tokenizer.pad_token_id = tokenizer.eos_token_id
-            model = AutoModelForCausalLM.from_pretrained(
-                mid, dtype="auto", device_map="auto", trust_remote_code=True, low_cpu_mem_usage=True,
-            )
-            model.eval()
+        if question.startswith("/topk "):
+            try:
+                top_k = max(1, int(question.split(maxsplit=1)[1]))
+                print(f"top_k set to {top_k}")
+            except ValueError:
+                print("Usage: /topk 5")
+            print("===DONE===", flush=True)
+            continue
+        if question == "/rag off":
+            runner.rag = None
+            print("RAG disabled for this session.")
+            print("===DONE===", flush=True)
+            continue
+        if question == "/rag on":
+            if runner.rag is None:
+                from inference.run_kg_grounded_inference import KGRAGInference
+                runner.rag = KGRAGInference(config)
+            print("RAG enabled for this session.")
+            print("===DONE===", flush=True)
+            continue
+        if question == "/clear":
+            history.clear()
+            print("Conversation history cleared.")
+            print("===DONE===", flush=True)
+            continue
+        if question == "/save":
+            print(f"History saved to {HISTORY_FILE} ({len(history) // 2} turns).")
             print("===DONE===", flush=True)
             continue
 
-        messages.append({"role": "user", "content": line})
-        prompt = tokenizer.apply_chat_template(
-            messages, tokenize=False, add_generation_prompt=True,
-        )
-        inputs = tokenizer(prompt, return_tensors="pt")
-
-        streamer = TextIteratorStreamer(
-            tokenizer, skip_prompt=True, skip_special_tokens=True,
-        )
-
-        gen_kw = dict(
-            input_ids=inputs.input_ids,
-            attention_mask=inputs.attention_mask,
-            streamer=streamer,
-            max_new_tokens=512,
-            do_sample=True,
-            temperature=0.7,
-            top_p=0.9,
-            repetition_penalty=1.05,
-            pad_token_id=tokenizer.eos_token_id,
-        )
-
-        Thread(target=model.generate, kwargs=gen_kw).start()
+        _save_turn("user", question)
 
-        reply = ""
-        for tok in streamer:
-            sys.stdout.write(tok)
-            sys.stdout.flush()
-            reply += tok
+        try:
+            t_gen = time.time()
+            print("===RESPONSE===", flush=True)
+            answer_parts = []
+            for token in runner.answer_stream(question, top_k=top_k, history=history):
+                print(token, end="", flush=True)
+                answer_parts.append(token)
+            answer = "".join(answer_parts)
+            gen_elapsed = time.time() - t_gen
+            print(f"\n===DONE===", flush=True)
+            print(f"[chat] Generated in {gen_elapsed:.1f}s", file=sys.stderr, flush=True)
 
-        messages.append({"role": "assistant", "content": reply})
-        print("\n===DONE===", flush=True)
+            history.append({"role": "user", "content": question})
+            history.append({"role": "assistant", "content": answer})
+            _save_turn("assistant", answer)
+        except Exception as exc:
+            print(f"\n[COFOS error] {exc}")
+            print("===DONE===", flush=True)
 
 
 if __name__ == "__main__":
@@ -87,6 +505,6 @@ if __name__ == "__main__":
         main()
     except KeyboardInterrupt:
         pass
-    except Exception as e:
-        print(f"[chat] Error: {e}", file=sys.stderr)
+    except Exception as exc:
+        print(f"[chat] Error: {exc}", file=sys.stderr, flush=True)
         sys.exit(1)