handrolled 0.1.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 (31) hide show
  1. handrolled-0.1.0/LICENSE +21 -0
  2. handrolled-0.1.0/PKG-INFO +99 -0
  3. handrolled-0.1.0/README.md +74 -0
  4. handrolled-0.1.0/pyproject.toml +42 -0
  5. handrolled-0.1.0/setup.cfg +4 -0
  6. handrolled-0.1.0/src/handrolled/__init__.py +26 -0
  7. handrolled-0.1.0/src/handrolled/deeprank.py +183 -0
  8. handrolled-0.1.0/src/handrolled/embed_onnx.py +168 -0
  9. handrolled-0.1.0/src/handrolled/importers/__init__.py +68 -0
  10. handrolled-0.1.0/src/handrolled/importers/basic_memory.py +67 -0
  11. handrolled-0.1.0/src/handrolled/importers/jsonl.py +35 -0
  12. handrolled-0.1.0/src/handrolled/importers/markdown.py +137 -0
  13. handrolled-0.1.0/src/handrolled/importers/mem0.py +41 -0
  14. handrolled-0.1.0/src/handrolled/importers/notion.py +28 -0
  15. handrolled-0.1.0/src/handrolled/importers/obsidian.py +24 -0
  16. handrolled-0.1.0/src/handrolled/importers/records.py +94 -0
  17. handrolled-0.1.0/src/handrolled/mounts.py +189 -0
  18. handrolled-0.1.0/src/handrolled/pages.py +274 -0
  19. handrolled-0.1.0/src/handrolled/py.typed +0 -0
  20. handrolled-0.1.0/src/handrolled/rerank.py +197 -0
  21. handrolled-0.1.0/src/handrolled/text.py +267 -0
  22. handrolled-0.1.0/src/handrolled/webdav.py +232 -0
  23. handrolled-0.1.0/src/handrolled.egg-info/PKG-INFO +99 -0
  24. handrolled-0.1.0/src/handrolled.egg-info/SOURCES.txt +29 -0
  25. handrolled-0.1.0/src/handrolled.egg-info/dependency_links.txt +1 -0
  26. handrolled-0.1.0/src/handrolled.egg-info/top_level.txt +1 -0
  27. handrolled-0.1.0/tests/test_import.py +246 -0
  28. handrolled-0.1.0/tests/test_importers.py +170 -0
  29. handrolled-0.1.0/tests/test_mounts.py +95 -0
  30. handrolled-0.1.0/tests/test_rerank_idle.py +61 -0
  31. handrolled-0.1.0/tests/test_text.py +151 -0
@@ -0,0 +1,21 @@
1
+ MIT License
2
+
3
+ Copyright (c) 2026 dsh-kb contributors
4
+
5
+ Permission is hereby granted, free of charge, to any person obtaining a copy
6
+ of this software and associated documentation files (the "Software"), to deal
7
+ in the Software without restriction, including without limitation the rights
8
+ to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
9
+ copies of the Software, and to permit persons to whom the Software is
10
+ furnished to do so, subject to the following conditions:
11
+
12
+ The above copyright notice and this permission notice shall be included in all
13
+ copies or substantial portions of the Software.
14
+
15
+ THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
16
+ IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
17
+ FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
18
+ AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
19
+ LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
20
+ OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
21
+ SOFTWARE.
@@ -0,0 +1,99 @@
1
+ Metadata-Version: 2.4
2
+ Name: handrolled
3
+ Version: 0.1.0
4
+ Summary: Hand-rolled, stdlib-only building blocks for agent tooling: CJK full-text search, WebDAV, network mounts, and import adapters.
5
+ Author: handrolled contributors
6
+ License-Expression: MIT
7
+ Project-URL: Homepage, https://github.com/LUgABeEEE/handrolled
8
+ Project-URL: Source, https://github.com/LUgABeEEE/handrolled
9
+ Project-URL: Issues, https://github.com/LUgABeEEE/handrolled/issues
10
+ Keywords: stdlib,zero-dependency,sqlite-fts5,cjk,webdav,markdown,agent
11
+ Classifier: Development Status :: 3 - Alpha
12
+ Classifier: Intended Audience :: Developers
13
+ Classifier: Operating System :: OS Independent
14
+ Classifier: Programming Language :: Python :: 3
15
+ Classifier: Programming Language :: Python :: 3.10
16
+ Classifier: Programming Language :: Python :: 3.11
17
+ Classifier: Programming Language :: Python :: 3.12
18
+ Classifier: Programming Language :: Python :: 3.13
19
+ Classifier: Topic :: Text Processing :: Indexing
20
+ Classifier: Topic :: Utilities
21
+ Requires-Python: >=3.10
22
+ Description-Content-Type: text/markdown
23
+ License-File: LICENSE
24
+ Dynamic: license-file
25
+
26
+ # handrolled
27
+
28
+ **Hand-rolled, stdlib-only replacements for things you would normally `pip install`.**
29
+
30
+ Every module here exists for the same reason: the obvious solution adds a dependency,
31
+ and the dependency is bigger than the problem.
32
+
33
+ ```bash
34
+ pip install handrolled # installs nothing else
35
+ ```
36
+
37
+ Python ≥ 3.10, **no runtime dependencies**. `import handrolled` never touches
38
+ third-party code — the optional pieces import inside functions and return `None`
39
+ when their dependency is missing.
40
+
41
+ ## What's in it
42
+
43
+ | Module | Instead of | What it actually does |
44
+ |---|---|---|
45
+ | `handrolled.text` | `jieba`, a query builder, `PyYAML` | CJK bigram expansion for SQLite FTS5, FTS5 match-expression building (quotes and neutralises user input), and a YAML-**subset** front matter reader |
46
+ | `handrolled.webdav` | `requests` + `webdavclient3` | WebDAV over the RFC 4918 subset that gets used: `PROPFIND`, `GET`, `PUT`, `MKCOL`, `DELETE`; Basic auth, mTLS, and connect-host/SNI separation for appliances whose certificate CN is not their address |
47
+ | `handrolled.mounts` | parsing `mount` output | Network mount inspection that reports the **real** filesystem rather than its `autofs` trigger, `rw`/`ro` state, and `ESTALE`/`EACCES`/`ENOENT` turned into sentences a human can act on |
48
+ | `handrolled.pages` + `handrolled.importers` | hand-written one-off scripts | Turn another tool's export into Markdown pages: `mem0`, Obsidian, Notion's Markdown export, basic-memory, plain JSON/JSONL, and a neutral `records/v1` format. Deterministic, additive, guarded |
49
+ | `handrolled.rerank` | `sentence-transformers` | Cross-encoder reranking of `(query, passage)` pairs from an int8 ONNX model, with idle unloading so the memory comes back |
50
+ | `handrolled.deeprank` | an SDK per vendor | Let a remote chat model pick which candidate pages answer a query; any failure degrades to the original order |
51
+ | `handrolled.embed_onnx` | `sentence-transformers` | ONNX embedding backend (last-token pooling, L2 normalised), kept for A/B work |
52
+
53
+ ## Design rules
54
+
55
+ 1. **Standard library only.** If a module cannot be written with the stdlib and a
56
+ documented protocol, it does not belong here.
57
+ 2. **Optional means optional.** Anything needing a model or an extra package
58
+ imports lazily and returns `None`/`False` when it is absent. Nothing raises at
59
+ import time, and nothing silently changes behaviour because a package happens
60
+ to be installed.
61
+ 3. **No machine assumptions.** Optional layers have no built-in paths: unset
62
+ environment variable means the layer is off. Nothing guesses where your models
63
+ or credentials live.
64
+ 4. **Deterministic output.** `handrolled.pages` renders the same records to the
65
+ same bytes, so a second run is a no-op and a diff means something.
66
+
67
+ ## Quick example
68
+
69
+ ```python
70
+ from handrolled import importers
71
+ from handrolled.pages import plan, apply
72
+
73
+ records = importers.load("obsidian", "/path/to/vault")
74
+ planned = plan(policy, records, prefix="notes", zone="private")
75
+ print(f"{planned.items} records → {len(planned.pages)} pages")
76
+ apply(source, planned, write=True)
77
+ ```
78
+
79
+ `policy` is any object with `zones` (things with a `.name`) and
80
+ `is_never(rel) -> bool` — the caller owns the privacy model, not this library.
81
+
82
+ ## Where this comes from
83
+
84
+ Extracted from [dsh-kb](https://github.com/LUgABeEEE/dsh-kb), a personal
85
+ knowledge base built to be queried cheaply by AI agents. Those modules were
86
+ written there, hardened against a real deployment, and turned out to have nothing
87
+ to do with knowledge bases in particular.
88
+
89
+ ## Tests
90
+
91
+ ```bash
92
+ PYTHONPATH=src python3 -m unittest discover -s tests -t .
93
+ ```
94
+
95
+ Offline, dependency-free, no fixtures to download.
96
+
97
+ ## License
98
+
99
+ MIT — see [LICENSE](LICENSE).
@@ -0,0 +1,74 @@
1
+ # handrolled
2
+
3
+ **Hand-rolled, stdlib-only replacements for things you would normally `pip install`.**
4
+
5
+ Every module here exists for the same reason: the obvious solution adds a dependency,
6
+ and the dependency is bigger than the problem.
7
+
8
+ ```bash
9
+ pip install handrolled # installs nothing else
10
+ ```
11
+
12
+ Python ≥ 3.10, **no runtime dependencies**. `import handrolled` never touches
13
+ third-party code — the optional pieces import inside functions and return `None`
14
+ when their dependency is missing.
15
+
16
+ ## What's in it
17
+
18
+ | Module | Instead of | What it actually does |
19
+ |---|---|---|
20
+ | `handrolled.text` | `jieba`, a query builder, `PyYAML` | CJK bigram expansion for SQLite FTS5, FTS5 match-expression building (quotes and neutralises user input), and a YAML-**subset** front matter reader |
21
+ | `handrolled.webdav` | `requests` + `webdavclient3` | WebDAV over the RFC 4918 subset that gets used: `PROPFIND`, `GET`, `PUT`, `MKCOL`, `DELETE`; Basic auth, mTLS, and connect-host/SNI separation for appliances whose certificate CN is not their address |
22
+ | `handrolled.mounts` | parsing `mount` output | Network mount inspection that reports the **real** filesystem rather than its `autofs` trigger, `rw`/`ro` state, and `ESTALE`/`EACCES`/`ENOENT` turned into sentences a human can act on |
23
+ | `handrolled.pages` + `handrolled.importers` | hand-written one-off scripts | Turn another tool's export into Markdown pages: `mem0`, Obsidian, Notion's Markdown export, basic-memory, plain JSON/JSONL, and a neutral `records/v1` format. Deterministic, additive, guarded |
24
+ | `handrolled.rerank` | `sentence-transformers` | Cross-encoder reranking of `(query, passage)` pairs from an int8 ONNX model, with idle unloading so the memory comes back |
25
+ | `handrolled.deeprank` | an SDK per vendor | Let a remote chat model pick which candidate pages answer a query; any failure degrades to the original order |
26
+ | `handrolled.embed_onnx` | `sentence-transformers` | ONNX embedding backend (last-token pooling, L2 normalised), kept for A/B work |
27
+
28
+ ## Design rules
29
+
30
+ 1. **Standard library only.** If a module cannot be written with the stdlib and a
31
+ documented protocol, it does not belong here.
32
+ 2. **Optional means optional.** Anything needing a model or an extra package
33
+ imports lazily and returns `None`/`False` when it is absent. Nothing raises at
34
+ import time, and nothing silently changes behaviour because a package happens
35
+ to be installed.
36
+ 3. **No machine assumptions.** Optional layers have no built-in paths: unset
37
+ environment variable means the layer is off. Nothing guesses where your models
38
+ or credentials live.
39
+ 4. **Deterministic output.** `handrolled.pages` renders the same records to the
40
+ same bytes, so a second run is a no-op and a diff means something.
41
+
42
+ ## Quick example
43
+
44
+ ```python
45
+ from handrolled import importers
46
+ from handrolled.pages import plan, apply
47
+
48
+ records = importers.load("obsidian", "/path/to/vault")
49
+ planned = plan(policy, records, prefix="notes", zone="private")
50
+ print(f"{planned.items} records → {len(planned.pages)} pages")
51
+ apply(source, planned, write=True)
52
+ ```
53
+
54
+ `policy` is any object with `zones` (things with a `.name`) and
55
+ `is_never(rel) -> bool` — the caller owns the privacy model, not this library.
56
+
57
+ ## Where this comes from
58
+
59
+ Extracted from [dsh-kb](https://github.com/LUgABeEEE/dsh-kb), a personal
60
+ knowledge base built to be queried cheaply by AI agents. Those modules were
61
+ written there, hardened against a real deployment, and turned out to have nothing
62
+ to do with knowledge bases in particular.
63
+
64
+ ## Tests
65
+
66
+ ```bash
67
+ PYTHONPATH=src python3 -m unittest discover -s tests -t .
68
+ ```
69
+
70
+ Offline, dependency-free, no fixtures to download.
71
+
72
+ ## License
73
+
74
+ MIT — see [LICENSE](LICENSE).
@@ -0,0 +1,42 @@
1
+ [build-system]
2
+ requires = ["setuptools>=77"]
3
+ build-backend = "setuptools.build_meta"
4
+
5
+ [project]
6
+ name = "handrolled"
7
+ description = "Hand-rolled, stdlib-only building blocks for agent tooling: CJK full-text search, WebDAV, network mounts, and import adapters."
8
+ readme = "README.md"
9
+ requires-python = ">=3.10"
10
+ license = "MIT"
11
+ license-files = ["LICENSE"]
12
+ authors = [{ name = "handrolled contributors" }]
13
+ keywords = ["stdlib", "zero-dependency", "sqlite-fts5", "cjk", "webdav", "markdown", "agent"]
14
+ classifiers = [
15
+ "Development Status :: 3 - Alpha",
16
+ "Intended Audience :: Developers",
17
+ "Operating System :: OS Independent",
18
+ "Programming Language :: Python :: 3",
19
+ "Programming Language :: Python :: 3.10",
20
+ "Programming Language :: Python :: 3.11",
21
+ "Programming Language :: Python :: 3.12",
22
+ "Programming Language :: Python :: 3.13",
23
+ "Topic :: Text Processing :: Indexing",
24
+ "Topic :: Utilities",
25
+ ]
26
+ # The whole point: nothing to install, nothing to keep patched.
27
+ dependencies = []
28
+ dynamic = ["version"]
29
+
30
+ [project.urls]
31
+ Homepage = "https://github.com/LUgABeEEE/handrolled"
32
+ Source = "https://github.com/LUgABeEEE/handrolled"
33
+ Issues = "https://github.com/LUgABeEEE/handrolled/issues"
34
+
35
+ [tool.setuptools.dynamic]
36
+ version = { attr = "handrolled.__version__" }
37
+
38
+ [tool.setuptools.packages.find]
39
+ where = ["src"]
40
+
41
+ [tool.setuptools.package-data]
42
+ handrolled = ["py.typed"]
@@ -0,0 +1,4 @@
1
+ [egg_info]
2
+ tag_build =
3
+ tag_date = 0
4
+
@@ -0,0 +1,26 @@
1
+ """Hand-rolled, stdlib-only replacements for things you would normally pip install.
2
+
3
+ Every module here exists for the same reason: the obvious solution adds a
4
+ dependency, and the dependency is bigger than the problem.
5
+
6
+ * :mod:`handrolled.text` — CJK bigrams, FTS5 match expressions and a
7
+ YAML-subset front matter reader (instead of jieba, a query builder and PyYAML)
8
+ * :mod:`handrolled.webdav` — a WebDAV client speaking the RFC 4918 subset that
9
+ actually gets used (instead of requests + webdavclient3)
10
+ * :mod:`handrolled.mounts` — network mount inspection that reports what is
11
+ really mounted, not its autofs trigger (instead of parsing ``mount`` output)
12
+ * :mod:`handrolled.pages` + :mod:`handrolled.importers` — turn another tool's
13
+ export (mem0, Obsidian, Notion, basic-memory, JSON/JSONL) into Markdown pages
14
+ * :mod:`handrolled.rerank`, :mod:`handrolled.deeprank`,
15
+ :mod:`handrolled.embed_onnx` — optional ranking helpers that degrade to
16
+ ``None`` when their model or dependency is absent
17
+
18
+ Nothing here imports a third-party package at module import time; the optional
19
+ ones import inside functions and return ``None`` when the dependency is missing.
20
+ """
21
+
22
+ from __future__ import annotations
23
+
24
+ __version__ = "0.1.0"
25
+
26
+ __all__ = ["__version__"]
@@ -0,0 +1,183 @@
1
+ """深池 LLM 重排(`deep=true` 慢车道)。
2
+
3
+ 为什么单独一层而不是改主链路:
4
+ * 主链路要快(p95 < 1s),LLM 一次 5~15s,只能按需触发;
5
+ * 需要「看得到更远的候选」——先把候选扩到 40 页,再让模型挑;
6
+ * 失败必须静默降级:API 挂了就退回普通排序,绝不影响检索可用性。
7
+
8
+ 一次 33 题评测:普通 25/33 → 深池 + 挑 1 页 = 26/33,代价 p50 8.4s。
9
+ 单次用量约 3.7k 输入 token / 120 输出 token。
10
+
11
+ 注意一个实测过的失败模式:候选多时模型会报出**超出范围的编号**(例如只有 30 页却答 31)。
12
+ 所以解析不到合法编号时,会用「前 10 页」重试一次,仍然失败才降级。
13
+ """
14
+
15
+ from __future__ import annotations
16
+
17
+ import json
18
+ import os
19
+ import re
20
+ import urllib.request
21
+ from typing import Sequence
22
+
23
+ #: This is an opt-in slow lane: nothing below runs unless the deployment says
24
+ #: where to send the request and with which credentials. No vendor endpoint,
25
+ #: no key file and no log path is assumed.
26
+ DEFAULT_MODEL = ""
27
+ DEFAULT_BASE = ""
28
+ CRED_PATH = os.environ.get("DSH_KB_DEEP_CREDENTIALS") or ""
29
+
30
+ LAST_RAW = ""
31
+ LAST_ERROR = ""
32
+ LOG_PATH = os.environ.get("DSH_KB_DEEP_LOG") or ""
33
+
34
+ PROMPT = """用户问题:%s
35
+
36
+ 下面是知识库候选页面(编号 1-%d,含该页最相关片段)。请挑出最可能回答了该问题的 %d 个页面,
37
+ 按可能性从高到低输出编号,逗号分隔,只输出编号(必须是 1 到 %d 之间的整数)。
38
+
39
+ %s"""
40
+
41
+
42
+ def _api_key() -> str | None:
43
+ """API key: ``DSH_KB_DEEP_API_KEY`` first, then an optional credentials file.
44
+
45
+ The file is whatever ``DSH_KB_DEEP_CREDENTIALS`` points at and only has to
46
+ contain ``DEEP_API_KEY`` (top level or under ``refs``). YAML is used when
47
+ PyYAML happens to be importable, otherwise one line is matched with a regex —
48
+ the MCP server interpreter may not have yaml installed.
49
+ """
50
+ global LAST_ERROR
51
+ key = (os.environ.get("DSH_KB_DEEP_API_KEY") or "").strip()
52
+ if key:
53
+ return key
54
+ if not CRED_PATH:
55
+ LAST_ERROR = "no-key"
56
+ return None
57
+ try:
58
+ raw = open(CRED_PATH, encoding="utf-8").read()
59
+ except Exception as exc:
60
+ LAST_ERROR = "cred-read: %s" % type(exc).__name__
61
+ return None
62
+ try:
63
+ import yaml
64
+
65
+ data = yaml.safe_load(raw) or {}
66
+ refs = data.get("refs") or data
67
+ value = refs.get("DEEP_API_KEY")
68
+ if value:
69
+ return str(value)
70
+ except Exception:
71
+ pass
72
+ m = re.search(r"^\s*DEEP_API_KEY\s*:\s*(\S+)", raw, re.M)
73
+ if m:
74
+ return m.group(1).strip().strip(chr(34)).strip(chr(39))
75
+ LAST_ERROR = "no-key"
76
+ return None
77
+
78
+
79
+ def _endpoint() -> str:
80
+ """OpenAI-compatible base URL, without the trailing slash."""
81
+ return (os.environ.get("DSH_KB_DEEP_ENDPOINT") or DEFAULT_BASE).rstrip("/")
82
+
83
+
84
+ def _parse_order(text: str, n: int, pick: int) -> list[int]:
85
+ """把「3,1,7」解析成 1 基下标;越界与重复丢掉。"""
86
+ out: list[int] = []
87
+ for raw in re.findall(r"\d+", text or ""):
88
+ try:
89
+ idx = int(raw)
90
+ except ValueError:
91
+ continue
92
+ if 1 <= idx <= n and idx not in out:
93
+ out.append(idx)
94
+ if len(out) >= pick:
95
+ break
96
+ return out
97
+
98
+
99
+ def _request(query: str, pages: Sequence[tuple[str, str]], *, key: str, model: str,
100
+ pick: int, timeout: float) -> tuple[list[int], str]:
101
+ blocks = ["[%d] %s" % (i, path.split("/")[-1]) + chr(10) + (body or "")[:320]
102
+ for i, (path, body) in enumerate(pages, 1)]
103
+ prompt = PROMPT % (query, len(pages), pick, len(pages), chr(10) + chr(10).join(blocks))
104
+ body = json.dumps({"model": model, "temperature": 0, "max_tokens": 400,
105
+ "messages": [{"role": "user", "content": prompt}]}).encode("utf-8")
106
+ req = urllib.request.Request(_endpoint() + "/chat/completions", data=body,
107
+ headers={"Content-Type": "application/json",
108
+ "Authorization": "Bearer " + key})
109
+ with urllib.request.urlopen(req, timeout=timeout) as resp:
110
+ data = json.loads(resp.read().decode("utf-8"))
111
+ text = ((data.get("choices") or [{}])[0].get("message") or {}).get("content") or ""
112
+ return _parse_order(text, len(pages), pick), text
113
+
114
+
115
+ def _log_failure(query: str, error: str, raw: str) -> None:
116
+ """降级时留一行证据:MCP 进程看不到 stdout,只有文件能查。
117
+
118
+ Nothing is written unless ``DSH_KB_DEEP_LOG`` is set — a library has no
119
+ business picking a log location on someone else's machine.
120
+ """
121
+ if not LOG_PATH:
122
+ return
123
+ try:
124
+ from datetime import datetime
125
+
126
+ line = "%s\t%s\t%s\t%s" % (
127
+ datetime.now().isoformat(timespec="seconds"),
128
+ error,
129
+ (query or "")[:40].replace("\t", " "),
130
+ (raw or "").replace("\t", " ")[:40],
131
+ )
132
+ with open(LOG_PATH, "a", encoding="utf-8") as fh:
133
+ fh.write(line + chr(10))
134
+ except Exception:
135
+ pass
136
+
137
+
138
+ def pick_pages(
139
+ query: str,
140
+ pages: Sequence[tuple[str, str]],
141
+ *,
142
+ model: str = DEFAULT_MODEL,
143
+ pick: int = 1,
144
+ timeout: float = 25.0,
145
+ ) -> list[int] | None:
146
+ """让模型从候选页面里挑 pick 个;失败返回 None。
147
+
148
+ ``pages`` 是 [(path, passage)],passage 由调用方切好(避免这里再读盘)。
149
+ """
150
+ global LAST_RAW, LAST_ERROR
151
+ LAST_RAW, LAST_ERROR = "", ""
152
+ if not pages:
153
+ return []
154
+ key = _api_key()
155
+ if not key:
156
+ LAST_ERROR = "no-key"
157
+ return None
158
+ if not _endpoint() or not model:
159
+ LAST_ERROR = "no-endpoint-or-model"
160
+ return None
161
+ for attempt, subset in enumerate((list(pages), list(pages)[:10])):
162
+ if attempt and len(subset) >= len(pages):
163
+ break
164
+ try:
165
+ order, text = _request(query, subset, key=key, model=model, pick=pick,
166
+ timeout=timeout if attempt == 0 else min(timeout, 15.0))
167
+ except Exception as exc: # noqa: BLE001 - 任何失败都降级
168
+ LAST_ERROR = "%s: %s" % (type(exc).__name__, str(exc)[:120])
169
+ continue
170
+ LAST_RAW = (text or "").strip()[:80]
171
+ if order:
172
+ return order
173
+ LAST_ERROR = "unparsable"
174
+ _log_failure(query, LAST_ERROR, LAST_RAW)
175
+ return None
176
+
177
+
178
+ def reorder(items: Sequence, order: Sequence[int]) -> list:
179
+ """按 1 基下标把挑中的提到最前,其余保序。"""
180
+ idx = [i - 1 for i in order if 0 < i <= len(items)]
181
+ picked = [items[i] for i in idx]
182
+ rest = [x for j, x in enumerate(items) if j not in set(idx)]
183
+ return picked + rest
@@ -0,0 +1,168 @@
1
+ """ONNX 嵌入后端(Qwen3-Embedding-0.6B int8)—— 换模型 A/B 实验用。
2
+
3
+ 为什么单独一个模块:
4
+ * 生产嵌入走 ollama(bge-m3)。要评估"换模型能不能把召回抬上去",必须能在
5
+ 不改生产代码路径的前提下换掉**嵌入实现**,同时仍复用 VectorArm 的矩阵缓存、
6
+ 页级去重、重排与融合逻辑 —— 那些是逻辑资产,不该跟着模型一起换。
7
+ * 模型是 int8 ONNX(约 610 MB),本仓库不附带、不下载;路径由
8
+ ``DSH_KB_EMBED_ONNX_ROOT`` 指定(其下 ``onnx/model_int8.onnx`` + ``tokenizer.json``),
9
+ 依赖库目录由 ``DSH_KB_ONNX_LIB`` 指定(可选,能 import 到即可)。
10
+
11
+ 池化方式(Qwen3-Embedding 官方口径):
12
+ * 取**最后一个 token** 的 hidden state(不是 CLS,也不是 mean);
13
+ * L2 归一化;
14
+ * 查询侧加指令前缀 ``Instruct: {task}\nQuery: {query}``,文档侧不加 —— 这是它
15
+ 与非对称检索模型拉开差距的地方,和 bge-m3 的用法不同,别混。
16
+ """
17
+
18
+ from __future__ import annotations
19
+
20
+ import os
21
+ import sys
22
+ import threading
23
+ from pathlib import Path
24
+
25
+ __all__ = ["available", "embed_texts", "embed_query", "dim", "MODEL_NAME"]
26
+
27
+ #: Both are unset by default: the ONNX backend is an experiment, not a shipped path.
28
+ _ROOT = os.environ.get("DSH_KB_EMBED_ONNX_ROOT")
29
+ _LIB = os.environ.get("DSH_KB_ONNX_LIB")
30
+ ROOT = Path(_ROOT) if _ROOT else None
31
+ LIB = Path(_LIB) if _LIB else None
32
+ MODEL_NAME = os.environ.get("DSH_KB_EMBED_ONNX_NAME", "qwen3-embedding-0.6b-int8")
33
+ INSTRUCTION = os.environ.get(
34
+ "DSH_KB_EMBED_INSTRUCTION",
35
+ "Given a web search query, retrieve relevant passages that answer the query",
36
+ )
37
+ DIM = 1024
38
+ MAX_LEN = int(os.environ.get("DSH_KB_EMBED_MAXLEN", "512"))
39
+ BATCH = int(os.environ.get("DSH_KB_EMBED_BATCH", "8"))
40
+
41
+ _LOCK = threading.Lock()
42
+ _SESSION = None
43
+ _TOKENIZER = None
44
+ _INPUT_NAMES = ()
45
+ _INPUT_META = ()
46
+
47
+
48
+ def _lib_path() -> None:
49
+ if LIB is not None and LIB.is_dir() and str(LIB) not in sys.path:
50
+ sys.path.insert(0, str(LIB))
51
+
52
+
53
+ def available() -> bool:
54
+ """模型文件齐全且能 import 依赖时才为 True(不加载模型,不看网络)。"""
55
+ if ROOT is None:
56
+ return False
57
+ need = [ROOT / "onnx" / "model_int8.onnx", ROOT / "tokenizer.json"]
58
+ if not all(p.is_file() and p.stat().st_size > 1000 for p in need):
59
+ return False
60
+ try:
61
+ _lib_path()
62
+ import numpy # noqa: F401
63
+ import onnxruntime # noqa: F401
64
+ import tokenizers # noqa: F401
65
+ except Exception:
66
+ return False
67
+ return True
68
+
69
+
70
+ def _load():
71
+ global _SESSION, _TOKENIZER, _INPUT_NAMES
72
+ if _SESSION is not None:
73
+ return _SESSION, _TOKENIZER
74
+ with _LOCK:
75
+ if _SESSION is not None:
76
+ return _SESSION, _TOKENIZER
77
+ _lib_path()
78
+ import onnxruntime as ort
79
+ from tokenizers import Tokenizer
80
+
81
+ opts = ort.SessionOptions()
82
+ opts.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL
83
+ opts.intra_op_num_threads = int(os.environ.get("DSH_KB_EMBED_THREADS", "4"))
84
+ opts.inter_op_num_threads = 1
85
+ # 内存:推理 arena 会一直占着,和别的常驻服务共用一台小机器时很痛。默认
86
+ # 中间张量池留到进程结束(实测 4 进程各占 3G),关掉后 RSS 降到 ~1.3G。
87
+ opts.enable_cpu_mem_arena = os.environ.get("DSH_KB_EMBED_ARENA", "0") == "1"
88
+ sess = ort.InferenceSession(
89
+ str(ROOT / "onnx" / "model_int8.onnx"), opts, providers=["CPUExecutionProvider"]
90
+ )
91
+ tok = Tokenizer.from_file(str(ROOT / "tokenizer.json"))
92
+ tok.enable_truncation(max_length=MAX_LEN)
93
+ tok.no_padding()
94
+ _SESSION, _TOKENIZER = sess, tok
95
+ global _INPUT_META
96
+ _INPUT_NAMES = tuple(i.name for i in sess.get_inputs())
97
+ _INPUT_META = tuple((i.name, list(i.shape or ())) for i in sess.get_inputs())
98
+ return sess, tok
99
+
100
+
101
+ def dim() -> int:
102
+ return DIM
103
+
104
+
105
+ def _postprocess(hidden, mask):
106
+ """last-token pooling(按 attention_mask 取每条最后一个真实 token)+ L2。"""
107
+ import numpy as np
108
+
109
+ if isinstance(hidden, (list, tuple)):
110
+ hidden = hidden[0]
111
+ arr = np.asarray(hidden)
112
+ if arr.ndim == 3:
113
+ idx = np.asarray(mask).sum(axis=1) - 1
114
+ idx = np.clip(idx, 0, arr.shape[1] - 1)
115
+ arr = arr[np.arange(arr.shape[0]), idx]
116
+ norms = np.linalg.norm(arr, axis=1, keepdims=True)
117
+ norms[norms == 0] = 1.0
118
+ return arr / norms
119
+
120
+
121
+ def embed_texts(texts, *, is_query: bool = False, batch_size: int | None = None):
122
+ """批量嵌入,返回 list[list[float]](float32 归一化向量)。"""
123
+ if isinstance(texts, str):
124
+ texts = [texts]
125
+ sess, tok = _load()
126
+ import numpy as np
127
+
128
+ bs = batch_size or BATCH
129
+ prepared = []
130
+ for t in texts:
131
+ body = str(t or "").strip()
132
+ if is_query:
133
+ body = "Instruct: " + INSTRUCTION + "\nQuery: " + body
134
+ prepared.append(body or " ")
135
+ out = []
136
+ for i in range(0, len(prepared), bs):
137
+ chunk = prepared[i : i + bs]
138
+ encs = tok.encode_batch(chunk)
139
+ maxlen = max(len(e.ids) for e in encs)
140
+ ids = np.zeros((len(encs), maxlen), dtype="int64")
141
+ mask = np.zeros((len(encs), maxlen), dtype="int64")
142
+ for r, e in enumerate(encs):
143
+ ids[r, : len(e.ids)] = e.ids
144
+ mask[r, : len(e.ids)] = 1
145
+ feeds = {
146
+ "input_ids": ids,
147
+ "attention_mask": mask,
148
+ "position_ids": np.tile(np.arange(maxlen, dtype="int64"), (len(encs), 1)),
149
+ }
150
+ # 这个 ONNX 导出是「带 KV cache」的解码器图:56 个 past 输入必须显式给空张量
151
+ # (past_sequence_length=0)。不给就直接 ValueError;给了等价于无缓存的整段前向,
152
+ # 我们要的正是 last_hidden_state 整段。
153
+ for name, meta in _INPUT_META:
154
+ if name in feeds:
155
+ continue
156
+ if name.startswith("past_key_values"):
157
+ shape = [len(encs) if str(s).startswith("batch") else int(s)
158
+ if not str(s).startswith("past") else 0
159
+ for s in (meta or [])]
160
+ feeds[name] = np.zeros(shape or (len(encs), 1, 0, 1), dtype="float32")
161
+ hidden = sess.run(None, feeds)
162
+ vecs = _postprocess(hidden, mask)
163
+ out.extend([[float(x) for x in row] for row in vecs])
164
+ return out
165
+
166
+
167
+ def embed_query(text: str):
168
+ return embed_texts([text], is_query=True)[0]