okfgraph 0.2.6__tar.gz → 0.2.11__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.
- okfgraph-0.2.11/PKG-INFO +359 -0
- {okfgraph-0.2.6 → okfgraph-0.2.11}/README.md +5 -1
- {okfgraph-0.2.6 → okfgraph-0.2.11}/okfgraph/components/__init__.py +2 -1
- {okfgraph-0.2.6 → okfgraph-0.2.11}/okfgraph/components/embedding.py +232 -13
- {okfgraph-0.2.6 → okfgraph-0.2.11}/okfgraph/router.py +77 -15
- okfgraph-0.2.11/okfgraph.egg-info/PKG-INFO +359 -0
- {okfgraph-0.2.6 → okfgraph-0.2.11}/okfgraph.egg-info/SOURCES.txt +4 -0
- {okfgraph-0.2.6 → okfgraph-0.2.11}/pyproject.toml +2 -1
- okfgraph-0.2.11/tests/test_explicit_files.py +147 -0
- {okfgraph-0.2.6 → okfgraph-0.2.11}/tests/test_gpu_integration.py +5 -1
- okfgraph-0.2.11/tests/test_lazy_encoder.py +202 -0
- okfgraph-0.2.11/tests/test_ort.py +196 -0
- okfgraph-0.2.11/tests/test_packaging.py +29 -0
- {okfgraph-0.2.6 → okfgraph-0.2.11}/tests/test_rust_backend.py +32 -1
- {okfgraph-0.2.6 → okfgraph-0.2.11}/tests/test_rust_e2e.py +4 -1
- okfgraph-0.2.6/PKG-INFO +0 -25
- okfgraph-0.2.6/okfgraph.egg-info/PKG-INFO +0 -25
- {okfgraph-0.2.6 → okfgraph-0.2.11}/LICENSE +0 -0
- {okfgraph-0.2.6 → okfgraph-0.2.11}/okfgraph/__init__.py +0 -0
- {okfgraph-0.2.6 → okfgraph-0.2.11}/okfgraph/cli.py +0 -0
- {okfgraph-0.2.6 → okfgraph-0.2.11}/okfgraph/components/converters.py +0 -0
- {okfgraph-0.2.6 → okfgraph-0.2.11}/okfgraph/components/delta.py +0 -0
- {okfgraph-0.2.6 → okfgraph-0.2.11}/okfgraph/components/diff.py +0 -0
- {okfgraph-0.2.6 → okfgraph-0.2.11}/okfgraph/components/doctor.py +0 -0
- {okfgraph-0.2.6 → okfgraph-0.2.11}/okfgraph/components/export.py +0 -0
- {okfgraph-0.2.6 → okfgraph-0.2.11}/okfgraph/components/image_assets.py +0 -0
- {okfgraph-0.2.6 → okfgraph-0.2.11}/okfgraph/components/import_.py +0 -0
- {okfgraph-0.2.6 → okfgraph-0.2.11}/okfgraph/components/ingest.py +0 -0
- {okfgraph-0.2.6 → okfgraph-0.2.11}/okfgraph/components/links.py +0 -0
- {okfgraph-0.2.6 → okfgraph-0.2.11}/okfgraph/components/lint.py +0 -0
- {okfgraph-0.2.6 → okfgraph-0.2.11}/okfgraph/components/purge.py +0 -0
- {okfgraph-0.2.6 → okfgraph-0.2.11}/okfgraph/components/ranking.py +0 -0
- {okfgraph-0.2.6 → okfgraph-0.2.11}/okfgraph/components/schema.py +0 -0
- {okfgraph-0.2.6 → okfgraph-0.2.11}/okfgraph/components/search.py +0 -0
- {okfgraph-0.2.6 → okfgraph-0.2.11}/okfgraph/config.py +0 -0
- {okfgraph-0.2.6 → okfgraph-0.2.11}/okfgraph/images.py +0 -0
- {okfgraph-0.2.6 → okfgraph-0.2.11}/okfgraph/mcp_server.py +0 -0
- {okfgraph-0.2.6 → okfgraph-0.2.11}/okfgraph/models.py +0 -0
- {okfgraph-0.2.6 → okfgraph-0.2.11}/okfgraph/security.py +0 -0
- {okfgraph-0.2.6 → okfgraph-0.2.11}/okfgraph/tools.py +0 -0
- {okfgraph-0.2.6 → okfgraph-0.2.11}/okfgraph.egg-info/dependency_links.txt +0 -0
- {okfgraph-0.2.6 → okfgraph-0.2.11}/okfgraph.egg-info/entry_points.txt +0 -0
- {okfgraph-0.2.6 → okfgraph-0.2.11}/okfgraph.egg-info/requires.txt +0 -0
- {okfgraph-0.2.6 → okfgraph-0.2.11}/okfgraph.egg-info/top_level.txt +0 -0
- {okfgraph-0.2.6 → okfgraph-0.2.11}/setup.cfg +0 -0
- {okfgraph-0.2.6 → okfgraph-0.2.11}/tests/test_chunk_search.py +0 -0
- {okfgraph-0.2.6 → okfgraph-0.2.11}/tests/test_chunking.py +0 -0
- {okfgraph-0.2.6 → okfgraph-0.2.11}/tests/test_cli.py +0 -0
- {okfgraph-0.2.6 → okfgraph-0.2.11}/tests/test_config.py +0 -0
- {okfgraph-0.2.6 → okfgraph-0.2.11}/tests/test_converter.py +0 -0
- {okfgraph-0.2.6 → okfgraph-0.2.11}/tests/test_delta.py +0 -0
- {okfgraph-0.2.6 → okfgraph-0.2.11}/tests/test_diff.py +0 -0
- {okfgraph-0.2.6 → okfgraph-0.2.11}/tests/test_directory_hash.py +0 -0
- {okfgraph-0.2.6 → okfgraph-0.2.11}/tests/test_doctor.py +0 -0
- {okfgraph-0.2.6 → okfgraph-0.2.11}/tests/test_export_compliance.py +0 -0
- {okfgraph-0.2.6 → okfgraph-0.2.11}/tests/test_graph_enrichment.py +0 -0
- {okfgraph-0.2.6 → okfgraph-0.2.11}/tests/test_images.py +0 -0
- {okfgraph-0.2.6 → okfgraph-0.2.11}/tests/test_ingest.py +0 -0
- {okfgraph-0.2.6 → okfgraph-0.2.11}/tests/test_ingest_tool.py +0 -0
- {okfgraph-0.2.6 → okfgraph-0.2.11}/tests/test_integration.py +0 -0
- {okfgraph-0.2.6 → okfgraph-0.2.11}/tests/test_lint.py +0 -0
- {okfgraph-0.2.6 → okfgraph-0.2.11}/tests/test_logging.py +0 -0
- {okfgraph-0.2.6 → okfgraph-0.2.11}/tests/test_mcp_server.py +0 -0
- {okfgraph-0.2.6 → okfgraph-0.2.11}/tests/test_models.py +0 -0
- {okfgraph-0.2.6 → okfgraph-0.2.11}/tests/test_obsidian.py +0 -0
- {okfgraph-0.2.6 → okfgraph-0.2.11}/tests/test_okf_ingest_tool.py +0 -0
- {okfgraph-0.2.6 → okfgraph-0.2.11}/tests/test_parity.py +0 -0
- {okfgraph-0.2.6 → okfgraph-0.2.11}/tests/test_pdf_e2e.py +0 -0
- {okfgraph-0.2.6 → okfgraph-0.2.11}/tests/test_ppr_search.py +0 -0
- {okfgraph-0.2.6 → okfgraph-0.2.11}/tests/test_ranking.py +0 -0
- {okfgraph-0.2.6 → okfgraph-0.2.11}/tests/test_reconstruction.py +0 -0
- {okfgraph-0.2.6 → okfgraph-0.2.11}/tests/test_reserved.py +0 -0
- {okfgraph-0.2.6 → okfgraph-0.2.11}/tests/test_roundup_cli.py +0 -0
- {okfgraph-0.2.6 → okfgraph-0.2.11}/tests/test_router.py +0 -0
- {okfgraph-0.2.6 → okfgraph-0.2.11}/tests/test_router_misc.py +0 -0
- {okfgraph-0.2.6 → okfgraph-0.2.11}/tests/test_sanitize.py +0 -0
- {okfgraph-0.2.6 → okfgraph-0.2.11}/tests/test_search_browser.py +0 -0
- {okfgraph-0.2.6 → okfgraph-0.2.11}/tests/test_security.py +0 -0
okfgraph-0.2.11/PKG-INFO
ADDED
|
@@ -0,0 +1,359 @@
|
|
|
1
|
+
Metadata-Version: 2.4
|
|
2
|
+
Name: okfgraph
|
|
3
|
+
Version: 0.2.11
|
|
4
|
+
Summary: Ladybug-backed OKF knowledge graph with ONNX + Jina v5 embeddings
|
|
5
|
+
License-Expression: Apache-2.0 OR MIT
|
|
6
|
+
Requires-Python: >=3.11
|
|
7
|
+
Description-Content-Type: text/markdown
|
|
8
|
+
License-File: LICENSE
|
|
9
|
+
Requires-Dist: ladybug==0.20.3
|
|
10
|
+
Requires-Dist: okf-embed>=0.1
|
|
11
|
+
Requires-Dist: pydantic>=2.0
|
|
12
|
+
Requires-Dist: python-frontmatter>=1.0
|
|
13
|
+
Requires-Dist: pyyaml>=6.0
|
|
14
|
+
Requires-Dist: numpy>=1.26
|
|
15
|
+
Requires-Dist: onnxruntime==1.29.0
|
|
16
|
+
Requires-Dist: mordant>=0.9
|
|
17
|
+
Requires-Dist: mcp>=2.0
|
|
18
|
+
Requires-Dist: fasteners>=0.19
|
|
19
|
+
Provides-Extra: dev
|
|
20
|
+
Requires-Dist: pytest>=8.0; extra == "dev"
|
|
21
|
+
Provides-Extra: pdf
|
|
22
|
+
Requires-Dist: bobine>=0.5; extra == "pdf"
|
|
23
|
+
Provides-Extra: omni
|
|
24
|
+
Requires-Dist: sentence-transformers>=3.0; extra == "omni"
|
|
25
|
+
Requires-Dist: Pillow>=10.0; extra == "omni"
|
|
26
|
+
Dynamic: license-file
|
|
27
|
+
|
|
28
|
+
# OKFgraph 0.2.7
|
|
29
|
+
|
|
30
|
+
[](https://pypi.org/project/okfgraph/)
|
|
31
|
+
[](https://www.python.org/)
|
|
32
|
+
[](LICENSE)
|
|
33
|
+
[](https://modelcontextprotocol.io/)
|
|
34
|
+
[](https://pypi.org/project/ladybug/)
|
|
35
|
+
|
|
36
|
+
**Ladybug-backed knowledge graph with Rust-driven Jina v5 embeddings, model-free
|
|
37
|
+
graph retrieval, and agent-first MCP + CLI surfaces.**
|
|
38
|
+
|
|
39
|
+
OKFgraph is a Python library, CLI (`okf`), and MCP server (`okf-mcp`) for
|
|
40
|
+
building, querying, and maintaining a knowledge graph from Markdown/OKF
|
|
41
|
+
documents. It combines hybrid semantic search (vector + FTS, RRF-fused),
|
|
42
|
+
chunk-level retrieval, **model-free PPR retrieval that works with no embedding
|
|
43
|
+
model loaded**, structural diffing, scored health checks, and Obsidian-vault
|
|
44
|
+
compatible wikilinks — in a single LadybugDB file.
|
|
45
|
+
|
|
46
|
+
Design stance: **slim dependencies, no legacy fallbacks.** Embeddings run
|
|
47
|
+
through the `okf-embed` Rust wheel (ONNX Runtime, no torch / transformers /
|
|
48
|
+
optimum anywhere). PDF conversion runs through the `bobine` Rust engine behind
|
|
49
|
+
a swappable `DocumentConverter` seam. What isn't needed isn't installed.
|
|
50
|
+
|
|
51
|
+
---
|
|
52
|
+
|
|
53
|
+
## Features
|
|
54
|
+
|
|
55
|
+
| Category | Features |
|
|
56
|
+
|---|---|
|
|
57
|
+
| **Embeddings** | Jina v5 (`jina-embeddings-v5-text-small-retrieval`) via the `okf-embed` Rust wheel; last-token pooling, Matryoshka truncation (32–1024, default 512); omni model (`…-omni-small-retrieval`) lazy-loaded for images only |
|
|
58
|
+
| **Search** | Hybrid RRF fusion (vector + FTS) at concept and chunk granularity; `rank=none\|hub\|ppr` — including **PPR**: lexical seeds → exact Personalized PageRank, zero model load, deterministic |
|
|
59
|
+
| **Read** | Body / chunks / rebuilt document / graph context, with optional **token budgets** (`max_tokens`): self first, then PPR-ranked neighbours, index-first for context |
|
|
60
|
+
| **Storage** | LadybugDB `==0.20.3` (pinned — newer 0.20.x segfaults index builds): graph + vector + FTS in one file |
|
|
61
|
+
| **Links** | Path links (`](doc.md)`) + `[[wikilinks]]` resolved by name (`uid` → `aliases` → `title` → filename stem); ambiguous names never resolve; broken links tracked and repairable |
|
|
62
|
+
| **Import** | Single files, whole bundles (delta-aware: only changed files re-embed, `--purge` drops deleted concepts), markdown / PDF / raw thoughts; mordant lint on the way in |
|
|
63
|
+
| **PDF** | `DocumentConverter` seam with `BobineConverter` default (routing `auto\|surgical\|always\|never`); bring your own converter, no code changes |
|
|
64
|
+
| **Diff** | `okf diff`: structural snapshot (dir vs dir, no model load) and drift (graph vs dir) modes; CI exit codes + `--json` |
|
|
65
|
+
| **Doctor** | `okf doctor`: 0–100 health score (broken/orphan/stale/duplicate-title/missing-description), safe `--fix` that never touches `reviewed: true`, `--strict` CI gate |
|
|
66
|
+
| **Export** | OKF round-trip with See Also / Cited By enrichment + index files, or `--flavor obsidian` (`[[Title]]` wikilinks, no index files, edge-lossless re-import) |
|
|
67
|
+
| **Images** | Modes `text` / `optional` / `omni`, shared text+image vector space, content-hash dedup, `okf-asset://` protocol |
|
|
68
|
+
| **MCP** | 5 tools (`search`, `read`, `traverse`, `ingest`, `export_bundle`), MCP ≥ 2.0 (`MCPServer` + `ToolAnnotations`), stdio transport |
|
|
69
|
+
| **CLI** | Same 5 verbs plus maintenance (`init`, `import`, `diff`, `doctor`, `shell`, `reindex`, `broken-links`, `deleted-*`, …), slim per-command help, `okfgraph.toml` config |
|
|
70
|
+
| **Skills** | 3 harness-neutral skills (`okfgraph-mcp`, `okfgraph-cli`, `okfgraph-ingest`), `.mcp.json` wiring included |
|
|
71
|
+
|
|
72
|
+
---
|
|
73
|
+
|
|
74
|
+
## Installation
|
|
75
|
+
|
|
76
|
+
Requires Python ≥ 3.11.
|
|
77
|
+
|
|
78
|
+
```bash
|
|
79
|
+
pip install "okfgraph[pdf,omni]" # PyPI (okf-embed ships platform wheels; published)
|
|
80
|
+
```
|
|
81
|
+
|
|
82
|
+
Or from source with `uv` (also builds the `okf-embed` Rust wheel from
|
|
83
|
+
`rust/okf-embed` via `[tool.uv.sources]`):
|
|
84
|
+
|
|
85
|
+
```bash
|
|
86
|
+
git clone <repo> && cd OKFgraph
|
|
87
|
+
uv sync # core: ladybug, okf-embed, onnxruntime, mordant, mcp, …
|
|
88
|
+
uv sync --extra pdf # bobine PDF converter
|
|
89
|
+
uv sync --extra omni # sentence-transformers + Pillow (image embeddings)
|
|
90
|
+
uv sync --extra dev # pytest
|
|
91
|
+
```
|
|
92
|
+
|
|
93
|
+
Core dependencies are deliberately few: `ladybug==0.20.3`, `okf-embed`,
|
|
94
|
+
`onnxruntime==1.29.0` (one pinned ORT binary shared by bobine + okf-embed;
|
|
95
|
+
`ORT_DYLIB_PATH`-overridable), `mordant`, `mcp>=2.0`, `pydantic`, `pyyaml`,
|
|
96
|
+
`numpy`, `python-frontmatter`, `fasteners`. There is no torch, no
|
|
97
|
+
transformers, no optimum, no PDF stack in the core install.
|
|
98
|
+
|
|
99
|
+
---
|
|
100
|
+
|
|
101
|
+
## Quick start
|
|
102
|
+
|
|
103
|
+
```bash
|
|
104
|
+
okf init --db kb.db --bundle kb/ # once; persists to okfgraph.toml
|
|
105
|
+
okf import --all --bundle kb/ # delta-aware bulk import
|
|
106
|
+
okf search "honey badger defense" # hybrid semantic search
|
|
107
|
+
okf search --rank ppr "honey badger" # same question, no model load
|
|
108
|
+
okf read savanna --include context --max-tokens 1500
|
|
109
|
+
okf traverse savanna --relationship LINKS_TO --direction BOTH
|
|
110
|
+
okf diff # drift: graph vs bundle dir
|
|
111
|
+
okf doctor # health score + findings
|
|
112
|
+
okf lint kb/ # pre-import gate: frontmatter + links, no model
|
|
113
|
+
```
|
|
114
|
+
|
|
115
|
+
### Programmatic usage
|
|
116
|
+
|
|
117
|
+
```python
|
|
118
|
+
from pathlib import Path
|
|
119
|
+
from okfgraph import OKFRouter
|
|
120
|
+
|
|
121
|
+
router = OKFRouter(db_path="kb.db", bundle_root="kb/", embedding_dim=512)
|
|
122
|
+
|
|
123
|
+
# Import: whole bundle (delta-aware), one file, or raw reasoning
|
|
124
|
+
ids = router.import_mgr.import_bundle(Path("kb/"), purge_deleted=False)
|
|
125
|
+
cid = router.import_from_okf(Path("kb/auth.md"), mode="text")
|
|
126
|
+
th = router.ingest_mgr.ingest_thoughts("Session decided X because Y", topic="auth-refactor")
|
|
127
|
+
|
|
128
|
+
# Search: hybrid (default), hub-blended, or model-free PPR
|
|
129
|
+
hits = router.search_hybrid("session auth decisions", limit=5)
|
|
130
|
+
hubbed = router.search_hybrid("session auth decisions", rank="hub")
|
|
131
|
+
cold = router.search_engine.search_with_ppr("auth refactor") # no ONNX load
|
|
132
|
+
|
|
133
|
+
# Read: full shapes, or a token-budgeted section list
|
|
134
|
+
doc = router.get_by_id("auth")
|
|
135
|
+
reading = router.search_engine.read_with_budget("auth", include="context", max_tokens=1500)
|
|
136
|
+
|
|
137
|
+
# Maintain: drift, health, links
|
|
138
|
+
print(router.diff_db_dir(Path("kb/"))["identical"]) # True when in sync
|
|
139
|
+
print(router.diagnose()["score"]) # 0-100
|
|
140
|
+
print(router.doctor_fix()) # safe repairs only
|
|
141
|
+
print(router.repair_links()) # re-point broken links
|
|
142
|
+
|
|
143
|
+
# Export: OKF bundle or Obsidian vault
|
|
144
|
+
router.export_mgr.export_bundle(Path("out-okf/"))
|
|
145
|
+
router.export_mgr.export_bundle(Path("out-vault/"), flavor="obsidian")
|
|
146
|
+
|
|
147
|
+
router.close()
|
|
148
|
+
```
|
|
149
|
+
|
|
150
|
+
### Ingestion kinds
|
|
151
|
+
|
|
152
|
+
| Kind | Entry point | Notes |
|
|
153
|
+
|---|---|---|
|
|
154
|
+
| `md` | `ingest_md(md_path, concept_id?, title?, tags?, mode?)` | mordant-linted; frontmatter wins over overrides |
|
|
155
|
+
| `pdf` | `ingest_pdf(pdf_path, auto_import?, routing_mode?, …)` | bobine converter; `md_path` is transient when auto-importing — verify by searching, not by reading the path |
|
|
156
|
+
| `thoughts` | `ingest_thoughts(text, topic, tags?)` | cheapest, highest value: persist session reasoning with a stable topic scheme (`auth-refactor`, `api-design`) |
|
|
157
|
+
|
|
158
|
+
Frontmatter that matters: `title`, `type`, `tags`, `aliases: [...]` (wikilink
|
|
159
|
+
names), `id:` (stable identity, preserved as `uid`, written back on export),
|
|
160
|
+
`reviewed: true` (immune to `doctor --fix`). Credentials in `resource:`
|
|
161
|
+
URIs are stripped to `***@` at parse time (graphs imported before 0.2.3
|
|
162
|
+
keep old values until re-import). Rule of thumb: thoughts > md >
|
|
163
|
+
pdf — reasoning you already hold beats re-extracting it from files.
|
|
164
|
+
|
|
165
|
+
---
|
|
166
|
+
|
|
167
|
+
## CLI reference
|
|
168
|
+
|
|
169
|
+
Five verbs mirror the MCP tools; maintenance commands cover the rest. Global
|
|
170
|
+
flags (`--db`, `--bundle`, `--dim`, …) are documented once in `okf --help`,
|
|
171
|
+
accepted everywhere, and usually live in `okfgraph.toml`.
|
|
172
|
+
|
|
173
|
+
| Command | Description |
|
|
174
|
+
|---|---|
|
|
175
|
+
| `okf search QUERY [--target concepts\|chunks\|images] [--rank none\|hub\|ppr] [--expand] [--hub-rerank]` | Concepts (RRF hybrid), chunks (passages), images; `--rank ppr` is model-free |
|
|
176
|
+
| `okf read ID [--include body\|chunks\|document\|context] [--max-tokens N]` | Full shapes uncapped; budgeted section list with `--max-tokens` |
|
|
177
|
+
| `okf traverse [ID] [--relationship CONTAINS\|LINKS_TO\|PART_OF\|INCLUDES_ASSET] [--direction …] [--target ID]` | Relationships, directory listing (empty ID = root), shortest path via `--target` |
|
|
178
|
+
| `okf ingest --kind md\|pdf\|thoughts …` | `--md-file`, `--pdf-file` (`--routing-mode`, `--auto-import`), `--thoughts --topic` |
|
|
179
|
+
| `okf export --all\|--concept-id ID --output DIR [--flavor okf\|obsidian]` | Bundle export; obsidian = `[[Title]]` links, no index files |
|
|
180
|
+
| `okf diff [OLD] [NEW] [--json]` | Snapshot (two dirs, no model) or drift (graph vs dir); exit 0 identical / 1 different |
|
|
181
|
+
| `okf lint [DIR] [--json]` | Pre-import gate (no DB, no model); exit 0 clean / 1 errors / 2 bad dir |
|
|
182
|
+
| `okf doctor [--fix] [--strict] [--stale-days N] [--json]` | Score + findings; `--fix` repairs safely, `--strict` exits 1 on any finding |
|
|
183
|
+
| `okf import [--all] [--purge] [--mode text\|optional\|omni]` | Bulk/single import, delta-aware |
|
|
184
|
+
| `okf init`, `okf model-info`, `okf shell`, `okf reindex`, `okf broken-links`, `okf repair-links`, `okf deleted-*` | Setup, cache inspection, REPL, index rebuild, link + soft-delete maintenance |
|
|
185
|
+
|
|
186
|
+
Every CLI call cold-boots the router (model load ~30s when the embedder is
|
|
187
|
+
needed) — batch reads, and reach for `--rank ppr` for topic queries in cold
|
|
188
|
+
sessions.
|
|
189
|
+
|
|
190
|
+
---
|
|
191
|
+
|
|
192
|
+
## MCP server
|
|
193
|
+
|
|
194
|
+
5 tools, MCP ≥ 2.0, stdio transport. Wire once per project with a **stable**
|
|
195
|
+
`--db-path` + `--bundle` (a temp dir means amnesia every session); first boot
|
|
196
|
+
creates the schema; pre-approve the `ingest` / `export_bundle` write tools for
|
|
197
|
+
knowledge workflows. `.mcp.json` ships a ready config (`uv run --project .
|
|
198
|
+
okf-mcp`).
|
|
199
|
+
|
|
200
|
+
```bash
|
|
201
|
+
okf-mcp --db-path ./kb.db --bundle-root ./kb --embedding-dim 1024
|
|
202
|
+
```
|
|
203
|
+
|
|
204
|
+
| Tool | Routing |
|
|
205
|
+
|---|---|
|
|
206
|
+
| `search(query, target?, rank?, …)` | `concepts` = open questions; `chunks` = exact passages (`expand`, `hub_rerank`); `images` = assets; `rank=ppr` = model-free topic search |
|
|
207
|
+
| `read(concept_id, include?, max_tokens?)` | Full shapes, or a budgeted section list |
|
|
208
|
+
| `traverse(start_id, relationship?, direction?, target?, …)` | Walk edges, list dirs, connect two concepts |
|
|
209
|
+
| `ingest(kind, …)` | `md` / `pdf` / `thoughts` with per-kind params |
|
|
210
|
+
| `export_bundle(output_dir, …, flavor?)` | `okf` or `obsidian` |
|
|
211
|
+
|
|
212
|
+
Verify wiring with any `search` — an empty graph returns `[]`, which still
|
|
213
|
+
proves the plumbing works.
|
|
214
|
+
|
|
215
|
+
### Skills
|
|
216
|
+
|
|
217
|
+
Harness-neutral skill sources live in `skills/` (Claude Code auto-loads
|
|
218
|
+
`.claude/skills/<name>/SKILL.md`; other harnesses have their own dirs — see
|
|
219
|
+
`docs/harness-integration.md`):
|
|
220
|
+
|
|
221
|
+
- **`okfgraph-mcp`** — finding knowledge via MCP tools (pick this *or* CLI).
|
|
222
|
+
- **`okfgraph-cli`** — same graph through `okf` shell commands.
|
|
223
|
+
- **`okfgraph-ingest`** — feeding discipline: what to store, which kind,
|
|
224
|
+
converter modes, wikilink/alias conventions. Install alongside either
|
|
225
|
+
finder when the agent should persist knowledge.
|
|
226
|
+
|
|
227
|
+
---
|
|
228
|
+
|
|
229
|
+
## Architecture
|
|
230
|
+
|
|
231
|
+
```
|
|
232
|
+
┌──────────── MCP (5 tools) / CLI (5 verbs + maintenance)
|
|
233
|
+
│ │
|
|
234
|
+
OKFRouter (facade: owns conn, encoder, lock; thin proxies)
|
|
235
|
+
│
|
|
236
|
+
┌────────────┼──────────────────────────────────────────┐
|
|
237
|
+
│ │ │
|
|
238
|
+
Import ──► Embed ──► Search ◄── ranking ──► Export ──► Doctor/Diff
|
|
239
|
+
│ │ │ (seeds+PPR) │ │
|
|
240
|
+
│ │ │ links.py │
|
|
241
|
+
Delta ──► Schema ──► ImageAssets ──► Purge ──► Ingest ──► Converters
|
|
242
|
+
│ (bobine seam)
|
|
243
|
+
└──────────────── LadybugDB (graph + vector + FTS, one file) ──┘
|
|
244
|
+
▲
|
|
245
|
+
Rust okf-embed (Jina v5, ORT) — tokenize, last-token
|
|
246
|
+
pool, L2-norm, Matryoshka truncate; numerics pinned
|
|
247
|
+
by tests/test_parity.py
|
|
248
|
+
```
|
|
249
|
+
|
|
250
|
+
Components live in `okfgraph/components/` (one concern each, dependencies
|
|
251
|
+
injected): `search`, `ranking` (pure seeds + exact PPR), `links` (pure
|
|
252
|
+
extraction + name index), `import_` (batch/single upsert, delta-aware),
|
|
253
|
+
`export` (okf/obsidian flavors), `diff` (structural compare), `doctor`
|
|
254
|
+
(scored scan + safe fixes), `embedding` (Rust bridge + chunking),
|
|
255
|
+
`converters` (`DocumentConverter` protocol + `BobineConverter`),
|
|
256
|
+
`image_assets`, `delta`, `purge`, `schema`, `ingest`.
|
|
257
|
+
|
|
258
|
+
Key design decisions:
|
|
259
|
+
|
|
260
|
+
- **Rust-only embeddings, fail-fast** — no Python fallback; a mid-run stack
|
|
261
|
+
switch would silently mix vector spaces in one index.
|
|
262
|
+
- **Lazy session init** — router construction never downloads the model or
|
|
263
|
+
builds the ONNX session; first encode opens once (thread-safe), token
|
|
264
|
+
counting uses a tokenizer-only handle, so PPR search, budgeted reads,
|
|
265
|
+
diff, and doctor stay cold.
|
|
266
|
+
- **Last-token pooling** (not mean) — required by Jina v5; mean pooling
|
|
267
|
+
breaks alignment with omni image embeddings.
|
|
268
|
+
- **Single pinned ORT** (`onnxruntime==1.29.0`, `ORT_DYLIB_PATH`-overridable)
|
|
269
|
+
shared by bobine + okf-embed.
|
|
270
|
+
- **Bobine is a plugin, not a dependency** — `DocumentConverter.convert()`
|
|
271
|
+
is the seam; provider owns its options (`routing_mode`,
|
|
272
|
+
`extract_images`); missing bobine → clear `RuntimeError`, never a silent
|
|
273
|
+
fallback path.
|
|
274
|
+
- **Deterministic by construction** — sorted traversal order, fixed PPR
|
|
275
|
+
constants and accumulation order, golden fixtures under `tests/fixtures/`
|
|
276
|
+
(PPR scores, diff reports, doctor scores, vault round-trips).
|
|
277
|
+
- **Bundle is source of truth, graph is the index** — edit markdown,
|
|
278
|
+
re-import; verify ingests by searching, since empty `[]` proves wiring
|
|
279
|
+
while errors prove broken setup.
|
|
280
|
+
- **MCP-first surface** — CLI mirrors the same 5 verbs; maintenance
|
|
281
|
+
(`diff`, `doctor`) is CLI-only to keep the agent tool surface at 5.
|
|
282
|
+
|
|
283
|
+
---
|
|
284
|
+
|
|
285
|
+
## API reference (`OKFRouter` proxies)
|
|
286
|
+
|
|
287
|
+
| Method | Description |
|
|
288
|
+
|---|---|
|
|
289
|
+
| `import_from_okf(file_path, mode?)` / `import_mgr.import_bundle(dir?, batch_size?, mode?, purge_deleted?)` | Single / bulk import |
|
|
290
|
+
| `ingest_mgr.ingest_md / ingest_pdf / ingest_thoughts` | Kind-based ingestion with lint + chunk + embed + link |
|
|
291
|
+
| `search_hybrid(query, …, rank="none"\|"hub"\|"ppr")` | RRF hybrid; `hub` blends authority, `ppr` is model-free |
|
|
292
|
+
| `search_engine.search_with_ppr / read_with_budget / search_chunks / traverse / find_path / get_chunks / get_by_id / list_directory` | Model-free PPR, budgeted reads, retrieval + navigation |
|
|
293
|
+
| `embed_engine.reconstruct_document(id)` / `count_tokens(text)` | ~98% chunk→markdown rebuild; Rust counter with chars/4 fallback |
|
|
294
|
+
| `export_mgr.export_bundle / export_to_okf(..., flavor="okf"\|"obsidian")` | Filtered export, both flavors |
|
|
295
|
+
| `diff_mgr.diff_dirs(old, new)` / `diff_db_dir(bundle_dir)` | Snapshot / drift structural reports |
|
|
296
|
+
| `diagnose(stale_days?)` / `doctor_fix()` | Health report / safe repairs |
|
|
297
|
+
| `list_broken_links()` / `repair_links(skip_sources?)` | Unresolved refs; exact-id + unique-name repair |
|
|
298
|
+
| `model_info(...)` / `default_cache_dir()` | Cache inspection without model load |
|
|
299
|
+
|
|
300
|
+
---
|
|
301
|
+
|
|
302
|
+
## Testing
|
|
303
|
+
|
|
304
|
+
```bash
|
|
305
|
+
uv run --project . pytest tests/ -q # full suite (model loads; takes a while)
|
|
306
|
+
uv run --project . pytest tests/test_ranking.py tests/test_mcp_server.py -q # fast subset
|
|
307
|
+
cd rust/okf-embed && cargo test --locked # Rust unit tests (pure, no model)
|
|
308
|
+
```
|
|
309
|
+
|
|
310
|
+
Golden fixtures under `tests/fixtures/` (`ppr_graph`, `diff_a`/`diff_b`,
|
|
311
|
+
`doctor_bundle`, `obsidian_vault`, `ppr_bundle`) lock deterministic behavior:
|
|
312
|
+
PPR scores byte-stable across runs, diff reports exact, doctor score pinned
|
|
313
|
+
(83 on the fixture bundle), obsidian export→import edge-identical. Live suites
|
|
314
|
+
reuse one class-scoped router each so the ONNX model loads once per class.
|
|
315
|
+
|
|
316
|
+
---
|
|
317
|
+
|
|
318
|
+
## Project structure
|
|
319
|
+
|
|
320
|
+
```text
|
|
321
|
+
okfgraph/
|
|
322
|
+
├── okfgraph/
|
|
323
|
+
│ ├── __init__.py # ConceptModel, OKFRouter, cli_main
|
|
324
|
+
│ ├── models.py # ConceptModel / ChunkModel / ImageAssetModel (extra frontmatter allowed)
|
|
325
|
+
│ ├── router.py # OKFRouter facade (owns resources, thin proxies)
|
|
326
|
+
│ ├── cli.py # okf: 5 verbs + maintenance, slim help, okfgraph.toml
|
|
327
|
+
│ ├── mcp_server.py # okf-mcp: 5 tools, MCP ≥ 2.0, lifespan-managed router
|
|
328
|
+
│ ├── config.py # okfgraph.toml + env + CLI merge
|
|
329
|
+
│ ├── images.py # IngestMode (text|optional|omni), planning helpers
|
|
330
|
+
│ ├── security.py # SSRF/domain guards for remote images
|
|
331
|
+
│ ├── tools.py # legacy tool definitions (superseded by mcp_server)
|
|
332
|
+
│ └── components/ # ranking, links, search, lint, import_, export, diff, doctor,
|
|
333
|
+
│ # embedding, converters, image_assets, delta, purge, schema, ingest
|
|
334
|
+
├── rust/okf-embed/ # Jina v5 Rust loader (ort, CUDA-opportunistic) + Python wheel
|
|
335
|
+
├── skills/ # okfgraph-mcp, okfgraph-cli, okfgraph-ingest (harness-neutral source)
|
|
336
|
+
├── tests/fixtures/ # conformance corpus: ppr, diff, doctor, obsidian, bundles
|
|
337
|
+
├── docs/ # converters, harness-integration, plan-retrieval-roundup, diagnostics…
|
|
338
|
+
├── .mcp.json # ready MCP wiring (uv run --project . okf-mcp)
|
|
339
|
+
├── architecture.md # long-form architecture spec
|
|
340
|
+
└── pyproject.toml # slim core deps + pdf/omni/dev extras
|
|
341
|
+
```
|
|
342
|
+
|
|
343
|
+
---
|
|
344
|
+
|
|
345
|
+
## Requirements
|
|
346
|
+
|
|
347
|
+
Core (`uv sync`): `ladybug==0.20.3`, `okf-embed` (built from `rust/okf-embed`),
|
|
348
|
+
`onnxruntime==1.29.0`, `mordant>=0.9`, `mcp>=2.0`, `pydantic>=2`, `pyyaml>=6`,
|
|
349
|
+
`numpy>=1.26`, `python-frontmatter>=1`, `fasteners>=0.19`. Python ≥ 3.11.
|
|
350
|
+
|
|
351
|
+
- `--extra pdf`: `bobine>=0.5` (default PDF converter).
|
|
352
|
+
- `--extra omni`: `sentence-transformers>=3`, `Pillow>=10` (image embeddings).
|
|
353
|
+
- `--extra dev`: `pytest>=8`.
|
|
354
|
+
|
|
355
|
+
---
|
|
356
|
+
|
|
357
|
+
## License
|
|
358
|
+
|
|
359
|
+
See LICENSE for details (Apache-2.0 OR MIT).
|
|
@@ -1,4 +1,4 @@
|
|
|
1
|
-
# OKFgraph 0.2.
|
|
1
|
+
# OKFgraph 0.2.7
|
|
2
2
|
|
|
3
3
|
[](https://pypi.org/project/okfgraph/)
|
|
4
4
|
[](https://www.python.org/)
|
|
@@ -232,6 +232,10 @@ Key design decisions:
|
|
|
232
232
|
|
|
233
233
|
- **Rust-only embeddings, fail-fast** — no Python fallback; a mid-run stack
|
|
234
234
|
switch would silently mix vector spaces in one index.
|
|
235
|
+
- **Lazy session init** — router construction never downloads the model or
|
|
236
|
+
builds the ONNX session; first encode opens once (thread-safe), token
|
|
237
|
+
counting uses a tokenizer-only handle, so PPR search, budgeted reads,
|
|
238
|
+
diff, and doctor stay cold.
|
|
235
239
|
- **Last-token pooling** (not mean) — required by Jina v5; mean pooling
|
|
236
240
|
breaks alignment with omni image embeddings.
|
|
237
241
|
- **Single pinned ORT** (`onnxruntime==1.29.0`, `ORT_DYLIB_PATH`-overridable)
|
|
@@ -13,7 +13,7 @@ stubs (``...``) until their respective phase moves the implementation over.
|
|
|
13
13
|
from okfgraph.components.schema import SchemaManager
|
|
14
14
|
from okfgraph.components.delta import DeltaDetector
|
|
15
15
|
from okfgraph.components.purge import PurgeManager
|
|
16
|
-
from okfgraph.components.embedding import EmbeddingEngine
|
|
16
|
+
from okfgraph.components.embedding import EmbeddingEngine, LazyRustEncoder
|
|
17
17
|
from okfgraph.components.image_assets import ImageAssetManager
|
|
18
18
|
from okfgraph.components.search import SearchEngine
|
|
19
19
|
from okfgraph.components.import_ import ImportManager, parse_source_file
|
|
@@ -37,6 +37,7 @@ __all__ = [
|
|
|
37
37
|
"DeltaDetector",
|
|
38
38
|
"PurgeManager",
|
|
39
39
|
"EmbeddingEngine",
|
|
40
|
+
"LazyRustEncoder",
|
|
40
41
|
"ImageAssetManager",
|
|
41
42
|
"SearchEngine",
|
|
42
43
|
"ImportManager",
|
|
@@ -6,33 +6,250 @@ here. Public callers reach these via router.<method> (component bridge).
|
|
|
6
6
|
"""
|
|
7
7
|
import logging
|
|
8
8
|
import math
|
|
9
|
+
import threading
|
|
9
10
|
from pathlib import Path
|
|
10
11
|
|
|
11
12
|
import mordant
|
|
12
13
|
from typing import Any, Dict, List, Optional
|
|
13
14
|
logger = logging.getLogger(__name__)
|
|
14
15
|
|
|
15
|
-
|
|
16
|
-
|
|
16
|
+
_ORT_MODULE_NAMES = ("onnxruntime", "onnxruntime-gpu")
|
|
17
|
+
|
|
18
|
+
|
|
19
|
+
def _candidate_ort_library_names(os_name=None, sys_platform=None):
|
|
20
|
+
"""Return the ORT library filenames for an OS/platform pair.
|
|
21
|
+
|
|
22
|
+
Takes explicit arguments so unit tests can cover Windows/macOS/Linux
|
|
23
|
+
from any host. Linux uses a versioned ``libonnxruntime.so.*`` glob
|
|
24
|
+
(handled by :func:`_find_runtime_in_package`); this helper returns the
|
|
25
|
+
unversioned fallback name for that platform.
|
|
26
|
+
"""
|
|
27
|
+
import os
|
|
28
|
+
import sys
|
|
29
|
+
if os_name is None:
|
|
30
|
+
os_name = os.name
|
|
31
|
+
if sys_platform is None:
|
|
32
|
+
sys_platform = sys.platform
|
|
33
|
+
if os_name == "nt":
|
|
34
|
+
return ("onnxruntime.dll",)
|
|
35
|
+
if sys_platform == "darwin":
|
|
36
|
+
return ("libonnxruntime.dylib",)
|
|
37
|
+
return ("libonnxruntime.so",)
|
|
38
|
+
|
|
39
|
+
|
|
40
|
+
def _find_runtime_in_package(package_dir, os_name=None, sys_platform=None):
|
|
41
|
+
"""Find an ORT shared library inside a pip-installed package directory.
|
|
42
|
+
|
|
43
|
+
Returns a :class:`Path` or ``None``. Prefers versioned Linux libraries
|
|
44
|
+
(``libonnxruntime.so.*``) over the unversioned ``libonnxruntime.so``.
|
|
45
|
+
"""
|
|
46
|
+
import os
|
|
47
|
+
import sys
|
|
48
|
+
if os_name is None:
|
|
49
|
+
os_name = os.name
|
|
50
|
+
if sys_platform is None:
|
|
51
|
+
sys_platform = sys.platform
|
|
52
|
+
capi_dir = Path(package_dir) / "capi"
|
|
53
|
+
if not capi_dir.is_dir():
|
|
54
|
+
return None
|
|
55
|
+
if os_name == "nt" or sys_platform == "darwin":
|
|
56
|
+
candidate = capi_dir / _candidate_ort_library_names(os_name, sys_platform)[0]
|
|
57
|
+
return candidate if candidate.is_file() else None
|
|
58
|
+
versioned = sorted(capi_dir.glob("libonnxruntime.so.*"))
|
|
59
|
+
if versioned:
|
|
60
|
+
return versioned[-1]
|
|
61
|
+
candidate = capi_dir / "libonnxruntime.so"
|
|
62
|
+
return candidate if candidate.is_file() else None
|
|
63
|
+
|
|
64
|
+
|
|
65
|
+
def _configure_windows_ort_dll_directory(runtime_path) -> bool:
|
|
66
|
+
"""Add the ORT ``capi`` directory to Windows DLL resolution.
|
|
67
|
+
|
|
68
|
+
Best-effort only: returns ``False`` (never raises) when unavailable.
|
|
69
|
+
"""
|
|
70
|
+
import os
|
|
71
|
+
if os.name != "nt":
|
|
72
|
+
return False
|
|
73
|
+
add_dll_directory = getattr(os, "add_dll_directory", None)
|
|
74
|
+
if not callable(add_dll_directory):
|
|
75
|
+
return False
|
|
76
|
+
try:
|
|
77
|
+
add_dll_directory(str(Path(runtime_path).parent))
|
|
78
|
+
return True
|
|
79
|
+
except OSError as exc:
|
|
80
|
+
logger.debug("ORT DLL directory configuration failed: %s", exc)
|
|
81
|
+
return False
|
|
82
|
+
|
|
83
|
+
|
|
84
|
+
def _warm_ort_gpu_dlls(module) -> bool:
|
|
85
|
+
"""Preload NVIDIA DLLs when the installed ORT package exposes CUDA.
|
|
86
|
+
|
|
87
|
+
This mirrors bobine's GPU bootstrap: only attempt the preload when the
|
|
88
|
+
package reports a CUDA execution provider, and never let it fail
|
|
89
|
+
runtime discovery.
|
|
90
|
+
"""
|
|
91
|
+
preload = getattr(module, "preload_dlls", None)
|
|
92
|
+
if not callable(preload):
|
|
93
|
+
return False
|
|
94
|
+
try:
|
|
95
|
+
providers = list(module.get_available_providers())
|
|
96
|
+
except Exception as exc:
|
|
97
|
+
logger.debug("ORT provider query failed: %s", exc)
|
|
98
|
+
return False
|
|
99
|
+
if "CUDAExecutionProvider" not in providers:
|
|
100
|
+
return False
|
|
101
|
+
try:
|
|
102
|
+
preload()
|
|
103
|
+
return True
|
|
104
|
+
except Exception as exc:
|
|
105
|
+
logger.debug("ORT GPU DLL preload failed: %s", exc)
|
|
106
|
+
return False
|
|
107
|
+
|
|
108
|
+
|
|
109
|
+
def resolve_ort_dylib(*, warm_gpu: bool = True, os_name=None, sys_platform=None) -> Optional[str]:
|
|
110
|
+
"""Point ``ORT_DYLIB_PATH`` at a pip-installed ORT build when unset.
|
|
17
111
|
|
|
18
112
|
Both bobine and okf-embed load ONNX Runtime dynamically; sharing one
|
|
19
113
|
binary avoids version/CUDA drift between the two runtimes. Explicit
|
|
20
114
|
user configuration always wins — this only fills the gap.
|
|
115
|
+
|
|
116
|
+
Resolution order:
|
|
117
|
+
|
|
118
|
+
1. Existing ``ORT_DYLIB_PATH``.
|
|
119
|
+
2. Pip-installed ``onnxruntime`` or ``onnxruntime-gpu`` package.
|
|
120
|
+
3. OS loader path (represented by returning ``None``).
|
|
121
|
+
|
|
122
|
+
Missing runtimes never raise here; session creation or encoding is the
|
|
123
|
+
fail-fast boundary. ``os_name``/``sys_platform`` are explicit so unit
|
|
124
|
+
tests can cover every platform from any host (patching ``os.name``
|
|
125
|
+
would break ``pathlib`` dispatch on Windows).
|
|
21
126
|
"""
|
|
22
127
|
import os
|
|
23
|
-
|
|
24
|
-
|
|
25
|
-
|
|
26
|
-
|
|
27
|
-
|
|
28
|
-
|
|
29
|
-
|
|
30
|
-
|
|
31
|
-
|
|
32
|
-
|
|
128
|
+
explicit = os.environ.get("ORT_DYLIB_PATH")
|
|
129
|
+
if explicit:
|
|
130
|
+
return explicit
|
|
131
|
+
for module_name in _ORT_MODULE_NAMES:
|
|
132
|
+
try:
|
|
133
|
+
module = __import__(module_name)
|
|
134
|
+
except ImportError:
|
|
135
|
+
continue
|
|
136
|
+
try:
|
|
137
|
+
package_dir = Path(str(module.__file__)).parent
|
|
138
|
+
except Exception:
|
|
139
|
+
continue
|
|
140
|
+
found = _find_runtime_in_package(package_dir, os_name, sys_platform)
|
|
141
|
+
if found is None:
|
|
142
|
+
continue
|
|
143
|
+
os.environ["ORT_DYLIB_PATH"] = str(found)
|
|
144
|
+
_configure_windows_ort_dll_directory(found)
|
|
145
|
+
if warm_gpu:
|
|
146
|
+
_warm_ort_gpu_dlls(module)
|
|
147
|
+
logger.debug("resolved ORT runtime: %s", found)
|
|
148
|
+
return str(found)
|
|
33
149
|
return None
|
|
34
150
|
|
|
35
151
|
|
|
152
|
+
class LazyRustEncoder:
|
|
153
|
+
"""Defers the ONNX session open until the first real encode.
|
|
154
|
+
|
|
155
|
+
Router construction stays cheap: the ``okf-embed`` wheel import is still
|
|
156
|
+
validated eagerly (fail fast on a missing install), but ``JinaV5.open``
|
|
157
|
+
— model download + session build — waits for the first ``encode`` /
|
|
158
|
+
``encode_batch`` / ``used_cuda`` access. Token counting uses the
|
|
159
|
+
separate lightweight ``JinaTokenizer`` handle, so budgeted reads and the
|
|
160
|
+
context-window guard stay cold too.
|
|
161
|
+
|
|
162
|
+
A failed session open is cached and re-raised: configuration errors stay
|
|
163
|
+
fail-fast (once, at first encode) instead of retrying network/model
|
|
164
|
+
acquisition on every call. Thread-safe: concurrent first encodes open
|
|
165
|
+
exactly one session.
|
|
166
|
+
"""
|
|
167
|
+
|
|
168
|
+
def __init__(self, *, model_id, truncate_dim, device,
|
|
169
|
+
session_factory, tokenizer_factory, on_open=None):
|
|
170
|
+
self._model_id = model_id
|
|
171
|
+
self._truncate_dim = truncate_dim
|
|
172
|
+
self._device = device
|
|
173
|
+
self._session_factory = session_factory
|
|
174
|
+
self._tokenizer_factory = tokenizer_factory
|
|
175
|
+
self._on_open = on_open
|
|
176
|
+
self._lock = threading.Lock()
|
|
177
|
+
self._encoder = None
|
|
178
|
+
self._encoder_error = None
|
|
179
|
+
self._open_reported = False
|
|
180
|
+
self._tokenizer = None
|
|
181
|
+
|
|
182
|
+
@property
|
|
183
|
+
def is_loaded(self) -> bool:
|
|
184
|
+
"""True once the ONNX session has been opened."""
|
|
185
|
+
return self._encoder is not None
|
|
186
|
+
|
|
187
|
+
@property
|
|
188
|
+
def model_id(self) -> str:
|
|
189
|
+
return self._model_id
|
|
190
|
+
|
|
191
|
+
@property
|
|
192
|
+
def dim(self) -> int:
|
|
193
|
+
return self._truncate_dim
|
|
194
|
+
|
|
195
|
+
@property
|
|
196
|
+
def used_cuda(self) -> bool:
|
|
197
|
+
"""Effective device — opens the session on first access."""
|
|
198
|
+
return bool(self._get_encoder().used_cuda)
|
|
199
|
+
|
|
200
|
+
def encode(self, text: str, task: str = "Document"):
|
|
201
|
+
return self._get_encoder().encode(text, task=task)
|
|
202
|
+
|
|
203
|
+
def encode_batch(self, texts, task: str = "Document"):
|
|
204
|
+
return self._get_encoder().encode_batch(texts, task=task)
|
|
205
|
+
|
|
206
|
+
def count_tokens(self, text: str) -> int:
|
|
207
|
+
"""Exact count via the tokenizer-only handle (never opens the session)."""
|
|
208
|
+
return int(self._get_tokenizer().count_tokens(text))
|
|
209
|
+
|
|
210
|
+
def _get_encoder(self):
|
|
211
|
+
encoder = self._encoder
|
|
212
|
+
if encoder is not None:
|
|
213
|
+
return encoder
|
|
214
|
+
with self._lock:
|
|
215
|
+
if self._encoder is not None:
|
|
216
|
+
return self._encoder
|
|
217
|
+
if self._encoder_error is not None:
|
|
218
|
+
raise self._encoder_error
|
|
219
|
+
try:
|
|
220
|
+
encoder = self._session_factory()
|
|
221
|
+
except Exception as exc:
|
|
222
|
+
self._encoder_error = exc
|
|
223
|
+
raise
|
|
224
|
+
self._encoder = encoder
|
|
225
|
+
if self._on_open is not None and not self._open_reported:
|
|
226
|
+
self._open_reported = True
|
|
227
|
+
self._on_open(encoder)
|
|
228
|
+
return encoder
|
|
229
|
+
|
|
230
|
+
def _get_tokenizer(self):
|
|
231
|
+
tokenizer = self._tokenizer
|
|
232
|
+
if tokenizer is not None:
|
|
233
|
+
return tokenizer
|
|
234
|
+
with self._lock:
|
|
235
|
+
if self._tokenizer is not None:
|
|
236
|
+
return self._tokenizer
|
|
237
|
+
try:
|
|
238
|
+
tokenizer = self._tokenizer_factory()
|
|
239
|
+
except Exception as exc:
|
|
240
|
+
logger.debug("tokenizer-only open failed: %s", exc)
|
|
241
|
+
raise
|
|
242
|
+
self._tokenizer = tokenizer
|
|
243
|
+
return tokenizer
|
|
244
|
+
|
|
245
|
+
def __repr__(self) -> str:
|
|
246
|
+
state = "loaded" if self._encoder is not None else "cold"
|
|
247
|
+
return (
|
|
248
|
+
f"LazyRustEncoder({self._model_id}, "
|
|
249
|
+
f"dim={self._truncate_dim}, {state})"
|
|
250
|
+
)
|
|
251
|
+
|
|
252
|
+
|
|
36
253
|
class EmbeddingEngine:
|
|
37
254
|
"""Owns the embedding model and chunking logic.
|
|
38
255
|
|
|
@@ -43,7 +260,8 @@ class EmbeddingEngine:
|
|
|
43
260
|
|
|
44
261
|
def __init__(self, rust_encoder, embedding_dim, device,
|
|
45
262
|
cache_dir, model_id, omni_model_id, omni,
|
|
46
|
-
chunk_size, chunk_overlap, enable_chunking, conn
|
|
263
|
+
chunk_size, chunk_overlap, enable_chunking, conn,
|
|
264
|
+
ort_dylib=None):
|
|
47
265
|
self.encoder = rust_encoder
|
|
48
266
|
self.embedding_dim = embedding_dim
|
|
49
267
|
self.device = device
|
|
@@ -55,6 +273,7 @@ class EmbeddingEngine:
|
|
|
55
273
|
self.chunk_overlap = chunk_overlap
|
|
56
274
|
self.enable_chunking = enable_chunking
|
|
57
275
|
self.conn = conn
|
|
276
|
+
self.ort_dylib = ort_dylib
|
|
58
277
|
|
|
59
278
|
def _encode(self, text: str, task: str = "Document") -> List[float]:
|
|
60
279
|
"""Encode text with the Rust Jina v5 encoder.
|