from mcp import types

similarity_tool_definitions = [
    types.Tool(
        name="node_similarity",
        description="The Node Similarity algorithm compares a set of nodes based on the nodes they are connected to. "
        "Two nodes are considered similar if they share many of the same neighbors. "
        "Node Similarity computes pair-wise similarities based on the Jaccard metric, also known as the Jaccard Similarity Score, the Overlap coefficient, also known as the Szymkiewicz–Simpson coefficient, and the Cosine Similarity score. "
        "The first two are most frequently associated with unweighted sets, whereas Cosine with weighted input.",
        inputSchema={
            "type": "object",
            "properties": {
                "similarityCutoff": {
                    "type": "number",
                    "description": "Lower limit for the similarity score to be present in the result. Values must be between 0 and 1.",
                },
                "degreeCutoff": {
                    "type": "integer",
                    "description": "Inclusive lower bound on the node degree for a node to be considered in the comparisons. This value can not be lower than 1.",
                },
                "upperDegreeCutoff": {
                    "type": "integer",
                    "description": "Inclusive upper bound on the node degree for a node to be considered in the comparisons. This value can not be lower than 1.",
                },
                "topK": {
                    "type": "integer",
                    "description": "Limit on the number of scores per node. The K largest results are returned. This value cannot be lower than 1.",
                },
                "bottomK": {
                    "type": "integer",
                    "description": "Limit on the number of scores per node. The K smallest results are returned. This value cannot be lower than 1.",
                },
                "topN": {
                    "type": "integer",
                    "description": "Global limit on the number of scores computed. The N largest total results are returned. This value cannot be negative, a value of 0 means no global limit.",
                },
                "bottomN": {
                    "type": "integer",
                    "description": "Global limit on the number of scores computed. The N smallest total results are returned. This value cannot be negative, a value of 0 means no global limit.",
                },
                "relationshipWeightProperty": {
                    "type": "string",
                    "description": "Name of the relationship property to use as weights. If unspecified, the algorithm runs unweighted.",
                },
                "similarityMetric": {
                    "type": "string",
                    "enum": ["JACCARD", "OVERLAP", "COSINE"],
                    "description": "The metric used to compute similarity.",
                },
                "useComponents": {
                    "type": "boolean",
                    "description": "If enabled, Node Similarity will use components to improve the performance of the computation, skipping comparisons of nodes in different components. Set to false (Default): the algorithm does not use components, but computes similarity across the entire graph. Set to true: the algorithm uses components, and will compute these components before computing similarity. Set to String: use pre-computed components stored in graph, String is the key for a node property representing components.",
                },
            },
        },
    ),
    types.Tool(
        name="filtered_node_similarity",
        description="The Filtered Node Similarity algorithm is an extension to the Node Similarity algorithm. "
        "It adds support for filtering on source nodes, target nodes, or both.",
        inputSchema={
            "type": "object",
            "properties": {
                "sourceNodeFilter": {
                    "type": ["integer", "array", "string"],
                    "description": "The source node filter to apply. Accepts a single node id, a List of node ids, or a single label.",
                },
                "targetNodeFilter": {
                    "type": ["integer", "array", "string"],
                    "description": "The target node filter to apply. Accepts a single node id, a List of node ids, or a single label.",
                },
                "similarityCutoff": {
                    "type": "number",
                    "description": "Lower limit for the similarity score to be present in the result. Values must be between 0 and 1.",
                },
                "degreeCutoff": {
                    "type": "integer",
                    "description": "Inclusive lower bound on the node degree for a node to be considered in the comparisons. This value can not be lower than 1.",
                },
                "upperDegreeCutoff": {
                    "type": "integer",
                    "description": "Inclusive upper bound on the node degree for a node to be considered in the comparisons. This value can not be lower than 1.",
                },
                "topK": {
                    "type": "integer",
                    "description": "Limit on the number of scores per node. The K largest results are returned. This value cannot be lower than 1.",
                },
                "bottomK": {
                    "type": "integer",
                    "description": "Limit on the number of scores per node. The K smallest results are returned. This value cannot be lower than 1.",
                },
                "topN": {
                    "type": "integer",
                    "description": "Global limit on the number of scores computed. The N largest total results are returned. This value cannot be negative, a value of 0 means no global limit.",
                },
                "bottomN": {
                    "type": "integer",
                    "description": "Global limit on the number of scores computed. The N smallest total results are returned. This value cannot be negative, a value of 0 means no global limit.",
                },
                "relationshipWeightProperty": {
                    "type": "string",
                    "description": "Name of the relationship property to use as weights. If unspecified, the algorithm runs unweighted.",
                },
                "similarityMetric": {
                    "type": "string",
                    "enum": ["JACCARD", "OVERLAP", "COSINE"],
                    "description": "The metric used to compute similarity.",
                },
                "useComponents": {
                    "type": "boolean",
                    "description": "If enabled, Node Similarity will use components to improve the performance of the computation, skipping comparisons of nodes in different components. Set to false (Default): the algorithm does not use components, but computes similarity across the entire graph. Set to true: the algorithm uses components, and will compute these components before computing similarity. Set to String: use pre-computed components stored in graph, String is the key for a node property representing components.",
                },
            },
            "required": ["sourceNodeFilter", "targetNodeFilter"],
        },
    ),
    types.Tool(
        name="k_nearest_neighbors",
        description="The K-Nearest Neighbors algorithm computes a distance value for all node pairs in the graph and creates new relationships between each node and its k nearest neighbors. "
        "The distance is calculated based on node properties. "
        "The input of this algorithm is a homogeneous graph; any node label or relationships type information in the graph is ignored. "
        "The graph does not need to be connected, in fact, existing relationships between nodes will be ignored - apart from random walk sampling if that that initial sampling option is used. "
        "New relationships are created between each node and its k nearest neighbors. "
        "The K-Nearest Neighbors algorithm compares given properties of each node. The k nodes where these properties are most similar are the k-nearest neighbors. "
        "The initial set of neighbors is picked at random and verified and refined in multiple iterations. The number of iterations is limited by the configuration parameter maxIterations. "
        "The algorithm may stop earlier if the neighbor lists only change by a small amount, which can be controlled by the configuration parameter deltaThreshold. "
        "The particular implementation is based on Efficient k-nearest neighbor graph construction for generic similarity measures by Wei Dong et al. "
        "Instead of comparing every node with every other node, the algorithm selects possible neighbors based on the assumption, that the neighbors-of-neighbors of a node are most likely already the nearest one. "
        "The algorithm scales quasi-linear with respect to the node count, instead of being quadratic. "
        "Furthermore, the algorithm only compares a sample of all possible neighbors on each iteration, assuming that eventually all possible neighbors will be seen. "
        "This can be controlled with the configuration parameter sampleRate: A valid sample rate must be in between 0 (exclusive) and 1 (inclusive). "
        "The default value is 0.5. The parameter is used to control the trade-off between accuracy and runtime-performance. "
        "A higher sample rate will increase the accuracy of the result. The algorithm will also require more memory and will take longer to compute. "
        "A lower sample rate will increase the runtime-performance. Some potential nodes may be missed in the comparison and may not be included in the result. "
        "When encountered neighbors have equal similarity to the least similar already known neighbor, randomly selecting which node to keep can reduce the risk of some neighborhoods not being explored. "
        "This behavior is controlled by the configuration parameter perturbationRate. "
        "The output of the algorithm are new relationships between nodes and their k-nearest neighbors. Similarity scores are expressed via relationship properties.",
        inputSchema={
            "type": "object",
            "properties": {
                "nodeProperties": {
                    "type": ["string", "object", "array"],
                    "description": "The node properties to use for similarity computation along with their selected similarity metrics. Accepts a single property key, a Map of property keys to metrics, or a List of property keys and/or Maps, as above.",
                },
                "topK": {
                    "type": "integer",
                    "description": "The number of neighbors to find for each node. The K-nearest neighbors are returned. This value cannot be lower than 1.",
                },
                "sampleRate": {
                    "type": "number",
                    "description": "Sample rate to limit the number of comparisons per node. Value must be between 0 (exclusive) and 1 (inclusive).",
                },
                "deltaThreshold": {
                    "type": "number",
                    "description": "Value as a percentage to determine when to stop early. If fewer updates than the configured value happen, the algorithm stops. Value must be between 0 (exclusive) and 1 (inclusive).",
                },
                "maxIterations": {
                    "type": "integer",
                    "description": "Hard limit to stop the algorithm after that many iterations.",
                },
                "randomJoins": {
                    "type": "integer",
                    "description": "The number of random attempts per node to connect new node neighbors based on random selection, for each iteration.",
                },
                "initialSampler": {
                    "type": "string",
                    "enum": ["uniform", "randomWalk"],
                    "description": 'The method used to sample the first k random neighbors for each node. "uniform" and "randomWalk", both case-insensitive, are valid inputs.',
                },
                "similarityCutoff": {
                    "type": "number",
                    "description": "Filter out from the list of K-nearest neighbors nodes with similarity below this threshold.",
                },
                "perturbationRate": {
                    "type": "number",
                    "description": "The probability of replacing the least similar known neighbor with an encountered neighbor of equal similarity.",
                },
            },
            "required": ["nodeProperties"],
        },
    ),
    types.Tool(
        name="filtered_k_nearest_neighbors",
        description="The Filtered K-Nearest Neighbors algorithm extends our popular K-Nearest Neighbors algorithm with filtering on source nodes, target nodes or both.",
        inputSchema={
            "type": "object",
            "properties": {
                "sourceNodeFilter": {
                    "type": ["integer", "array", "string"],
                    "description": "The source node filter to apply. Accepts a single node id, a List of node ids, or a single label.",
                },
                "targetNodeFilter": {
                    "type": ["integer", "array", "string"],
                    "description": "The target node filter to apply. Accepts a single node id, a List of node ids, or a single label.",
                },
                "nodeProperties": {
                    "type": ["string", "object", "array"],
                    "description": "The node properties to use for similarity computation along with their selected similarity metrics. Accepts a single property key, a Map of property keys to metrics, or a List of property keys and/or Maps, as above.",
                },
                "topK": {
                    "type": "integer",
                    "description": "The number of neighbors to find for each node. The K-nearest neighbors are returned. This value cannot be lower than 1.",
                },
                "sampleRate": {
                    "type": "number",
                    "description": "Sample rate to limit the number of comparisons per node. Value must be between 0 (exclusive) and 1 (inclusive).",
                },
                "deltaThreshold": {
                    "type": "number",
                    "description": "Value as a percentage to determine when to stop early. If fewer updates than the configured value happen, the algorithm stops. Value must be between 0 (exclusive) and 1 (inclusive).",
                },
                "maxIterations": {
                    "type": "integer",
                    "description": "Hard limit to stop the algorithm after that many iterations.",
                },
                "randomJoins": {
                    "type": "integer",
                    "description": "The number of random attempts per node to connect new node neighbors based on random selection, for each iteration.",
                },
                "initialSampler": {
                    "type": "string",
                    "enum": ["uniform", "randomWalk"],
                    "description": 'The method used to sample the first k random neighbors for each node. "uniform" and "randomWalk", both case-insensitive, are valid inputs.',
                },
                "similarityCutoff": {
                    "type": "number",
                    "description": "Filter out from the list of K-nearest neighbors nodes with similarity below this threshold.",
                },
                "perturbationRate": {
                    "type": "number",
                    "description": "The probability of replacing the least similar known neighbor with an encountered neighbor of equal similarity.",
                },
                "seedTargetNodes": {
                    "type": "boolean",
                    "description": "Enable seeding of target nodes.",
                },
            },
            "required": ["sourceNodeFilter", "targetNodeFilter", "nodeProperties"],
        },
    ),
]
