diff --git a/src/couchdb/collections.json b/src/couchdb/collections.json index e99f367db..e03ab8613 100644 --- a/src/couchdb/collections.json +++ b/src/couchdb/collections.json @@ -24,6 +24,23 @@ ] ] }, + "catalog": { + "format": "csv", + "indexes": [ + [ + "category" + ], + [ + "asset" + ], + [ + "failure_mode" + ], + [ + "sensor" + ] + ] + }, "iot": { "format": "json", "primary_key": [ @@ -111,4 +128,4 @@ "id_prefix": "fc", "indexes": [] } -} \ No newline at end of file +} diff --git a/src/servers/utilities/main.py b/src/servers/utilities/main.py index 42858783c..dd940f7cf 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 CouchDB catalog database: %s", CATALOG_DBNAME) +except Exception as e: + logger.error("Failed to connect to CouchDB 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 from CouchDB." + ), ) @@ -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,55 @@ 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 sensor catalog entries. 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 asset catalog entries. 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 Model Catalog") +def get_failure_model_catalog( + failure_mode: Optional[str] = None, + category: Optional[str] = None, +) -> Union[CatalogResult, ErrorResult]: + """Return failure-mode catalog entries. + + 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..c60e0ff91 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_model_catalog_queries_failure_mode( + self, fake_catalog_db + ): + data = await call_tool( + mcp, + "get_failure_model_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" + }