"""
MCP server for aare-core Z3 SMT verification.
Provides tools for:
- verify: Verify LLM output against compliance ontology
- list_ontologies: List available compliance ontologies
- parse_claims: Extract structured data from LLM output
"""
import json
import logging
from typing import Any
from mcp.server import Server
from mcp.server.stdio import stdio_server
from mcp.types import Tool, TextContent
from aare_core import SMTVerifier, LLMParser, OntologyLoader
logger = logging.getLogger(__name__)
# Initialize aare-core components
loader = OntologyLoader()
parser = LLMParser()
verifier = SMTVerifier()
# Create MCP server
server = Server("aare-verify")
@server.list_tools()
async def list_tools() -> list[Tool]:
"""Return list of available tools."""
return [
Tool(
name="verify",
description="Verify LLM output against a compliance ontology using Z3 SMT solver. "
"Returns verification result with any constraint violations and formal proof.",
inputSchema={
"type": "object",
"properties": {
"llm_output": {
"type": "string",
"description": "The LLM-generated text to verify for compliance"
},
"ontology": {
"type": "string",
"description": "Name of the ontology to verify against (e.g., 'mortgage-compliance-v1', 'hipaa-v1')"
}
},
"required": ["llm_output", "ontology"]
}
),
Tool(
name="list_ontologies",
description="List all available compliance ontologies that can be used for verification.",
inputSchema={
"type": "object",
"properties": {}
}
),
Tool(
name="parse_claims",
description="Extract structured data from LLM output using ontology-defined extractors. "
"Useful for debugging or understanding what data will be verified.",
inputSchema={
"type": "object",
"properties": {
"llm_output": {
"type": "string",
"description": "The LLM-generated text to parse"
},
"ontology": {
"type": "string",
"description": "Name of the ontology containing extractor definitions"
}
},
"required": ["llm_output", "ontology"]
}
)
]
@server.call_tool()
async def call_tool(name: str, arguments: dict[str, Any]) -> list[TextContent]:
"""Handle tool calls."""
if name == "verify":
return await _handle_verify(arguments)
elif name == "list_ontologies":
return await _handle_list_ontologies()
elif name == "parse_claims":
return await _handle_parse_claims(arguments)
else:
return [TextContent(
type="text",
text=json.dumps({"error": f"Unknown tool: {name}"}, indent=2)
)]
async def _handle_verify(arguments: dict[str, Any]) -> list[TextContent]:
"""Verify LLM output against ontology."""
try:
llm_output = arguments["llm_output"]
ontology_name = arguments["ontology"]
# Load ontology
ontology = loader.load(ontology_name)
# Parse LLM output to extract structured data
data = parser.parse(llm_output, ontology)
# Verify against constraints
result = verifier.verify(data, ontology)
# Build response
response = {
"verified": result["verified"],
"ontology": ontology_name,
"violations": result["violations"],
"parsed_data": data,
"proof": result["proof"],
"execution_time_ms": result["execution_time_ms"]
}
if result.get("warnings"):
response["warnings"] = result["warnings"]
return [TextContent(
type="text",
text=json.dumps(response, indent=2, default=str)
)]
except Exception as e:
logger.exception("Verification failed")
return [TextContent(
type="text",
text=json.dumps({
"error": str(e),
"type": type(e).__name__
}, indent=2)
)]
async def _handle_list_ontologies() -> list[TextContent]:
"""List available ontologies."""
try:
ontologies = loader.list_available()
return [TextContent(
type="text",
text=json.dumps({
"ontologies": ontologies,
"count": len(ontologies)
}, indent=2)
)]
except Exception as e:
logger.exception("Failed to list ontologies")
return [TextContent(
type="text",
text=json.dumps({
"error": str(e),
"type": type(e).__name__
}, indent=2)
)]
async def _handle_parse_claims(arguments: dict[str, Any]) -> list[TextContent]:
"""Parse claims from LLM output."""
try:
llm_output = arguments["llm_output"]
ontology_name = arguments["ontology"]
# Load ontology
ontology = loader.load(ontology_name)
# Parse LLM output
data = parser.parse(llm_output, ontology, include_confidence=True)
# Extract confidence scores and clean up response
confidence_scores = data.pop("_confidence_scores", {})
# Convert ExtractionResult objects to serializable format
parsed_data = {}
extraction_details = {}
for key, value in data.items():
if hasattr(value, 'value'): # ExtractionResult
parsed_data[key] = value.value
extraction_details[key] = {
"value": value.value,
"confidence": value.confidence,
"source": value.source,
"extractor_type": value.extractor_type
}
else:
parsed_data[key] = value
return [TextContent(
type="text",
text=json.dumps({
"ontology": ontology_name,
"parsed_data": parsed_data,
"extraction_details": extraction_details
}, indent=2, default=str)
)]
except Exception as e:
logger.exception("Parsing failed")
return [TextContent(
type="text",
text=json.dumps({
"error": str(e),
"type": type(e).__name__
}, indent=2)
)]
async def run_server():
"""Run the MCP server."""
async with stdio_server() as (read_stream, write_stream):
await server.run(
read_stream,
write_stream,
server.create_initialization_options()
)
def main():
"""Entry point for the MCP server."""
import asyncio
asyncio.run(run_server())
if __name__ == "__main__":
main()
