finos-mcp-core 0.1.1__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.
- finos_mcp_core-0.1.1/.gitignore +16 -0
- finos_mcp_core-0.1.1/PKG-INFO +21 -0
- finos_mcp_core-0.1.1/README.md +1 -0
- finos_mcp_core-0.1.1/pyproject.toml +31 -0
- finos_mcp_core-0.1.1/src/finos_mcp/core/__init__.py +80 -0
- finos_mcp_core-0.1.1/src/finos_mcp/core/audit.py +82 -0
- finos_mcp_core-0.1.1/src/finos_mcp/core/catalog.py +463 -0
- finos_mcp_core-0.1.1/src/finos_mcp/core/errors.py +121 -0
- finos_mcp_core-0.1.1/src/finos_mcp/core/metrics.py +62 -0
- finos_mcp_core-0.1.1/src/finos_mcp/core/models.py +97 -0
- finos_mcp_core-0.1.1/src/finos_mcp/core/policy.py +102 -0
- finos_mcp_core-0.1.1/src/finos_mcp/core/ratelimit.py +62 -0
- finos_mcp_core-0.1.1/src/finos_mcp/core/schema.py +153 -0
- finos_mcp_core-0.1.1/src/finos_mcp/core/server.py +328 -0
- finos_mcp_core-0.1.1/src/finos_mcp/core/vendor.py +107 -0
- finos_mcp_core-0.1.1/tests/test_catalog_schema_vendor.py +169 -0
- finos_mcp_core-0.1.1/tests/test_semantic.py +158 -0
- finos_mcp_core-0.1.1/tests/test_server_safety.py +232 -0
|
@@ -0,0 +1,21 @@
|
|
|
1
|
+
Metadata-Version: 2.5
|
|
2
|
+
Name: finos-mcp-core
|
|
3
|
+
Version: 0.1.1
|
|
4
|
+
Summary: Shared safety layer for finos-mcp servers: read-only enforcement, rate limits, audit, structured errors
|
|
5
|
+
License-Expression: Apache-2.0
|
|
6
|
+
Requires-Python: >=3.12
|
|
7
|
+
Requires-Dist: bm25s>=0.2
|
|
8
|
+
Requires-Dist: jsonschema<5,>=4.23
|
|
9
|
+
Requires-Dist: mcp==2.1.1
|
|
10
|
+
Requires-Dist: pydantic<3,>=2.10
|
|
11
|
+
Requires-Dist: pystemmer>=2.2
|
|
12
|
+
Requires-Dist: python-frontmatter>=1.1
|
|
13
|
+
Requires-Dist: pyyaml>=6
|
|
14
|
+
Requires-Dist: rapidfuzz>=3.9
|
|
15
|
+
Requires-Dist: referencing>=0.35
|
|
16
|
+
Provides-Extra: semantic
|
|
17
|
+
Requires-Dist: model2vec>=0.4; extra == 'semantic'
|
|
18
|
+
Requires-Dist: numpy>=2.5.2; extra == 'semantic'
|
|
19
|
+
Description-Content-Type: text/markdown
|
|
20
|
+
|
|
21
|
+
# finos-mcp-core
|
|
@@ -0,0 +1 @@
|
|
|
1
|
+
# finos-mcp-core
|
|
@@ -0,0 +1,31 @@
|
|
|
1
|
+
[project]
|
|
2
|
+
name = "finos-mcp-core"
|
|
3
|
+
version = "0.1.1"
|
|
4
|
+
description = "Shared safety layer for finos-mcp servers: read-only enforcement, rate limits, audit, structured errors"
|
|
5
|
+
readme = "README.md"
|
|
6
|
+
requires-python = ">=3.12"
|
|
7
|
+
license = "Apache-2.0"
|
|
8
|
+
dependencies = [
|
|
9
|
+
"mcp==2.1.1",
|
|
10
|
+
"pydantic>=2.10,<3",
|
|
11
|
+
"jsonschema>=4.23,<5",
|
|
12
|
+
"referencing>=0.35",
|
|
13
|
+
"bm25s>=0.2",
|
|
14
|
+
"rapidfuzz>=3.9",
|
|
15
|
+
"python-frontmatter>=1.1",
|
|
16
|
+
"pyyaml>=6",
|
|
17
|
+
"pystemmer>=2.2",
|
|
18
|
+
]
|
|
19
|
+
|
|
20
|
+
[project.optional-dependencies]
|
|
21
|
+
semantic = [
|
|
22
|
+
"model2vec>=0.4",
|
|
23
|
+
"numpy>=2.5.2",
|
|
24
|
+
]
|
|
25
|
+
|
|
26
|
+
[build-system]
|
|
27
|
+
requires = ["hatchling"]
|
|
28
|
+
build-backend = "hatchling.build"
|
|
29
|
+
|
|
30
|
+
[tool.hatch.build.targets.wheel]
|
|
31
|
+
packages = ["src/finos_mcp"]
|
|
@@ -0,0 +1,80 @@
|
|
|
1
|
+
"""finos-mcp-core: the shared safety layer for finos-mcp servers."""
|
|
2
|
+
|
|
3
|
+
from .audit import AuditLog
|
|
4
|
+
from .catalog import (
|
|
5
|
+
DEFAULT_EMBEDDING_MODEL,
|
|
6
|
+
Catalog,
|
|
7
|
+
Document,
|
|
8
|
+
SearchMode,
|
|
9
|
+
SemanticIndex,
|
|
10
|
+
normalize_id,
|
|
11
|
+
slugify,
|
|
12
|
+
split_sections,
|
|
13
|
+
)
|
|
14
|
+
from .errors import (
|
|
15
|
+
ErrorEnvelope,
|
|
16
|
+
FinosToolError,
|
|
17
|
+
ambiguous,
|
|
18
|
+
invalid_input,
|
|
19
|
+
not_found,
|
|
20
|
+
unsupported_format,
|
|
21
|
+
)
|
|
22
|
+
from .metrics import Metrics
|
|
23
|
+
from .models import (
|
|
24
|
+
Citation,
|
|
25
|
+
Page,
|
|
26
|
+
SearchHit,
|
|
27
|
+
Section,
|
|
28
|
+
ServerInfo,
|
|
29
|
+
ValidationIssue,
|
|
30
|
+
ValidationReport,
|
|
31
|
+
paginate,
|
|
32
|
+
)
|
|
33
|
+
from .policy import RateLimit, SafetyPolicy
|
|
34
|
+
from .ratelimit import RateLimiter
|
|
35
|
+
from .schema import SchemaRegistry
|
|
36
|
+
from .server import READ_ONLY, build_server, register_server_info, register_tool, runtime_for
|
|
37
|
+
from .vendor import SourceManifest, VendorIntegrityError, build_manifest, verify
|
|
38
|
+
|
|
39
|
+
__version__ = "0.1.1"
|
|
40
|
+
|
|
41
|
+
__all__ = [
|
|
42
|
+
"DEFAULT_EMBEDDING_MODEL",
|
|
43
|
+
"READ_ONLY",
|
|
44
|
+
"AuditLog",
|
|
45
|
+
"Catalog",
|
|
46
|
+
"Citation",
|
|
47
|
+
"Document",
|
|
48
|
+
"ErrorEnvelope",
|
|
49
|
+
"FinosToolError",
|
|
50
|
+
"Metrics",
|
|
51
|
+
"Page",
|
|
52
|
+
"RateLimit",
|
|
53
|
+
"RateLimiter",
|
|
54
|
+
"SafetyPolicy",
|
|
55
|
+
"SchemaRegistry",
|
|
56
|
+
"SearchHit",
|
|
57
|
+
"SearchMode",
|
|
58
|
+
"Section",
|
|
59
|
+
"SemanticIndex",
|
|
60
|
+
"ServerInfo",
|
|
61
|
+
"SourceManifest",
|
|
62
|
+
"ValidationIssue",
|
|
63
|
+
"ValidationReport",
|
|
64
|
+
"VendorIntegrityError",
|
|
65
|
+
"__version__",
|
|
66
|
+
"ambiguous",
|
|
67
|
+
"build_manifest",
|
|
68
|
+
"build_server",
|
|
69
|
+
"invalid_input",
|
|
70
|
+
"normalize_id",
|
|
71
|
+
"not_found",
|
|
72
|
+
"paginate",
|
|
73
|
+
"register_server_info",
|
|
74
|
+
"register_tool",
|
|
75
|
+
"runtime_for",
|
|
76
|
+
"slugify",
|
|
77
|
+
"split_sections",
|
|
78
|
+
"unsupported_format",
|
|
79
|
+
"verify",
|
|
80
|
+
]
|
|
@@ -0,0 +1,82 @@
|
|
|
1
|
+
"""JSON-lines audit log of every request.
|
|
2
|
+
|
|
3
|
+
Under stdio transport, stdout is the protocol channel, so the log goes to stderr
|
|
4
|
+
(default) or to the file named by ``FINOS_MCP_AUDIT_PATH``. Tool arguments are
|
|
5
|
+
hashed by default so the log never becomes a second copy of the agent's inputs.
|
|
6
|
+
"""
|
|
7
|
+
|
|
8
|
+
from __future__ import annotations
|
|
9
|
+
|
|
10
|
+
import hashlib
|
|
11
|
+
import json
|
|
12
|
+
import os
|
|
13
|
+
import sys
|
|
14
|
+
import threading
|
|
15
|
+
from collections.abc import Mapping
|
|
16
|
+
from datetime import UTC, datetime
|
|
17
|
+
from pathlib import Path
|
|
18
|
+
from typing import Any, TextIO
|
|
19
|
+
|
|
20
|
+
ENV_AUDIT_PATH = "FINOS_MCP_AUDIT_PATH"
|
|
21
|
+
|
|
22
|
+
|
|
23
|
+
def canonical_json(value: Any) -> str:
|
|
24
|
+
return json.dumps(value, sort_keys=True, separators=(",", ":"), ensure_ascii=False, default=str)
|
|
25
|
+
|
|
26
|
+
|
|
27
|
+
def sha256_text(text: str) -> str:
|
|
28
|
+
return hashlib.sha256(text.encode("utf-8")).hexdigest()
|
|
29
|
+
|
|
30
|
+
|
|
31
|
+
class AuditLog:
|
|
32
|
+
def __init__(
|
|
33
|
+
self,
|
|
34
|
+
server: str,
|
|
35
|
+
*,
|
|
36
|
+
stream: TextIO | None = None,
|
|
37
|
+
path: Path | None = None,
|
|
38
|
+
hash_inputs: bool = True,
|
|
39
|
+
) -> None:
|
|
40
|
+
self.server = server
|
|
41
|
+
self.hash_inputs = hash_inputs
|
|
42
|
+
self._path = path
|
|
43
|
+
self._stream: TextIO | None = stream
|
|
44
|
+
self._lock = threading.Lock()
|
|
45
|
+
self.records_written = 0
|
|
46
|
+
|
|
47
|
+
@classmethod
|
|
48
|
+
def from_env(cls, server: str, *, hash_inputs: bool = True) -> AuditLog:
|
|
49
|
+
raw = os.environ.get(ENV_AUDIT_PATH)
|
|
50
|
+
if raw:
|
|
51
|
+
p = Path(raw).expanduser()
|
|
52
|
+
p.parent.mkdir(parents=True, exist_ok=True)
|
|
53
|
+
return cls(server, path=p, hash_inputs=hash_inputs)
|
|
54
|
+
return cls(server, stream=sys.stderr, hash_inputs=hash_inputs)
|
|
55
|
+
|
|
56
|
+
def describe_args(self, arguments: Mapping[str, Any] | None) -> dict[str, Any]:
|
|
57
|
+
text = canonical_json(arguments or {})
|
|
58
|
+
out: dict[str, Any] = {"args_bytes": len(text.encode("utf-8"))}
|
|
59
|
+
if self.hash_inputs:
|
|
60
|
+
out["args_sha256"] = sha256_text(text)
|
|
61
|
+
else:
|
|
62
|
+
out["args"] = arguments or {}
|
|
63
|
+
return out
|
|
64
|
+
|
|
65
|
+
def write(self, record: dict[str, Any]) -> None:
|
|
66
|
+
line = json.dumps(
|
|
67
|
+
{
|
|
68
|
+
"ts": datetime.now(UTC).isoformat(timespec="milliseconds"),
|
|
69
|
+
"server": self.server,
|
|
70
|
+
**record,
|
|
71
|
+
},
|
|
72
|
+
ensure_ascii=False,
|
|
73
|
+
default=str,
|
|
74
|
+
)
|
|
75
|
+
with self._lock:
|
|
76
|
+
if self._path is not None:
|
|
77
|
+
with self._path.open("a", encoding="utf-8", newline="\n") as fh:
|
|
78
|
+
fh.write(line + "\n")
|
|
79
|
+
elif self._stream is not None:
|
|
80
|
+
self._stream.write(line + "\n")
|
|
81
|
+
self._stream.flush()
|
|
82
|
+
self.records_written += 1
|
|
@@ -0,0 +1,463 @@
|
|
|
1
|
+
"""Generic document catalog with forgiving id resolution and BM25 search.
|
|
2
|
+
|
|
3
|
+
Servers load their vendored records into `Document`s and get two things back:
|
|
4
|
+
|
|
5
|
+
* `resolve(query)` normalises every id form a model might emit and returns
|
|
6
|
+
exactly one document, or raises a structured `not_found` / `ambiguous_id`.
|
|
7
|
+
* `search(query, k)` is BM25 over section-level chunks with a fuzzy title boost,
|
|
8
|
+
optionally fused (reciprocal rank fusion) with a dense `SemanticIndex` when the
|
|
9
|
+
`semantic` extra is installed. By default -- no extra installed -- there is no
|
|
10
|
+
embeddings, no model download, sub-second startup.
|
|
11
|
+
"""
|
|
12
|
+
|
|
13
|
+
from __future__ import annotations
|
|
14
|
+
|
|
15
|
+
import os
|
|
16
|
+
import re
|
|
17
|
+
from collections.abc import Callable, Iterable
|
|
18
|
+
from dataclasses import dataclass
|
|
19
|
+
from functools import lru_cache
|
|
20
|
+
from pathlib import Path
|
|
21
|
+
from typing import Any, Literal
|
|
22
|
+
|
|
23
|
+
import bm25s
|
|
24
|
+
import Stemmer
|
|
25
|
+
from pydantic import BaseModel, Field
|
|
26
|
+
from rapidfuzz import fuzz, process
|
|
27
|
+
|
|
28
|
+
from .errors import FinosToolError, ambiguous, not_found
|
|
29
|
+
from .models import SearchHit, Section
|
|
30
|
+
|
|
31
|
+
_STEMMER = Stemmer.Stemmer("english")
|
|
32
|
+
|
|
33
|
+
DEFAULT_EMBEDDING_MODEL = "minishlab/potion-base-4M"
|
|
34
|
+
_RRF_K = 60
|
|
35
|
+
_SEARCH_MODES = ("hybrid", "lexical", "dense")
|
|
36
|
+
SearchMode = Literal["hybrid", "lexical", "dense"]
|
|
37
|
+
# Embedding every document eagerly at Catalog construction is fine for a few hundred
|
|
38
|
+
# short records (AIGF: ~46) but not for a catalog with thousands of entries (e.g. CDM's
|
|
39
|
+
# type registry): skip semantic indexing there rather than making every server that
|
|
40
|
+
# happens to share this module pay an unbounded embedding pass at startup.
|
|
41
|
+
_MAX_EAGER_EMBEDDING_DOCS = 500
|
|
42
|
+
|
|
43
|
+
|
|
44
|
+
def _tokenize(texts: list[str]) -> Any:
|
|
45
|
+
"""BM25 tokens with English stopwords removed and Snowball stemming applied."""
|
|
46
|
+
return bm25s.tokenize(texts, stopwords="en", stemmer=_STEMMER, show_progress=False)
|
|
47
|
+
|
|
48
|
+
|
|
49
|
+
_NORM_RE = re.compile(r"[\s_\-./:]+")
|
|
50
|
+
|
|
51
|
+
|
|
52
|
+
def normalize_id(text: str) -> str:
|
|
53
|
+
"""Case-fold and strip separators: 'AIR-SEC-010', 'air sec 10', 'air_sec_010' all align."""
|
|
54
|
+
t = _NORM_RE.sub("", text.strip().casefold())
|
|
55
|
+
# collapse zero-padded trailing numbers: airsec010 -> airsec10
|
|
56
|
+
return re.sub(r"(?<=[a-z])0+(\d)", r"\1", t)
|
|
57
|
+
|
|
58
|
+
|
|
59
|
+
def slugify(text: str) -> str:
|
|
60
|
+
return re.sub(r"[^a-z0-9]+", "-", text.casefold()).strip("-")
|
|
61
|
+
|
|
62
|
+
|
|
63
|
+
def split_sections(markdown: str, *, min_level: int = 2) -> list[Section]:
|
|
64
|
+
"""Split markdown into sections on '##'-style headings (level >= min_level)."""
|
|
65
|
+
sections: list[Section] = []
|
|
66
|
+
heading: str | None = None
|
|
67
|
+
level = min_level
|
|
68
|
+
buf: list[str] = []
|
|
69
|
+
fence = False
|
|
70
|
+
|
|
71
|
+
def flush() -> None:
|
|
72
|
+
body = "\n".join(buf).strip()
|
|
73
|
+
if heading is not None or body:
|
|
74
|
+
sections.append(
|
|
75
|
+
Section(
|
|
76
|
+
heading=heading or "",
|
|
77
|
+
slug=slugify(heading) if heading else "preamble",
|
|
78
|
+
level=level,
|
|
79
|
+
body=body,
|
|
80
|
+
)
|
|
81
|
+
)
|
|
82
|
+
|
|
83
|
+
for line in markdown.splitlines():
|
|
84
|
+
if line.startswith("```"):
|
|
85
|
+
fence = not fence
|
|
86
|
+
m = None if fence else re.match(r"^(#{1,6})\s+(.*?)\s*#*\s*$", line)
|
|
87
|
+
if m and len(m.group(1)) >= min_level:
|
|
88
|
+
flush()
|
|
89
|
+
heading, level, buf = m.group(2).strip(), len(m.group(1)), []
|
|
90
|
+
else:
|
|
91
|
+
buf.append(line)
|
|
92
|
+
flush()
|
|
93
|
+
return sections
|
|
94
|
+
|
|
95
|
+
|
|
96
|
+
class Document(BaseModel):
|
|
97
|
+
id: str
|
|
98
|
+
title: str
|
|
99
|
+
aliases: list[str] = Field(default_factory=list)
|
|
100
|
+
kind: str = "document"
|
|
101
|
+
body: str = ""
|
|
102
|
+
sections: list[Section] = Field(default_factory=list)
|
|
103
|
+
meta: dict[str, Any] = Field(default_factory=dict)
|
|
104
|
+
|
|
105
|
+
|
|
106
|
+
@dataclass(slots=True)
|
|
107
|
+
class _Chunk:
|
|
108
|
+
doc_index: int
|
|
109
|
+
section: str | None
|
|
110
|
+
text: str
|
|
111
|
+
|
|
112
|
+
|
|
113
|
+
@lru_cache(maxsize=8)
|
|
114
|
+
def _load_static_model(model_source: str) -> Any:
|
|
115
|
+
"""Load and cache a `model2vec.StaticModel` by name or local path.
|
|
116
|
+
|
|
117
|
+
Cached so multiple `Catalog`s (risks, controls, combined, ...) in one process
|
|
118
|
+
share a single loaded model instead of re-downloading/re-loading it each time.
|
|
119
|
+
A failed load is never cached (the exception simply propagates to the caller,
|
|
120
|
+
which retries on the next `Catalog` build)."""
|
|
121
|
+
from model2vec import StaticModel
|
|
122
|
+
|
|
123
|
+
# model2vec defaults to force_download=True, which would re-fetch the model from
|
|
124
|
+
# the Hub on every process start (every server, every test subprocess) even when
|
|
125
|
+
# it is already cached locally -- exactly the network dependency the `semantic`
|
|
126
|
+
# extra's docs promise not to add after the first download.
|
|
127
|
+
#
|
|
128
|
+
# Do NOT try to make this offline-first by setting HF_HUB_OFFLINE around the call:
|
|
129
|
+
# huggingface_hub freezes that env var into a module constant when it is first
|
|
130
|
+
# imported, which happens *inside* this call, so the setting outlives the block and
|
|
131
|
+
# a later retry stays offline for the life of the process -- a machine with no
|
|
132
|
+
# cached model could then never download one. Point FINOS_MCP_EMBEDDING_MODEL at a
|
|
133
|
+
# local directory for a genuinely offline deployment.
|
|
134
|
+
return StaticModel.from_pretrained(model_source, force_download=False)
|
|
135
|
+
|
|
136
|
+
|
|
137
|
+
class SemanticIndex:
|
|
138
|
+
"""Dense embedding index over whole-document text, built once per `Catalog`.
|
|
139
|
+
|
|
140
|
+
Uses `model2vec` static embeddings (pure numpy at inference time, no torch).
|
|
141
|
+
Vectors are L2-normalised so a dot product is a cosine similarity."""
|
|
142
|
+
|
|
143
|
+
def __init__(self, model: Any, vectors: Any) -> None:
|
|
144
|
+
self._model = model
|
|
145
|
+
self._vectors = vectors # shape (n_docs, dim); row i == Catalog._docs[i]
|
|
146
|
+
|
|
147
|
+
@classmethod
|
|
148
|
+
def build(cls, docs: list[Document], *, model_name: str) -> SemanticIndex:
|
|
149
|
+
import numpy as np
|
|
150
|
+
|
|
151
|
+
model = _load_static_model(model_name)
|
|
152
|
+
if not docs:
|
|
153
|
+
return cls(model, np.zeros((0, 1), dtype=np.float32))
|
|
154
|
+
texts = [f"{d.title} {d.body[:4000]}" for d in docs]
|
|
155
|
+
vectors = np.asarray(model.encode(texts, normalize=True), dtype=np.float32)
|
|
156
|
+
return cls(model, vectors)
|
|
157
|
+
|
|
158
|
+
def rank(self, query: str) -> list[tuple[int, float]]:
|
|
159
|
+
"""Return every doc index with its cosine similarity to `query`, best first."""
|
|
160
|
+
import numpy as np
|
|
161
|
+
|
|
162
|
+
if self._vectors.shape[0] == 0:
|
|
163
|
+
return []
|
|
164
|
+
qv = np.asarray(self._model.encode([query], normalize=True)[0], dtype=np.float32)
|
|
165
|
+
sims = self._vectors @ qv
|
|
166
|
+
order = np.argsort(-sims)
|
|
167
|
+
return [(int(i), float(sims[i])) for i in order]
|
|
168
|
+
|
|
169
|
+
|
|
170
|
+
def _reciprocal_rank_fusion(*rankings: list[int], k: int = _RRF_K) -> list[int]:
|
|
171
|
+
scores: dict[int, float] = {}
|
|
172
|
+
for ranking in rankings:
|
|
173
|
+
for rank, doc_index in enumerate(ranking):
|
|
174
|
+
scores[doc_index] = scores.get(doc_index, 0.0) + 1.0 / (k + rank + 1)
|
|
175
|
+
return [doc_index for doc_index, _ in sorted(scores.items(), key=lambda kv: -kv[1])]
|
|
176
|
+
|
|
177
|
+
|
|
178
|
+
class Catalog:
|
|
179
|
+
def __init__(
|
|
180
|
+
self,
|
|
181
|
+
documents: Iterable[Document],
|
|
182
|
+
*,
|
|
183
|
+
kind: str,
|
|
184
|
+
citation_uri: Callable[[Document], str],
|
|
185
|
+
) -> None:
|
|
186
|
+
self.kind = kind
|
|
187
|
+
self._docs: list[Document] = list(documents)
|
|
188
|
+
self._citation_uri = citation_uri
|
|
189
|
+
self._by_norm: dict[str, list[int]] = {}
|
|
190
|
+
for i, d in enumerate(self._docs):
|
|
191
|
+
for key in (d.id, *d.aliases):
|
|
192
|
+
self._by_norm.setdefault(normalize_id(key), []).append(i)
|
|
193
|
+
self._title_norm: dict[str, int] = {
|
|
194
|
+
normalize_id(d.title): i for i, d in enumerate(self._docs)
|
|
195
|
+
}
|
|
196
|
+
self._titles = [d.title for d in self._docs]
|
|
197
|
+
self._chunks: list[_Chunk] = []
|
|
198
|
+
for i, d in enumerate(self._docs):
|
|
199
|
+
self._chunks.append(_Chunk(i, None, f"{d.title}\n{d.body}"))
|
|
200
|
+
for s in d.sections:
|
|
201
|
+
if s.body:
|
|
202
|
+
self._chunks.append(
|
|
203
|
+
_Chunk(i, s.heading or None, f"{d.title}\n{s.heading}\n{s.body}")
|
|
204
|
+
)
|
|
205
|
+
self._retriever: bm25s.BM25 | None = None
|
|
206
|
+
if self._chunks:
|
|
207
|
+
tokens = _tokenize([c.text for c in self._chunks])
|
|
208
|
+
self._retriever = bm25s.BM25()
|
|
209
|
+
self._retriever.index(tokens, show_progress=False)
|
|
210
|
+
self.semantic: SemanticIndex | None = None
|
|
211
|
+
self.semantic_status: str = "disabled: not attempted"
|
|
212
|
+
self._build_semantic()
|
|
213
|
+
|
|
214
|
+
# ---------------------------------------------------------------- semantic
|
|
215
|
+
def _build_semantic(self) -> None:
|
|
216
|
+
"""Best-effort dense index build. Never raises: any failure -- the
|
|
217
|
+
`semantic` extra not installed, no network/no cached model, a bad
|
|
218
|
+
`FINOS_MCP_EMBEDDING_MODEL` -- leaves `self.semantic` as `None` and
|
|
219
|
+
records why in `self.semantic_status`."""
|
|
220
|
+
if len(self._docs) > _MAX_EAGER_EMBEDDING_DOCS:
|
|
221
|
+
self.semantic_status = (
|
|
222
|
+
f"disabled: catalog too large for eager embedding "
|
|
223
|
+
f"({len(self._docs)} docs > {_MAX_EAGER_EMBEDDING_DOCS})"
|
|
224
|
+
)
|
|
225
|
+
return
|
|
226
|
+
model_name = os.environ.get("FINOS_MCP_EMBEDDING_MODEL") or DEFAULT_EMBEDDING_MODEL
|
|
227
|
+
try:
|
|
228
|
+
import model2vec # noqa: F401
|
|
229
|
+
except ImportError:
|
|
230
|
+
self.semantic_status = "disabled: model2vec not installed"
|
|
231
|
+
return
|
|
232
|
+
source = str(Path(model_name)) if Path(model_name).is_dir() else model_name
|
|
233
|
+
try:
|
|
234
|
+
self.semantic = SemanticIndex.build(self._docs, model_name=source)
|
|
235
|
+
self.semantic_status = f"enabled: {model_name}"
|
|
236
|
+
except Exception as exc: # loading must never break the server
|
|
237
|
+
self.semantic = None
|
|
238
|
+
self.semantic_status = f"disabled: model load failed: {exc}"
|
|
239
|
+
|
|
240
|
+
# ------------------------------------------------------------------ access
|
|
241
|
+
def __len__(self) -> int:
|
|
242
|
+
return len(self._docs)
|
|
243
|
+
|
|
244
|
+
def __iter__(self) -> Iterable[Document]:
|
|
245
|
+
return iter(self._docs)
|
|
246
|
+
|
|
247
|
+
@property
|
|
248
|
+
def documents(self) -> list[Document]:
|
|
249
|
+
return list(self._docs)
|
|
250
|
+
|
|
251
|
+
def get(self, id: str) -> Document | None:
|
|
252
|
+
hits = self._by_norm.get(normalize_id(id), [])
|
|
253
|
+
return self._docs[hits[0]] if len(hits) == 1 else None
|
|
254
|
+
|
|
255
|
+
def citation_uri(self, doc: Document) -> str:
|
|
256
|
+
return self._citation_uri(doc)
|
|
257
|
+
|
|
258
|
+
# ----------------------------------------------------------------- resolve
|
|
259
|
+
def resolve(self, query: str) -> Document:
|
|
260
|
+
"""Return exactly one document for an id, alias, or title; else raise a structured error."""
|
|
261
|
+
q = query.strip()
|
|
262
|
+
if not q:
|
|
263
|
+
raise not_found(
|
|
264
|
+
self.kind, query, hint="Pass an id such as the ones returned by list_* tools."
|
|
265
|
+
)
|
|
266
|
+
norm = normalize_id(q)
|
|
267
|
+
exact = self._by_norm.get(norm)
|
|
268
|
+
if exact and len(set(exact)) == 1:
|
|
269
|
+
return self._docs[exact[0]]
|
|
270
|
+
if norm in self._title_norm:
|
|
271
|
+
return self._docs[self._title_norm[norm]]
|
|
272
|
+
# fuzzy title match
|
|
273
|
+
matches = process.extract(q, self._titles, scorer=fuzz.WRatio, limit=5, score_cutoff=60)
|
|
274
|
+
if matches:
|
|
275
|
+
best_score = matches[0][1]
|
|
276
|
+
runner_up = matches[1][1] if len(matches) > 1 else 0.0
|
|
277
|
+
if best_score >= 90 and best_score - runner_up >= 5:
|
|
278
|
+
return self._docs[matches[0][2]]
|
|
279
|
+
candidates = [self._docs[m[2]].id for m in matches]
|
|
280
|
+
raise ambiguous(self.kind, query, candidates)
|
|
281
|
+
raise not_found(
|
|
282
|
+
self.kind,
|
|
283
|
+
query,
|
|
284
|
+
hint=f"Use search to find a {self.kind} by keyword, then call again with its id.",
|
|
285
|
+
)
|
|
286
|
+
|
|
287
|
+
def try_resolve(self, query: str) -> Document | None:
|
|
288
|
+
try:
|
|
289
|
+
return self.resolve(query)
|
|
290
|
+
except FinosToolError:
|
|
291
|
+
return None
|
|
292
|
+
|
|
293
|
+
# ------------------------------------------------------------------ search
|
|
294
|
+
def search(
|
|
295
|
+
self,
|
|
296
|
+
query: str,
|
|
297
|
+
*,
|
|
298
|
+
k: int = 10,
|
|
299
|
+
predicate: Callable[[Document], bool] | None = None,
|
|
300
|
+
title_boost: float = 0.35,
|
|
301
|
+
mode: SearchMode = "hybrid",
|
|
302
|
+
) -> list[SearchHit]:
|
|
303
|
+
"""Search `query` and return up to `k` hits, best first.
|
|
304
|
+
|
|
305
|
+
`mode="lexical"` is exactly today's BM25 + fuzzy-title-boost algorithm.
|
|
306
|
+
`mode="dense"` uses only the `SemanticIndex` (empty if unavailable).
|
|
307
|
+
`mode="hybrid"` (default) reciprocal-rank-fuses the two, except that a
|
|
308
|
+
clearly-confident lexical top hit (its BM25-only relevance, before the
|
|
309
|
+
title boost, is >= 90% of the corpus max) is kept in first place -- a
|
|
310
|
+
strong keyword/title match should not be bumped by a dense disagreement.
|
|
311
|
+
`FINOS_MCP_SEARCH_MODE` (`hybrid`/`lexical`/`dense`) overrides `mode`
|
|
312
|
+
for the whole process, e.g. to force lexical-only for an eval run.
|
|
313
|
+
"""
|
|
314
|
+
if not query.strip():
|
|
315
|
+
return []
|
|
316
|
+
env_mode = os.environ.get("FINOS_MCP_SEARCH_MODE")
|
|
317
|
+
if env_mode in _SEARCH_MODES:
|
|
318
|
+
mode = env_mode # type: ignore[assignment]
|
|
319
|
+
|
|
320
|
+
if mode == "dense":
|
|
321
|
+
return self._dense_hits(query, k=k, predicate=predicate)
|
|
322
|
+
|
|
323
|
+
lex_n = (
|
|
324
|
+
min(len(self._chunks), max(k * 4, 20))
|
|
325
|
+
if mode == "lexical"
|
|
326
|
+
else min(len(self._chunks), max(k * 8, 50))
|
|
327
|
+
)
|
|
328
|
+
lexical = self._lexical_rank(query, title_boost=title_boost, n=lex_n)
|
|
329
|
+
lex_lookup = {d: (rel, section, text) for d, rel, _raw, section, text in lexical}
|
|
330
|
+
|
|
331
|
+
if mode == "lexical" or self.semantic is None:
|
|
332
|
+
order = [d for d, *_ in lexical]
|
|
333
|
+
return self._hits_from_order(order, lex_lookup, query, k=k, predicate=predicate)
|
|
334
|
+
|
|
335
|
+
dense_ranked = self.semantic.rank(query)
|
|
336
|
+
# RRF only cares about rank position, not magnitude, so an unfiltered dense
|
|
337
|
+
# ranking would drag in every document in the corpus (cosine similarity is
|
|
338
|
+
# rarely exactly zero for unrelated text the way BM25's term overlap is).
|
|
339
|
+
# Keep only docs at least half as similar as the best dense match.
|
|
340
|
+
dense_floor = 0.5 * dense_ranked[0][1] if dense_ranked else 0.0
|
|
341
|
+
dense_order = [d for d, sim in dense_ranked if sim >= dense_floor]
|
|
342
|
+
lexical_order = [d for d, *_ in lexical]
|
|
343
|
+
fused = _reciprocal_rank_fusion(lexical_order, dense_order)
|
|
344
|
+
if lexical and lexical[0][2] >= 0.9: # raw (pre-title-boost) BM25 relevance
|
|
345
|
+
top_doc = lexical[0][0]
|
|
346
|
+
fused = [top_doc, *(d for d in fused if d != top_doc)]
|
|
347
|
+
return self._hits_from_order(fused, lex_lookup, query, k=k, predicate=predicate)
|
|
348
|
+
|
|
349
|
+
def _lexical_rank(
|
|
350
|
+
self, query: str, *, title_boost: float, n: int
|
|
351
|
+
) -> list[tuple[int, float, float, str | None, str]]:
|
|
352
|
+
"""The BM25 + fuzzy-title-boost algorithm `search(mode="lexical")` has always
|
|
353
|
+
used. Returns `(doc_index, rel, raw_rel, section, text)` sorted by `rel`
|
|
354
|
+
descending; `raw_rel` is the BM25-only component (no title boost)."""
|
|
355
|
+
if self._retriever is None:
|
|
356
|
+
return []
|
|
357
|
+
q_tokens = _tokenize([query])
|
|
358
|
+
n = min(len(self._chunks), max(n, 1))
|
|
359
|
+
try:
|
|
360
|
+
ids, scores = self._retriever.retrieve(q_tokens, k=n, show_progress=False)
|
|
361
|
+
except (ValueError, IndexError):
|
|
362
|
+
return []
|
|
363
|
+
best: dict[int, tuple[float, float, str | None, str]] = {}
|
|
364
|
+
max_score = float(scores[0][0]) if len(scores[0]) and float(scores[0][0]) > 0 else 1.0
|
|
365
|
+
for chunk_id, score in zip(ids[0], scores[0], strict=True):
|
|
366
|
+
s = float(score)
|
|
367
|
+
if s <= 0:
|
|
368
|
+
continue
|
|
369
|
+
chunk = self._chunks[int(chunk_id)]
|
|
370
|
+
raw_rel = s / max_score
|
|
371
|
+
doc = self._docs[chunk.doc_index]
|
|
372
|
+
rel = raw_rel + title_boost * (
|
|
373
|
+
fuzz.partial_ratio(query.casefold(), doc.title.casefold()) / 100.0
|
|
374
|
+
)
|
|
375
|
+
prev = best.get(chunk.doc_index)
|
|
376
|
+
if prev is None or rel > prev[0]:
|
|
377
|
+
best[chunk.doc_index] = (rel, raw_rel, chunk.section, chunk.text)
|
|
378
|
+
# title-only matches for docs BM25 missed entirely
|
|
379
|
+
for title, score, idx in process.extract(
|
|
380
|
+
query, self._titles, scorer=fuzz.partial_ratio, limit=5
|
|
381
|
+
):
|
|
382
|
+
if idx not in best and score >= 85:
|
|
383
|
+
best[idx] = (
|
|
384
|
+
title_boost * score / 100.0,
|
|
385
|
+
0.0,
|
|
386
|
+
None,
|
|
387
|
+
f"{title}\n{self._docs[idx].body[:300]}",
|
|
388
|
+
)
|
|
389
|
+
return [
|
|
390
|
+
(doc_index, rel, raw_rel, section, text)
|
|
391
|
+
for doc_index, (rel, raw_rel, section, text) in sorted(
|
|
392
|
+
best.items(), key=lambda kv: -kv[1][0]
|
|
393
|
+
)
|
|
394
|
+
]
|
|
395
|
+
|
|
396
|
+
def _dense_hits(
|
|
397
|
+
self, query: str, *, k: int, predicate: Callable[[Document], bool] | None
|
|
398
|
+
) -> list[SearchHit]:
|
|
399
|
+
if self.semantic is None:
|
|
400
|
+
return []
|
|
401
|
+
hits: list[SearchHit] = []
|
|
402
|
+
for doc_index, sim in self.semantic.rank(query):
|
|
403
|
+
doc = self._docs[doc_index]
|
|
404
|
+
if predicate is not None and not predicate(doc):
|
|
405
|
+
continue
|
|
406
|
+
text = f"{doc.title}\n{doc.body[:300]}"
|
|
407
|
+
hits.append(
|
|
408
|
+
SearchHit(
|
|
409
|
+
id=doc.id,
|
|
410
|
+
title=doc.title,
|
|
411
|
+
section=None,
|
|
412
|
+
score=round(sim, 4),
|
|
413
|
+
snippet=_snippet(text, query),
|
|
414
|
+
citation_uri=self._citation_uri(doc),
|
|
415
|
+
)
|
|
416
|
+
)
|
|
417
|
+
if len(hits) >= k:
|
|
418
|
+
break
|
|
419
|
+
return hits
|
|
420
|
+
|
|
421
|
+
def _hits_from_order(
|
|
422
|
+
self,
|
|
423
|
+
order: list[int],
|
|
424
|
+
lex_lookup: dict[int, tuple[float, str | None, str]],
|
|
425
|
+
query: str,
|
|
426
|
+
*,
|
|
427
|
+
k: int,
|
|
428
|
+
predicate: Callable[[Document], bool] | None,
|
|
429
|
+
) -> list[SearchHit]:
|
|
430
|
+
hits: list[SearchHit] = []
|
|
431
|
+
for doc_index in order:
|
|
432
|
+
doc = self._docs[doc_index]
|
|
433
|
+
if predicate is not None and not predicate(doc):
|
|
434
|
+
continue
|
|
435
|
+
if doc_index in lex_lookup:
|
|
436
|
+
rel, section, text = lex_lookup[doc_index]
|
|
437
|
+
else:
|
|
438
|
+
rel, section, text = 0.0, None, f"{doc.title}\n{doc.body[:300]}"
|
|
439
|
+
hits.append(
|
|
440
|
+
SearchHit(
|
|
441
|
+
id=doc.id,
|
|
442
|
+
title=doc.title,
|
|
443
|
+
section=section,
|
|
444
|
+
score=round(rel, 4),
|
|
445
|
+
snippet=_snippet(text, query),
|
|
446
|
+
citation_uri=self._citation_uri(doc),
|
|
447
|
+
)
|
|
448
|
+
)
|
|
449
|
+
if len(hits) >= k:
|
|
450
|
+
break
|
|
451
|
+
return hits
|
|
452
|
+
|
|
453
|
+
|
|
454
|
+
def _snippet(text: str, query: str, width: int = 240) -> str:
|
|
455
|
+
body = " ".join(text.split())
|
|
456
|
+
terms = [t for t in re.findall(r"\w+", query.casefold()) if len(t) > 2]
|
|
457
|
+
low = body.casefold()
|
|
458
|
+
pos = min((low.find(t) for t in terms if low.find(t) >= 0), default=-1)
|
|
459
|
+
if pos < 0:
|
|
460
|
+
return body[:width] + ("…" if len(body) > width else "")
|
|
461
|
+
start = max(0, pos - width // 3)
|
|
462
|
+
end = min(len(body), start + width)
|
|
463
|
+
return ("…" if start > 0 else "") + body[start:end] + ("…" if end < len(body) else "")
|