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.
- {confluence_knowledge-0.4.0 → confluence_knowledge-0.5.0}/PKG-INFO +7 -5
- {confluence_knowledge-0.4.0 → confluence_knowledge-0.5.0}/README.md +3 -3
- {confluence_knowledge-0.4.0 → confluence_knowledge-0.5.0}/pyproject.toml +5 -2
- {confluence_knowledge-0.4.0 → confluence_knowledge-0.5.0}/src/confluence_knowledge/cli.py +30 -0
- {confluence_knowledge-0.4.0 → confluence_knowledge-0.5.0}/src/confluence_knowledge/crawl.py +12 -5
- {confluence_knowledge-0.4.0 → confluence_knowledge-0.5.0}/src/confluence_knowledge/index.py +35 -1
- confluence_knowledge-0.5.0/src/confluence_knowledge/models.lock.json +30 -0
- confluence_knowledge-0.5.0/src/confluence_knowledge/retrieval.py +347 -0
- {confluence_knowledge-0.4.0 → confluence_knowledge-0.5.0}/src/confluence_knowledge/server.py +36 -8
- confluence_knowledge-0.5.0/src/confluence_knowledge/update.py +55 -0
- {confluence_knowledge-0.4.0 → confluence_knowledge-0.5.0}/src/confluence_knowledge.egg-info/PKG-INFO +7 -5
- {confluence_knowledge-0.4.0 → confluence_knowledge-0.5.0}/src/confluence_knowledge.egg-info/SOURCES.txt +5 -1
- {confluence_knowledge-0.4.0 → confluence_knowledge-0.5.0}/src/confluence_knowledge.egg-info/requires.txt +3 -1
- confluence_knowledge-0.5.0/tests/test_prepare_models.py +114 -0
- {confluence_knowledge-0.4.0 → confluence_knowledge-0.5.0}/tests/test_server_flow.py +1 -0
- confluence_knowledge-0.5.0/tests/test_sparse_retrieval.py +78 -0
- confluence_knowledge-0.5.0/tests/test_stale_model.py +58 -0
- confluence_knowledge-0.5.0/tests/test_update.py +40 -0
- confluence_knowledge-0.4.0/src/confluence_knowledge/models.lock.json +0 -15
- confluence_knowledge-0.4.0/src/confluence_knowledge/retrieval.py +0 -206
- confluence_knowledge-0.4.0/tests/test_prepare_models.py +0 -122
- {confluence_knowledge-0.4.0 → confluence_knowledge-0.5.0}/setup.cfg +0 -0
- {confluence_knowledge-0.4.0 → confluence_knowledge-0.5.0}/src/confluence_knowledge/__init__.py +0 -0
- {confluence_knowledge-0.4.0 → confluence_knowledge-0.5.0}/src/confluence_knowledge/auth.py +0 -0
- {confluence_knowledge-0.4.0 → confluence_knowledge-0.5.0}/src/confluence_knowledge/client.py +0 -0
- {confluence_knowledge-0.4.0 → confluence_knowledge-0.5.0}/src/confluence_knowledge/extract.py +0 -0
- {confluence_knowledge-0.4.0 → confluence_knowledge-0.5.0}/src/confluence_knowledge/extract_files.py +0 -0
- {confluence_knowledge-0.4.0 → confluence_knowledge-0.5.0}/src/confluence_knowledge/registration.py +0 -0
- {confluence_knowledge-0.4.0 → confluence_knowledge-0.5.0}/src/confluence_knowledge/settings.py +0 -0
- {confluence_knowledge-0.4.0 → confluence_knowledge-0.5.0}/src/confluence_knowledge.egg-info/dependency_links.txt +0 -0
- {confluence_knowledge-0.4.0 → confluence_knowledge-0.5.0}/src/confluence_knowledge.egg-info/entry_points.txt +0 -0
- {confluence_knowledge-0.4.0 → confluence_knowledge-0.5.0}/src/confluence_knowledge.egg-info/top_level.txt +0 -0
- {confluence_knowledge-0.4.0 → confluence_knowledge-0.5.0}/tests/test_agent_doc_consistency.py +0 -0
- {confluence_knowledge-0.4.0 → confluence_knowledge-0.5.0}/tests/test_client_get_only.py +0 -0
- {confluence_knowledge-0.4.0 → confluence_knowledge-0.5.0}/tests/test_crawl.py +0 -0
- {confluence_knowledge-0.4.0 → confluence_knowledge-0.5.0}/tests/test_extract.py +0 -0
- {confluence_knowledge-0.4.0 → confluence_knowledge-0.5.0}/tests/test_index.py +0 -0
- {confluence_knowledge-0.4.0 → confluence_knowledge-0.5.0}/tests/test_registration.py +0 -0
- {confluence_knowledge-0.4.0 → confluence_knowledge-0.5.0}/tests/test_retrieval.py +0 -0
- {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.
|
|
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:
|
|
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 (`
|
|
42
|
-
(`
|
|
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 (`
|
|
9
|
-
(`
|
|
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.
|
|
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
|
-
|
|
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
|
-
|
|
214
|
-
|
|
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
|
|
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}
|