confluence-knowledge 0.4.0__tar.gz → 0.5.0__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (40) hide show
  1. {confluence_knowledge-0.4.0 → confluence_knowledge-0.5.0}/PKG-INFO +7 -5
  2. {confluence_knowledge-0.4.0 → confluence_knowledge-0.5.0}/README.md +3 -3
  3. {confluence_knowledge-0.4.0 → confluence_knowledge-0.5.0}/pyproject.toml +5 -2
  4. {confluence_knowledge-0.4.0 → confluence_knowledge-0.5.0}/src/confluence_knowledge/cli.py +30 -0
  5. {confluence_knowledge-0.4.0 → confluence_knowledge-0.5.0}/src/confluence_knowledge/crawl.py +12 -5
  6. {confluence_knowledge-0.4.0 → confluence_knowledge-0.5.0}/src/confluence_knowledge/index.py +35 -1
  7. confluence_knowledge-0.5.0/src/confluence_knowledge/models.lock.json +30 -0
  8. confluence_knowledge-0.5.0/src/confluence_knowledge/retrieval.py +347 -0
  9. {confluence_knowledge-0.4.0 → confluence_knowledge-0.5.0}/src/confluence_knowledge/server.py +36 -8
  10. confluence_knowledge-0.5.0/src/confluence_knowledge/update.py +55 -0
  11. {confluence_knowledge-0.4.0 → confluence_knowledge-0.5.0}/src/confluence_knowledge.egg-info/PKG-INFO +7 -5
  12. {confluence_knowledge-0.4.0 → confluence_knowledge-0.5.0}/src/confluence_knowledge.egg-info/SOURCES.txt +5 -1
  13. {confluence_knowledge-0.4.0 → confluence_knowledge-0.5.0}/src/confluence_knowledge.egg-info/requires.txt +3 -1
  14. confluence_knowledge-0.5.0/tests/test_prepare_models.py +114 -0
  15. {confluence_knowledge-0.4.0 → confluence_knowledge-0.5.0}/tests/test_server_flow.py +1 -0
  16. confluence_knowledge-0.5.0/tests/test_sparse_retrieval.py +78 -0
  17. confluence_knowledge-0.5.0/tests/test_stale_model.py +58 -0
  18. confluence_knowledge-0.5.0/tests/test_update.py +40 -0
  19. confluence_knowledge-0.4.0/src/confluence_knowledge/models.lock.json +0 -15
  20. confluence_knowledge-0.4.0/src/confluence_knowledge/retrieval.py +0 -206
  21. confluence_knowledge-0.4.0/tests/test_prepare_models.py +0 -122
  22. {confluence_knowledge-0.4.0 → confluence_knowledge-0.5.0}/setup.cfg +0 -0
  23. {confluence_knowledge-0.4.0 → confluence_knowledge-0.5.0}/src/confluence_knowledge/__init__.py +0 -0
  24. {confluence_knowledge-0.4.0 → confluence_knowledge-0.5.0}/src/confluence_knowledge/auth.py +0 -0
  25. {confluence_knowledge-0.4.0 → confluence_knowledge-0.5.0}/src/confluence_knowledge/client.py +0 -0
  26. {confluence_knowledge-0.4.0 → confluence_knowledge-0.5.0}/src/confluence_knowledge/extract.py +0 -0
  27. {confluence_knowledge-0.4.0 → confluence_knowledge-0.5.0}/src/confluence_knowledge/extract_files.py +0 -0
  28. {confluence_knowledge-0.4.0 → confluence_knowledge-0.5.0}/src/confluence_knowledge/registration.py +0 -0
  29. {confluence_knowledge-0.4.0 → confluence_knowledge-0.5.0}/src/confluence_knowledge/settings.py +0 -0
  30. {confluence_knowledge-0.4.0 → confluence_knowledge-0.5.0}/src/confluence_knowledge.egg-info/dependency_links.txt +0 -0
  31. {confluence_knowledge-0.4.0 → confluence_knowledge-0.5.0}/src/confluence_knowledge.egg-info/entry_points.txt +0 -0
  32. {confluence_knowledge-0.4.0 → confluence_knowledge-0.5.0}/src/confluence_knowledge.egg-info/top_level.txt +0 -0
  33. {confluence_knowledge-0.4.0 → confluence_knowledge-0.5.0}/tests/test_agent_doc_consistency.py +0 -0
  34. {confluence_knowledge-0.4.0 → confluence_knowledge-0.5.0}/tests/test_client_get_only.py +0 -0
  35. {confluence_knowledge-0.4.0 → confluence_knowledge-0.5.0}/tests/test_crawl.py +0 -0
  36. {confluence_knowledge-0.4.0 → confluence_knowledge-0.5.0}/tests/test_extract.py +0 -0
  37. {confluence_knowledge-0.4.0 → confluence_knowledge-0.5.0}/tests/test_index.py +0 -0
  38. {confluence_knowledge-0.4.0 → confluence_knowledge-0.5.0}/tests/test_registration.py +0 -0
  39. {confluence_knowledge-0.4.0 → confluence_knowledge-0.5.0}/tests/test_retrieval.py +0 -0
  40. {confluence_knowledge-0.4.0 → confluence_knowledge-0.5.0}/tests/test_secrets.py +0 -0
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: confluence-knowledge
3
- Version: 0.4.0
3
+ Version: 0.5.0
4
4
  Summary: Local, GET-only Confluence subtree indexer + grounded-Q&A MCP service
5
5
  License: MIT
6
6
  Project-URL: Repository, https://github.com/nkunal1992/agents
@@ -19,7 +19,9 @@ Requires-Dist: keyring>=25.2
19
19
  Requires-Dist: mcp<2,>=1.6
20
20
  Requires-Dist: pip-system-certs>=5.0
21
21
  Provides-Extra: models
22
- Requires-Dist: sentence-transformers<4.0,>=3.0; extra == "models"
22
+ Requires-Dist: onnxruntime>=1.17; extra == "models"
23
+ Requires-Dist: tokenizers>=0.15; extra == "models"
24
+ Requires-Dist: fastembed>=0.9; extra == "models"
23
25
  Requires-Dist: huggingface-hub>=0.26; extra == "models"
24
26
  Provides-Extra: files
25
27
  Requires-Dist: python-docx>=1.1; extra == "files"
@@ -38,9 +40,9 @@ See `../../docs/confluence-knowledge.md` for the full guide.
38
40
 
39
41
  - GET-only `ConfluenceClient` — non-GET raises `SecurityException` pre-flight.
40
42
  - Recursive subtree crawl (Cloud v2 + Data Center), resumable.
41
- - Local embeddings (`all-MiniLM-L6-v2`) + local reranker
42
- (`ms-marco-MiniLM-L6-v2`) — Apache-2.0, pinned in `models.lock.json`,
43
- never a cloud embedding call.
43
+ - Local ONNX embeddings (`bge-m3`, quantized) + sparse leg (`Qdrant/bm25`)
44
+ + reranker (`bge-reranker-v2-m3`, quantized) — MIT/Apache-2.0, pinned in
45
+ `models.lock.json`, no torch, never a cloud embedding call.
44
46
  - SQLite + FTS5 + vector-blob index under `.conf_knowledge_index/`.
45
47
  - Credentials via OS keyring / env var / chmod-600 `.conf_auth.json`.
46
48
 
@@ -5,9 +5,9 @@ See `../../docs/confluence-knowledge.md` for the full guide.
5
5
 
6
6
  - GET-only `ConfluenceClient` — non-GET raises `SecurityException` pre-flight.
7
7
  - Recursive subtree crawl (Cloud v2 + Data Center), resumable.
8
- - Local embeddings (`all-MiniLM-L6-v2`) + local reranker
9
- (`ms-marco-MiniLM-L6-v2`) — Apache-2.0, pinned in `models.lock.json`,
10
- never a cloud embedding call.
8
+ - Local ONNX embeddings (`bge-m3`, quantized) + sparse leg (`Qdrant/bm25`)
9
+ + reranker (`bge-reranker-v2-m3`, quantized) — MIT/Apache-2.0, pinned in
10
+ `models.lock.json`, no torch, never a cloud embedding call.
11
11
  - SQLite + FTS5 + vector-blob index under `.conf_knowledge_index/`.
12
12
  - Credentials via OS keyring / env var / chmod-600 `.conf_auth.json`.
13
13
 
@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
4
4
 
5
5
  [project]
6
6
  name = "confluence-knowledge"
7
- version = "0.4.0"
7
+ version = "0.5.0"
8
8
  description = "Local, GET-only Confluence subtree indexer + grounded-Q&A MCP service"
9
9
  readme = "README.md"
10
10
  requires-python = ">=3.11"
@@ -31,7 +31,10 @@ Issues = "https://github.com/nkunal1992/agents/issues"
31
31
 
32
32
  [project.optional-dependencies]
33
33
  models = [
34
- "sentence-transformers>=3.0,<4.0",
34
+ # ONNX stack — no torch: ~40 MB runtime vs 2.4 GB, native ARM64 Windows
35
+ "onnxruntime>=1.17",
36
+ "tokenizers>=0.15",
37
+ "fastembed>=0.9", # sparse leg (Qdrant/bm25 ONNX)
35
38
  "huggingface-hub>=0.26",
36
39
  ]
37
40
  files = [
@@ -94,9 +94,21 @@ def _wizard() -> int:
94
94
 
95
95
  def _doctor() -> int:
96
96
  from .server import status
97
+ from .update import check
97
98
  try:
98
99
  s = status()
99
100
  print(json.dumps(s, indent=2))
101
+ upd = check("confluence-knowledge")
102
+ if upd["installed"]:
103
+ line = f"version: {upd['installed']}"
104
+ if upd["latest"]:
105
+ line += f" (latest: {upd['latest']})"
106
+ elif upd["latest"] is None:
107
+ line += " (latest: unchecked — offline or internal mirror)"
108
+ print(line)
109
+ if upd["update_available"]:
110
+ print("update available: run `confluence-knowledge update` "
111
+ "(or pip install -U confluence-knowledge[models])")
100
112
  if s.get("configured"):
101
113
  from .server import _client_for # internal but fine in CLI
102
114
  profile = SiteProfile.load(default_home())
@@ -122,6 +134,7 @@ def main() -> None:
122
134
  sub.add_parser("setup", help="guided credential + config setup")
123
135
  sub.add_parser("doctor", help="verify config, credentials, models")
124
136
  sub.add_parser("prepare-models", help="download the pinned local models")
137
+ sub.add_parser("update", help="upgrade this package via pip -U")
125
138
  reg = sub.add_parser(
126
139
  "register",
127
140
  help="write this server into the IDE's MCP config (no manual paste)")
@@ -145,6 +158,23 @@ def main() -> None:
145
158
  for line in prepare_models(default_home()):
146
159
  print(line)
147
160
  return
161
+ if args.cmd == "update":
162
+ import subprocess
163
+ cmd = [sys.executable, "-m", "pip", "install", "-U",
164
+ "confluence-knowledge[models]"]
165
+ if os.name == "nt":
166
+ # this exe is the running process on Windows — pip can't replace
167
+ # a file that's in use; hand the user the equivalent python -m
168
+ # pip command instead (same mechanism `pip install -U pip` uses)
169
+ print("on Windows, upgrade via the venv python directly:\n "
170
+ + " ".join(f'"{c}"' if " " in c else c for c in cmd))
171
+ return
172
+ r = subprocess.run(cmd, check=False)
173
+ if r.returncode == 0:
174
+ print("\nupdated — run `confluence-knowledge prepare-models` "
175
+ "(new releases may pin new model artifacts) then "
176
+ "`confluence-knowledge doctor`")
177
+ sys.exit(r.returncode)
148
178
  if args.cmd == "register":
149
179
  from .registration import register
150
180
  try:
@@ -111,13 +111,17 @@ def _fetch_page_dc(client: ConfluenceClient, page_id: str) -> dict:
111
111
  }
112
112
 
113
113
 
114
+ from .retrieval import EMBED_MODEL_ID as _MODEL_ID
115
+
116
+
114
117
  class Crawler:
115
118
  def __init__(self, client: ConfluenceClient, index: Index, env: str,
116
- embedder=None, site_url: str | None = None):
119
+ embedder=None, sparse=None, site_url: str | None = None):
117
120
  self.client = client
118
121
  self.index = index
119
122
  self.env = env
120
123
  self.embedder = embedder # retrieval.LocalEmbedder or None
124
+ self.sparse = sparse # retrieval.LocalSparseEmbedder or None
121
125
  # fallback base for page links when a response carries no _links.base
122
126
  if site_url and env == "cloud" and not site_url.rstrip("/").endswith("/wiki"):
123
127
  site_url = site_url.rstrip("/") + "/wiki"
@@ -208,15 +212,17 @@ class Crawler:
208
212
  chunks, warns = extract_chunks(page["storage"])
209
213
  for c in chunks:
210
214
  c.chunk_id = f"{pid}:{content_hash[:8]}:{c.ordinal}"
211
- vecs = None
215
+ vecs, sparse_vecs = None, None
212
216
  if self.embedder is not None and chunks:
213
- vecs = self.embedder.embed(
214
- [f"{c.heading_path}\n{c.text}" for c in chunks])
217
+ texts = [f"{c.heading_path}\n{c.text}" for c in chunks]
218
+ vecs = self.embedder.embed(texts)
219
+ if self.sparse is not None:
220
+ sparse_vecs = self.sparse.embed(texts)
215
221
  self.index.upsert_page({
216
222
  "page_id": pid, "parent_id": None, "title": page["title"],
217
223
  "url": page["url"], "level": levels.get(pid, 0),
218
224
  "version": page["version"], "content_hash": content_hash,
219
- "warnings": warns}, chunks, vecs)
225
+ "warnings": warns}, chunks, vecs, sparse_vecs)
220
226
  seen.add(pid)
221
227
  for cid, ctype in self._children(pid, is_folder=False):
222
228
  if cid not in seen:
@@ -230,6 +236,7 @@ class Crawler:
230
236
  state = "running"
231
237
  else:
232
238
  self.index.publish(seen)
239
+ self.index.set_model_id(_MODEL_ID)
233
240
  self.index.save_job(job_id, "done", [], errors)
234
241
  state = "done"
235
242
  return {"job_id": job_id, "state": state, "remaining": len(queue),
@@ -27,6 +27,7 @@ CREATE TABLE IF NOT EXISTS chunks (
27
27
  chunk_id TEXT PRIMARY KEY, page_id TEXT, heading TEXT, kind TEXT,
28
28
  ordinal INTEGER, text TEXT, active INTEGER DEFAULT 0);
29
29
  CREATE TABLE IF NOT EXISTS vectors (chunk_id TEXT PRIMARY KEY, vec BLOB);
30
+ CREATE TABLE IF NOT EXISTS sparse_vectors (chunk_id TEXT PRIMARY KEY, payload BLOB);
30
31
  CREATE VIRTUAL TABLE IF NOT EXISTS chunks_fts USING fts5(
31
32
  chunk_id UNINDEXED, text);
32
33
  CREATE TABLE IF NOT EXISTS jobs (
@@ -110,7 +111,8 @@ class Index:
110
111
  "errors": job["errors"]}
111
112
 
112
113
  # ---- content ----
113
- def upsert_page(self, page: dict, chunks: list, vecs: np.ndarray | None) -> None:
114
+ def upsert_page(self, page: dict, chunks: list, vecs: np.ndarray | None,
115
+ sparse: list[dict[int, float]] | None = None) -> None:
114
116
  with self.lock, self.db:
115
117
  self.db.execute(
116
118
  "INSERT OR REPLACE INTO pages VALUES (?,?,?,?,?,?,?,?,?,0)",
@@ -123,6 +125,7 @@ class Index:
123
125
  for cid in old:
124
126
  self.db.execute("DELETE FROM chunks_fts WHERE chunk_id=?", (cid,))
125
127
  self.db.execute("DELETE FROM vectors WHERE chunk_id=?", (cid,))
128
+ self.db.execute("DELETE FROM sparse_vectors WHERE chunk_id=?", (cid,))
126
129
  self.db.execute("DELETE FROM chunks WHERE page_id=?", (page["page_id"],))
127
130
  for i, c in enumerate(chunks):
128
131
  self.db.execute(
@@ -134,6 +137,37 @@ class Index:
134
137
  self.db.execute(
135
138
  "INSERT INTO vectors VALUES (?,?)",
136
139
  (c.chunk_id, vecs[i].astype(np.float32).tobytes()))
140
+ if sparse is not None and i < len(sparse) and sparse[i]:
141
+ idx_arr = np.fromiter(sparse[i].keys(), dtype=np.int32)
142
+ val_arr = np.fromiter(sparse[i].values(), dtype=np.float32)
143
+ self.db.execute(
144
+ "INSERT INTO sparse_vectors VALUES (?,?)",
145
+ (c.chunk_id,
146
+ idx_arr.tobytes() + val_arr.tobytes()))
147
+
148
+ def sparse_for(self, chunk_ids: list[str]) -> dict[str, dict[int, float]]:
149
+ if not chunk_ids:
150
+ return {}
151
+ q = ",".join("?" * len(chunk_ids))
152
+ out = {}
153
+ with self.lock:
154
+ for cid, blob in self.db.execute(
155
+ f"SELECT chunk_id,payload FROM sparse_vectors "
156
+ f"WHERE chunk_id IN ({q})", chunk_ids):
157
+ # layout: int32[n] indices + float32[n] values → n = len/8
158
+ k = len(blob) // 8
159
+ idx = np.frombuffer(blob[:4 * k], dtype=np.int32)
160
+ vals = np.frombuffer(blob[4 * k:], dtype=np.float32)
161
+ out[cid] = dict(zip(idx.tolist(), vals.tolist()))
162
+ return out
163
+
164
+ def model_id(self) -> str | None:
165
+ return self.meta_get("embed_model")
166
+
167
+ def set_model_id(self, model_id: str) -> None:
168
+ with self.lock, self.db:
169
+ self.db.execute("INSERT OR REPLACE INTO meta VALUES ('embed_model', ?)",
170
+ (model_id,))
137
171
 
138
172
  def publish(self, seen_page_ids: set[str]) -> None:
139
173
  """Atomically make exactly the crawled page set the active snapshot."""
@@ -0,0 +1,30 @@
1
+ {
2
+ "comment": "Pinned ONNX model artifacts. Retrieval is local-only — torch and cloud embedding APIs are never used. 'files' limits the snapshot to what the runtime needs.",
3
+ "models": {
4
+ "embed": {
5
+ "repo": "Xenova/bge-m3",
6
+ "revision": "4de13258303883538bd53b696b452bf8099f0858",
7
+ "license": "MIT",
8
+ "files": ["onnx/sentence_transformers_quantized.onnx", "tokenizer.json", "tokenizer_config.json", "config.json", "sentencepiece.bpe.model", "special_tokens_map.json", "README.md"]
9
+ },
10
+ "sparse": {
11
+ "repo": "Qdrant/bm25",
12
+ "revision": "22b8d2af71a76161e18dd432d2cee0eefa66e412",
13
+ "license": "Apache-2.0",
14
+ "note": "downloaded through fastembed's own cache (models/fastembed)"
15
+ },
16
+ "rerank": {
17
+ "repo": "onnx-community/bge-reranker-v2-m3-ONNX",
18
+ "revision": "6f5ff65298512715a1e669753bc754d2bc8f367b",
19
+ "license": "MIT",
20
+ "files": ["onnx/model_quantized.onnx", "tokenizer.json", "tokenizer_config.json", "config.json", "special_tokens_map.json", "quantize_config.json", "README.md"]
21
+ },
22
+ "rerank_light": {
23
+ "repo": "onnx-community/bge-reranker-base-ONNX",
24
+ "revision": "f131ae72114cb18333648a2ee6b0942849e8ab95",
25
+ "license": "MIT",
26
+ "files": ["onnx/model_quantized.onnx", "tokenizer.json", "tokenizer_config.json", "config.json", "special_tokens_map.json", "quantize_config.json", "README.md"],
27
+ "note": "downloaded only when *_RERANKER env selects it"
28
+ }
29
+ }
30
+ }
@@ -0,0 +1,347 @@
1
+ """Local hybrid retrieval: dense + sparse + FTS → RRF → ONNX cross-encoder.
2
+
3
+ Everything runs on this machine via ONNX Runtime — no torch, no cloud.
4
+ Models load local-only from pinned artifacts; a missing model is a setup
5
+ error, never a silent fallback.
6
+
7
+ Backend layout (all under <home>/models/):
8
+ embed — Xenova/bge-m3 onnx/sentence_transformers_quantized.onnx
9
+ (dense 1024-dim, normalized, multilingual, 8k ctx)
10
+ sparse — Qdrant/bm25 via fastembed SparseTextEmbedding
11
+ (learned term-weight leg; BGE-M3's lexical head has no ONNX
12
+ export, so a purpose-built sparse model fills that slot)
13
+ rerank — onnx-community/bge-reranker-v2-m3-ONNX onnx/model_quantized.onnx
14
+ (sequence-classification logits → relevance score)
15
+ """
16
+
17
+ from __future__ import annotations
18
+
19
+ import os
20
+ import shutil
21
+ from pathlib import Path
22
+
23
+ import numpy as np
24
+
25
+ from .index import Index
26
+ from .settings import SetupError
27
+
28
+ EMBED_REPO = "Xenova/bge-m3" # Apache-2.0/MIT
29
+ EMBED_ONNX = "onnx/sentence_transformers_quantized.onnx"
30
+ SPARSE_MODEL = "Qdrant/bm25" # MIT — via fastembed
31
+ RERANK_REPO = "onnx-community/bge-reranker-v2-m3-ONNX" # MIT
32
+ RERANK_ONNX = "onnx/model_quantized.onnx"
33
+ RERANK_BASE_REPO = "onnx-community/bge-reranker-base-ONNX" # light override
34
+
35
+ EMBED_MODEL_ID = "bge-m3-onnx-quantized"
36
+
37
+ TOP_K_RETRIEVE = 30
38
+ TOP_K_EVIDENCE = 8
39
+ RERANK_MAX = 25 # bounds cross-encoder latency on CPU
40
+ EMBED_MAX_TOKENS = 512 # quantized model input cap (chunks stay under)
41
+ RERANK_MAX_TOKENS = 512
42
+
43
+
44
+ def _models_dir(home: Path) -> Path:
45
+ return home / "models"
46
+
47
+
48
+ def _model_home(home: Path, repo: str) -> Path:
49
+ return _models_dir(home) / repo.split("/")[-1]
50
+
51
+
52
+ def _require(dest: Path, repo: str, cli: str, needed: list[str]) -> Path:
53
+ if not all((dest / f).exists() for f in needed):
54
+ raise SetupError(
55
+ f"local model '{repo}' not prepared — run `{cli} prepare-models`. "
56
+ "Cloud embeddings are never used.")
57
+ return dest
58
+
59
+
60
+ def _ort_session(onnx_path: Path, cli: str):
61
+ try:
62
+ import onnxruntime as ort
63
+ except ImportError as e:
64
+ raise SetupError("onnxruntime isn't installed — reinstall the "
65
+ f"service ({cli} bootstrap reinstalls deps)") from e
66
+ return ort.InferenceSession(
67
+ str(onnx_path), providers=["CPUExecutionProvider"],
68
+ sess_options=_ort_opts(ort))
69
+
70
+
71
+ def _ort_opts(ort):
72
+ so = ort.SessionOptions()
73
+ so.intra_op_num_threads = max(1, (os.cpu_count() or 4) // 2)
74
+ so.inter_op_num_threads = 1
75
+ so.log_severity_level = 3 # errors only
76
+ return so
77
+
78
+
79
+ def _tokenizer(dest: Path):
80
+ from tokenizers import Tokenizer
81
+ return Tokenizer.from_file(str(dest / "tokenizer.json"))
82
+
83
+
84
+ class LocalEmbedder:
85
+ """BGE-M3 dense embeddings via ONNX (no torch). Same interface."""
86
+
87
+ def __init__(self, home: Path, cli: str = "confluence-knowledge"):
88
+ dest = _require(_model_home(home, EMBED_REPO), EMBED_REPO, cli,
89
+ [EMBED_ONNX, "tokenizer.json"])
90
+ self._sess = _ort_session(dest / EMBED_ONNX, cli)
91
+ self._tok = _tokenizer(dest)
92
+ self._tok.enable_padding()
93
+ self._tok.enable_truncation(EMBED_MAX_TOKENS)
94
+
95
+ def embed(self, texts: list[str]) -> np.ndarray:
96
+ enc = self._tok.encode_batch(texts, add_special_tokens=True)
97
+ ids = np.asarray([e.ids for e in enc], dtype=np.int64)
98
+ mask = np.asarray([e.attention_mask for e in enc], dtype=np.int64)
99
+ out = self._sess.run(
100
+ ["sentence_embedding"],
101
+ {"input_ids": ids, "attention_mask": mask})[0]
102
+ v = np.asarray(out, dtype=np.float32)
103
+ n = np.linalg.norm(v, axis=1, keepdims=True)
104
+ return v / (n + 1e-9) # defensive normalize — quantized graphs drift
105
+
106
+
107
+ class LocalSparseEmbedder:
108
+ """Learned term-weight sparse vectors via fastembed (ONNX, no torch).
109
+
110
+ Fills the role BGE-M3's lexical head plays on torch — no ONNX export of
111
+ that head exists, so a purpose-built sparse model covers the same leg.
112
+ """
113
+
114
+ def __init__(self, home: Path, cli: str = "confluence-knowledge"):
115
+ try:
116
+ from fastembed import SparseTextEmbedding
117
+ except ImportError as e:
118
+ raise SetupError("fastembed isn't installed — reinstall the "
119
+ f"service ({cli} bootstrap reinstalls deps)") from e
120
+ os.environ.setdefault("HF_HUB_OFFLINE", "1")
121
+ os.environ.setdefault("TRANSFORMERS_OFFLINE", "1")
122
+ dest = _models_dir(home) / "fastembed"
123
+ try:
124
+ self._m = SparseTextEmbedding(SPARSE_MODEL, cache_dir=str(dest),
125
+ local_files_only=True)
126
+ except Exception as e:
127
+ raise SetupError(
128
+ f"sparse model '{SPARSE_MODEL}' not prepared — run "
129
+ f"`{cli} prepare-models`.") from e
130
+
131
+ def embed(self, texts: list[str]) -> list[dict[int, float]]:
132
+ """One {token_idx: weight} dict per input — the sparse payload."""
133
+ out = []
134
+ for emb in self._m.embed(texts):
135
+ out.append(dict(zip(emb.indices.tolist(),
136
+ emb.values.tolist())))
137
+ return out
138
+
139
+
140
+ class LocalReranker:
141
+ """bge-reranker-v2-m3 via ONNX sequence-classification logits.
142
+
143
+ Override with env `*_EMBED_RERANKER` = "onnx-community/bge-reranker-base-ONNX"
144
+ (lighter, English) for slow machines — doctor surfaces the hint.
145
+ """
146
+
147
+ def __init__(self, home: Path, cli: str = "confluence-knowledge"):
148
+ # per-app override → "<APP>_RERANKER" (e.g. SP_KNOWLEDGE_RERANKER)
149
+ # plus a global KNOWLEDGE_RERANKER fallback
150
+ env_key = cli.upper().replace("-", "_") + "_RERANKER"
151
+ repo = os.environ.get(env_key, os.environ.get("KNOWLEDGE_RERANKER",
152
+ RERANK_REPO))
153
+ dest = _require(_model_home(home, repo), repo, cli,
154
+ [RERANK_ONNX, "tokenizer.json"])
155
+ self._sess = _ort_session(dest / RERANK_ONNX, cli)
156
+ self._tok = _tokenizer(dest)
157
+ self._tok.no_padding()
158
+ self._tok.enable_truncation(RERANK_MAX_TOKENS)
159
+ self.repo = repo
160
+
161
+ def score(self, query: str, docs: list[str]) -> np.ndarray:
162
+ pairs = [[query, d] for d in docs]
163
+ enc = self._tok.encode_batch(pairs, add_special_tokens=True)
164
+ maxlen = max(len(e.ids) for e in enc)
165
+ ids = np.zeros((len(enc), maxlen), dtype=np.int64)
166
+ mask = np.zeros((len(enc), maxlen), dtype=np.int64)
167
+ for i, e in enumerate(enc):
168
+ ids[i, :len(e.ids)] = e.ids
169
+ mask[i, :len(e.attention_mask)] = e.attention_mask
170
+ logits = self._sess.run(
171
+ ["logits"], {"input_ids": ids, "attention_mask": mask})[0]
172
+ s = logits.reshape(-1).astype(np.float32)
173
+ return 1.0 / (1.0 + np.exp(-s)) # sigmoid → 0..1 relevance
174
+
175
+
176
+ def prepare_models(home: Path) -> list[str]:
177
+ """Download pinned ONNX artifacts once (public; no document text sent)."""
178
+ from huggingface_hub import snapshot_download
179
+
180
+ out = []
181
+ lock = _load_lock()
182
+ models = _models_dir(home)
183
+ models.mkdir(parents=True, exist_ok=True)
184
+
185
+ # dense + rerankers: direct HF snapshot (repo + revision pinned in the
186
+ # lock); rerank_light is fetched too so *_RERANKER overrides always work
187
+ for key in ("embed", "rerank", "rerank_light"):
188
+ spec = lock["models"][key]
189
+ dest = _models_dir(home) / spec["repo"].split("/")[-1]
190
+ if dest.exists():
191
+ complete = all((dest / f).exists() for f in spec["files"])
192
+ if complete:
193
+ out.append(f"already prepared: {spec['repo']}")
194
+ continue
195
+ shutil.rmtree(dest) # partial/interrupted → refetch
196
+ snapshot_download(repo_id=spec["repo"], revision=spec["revision"],
197
+ local_dir=dest, allow_patterns=spec["files"])
198
+ if not _license_ok(dest):
199
+ _license_ok_upstream(dest, spec, snapshot_download)
200
+ if not _license_ok(dest):
201
+ raise SetupError(
202
+ f"{spec['repo']} license artifact missing — refusing to ship")
203
+ out.append(f"prepared {spec['repo']}@{spec['revision'][:8]} "
204
+ f"({spec['license']})")
205
+
206
+ # sparse: fastembed owns its cache layout — let it download into our dir
207
+ os.environ.pop("HF_HUB_OFFLINE", None)
208
+ spec = lock["models"]["sparse"]
209
+ try:
210
+ from fastembed import SparseTextEmbedding
211
+ SparseTextEmbedding(spec["repo"], cache_dir=str(models / "fastembed"))
212
+ out.append(f"prepared {spec['repo']} ({spec['license']})")
213
+ except Exception as e:
214
+ raise SetupError(f"couldn't fetch sparse model {spec['repo']}: {e}") from e
215
+ return out
216
+
217
+
218
+ def _license_ok_upstream(dest: Path, spec: dict, snapshot_download) -> None:
219
+ """ONNX-community repackages sometimes drop the license line — if the card
220
+ points at a `base_model`, fetch that model's card and judge IT. The
221
+ upstream card is saved next to ours for audit."""
222
+ import re
223
+ for md in dest.glob("*.md"):
224
+ m = re.search(r"base_model:\s*(?:-\s*)?([\w.-]+/[\w.-]+)",
225
+ md.read_text(errors="ignore")[:6000])
226
+ if not m:
227
+ continue
228
+ try:
229
+ path = snapshot_download(repo_id=m.group(1),
230
+ allow_patterns=["README.md", "LICENSE*"],
231
+ local_dir=dest / "_upstream")
232
+ except Exception: # noqa: BLE001 — upstream card optional; license judge stays on local files
233
+ return
234
+ up = Path(path)
235
+ for f in list(up.glob("README.md")) + list(up.glob("LICENSE*")):
236
+ f.rename(dest / f"_upstream_{f.name}")
237
+ shutil.rmtree(up, ignore_errors=True)
238
+ return
239
+
240
+
241
+ def _license_ok(dest: Path) -> bool:
242
+ for md in list(dest.glob("*.md")) + list(dest.glob("LICENSE*")):
243
+ head = md.read_text(errors="ignore")[:6000].lower()
244
+ if ("apache" in head or "mit license" in head or "license: mit" in head
245
+ or "cc-by" in head):
246
+ return True
247
+ return any(dest.glob("LICENSE*"))
248
+
249
+
250
+ def _load_lock() -> dict:
251
+ import json
252
+
253
+ p = Path(__file__).parent / "models.lock.json" # shipped inside the package
254
+ if p.exists():
255
+ return json.loads(p.read_text())
256
+ raise SetupError("models.lock.json missing from package — reinstall the service")
257
+
258
+
259
+ def _cosine(q: np.ndarray, mat: dict[str, np.ndarray]) -> list[tuple[str, float]]:
260
+ if not mat:
261
+ return []
262
+ ids = list(mat)
263
+ m = np.stack([mat[i] for i in ids])
264
+ m = m / (np.linalg.norm(m, axis=1, keepdims=True) + 1e-9)
265
+ q = q / (np.linalg.norm(q) + 1e-9)
266
+ return sorted(zip(ids, (m @ q).tolist()), key=lambda t: -t[1])
267
+
268
+
269
+ def _sparse_rank(q: dict[int, float],
270
+ docs: dict[str, dict[int, float]]) -> list[tuple[str, float]]:
271
+ """Score = Σ wq[i]·wd[i] over shared token indices — plain dot product."""
272
+ if not docs or not q:
273
+ return []
274
+ scores = [(cid, sum(q[i] * w for i, w in d.items() if i in q))
275
+ for cid, d in docs.items()]
276
+ scores.sort(key=lambda t: -t[1])
277
+ return scores
278
+
279
+
280
+ def retrieve(index: Index, query: str, embedder: LocalEmbedder,
281
+ reranker: LocalReranker | None = None,
282
+ sparse: LocalSparseEmbedder | None = None) -> list[dict]:
283
+ """Dense + sparse + FTS → RRF fuse → rerank → bounded evidence."""
284
+ chunks = index.active_chunks()
285
+ if not chunks:
286
+ return []
287
+
288
+ # dense leg
289
+ qvec = embedder.embed([query])[0]
290
+ vecs = index.vectors_for([c["chunk_id"] for c in chunks])
291
+ dense_rank = [cid for cid, _ in _cosine(qvec, vecs)][:TOP_K_RETRIEVE]
292
+
293
+ # keyword leg
294
+ fts_rank = [cid for cid, _ in index.fts_search(query)][:TOP_K_RETRIEVE]
295
+
296
+ # sparse leg (optional — learned term weights)
297
+ sparse_rank: list[str] = []
298
+ if sparse is not None:
299
+ q_sparse = sparse.embed([query])[0]
300
+ doc_sparse = index.sparse_for([c["chunk_id"] for c in chunks])
301
+ sparse_rank = [cid for cid, _ in
302
+ _sparse_rank(q_sparse, doc_sparse)][:TOP_K_RETRIEVE]
303
+
304
+ # reciprocal-rank fusion over 2-3 legs
305
+ rrf: dict[str, float] = {}
306
+ for leg in (dense_rank, fts_rank, sparse_rank):
307
+ for rank, cid in enumerate(leg):
308
+ rrf[cid] = rrf.get(cid, 0.0) + 1.0 / (60 + rank + 1)
309
+ fused = sorted(rrf, key=lambda c: -rrf[c])[:TOP_K_RETRIEVE]
310
+ by_id = {c["chunk_id"]: c for c in chunks}
311
+ candidates = [by_id[c] for c in fused if c in by_id]
312
+
313
+ if reranker is not None and candidates:
314
+ scored = candidates[:RERANK_MAX]
315
+ scores = reranker.score(query, [c["text"] for c in scored])
316
+ order = np.argsort(-scores)[:TOP_K_EVIDENCE]
317
+ evidence = []
318
+ for i in order:
319
+ c = dict(scored[i])
320
+ c["rerank_score"] = float(scores[i])
321
+ evidence.append(c)
322
+ else:
323
+ evidence = [dict(c, rerank_score=None) for c in candidates[:TOP_K_EVIDENCE]]
324
+
325
+ pages = index.active_pages()
326
+ for e in evidence:
327
+ p = pages.get(e["page_id"], {})
328
+ e["page_title"], e["page_url"], e["level"] = (
329
+ p.get("title", ""), p.get("url", ""), p.get("level"))
330
+ return evidence
331
+
332
+
333
+ def check_citations(evidence: list[dict], mappings: list[dict]) -> dict:
334
+ """Verify claim→chunk citations: chunk exists in this evidence set and the
335
+ claimed quote is literally present in it. Provenance, not entailment."""
336
+ by_id = {e["chunk_id"]: e for e in evidence}
337
+ bad = []
338
+ for m in mappings:
339
+ ch = by_id.get(m.get("chunk_id", ""))
340
+ if ch is None:
341
+ bad.append({"claim": m.get("claim", "")[:80], "reason": "chunk not in evidence set"})
342
+ continue
343
+ quote = m.get("quote", "")
344
+ if quote and quote not in ch["text"]:
345
+ bad.append({"claim": m.get("claim", "")[:80],
346
+ "reason": "quote not present verbatim in cited chunk"})
347
+ return {"ok": not bad, "violations": bad}