Add Python bindings via PyO3/maturin

Cargo.toml

@@ -36,6 +36,7 @@ tempfile = "3.20.0"
 
 napi = { version = "3.2.2", optional = true }
 napi-derive = { version = "3.2.2", optional = true }
+pyo3 = { version = "0.23", optional = true, features = ["extension-module", "abi3-py38"] }
 once_cell = "1.21.3"
 lru = "0.18.0"
 tracing = "0.1.41"
@@ -84,6 +85,7 @@ napi-build = "2.2.3"
 #default = ["t5-quantized"]
 #default = ["t5-openvino"]
 napi = ["dep:napi", "dep:napi-derive"]
+python = ["dep:pyo3"]
 t5-quantized = []
 t5-openvino = ["dep:openvino"]
 embed-assets = []

Makefile

@@ -11,12 +11,14 @@ ifeq ($(UNAME_S),Darwin)
     # Apple Silicon: Metal GPU + Accelerate BLAS
     CLI_FEATURES := t5-quantized,metal,progress
     NAPI_FEATURES := t5-quantized,metal,napi
+    PYTHON_FEATURES := t5-quantized,metal,python
     RUSTFLAGS_EXTRA :=
     TARGET := aarch64-apple-darwin
   else
     # Intel Mac: CPU-only with FBGEMM + hybrid-dequant
     CLI_FEATURES := t5-quantized,fbgemm,hybrid-dequant,progress
     NAPI_FEATURES := t5-quantized,fbgemm,hybrid-dequant,napi
+    PYTHON_FEATURES := t5-quantized,fbgemm,hybrid-dequant,python
     RUSTFLAGS_EXTRA := -C target-feature=+avx2,+fma
     TARGET := x86_64-apple-darwin
   endif
@@ -26,9 +28,17 @@ else ifeq ($(UNAME_S),Linux)
   ifneq ($(NVCC),)
     CLI_FEATURES := t5-quantized,cuda,progress
     NAPI_FEATURES := t5-quantized,cuda,napi
+    PYTHON_FEATURES := t5-quantized,cuda,python
+  else ifeq ($(UNAME_M),aarch64)
+    # Linux ARM (Graviton, Pi, Ampere): fbgemm/hybrid-dequant are x86-only
+    CLI_FEATURES := t5-quantized,progress
+    NAPI_FEATURES := t5-quantized,napi
+    PYTHON_FEATURES := t5-quantized,python
   else
+    # Linux x86_64 CPU-only
     CLI_FEATURES := t5-quantized,fbgemm,hybrid-dequant,progress
     NAPI_FEATURES := t5-quantized,fbgemm,hybrid-dequant,napi
+    PYTHON_FEATURES := t5-quantized,fbgemm,hybrid-dequant,python
   endif
   PICKBRAIN_FEATURES := $(CLI_FEATURES),embed-assets
   RUSTFLAGS_EXTRA :=
@@ -163,6 +173,14 @@ run: module
 	ln -sf target/release/warp-macos-universal.node warp.node
 	node index.cjs
 
+python-wheel: prereqs download
+	@command -v maturin >/dev/null 2>&1 || { echo "maturin not found — run: pip install maturin" >&2; exit 1; }
+	maturin build --release $(BUILD_TARGET) --features $(PYTHON_FEATURES)
+
+python-dev: prereqs download
+	@command -v maturin >/dev/null 2>&1 || { echo "maturin not found — run: pip install maturin" >&2; exit 1; }
+	maturin develop --features $(PYTHON_FEATURES)
+
 distclean:
 	rm -rf target env html xtr-base-en openvino_model
 	rm -f warp-cli warp.node pickbrain Cargo.lock lcov.info output.txt
@@ -179,4 +197,4 @@ distclean:
 	fi
 	rm -rf .make-stamps .stamp* stamp-* *.stamp */.stamp* */*.stamp */stamp-*
 
-.PHONY: prereqs download ovdownload build buildemb warp-cli pickbrain pickbrain-install module macintel winintel win test bench nfcorpus nfcorpus-score run distclean
+.PHONY: prereqs download ovdownload build buildemb warp-cli pickbrain pickbrain-install module macintel winintel win test bench nfcorpus nfcorpus-score run python-wheel python-dev distclean

README.md

@@ -97,13 +97,82 @@ Other flags:
 - `fbgemm` -- fbgemm-rs packed GEMM (bf16 weights, faster on x86)
 - `hybrid-dequant` -- F32 attention + Q4K FFN with fused gated-gelu (x86, requires `fbgemm`)
 - `napi` -- Node.js native module via napi-rs
+- `python` -- Python extension module via PyO3/maturin
 - `embed-assets` -- bake weights into binary
 - `progress` -- progress bars for CLI
 
 Platform-specific recommended features (these are what `make` uses automatically):
 - **Apple Silicon**: `t5-quantized,metal`
 - **Intel Mac (x86_64)**: `t5-quantized,fbgemm,hybrid-dequant`
 - **Intel Windows (x86_64)**: `t5-openvino,fbgemm`
+- **Linux x86_64 (CPU)**: `t5-quantized,fbgemm,hybrid-dequant`
+- **Linux x86_64 (CUDA)**: `t5-quantized,cuda`
+- **Linux ARM (Graviton, Pi, Ampere)**: `t5-quantized` (`fbgemm`/`hybrid-dequant` are x86-only)
+
+## Using as a Python module ##
+
+Requires [maturin](https://github.com/PyO3/maturin):
+
+```
+pip install maturin
+```
+
+Build a wheel for your current platform (auto-selects the right backend features):
+
+```
+make python-wheel
+pip install target/wheels/witchcraft-*.whl
+```
+
+Or, inside a virtualenv, install directly for development:
+
+```
+make python-dev
+```
+
+The Makefile picks the correct feature set automatically (`metal` on Apple Silicon,
+`fbgemm,hybrid-dequant` on Intel, `cuda` on Linux with a GPU). The wheel works on
+Python 3.8+, including 3.14+.
+
+To run the Python test suite after installing the wheel:
+
+```
+pip install pytest
+pytest tests/test_witchcraft.py
+```
+
+Tests that exercise the embedder (search, score, add, etc.) are skipped unless
+model assets are present. Run `make download` first to enable them.
+
+Then in Python:
+
+```python
+import witchcraft
+
+wc = witchcraft.Witchcraft('/path/to/db.sqlite', '/path/to/assets')
+
+# Add documents (fire-and-forget, processed by background thread)
+wc.add('550e8400-e29b-41d4-a716-446655440000',
+       '2024-01-15T10:00:00Z',
+       '{"source": "dropbox"}',
+       'The document text goes here')
+
+# Trigger embedding + index build
+wc.index()
+
+# Hybrid semantic + BM25 search
+results = wc.search('does milk intake cause acne?', threshold=0.3, top_k=5)
+for r in results:
+    print(r['score'], r['body'])
+
+# Score individual sentences against a query
+scores = wc.score('acne and diet', ['milk causes acne', 'exercise helps skin'])
+
+# Shut down the background indexer cleanly
+wc.shutdown()
+```
+
+`search` returns a list of dicts with keys: `score`, `metadata`, `body`, `idx`, `date`.
 
 ## Using as Node module ##
 

build.rs

@@ -1,3 +1,5 @@
 fn main() {
-    napi_build::setup(); // enables N-API compatibility
+    if std::env::var("CARGO_FEATURE_NAPI").is_ok() {
+        napi_build::setup();
+    }
 }

pyproject.toml

@@ -0,0 +1,14 @@
+[build-system]
+requires = ["maturin>=1.0,<2.0"]
+build-backend = "maturin"
+
+[project]
+name = "witchcraft"
+version = "0.1.2"
+description = "Semantic search with hybrid dense/sparse retrieval"
+requires-python = ">=3.8"
+license = {text = "Apache-2.0"}
+
+[tool.maturin]
+features = ["python"]
+module-name = "witchcraft"

src/lib.rs

@@ -64,6 +64,9 @@ use sql_generator::build_filter_sql_and_params;
 #[allow(dead_code)]
 mod napi;
 
+#[cfg(feature = "python")]
+mod python;
+
 use anyhow::Result;
 use candle_core::{DType, Device, IndexOp, Tensor, D};