server.py

from mcp.server.fastmcp import FastMCP
from tools.validate_math import validate_math
from tools.check_consistency import check_consistency
from auth.jwt_middleware import extract_tenant_id
from logger.failure_logger import log_failure
from models.verification_result import ValidationResult

mcp = FastMCP("lemma")


@mcp.tool()
def verify_output(
    output: str,
    spec: str = "",
    context: str = "",
    token: str = ""
) -> dict:
    """
    Master verification tool. Runs all validators on LLM output.
    Accepts raw LLM output + optional spec and context strings.
    Returns aggregated validation result with violations.
    """
    tenant_id = extract_tenant_id(token) if token else "anonymous"

    math_result = validate_math(output)
    consistency_result = check_consistency(output, context)

    all_violations = math_result.violations + consistency_result.violations
    is_valid = math_result.valid and consistency_result.valid

    scores = {
        "math_validity": 1.0 if math_result.valid else 0.0,
        "consistency": 1.0 if consistency_result.valid else 0.5,
    }

    if not is_valid:
        log_failure(
            tenant_id=tenant_id,
            raw_output=output,
            context=context,
            violations=all_violations,
            scores=scores
        )

    return {
        "verified": is_valid,
        "violations": all_violations,
        "scores": scores,
        "tenant_id": tenant_id
    }


@mcp.tool()
def validate_math_tool(output: str) -> dict:
    """
    Validate mathematical expressions in LLM output using SymPy.
    Checks arithmetic consistency and flags incorrect calculations.
    """
    result = validate_math(output)
    return result.model_dump()


@mcp.tool()
def check_consistency_tool(output: str, context: str = "") -> dict:
    """
    Check logical consistency of LLM output against provided context.
    Detects contradictions between output claims and prior context.
    """
    result = check_consistency(output, context)
    return result.model_dump()


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