VI MCP Server

vi_mcp_server.py

import httpx
import os
from contextlib import asynccontextmanager
from fastmcp import FastMCP
from fastmcp.server.auth import OAuthProxy
from fastmcp.server.auth.providers.jwt import JWTVerifier
from fastmcp.server.dependencies import get_access_token

endpoint = os.getenv("VIYA_ENDPOINT", "https://example.com")

JWKS_URI = f"{endpoint}/SASLogon/token_keys"
VI_BASE_URL = f"{endpoint}/SASVisualInvestigator"
SVI_BASE_URL = f"{endpoint}/svi-sand"
SVI_DATAHUB_URL = f"{endpoint}/svi-datahub"


token_verifier = JWTVerifier(jwks_uri=JWKS_URI, audience=[])

viya_auth = OAuthProxy(
    upstream_authorization_endpoint=f"{endpoint}/SASLogon/oauth/authorize",
    upstream_token_endpoint=f"{endpoint}/SASLogon/oauth/token",
    upstream_client_id="vi-mcp",
    upstream_client_secret="mackle",
    jwt_signing_key="default",
    base_url="http://localhost:8000",
    forward_pkce=True,
    token_verifier=token_verifier,
    valid_scopes=["openid", "uaa.user"],
)

mcp = FastMCP("My MCP Server", auth=viya_auth)

@asynccontextmanager
async def authenticated_client():
    token = get_access_token().token
    async with httpx.AsyncClient(headers={"Authorization": f"Bearer {token}"}) as client:
        yield client

@mcp.tool
async def get_vi_page(page_type: str, entity_id: str) -> dict:
    """Fetch a SAS Visual Investigator page by page type and entity ID."""
    async with authenticated_client() as client:
        response = await client.get(f"{VI_BASE_URL}/pages/{page_type}/{entity_id}")
        response.raise_for_status()
        return response.json()


@mcp.tool
async def get_entity_metadata(entity_type: str) -> dict:
    """Fetch metadata for a SAS Visual Investigator entity type."""
    async with authenticated_client() as client:
        response = await client.get(f"{SVI_DATAHUB_URL}/metadata/documents/{entity_type}")
        response.raise_for_status()
        return response.json()


@mcp.tool
async def get_entity_fields(entity_type: str) -> list[dict]:
    """Return the name and dataType of each field for a SAS Visual Investigator entity type."""
    async with authenticated_client() as client:
        response = await client.get(f"{SVI_DATAHUB_URL}/metadata/documents/{entity_type}")
        response.raise_for_status()
        data = response.json()
        return [{"name": f["name"], "dataType": f["dataType"]} for f in data.get("fields", [])]


@mcp.tool
async def get_searchable_types() -> list[str]:
    """Return a list of all searchable entity type names in SAS Visual Investigator."""
    async with authenticated_client() as client:
        response = await client.get(f"{SVI_BASE_URL}/admin/config/types", params={"start": 0, "limit": 200})
        response.raise_for_status()
        data = response.json()
        return [item["name"] for item in data.get("items", [])]


@mcp.tool
async def get_reference_list(name: str) -> dict:
    """Retrieve a reference list's definition from SAS Visual Investigator by name."""
    async with authenticated_client() as client:
        response = await client.get(f"{SVI_DATAHUB_URL}/lists", params={"name": name})
        response.raise_for_status()
        return response.json()


@mcp.tool
async def create_document(
    object_type_name: str,
    object_type_id: int,
    field_values: dict,
) -> dict:
    """Create a new document record in SAS Visual Investigator.

    Args:
        object_type_name: The entity type name (e.g. 'lab_test').
        object_type_id: The numeric ID of the entity type.
        field_values: A dict of field name to value pairs to set on the record.
    """
    payload = {
        "objectTypeName": object_type_name,
        "objectTypeId": object_type_id,
        "fieldValues": field_values,
        "comments": [],
    }
    async with authenticated_client() as client:
        response = await client.post(f"{SVI_DATAHUB_URL}/documents", json=payload)
        response.raise_for_status()
        return response.json()



@mcp.tool
async def get_link_relationship_types() -> dict:
    """Return all link relationship types defined in SAS Visual Investigator."""
    async with authenticated_client() as client:
        response = await client.get(f"{SVI_DATAHUB_URL}/admin/relationships", params={"type": "LINK"})
        response.raise_for_status()
        return response.json()


@mcp.tool
async def create_link(
    relationship_type_name: str,
    from_object_type_name: str,
    from_object_id: str,
    to_object_type_name: str,
    to_object_id: str,
    field_values: dict | None = None,
) -> dict:
    """Create a link (relationship) between two existing records in SAS Visual Investigator.

    Args:
        relationship_type_name: The relationship type name (e.g. 'lab_test').
        from_object_type_name: Entity type of the source record (e.g. 'disease_event').
        from_object_id: ID of the source record (e.g. 'DIS-2026-29').
        to_object_type_name: Entity type of the target record (e.g. 'lab_test').
        to_object_id: ID of the target record (e.g. 'LBT-2026-36').
        field_values: Optional dict of field values to set on the link itself.
    """
    payload = {
        "@type": "DocumentLink",
        "relationshipTypeName": relationship_type_name,
        "fieldValues": field_values or {},
        "fromObjectTypeName": from_object_type_name,
        "fromObjectId": from_object_id,
        "toObjectTypeName": to_object_type_name,
        "toObjectId": to_object_id,
    }
    async with authenticated_client() as client:
        response = await client.post(f"{SVI_DATAHUB_URL}/links", json=payload)
        response.raise_for_status()
        return response.json()


@mcp.tool
async def delete_document(object_type_name: str, document_id: str) -> dict:
    """Lock then delete a document record from SAS Visual Investigator.

    Args:
        object_type_name: The entity type name (e.g. 'disease_event').
        document_id: The record ID to delete (e.g. 'DIS-2026-28').
    """
    async with authenticated_client() as client:
        lock = await client.post(
            f"{SVI_DATAHUB_URL}/locks/documents",
            params={"type": object_type_name, "id": document_id},
        )
        lock.raise_for_status()
        response = await client.delete(f"{SVI_DATAHUB_URL}/documents/{object_type_name}/{document_id}")
        response.raise_for_status()
        return {"deleted": True, "object_type_name": object_type_name, "document_id": document_id}


@mcp.tool
async def search_records(lucene_query: str, start: int = 1, limit: int = 40) -> dict:
    """Search SAS Visual Investigator records using a Lucene query string.

    Returns summary results and facets for the matching records.
    """
    payload = {
        "query": {
            "type": "text",
            "language": "lucene",
            "text": lucene_query,
        },
        "visualizations": {
            "summary": {"type": "summary", "start": start, "limit": limit},
            "facets": {"type": "facets"},
            "selection": {"type": "facets", "mode": "selection"},
        },
    }
    async with authenticated_client() as client:
        response = await client.post(f"{SVI_BASE_URL}/searches", json=payload)
        response.raise_for_status()
        return response.json()


if __name__ == "__main__":
    mcp.run()