passageway 0.3.0__py3-none-any.whl
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.
- passageway/__init__.py +0 -0
- passageway/chunk.py +68 -0
- passageway/config.py +315 -0
- passageway/embed.py +213 -0
- passageway/embed_command.py +258 -0
- passageway/eval/__init__.py +1 -0
- passageway/eval/__main__.py +5 -0
- passageway/eval/cases.py +93 -0
- passageway/eval/cli.py +123 -0
- passageway/eval/compare.py +138 -0
- passageway/eval/measure.py +81 -0
- passageway/eval/report.py +282 -0
- passageway/eval/run.py +121 -0
- passageway/index.py +248 -0
- passageway/init.py +60 -0
- passageway/models.py +202 -0
- passageway/protocol.py +229 -0
- passageway/rank.py +153 -0
- passageway/render.py +67 -0
- passageway/server.py +401 -0
- passageway/sessions.py +39 -0
- passageway/skip.py +34 -0
- passageway/sync.py +128 -0
- passageway/testing/__init__.py +0 -0
- passageway/testing/contract.py +17 -0
- passageway/testing/embedder_contract.py +37 -0
- passageway/testing/fake_connector.py +37 -0
- passageway/testing/fake_embedder.py +102 -0
- passageway-0.3.0.dist-info/METADATA +73 -0
- passageway-0.3.0.dist-info/RECORD +34 -0
- passageway-0.3.0.dist-info/WHEEL +4 -0
- passageway-0.3.0.dist-info/entry_points.txt +2 -0
- passageway-0.3.0.dist-info/licenses/LICENSE +75 -0
- passageway-0.3.0.dist-info/licenses/LICENSE-FREE-TRIAL +57 -0
passageway/__init__.py
ADDED
|
File without changes
|
passageway/chunk.py
ADDED
|
@@ -0,0 +1,68 @@
|
|
|
1
|
+
"""Turns an item into the pieces that are embedded and returned as passages."""
|
|
2
|
+
|
|
3
|
+
import re
|
|
4
|
+
from dataclasses import dataclass
|
|
5
|
+
|
|
6
|
+
from passageway.protocol import Item
|
|
7
|
+
|
|
8
|
+
MAX_CHARS = 1500 # a section over this is split; the embedding model cuts at a few hundred tokens anyway
|
|
9
|
+
|
|
10
|
+
|
|
11
|
+
@dataclass(frozen=True)
|
|
12
|
+
class Chunk:
|
|
13
|
+
heading: str
|
|
14
|
+
text: str
|
|
15
|
+
ordinal: int
|
|
16
|
+
|
|
17
|
+
|
|
18
|
+
def _cut_words(paragraph: str, cap: int) -> list[str]:
|
|
19
|
+
out, line = [], ""
|
|
20
|
+
for word in paragraph.split():
|
|
21
|
+
while len(word) > cap:
|
|
22
|
+
if line:
|
|
23
|
+
out.append(line)
|
|
24
|
+
line = ""
|
|
25
|
+
out.append(word[:cap])
|
|
26
|
+
word = word[cap:]
|
|
27
|
+
if line and len(line) + 1 + len(word) > cap:
|
|
28
|
+
out.append(line)
|
|
29
|
+
line = word
|
|
30
|
+
else:
|
|
31
|
+
line = f"{line} {word}" if line else word
|
|
32
|
+
if line:
|
|
33
|
+
out.append(line)
|
|
34
|
+
return out
|
|
35
|
+
|
|
36
|
+
|
|
37
|
+
def split_text(text: str, cap: int = MAX_CHARS) -> list[str]:
|
|
38
|
+
"""Pieces of at most `cap` characters, cut at blank lines first, then between words."""
|
|
39
|
+
text = text.strip()
|
|
40
|
+
if len(text) <= cap:
|
|
41
|
+
return [text] if text else []
|
|
42
|
+
pieces, current = [], ""
|
|
43
|
+
for paragraph in re.split(r"\n\s*\n", text):
|
|
44
|
+
paragraph = paragraph.strip()
|
|
45
|
+
for part in [paragraph] if len(paragraph) <= cap else _cut_words(paragraph, cap):
|
|
46
|
+
if not part:
|
|
47
|
+
continue
|
|
48
|
+
if current and len(current) + 2 + len(part) > cap:
|
|
49
|
+
pieces.append(current)
|
|
50
|
+
current = part
|
|
51
|
+
else:
|
|
52
|
+
current = f"{current}\n\n{part}" if current else part
|
|
53
|
+
if current:
|
|
54
|
+
pieces.append(current)
|
|
55
|
+
return pieces
|
|
56
|
+
|
|
57
|
+
|
|
58
|
+
def chunks(item: Item) -> list[Chunk]:
|
|
59
|
+
out: list[Chunk] = []
|
|
60
|
+
for section in item.sections:
|
|
61
|
+
for piece in split_text(section.text):
|
|
62
|
+
out.append(Chunk(section.heading or "", piece, len(out)))
|
|
63
|
+
return out
|
|
64
|
+
|
|
65
|
+
|
|
66
|
+
def passage_text(item: Item, chunk: Chunk) -> str:
|
|
67
|
+
trail = " › ".join([*item.path, *([chunk.heading] if chunk.heading else [])])
|
|
68
|
+
return f"{item.title} › {trail}: {chunk.text}" if trail else f"{item.title}: {chunk.text}"
|
passageway/config.py
ADDED
|
@@ -0,0 +1,315 @@
|
|
|
1
|
+
"""Configuration: environment variables for the core, one TOML file for connectors (docs/install.md)."""
|
|
2
|
+
|
|
3
|
+
import math
|
|
4
|
+
import re
|
|
5
|
+
import sys
|
|
6
|
+
import threading
|
|
7
|
+
import tomllib
|
|
8
|
+
from collections.abc import Mapping
|
|
9
|
+
from dataclasses import dataclass, field
|
|
10
|
+
from pathlib import Path
|
|
11
|
+
|
|
12
|
+
from mcp.server.auth.routes import validate_issuer_url
|
|
13
|
+
from pydantic import AnyHttpUrl
|
|
14
|
+
|
|
15
|
+
from passageway.protocol import MAX_LINE_BYTES, MAX_OUTPUT_BYTES
|
|
16
|
+
|
|
17
|
+
|
|
18
|
+
class ConfigError(Exception):
|
|
19
|
+
pass
|
|
20
|
+
|
|
21
|
+
|
|
22
|
+
DEFAULT_TOOL_DESCRIPTION = (
|
|
23
|
+
"Search the user's indexed sources (documents, notes, decisions, how their systems and tools actually "
|
|
24
|
+
"behave) and get up to 3 short passages with links. Call it on your own before a task the sources may "
|
|
25
|
+
"cover, when you hit an unfamiliar error, host or service, and before making a decision they may already "
|
|
26
|
+
"have made. Returns nothing-relevant when there is no match. Open the link for the full document.")
|
|
27
|
+
|
|
28
|
+
DEFAULT_INSTRUCTIONS = (
|
|
29
|
+
"This server searches the user's indexed sources: their documents, notes, decisions and how their systems "
|
|
30
|
+
"and tools actually behave. Call `recall` before starting a task the sources may cover, when you hit an "
|
|
31
|
+
"unfamiliar error, host or service, and before making a decision they may already have made, even when you "
|
|
32
|
+
"think you know the answer: the sources hold what is true for this user. A hit is a pointer: open its link "
|
|
33
|
+
"for the full document before relying on it. If the context already holds passages headed \"Recall results\", "
|
|
34
|
+
"those came from this server; don't search for the same thing again.")
|
|
35
|
+
|
|
36
|
+
_NAME = re.compile(r"[a-z0-9][a-z0-9_-]*")
|
|
37
|
+
|
|
38
|
+
|
|
39
|
+
@dataclass(frozen=True)
|
|
40
|
+
class ConnectorConfig:
|
|
41
|
+
name: str
|
|
42
|
+
command: tuple[str, ...]
|
|
43
|
+
env: dict[str, str]
|
|
44
|
+
secrets: dict[str, str] # already resolved from env:NAME references
|
|
45
|
+
timeout: float = 300.0
|
|
46
|
+
max_line_bytes: int = MAX_LINE_BYTES
|
|
47
|
+
max_output_bytes: int = MAX_OUTPUT_BYTES
|
|
48
|
+
max_delete_fraction: float = 0.3 # a reconcile that would delete more than this share of a source is refused
|
|
49
|
+
|
|
50
|
+
|
|
51
|
+
EMBEDDER_KINDS = ("onnx", "command")
|
|
52
|
+
_ONNX_KEYS = {"kind", "min_similarity", "model", "model_dir", "device", "download"}
|
|
53
|
+
_COMMAND_KEYS = {"kind", "command", "env", "secrets", "min_similarity", "batch_size", "timeout"}
|
|
54
|
+
|
|
55
|
+
|
|
56
|
+
@dataclass(frozen=True)
|
|
57
|
+
class EmbedderConfig:
|
|
58
|
+
"""How texts become vectors: the built-in ONNX model, or any program speaking the embedder protocol."""
|
|
59
|
+
kind: str = "onnx"
|
|
60
|
+
command: tuple[str, ...] = ()
|
|
61
|
+
env: dict[str, str] = field(default_factory=dict)
|
|
62
|
+
secrets: dict[str, str] = field(default_factory=dict) # already resolved from env:NAME references
|
|
63
|
+
min_similarity: float | None = None # the relevance gate of this model; None keeps the built-in default
|
|
64
|
+
batch_size: int = 32 # texts per request
|
|
65
|
+
timeout: float = 60.0 # seconds per request
|
|
66
|
+
model_dir: Path | None = None # onnx only; None: <data dir>/models/<model>
|
|
67
|
+
model: str = "multilingual-e5-small" # onnx only: a preset of passageway.models
|
|
68
|
+
device: str = "auto" # onnx only: auto | cpu | coreml
|
|
69
|
+
download: bool = True # onnx only: fetch the pinned, checksummed files when they are missing
|
|
70
|
+
|
|
71
|
+
|
|
72
|
+
@dataclass
|
|
73
|
+
class Config:
|
|
74
|
+
password: str # PASSAGEWAY_LOGIN_PASSWORD: static Bearer token and the OAuth login page's password
|
|
75
|
+
connectors: list[ConnectorConfig]
|
|
76
|
+
data_dir: Path
|
|
77
|
+
model_dir: Path
|
|
78
|
+
base_url: str | None = None # public URL; set, it turns on OAuth for /mcp
|
|
79
|
+
host: str = "127.0.0.1"
|
|
80
|
+
port: int = 8082
|
|
81
|
+
poll_seconds: int = 60
|
|
82
|
+
reconcile_seconds: int = 1800
|
|
83
|
+
log_level: str = "INFO"
|
|
84
|
+
tool_description: str = DEFAULT_TOOL_DESCRIPTION
|
|
85
|
+
instructions: str = DEFAULT_INSTRUCTIONS
|
|
86
|
+
instance: str | None = None # PASSAGEWAY_INSTANCE: one of several servers on this machine
|
|
87
|
+
models_root: Path | None = None # where the built-in models live; None: data_dir (instances share the default)
|
|
88
|
+
min_similarity: float | None = None # the relevance gate; None keeps the embedder preset's value
|
|
89
|
+
embedder: EmbedderConfig = field(default_factory=EmbedderConfig) # the one the server uses
|
|
90
|
+
embedders: dict[str, EmbedderConfig] = field(default_factory=dict) # candidates for `passageway eval compare`
|
|
91
|
+
|
|
92
|
+
|
|
93
|
+
_INSTANCE = r"[a-z0-9][a-z0-9-]*"
|
|
94
|
+
|
|
95
|
+
|
|
96
|
+
def instance_name(env: Mapping[str, str]) -> str | None:
|
|
97
|
+
"""PASSAGEWAY_INSTANCE names one of several servers on this machine; unset, the one and only."""
|
|
98
|
+
name = env.get("PASSAGEWAY_INSTANCE", "").strip()
|
|
99
|
+
if name and not re.fullmatch(_INSTANCE, name):
|
|
100
|
+
raise ConfigError(f"PASSAGEWAY_INSTANCE must be lowercase letters, digits and hyphens, got {name!r}")
|
|
101
|
+
return name or None
|
|
102
|
+
|
|
103
|
+
|
|
104
|
+
def take_instance_argument(argv: list[str]) -> tuple[list[str], str | None]:
|
|
105
|
+
"""`--instance NAME` or `--instance=NAME`, anywhere in the command line: the rest of it and the name."""
|
|
106
|
+
rest, name, args = [], None, iter(argv)
|
|
107
|
+
for arg in args:
|
|
108
|
+
if arg == "--instance":
|
|
109
|
+
name = next(args, "")
|
|
110
|
+
elif arg.startswith("--instance="):
|
|
111
|
+
name = arg.split("=", 1)[1]
|
|
112
|
+
else:
|
|
113
|
+
rest.append(arg)
|
|
114
|
+
continue
|
|
115
|
+
if not re.fullmatch(_INSTANCE, name):
|
|
116
|
+
raise ConfigError(f"--instance needs a name of lowercase letters, digits and hyphens, got {name!r}")
|
|
117
|
+
return rest, name
|
|
118
|
+
|
|
119
|
+
|
|
120
|
+
def existing_instances() -> list[str]:
|
|
121
|
+
return sorted(p.parent.name.removeprefix("passageway-") for p in (Path.home() / ".config").glob("passageway-*/passageway.toml"))
|
|
122
|
+
|
|
123
|
+
|
|
124
|
+
def instance_problem(env: Mapping[str, str]) -> str | None:
|
|
125
|
+
"""Why this command would run against an instance that is not there, or None. A named instance must have
|
|
126
|
+
been created by init; with none named and no default config, the named ones are what the user meant."""
|
|
127
|
+
if env.get("PASSAGEWAY_CONFIG"):
|
|
128
|
+
return None
|
|
129
|
+
name, known = instance_name(env), existing_instances()
|
|
130
|
+
listed = ", ".join(known) or "none"
|
|
131
|
+
if name and not default_config_path(name).exists():
|
|
132
|
+
return (f"there is no passageway instance {name!r} (instances on this machine: {listed}); "
|
|
133
|
+
f"create it with: passageway --instance {name} init")
|
|
134
|
+
if not name and known and not default_config_path().exists():
|
|
135
|
+
return f"this machine has named passageway instances ({listed}): pick one with --instance NAME"
|
|
136
|
+
return None
|
|
137
|
+
|
|
138
|
+
|
|
139
|
+
def _suffix(instance: str | None) -> str:
|
|
140
|
+
return f"-{instance}" if instance else ""
|
|
141
|
+
|
|
142
|
+
|
|
143
|
+
def default_config_path(instance: str | None = None) -> Path:
|
|
144
|
+
return Path.home() / ".config" / f"passageway{_suffix(instance)}" / "passageway.toml"
|
|
145
|
+
|
|
146
|
+
|
|
147
|
+
def default_data_dir(instance: str | None = None) -> Path:
|
|
148
|
+
"""The data folder; an instance gets its own next to the shared one, which keeps the built-in models."""
|
|
149
|
+
base = Path.home() / "Library" / "Application Support" if sys.platform == "darwin" else Path.home() / ".local" / "share"
|
|
150
|
+
return base / f"passageway{_suffix(instance)}"
|
|
151
|
+
|
|
152
|
+
|
|
153
|
+
def _int(env: Mapping[str, str], key: str, default: int) -> int:
|
|
154
|
+
raw = env.get(key)
|
|
155
|
+
if not raw:
|
|
156
|
+
return default
|
|
157
|
+
try:
|
|
158
|
+
return int(raw)
|
|
159
|
+
except ValueError:
|
|
160
|
+
raise ConfigError(f"{key} must be an integer, got {raw!r}") from None
|
|
161
|
+
|
|
162
|
+
|
|
163
|
+
def _resolve(connector: str, ref: object, env: Mapping[str, str], what: str = "connector") -> str:
|
|
164
|
+
if not isinstance(ref, str) or not ref.startswith("env:") or len(ref) == 4:
|
|
165
|
+
raise ConfigError(f"{what} {connector!r}: a secret must be a reference like \"env:NAME\" (in 'secrets')")
|
|
166
|
+
var = ref[4:]
|
|
167
|
+
if not env.get(var):
|
|
168
|
+
raise ConfigError(f"{what} {connector!r}: the environment variable {var} is not set")
|
|
169
|
+
return env[var]
|
|
170
|
+
|
|
171
|
+
|
|
172
|
+
def _table(connector: str, raw: dict, key: str, what: str = "connector") -> dict[str, str]:
|
|
173
|
+
table = raw.get(key, {})
|
|
174
|
+
if not isinstance(table, dict) or not all(isinstance(v, str) for v in table.values()):
|
|
175
|
+
raise ConfigError(f"{what} {connector!r}: '{key}' must be a table of strings")
|
|
176
|
+
return table
|
|
177
|
+
|
|
178
|
+
|
|
179
|
+
def _embedder(name: str, raw: object, env: Mapping[str, str]) -> EmbedderConfig:
|
|
180
|
+
where = f"embedder {name!r}"
|
|
181
|
+
if not isinstance(raw, dict):
|
|
182
|
+
raise ConfigError(f"{where}: must be a table")
|
|
183
|
+
kind = raw.get("kind", "onnx")
|
|
184
|
+
if kind not in EMBEDDER_KINDS:
|
|
185
|
+
raise ConfigError(f"{where}: 'kind' must be one of {', '.join(EMBEDDER_KINDS)}, got {kind!r}")
|
|
186
|
+
allowed = _ONNX_KEYS if kind == "onnx" else _COMMAND_KEYS
|
|
187
|
+
for key in raw:
|
|
188
|
+
if key not in allowed:
|
|
189
|
+
raise ConfigError(f"{where}: unknown key {key!r} for kind {kind!r} (known: {', '.join(sorted(allowed))})")
|
|
190
|
+
gate = raw.get("min_similarity")
|
|
191
|
+
if gate is not None:
|
|
192
|
+
if isinstance(gate, bool) or not isinstance(gate, (int, float)) or not 0 <= gate <= 1:
|
|
193
|
+
raise ConfigError(f"{where}: 'min_similarity' must be a number from 0 to 1")
|
|
194
|
+
gate = float(gate)
|
|
195
|
+
if kind == "onnx":
|
|
196
|
+
from passageway.models import DEVICES, PRESETS
|
|
197
|
+
|
|
198
|
+
model_dir = raw.get("model_dir")
|
|
199
|
+
if model_dir is not None and not isinstance(model_dir, str):
|
|
200
|
+
raise ConfigError(f"{where}: 'model_dir' must be a path")
|
|
201
|
+
model = raw.get("model", "multilingual-e5-small")
|
|
202
|
+
if not isinstance(model, str) or model not in PRESETS:
|
|
203
|
+
raise ConfigError(f"{where}: 'model' must be one of the presets {', '.join(PRESETS)}, got {model!r} "
|
|
204
|
+
"(another model runs as kind = \"command\")")
|
|
205
|
+
device = raw.get("device", "auto")
|
|
206
|
+
if device not in DEVICES:
|
|
207
|
+
raise ConfigError(f"{where}: 'device' must be one of {', '.join(DEVICES)}, got {device!r}")
|
|
208
|
+
auto = raw.get("download", True)
|
|
209
|
+
if not isinstance(auto, bool):
|
|
210
|
+
raise ConfigError(f"{where}: 'download' must be true or false")
|
|
211
|
+
return EmbedderConfig("onnx", min_similarity=gate, model_dir=Path(model_dir) if model_dir else None, model=model,
|
|
212
|
+
device=device, download=auto)
|
|
213
|
+
command = raw.get("command")
|
|
214
|
+
if not isinstance(command, list) or not command or not all(isinstance(c, str) for c in command):
|
|
215
|
+
raise ConfigError(f"{where}: 'command' must be a non-empty list of strings")
|
|
216
|
+
secrets = {k: _resolve(name, v, env, "embedder") for k, v in _table(name, raw, "secrets", "embedder").items()}
|
|
217
|
+
batch = raw.get("batch_size", 32)
|
|
218
|
+
if isinstance(batch, bool) or not isinstance(batch, int) or not 1 <= batch <= 1024:
|
|
219
|
+
raise ConfigError(f"{where}: 'batch_size' must be a whole number from 1 to 1024")
|
|
220
|
+
timeout = raw.get("timeout", 60.0)
|
|
221
|
+
if (isinstance(timeout, bool) or not isinstance(timeout, (int, float)) or not math.isfinite(timeout) or timeout <= 0
|
|
222
|
+
or timeout > threading.TIMEOUT_MAX):
|
|
223
|
+
raise ConfigError(f"{where}: 'timeout' must be a number of seconds greater than 0 (and not absurdly large)")
|
|
224
|
+
return EmbedderConfig("command", tuple(command), _table(name, raw, "env", "embedder"), secrets, gate, batch, float(timeout))
|
|
225
|
+
|
|
226
|
+
|
|
227
|
+
def _number(connector: str, raw: dict, key: str, default, *, kind=float, low=None, high=None):
|
|
228
|
+
value = raw.get(key, default)
|
|
229
|
+
ok = (isinstance(value, (int, float)) and not isinstance(value, bool) and (kind is float or isinstance(value, int))
|
|
230
|
+
and math.isfinite(value))
|
|
231
|
+
if not ok or (low is not None and value <= low) or (high is not None and value > high):
|
|
232
|
+
bound = f" (greater than {low}{'' if high is None else f' and at most {high}'})" if low is not None else ""
|
|
233
|
+
raise ConfigError(f"connector {connector!r}: '{key}' must be a number{bound}")
|
|
234
|
+
return kind(value)
|
|
235
|
+
|
|
236
|
+
|
|
237
|
+
def _load_toml(path: Path, env: Mapping[str, str]) -> tuple[list[ConnectorConfig], EmbedderConfig, dict[str, EmbedderConfig]]:
|
|
238
|
+
try:
|
|
239
|
+
data = tomllib.loads(path.read_text())
|
|
240
|
+
except FileNotFoundError:
|
|
241
|
+
raise ConfigError(f"{path} not found: list your connectors in a passageway.toml (see docs/install.md)") from None
|
|
242
|
+
except tomllib.TOMLDecodeError as e:
|
|
243
|
+
raise ConfigError(f"{path}: {e}") from None
|
|
244
|
+
connectors: list[ConnectorConfig] = []
|
|
245
|
+
entries = data.get("connector", [])
|
|
246
|
+
if not isinstance(entries, list) or not all(isinstance(e, dict) for e in entries):
|
|
247
|
+
raise ConfigError(f"{path}: connectors are written as [[connector]] tables, one per source")
|
|
248
|
+
for raw in entries:
|
|
249
|
+
name = raw.get("name", "")
|
|
250
|
+
if not isinstance(name, str) or not _NAME.fullmatch(name):
|
|
251
|
+
raise ConfigError(f"connector name {name!r} must match [a-z0-9][a-z0-9_-]* (it names the index file)")
|
|
252
|
+
if any(c.name == name for c in connectors):
|
|
253
|
+
raise ConfigError(f"duplicate connector name {name!r}")
|
|
254
|
+
command = raw.get("command")
|
|
255
|
+
if not isinstance(command, list) or not command or not all(isinstance(c, str) for c in command):
|
|
256
|
+
raise ConfigError(f"connector {name!r}: 'command' must be a non-empty list of strings")
|
|
257
|
+
secrets = {k: _resolve(name, v, env) for k, v in _table(name, raw, "secrets").items()}
|
|
258
|
+
connectors.append(ConnectorConfig(
|
|
259
|
+
name, tuple(command), _table(name, raw, "env"), secrets,
|
|
260
|
+
_number(name, raw, "timeout", 300.0, low=0, high=threading.TIMEOUT_MAX),
|
|
261
|
+
_number(name, raw, "max_line_bytes", MAX_LINE_BYTES, kind=int, low=0),
|
|
262
|
+
_number(name, raw, "max_output_bytes", MAX_OUTPUT_BYTES, kind=int, low=0),
|
|
263
|
+
_number(name, raw, "max_delete_fraction", 0.3, low=0, high=1)))
|
|
264
|
+
if not connectors:
|
|
265
|
+
raise ConfigError(f"no connectors in {path}: add at least one [[connector]]")
|
|
266
|
+
embedder = _embedder("default", data.get("embedder", {}), env)
|
|
267
|
+
candidates = data.get("embedders", {})
|
|
268
|
+
if not isinstance(candidates, dict):
|
|
269
|
+
raise ConfigError(f"{path}: 'embedders' holds the candidates for `passageway eval compare`, one [embedders.NAME] table each")
|
|
270
|
+
embedders: dict[str, EmbedderConfig] = {}
|
|
271
|
+
for name, raw in candidates.items():
|
|
272
|
+
if not _NAME.fullmatch(name):
|
|
273
|
+
raise ConfigError(f"embedder name {name!r} must match [a-z0-9][a-z0-9_-]*")
|
|
274
|
+
if name == "default":
|
|
275
|
+
raise ConfigError("the embedder name 'default' is taken: it stands for the [embedder] the server uses")
|
|
276
|
+
embedders[name] = _embedder(name, raw, env)
|
|
277
|
+
return connectors, embedder, embedders
|
|
278
|
+
|
|
279
|
+
|
|
280
|
+
def load_config(env: Mapping[str, str], *, require_password: bool = True) -> Config:
|
|
281
|
+
"""require_password=False is for commands that do not serve (the evaluation): nobody logs in to them."""
|
|
282
|
+
password = env.get("PASSAGEWAY_LOGIN_PASSWORD", "").strip()
|
|
283
|
+
password_file = env.get("PASSAGEWAY_LOGIN_PASSWORD_FILE", "").strip()
|
|
284
|
+
if not password and password_file: # a file the service manager can point at (passageway init writes one)
|
|
285
|
+
try:
|
|
286
|
+
password = Path(password_file).read_text().strip()
|
|
287
|
+
except OSError as e:
|
|
288
|
+
raise ConfigError(f"PASSAGEWAY_LOGIN_PASSWORD_FILE={password_file}: {e.strerror or e}") from None
|
|
289
|
+
if not password:
|
|
290
|
+
raise ConfigError(f"PASSAGEWAY_LOGIN_PASSWORD_FILE={password_file} is empty")
|
|
291
|
+
if require_password and not password:
|
|
292
|
+
raise ConfigError("PASSAGEWAY_LOGIN_PASSWORD is not set (nor PASSAGEWAY_LOGIN_PASSWORD_FILE): clients authenticate "
|
|
293
|
+
"with it, so the engine refuses to start without it; `passageway init` generates one")
|
|
294
|
+
base_url = env.get("PASSAGEWAY_BASE_URL", "").strip().rstrip("/") or None
|
|
295
|
+
if base_url:
|
|
296
|
+
try:
|
|
297
|
+
validate_issuer_url(AnyHttpUrl(base_url))
|
|
298
|
+
except ValueError as e:
|
|
299
|
+
raise ConfigError(f"PASSAGEWAY_BASE_URL={base_url!r} is not a valid OAuth issuer "
|
|
300
|
+
f"(https, no query or fragment; plain http only for localhost): {e}") from None
|
|
301
|
+
instance = instance_name(env)
|
|
302
|
+
config_path = Path(env.get("PASSAGEWAY_CONFIG") or default_config_path(instance))
|
|
303
|
+
connectors, embedder, embedders = _load_toml(config_path, env)
|
|
304
|
+
data_dir = Path(env.get("PASSAGEWAY_DATA_DIR") or default_data_dir(instance))
|
|
305
|
+
models_root = Path(env.get("PASSAGEWAY_DATA_DIR") or default_data_dir())
|
|
306
|
+
return Config(
|
|
307
|
+
password=password, connectors=connectors, data_dir=data_dir, instance=instance, models_root=models_root,
|
|
308
|
+
model_dir=Path(env.get("PASSAGEWAY_MODEL_DIR") or models_root / "models" / "e5-small"),
|
|
309
|
+
base_url=base_url, host=env.get("PASSAGEWAY_HOST", "127.0.0.1"),
|
|
310
|
+
port=_int(env, "PASSAGEWAY_PORT", 8082), poll_seconds=_int(env, "PASSAGEWAY_POLL_SECONDS", 60),
|
|
311
|
+
reconcile_seconds=_int(env, "PASSAGEWAY_RECONCILE_SECONDS", 1800),
|
|
312
|
+
log_level=env.get("PASSAGEWAY_LOG_LEVEL", "INFO"),
|
|
313
|
+
tool_description=env.get("PASSAGEWAY_TOOL_DESCRIPTION", "").strip() or DEFAULT_TOOL_DESCRIPTION,
|
|
314
|
+
instructions=env.get("PASSAGEWAY_INSTRUCTIONS", "").strip() or DEFAULT_INSTRUCTIONS,
|
|
315
|
+
min_similarity=embedder.min_similarity, embedder=embedder, embedders=embedders)
|
passageway/embed.py
ADDED
|
@@ -0,0 +1,213 @@
|
|
|
1
|
+
"""Text embedders: the built-in ONNX model, any program speaking the embedder protocol, and a hashing stand-in for tests."""
|
|
2
|
+
|
|
3
|
+
import hashlib
|
|
4
|
+
import logging
|
|
5
|
+
import re
|
|
6
|
+
from pathlib import Path
|
|
7
|
+
from typing import Protocol
|
|
8
|
+
|
|
9
|
+
import numpy as np
|
|
10
|
+
|
|
11
|
+
from passageway.config import ConfigError, EmbedderConfig
|
|
12
|
+
from passageway.models import DEFAULT_PRESET, PRESETS, ModelError, download, files_missing, model_folder_for, verify
|
|
13
|
+
|
|
14
|
+
log = logging.getLogger(__name__)
|
|
15
|
+
|
|
16
|
+
|
|
17
|
+
class EmbedderError(Exception):
|
|
18
|
+
"""An embedder failed: a program of its own died or answered badly, or the model raised."""
|
|
19
|
+
|
|
20
|
+
|
|
21
|
+
class EmbedderStartError(ConfigError):
|
|
22
|
+
"""An embedder could not start. `public` is what an open endpoint may show; the message is for the log."""
|
|
23
|
+
|
|
24
|
+
def __init__(self, message: str, public: str):
|
|
25
|
+
super().__init__(message)
|
|
26
|
+
self.public = public
|
|
27
|
+
|
|
28
|
+
|
|
29
|
+
class Embedder(Protocol):
|
|
30
|
+
dim: int
|
|
31
|
+
name: str
|
|
32
|
+
|
|
33
|
+
def embed_queries(self, texts: list[str]) -> np.ndarray: ...
|
|
34
|
+
|
|
35
|
+
def embed_passages(self, texts: list[str]) -> np.ndarray: ...
|
|
36
|
+
|
|
37
|
+
|
|
38
|
+
def _normalise(v: np.ndarray) -> np.ndarray:
|
|
39
|
+
norms = np.linalg.norm(v, axis=1, keepdims=True)
|
|
40
|
+
norms[norms == 0] = 1
|
|
41
|
+
return (v / norms).astype(np.float32)
|
|
42
|
+
|
|
43
|
+
|
|
44
|
+
def token_batches(lengths: list[int], budget: int) -> list[list[int]]:
|
|
45
|
+
"""Group indices so each batch's padded size (count × longest) stays within budget.
|
|
46
|
+
|
|
47
|
+
Memory at inference grows with the padded batch, so long passages go alone or in pairs
|
|
48
|
+
while short ones share a batch. Sorting by length keeps padding small.
|
|
49
|
+
"""
|
|
50
|
+
batches: list[list[int]] = []
|
|
51
|
+
current: list[int] = []
|
|
52
|
+
for i in sorted(range(len(lengths)), key=lambda i: lengths[i]):
|
|
53
|
+
if current and (len(current) + 1) * lengths[i] > budget:
|
|
54
|
+
batches.append(current)
|
|
55
|
+
current = []
|
|
56
|
+
current.append(i)
|
|
57
|
+
if current:
|
|
58
|
+
batches.append(current)
|
|
59
|
+
return batches
|
|
60
|
+
|
|
61
|
+
|
|
62
|
+
class OnnxEmbedder:
|
|
63
|
+
"""A built-in model: an ONNX encoder and its tokenizer from a preset, mean pooling over the attention mask, L2-normalised."""
|
|
64
|
+
|
|
65
|
+
max_tokens = 512
|
|
66
|
+
# Padded tokens per inference call: 1024 keeps peak RSS near 310 MB for e5-small (8 × 512 reached 550 MB).
|
|
67
|
+
token_budget = 1024
|
|
68
|
+
|
|
69
|
+
def __init__(self, model_dir: Path, preset=None, device: str = "auto", threads: int = 4):
|
|
70
|
+
import onnxruntime as ort
|
|
71
|
+
|
|
72
|
+
from passageway.models import PRESETS, choose_device
|
|
73
|
+
|
|
74
|
+
self.preset = preset or PRESETS["multilingual-e5-small"]
|
|
75
|
+
self.name, self.dim, self.min_similarity = self.preset.label, self.preset.dim, self.preset.min_similarity
|
|
76
|
+
self.max_tokens = self.preset.max_tokens
|
|
77
|
+
options = ort.SessionOptions()
|
|
78
|
+
options.intra_op_num_threads = threads
|
|
79
|
+
options.inter_op_num_threads = 1
|
|
80
|
+
# The arena would keep the largest batch's peak reserved for the life of the process.
|
|
81
|
+
options.enable_cpu_mem_arena = False
|
|
82
|
+
providers = choose_device(device, ort.get_available_providers())
|
|
83
|
+
self._session = ort.InferenceSession(str(Path(model_dir) / "model.onnx"), options, providers=providers)
|
|
84
|
+
self._inputs = {i.name for i in self._session.get_inputs()}
|
|
85
|
+
self._encode_batch = self._tokenizer(self.preset, Path(model_dir))
|
|
86
|
+
|
|
87
|
+
@staticmethod
|
|
88
|
+
def _tokenizer(preset, model_dir: Path):
|
|
89
|
+
"""A function from texts to (ids, mask) arrays, padded to the longest text and cut at max_tokens."""
|
|
90
|
+
if preset.tokenizer == "sentencepiece":
|
|
91
|
+
import sentencepiece as spm
|
|
92
|
+
|
|
93
|
+
# The SentencePiece model gives the same ids as the HF tokenizer.json at ~60 MB instead of ~280 MB.
|
|
94
|
+
sp = spm.SentencePieceProcessor(model_file=str(model_dir / "sentencepiece.bpe.model"))
|
|
95
|
+
|
|
96
|
+
def encode(text: str) -> list[int]:
|
|
97
|
+
# XLM-R ids: fairseq offset of 1 over SentencePiece, its <unk> (0) is 3, wrapped in <s> (0) … </s> (2).
|
|
98
|
+
ids = [3 if i == 0 else i + 1 for i in sp.encode(text)][: preset.max_tokens - 2]
|
|
99
|
+
return [0, *ids, 2]
|
|
100
|
+
|
|
101
|
+
pad = 1
|
|
102
|
+
else:
|
|
103
|
+
from tokenizers import Tokenizer
|
|
104
|
+
|
|
105
|
+
tok = Tokenizer.from_file(str(model_dir / "tokenizer.json"))
|
|
106
|
+
tok.enable_truncation(preset.max_tokens)
|
|
107
|
+
tok.no_padding()
|
|
108
|
+
pad = next((i for i in (tok.token_to_id(t) for t in ("[PAD]", "<pad>")) if i is not None), 0)
|
|
109
|
+
|
|
110
|
+
def encode(text: str) -> list[int]:
|
|
111
|
+
return tok.encode(text).ids
|
|
112
|
+
|
|
113
|
+
def batch(texts: list[str]) -> tuple[np.ndarray, np.ndarray]:
|
|
114
|
+
encoded = [encode(t) for t in texts]
|
|
115
|
+
width = max(len(e) for e in encoded)
|
|
116
|
+
ids = np.full((len(encoded), width), pad, dtype=np.int64)
|
|
117
|
+
mask = np.zeros((len(encoded), width), dtype=np.int64)
|
|
118
|
+
for row, e in enumerate(encoded):
|
|
119
|
+
ids[row, : len(e)] = e
|
|
120
|
+
mask[row, : len(e)] = 1
|
|
121
|
+
return ids, mask
|
|
122
|
+
|
|
123
|
+
return batch
|
|
124
|
+
|
|
125
|
+
def _embed(self, texts: list[str]) -> np.ndarray:
|
|
126
|
+
if not texts:
|
|
127
|
+
return np.zeros((0, self.dim), dtype=np.float32)
|
|
128
|
+
lengths = [len(self._encode_batch([t])[0][0]) for t in texts]
|
|
129
|
+
out = np.zeros((len(texts), self.dim), dtype=np.float32)
|
|
130
|
+
for batch in token_batches(lengths, self.token_budget):
|
|
131
|
+
ids, mask = self._encode_batch([texts[i] for i in batch])
|
|
132
|
+
feed = {"input_ids": ids, "attention_mask": mask}
|
|
133
|
+
if "token_type_ids" in self._inputs:
|
|
134
|
+
feed["token_type_ids"] = np.zeros_like(ids)
|
|
135
|
+
hidden = self._session.run(None, feed)[0]
|
|
136
|
+
weights = mask[..., None].astype(np.float32)
|
|
137
|
+
out[batch] = (hidden * weights).sum(axis=1) / np.maximum(weights.sum(axis=1), 1e-9)
|
|
138
|
+
return _normalise(out)
|
|
139
|
+
|
|
140
|
+
def embed_queries(self, texts: list[str]) -> np.ndarray:
|
|
141
|
+
return self._embed([self.preset.query_prefix + t for t in texts])
|
|
142
|
+
|
|
143
|
+
def embed_passages(self, texts: list[str]) -> np.ndarray:
|
|
144
|
+
return self._embed([self.preset.passage_prefix + t for t in texts])
|
|
145
|
+
|
|
146
|
+
|
|
147
|
+
class HashEmbedder:
|
|
148
|
+
"""Bag of lower-cased words hashed into buckets: overlap in words means similarity. Tests only."""
|
|
149
|
+
|
|
150
|
+
name = "hash-test"
|
|
151
|
+
|
|
152
|
+
def __init__(self, dim: int = 384):
|
|
153
|
+
self.dim = dim
|
|
154
|
+
|
|
155
|
+
def _embed(self, texts: list[str]) -> np.ndarray:
|
|
156
|
+
v = np.zeros((len(texts), self.dim), dtype=np.float32)
|
|
157
|
+
for row, text in enumerate(texts):
|
|
158
|
+
for word in re.findall(r"\w+", text.lower()):
|
|
159
|
+
bucket = int.from_bytes(hashlib.blake2b(word.encode(), digest_size=4).digest(), "big")
|
|
160
|
+
v[row, bucket % self.dim] += 1
|
|
161
|
+
return _normalise(v)
|
|
162
|
+
|
|
163
|
+
embed_queries = _embed
|
|
164
|
+
embed_passages = _embed
|
|
165
|
+
|
|
166
|
+
|
|
167
|
+
def model_folder(cfg: EmbedderConfig, config) -> Path:
|
|
168
|
+
"""Where a preset's files live: one rule with `passageway model` (models.model_folder_for); config.model_dir is the
|
|
169
|
+
server's resolved PASSAGEWAY_MODEL_DIR for the default preset."""
|
|
170
|
+
if cfg.model_dir:
|
|
171
|
+
return cfg.model_dir
|
|
172
|
+
if cfg.model == DEFAULT_PRESET:
|
|
173
|
+
return config.model_dir
|
|
174
|
+
return model_folder_for(cfg.model, config.models_root or config.data_dir, None)
|
|
175
|
+
|
|
176
|
+
|
|
177
|
+
def build_embedder(cfg: EmbedderConfig, config, label: str = "default") -> Embedder:
|
|
178
|
+
"""The embedder a configuration describes. A model or program that is not there is a ConfigError, so that a server
|
|
179
|
+
can start in its error mode and say what is wrong instead of crashing."""
|
|
180
|
+
if cfg.kind == "command":
|
|
181
|
+
from passageway.embed_command import CommandEmbedder
|
|
182
|
+
|
|
183
|
+
if cfg.min_similarity is None:
|
|
184
|
+
log.warning("embedder %r has no min_similarity: the relevance gate of the built-in model applies, which is "
|
|
185
|
+
"probably wrong for this one; find a value with `passageway eval --sweep`", label)
|
|
186
|
+
try:
|
|
187
|
+
return CommandEmbedder(label, cfg.command, env=cfg.env, secrets=cfg.secrets, batch_size=cfg.batch_size,
|
|
188
|
+
timeout=cfg.timeout)
|
|
189
|
+
except EmbedderError as e: # its output can hold hosts and paths: only the log gets it
|
|
190
|
+
raise EmbedderStartError(str(e), public=f"the embedder {label!r} could not start; see the server log") from None
|
|
191
|
+
preset = PRESETS[cfg.model]
|
|
192
|
+
folder = model_folder(cfg, config)
|
|
193
|
+
missing, wrong = files_missing(preset, folder), verify(preset, folder)
|
|
194
|
+
if missing or wrong:
|
|
195
|
+
if not cfg.download:
|
|
196
|
+
what = f"missing {', '.join(missing)}" if missing else f"a wrong checksum for {', '.join(wrong)}"
|
|
197
|
+
raise ConfigError(f"the model {cfg.model!r} in {folder} has {what}: run `passageway model download {cfg.model}` "
|
|
198
|
+
"(download = true in [embedder] does it at start)")
|
|
199
|
+
log.info("fetching the model %s (%d MB) from huggingface.co into %s, pinned and checksummed; the server answers "
|
|
200
|
+
"once it is there", cfg.model, preset.size_mb, folder)
|
|
201
|
+
try:
|
|
202
|
+
download(preset, folder, report=lambda m: log.info("%s", m))
|
|
203
|
+
except (ModelError, OSError) as e:
|
|
204
|
+
raise ConfigError(f"could not fetch the model {cfg.model!r}: {e}") from None
|
|
205
|
+
try:
|
|
206
|
+
embedder = OnnxEmbedder(folder, preset, cfg.device)
|
|
207
|
+
except ModelError as e:
|
|
208
|
+
raise ConfigError(str(e)) from None
|
|
209
|
+
if cfg.min_similarity is None and getattr(embedder, "min_similarity", None) is None:
|
|
210
|
+
log.warning("the model %r has no measured min_similarity yet, so the default gate of the built-in model (0.845) "
|
|
211
|
+
"applies, and on another similarity scale almost nothing will pass it: set min_similarity under "
|
|
212
|
+
"[embedder] from a run of `passageway eval --sweep`", cfg.model)
|
|
213
|
+
return embedder
|