diff --git a/README.md b/README.md index a333da839..6ae9bd2e1 100644 --- a/README.md +++ b/README.md @@ -83,7 +83,7 @@ AssetOpsBench is a **unified framework for developing, orchestrating, and evalua | MCP Servers | Important tools | |---|---| | **IoT** | `get_sites`, `get_history`, `get_assets`, `get_sensors` | -| **FMSR** | `get_sensors`, `get_failure_modes`, `get_failure_sensor_mapping` | +| **FMSR** | `get_failure_modes`, `generate_failure_modes`, `add_failure_modes`, `generate_failure_mode_sensor_mapping` | | **TSFM** | `forecasting`, `timeseries_anomaly_detection` | | **WO** | `get_work_order_distribution`, `predict_next_work_order`, ... | | **Vibration** | `compute_fft_spectrum`, `compute_envelope_spectrum`, ... | @@ -109,7 +109,7 @@ The `src/` directory contains MCP servers and a plan-execute runner built on the | Domain | Example Task | |---|---| | **IoT** | "List all sensors of Chiller 6 in MAIN site" | -| **FMSR** | "Identify failure modes detected by Chiller 6 Supply Temperature" | +| **FMSR** | "List known failure modes for asset class pump" | | **TSFM** | "Forecast Chiller 9 Condenser Water Flow for the week of 2020-04-27" | | **WO** | "Generate a work order for Chiller 6 anomaly detection" | diff --git a/docs/data.md b/docs/data.md index 52ab82767..259d2c6ec 100644 --- a/docs/data.md +++ b/docs/data.md @@ -21,7 +21,7 @@ docker compose -f src/couchdb/docker-compose.yaml up -d # 2. Verify (Fauxton UI: http://localhost:5984/_utils — admin / password) curl -s -u admin:password http://localhost:5984/_all_dbs -# → ["workorder","iot","vibration","failurecode", ...] +# → ["workorder","iot","vibration","failure_code", ...] curl -s -u admin:password http://localhost:5984/workorder/wo:MAIN:1000045 ``` @@ -133,7 +133,7 @@ init_data() # back to default "workorder": "shared/work_order/workorders.csv", "iot": ["shared/iot/chiller_6.json", "shared/iot/metro_pump_1.json"], "vibration": "shared/iot/motor_01.json", - "failurecode": "shared/failure_code/failure_code_sample.csv" + "failure_code": "shared/failure_code/failure_code_sample.csv" } ``` @@ -154,7 +154,7 @@ Edit this file and re-run `init_data.py` — that's the whole change. ### 2. Change *how* a collection is parsed — `collections.json` ```json -"failurecode": { +"failure_code": { "format": "csv", "primary_key": ["code"], "id_prefix": "fc", @@ -283,4 +283,4 @@ folder's parent); data unique to a private scenario goes in that scenario's fold Consumers should sort client-side or add the index to `collections.json`. - **`--reuse` keeps history; default reload wipes it.** Anything written at runtime (e.g. work orders created by agent tools) is deleted on the next default load — - that's the deterministic-benchmark behavior. \ No newline at end of file + that's the deterministic-benchmark behavior. diff --git a/docs/mcp-servers.md b/docs/mcp-servers.md index cf79790ce..e1ffa6b9b 100644 --- a/docs/mcp-servers.md +++ b/docs/mcp-servers.md @@ -59,13 +59,15 @@ Synthetic motor vibration data (`asset_id: Motor_01`, from `motor_01.json`) ship ## fmsr — Failure Mode and Sensor Relations **Path:** `src/servers/fmsr/main.py` -**Requires:** `WATSONX_APIKEY`, `WATSONX_PROJECT_ID`, `WATSONX_URL` for unknown assets; curated lists for `chiller` and `ahu` work without credentials. -**Failure-mode data:** `src/servers/fmsr/failure_modes.yaml` (edit to add/change asset entries) +**Requires:** LLM credentials for `generate_failure_modes` and `generate_failure_mode_sensor_mapping`; `get_failure_modes` reads the database. +**Failure-mode data:** `src/couchdb/scenarios_data/shared/fmea/failure_modes_sample.json` loaded into the `failure_mode` database collection. | Tool | Category | Arguments | Description | | --------------------------------- | ------------- | ---------------------------------------- | ------------------------------------------------------------------------------------------------------------------------------------------------------- | -| `get_failure_modes` | read, LLM-use | `asset_name` | Return known failure modes for an asset. Uses a curated YAML list for chillers and AHUs; falls back to the LLM for other types. | -| `get_failure_mode_sensor_mapping` | read, LLM-use | `asset_name`, `failure_modes`, `sensors` | For each (failure mode, sensor) pair, determine relevancy via LLM. Returns bidirectional `fm→sensors` and `sensor→fms` maps plus full per-pair details. | +| `get_failure_modes` | read | `asset_class` | Return known failure modes for an asset class from the database. Returns `asset_class`, `failure_modes`, `exhaustive`, and `source`. | +| `generate_failure_modes` | read, LLM-use | `asset_class`, `max_modes?` | Generate or extend a failure-mode list without writing the database. | +| `add_failure_modes` | write | `asset_class`, `failure_modes`, `exhaustive?`, `source?` | Persist failure modes for an asset class. | +| `generate_failure_mode_sensor_mapping` | read, LLM-use | `asset_class`, `failure_modes`, `sensors` | Score failure-mode/sensor relevancy via LLM and return bidirectional mappings. | ## wo — Work Order diff --git a/docs/opencode-agent.md b/docs/opencode-agent.md index 7d5d98f1d..3cf86d3e4 100644 --- a/docs/opencode-agent.md +++ b/docs/opencode-agent.md @@ -199,6 +199,18 @@ The custom provider route is important because OpenCode's built-in `openai/*` provider validates model names against its own model registry. Router-hosted models such as `MiniMax-M3` must be registered explicitly. +### FMSR tool model + +FMSR tools take `asset_class`. `--model-id` controls OpenCode; LLM-backed FMSR +tools use `FMSR_MODEL_ID`. + +```bash +MODEL_ID=tokenrouter/MiniMax-M3 +FMSR_MODEL_ID="$MODEL_ID" \ +uv run opencode-agent --model-id "$MODEL_ID" --show-trajectory \ + "Use generate_failure_mode_sensor_mapping for asset_class pump with failure_modes ['seal leakage'] and sensors ['Pressure sensor']." +``` + --- ## Permissions and web access @@ -383,7 +395,31 @@ uv run opencode-agent --run-id opencode-smoke --scenario-id smoke \ --model-id tokenrouter/MiniMax-M3 \ "What sites are available?" -# 5. benchmark-suite dry run +# 5. FMSR get_failure_modes smoke test +uv run opencode-agent --show-trajectory \ + --model-id tokenrouter/MiniMax-M3 \ + "Use the FMSR get_failure_modes tool for asset_class pump. Return only the failure_modes list." + +# 6. FMSR generate_failure_modes smoke test +MODEL_ID=tokenrouter/MiniMax-M3 +FMSR_MODEL_ID="$MODEL_ID" \ +uv run opencode-agent --show-trajectory \ + --model-id "$MODEL_ID" \ + "Use generate_failure_modes for asset_class pump with max_modes 3. Return known, generated, and message." + +# 7. FMSR add_failure_modes smoke test +uv run opencode-agent --show-trajectory \ + --model-id tokenrouter/MiniMax-M3 \ + "Use add_failure_modes for asset_class pump with failure_modes ['bearing wear'], exhaustive false, and source 'manual smoke test'. Return added, total, source, and message." + +# 8. FMSR mapping smoke test +MODEL_ID=tokenrouter/MiniMax-M3 +FMSR_MODEL_ID="$MODEL_ID" \ +uv run opencode-agent --show-trajectory \ + --model-id "$MODEL_ID" \ + "Use generate_failure_mode_sensor_mapping for asset_class pump with failure_modes ['seal leakage'] and sensors ['Pressure sensor', 'Vibration sensor']." + +# 9. benchmark-suite dry run uv run python -m benchmark.scenario_suite_runner \ --scenario-ids benchmarks/scenario_suite/scenarios.txt \ --scenario-root /path/to/scenarios_data \ @@ -391,7 +427,7 @@ uv run python -m benchmark.scenario_suite_runner \ --model-id tokenrouter/MiniMax-M3 \ --dry-run -# 6. benchmark-suite CLI workspace dry run +# 10. benchmark-suite CLI workspace dry run uv run python -m benchmark.scenario_suite_runner \ --scenario-ids benchmarks/scenario_suite/scenarios.txt \ --scenario-root /path/to/scenarios_data \ @@ -435,4 +471,4 @@ uv run python -m benchmark.scenario_suite_runner \ documentation. No MCP servers are modified. The runner generates OpenCode MCP/provider config -at runtime. \ No newline at end of file +at runtime. diff --git a/src/couchdb/scenarios_data/shared/fmea/failure_modes_sample.json b/src/couchdb/scenarios_data/shared/fmea/failure_modes_sample.json index 1bae9df9d..6a29c2211 100644 --- a/src/couchdb/scenarios_data/shared/fmea/failure_modes_sample.json +++ b/src/couchdb/scenarios_data/shared/fmea/failure_modes_sample.json @@ -1,8 +1,6 @@ [ { - "_id": "fm:sample-pump", - "doctype": "failure_mode", - "asset_class": "sample pump", + "asset_class": "pump", "failure_modes": [ "seal leakage", "impeller wear" @@ -10,4 +8,4 @@ "source": "synthetic sample", "exhaustive": false } -] +] \ No newline at end of file diff --git a/src/couchdb/tests/test_loader.py b/src/couchdb/tests/test_loader.py index 9b3780d93..4cf56e526 100644 --- a/src/couchdb/tests/test_loader.py +++ b/src/couchdb/tests/test_loader.py @@ -53,5 +53,5 @@ def test_failure_mode_collection_parses_shared_json() -> None: assert cfg["format"] == "json" assert cfg["primary_key"] == ["asset_class"] assert len(docs) == 1 - assert docs[0]["asset_class"] == "sample pump" + assert docs[0]["asset_class"] == "pump" assert docs[0]["failure_modes"] == ["seal leakage", "impeller wear"] diff --git a/src/mcphub/tests/test_workflows.py b/src/mcphub/tests/test_workflows.py new file mode 100644 index 000000000..e9a4c08c1 --- /dev/null +++ b/src/mcphub/tests/test_workflows.py @@ -0,0 +1,38 @@ +from mcphub.workflows import chiller_triage + + +class FakeToolUniverse: + def __init__(self): + self.calls = [] + + def run(self, name, arguments): + self.calls.append((name, arguments)) + if name == "iot.sensors": + return {"sensors": ["Pressure sensor", "Vibration sensor"]} + if name == "fmsr.get_failure_modes": + return { + "asset_class": "pump", + "failure_modes": ["seal leakage", "impeller wear"], + } + if name == "fmsr.generate_failure_mode_sensor_mapping": + return {"metadata": arguments} + raise AssertionError(f"unexpected tool call: {name}") + + +def test_chiller_triage_passes_lists_to_mapping_tool(): + tu = FakeToolUniverse() + + result = chiller_triage(tu, asset_id="Pump-1", raise_work_order=False) + + mapping_call = [ + arguments + for name, arguments in tu.calls + if name == "fmsr.generate_failure_mode_sensor_mapping" + ][0] + assert mapping_call == { + "asset_class": "pump", + "failure_modes": ["seal leakage", "impeller wear"], + "sensors": ["Pressure sensor", "Vibration sensor"], + } + assert result["failure_modes_result"]["asset_class"] == "pump" + assert result["failure_modes"] == ["seal leakage", "impeller wear"] diff --git a/src/mcphub/workflows.py b/src/mcphub/workflows.py index 7ea6be99c..141f372d2 100644 --- a/src/mcphub/workflows.py +++ b/src/mcphub/workflows.py @@ -10,23 +10,39 @@ ``REGISTERED``. """ +import re + REGISTERED = ["chiller_triage", "sensor_inventory_gap"] +def _asset_class_from_asset_id(asset_id): + key = re.sub(r"\d+", "", asset_id or "") + key = re.sub(r"[_\-]+", " ", key) + return re.sub(r"\s+", " ", key).strip().lower() + + def chiller_triage(tu, asset_id, site="MAIN", raise_work_order=True, priority="2"): """Sensors -> failure modes -> mode/sensor mapping -> (optional) work order.""" - sensors = tu.run("iot.sensors", {"site_name": site, "asset_id": asset_id}) - failure_modes = tu.run("fmsr.get_failure_modes", {"asset_name": asset_id}) - mapping = tu.run("fmsr.get_failure_mode_sensor_mapping", { - "asset_name": asset_id, + asset_class = _asset_class_from_asset_id(asset_id) + sensors = _names(tu.run("iot.sensors", {"site_name": site, "asset_id": asset_id})) + failure_modes_result = tu.run("fmsr.get_failure_modes", {"asset_class": asset_class}) + failure_modes = ( + failure_modes_result.get("failure_modes", []) + if isinstance(failure_modes_result, dict) + else failure_modes_result + ) + mapping = tu.run("fmsr.generate_failure_mode_sensor_mapping", { + "asset_class": asset_class, "failure_modes": failure_modes, "sensors": sensors, }) out = { "asset_id": asset_id, + "asset_class": asset_class, "site_name": site, "sensors": sensors, + "failure_modes_result": failure_modes_result, "failure_modes": failure_modes, "failure_mode_sensor_mapping": mapping, } @@ -63,4 +79,4 @@ def _names(value): return list(value.keys()) if isinstance(value, list): return [v if isinstance(v, str) else v.get("name", str(v)) for v in value] - return [value] \ No newline at end of file + return [value] diff --git a/src/servers/fmsr/failure_modes.yaml b/src/servers/fmsr/failure_modes.yaml deleted file mode 100644 index 16712b8ee..000000000 --- a/src/servers/fmsr/failure_modes.yaml +++ /dev/null @@ -1,15 +0,0 @@ -chiller: - - "Compressor Overheating: Failed due to Normal wear, overheating" - - "Heat Exchangers: Fans: Degraded motor or worn bearing due to Normal use" - - "Evaporator Water side fouling" - - "Condenser Water side fouling" - - "Condenser Improper water side flow rate" - - "Purge Unit Excessive purge" - - "Refrigerant Operated Control Valve Failed spring" - -ahu: - - "Pressure Regulators Diaphragm failure" - - "Steam Heating Coils Air side fouling" - - "Belts or sheaves Wear" - - "Improper switch position" - - "Solenoid Valves Bound due to hardened grease" diff --git a/src/servers/fmsr/main.py b/src/servers/fmsr/main.py index d858f1587..f91c7da99 100644 --- a/src/servers/fmsr/main.py +++ b/src/servers/fmsr/main.py @@ -1,29 +1,16 @@ -"""FMSR (Failure Mode and Sensor Reasoning) MCP Server. - -Exposes two tools: - get_failure_modes – lists failure modes for an asset - get_failure_mode_sensor_mapping – returns bidirectional FM↔sensor relevancy mapping - -For chillers and AHUs get_failure_modes returns a curated hardcoded list. -For any other asset type the LLM is queried as a fallback. -The mapping tool always calls the LLM to determine per-pair relevancy. - -LLM backend is configured via the FMSR_MODEL_ID environment variable -(default: ``watsonx/meta-llama/llama-3-3-70b-instruct``). Any model string -supported by litellm works — the provider is encoded in the prefix. -""" +"""Failure mode and sensor reasoning for industrial asset classes.""" from __future__ import annotations import logging import os import re -from pathlib import Path -from typing import Dict, List, Union +from typing import Dict, List, Optional, Union from concurrent.futures import ThreadPoolExecutor, as_completed -import yaml +import couchdb3 +from couchdb3.exceptions import NotFoundError from dotenv import load_dotenv from mcp.server.fastmcp import FastMCP from pydantic import BaseModel @@ -37,46 +24,113 @@ logger = logging.getLogger("fmsr-mcp-server") -# ── Hardcoded asset data ────────────────────────────────────────────────────── +# ── Database stores ─────────────────────────────────────────────────────────── +# Under AssetOpsBench's loader, database name == collection key. Failure modes +# live in the 'failure_mode' database; generic catalog lookups live in the +# utilities MCP server. -_FAILURE_MODES_FILE = Path(__file__).parent / "failure_modes.yaml" -with _FAILURE_MODES_FILE.open() as _f: - _ASSET_FAILURE_MODES: dict[str, list[str]] = yaml.safe_load(_f) +COUCHDB_URL = os.environ.get("COUCHDB_URL", "http://localhost:5984") +COUCHDB_USERNAME = os.environ.get("COUCHDB_USERNAME", "admin") +COUCHDB_PASSWORD = os.environ.get("COUCHDB_PASSWORD", "password") +FAILURE_MODE_DBNAME = os.environ.get("FAILURE_MODE_DBNAME", "failure_mode") -# ── Prompt templates ────────────────────────────────────────────────────────── +def _connect(dbname): + try: + h = couchdb3.Database( + dbname, + url=COUCHDB_URL, + user=COUCHDB_USERNAME, + password=COUCHDB_PASSWORD, + ) + logger.info("Connected to database: %s", dbname) + return h + except Exception as e: # noqa: BLE001 + logger.error("Failed to connect to database '%s': %s", dbname, e) + return None -_ASSET2FM_PROMPT = ( - "What are different failure modes for asset {asset_name}?\n" - "Your response should be a numbered list with each failure mode on a new line. " - "Please only list the failure mode name.\n" - "For example: \n\n1. foo\n\n2. bar\n\n3. baz" -) + +fm_db = _connect(FAILURE_MODE_DBNAME) + + +def _asset_class_key(asset_class: str) -> str: + """Normalise an asset class to the database key format ('Hydraulic_Pump' -> 'hydraulic pump').""" + key = re.sub(r"\d+", "", asset_class or "") + key = re.sub(r"[_\-]+", " ", key) + return re.sub(r"\s+", " ", key).strip().lower() + + +def _known_asset_classes(limit: int = 10) -> List[str]: + """Return known asset classes from the failure_mode collection for error guidance.""" + if not fm_db: + return [] + try: + res = fm_db.find({}, fields=["asset_class"], limit=limit) + except Exception: # noqa: BLE001 + return [] + classes = [ + doc.get("asset_class") + for doc in res.get("docs", []) + if isinstance(doc.get("asset_class"), str) and doc.get("asset_class") + ] + return sorted(dict.fromkeys(classes)) + + +def _missing_asset_class_error(original: str, normalized: str) -> ErrorResult: + message = ( + f"no failure_mode record for asset_class '{normalized}' in database. " + f"Input was normalized from {original!r}; check that asset_class matches a stored class." + ) + known = _known_asset_classes() + if known: + message += f" Available asset_class values include: {', '.join(known)}." + return ErrorResult(error=message) + + +def _is_not_found_error(exc: Exception) -> bool: + if isinstance(exc, (KeyError, NotFoundError)): + return True + response = getattr(exc, "response", None) + status_code = getattr(response, "status_code", None) or getattr( + exc, "status_code", None + ) + return status_code == 404 + + +# ── Prompt templates ────────────────────────────────────────────────────────── _RELEVANCY_PROMPT = ( - "For the asset {asset_name}, if the failure {failure_mode} occurs, " - "can sensor {sensor} help monitor or detect the failure for {asset_name}?\n" - "Provide the answer in the first line and reason in the second line. " - "If the answer is Yes, provide the temporal behaviour of the sensor " - "when the failure occurs in the third line." + "For asset class {asset_class}, if the failure {failure_mode} occurs, " + "can sensor {sensor} help monitor or detect the failure for assets of class {asset_class}?\n" + "Provide the answer in the first line and reason in the second line." ) +_FAILURE_MODE_PROMPT = ( + "List up to {max_modes} common failure modes for asset class {asset_class}.\n" + "Return only failure mode names, one per line." +) -# ── Output parsers ──────────────────────────────────────────────────────────── +_FAILURE_MODE_EXTEND_PROMPT = ( + "Asset class {asset_class} already has these stored failure modes:\n{stored_modes}\n\n" + "List up to {max_modes} additional failure modes for asset class {asset_class} " + "that are not already in the stored list.\n" + "Return only failure mode names, one per line." +) -def _parse_numbered_list(text: str) -> list[str]: - """Parse a numbered list response into a plain list of strings.""" - items = [] +def _parse_failure_mode_list(text: str) -> List[str]: + items: List[str] = [] for line in text.splitlines(): - m = re.match(r"^\d+[\.\)]\s*(.+)", line.strip()) - if m: - items.append(m.group(1).strip()) + item = line.strip() + if not item: + continue + item = re.sub(r"^\s*(?:[-*•]|\d+[\.\)])\s*", "", item).strip() + if item: + items.append(item) return items def _parse_relevancy(text: str) -> dict: - """Parse a 3-line relevancy response into {answer, reason, temporal_behavior}.""" lines = [ln for ln in text.strip().splitlines() if ln.strip()] if lines and lines[0].lower().startswith("yes"): answer = "Yes" @@ -85,27 +139,26 @@ def _parse_relevancy(text: str) -> dict: else: answer = "Unknown" reason = lines[1] if len(lines) >= 2 else "Unknown" - temporal = lines[2] if (answer == "Yes" and len(lines) >= 3) else "Unknown" - return {"answer": answer, "reason": reason, "temporal_behavior": temporal} + return {"answer": answer, "reason": reason} # ── LLM backend (lazy init; graceful degradation if creds are absent) ───────── _DEFAULT_MODEL_ID = "watsonx/meta-llama/llama-3-3-70b-instruct" _MAX_RETRIES = 3 +_MODEL_ID = os.environ.get("FMSR_MODEL_ID", _DEFAULT_MODEL_ID) def _build_llm(): from llm import make_backend - model_id = os.environ.get("FMSR_MODEL_ID", _DEFAULT_MODEL_ID) - if model_id.startswith("watsonx/"): + if _MODEL_ID.startswith("watsonx/"): missing = [ v for v in ("WATSONX_APIKEY", "WATSONX_PROJECT_ID") if not os.environ.get(v) ] if missing: raise RuntimeError(f"Missing env vars for WatsonX: {missing}") - elif model_id.startswith("tokenrouter/"): + elif _MODEL_ID.startswith("tokenrouter/"): missing = [ v for v in ("TOKENROUTER_API_KEY", "TOKENROUTER_BASE_URL") @@ -119,51 +172,48 @@ def _build_llm(): ] if missing: raise RuntimeError(f"Missing env vars for LiteLLM: {missing}") - return make_backend(model_id) + return make_backend(_MODEL_ID) try: _llm = _build_llm() _llm_available = True -except Exception as _e: - logger.warning("LLM unavailable (FMSR will use curated data only): %s", _e) +except Exception as _e: # noqa: BLE001 + logger.warning("LLM unavailable (generate_* tools disabled): %s", _e) _llm = None _llm_available = False -# ── LLM call helpers with retry ─────────────────────────────────────────────── - -_asset2fm_cache: dict[str, list[str]] = {} - - -def _call_asset2fm(asset_name: str) -> list[str]: - """Query the LLM for failure modes of an asset. Retries up to _MAX_RETRIES times. - Results are cached to avoid redundant LLM calls for the same asset.""" - if asset_name in _asset2fm_cache: - return _asset2fm_cache[asset_name] - - prompt = _ASSET2FM_PROMPT.format(asset_name=asset_name) +def _call_relevancy(asset_class: str, failure_mode: str, sensor: str) -> dict: + prompt = _RELEVANCY_PROMPT.format( + asset_class=asset_class, failure_mode=failure_mode, sensor=sensor + ) last_exc: Exception | None = None for _ in range(_MAX_RETRIES): try: - result = _parse_numbered_list(_llm.generate(prompt)) - _asset2fm_cache[asset_name] = result - return result - except Exception as exc: + return _parse_relevancy(_llm.generate(prompt)) + except Exception as exc: # noqa: BLE001 last_exc = exc raise last_exc -def _call_relevancy(asset_name: str, failure_mode: str, sensor: str) -> dict: - """Query the LLM for FM↔sensor relevancy. Retries up to _MAX_RETRIES times.""" - prompt = _RELEVANCY_PROMPT.format( - asset_name=asset_name, failure_mode=failure_mode, sensor=sensor +def _call_failure_mode_generation( + asset_class: str, known: List[str], max_modes: int +) -> List[str]: + prompt = ( + _FAILURE_MODE_EXTEND_PROMPT.format( + asset_class=asset_class, + stored_modes="\n".join(f"- {mode}" for mode in known), + max_modes=max_modes, + ) + if known + else _FAILURE_MODE_PROMPT.format(asset_class=asset_class, max_modes=max_modes) ) last_exc: Exception | None = None for _ in range(_MAX_RETRIES): try: - return _parse_relevancy(_llm.generate(prompt)) - except Exception as exc: + return _parse_failure_mode_list(_llm.generate(prompt)) + except Exception as exc: # noqa: BLE001 last_exc = exc raise last_exc @@ -176,21 +226,41 @@ class ErrorResult(BaseModel): class FailureModesResult(BaseModel): - asset_name: str + asset_class: str failure_modes: List[str] + exhaustive: bool = False # the stored list is not claimed to be complete + source: Optional[str] = None # provenance: ISO / curated / LLM: + + +class GenerateFailureModesResult(BaseModel): + asset_class: str + known: List[str] + generated: List[str] + failure_modes: List[str] + source: str + message: str + + +class AddFailureModesResult(BaseModel): + asset_class: str + added: List[str] + failure_modes: List[str] + total: int + exhaustive: bool + source: Optional[str] = None + message: str class RelevancyEntry(BaseModel): - asset_name: str + asset_class: str failure_mode: str sensor: str relevancy_answer: str relevancy_reason: str - temporal_behavior: str class MappingMetadata(BaseModel): - asset_name: str + asset_class: str failure_modes: List[str] sensors: List[str] @@ -206,51 +276,238 @@ class FailureModeSensorMappingResult(BaseModel): mcp = FastMCP( "fmsr", - instructions="Failure mode and sensor reasoning: get failure modes for assets and determine which sensors can detect each failure.", + instructions="Failure mode and sensor reasoning for industrial asset classes.", ) @mcp.tool(title="Get Failure Modes") -def get_failure_modes(asset_name: str) -> Union[FailureModesResult, ErrorResult]: - """Returns a list of known failure modes for the given asset. - For chillers and AHUs returns a curated list. For other assets queries the LLM.""" - asset_key = re.sub(r"\d+", "", asset_name).strip().lower() - if not asset_key or asset_key == "none": - return ErrorResult(error="asset_name is required") - - if asset_key in _ASSET_FAILURE_MODES: +def get_failure_modes(asset_class: str) -> Union[FailureModesResult, ErrorResult]: + """READ the known failure modes for an asset class. + + Args: + asset_class: Asset class to look up, such as "pump". Case, whitespace, + digits, underscores, and hyphens are normalized before querying. + """ + raw_asset_class = asset_class + key = _asset_class_key(asset_class) + if not key or key == "none": + return ErrorResult(error="asset_class is required") + try: + d = _find_failure_mode_doc(key) + if d is None: + return _missing_asset_class_error(raw_asset_class, key) return FailureModesResult( - asset_name=asset_name, - failure_modes=_ASSET_FAILURE_MODES[asset_key], + asset_class=d.get("asset_class", key), + failure_modes=d.get("failure_modes", []), + exhaustive=d.get("exhaustive", False), + source=d.get("source"), ) + except Exception as exc: # noqa: BLE001 + logger.error("get_failure_modes failed: %s", exc) + return ErrorResult(error=str(exc)) + +def _find_failure_mode_doc(asset_class: str) -> Optional[dict]: + """Return the stored failure-mode doc for an asset class, or None.""" + if not fm_db: + raise RuntimeError("database not connected") + key = _asset_class_key(asset_class) + try: + d = fm_db.get(f"fm:{key}", check=True) + except Exception as exc: # noqa: BLE001 + if _is_not_found_error(exc): + d = None + else: + raise RuntimeError( + f"database lookup failed for asset_class '{key}': {exc}" + ) from exc + try: + if d is None: + res = fm_db.find({"asset_class": key}, limit=1) + docs = res["docs"] + if docs: + d = docs[0] + return d + except Exception as exc: # noqa: BLE001 + raise RuntimeError( + f"database lookup failed for asset_class '{key}': {exc}" + ) from exc + + +def _known_failure_modes(asset_class: str) -> List[str]: + """Return stored failure modes for an asset class, or [] if none are available.""" + d = _find_failure_mode_doc(asset_class) + if d is None: + return [] + return [ + mode.strip() + for mode in d.get("failure_modes", []) + if isinstance(mode, str) and mode.strip() + ] + + +@mcp.tool(title="Generate Failure Modes") +def generate_failure_modes( + asset_class: str, + max_modes: int = 10, +) -> Union[GenerateFailureModesResult, ErrorResult]: + """GENERATE a new or extended failure-mode list for an asset class. + + This tool does not write to the database. If the normalized `asset_class` + exists in the database, the current stored failure modes are used as context + and the LLM generates additional modes. If no stored modes exist, the LLM + generates a new list from scratch. + + Args: + asset_class: Asset class to reason about, such as "pump". Case, + whitespace, digits, underscores, and hyphens are normalized before + prompting the LLM. + max_modes: Maximum number of new failure modes to request from the LLM. + """ + key = _asset_class_key(asset_class) + if not key or key == "none": + return ErrorResult(error="asset_class is required") + if max_modes <= 0: + return ErrorResult(error="max_modes must be greater than 0") if not _llm_available: - return ErrorResult(error="LLM unavailable and asset not in local database") + return ErrorResult(error="LLM unavailable") try: - result = _call_asset2fm(asset_name) - return FailureModesResult(asset_name=asset_name, failure_modes=result) - except Exception as exc: - logger.error("_call_asset2fm failed: %s", exc) + base = _known_failure_modes(key) + base = [mode.strip() for mode in base if mode and mode.strip()] + raw = _call_failure_mode_generation(key, base, max_modes) + seen = {mode.lower() for mode in base} + generated: List[str] = [] + for mode in raw: + candidate = mode.strip() + normalized = candidate.lower() + if candidate and normalized not in seen: + seen.add(normalized) + generated.append(candidate) + if len(generated) > max_modes: + generated = generated[:max_modes] + return GenerateFailureModesResult( + asset_class=key, + known=base, + generated=generated, + failure_modes=base + generated, + source=f"LLM:{_MODEL_ID}", + message=( + f"generated {len(generated)} new failure mode(s) for asset_class '{key}' " + f"using {len(base)} stored mode(s) as context; nothing was persisted." + ), + ) + except Exception as exc: # noqa: BLE001 + logger.error("generate_failure_modes failed: %s", exc) return ErrorResult(error=str(exc)) -@mcp.tool(title="Get Failure Mode Sensor Mapping") -def get_failure_mode_sensor_mapping( - asset_name: str, +@mcp.tool(title="Add Failure Modes") +def add_failure_modes( + asset_class: str, + failure_modes: List[str], + exhaustive: Optional[bool] = None, + source: Optional[str] = None, +) -> Union[AddFailureModesResult, ErrorResult]: + """WRITE failure modes for an asset class to the database. + + Existing modes are preserved, incoming modes are merged case-insensitively, + and only newly added modes are reported. Use this after curated or generated + modes should become available to future `get_failure_modes` calls. + + Args: + asset_class: Asset class to update, such as "pump". Case, whitespace, + digits, underscores, and hyphens are normalized before writing. + failure_modes: Failure modes to add for the asset class. + exhaustive: Set true only when the stored list is believed complete. If + omitted, the existing value is preserved; new records default false. + source: Optional provenance for the stored list. + """ + key = _asset_class_key(asset_class) + if not key or key == "none": + return ErrorResult(error="asset_class is required") + incoming = [ + mode.strip() + for mode in (failure_modes or []) + if isinstance(mode, str) and mode.strip() + ] + if not incoming: + return ErrorResult(error="failure_modes list is required") + if not fm_db: + return ErrorResult(error="database not connected") + + try: + doc_id = f"fm:{key}" + doc = _find_failure_mode_doc(key) + existing = [ + mode.strip() + for mode in (doc or {}).get("failure_modes", []) + if isinstance(mode, str) and mode.strip() + ] + seen = set() + merged: List[str] = [] + for mode in existing: + normalized = mode.lower() + if normalized not in seen: + seen.add(normalized) + merged.append(mode) + + added: List[str] = [] + for mode in incoming: + normalized = mode.lower() + if normalized not in seen: + seen.add(normalized) + merged.append(mode) + added.append(mode) + + if doc is None: + doc = {"_id": doc_id, "asset_class": key} + stored_exhaustive = False + else: + doc.setdefault("_id", doc_id) + doc["asset_class"] = key + stored_exhaustive = bool(doc.get("exhaustive", False)) + doc["failure_modes"] = merged + doc["exhaustive"] = stored_exhaustive if exhaustive is None else exhaustive + doc["source"] = source or doc.get("source") or "user" + fm_db.save(doc) + + return AddFailureModesResult( + asset_class=key, + added=added, + failure_modes=merged, + total=len(merged), + exhaustive=doc["exhaustive"], + source=doc.get("source"), + message=( + f"added {len(added)} new failure mode(s) to asset_class '{key}' " + f"({len(merged)} total)." + ), + ) + except Exception as exc: # noqa: BLE001 + logger.error("add_failure_modes failed: %s", exc) + return ErrorResult(error=str(exc)) + + +@mcp.tool(title="Generate Failure Mode Sensor Mapping") +def generate_failure_mode_sensor_mapping( + asset_class: str, failure_modes: List[str], sensors: List[str], ) -> Union[FailureModeSensorMappingResult, ErrorResult]: - """For each (failure_mode, sensor) pair determines whether the sensor can detect - the failure. Returns a bidirectional mapping (fm→sensors, sensor→fms) plus - the full per-pair relevancy details. - - Note: one LLM call is made per (failure_mode, sensor) pair sequentially. - Keep both lists small (e.g. ≤5 failure modes, ≤10 sensors) to avoid long - runtimes. For a chiller with 7 failure modes and 20+ sensors the call will - take several minutes.""" - if not asset_name: - return ErrorResult(error="asset_name is required") + """GENERATE, for each (failure_mode, sensor) pair, whether the sensor can detect the failure + (one LLM call per pair). Returns a bidirectional mapping (fm→sensors, sensor→fms) plus per-pair + relevancy details. Keep both lists small (e.g. ≤5 failure modes, ≤10 sensors) to bound runtime. + + Args: + asset_class: Asset class to reason about, such as "pump". Case, whitespace, + digits, underscores, and hyphens are normalized before prompting the LLM. + failure_modes: Failure modes for the asset class. + sensors: Sensor names to evaluate for detection relevance. + """ + key = _asset_class_key(asset_class) + if not key or key == "none": + return ErrorResult(error="asset_class is required") if not failure_modes: return ErrorResult(error="failure_modes list is required") if not sensors: @@ -261,38 +518,35 @@ def get_failure_mode_sensor_mapping( full_relevancy: List[RelevancyEntry] = [] fm2sensor: Dict[str, List[str]] = {} sensor2fm: Dict[str, List[str]] = {} - try: pairs = [(s, fm) for s in sensors for fm in failure_modes] with ThreadPoolExecutor() as executor: futures = { - executor.submit(_call_relevancy, asset_name, fm, s): (s, fm) + executor.submit(_call_relevancy, key, fm, s): (s, fm) for s, fm in pairs } for future in as_completed(futures): s, fm = futures[future] gen = future.result() - entry = RelevancyEntry( - asset_name=asset_name, - failure_mode=fm, - sensor=s, - relevancy_answer=gen["answer"], - relevancy_reason=gen["reason"], - temporal_behavior=gen["temporal_behavior"], + full_relevancy.append( + RelevancyEntry( + asset_class=key, + failure_mode=fm, + sensor=s, + relevancy_answer=gen["answer"], + relevancy_reason=gen["reason"], + ) ) - full_relevancy.append(entry) if "yes" in gen["answer"].lower(): fm2sensor.setdefault(fm, []).append(s) sensor2fm.setdefault(s, []).append(fm) - except Exception as exc: + except Exception as exc: # noqa: BLE001 logger.error("_call_relevancy failed: %s", exc) return ErrorResult(error=str(exc)) return FailureModeSensorMappingResult( metadata=MappingMetadata( - asset_name=asset_name, - failure_modes=failure_modes, - sensors=sensors, + asset_class=key, failure_modes=failure_modes, sensors=sensors ), fm2sensor=fm2sensor, sensor2fm=sensor2fm, diff --git a/src/servers/fmsr/tests/conftest.py b/src/servers/fmsr/tests/conftest.py index b2b10e0db..37899e2fb 100644 --- a/src/servers/fmsr/tests/conftest.py +++ b/src/servers/fmsr/tests/conftest.py @@ -16,6 +16,53 @@ async def call_tool(mcp_instance, tool_name: str, args: dict) -> dict: return json.loads(contents[0].text) +class FakeDatabase: + def __init__(self, docs=None): + self.docs = {doc["_id"]: dict(doc) for doc in docs or []} + + def find(self, selector, fields=None, limit=None): + docs = [doc for doc in self.docs.values() if self._matches(doc, selector)] + if limit is not None: + docs = docs[:limit] + if fields is not None: + docs = [ + {field: doc[field] for field in fields if field in doc} for doc in docs + ] + return {"docs": docs} + + def get(self, doc_id, *, check=False, default_value=None): + if doc_id not in self.docs: + if check: + raise KeyError(doc_id) + return default_value + return dict(self.docs[doc_id]) + + def save(self, doc): + self.docs[doc["_id"]] = dict(doc) + + @staticmethod + def _matches(doc, selector): + for key, expected in selector.items(): + if isinstance(expected, dict) and "$exists" in expected: + exists = key in doc + if exists != expected["$exists"]: + return False + elif doc.get(key) != expected: + return False + return True + + +class BrokenDatabase(FakeDatabase): + def find(self, selector, fields=None, limit=None): + raise RuntimeError("database read failed") + + def get(self, doc_id, *, check=False, default_value=None): + raise RuntimeError("database read failed") + + def save(self, doc): + raise RuntimeError("database write failed") + + @pytest.fixture def no_llm(): """Simulate missing WatsonX credentials.""" @@ -23,6 +70,37 @@ def no_llm(): yield +@pytest.fixture +def fake_fm_db(): + db = FakeDatabase( + [ + { + "_id": "fm:pump", + "asset_class": "pump", + "failure_modes": ["seal leakage", "impeller wear"], + "exhaustive": False, + "source": "synthetic sample", + }, + ] + ) + with patch("servers.fmsr.main.fm_db", db): + yield db + + +@pytest.fixture +def empty_fm_db(): + db = FakeDatabase() + with patch("servers.fmsr.main.fm_db", db): + yield db + + +@pytest.fixture +def broken_fm_db(): + db = BrokenDatabase() + with patch("servers.fmsr.main.fm_db", db): + yield db + + @pytest.fixture def mock_relevancy_chain(): """Patch _call_relevancy so it always returns 'Yes' without calling the LLM.""" @@ -30,7 +108,6 @@ def mock_relevancy_chain(): return_value={ "answer": "Yes", "reason": "Relevant sensor", - "temporal_behavior": "Increases", } ) with patch("servers.fmsr.main._call_relevancy", mock): @@ -39,9 +116,15 @@ def mock_relevancy_chain(): @pytest.fixture -def mock_asset2fm_chain(): - """Patch _call_asset2fm to return a fixed failure mode list.""" - mock = MagicMock(return_value=["Fan Failure", "Belt Wear"]) - with patch("servers.fmsr.main._call_asset2fm", mock): +def mock_failure_mode_generation(): + """Patch failure-mode generation so tests do not call the LLM.""" + mock = MagicMock( + return_value=[ + "bearing wear", + "seal leakage", + "motor overheating", + ] + ) + with patch("servers.fmsr.main._call_failure_mode_generation", mock): with patch("servers.fmsr.main._llm_available", True): yield mock diff --git a/src/servers/fmsr/tests/test_tools.py b/src/servers/fmsr/tests/test_tools.py index cece28473..352fd50be 100644 --- a/src/servers/fmsr/tests/test_tools.py +++ b/src/servers/fmsr/tests/test_tools.py @@ -1,148 +1,388 @@ -"""Tests for FMSR MCP server tools. - -Unit tests use mocked LLM chains; integration tests require live WatsonX -credentials (skipped unless WATSONX_APIKEY is set). -""" +"""Tests for FMSR MCP server tools.""" import pytest + from servers.fmsr.main import mcp + from .conftest import call_tool, requires_watsonx -# --------------------------------------------------------------------------- -# get_failure_modes -# --------------------------------------------------------------------------- +class TestGetFailureModes: + @pytest.mark.anyio + async def test_reads_failure_modes_from_db(self, fake_fm_db): + data = await call_tool(mcp, "get_failure_modes", {"asset_class": "pump"}) + assert data["asset_class"] == "pump" + assert data["failure_modes"] == ["seal leakage", "impeller wear"] + assert data["exhaustive"] is False + assert data["source"] == "synthetic sample" -class TestGetFailureModes: @pytest.mark.anyio - async def test_chiller_returns_hardcoded(self): - data = await call_tool(mcp, "get_failure_modes", {"asset_name": "chiller"}) - assert "failure_modes" in data - assert len(data["failure_modes"]) == 7 - assert any("Compressor" in fm for fm in data["failure_modes"]) + async def test_asset_class_case_normalized(self, fake_fm_db): + data = await call_tool(mcp, "get_failure_modes", {"asset_class": "PUMP"}) + + assert data["asset_class"] == "pump" + assert data["failure_modes"] == ["seal leakage", "impeller wear"] @pytest.mark.anyio - async def test_chiller_number_stripped(self): - """'Chiller 6' normalises to 'chiller' for the lookup.""" - data = await call_tool(mcp, "get_failure_modes", {"asset_name": "Chiller 6"}) - assert "failure_modes" in data - assert len(data["failure_modes"]) == 7 + async def test_asset_class_spacing_normalized(self, fake_fm_db): + data = await call_tool(mcp, "get_failure_modes", {"asset_class": " PUMP "}) + + assert data["asset_class"] == "pump" + assert data["failure_modes"] == ["seal leakage", "impeller wear"] @pytest.mark.anyio - async def test_ahu_returns_hardcoded(self): - data = await call_tool(mcp, "get_failure_modes", {"asset_name": "ahu"}) - assert "failure_modes" in data - assert len(data["failure_modes"]) == 5 + async def test_empty_asset_class_returns_error(self, fake_fm_db): + data = await call_tool(mcp, "get_failure_modes", {"asset_class": ""}) + + assert data == {"error": "asset_class is required"} @pytest.mark.anyio - async def test_empty_asset_name_returns_error(self): - data = await call_tool(mcp, "get_failure_modes", {"asset_name": ""}) + async def test_missing_asset_class_returns_guidance(self, fake_fm_db): + data = await call_tool(mcp, "get_failure_modes", {"asset_class": "bad-pump-1"}) + assert "error" in data + assert "no failure_mode record for asset_class 'bad pump'" in data["error"] + assert "Input was normalized from 'bad-pump-1'" in data["error"] + assert "Available asset_class values include: pump" in data["error"] @pytest.mark.anyio - async def test_unknown_asset_no_llm(self, no_llm): - data = await call_tool(mcp, "get_failure_modes", {"asset_name": "Pump"}) - assert "error" in data + async def test_db_unavailable_returns_error(self, monkeypatch): + monkeypatch.setattr("servers.fmsr.main.fm_db", None) + + data = await call_tool(mcp, "get_failure_modes", {"asset_class": "pump"}) + + assert data == {"error": "database not connected"} + + @pytest.mark.anyio + async def test_database_read_error_returns_error(self, broken_fm_db): + data = await call_tool(mcp, "get_failure_modes", {"asset_class": "pump"}) + + assert data == { + "error": "database lookup failed for asset_class 'pump': database read failed" + } + + +class TestGenerateFailureModes: + @pytest.mark.anyio + async def test_extends_failure_modes_from_db( + self, fake_fm_db, mock_failure_mode_generation + ): + data = await call_tool( + mcp, + "generate_failure_modes", + {"asset_class": "Pump", "max_modes": 5}, + ) + + assert data["asset_class"] == "pump" + assert data["known"] == ["seal leakage", "impeller wear"] + assert data["generated"] == ["bearing wear", "motor overheating"] + assert data["failure_modes"] == [ + "seal leakage", + "impeller wear", + "bearing wear", + "motor overheating", + ] + assert data["source"].startswith("LLM:") + assert "nothing was persisted" in data["message"] + mock_failure_mode_generation.assert_called_once_with( + "pump", ["seal leakage", "impeller wear"], 5 + ) @pytest.mark.anyio - async def test_unknown_asset_llm_fallback(self, mock_asset2fm_chain): - data = await call_tool(mcp, "get_failure_modes", {"asset_name": "Pump"}) - assert "failure_modes" in data - assert data["failure_modes"] == ["Fan Failure", "Belt Wear"] - mock_asset2fm_chain.assert_called_once_with("Pump") + async def test_generates_from_scratch_for_missing_db_record( + self, empty_fm_db, mock_failure_mode_generation + ): + data = await call_tool( + mcp, + "generate_failure_modes", + {"asset_class": "compressor", "max_modes": 3}, + ) + + assert data["asset_class"] == "compressor" + assert data["known"] == [] + assert data["generated"] == [ + "bearing wear", + "seal leakage", + "motor overheating", + ] + mock_failure_mode_generation.assert_called_once_with("compressor", [], 3) + + @pytest.mark.anyio + async def test_database_read_error_returns_error( + self, broken_fm_db, mock_failure_mode_generation + ): + data = await call_tool( + mcp, + "generate_failure_modes", + {"asset_class": "pump", "max_modes": 3}, + ) + + assert data == { + "error": "database lookup failed for asset_class 'pump': database read failed" + } + mock_failure_mode_generation.assert_not_called() + + @pytest.mark.anyio + async def test_empty_asset_class_returns_error(self, mock_failure_mode_generation): + data = await call_tool( + mcp, + "generate_failure_modes", + {"asset_class": "", "max_modes": 3}, + ) + + assert data == {"error": "asset_class is required"} + + @pytest.mark.anyio + async def test_invalid_max_modes_returns_error(self, mock_failure_mode_generation): + data = await call_tool( + mcp, + "generate_failure_modes", + {"asset_class": "pump", "max_modes": 0}, + ) + + assert data == {"error": "max_modes must be greater than 0"} + + @pytest.mark.anyio + async def test_llm_unavailable_returns_error(self, no_llm): + data = await call_tool( + mcp, + "generate_failure_modes", + {"asset_class": "pump", "max_modes": 3}, + ) + + assert data == {"error": "LLM unavailable"} @requires_watsonx @pytest.mark.anyio - async def test_unknown_asset_integration(self): - data = await call_tool(mcp, "get_failure_modes", {"asset_name": "Boiler"}) - assert "failure_modes" in data - assert len(data["failure_modes"]) > 0 + async def test_integration(self): + data = await call_tool( + mcp, + "generate_failure_modes", + { + "asset_class": "pump", + "max_modes": 2, + }, + ) + + assert "generated" in data + assert len(data["generated"]) <= 2 + + +class TestAddFailureModes: + @pytest.mark.anyio + async def test_merges_with_existing_modes(self, fake_fm_db): + data = await call_tool( + mcp, + "add_failure_modes", + { + "asset_class": "Pump-1", + "failure_modes": ["impeller wear", "bearing wear"], + "exhaustive": True, + "source": "unit-test", + }, + ) + + assert data["asset_class"] == "pump" + assert data["added"] == ["bearing wear"] + assert data["failure_modes"] == [ + "seal leakage", + "impeller wear", + "bearing wear", + ] + assert data["total"] == 3 + assert data["exhaustive"] is True + assert data["source"] == "unit-test" + assert fake_fm_db.docs["fm:pump"]["failure_modes"] == [ + "seal leakage", + "impeller wear", + "bearing wear", + ] + assert fake_fm_db.docs["fm:pump"]["exhaustive"] is True + assert fake_fm_db.docs["fm:pump"]["source"] == "unit-test" + + @pytest.mark.anyio + async def test_omitted_exhaustive_preserves_existing_value(self, fake_fm_db): + fake_fm_db.docs["fm:pump"]["exhaustive"] = True + + data = await call_tool( + mcp, + "add_failure_modes", + { + "asset_class": "pump", + "failure_modes": ["bearing wear"], + }, + ) + + assert data["exhaustive"] is True + assert fake_fm_db.docs["fm:pump"]["exhaustive"] is True + + @pytest.mark.anyio + async def test_creates_new_asset_class_record(self, empty_fm_db): + data = await call_tool( + mcp, + "add_failure_modes", + { + "asset_class": "Gearbox-1", + "failure_modes": ["gear tooth wear", "bearing wear"], + }, + ) + + assert data["asset_class"] == "gearbox" + assert data["added"] == ["gear tooth wear", "bearing wear"] + assert data["failure_modes"] == ["gear tooth wear", "bearing wear"] + assert data["total"] == 2 + assert data["exhaustive"] is False + assert data["source"] == "user" + assert empty_fm_db.docs["fm:gearbox"]["asset_class"] == "gearbox" + assert empty_fm_db.docs["fm:gearbox"]["failure_modes"] == [ + "gear tooth wear", + "bearing wear", + ] + + @pytest.mark.anyio + async def test_empty_asset_class_returns_error(self, fake_fm_db): + data = await call_tool( + mcp, + "add_failure_modes", + {"asset_class": "", "failure_modes": ["bearing wear"]}, + ) + assert data == {"error": "asset_class is required"} -# --------------------------------------------------------------------------- -# get_failure_mode_sensor_mapping -# --------------------------------------------------------------------------- + @pytest.mark.anyio + async def test_empty_failure_modes_returns_error(self, fake_fm_db): + data = await call_tool( + mcp, + "add_failure_modes", + {"asset_class": "pump", "failure_modes": []}, + ) + + assert data == {"error": "failure_modes list is required"} + + @pytest.mark.anyio + async def test_db_unavailable_returns_error(self, monkeypatch): + monkeypatch.setattr("servers.fmsr.main.fm_db", None) + + data = await call_tool( + mcp, + "add_failure_modes", + {"asset_class": "pump", "failure_modes": ["bearing wear"]}, + ) + + assert data == {"error": "database not connected"} + + @pytest.mark.anyio + async def test_database_read_error_returns_error(self, broken_fm_db): + data = await call_tool( + mcp, + "add_failure_modes", + {"asset_class": "pump", "failure_modes": ["bearing wear"]}, + ) + + assert data == { + "error": "database lookup failed for asset_class 'pump': database read failed" + } _FAILURE_MODES = ["Compressor Overheating", "Condenser Water side fouling"] _SENSORS = ["Chiller 6 Power Input", "Chiller 6 Supply Temperature"] -class TestGetFailureModeSensorMapping: +class TestGenerateFailureModeSensorMapping: @pytest.mark.anyio async def test_returns_expected_keys(self, mock_relevancy_chain): data = await call_tool( mcp, - "get_failure_mode_sensor_mapping", + "generate_failure_mode_sensor_mapping", { - "asset_name": "Chiller 6", + "asset_class": "chiller", "failure_modes": _FAILURE_MODES, "sensors": _SENSORS, }, ) + assert "fm2sensor" in data assert "sensor2fm" in data assert "full_relevancy" in data - assert data["metadata"]["asset_name"] == "Chiller 6" + assert data["metadata"]["asset_class"] == "chiller" + assert data["full_relevancy"][0]["asset_class"] == "chiller" @pytest.mark.anyio async def test_full_relevancy_count(self, mock_relevancy_chain): - """2 sensors × 2 failure modes = 4 pairs.""" data = await call_tool( mcp, - "get_failure_mode_sensor_mapping", + "generate_failure_mode_sensor_mapping", { - "asset_name": "Chiller 6", + "asset_class": "Chiller", "failure_modes": _FAILURE_MODES, "sensors": _SENSORS, }, ) + assert len(data["full_relevancy"]) == 4 @pytest.mark.anyio async def test_empty_failure_modes_returns_error(self, mock_relevancy_chain): data = await call_tool( mcp, - "get_failure_mode_sensor_mapping", - {"asset_name": "Chiller 6", "failure_modes": [], "sensors": _SENSORS}, + "generate_failure_mode_sensor_mapping", + {"asset_class": "chiller", "failure_modes": [], "sensors": _SENSORS}, ) + assert "error" in data + @pytest.mark.anyio + async def test_empty_asset_class_returns_error(self, mock_relevancy_chain): + data = await call_tool( + mcp, + "generate_failure_mode_sensor_mapping", + { + "asset_class": "", + "failure_modes": _FAILURE_MODES, + "sensors": _SENSORS, + }, + ) + + assert data == {"error": "asset_class is required"} + @pytest.mark.anyio async def test_empty_sensors_returns_error(self, mock_relevancy_chain): data = await call_tool( mcp, - "get_failure_mode_sensor_mapping", - {"asset_name": "Chiller 6", "failure_modes": _FAILURE_MODES, "sensors": []}, + "generate_failure_mode_sensor_mapping", + {"asset_class": "chiller", "failure_modes": _FAILURE_MODES, "sensors": []}, ) + assert "error" in data @pytest.mark.anyio async def test_llm_unavailable_returns_error(self, no_llm): data = await call_tool( mcp, - "get_failure_mode_sensor_mapping", + "generate_failure_mode_sensor_mapping", { - "asset_name": "Chiller 6", + "asset_class": "chiller", "failure_modes": _FAILURE_MODES, "sensors": _SENSORS, }, ) - assert "error" in data + + assert data == {"error": "LLM unavailable"} @requires_watsonx @pytest.mark.anyio async def test_integration(self): data = await call_tool( mcp, - "get_failure_mode_sensor_mapping", + "generate_failure_mode_sensor_mapping", { - "asset_name": "Chiller 6", + "asset_class": "chiller", "failure_modes": ["Compressor Overheating"], "sensors": ["Chiller 6 Power Input"], }, ) + assert "full_relevancy" in data assert len(data["full_relevancy"]) == 1 assert data["full_relevancy"][0]["relevancy_answer"] in ("Yes", "No", "Unknown") diff --git a/src/servers/utilities/main.py b/src/servers/utilities/main.py index 42858783c..ae0ac2e25 100644 --- a/src/servers/utilities/main.py +++ b/src/servers/utilities/main.py @@ -1,15 +1,19 @@ import json import logging +import os import tempfile from datetime import datetime, timezone from pathlib import Path +from typing import Any, Optional, Union from uuid import uuid4 +import couchdb3 import pendulum +from dotenv import load_dotenv from mcp.server.fastmcp import FastMCP from pydantic import BaseModel -import os +load_dotenv() # Setup logging — default WARNING so stderr stays quiet when used as MCP server; # set LOG_LEVEL=INFO (or DEBUG) in the environment to see verbose output. @@ -19,9 +23,31 @@ logging.basicConfig(level=_log_level) logger = logging.getLogger("utilities-mcp-server") +# Configuration from environment +COUCHDB_URL = os.environ.get("COUCHDB_URL", "http://localhost:5984") +COUCHDB_USERNAME = os.environ.get("COUCHDB_USERNAME", "admin") +COUCHDB_PASSWORD = os.environ.get("COUCHDB_PASSWORD", "password") +CATALOG_DBNAME = os.environ.get("CATALOG_DBNAME", "catalog") +CATALOG_QUERY_LIMIT = 1000 + +try: + catalog_db = couchdb3.Database( + CATALOG_DBNAME, + url=COUCHDB_URL, + user=COUCHDB_USERNAME, + password=COUCHDB_PASSWORD, + ) + logger.info("Connected to catalog database: %s", CATALOG_DBNAME) +except Exception as e: + logger.error("Failed to connect to catalog database: %s", e) + catalog_db = None + mcp = FastMCP( "utilities", - instructions="General utilities: read JSON files and get current date/time.", + instructions=( + "General utilities: read JSON files, get current date/time, and query " + "asset, sensor, and failure-mode catalog data." + ), ) @@ -35,6 +61,18 @@ class TimeEnglishResult(BaseModel): iso: str +class ErrorResult(BaseModel): + error: str + + +class CatalogResult(BaseModel): + catalog_type: str + query: Optional[str] + total: int + entries: list[dict[str, Any]] + message: str + + # --- Helper Functions --- @@ -51,6 +89,57 @@ def get_temp_filename() -> str: return str(filepath) +def _clean_filter(value: Optional[str]) -> Optional[str]: + if value is None: + return None + value = value.strip() + return value or None + + +def _find_catalog( + *, + catalog_type: str, + field: str, + value: Optional[str], + fields: list[str], + category: Optional[str] = None, +) -> Union[CatalogResult, ErrorResult]: + if catalog_db is None: + return ErrorResult(error="catalog database is not available") + + query_value = _clean_filter(value) + category_value = _clean_filter(category) + selector: dict[str, Any] = {field: query_value or {"$exists": True}} + if category_value is not None: + selector["category"] = category_value + + try: + res = catalog_db.find( + selector, + fields=fields, + limit=CATALOG_QUERY_LIMIT, + ) + docs = res.get("docs", []) + except Exception as e: + logger.error("Error querying %s catalog: %s", catalog_type, e) + return ErrorResult(error=str(e)) + + query_parts = [] + if query_value is not None: + query_parts.append(query_value) + if category_value is not None: + query_parts.append(f"category={category_value}") + query = ", ".join(query_parts) if query_parts else None + + return CatalogResult( + catalog_type=catalog_type, + query=query, + total=len(docs), + entries=docs, + message=f"found {len(docs)} {catalog_type} catalog entries", + ) + + # --- JSON Tools --- @@ -66,6 +155,66 @@ def json_reader(file_name: str) -> str: return json.dumps({"error": str(e)}) +# --- Catalog Tools --- + + +@mcp.tool(title="Get Sensor Catalog") +def get_sensor_catalog( + sensor: Optional[str] = None, +) -> Union[CatalogResult, ErrorResult]: + """Return cataloged sensor types. + + Entries contain `sensor` and `description`. Omit sensor to list all cataloged + sensor types, or pass sensor for an exact sensor-name lookup. + """ + return _find_catalog( + catalog_type="sensor", + field="sensor", + value=sensor, + fields=["sensor", "description"], + ) + + +@mcp.tool(title="Get Asset Catalog") +def get_asset_catalog( + asset: Optional[str] = None, + category: Optional[str] = None, +) -> Union[CatalogResult, ErrorResult]: + """Return cataloged asset classes and categories. + + Entries contain `category`, `category_description`, `asset`, and + `description`. Omit filters to list all cataloged asset classes, or pass + asset and/or category for exact lookups. + """ + return _find_catalog( + catalog_type="asset", + field="asset", + value=asset, + category=category, + fields=["category", "category_description", "asset", "description"], + ) + + +@mcp.tool(title="Get Failure Mode Catalog") +def get_failure_mode_catalog( + failure_mode: Optional[str] = None, + category: Optional[str] = None, +) -> Union[CatalogResult, ErrorResult]: + """Return cataloged failure modes by asset category. + + Entries contain `category`, `failure_mode`, and `description`. Omit filters + to list all cataloged failure modes, or pass failure_mode and/or category + for exact lookups. + """ + return _find_catalog( + catalog_type="failure_mode", + field="failure_mode", + value=failure_mode, + category=category, + fields=["category", "failure_mode", "description"], + ) + + # --- Time Tools --- diff --git a/src/servers/utilities/tests/test_utilities.py b/src/servers/utilities/tests/test_utilities.py index 0940ea266..20e7c9d8d 100644 --- a/src/servers/utilities/tests/test_utilities.py +++ b/src/servers/utilities/tests/test_utilities.py @@ -5,10 +5,61 @@ import tempfile import pytest +from servers.utilities import main as utilities from servers.utilities.main import mcp from .conftest import call_tool +class FakeCatalogDB: + def __init__(self, docs): + self.docs = docs + self.calls = [] + + def find(self, selector, fields=None, limit=200): + self.calls.append({"selector": selector, "fields": fields, "limit": limit}) + matched = [] + for doc in self.docs: + if self._matches(doc, selector): + if fields: + matched.append( + {field: doc[field] for field in fields if field in doc} + ) + else: + matched.append(doc) + return {"docs": matched[:limit]} + + def _matches(self, doc, selector): + for field, expected in selector.items(): + if isinstance(expected, dict) and "$exists" in expected: + if (field in doc) != expected["$exists"]: + return False + elif doc.get(field) != expected: + return False + return True + + +@pytest.fixture +def fake_catalog_db(monkeypatch): + fake = FakeCatalogDB( + [ + {"sensor": "air flow", "description": "Airflow sensor"}, + { + "category": "rotating equipment", + "category_description": "Rotating equipment", + "asset": "Electric motor", + "description": "Converts electricity into motion.", + }, + { + "category": "rotating equipment", + "failure_mode": "Air inlet blockage", + "description": "Air intake is blocked.", + }, + ] + ) + monkeypatch.setattr(utilities, "catalog_db", fake) + return fake + + # --------------------------------------------------------------------------- # current_date_time # --------------------------------------------------------------------------- @@ -104,3 +155,64 @@ async def test_invalid_json_content(self): assert "error" in data finally: os.remove(tmp_name) + + +# --------------------------------------------------------------------------- +# catalog tools +# --------------------------------------------------------------------------- + + +class TestCatalogTools: + @pytest.mark.anyio + async def test_get_sensor_catalog_lists_sensor_entries(self, fake_catalog_db): + data = await call_tool(mcp, "get_sensor_catalog", {}) + + assert data["catalog_type"] == "sensor" + assert data["total"] == 1 + assert data["entries"] == [ + {"sensor": "air flow", "description": "Airflow sensor"} + ] + assert fake_catalog_db.calls[-1]["selector"] == { + "sensor": {"$exists": True} + } + assert fake_catalog_db.calls[-1]["limit"] == utilities.CATALOG_QUERY_LIMIT + + @pytest.mark.anyio + async def test_get_asset_catalog_filters_asset_and_category(self, fake_catalog_db): + data = await call_tool( + mcp, + "get_asset_catalog", + {"asset": "Electric motor", "category": "rotating equipment"}, + ) + + assert data["catalog_type"] == "asset" + assert data["query"] == "Electric motor, category=rotating equipment" + assert data["total"] == 1 + assert data["entries"][0]["asset"] == "Electric motor" + assert fake_catalog_db.calls[-1]["selector"] == { + "asset": "Electric motor", + "category": "rotating equipment", + } + + @pytest.mark.anyio + async def test_get_failure_mode_catalog_queries_failure_mode( + self, fake_catalog_db + ): + data = await call_tool( + mcp, + "get_failure_mode_catalog", + {"failure_mode": "Air inlet blockage"}, + ) + + assert data["catalog_type"] == "failure_mode" + assert data["total"] == 1 + assert data["entries"] == [ + { + "category": "rotating equipment", + "failure_mode": "Air inlet blockage", + "description": "Air intake is blocked.", + } + ] + assert fake_catalog_db.calls[-1]["selector"] == { + "failure_mode": "Air inlet blockage" + }