Refactor functions to use Pydantic models for input validation and enhance the application with additional NumPy-based mathematical tools.

app.py

@@ -1,26 +1,214 @@
 from fastmcp import FastMCP
 import numpy as np
+from pydantic import BaseModel
+from typing import List, Tuple, Optional
 
 mcp = FastMCP("Demo 🚀")
 
+class HelloInput(BaseModel):
+    name: str
+
 @mcp.tool()
-def hello(name: str) -> str:
-    return f"Hello, {name}!"
+def hello(input: HelloInput) -> str:
+    return f"Hello, {input.name}!"
+
+class MultiplyInput(BaseModel):
+    a: float
+    b: float
 
 @mcp.tool()
-def multiply(a: float, b: float) -> float:
+def multiply(input: MultiplyInput) -> float:
     """Multiplies two numbers."""
-    return a * b
+    return input.a * input.b
+
+class InnerProductInput(BaseModel):
+    a: List[float]
+    b: List[float]
 
 @mcp.tool()
-def inner_product(a: list[float], b: list[float]) -> float:
+def inner_product(input: InnerProductInput) -> float:
     """Calculates the inner product of two vectors."""
-    return np.dot(a, b)
+    return np.dot(input.a, input.b)
+
+class MatrixMultiplyInput(BaseModel):
+    a: List[List[float]]
+    b: List[List[float]]
 
 @mcp.tool()
-def matrix_multiply(a: list[list[float]], b: list[list[float]]) -> list[list[float]]:
+def matrix_multiply(input: MatrixMultiplyInput) -> List[List[float]]:
     """Multiplies two matrices."""
-    return np.matmul(a, b)
+    return np.matmul(input.a, input.b)
+
+class NumpyDotInput(BaseModel):
+    a: List[float]
+    b: List[float]
+
+@mcp.tool()
+def numpy_dot(input: NumpyDotInput) -> float:
+    """Calculates the dot product of two vectors."""
+    return np.dot(input.a, input.b)
+
+class NumpyMatmulInput(BaseModel):
+    a: List[List[float]]
+    b: List[List[float]]
+
+@mcp.tool()
+def numpy_matmul(input: NumpyMatmulInput) -> List[List[float]]:
+    """Multiplies two matrices using matmul."""
+    return np.matmul(input.a, input.b)
+
+class NumpyInvInput(BaseModel):
+    a: List[List[float]]
+
+@mcp.tool()
+def numpy_inv(input: NumpyInvInput) -> List[List[float]]:
+    """Calculates the inverse of a matrix."""
+    return np.linalg.inv(input.a)
+
+class NumpyDetInput(BaseModel):
+    a: List[List[float]]
+
+@mcp.tool()
+def numpy_det(input: NumpyDetInput) -> float:
+    """Calculates the determinant of a matrix."""
+    return np.linalg.det(input.a)
+
+class NumpyEigInput(BaseModel):
+    a: List[List[float]]
+
+@mcp.tool()
+def numpy_eig(input: NumpyEigInput) -> Tuple:
+    """Calculates the eigenvalues and eigenvectors of a matrix."""
+    return np.linalg.eig(input.a)
+
+class NumpySvdInput(BaseModel):
+    a: List[List[float]]
+
+@mcp.tool()
+def numpy_svd(input: NumpySvdInput) -> Tuple:
+    """Performs singular value decomposition on a matrix."""
+    return np.linalg.svd(input.a)
+
+class NumpyNormInput(BaseModel):
+    a: List[float]
+    ord: Optional[int] = None
+
+@mcp.tool()
+def numpy_norm(input: NumpyNormInput) -> float:
+    """Calculates the norm of a vector or matrix."""
+    return np.linalg.norm(input.a, input.ord)
+
+class NumpyCrossInput(BaseModel):
+    a: List[float]
+    b: List[float]
+
+@mcp.tool()
+def numpy_cross(input: NumpyCrossInput) -> List[float]:
+    """Calculates the cross product of two vectors."""
+    return np.cross(input.a, input.b)
+
+class NumpyInnerInput(BaseModel):
+    a: List[float]
+    b: List[float]
+
+@mcp.tool()
+def numpy_inner(input: NumpyInnerInput) -> float:
+    """Calculates the inner product of two vectors."""
+    return np.inner(input.a, input.b)
+
+class NumpyOuterInput(BaseModel):
+    a: List[float]
+    b: List[float]
+
+@mcp.tool()
+def numpy_outer(input: NumpyOuterInput) -> List[List[float]]:
+    """Calculates the outer product of two vectors."""
+    return np.outer(input.a, input.b)
+
+class NumpyTensordotInput(BaseModel):
+    a: List
+    b: List
+    axes: int = 2
+
+@mcp.tool()
+def numpy_tensordot(input: NumpyTensordotInput) -> float:
+    """Calculates the tensor dot product of two arrays."""
+    return np.tensordot(input.a, input.b, input.axes)
+
+class NumpyTraceInput(BaseModel):
+    a: List[List[float]]
+
+@mcp.tool()
+def numpy_trace(input: NumpyTraceInput) -> float:
+    """Calculates the trace of a matrix."""
+    return np.trace(input.a)
+
+class NumpyQrInput(BaseModel):
+    a: List[List[float]]
+
+@mcp.tool()
+def numpy_qr(input: NumpyQrInput) -> Tuple:
+    """Performs QR decomposition on a matrix."""
+    return np.linalg.qr(input.a)
+
+class NumpyCholeskyInput(BaseModel):
+    a: List[List[float]]
+
+@mcp.tool()
+def numpy_cholesky(input: NumpyCholeskyInput) -> List[List[float]]:
+    """Performs Cholesky decomposition on a matrix."""
+    return np.linalg.cholesky(input.a)
+
+class NumpySolveInput(BaseModel):
+    a: List[List[float]]
+    b: List[float]
+
+@mcp.tool()
+def numpy_solve(input: NumpySolveInput) -> List[float]:
+    """Solves a linear matrix equation."""
+    return np.linalg.solve(input.a, input.b)
+
+class NumpyLstsqInput(BaseModel):
+    a: List[List[float]]
+    b: List[float]
+
+@mcp.tool()
+def numpy_lstsq(input: NumpyLstsqInput) -> Tuple:
+    """Solves a linear least squares problem."""
+    return np.linalg.lstsq(input.a, input.b, rcond=None)
+
+class NumpyPinvInput(BaseModel):
+    a: List[List[float]]
+
+@mcp.tool()
+def numpy_pinv(input: NumpyPinvInput) -> List[List[float]]:
+    """Calculates the Moore-Penrose pseudo-inverse of a matrix."""
+    return np.linalg.pinv(input.a)
+
+class NumpyCondInput(BaseModel):
+    a: List[List[float]]
+    p: Optional[int] = None
+
+@mcp.tool()
+def numpy_cond(input: NumpyCondInput) -> float:
+    """Calculates the condition number of a matrix."""
+    return np.linalg.cond(input.a, input.p)
+
+class NumpyMatrixRankInput(BaseModel):
+    a: List[List[float]]
+
+@mcp.tool()
+def numpy_matrix_rank(input: NumpyMatrixRankInput) -> int:
+    """Calculates the rank of a matrix."""
+    return np.linalg.matrix_rank(input.a)
+
+class NumpyMultiDotInput(BaseModel):
+    arrays: List[List[List[float]]]
+
+@mcp.tool()
+def numpy_multi_dot(input: NumpyMultiDotInput) -> List[List[float]]:
+    """Computes the dot product of two or more arrays in a single function call, while automatically selecting the fastest evaluation order."""
+    return np.linalg.multi_dot(input.arrays)
 
 # Static resource
 @mcp.resource("config://version")
@@ -28,15 +216,21 @@ def get_version():
     return "2.0.1"
 
 # Dynamic resource template
+class GetProfileInput(BaseModel):
+    user_id: int
+
 @mcp.resource("users://{user_id}/profile")
-def get_profile(user_id: int):
+def get_profile(input: GetProfileInput):
     # Fetch profile for user_id...
-    return {"name": f"User {user_id}", "status": "active"}
+    return {"name": f"User {input.user_id}", "status": "active"}
+
+class SummarizeRequestInput(BaseModel):
+    text: str
 
 @mcp.prompt()
-def summarize_request(text: str) -> str:
+def summarize_request(input: SummarizeRequestInput) -> str:
     """Generate a prompt asking for a summary."""
-    return f"Please summarize the following text:\n\n{text}"
+    return f"Please summarize the following text:\n\n{input.text}"
 
 if __name__ == "__main__":
     mcp.run(transport="sse", host="0.0.0.0", port=7860)