@lzhzzzzwill/cofos 1.0.4 → 1.1.0

package/runtime/src/retrieval/retrieve_evidence.py

@@ -8,7 +8,6 @@ COF-ROS-Reasoner: 检索证据脚本
 - 返回证据列表
 """
 
-import os
 import json
 import yaml
 import logging
@@ -16,11 +15,8 @@ import pickle
 import re
 from pathlib import Path
 from typing import Dict, List, Any, Optional, Tuple
-from datetime import datetime
 from collections import defaultdict
 
-import numpy as np
-
 # 添加 src 路径
 
 try:
@@ -38,13 +34,16 @@ try:
 except ImportError:
     BM25Okapi = None
 
-from data_processing.utils import setup_logging, ensure_directory
+from data_processing.utils import setup_logging
 from retrieval.text_search import tokenize_search_text
 
 
 class EvidenceRetriever:
     """检索证据"""
 
+    # Upper bound on the per-query ranked-result cache (see ``_ranked_results``).
+    _RANKED_CACHE_MAX = 64
+
     def __init__(self, config: Dict[str, Any]):
         self.config = config
         self.faiss_dir = Path(config['paths']['retrieval_store'])
@@ -128,11 +127,39 @@ class EvidenceRetriever:
             logging.info(f"Loaded persisted BM25 index from {index_path}")
 
 
-    def _load_kg_triples(self) -> None:
-        """Load real KG triples and build lightweight in-memory indexes."""
+    def _reset_kg_indexes(self) -> None:
         self.kg_triples = []
         self.kg_by_head = defaultdict(list)
         self.kg_by_tail = defaultdict(list)
+        self._kg_nodes_by_type = None
+
+    @property
+    def kg_nodes_by_type(self) -> Dict[str, set]:
+        """Map node id prefix ("Material:", "ROS_or_Product:") -> node ids.
+
+        Built lazily from ``kg_triples`` and cached, so node lookup no longer
+        rescans every triple once per prefix on every query.  Derived rather
+        than populated during load so that callers which assemble a retriever
+        by hand (tests) get a consistent index too.
+        """
+        cached = getattr(self, '_kg_nodes_by_type', None)
+        if cached is not None:
+            return cached
+
+        nodes_by_type: Dict[str, set] = defaultdict(set)
+        for triple in getattr(self, 'kg_triples', []):
+            for field in ('head', 'tail'):
+                node = triple.get(field, '')
+                node_type, sep, _ = node.partition(':')
+                if sep and node_type:
+                    nodes_by_type[f"{node_type}:"].add(node)
+
+        self._kg_nodes_by_type = nodes_by_type
+        return nodes_by_type
+
+    def _load_kg_triples(self) -> None:
+        """Load real KG triples and build lightweight in-memory indexes."""
+        self._reset_kg_indexes()
 
         kg_path = Path(self.config['paths']['kg_dir']) / 'cofros_kg_triples.jsonl'
         if not kg_path.exists():
@@ -152,12 +179,11 @@ class EvidenceRetriever:
                     self.kg_triples.append(triple)
                     self.kg_by_head[head].append(triple)
                     self.kg_by_tail[tail].append(triple)
+            self._kg_nodes_by_type = None  # rebuilt on first lookup
             logging.info("Loaded %s KG triples for graph retrieval", len(self.kg_triples))
         except Exception as exc:
             logging.error("Failed to load KG triples from %s: %s", kg_path, exc)
-            self.kg_triples = []
-            self.kg_by_head = defaultdict(list)
-            self.kg_by_tail = defaultdict(list)
+            self._reset_kg_indexes()
 
     @staticmethod
     def _kg_label(node_id: str) -> str:
@@ -229,12 +255,7 @@ class EvidenceRetriever:
         return [node for _, node in self._find_scored_kg_nodes(query, node_prefix)]
 
     def _find_scored_kg_nodes(self, query: str, node_prefix: str) -> List[Tuple[float, str]]:
-        nodes = set()
-        for triple in self.kg_triples:
-            for field in ('head', 'tail'):
-                node = triple.get(field, '')
-                if node.startswith(node_prefix):
-                    nodes.add(node)
+        nodes = self.kg_nodes_by_type.get(node_prefix, set())
 
         scored = []
         for node in nodes:
@@ -551,11 +572,27 @@ class EvidenceRetriever:
         return self.is_confident_result(result)
 
     def hybrid_retrieve(self, query: str, top_k: int = 10) -> List[Dict[str, Any]]:
-        """混合检索"""
-        if not self.is_domain_query(query):
-            return []
+        """混合检索（按 query 缓存，避免同一问题重复打分）。
+
+        A single question triggers ``retrieve_kg_facts`` *and*
+        ``retrieve_with_sources``; both funnel into this method, so without the
+        cache every question scored the whole BM25 corpus twice.  Scoring and
+        reranking depend only on the query, so we cache the ranked list and just
+        slice it per ``top_k``.
+        """
+        ranked = self._ranked_results(query)
+        return ranked[:top_k]
 
-        results = []
+    def _ranked_results(self, query: str) -> List[Dict[str, Any]]:
+        cache = getattr(self, '_ranked_cache', None)
+        if cache is None:
+            cache = self._ranked_cache = {}
+        if query in cache:
+            return cache[query]
+
+        if not self.is_domain_query(query):
+            cache[query] = []
+            return cache[query]
 
         # 1. 向量检索
         vector_results = self._vector_search(query, self.vector_top_k)
@@ -571,7 +608,31 @@ class EvidenceRetriever:
 
         # 5. 丢弃低置信度结果，避免无关问题也被强制注入 RAG 上下文
         confident = [result for result in reranked if self.is_confident_result(result)]
-        return confident[:top_k]
+
+        # Bound the cache so a long chat session cannot grow it without limit.
+        if len(cache) >= self._RANKED_CACHE_MAX:
+            cache.pop(next(iter(cache)))
+        cache[query] = confident
+        return confident
+
+    def _get_embedding_model(self):
+        """Load the sentence encoder once and reuse it across queries.
+
+        This used to be constructed inside ``_vector_search``, so every single
+        question paid a full SentenceTransformer load.
+        """
+        if hasattr(self, '_embedding_model_cache'):
+            return self._embedding_model_cache
+
+        self._embedding_model_cache = None
+        embedding_model = self.config.get('models', {}).get('embedding_model', '')
+        if embedding_model and SentenceTransformer is not None:
+            try:
+                self._embedding_model_cache = SentenceTransformer(embedding_model)
+                logging.info("Loaded embedding model %s", embedding_model)
+            except Exception as exc:
+                logging.error("Failed to load embedding model %s: %s", embedding_model, exc)
+        return self._embedding_model_cache
 
     def _vector_search(self, query: str, top_k: int) -> List[Tuple[int, float]]:
         """向量搜索"""
@@ -579,12 +640,10 @@ class EvidenceRetriever:
             return []
 
         try:
-            # 嵌入查询
-            embedding_model = self.config.get('models', {}).get('embedding_model', '')
-            if not embedding_model:
+            model = self._get_embedding_model()
+            if model is None:
                 return []
 
-            model = SentenceTransformer(embedding_model)
             query_vec = model.encode([query], normalize_embeddings=True).astype('float32')
 
             # 搜索

package/scripts/chat.py

@@ -18,6 +18,13 @@ from pathlib import Path
 
 import re
 
+# Prefer the regular Hugging Face download path over Xet for better CLI
+# reliability on networks where xet-bridge CDN requests time out.
+os.environ.setdefault("HF_HUB_DISABLE_XET", "1")
+os.environ.setdefault("HF_HUB_ENABLE_HF_TRANSFER", "0")
+os.environ.setdefault("HF_HUB_DOWNLOAD_TIMEOUT", "120")
+os.environ.setdefault("HF_HUB_ETAG_TIMEOUT", "30")
+
 HISTORY_FILE = os.path.join(os.path.expanduser("~"), ".cofos_history.jsonl")
 HISTORY_MAX_TURNS = 20
 
@@ -26,7 +33,7 @@ DEFAULT_MODEL_PATH = os.environ.get(
     "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()
+CACHE_ROOT = Path(os.environ.get("COFOS_CACHE_DIR", Path.cwd() / ".cofos")).expanduser()
 
 
 def package_root() -> Path:
@@ -37,16 +44,8 @@ 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 "
+        "Could not find bundled COFOS runtime files. The npm package should include "
         "runtime/src and runtime/config."
     )
 
@@ -357,7 +356,12 @@ def ensure_runtime_data(project_root: Path, config: dict) -> None:
 
 
 def _load_history() -> list:
-    """Load recent conversation turns from the persistent history file."""
+    """Load recent conversation turns from the persistent history file.
+
+    The chat template requires strictly alternating user/assistant turns, so a
+    truncated or interleaved file must be repaired here rather than blowing up
+    inside ``apply_chat_template``.
+    """
     if not os.path.exists(HISTORY_FILE):
         return []
     turns = []
@@ -368,14 +372,33 @@ def _load_history() -> list:
                 if not line:
                     continue
                 try:
-                    turns.append(json.loads(line))
+                    turn = json.loads(line)
                 except json.JSONDecodeError:
                     continue
+                if (
+                    isinstance(turn, dict)
+                    and turn.get("role") in {"user", "assistant"}
+                    and isinstance(turn.get("content"), str)
+                ):
+                    turns.append(turn)
     except OSError:
         return []
-    # Keep only the last N turns
+
+    # Drop anything that breaks user→assistant alternation (e.g. a user turn
+    # whose answer never completed because generation was interrupted).
+    paired = []
+    for turn in turns:
+        expected = "user" if len(paired) % 2 == 0 else "assistant"
+        if turn["role"] == expected:
+            paired.append(turn)
+        elif turn["role"] == "user" and expected == "assistant":
+            # Previous user turn never got an answer; replace it.
+            paired[-1] = turn
+    if len(paired) % 2:
+        paired.pop()
+
     max_items = HISTORY_MAX_TURNS * 2
-    return turns[-max_items:] if len(turns) > max_items else turns
+    return paired[-max_items:] if len(paired) > max_items else paired
 
 
 def _save_turn(role: str, content: str) -> None:
@@ -388,8 +411,29 @@ def _save_turn(role: str, content: str) -> None:
         pass  # best-effort
 
 
+def _rewrite_history(turns: list) -> bool:
+    """Replace the on-disk history with *turns*. Returns True on success."""
+    try:
+        with open(HISTORY_FILE, "w", encoding="utf-8") as fh:
+            for turn in turns:
+                json.dump(turn, fh, ensure_ascii=False)
+                fh.write("\n")
+        return True
+    except OSError:
+        return False
+
+
 def main() -> None:
     args = parse_args()
+
+    # Resolve user-supplied paths against the shell's cwd *before* chdir below,
+    # otherwise a relative --pdf-dir would be looked up inside the installed
+    # npm package directory instead of where the user actually keeps the PDFs.
+    if args.pdf_dir:
+        args.pdf_dir = str(Path(args.pdf_dir).expanduser().resolve())
+    if args.config and args.config != DEFAULT_CONFIG and not is_hf_repo_id(args.config):
+        args.config = str(Path(args.config).expanduser().resolve())
+
     runtime_root = find_runtime_root()
     os.chdir(runtime_root)
     src_path = runtime_root / "src"
@@ -399,7 +443,7 @@ def main() -> None:
         sys.path.insert(0, str(scripts_path))
 
     import yaml
-    from inference.student_adapter_inference import StudentAdapterInference
+    from inference.student_adapter_inference import StudentAdapterInference, StudentRAGContext
 
     logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s")
 
@@ -448,38 +492,53 @@ def main() -> None:
         question = raw.strip("\r\n")
         if not question or question == "/exit":
             break
-        if question.startswith("/topk "):
+        if question == "/topk" or question.startswith("/topk "):
+            parts = question.split(maxsplit=1)
             try:
-                top_k = max(1, int(question.split(maxsplit=1)[1]))
+                top_k = max(1, int(parts[1]))
                 print(f"top_k set to {top_k}")
-            except ValueError:
+            except (IndexError, ValueError):
                 print("Usage: /topk 5")
             print("===DONE===", flush=True)
             continue
         if question == "/rag off":
             runner.rag = None
+            runner.use_rag = False
             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.")
+                # Started with --no-rag, so the indexes may never have been
+                # fetched. Do it now, and keep a failure from killing the
+                # backend the way an uncaught exception here used to.
+                try:
+                    ensure_runtime_data(runtime_root, config)
+                    runner.rag = StudentRAGContext(config)
+                    runner.use_rag = True
+                    print("RAG enabled for this session.")
+                except Exception as exc:
+                    print(f"Could not enable RAG: {exc}")
+            else:
+                print("RAG is already enabled.")
             print("===DONE===", flush=True)
             continue
         if question == "/clear":
             history.clear()
+            # Also clear the file; otherwise the next launch silently restores
+            # the turns the user just asked to forget.
+            _rewrite_history([])
             print("Conversation history cleared.")
             print("===DONE===", flush=True)
             continue
         if question == "/save":
-            print(f"History saved to {HISTORY_FILE} ({len(history) // 2} turns).")
+            if _rewrite_history(history):
+                print(f"History saved to {HISTORY_FILE} ({len(history) // 2} turns).")
+            else:
+                print(f"Could not write history to {HISTORY_FILE}.")
             print("===DONE===", flush=True)
             continue
 
-        _save_turn("user", question)
-
         try:
             t_gen = time.time()
             print("===RESPONSE===", flush=True)
@@ -489,13 +548,17 @@ def main() -> None:
                 answer_parts.append(token)
             answer = "".join(answer_parts)
             gen_elapsed = time.time() - t_gen
-            print(f"\n===DONE===", flush=True)
+            print("\n===DONE===", flush=True)
             print(f"[chat] Generated in {gen_elapsed:.1f}s", file=sys.stderr, flush=True)
 
+            # Persist both turns only once the answer completed, so an
+            # interrupted generation cannot leave a dangling user turn behind.
             history.append({"role": "user", "content": question})
             history.append({"role": "assistant", "content": answer})
+            _save_turn("user", question)
             _save_turn("assistant", answer)
         except Exception as exc:
+            logging.exception("Generation failed")
             print(f"\n[COFOS error] {exc}")
             print("===DONE===", flush=True)
 

package/scripts/sync_runtime.py

@@ -0,0 +1,45 @@
+#!/usr/bin/env python3
+"""Package-local runtime sync wrapper.
+
+When run from the source checkout, delegate to the project-level generator.
+When run from an installed/published package, validate that the generated
+runtime is already bundled and exit without needing repository parent folders.
+"""
+
+from __future__ import annotations
+
+import runpy
+import sys
+from pathlib import Path
+
+
+def main() -> int:
+    package_root = Path(__file__).resolve().parents[1]
+    project_sync = package_root.parent / "scripts" / "sync_fwdemo_runtime.py"
+    if project_sync.exists():
+        sys.argv = [str(project_sync), *sys.argv[1:]]
+        runpy.run_path(str(project_sync), run_name="__main__")
+        return 0
+
+    runtime_root = package_root / "runtime"
+    required = [
+        runtime_root / "config" / "config.yaml",
+        runtime_root / "config" / "prompt_templates.yaml",
+        runtime_root / "src" / "inference" / "student_adapter_inference.py",
+        runtime_root / "src" / "retrieval" / "retrieve_evidence.py",
+    ]
+    missing = [path.relative_to(package_root) for path in required if not path.exists()]
+    if missing:
+        print(
+            "Packaged COFOS runtime is incomplete; missing: "
+            + ", ".join(str(path) for path in missing),
+            file=sys.stderr,
+        )
+        return 1
+
+    print("Packaged COFOS runtime is already bundled; no source checkout sync needed.")
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())

package/runtime/config/qwen3_5_config.md

@@ -1,46 +0,0 @@
-# COFOS fwdemo runtime model configuration
-
-This file describes the bundled runtime used by the standalone `@lzhzzzzwill/cofos` npm package.
-
-## Standalone QA path
-
-The standalone CLI does use Hugging Face by default:
-
-- Student QA model: `Willlzh/COFOS`
-- Runtime KG/BM25 data: `Willlzh/COFOS_data` under `runtime/`
-- Local cache/write location: `~/.cache/cofos/`
-
-On startup, fwdemo checks local cache first. It downloads only missing model/data files.
-
-```bash
-cofos
-cofos --pdf-dir ./new_pdfs
-cofos --model /path/to/local/merged_model
-```
-
-PDFs passed with `--pdf-dir` are parsed into local chunks and merged into BM25 retrieval data under `~/.cache/cofos/`. They are not uploaded anywhere.
-
-## Legacy full-project Ollama settings
-
-The `llm:` block in `config.yaml` is kept for full COFOS repository workflows such as Silver JSON extraction, teacher distillation, and teacher KG-RAG experiments. Those workflows can call an internal Ollama-style endpoint, for example:
-
-```yaml
-llm:
-  provider: ollama
-  generate_url: http://10.161.138.150:11434/api/generate
-  teacher_model: qwen3.6:35b
-  extraction_model: qwen3.6:35b
-  inference_model: qwen3.6:35b
-```
-
-That legacy block is not the default standalone fwdemo QA model path. The standalone QA model is loaded through Transformers from `Willlzh/COFOS` unless the user passes a local `--model` path.
-
-## Retrieval
-
-BM25 is enabled by default. FAISS remains disabled unless a local embedding model is configured explicitly.
-
-```yaml
-retrieval:
-  enable_faiss: false
-  enable_bm25: true
-```